diff --git a/cuda_core/tests/example_tests/test_basic_examples.py b/cuda_core/tests/example_tests/test_basic_examples.py index 31b9f86e0a..b03112ba58 100644 --- a/cuda_core/tests/example_tests/test_basic_examples.py +++ b/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) diff --git a/cuda_core/tests/example_tests/utils.py b/cuda_core/tests/example_tests/utils.py new file mode 100644 index 0000000000..07dfefe4cd --- /dev/null +++ b/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()