Fix tests

Author: mdboom

.github/workflows/test-wheel-linux.yml

@@ -101,7 +101,7 @@ jobs:
         continue-on-error: false
         with:
           # for artifact fetching, graphics libs
-          dependencies: "jq wget libgl1 libegl1"
+          dependencies: "jq wget libgl1 libegl1 g++"
           dependent_exes: "jq wget"
 
       - name: Set environment variables

cuda_core/tests/example_tests/test_basic_examples.py

@@ -10,21 +10,42 @@
 
 import pytest
 
-DISALLOW_LIST = {
-    "gl_interop_plasma.py",  # requires a display
+from cuda.core import Device
+
+
+def pytorch_installed():
+    try:
+        import torch  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def has_compute_capability_9_or_higher():
+    dev = Device()
+    arch = dev.compute_capability
+    return arch > (9, 0)
+
+
+REQUIREMENTS = {
+    "gl_interop_plasma.py": (lambda: False),  # requires a display
+    "pytorch_example.py": pytorch_installed,  # requires PyTorch
+    "thread_block_cluster.py": has_compute_capability_9_or_higher,  # requires CC 9.0+
 }
 
 samples_path = os.path.join(os.path.dirname(__file__), "..", "..", "examples")
 sample_files = [
-    x
+    (x, REQUIREMENTS.get(x, lambda: True))
     for x in (os.path.basename(x) for x in glob.glob(samples_path + "**/*.py", recursive=True))
-    if x not in DISALLOW_LIST
 ]
 
 
-@pytest.mark.parametrize("example", sample_files)
+@pytest.mark.parametrize("example,requirement", sample_files)
 class TestExamples:
-    def test_example(self, example):
+    def test_example(self, example, requirement):
+        if not requirement():
+            pytest.skip(f"Skipping {example} due to unmet requirements")
         example_path = os.path.join(samples_path, example)
         process = subprocess.run([sys.executable, example_path], capture_output=True)  # noqa: S603
         if process.returncode != 0: