Add IS_PYTORCH_CI flag for testing (#16006)

Summary:
Use case:
Some data loader tests rely on `psutil` (a third party lib). So they are guarded by `skipIf`. But we want to always test them on CI envs. With `IS_PYTORCH_CI`, we can raise if `psutil` is not found.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/16006

Reviewed By: ezyang

Differential Revision: D13673957

Pulled By: yf225

fbshipit-source-id: c63a7138093f45333c0b371fed0bcc88b67f2a22

diff --git a/.jenkins/pytorch/common.sh b/.jenkins/pytorch/common.sh
index ca728df..357c637 100644 (file)
--- a/.jenkins/pytorch/common.sh
+++ b/.jenkins/pytorch/common.sh
@@ -25,6 +25,8 @@ set -ex
 # system; to find out more, grep for this string in ossci-job-dsl.
 echo "ENTERED_USER_LAND"
 
+export IS_PYTORCH_CI=1
+
 # compositional trap taken from https://stackoverflow.com/a/7287873/23845
 
 # note: printf is used instead of echo to avoid backslash

diff --git a/test/common_utils.py b/test/common_utils.py
index 940e55d..e6aa375 100644 (file)
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -63,6 +63,9 @@ PY34 = sys.version_info >= (3, 4)
 IS_WINDOWS = sys.platform == "win32"
 IS_PPC = platform.machine() == "ppc64le"
 
+# Environment variable `IS_PYTORCH_CI` is set in `.jenkins/common.sh`.
+IS_PYTORCH_CI = bool(os.environ.get('IS_PYTORCH_CI', 0))
+
 
 def _check_module_exists(name):
     r"""Returns if a top-level module with :attr:`name` exists *without**

diff --git a/test/run_test.py b/test/run_test.py
index cd11273..93263ee 100644 (file)
--- a/test/run_test.py
+++ b/test/run_test.py
@@ -340,9 +340,9 @@ def get_executable_command(options):
 
 
 def find_test_index(test, selected_tests, find_last_index=False):
-    """Find the index of the first or last occurrence of a given test/test module in the list of seleceted tests.
+    """Find the index of the first or last occurrence of a given test/test module in the list of selected tests.
 
-    This function is used to determine the indexes when slicing the list of selected tests when
+    This function is used to determine the indices when slicing the list of selected tests when
     ``options.first``(:attr:`find_last_index`=False) and/or ``options.last``(:attr:`find_last_index`=True) are used.
 
     :attr:`selected_tests` can be a list that contains multiple consequent occurrences of tests

diff --git a/test/test_dataloader.py b/test/test_dataloader.py
index 6d2701d..d89a1ac 100644 (file)
--- a/test/test_dataloader.py
+++ b/test/test_dataloader.py
@@ -16,17 +16,21 @@ from torch import multiprocessing as mp
 from torch.utils.data import _utils, Dataset, TensorDataset, DataLoader, ConcatDataset
 from torch.utils.data._utils import ExceptionWrapper, MP_STATUS_CHECK_INTERVAL
 from torch.utils.data.dataset import random_split
-from common_utils import (TestCase, run_tests, TEST_NUMPY, IS_WINDOWS, IS_PPC, NO_MULTIPROCESSING_SPAWN,
-                          skipIfRocm, load_tests)
+from common_utils import (TestCase, run_tests, TEST_NUMPY, IS_WINDOWS, IS_PPC,
+                          IS_PYTORCH_CI, NO_MULTIPROCESSING_SPAWN, skipIfRocm,
+                          load_tests)
 
 try:
     import psutil
     HAS_PSUTIL = True
 except ImportError:
     HAS_PSUTIL = False
-    warnings.warn(
-        "psutil not found. Some crucial data loader tests relying on it (e.g., "
-        "TestDataLoader.test_proper_exit) will not run.")
+    err_msg = ("psutil not found. Some critical data loader tests relying on it "
+               "(e.g., TestDataLoader.test_proper_exit) will not run.")
+    if IS_PYTORCH_CI:
+        raise ImportError(err_msg)
+    else:
+        warnings.warn(err_msg)
 
 
 # load_tests from common_utils is used to automatically filter tests for