refactor(tests): replace exec() with importlib in example runner

cuda_core/tests/example_tests/test_basic_examples.py

@@ -3,17 +3,18 @@
 
 # If we have subcategories of examples in the future, this file can be split along those lines
 
-import glob
 import os
 import platform
-import subprocess
 import sys
 import warnings
+from pathlib import Path
 
 import pytest
 
 from cuda.core import Device, ManagedMemoryResource, system
 
+from .utils import run_example
+
 try:
     from cuda.bindings._test_helpers.pep723 import has_package_requirements_or_skip
 except ImportError:
@@ -93,23 +94,24 @@ def has_recent_memory_pool_support() -> bool:
 }
 
 
-samples_path = os.path.join(os.path.dirname(__file__), "..", "..", "examples")
-sample_files = [os.path.basename(x) for x in glob.glob(samples_path + "**/*.py", recursive=True)]
+# not dividing, but navigating into the "examples" directory.
+EXAMPLES_DIR = Path(__file__).resolve().parents[2] / "examples"
+
+# recursively glob for test files in examples directory, sort for deterministic
+# test runs. Relative paths offer cleaner output when tests fail.
+SAMPLE_FILES = sorted([str(p.relative_to(EXAMPLES_DIR)) for p in EXAMPLES_DIR.glob("**/*.py")])
 
 
-@pytest.mark.parametrize("example", sample_files)
-def test_example(example):
-    example_path = os.path.join(samples_path, example)
+@pytest.mark.parametrize("example_rel_path", SAMPLE_FILES)
+# deinit_cuda is defined in conftest.py and pops the cuda context automatically.
+def test_example(example_rel_path: str, deinit_cuda) -> None:
+    example_path = str(EXAMPLES_DIR / example_rel_path)
     has_package_requirements_or_skip(example_path)
 
-    system_requirement = SYSTEM_REQUIREMENTS.get(example, lambda: True)
+    system_requirement = SYSTEM_REQUIREMENTS.get(example_rel_path, lambda: True)
     if not system_requirement():
-        pytest.skip(f"Skipping {example} due to unmet system requirement")
-
-    process = subprocess.run([sys.executable, example_path], capture_output=True)  # noqa: S603
-    if process.returncode != 0:
-        if process.stdout:
-            print(process.stdout.decode(errors="replace"))
-        if process.stderr:
-            print(process.stderr.decode(errors="replace"), file=sys.stderr)
-        raise AssertionError(f"`{example}` failed ({process.returncode})")
+        pytest.skip(f"Skipping {example_rel_path} due to unmet system requirement")
+
+    # redundantly set current device to 0 in case previous example was multi-GPU
+    Device(0).set_current()
+    run_example(str(EXAMPLES_DIR), example_rel_path)

cuda_core/tests/example_tests/utils.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+
+
+class SampleTestError(Exception):
+    pass
+
+
+def run_example(parent_dir: str, rel_path_to_example: str, env=None) -> None:
+    fullpath = Path(parent_dir) / rel_path_to_example
+    module_name = fullpath.stem
+
+    old_sys_path = sys.path.copy()
+    old_argv = sys.argv
+
+    try:
+        sys.path.append(parent_dir)
+        sys.argv = [str(fullpath)]
+
+        # Collect metadata for file 'module_name' located at 'fullpath'.
+        spec = importlib.util.spec_from_file_location(module_name, fullpath)
+
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Failed to load spec for {rel_path_to_example}")
+
+        # Otherwise convert the spec to a module, then run the module.
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+
+        # This runs top-level code.
+        spec.loader.exec_module(module)
+
+        # If the module has a main() function, call it.
+        if hasattr(module, "main"):
+            module.main()
+
+    except ImportError as e:
+        # for samples requiring any of optional dependencies
+        for m in ("cupy", "torch"):
+            if f"No module named '{m}'" in str(e):
+                pytest.skip(f"{m} not installed, skipping related tests")
+                break
+        else:
+            raise
+    except SystemExit:
+        # for samples that early return due to any missing requirements
+        pytest.skip(f"skip {rel_path_to_example}")
+    except Exception as e:
+        msg = "\n"
+        msg += f"Got error ({rel_path_to_example}):\n"
+        msg += str(e)
+        raise SampleTestError(msg) from e
+    finally:
+        sys.path = old_sys_path
+        sys.argv = old_argv
+
+        # further reduce the memory watermark
+        sys.modules.pop(module_name, None)
+        gc.collect()