Handle metadata automatically

Author: mdboom

cuda_core/examples/cuda_graphs.py

@@ -11,7 +11,7 @@
 # ################################################################################
 
 # /// script
-# dependencies = ["cuda_bindings", "cuda_core", "cupy-cuda13x"]
+# dependencies = ["cuda_bindings", "cuda_core", "nvidia-cuda-nvrtc", "cupy-cuda13x"]
 # ///
 
 import sys

cuda_core/examples/jit_lto_fractal.py

@@ -13,7 +13,7 @@
 # ################################################################################
 
 # /// script
-# dependencies = ["cuda_bindings", "cuda_core", "cupy-cuda13x"]
+# dependencies = ["cuda_bindings", "cuda_core", "nvidia-cuda-nvrtc", "cupy-cuda13x"]
 # ///
 
 import argparse

cuda_core/examples/memory_ops.py

@@ -11,7 +11,7 @@
 # ################################################################################
 
 # /// script
-# dependencies = ["cuda_bindings", "cuda_core", "cupy-cuda13x"]
+# dependencies = ["cuda_bindings", "cuda_core", "nvidia-cuda-nvrtc", "cupy-cuda13x"]
 # ///
 
 import sys

cuda_core/examples/strided_memory_view_gpu.py

@@ -10,7 +10,7 @@
 # ################################################################################
 
 # /// script
-# dependencies = ["cuda_bindings", "cuda_core", "cupy-cuda13x"]
+# dependencies = ["cuda_bindings", "cuda_core", "nvidia-cuda-nvrtc", "cupy-cuda13x"]
 # ///
 
 import string

cuda_core/examples/vector_add.py

@@ -10,7 +10,7 @@
 # ################################################################################
 
 # /// script
-# dependencies = ["cuda_bindings", "cuda_core", "cupy-cuda13x"]
+# dependencies = ["cuda_bindings", "cuda_core", "nvidia-cuda-nvrtc", "cupy-cuda13x"]
 # ///
 
 import cupy as cp

cuda_core/tests/example_tests/test_basic_examples.py

@@ -4,6 +4,7 @@
 # If we have subcategories of examples in the future, this file can be split along those lines
 
 import glob
+import importlib.metadata
 import os
 import platform
 import re
@@ -57,23 +58,15 @@ def has_cuda_path() -> bool:
     return os.environ.get("CUDA_PATH", os.environ.get("CUDA_HOME")) is not None
 
 
-PACKAGE_REQUIREMENTS = {
-    "cuda_graphs.py": ["cupy"],
-    "jit_lto_fractal.py": ["cupy"],
-    "memory_ops.py": ["cupy"],
-    "pytorch_example.py": ["torch"],
-    "saxpy.py": ["cupy"],
-    "simple_multi_gpu_example.py": ["cupy"],
-    "strided_memory_view_cpu.py": ["cffi"],
-    "strided_memory_view_gpu.py": ["cupy"],
-    "tma_tensor_map.py": ["cupy"],
-    "vector_add.py": ["cupy"],
-}
+# Specific system requirements for each of the examples.
 
 
 SYSTEM_REQUIREMENTS = {
     "gl_interop_plasma.py": has_display,
-    "pytorch_example.py": is_x86_64,  # PyTorch only provides CUDA support for x86_64
+    "pytorch_example.py": lambda: (
+        has_compute_capability_9_or_higher() and is_x86_64()
+    ),  # PyTorch only provides CUDA support for x86_64
+    "saxpy.py": has_compute_capability_9_or_higher,
     "simple_multi_gpu_example.py": has_multiple_devices,
     "strided_memory_view_cpu.py": is_not_windows,
     "thread_block_cluster.py": lambda: has_compute_capability_9_or_higher() and has_cuda_path(),
@@ -85,17 +78,44 @@ def has_cuda_path() -> bool:
 sample_files = [os.path.basename(x) for x in glob.glob(samples_path + "**/*.py", recursive=True)]
 
 
+def has_package_requirements_or_skip(example):
+    with open(example, encoding="utf-8") as f:
+        content = f.read()
+
+    # The canonical regex as defined in PEP 723
+    pep723 = re.search(r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$", content)
+    if not pep723:
+        return
+
+    metadata = {}
+    for line in pep723.group("content").splitlines():
+        line = line.lstrip("# ").rstrip()
+        if not line:
+            continue
+        key, value = line.split("=", 1)
+        key = key.strip()
+        value = value.strip()
+        metadata[key] = value
+
+    if "dependencies" in metadata:
+        dependencies = eval(metadata["dependencies"])  # noqa: S307
+        for dependency in dependencies:
+            name = re.match("[a-zA-Z0-9_-]+", dependency)
+            try:
+                importlib.metadata.distribution(name.string)
+            except importlib.metadata.PackageNotFoundError:
+                pytest.skip(f"Skipping {example} due to missing package requirement: {name}")
+
+
 @pytest.mark.parametrize("example", sample_files)
 def test_example(example):
-    package_requirements = PACKAGE_REQUIREMENTS.get(example, [])
-    for package in package_requirements:
-        pytest.importorskip(package, reason=f"Skipping {example} due to missing package requirement: {package}")
+    example_path = os.path.join(samples_path, example)
+    has_package_requirements_or_skip(example_path)
 
     system_requirement = SYSTEM_REQUIREMENTS.get(example, lambda: True)
     if not system_requirement():
         pytest.skip(f"Skipping {example} due to unmet system requirement")
 
-    example_path = os.path.join(samples_path, example)
     process = subprocess.run([sys.executable, example_path], capture_output=True)  # noqa: S603
     if process.returncode != 0:
         if process.stdout:
@@ -106,21 +126,19 @@ def test_example(example):
 
 
 @pytest.mark.parametrize("example", sample_files)
-@pytest.mark.skipif(not uv_installed(), reason="uv is required to test PEP 723 metadata installation")
+@pytest.mark.skipif(not uv_installed(), reason="uv is required to test PEP 723 metadata")
 def test_example_pep723(example):
+    example_path = os.path.join(samples_path, example)
+
     system_requirement = SYSTEM_REQUIREMENTS.get(example, lambda: True)
     if not system_requirement():
         pytest.skip(f"Skipping {example} due to unmet system requirement")
 
-    example_path = os.path.join(samples_path, example)
-
     # Have uv use the same version of Python that is running the test suite,
     # not because they have to match but to give Python version coverage in CI.
     version_info = sys.version_info
     py_version = f"{version_info.major}.{version_info.minor}"
 
-    print("Parent process environment:", os.environ)
-
     process = subprocess.run(["uv", "run", "--python", py_version, example_path], capture_output=True)  # noqa: S603, S607
     if process.returncode != 0:
         # This example requires a development version of cuda_core, so requirements can't be met.