Unified skipping logic in pytest_collection_modifyitems

Author: NicolasHug

test/conftest.py

@@ -1,9 +1,52 @@
+import os
 import random
 
 import pytest
 import torch
 
 
+def pytest_configure(config):
+    # register an additional marker (see pytest_collection_modifyitems)
+    config.addinivalue_line(
+        "markers", "needs_cuda: mark for tests that rely on a CUDA device"
+    )
+
+
+def pytest_collection_modifyitems(items):
+    # This hook is called by pytest after it has collected the tests (google its
+    # name to check out its doc!). We can ignore some tests as we see fit here,
+    # or add marks, such as a skip mark.
+
+    out_items = []
+    for item in items:
+        # The needs_cuda mark will exist if the test was explicitly decorated
+        # with the @needs_cuda decorator. It will also exist if it was
+        # parametrized with a parameter that has the mark: for example if a test
+        # is parametrized with
+        # @pytest.mark.parametrize('device', cpu_and_cuda())
+        # the "instances" of the tests where device == 'cuda' will have the
+        # 'needs_cuda' mark, and the ones with device == 'cpu' won't have the
+        # mark.
+        needs_cuda = item.get_closest_marker("needs_cuda") is not None
+
+        if (
+            needs_cuda
+            and not torch.cuda.is_available()
+            and os.environ.get("FAIL_WITHOUT_CUDA") is None
+        ):
+            # We skip CUDA tests on non-CUDA machines, but only if the
+            # FAIL_WITHOUT_CUDA env var wasn't set. If it's set, the test will
+            # typically fail with a "Unsupported device: cuda" error. This is
+            # normal and desirable: this env var is set on CI jobs that are
+            # supposed to run the CUDA tests, so if CUDA isn't available on
+            # those for whatever reason, we need to know.
+            item.add_marker(pytest.mark.skip(reason="CUDA not available."))
+
+        out_items.append(item)
+
+    items[:] = out_items
+
+
 @pytest.fixture(autouse=True)
 def prevent_leaking_rng():
     # Prevent each test from leaking the rng to all other test when they call
@@ -20,10 +63,3 @@ def prevent_leaking_rng():
     random.setstate(builtin_rng_state)
     if torch.cuda.is_available():
         torch.cuda.set_rng_state(cuda_rng_state)
-
-
-def pytest_configure(config):
-    # register an additional marker (see pytest_collection_modifyitems)
-    config.addinivalue_line(
-        "markers", "needs_cuda: mark for tests that rely on a CUDA device"
-    )

test/utils.py

@@ -12,21 +12,14 @@
 import torch
 
 
-# Decorator for skipping CUDA tests when CUDA isn't available
+# Decorator for skipping CUDA tests when CUDA isn't available. The tests are
+# effectively marked to be skipped in pytest_collection_modifyitems() of
+# conftest.py
 def needs_cuda(test_item):
-    if not torch.cuda.is_available():
-        if os.environ.get("FAIL_WITHOUT_CUDA") == "1":
-            raise RuntimeError("CUDA is required for this test")
-        return pytest.mark.skip(reason="CUDA not available")(test_item)
-    return test_item
+    return pytest.mark.needs_cuda(test_item)
 
 
 def cpu_and_cuda():
-    # TODO: Figure out how to dedupe this logic in needs_cuda() and here.
-    if not torch.cuda.is_available():
-        if os.environ.get("FAIL_WITHOUT_CUDA") == "1":
-            raise RuntimeError("CUDA is required for this test")
-        return ("cpu",)
     return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))