From 4ec387e5bd91cb3054f3aab3dc50f9ffc3ca078d Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Fri, 9 Jan 2026 14:39:09 -0800 Subject: [PATCH 01/11] initial commit --- cuda_pathfinder/cuda/pathfinder/__init__.py | 6 + .../_binaries/find_nvidia_binaries.py | 120 ++++++++++++ .../_binaries/supported_nvidia_binaries.py | 24 +++ .../_static_libs/find_nvidia_static_libs.py | 184 ++++++++++++++++++ .../supported_nvidia_static_libs.py | 26 +++ .../tests/test_find_nvidia_binaries.py | 49 +++++ .../tests/test_find_nvidia_static_libs.py | 68 +++++++ 7 files changed, 477 insertions(+) create mode 100644 cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py create mode 100644 cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py create mode 100644 cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py create mode 100644 cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py create mode 100644 cuda_pathfinder/tests/test_find_nvidia_binaries.py create mode 100644 cuda_pathfinder/tests/test_find_nvidia_static_libs.py diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index afffd54263..d227c94740 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/cuda_pathfinder/cuda/pathfinder/__init__.py @@ -5,6 +5,8 @@ from cuda.pathfinder._version import __version__ # noqa: F401 +from cuda.pathfinder._binaries.find_nvidia_binaries import find_nvidia_binary as find_nvidia_binary +from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES as SUPPORTED_NVIDIA_BINARIES from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError as DynamicLibNotFoundError from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL as LoadedDL from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib @@ -13,6 +15,10 @@ ) from cuda.pathfinder._headers.find_nvidia_headers import find_nvidia_header_directory as find_nvidia_header_directory from cuda.pathfinder._headers.supported_nvidia_headers import SUPPORTED_HEADERS_CTK as _SUPPORTED_HEADERS_CTK +from cuda.pathfinder._static_libs.find_nvidia_static_libs import find_nvidia_static_lib as find_nvidia_static_lib +from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( + SUPPORTED_STATIC_LIBS as SUPPORTED_NVIDIA_STATIC_LIBS, +) # Indirections to help Sphinx find the docstrings. #: Mapping from short CUDA Toolkit (CTK) library names to their canonical diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py new file mode 100644 index 0000000000..f11597a606 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import functools +import os + +from cuda.pathfinder._binaries.supported_nvidia_binaries import SITE_PACKAGES_BINDIRS, SUPPORTED_BINARIES +from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path +from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +def _find_binary_under_site_packages(binary_name: str) -> str | None: + """Search for a binary in site-packages directories.""" + rel_dirs = SITE_PACKAGES_BINDIRS.get(binary_name) + if rel_dirs is None: + return None + + if IS_WINDOWS: + binary_filename = f"{binary_name}.exe" + else: + binary_filename = binary_name + + for rel_dir in rel_dirs: + for bin_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): + binary_path = os.path.join(bin_dir, binary_filename) + if os.path.isfile(binary_path): + return binary_path + return None + + +def _find_binary_in_conda(binary_name: str) -> str | None: + """Search for a binary in conda prefix.""" + conda_prefix = os.environ.get("CONDA_PREFIX") + if not conda_prefix: + return None + + if IS_WINDOWS: + binary_filename = f"{binary_name}.exe" + bin_dir = os.path.join(conda_prefix, "Library", "bin") + else: + binary_filename = binary_name + bin_dir = os.path.join(conda_prefix, "bin") + + binary_path = os.path.join(bin_dir, binary_filename) + if os.path.isfile(binary_path): + return binary_path + return None + + +def _find_binary_in_cuda_home(binary_name: str) -> str | None: + """Search for a binary in CUDA_HOME or CUDA_PATH.""" + cuda_home = get_cuda_home_or_path() + if cuda_home is None: + return None + + if IS_WINDOWS: + binary_filename = f"{binary_name}.exe" + else: + binary_filename = binary_name + + bin_dir = os.path.join(cuda_home, "bin") + binary_path = os.path.join(bin_dir, binary_filename) + if os.path.isfile(binary_path): + return binary_path + return None + + +@functools.cache +def find_nvidia_binary(binary_name: str) -> str | None: + """Locate a CUDA binary executable. + + Args: + binary_name (str): The name of the binary to find (e.g., ``"nvdisasm"``, + ``"cuobjdump"``). + + Returns: + str or None: Absolute path to the discovered binary, or ``None`` if the + binary cannot be found. + + Raises: + RuntimeError: If ``binary_name`` is not in the supported set. + + Search order: + 1. **NVIDIA Python wheels** + + - Scan installed distributions (``site-packages``) for binaries + shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvcc]``). + + 2. **Conda environments** + + - Check Conda-style installation prefixes (``$CONDA_PREFIX/bin`` on + Linux/Mac or ``$CONDA_PREFIX/Library/bin`` on Windows). + + 3. **CUDA Toolkit environment variables** + + - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in the + ``bin`` subdirectory. + """ + def _abs_norm(path: str | None) -> str | None: + if path: + return os.path.normpath(os.path.abspath(path)) + return None + + if binary_name not in SUPPORTED_BINARIES: + raise RuntimeError(f"UNKNOWN {binary_name=}") + + # Try site-packages first + if binary_path := _find_binary_under_site_packages(binary_name): + return _abs_norm(binary_path) + + # Try conda prefix + if binary_path := _find_binary_in_conda(binary_name): + return _abs_norm(binary_path) + + # Try CUDA_HOME/CUDA_PATH + if binary_path := _find_binary_in_cuda_home(binary_name): + return _abs_norm(binary_path) + + return None diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py new file mode 100644 index 0000000000..0e9bc6ca8a --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py @@ -0,0 +1,24 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# THIS FILE NEEDS TO BE REVIEWED/UPDATED FOR EACH CTK RELEASE +# Likely candidates for updates are: +# SUPPORTED_BINARIES +# SITE_PACKAGES_BINDIRS + +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + +# Supported CUDA binaries that can be found +SUPPORTED_BINARIES_COMMON = ( + "nvdisasm", + "cuobjdump", +) + +SUPPORTED_BINARIES = SUPPORTED_BINARIES_COMMON + +# Map from binary name to relative paths under site-packages +# These are typically from cuda-toolkit[nvcc] wheels +SITE_PACKAGES_BINDIRS = { + "nvdisasm": ["nvidia/cuda_nvcc/bin"], + "cuobjdump": ["nvidia/cuda_nvcc/bin"], +} diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py new file mode 100644 index 0000000000..fc8a320f65 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py @@ -0,0 +1,184 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import functools +import glob +import os + +from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( + SITE_PACKAGES_STATIC_LIBDIRS, + SUPPORTED_STATIC_LIBS, +) +from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path +from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +def _find_artifact_under_site_packages(artifact_name: str) -> str | None: + """Search for an artifact in site-packages directories.""" + rel_dirs = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name) + if rel_dirs is None: + return None + + for rel_dir in rel_dirs: + for lib_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): + artifact_path = os.path.join(lib_dir, artifact_name) + if os.path.isfile(artifact_path): + return artifact_path + return None + + +def _find_libdevice_in_conda() -> str | None: + """Search for libdevice.10.bc in conda prefix.""" + conda_prefix = os.environ.get("CONDA_PREFIX") + if not conda_prefix: + return None + + # Check multiple possible locations + if IS_WINDOWS: + possible_paths = [ + os.path.join(conda_prefix, "Library", "nvvm", "libdevice", "libdevice.10.bc"), + ] + else: + possible_paths = [ + os.path.join(conda_prefix, "nvvm", "libdevice", "libdevice.10.bc"), + ] + + for path in possible_paths: + if os.path.isfile(path): + return path + return None + + +def _find_libdevice_in_cuda_home() -> str | None: + """Search for libdevice.10.bc in CUDA_HOME or CUDA_PATH.""" + cuda_home = get_cuda_home_or_path() + if cuda_home is None: + return None + + libdevice_path = os.path.join(cuda_home, "nvvm", "libdevice", "libdevice.10.bc") + if os.path.isfile(libdevice_path): + return libdevice_path + return None + + +def _find_libcudadevrt_in_conda() -> str | None: + """Search for libcudadevrt.a in conda prefix.""" + conda_prefix = os.environ.get("CONDA_PREFIX") + if not conda_prefix: + return None + + if IS_WINDOWS: + # On Windows, it might be cudadevrt.lib + possible_paths = [ + os.path.join(conda_prefix, "Library", "lib", "libcudadevrt.a"), + os.path.join(conda_prefix, "Library", "lib", "x64", "cudadevrt.lib"), + ] + else: + possible_paths = [ + os.path.join(conda_prefix, "lib", "libcudadevrt.a"), + ] + + for path in possible_paths: + if os.path.isfile(path): + return path + return None + + +def _find_libcudadevrt_in_cuda_home() -> str | None: + """Search for libcudadevrt.a in CUDA_HOME or CUDA_PATH.""" + cuda_home = get_cuda_home_or_path() + if cuda_home is None: + return None + + if IS_WINDOWS: + # On Windows, check for cudadevrt.lib in various locations + possible_paths = [ + os.path.join(cuda_home, "lib", "x64", "cudadevrt.lib"), + os.path.join(cuda_home, "lib", "cudadevrt.lib"), + ] + else: + # On Linux, check lib64 and lib + possible_paths = [ + os.path.join(cuda_home, "lib64", "libcudadevrt.a"), + os.path.join(cuda_home, "lib", "libcudadevrt.a"), + ] + + for path in possible_paths: + if os.path.isfile(path): + return path + + # Also check targets subdirectories (for cross-compilation setups) + if not IS_WINDOWS: + targets_pattern = os.path.join(cuda_home, "targets", "*", "lib", "libcudadevrt.a") + for path in sorted(glob.glob(targets_pattern), reverse=True): + if os.path.isfile(path): + return path + + return None + + +@functools.cache +def find_nvidia_static_lib(artifact_name: str) -> str | None: + """Locate a CUDA static library or artifact file. + + Args: + artifact_name (str): The name of the artifact to find (e.g., + ``"libdevice.10.bc"``, ``"libcudadevrt.a"``). + + Returns: + str or None: Absolute path to the discovered artifact, or ``None`` if the + artifact cannot be found. + + Raises: + RuntimeError: If ``artifact_name`` is not in the supported set. + + Search order: + 1. **NVIDIA Python wheels** + + - Scan installed distributions (``site-packages``) for artifacts + shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvvm]``, + ``cuda-toolkit[cudart]``). + + 2. **Conda environments** + + - Check Conda-style installation prefixes for the specific artifact + layout. + + 3. **CUDA Toolkit environment variables** + + - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in + standard CUDA Toolkit directory layouts. + """ + if artifact_name not in SUPPORTED_STATIC_LIBS: + raise RuntimeError(f"UNKNOWN {artifact_name=}") + + def _abs_norm(path: str | None) -> str | None: + if path: + return os.path.normpath(os.path.abspath(path)) + return None + + # Try site-packages first + if artifact_path := _find_artifact_under_site_packages(artifact_name): + return _abs_norm(artifact_path) + + # Handle specific artifacts with custom search logic + if artifact_name == "libdevice.10.bc": + # Try conda prefix + if artifact_path := _find_libdevice_in_conda(): + return _abs_norm(artifact_path) + + # Try CUDA_HOME/CUDA_PATH + if artifact_path := _find_libdevice_in_cuda_home(): + return _abs_norm(artifact_path) + + elif artifact_name == "libcudadevrt.a": + # Try conda prefix + if artifact_path := _find_libcudadevrt_in_conda(): + return _abs_norm(artifact_path) + + # Try CUDA_HOME/CUDA_PATH + if artifact_path := _find_libcudadevrt_in_cuda_home(): + return _abs_norm(artifact_path) + + return None diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py new file mode 100644 index 0000000000..0ad2e83e58 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# THIS FILE NEEDS TO BE REVIEWED/UPDATED FOR EACH CTK RELEASE +# Likely candidates for updates are: +# SUPPORTED_STATIC_LIBS +# SITE_PACKAGES_STATIC_LIBDIRS + +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + +# Supported CUDA static libraries and artifacts that can be found +SUPPORTED_STATIC_LIBS_COMMON = ( + "libdevice.10.bc", + "libcudadevrt.a", +) + +SUPPORTED_STATIC_LIBS = SUPPORTED_STATIC_LIBS_COMMON + +# Map from artifact name to relative paths under site-packages +SITE_PACKAGES_STATIC_LIBDIRS = { + "libdevice.10.bc": ["nvidia/cuda_nvvm/nvvm/libdevice"], + "libcudadevrt.a": [ + "nvidia/cuda_cudart/lib", # Linux + "nvidia/cuda_cudart/lib/x64", # Windows (if present) + ], +} diff --git a/cuda_pathfinder/tests/test_find_nvidia_binaries.py b/cuda_pathfinder/tests/test_find_nvidia_binaries.py new file mode 100644 index 0000000000..15ba94c361 --- /dev/null +++ b/cuda_pathfinder/tests/test_find_nvidia_binaries.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest + +from cuda.pathfinder import find_nvidia_binary +from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES + +STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_BINARIES_STRICTNESS", "see_what_works") +assert STRICTNESS in ("see_what_works", "all_must_work") + + +def test_unknown_binary(): + with pytest.raises(RuntimeError, match=r"^UNKNOWN binary_name='unknown-binary'$"): + find_nvidia_binary("unknown-binary") + + +@pytest.mark.parametrize("binary_name", SUPPORTED_BINARIES) +def test_find_binaries(info_summary_append, binary_name): + binary_path = find_nvidia_binary(binary_name) + info_summary_append(f"{binary_path=!r}") + if binary_path: + assert os.path.isfile(binary_path) + # Verify the binary name is in the path + assert binary_name in os.path.basename(binary_path) + if STRICTNESS == "all_must_work": + assert binary_path is not None + + +def test_nvdisasm_specific(info_summary_append): + """Specific test for nvdisasm to ensure it's working.""" + binary_path = find_nvidia_binary("nvdisasm") + info_summary_append(f"nvdisasm path: {binary_path!r}") + # Only assert if we're in strict mode or if cuda-toolkit is installed + if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): + if binary_path: + assert os.path.isfile(binary_path) + + +def test_cuobjdump_specific(info_summary_append): + """Specific test for cuobjdump to ensure it's working.""" + binary_path = find_nvidia_binary("cuobjdump") + info_summary_append(f"cuobjdump path: {binary_path!r}") + # Only assert if we're in strict mode or if cuda-toolkit is installed + if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): + if binary_path: + assert os.path.isfile(binary_path) diff --git a/cuda_pathfinder/tests/test_find_nvidia_static_libs.py b/cuda_pathfinder/tests/test_find_nvidia_static_libs.py new file mode 100644 index 0000000000..9d6c14699f --- /dev/null +++ b/cuda_pathfinder/tests/test_find_nvidia_static_libs.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest + +from cuda.pathfinder import find_nvidia_static_lib +from cuda.pathfinder._static_libs.supported_nvidia_static_libs import SUPPORTED_STATIC_LIBS + +STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_STATIC_LIBS_STRICTNESS", "see_what_works") +assert STRICTNESS in ("see_what_works", "all_must_work") + + +def test_unknown_artifact(): + with pytest.raises(RuntimeError, match=r"^UNKNOWN artifact_name='unknown-artifact'$"): + find_nvidia_static_lib("unknown-artifact") + + +@pytest.mark.parametrize("artifact_name", SUPPORTED_STATIC_LIBS) +def test_find_static_libs(info_summary_append, artifact_name): + artifact_path = find_nvidia_static_lib(artifact_name) + info_summary_append(f"{artifact_path=!r}") + if artifact_path: + assert os.path.isfile(artifact_path) + # Verify the artifact name (or its base) is in the path + base_name = artifact_name.replace(".10", "") # Handle libdevice.10.bc -> libdevice + assert base_name.split(".")[0] in artifact_path.lower() + if STRICTNESS == "all_must_work": + assert artifact_path is not None + + +def test_libdevice_specific(info_summary_append): + """Specific test for libdevice.10.bc to ensure it's working.""" + artifact_path = find_nvidia_static_lib("libdevice.10.bc") + info_summary_append(f"libdevice.10.bc path: {artifact_path!r}") + if artifact_path: + assert os.path.isfile(artifact_path) + assert "libdevice" in artifact_path + # Should end with .bc + assert artifact_path.endswith(".bc") + # Only assert existence if we're in strict mode or if cuda is installed + if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): + if artifact_path: + assert os.path.isfile(artifact_path) + + +def test_libcudadevrt_specific(info_summary_append): + """Specific test for libcudadevrt.a to ensure it's working.""" + artifact_path = find_nvidia_static_lib("libcudadevrt.a") + info_summary_append(f"libcudadevrt.a path: {artifact_path!r}") + if artifact_path: + assert os.path.isfile(artifact_path) + # On Linux it should be .a, on Windows it might be .lib + assert artifact_path.endswith((".a", ".lib")) + assert "cudadevrt" in artifact_path.lower() + # Only assert existence if we're in strict mode or if cuda is installed + if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): + if artifact_path: + assert os.path.isfile(artifact_path) + + +def test_caching(): + """Test that the find functions are properly cached.""" + # Call twice and ensure we get the same object (due to functools.cache) + path1 = find_nvidia_static_lib("libdevice.10.bc") + path2 = find_nvidia_static_lib("libdevice.10.bc") + assert path1 is path2 # Should be the exact same object due to caching From 624213be2fc69c4fd419aa791ffb635509b61973 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Fri, 9 Jan 2026 14:59:03 -0800 Subject: [PATCH 02/11] one implementation of _abs_norm --- .../cuda/pathfinder/_binaries/find_nvidia_binaries.py | 6 +----- .../cuda/pathfinder/_headers/find_nvidia_headers.py | 7 +------ .../pathfinder/_static_libs/find_nvidia_static_libs.py | 6 +----- 3 files changed, 3 insertions(+), 16 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py index f11597a606..3057f76a47 100644 --- a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py @@ -7,6 +7,7 @@ from cuda.pathfinder._binaries.supported_nvidia_binaries import SITE_PACKAGES_BINDIRS, SUPPORTED_BINARIES from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.path_utils import _abs_norm from cuda.pathfinder._utils.platform_aware import IS_WINDOWS @@ -97,11 +98,6 @@ def find_nvidia_binary(binary_name: str) -> str | None: - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in the ``bin`` subdirectory. """ - def _abs_norm(path: str | None) -> str | None: - if path: - return os.path.normpath(os.path.abspath(path)) - return None - if binary_name not in SUPPORTED_BINARIES: raise RuntimeError(f"UNKNOWN {binary_name=}") diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py index 63f8a627fd..f0755034ab 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py @@ -8,15 +8,10 @@ from cuda.pathfinder._headers import supported_nvidia_headers from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.path_utils import _abs_norm from cuda.pathfinder._utils.platform_aware import IS_WINDOWS -def _abs_norm(path: str | None) -> str | None: - if path: - return os.path.normpath(os.path.abspath(path)) - return None - - def _joined_isfile(dirpath: str, basename: str) -> bool: return os.path.isfile(os.path.join(dirpath, basename)) diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py index fc8a320f65..2cb5788d0a 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py @@ -11,6 +11,7 @@ ) from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.path_utils import _abs_norm from cuda.pathfinder._utils.platform_aware import IS_WINDOWS @@ -153,11 +154,6 @@ def find_nvidia_static_lib(artifact_name: str) -> str | None: if artifact_name not in SUPPORTED_STATIC_LIBS: raise RuntimeError(f"UNKNOWN {artifact_name=}") - def _abs_norm(path: str | None) -> str | None: - if path: - return os.path.normpath(os.path.abspath(path)) - return None - # Try site-packages first if artifact_path := _find_artifact_under_site_packages(artifact_name): return _abs_norm(artifact_path) From 24fefb9b3162e727ece38a23b61ad8024215b732 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Fri, 9 Jan 2026 14:59:16 -0800 Subject: [PATCH 03/11] missing file --- .../cuda/pathfinder/_utils/path_utils.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/path_utils.py diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/path_utils.py b/cuda_pathfinder/cuda/pathfinder/_utils/path_utils.py new file mode 100644 index 0000000000..b5105d1950 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/path_utils.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + + +def _abs_norm(path: str | None) -> str | None: + """Normalize and convert a path to an absolute path. + + Args: + path (str or None): The path to normalize and make absolute. + + Returns: + str or None: The normalized absolute path, or None if the input is None + or empty. + """ + if path: + return os.path.normpath(os.path.abspath(path)) + return None From 6985ee8e3dbcc17c38636c1b839cac58c475b7f2 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Fri, 9 Jan 2026 15:27:37 -0800 Subject: [PATCH 04/11] removing file specific lookup --- .../_static_libs/artifact_search_config.py | 74 ++++++++ .../_static_libs/find_nvidia_static_libs.py | 170 +++++++++--------- 2 files changed, 159 insertions(+), 85 deletions(-) create mode 100644 cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py new file mode 100644 index 0000000000..43fefa0e8e --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Configuration for artifact search patterns across platforms and environments.""" + +from dataclasses import dataclass + +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +@dataclass(frozen=True) +class ArtifactSearchConfig: + """Configuration for searching a specific artifact (platform-resolved at init). + + Attributes: + canonical_name: Platform-agnostic identifier for the artifact. + filenames: Filenames to search for on the current platform. + conda_dirs: Directory paths relative to CONDA_PREFIX. + cuda_home_dirs: Directory paths relative to CUDA_HOME. + search_targets_subdirs: Whether to search targets/*/{lib,lib64} subdirs + for cross-compilation setups (Linux only). + """ + + canonical_name: str + filenames: tuple[str, ...] + conda_dirs: tuple[str, ...] + cuda_home_dirs: tuple[str, ...] + search_targets_subdirs: bool = False + + +def _create_config( + canonical_name: str, + linux_filenames: tuple[str, ...], + windows_filenames: tuple[str, ...], + conda_linux_dirs: tuple[str, ...], + conda_windows_dirs: tuple[str, ...], + cuda_home_linux_dirs: tuple[str, ...], + cuda_home_windows_dirs: tuple[str, ...], + search_targets_subdirs: bool = False, +) -> ArtifactSearchConfig: + """Create a platform-specific config by selecting appropriate values at init time.""" + return ArtifactSearchConfig( + canonical_name=canonical_name, + filenames=windows_filenames if IS_WINDOWS else linux_filenames, + conda_dirs=conda_windows_dirs if IS_WINDOWS else conda_linux_dirs, + cuda_home_dirs=cuda_home_windows_dirs if IS_WINDOWS else cuda_home_linux_dirs, + search_targets_subdirs=search_targets_subdirs, + ) + + +# Registry of all supported artifacts with their search configurations +# Platform selection happens once at module import time via _create_config +ARTIFACT_CONFIGS = { + "libcudadevrt.a": _create_config( + canonical_name="cudadevrt", + linux_filenames=("libcudadevrt.a",), + windows_filenames=("cudadevrt.lib",), + conda_linux_dirs=("lib",), + conda_windows_dirs=("Library/lib", "Library/lib/x64"), + cuda_home_linux_dirs=("lib64", "lib"), + cuda_home_windows_dirs=("lib/x64", "lib"), + search_targets_subdirs=True, + ), + "libdevice.10.bc": _create_config( + canonical_name="libdevice", + linux_filenames=("libdevice.10.bc",), + windows_filenames=("libdevice.10.bc",), + conda_linux_dirs=("nvvm/libdevice",), + conda_windows_dirs=("Library/nvvm/libdevice",), + cuda_home_linux_dirs=("nvvm/libdevice",), + cuda_home_windows_dirs=("nvvm/libdevice",), + search_targets_subdirs=False, + ), +} diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py index 2cb5788d0a..a7be94a88f 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py @@ -4,7 +4,9 @@ import functools import glob import os +from typing import Sequence +from cuda.pathfinder._static_libs.artifact_search_config import ARTIFACT_CONFIGS from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( SITE_PACKAGES_STATIC_LIBDIRS, SUPPORTED_STATIC_LIBS, @@ -29,92 +31,101 @@ def _find_artifact_under_site_packages(artifact_name: str) -> str | None: return None -def _find_libdevice_in_conda() -> str | None: - """Search for libdevice.10.bc in conda prefix.""" - conda_prefix = os.environ.get("CONDA_PREFIX") - if not conda_prefix: - return None +def _search_paths( + base_dir: str, + subdirs: Sequence[str], + filenames: Sequence[str], +) -> str | None: + """Search for a file in multiple subdirectories of a base directory. + + Args: + base_dir: The base directory to search in. + subdirs: Subdirectories to check (relative to base_dir). + filenames: Filenames to look for in each subdirectory. - # Check multiple possible locations - if IS_WINDOWS: - possible_paths = [ - os.path.join(conda_prefix, "Library", "nvvm", "libdevice", "libdevice.10.bc"), - ] - else: - possible_paths = [ - os.path.join(conda_prefix, "nvvm", "libdevice", "libdevice.10.bc"), - ] - - for path in possible_paths: - if os.path.isfile(path): - return path + Returns: + First matching file path, or None if not found. + """ + for subdir in subdirs: + dir_path = os.path.join(base_dir, subdir) + for filename in filenames: + file_path = os.path.join(dir_path, filename) + if os.path.isfile(file_path): + return file_path return None -def _find_libdevice_in_cuda_home() -> str | None: - """Search for libdevice.10.bc in CUDA_HOME or CUDA_PATH.""" - cuda_home = get_cuda_home_or_path() - if cuda_home is None: - return None +def _search_targets_subdirs( + cuda_home: str, + filenames: Sequence[str], +) -> str | None: + """Search in targets/*/lib and targets/*/lib64 subdirectories. + + For cross-compilation setups. Returns the first match, preferring + more recently modified targets. - libdevice_path = os.path.join(cuda_home, "nvvm", "libdevice", "libdevice.10.bc") - if os.path.isfile(libdevice_path): - return libdevice_path + Args: + cuda_home: The CUDA home directory. + filenames: Filenames to search for. + + Returns: + First matching file path, or None if not found. + """ + for lib_subdir in ("lib64", "lib"): + pattern = os.path.join(cuda_home, "targets", "*", lib_subdir) + for dir_path in sorted(glob.glob(pattern), reverse=True): + for filename in filenames: + file_path = os.path.join(dir_path, filename) + if os.path.isfile(file_path): + return file_path return None -def _find_libcudadevrt_in_conda() -> str | None: - """Search for libcudadevrt.a in conda prefix.""" +def _find_artifact_in_conda(artifact_name: str) -> str | None: + """Generic conda prefix search for any configured artifact. + + Args: + artifact_name: The name of the artifact to find. + + Returns: + Path to the artifact if found, None otherwise. + """ conda_prefix = os.environ.get("CONDA_PREFIX") if not conda_prefix: return None - if IS_WINDOWS: - # On Windows, it might be cudadevrt.lib - possible_paths = [ - os.path.join(conda_prefix, "Library", "lib", "libcudadevrt.a"), - os.path.join(conda_prefix, "Library", "lib", "x64", "cudadevrt.lib"), - ] - else: - possible_paths = [ - os.path.join(conda_prefix, "lib", "libcudadevrt.a"), - ] - - for path in possible_paths: - if os.path.isfile(path): - return path - return None + config = ARTIFACT_CONFIGS.get(artifact_name) + if not config: + return None + + return _search_paths(conda_prefix, config.conda_dirs, config.filenames) -def _find_libcudadevrt_in_cuda_home() -> str | None: - """Search for libcudadevrt.a in CUDA_HOME or CUDA_PATH.""" +def _find_artifact_in_cuda_home(artifact_name: str) -> str | None: + """Generic CUDA_HOME/CUDA_PATH search for any configured artifact. + + Args: + artifact_name: The name of the artifact to find. + + Returns: + Path to the artifact if found, None otherwise. + """ cuda_home = get_cuda_home_or_path() if cuda_home is None: return None - if IS_WINDOWS: - # On Windows, check for cudadevrt.lib in various locations - possible_paths = [ - os.path.join(cuda_home, "lib", "x64", "cudadevrt.lib"), - os.path.join(cuda_home, "lib", "cudadevrt.lib"), - ] - else: - # On Linux, check lib64 and lib - possible_paths = [ - os.path.join(cuda_home, "lib64", "libcudadevrt.a"), - os.path.join(cuda_home, "lib", "libcudadevrt.a"), - ] - - for path in possible_paths: - if os.path.isfile(path): - return path - - # Also check targets subdirectories (for cross-compilation setups) - if not IS_WINDOWS: - targets_pattern = os.path.join(cuda_home, "targets", "*", "lib", "libcudadevrt.a") - for path in sorted(glob.glob(targets_pattern), reverse=True): - if os.path.isfile(path): - return path + config = ARTIFACT_CONFIGS.get(artifact_name) + if not config: + return None + + # Try standard directories first + if result := _search_paths(cuda_home, config.cuda_home_dirs, config.filenames): + return result + + # Try targets subdirectories for cross-compilation + if config.search_targets_subdirs and not IS_WINDOWS: + if result := _search_targets_subdirs(cuda_home, config.filenames): + return result return None @@ -158,23 +169,12 @@ def find_nvidia_static_lib(artifact_name: str) -> str | None: if artifact_path := _find_artifact_under_site_packages(artifact_name): return _abs_norm(artifact_path) - # Handle specific artifacts with custom search logic - if artifact_name == "libdevice.10.bc": - # Try conda prefix - if artifact_path := _find_libdevice_in_conda(): - return _abs_norm(artifact_path) - - # Try CUDA_HOME/CUDA_PATH - if artifact_path := _find_libdevice_in_cuda_home(): - return _abs_norm(artifact_path) - - elif artifact_name == "libcudadevrt.a": - # Try conda prefix - if artifact_path := _find_libcudadevrt_in_conda(): - return _abs_norm(artifact_path) + # Try conda prefix (generic, configuration-driven) + if artifact_path := _find_artifact_in_conda(artifact_name): + return _abs_norm(artifact_path) - # Try CUDA_HOME/CUDA_PATH - if artifact_path := _find_libcudadevrt_in_cuda_home(): - return _abs_norm(artifact_path) + # Try CUDA_HOME/CUDA_PATH (generic, configuration-driven) + if artifact_path := _find_artifact_in_cuda_home(artifact_name): + return _abs_norm(artifact_path) return None From 5ae7580f12b5a0ab150016b7aa023e76568d9f55 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Fri, 9 Jan 2026 15:37:17 -0800 Subject: [PATCH 05/11] removing platform specific filenames and path logic from search algorithms --- .../_binaries/find_nvidia_binaries.py | 83 +++++++++++++------ 1 file changed, 56 insertions(+), 27 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py index 3057f76a47..5fefc9103d 100644 --- a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py @@ -3,12 +3,46 @@ import functools import os +from typing import Sequence from cuda.pathfinder._binaries.supported_nvidia_binaries import SITE_PACKAGES_BINDIRS, SUPPORTED_BINARIES from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages from cuda.pathfinder._utils.path_utils import _abs_norm -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +def _get_binary_filename_candidates(binary_name: str) -> tuple[str, ...]: + """Generate possible binary filename variations. + + Returns multiple candidates to support fuzzy search across platforms. + The filesystem will naturally filter to what exists. + + Args: + binary_name: Base name of the binary (e.g., "nvdisasm"). + + Returns: + Tuple of possible filenames to try (exact name, with .exe extension). + """ + # Try exact name first, then with .exe extension + # This works across platforms - non-existent files simply won't be found + return (binary_name, f"{binary_name}.exe") + + +def _find_file_in_dir(directory: str, filename_candidates: Sequence[str]) -> str | None: + """Search for the first existing file from candidates in a directory. + + Args: + directory: Directory to search in. + filename_candidates: Possible filenames to try. + + Returns: + Path to first matching file, or None if none found. + """ + for filename in filename_candidates: + file_path = os.path.join(directory, filename) + if os.path.isfile(file_path): + return file_path + return None def _find_binary_under_site_packages(binary_name: str) -> str | None: @@ -17,35 +51,36 @@ def _find_binary_under_site_packages(binary_name: str) -> str | None: if rel_dirs is None: return None - if IS_WINDOWS: - binary_filename = f"{binary_name}.exe" - else: - binary_filename = binary_name + filename_candidates = _get_binary_filename_candidates(binary_name) for rel_dir in rel_dirs: for bin_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): - binary_path = os.path.join(bin_dir, binary_filename) - if os.path.isfile(binary_path): - return binary_path + if found := _find_file_in_dir(bin_dir, filename_candidates): + return found return None def _find_binary_in_conda(binary_name: str) -> str | None: - """Search for a binary in conda prefix.""" + """Search for a binary in conda prefix. + + Searches common conda bin directory locations across platforms. + """ conda_prefix = os.environ.get("CONDA_PREFIX") if not conda_prefix: return None - if IS_WINDOWS: - binary_filename = f"{binary_name}.exe" - bin_dir = os.path.join(conda_prefix, "Library", "bin") - else: - binary_filename = binary_name - bin_dir = os.path.join(conda_prefix, "bin") + filename_candidates = _get_binary_filename_candidates(binary_name) + + # Try both Windows and Unix bin directory layouts + # The filesystem will naturally tell us what exists + bin_dirs = ( + os.path.join(conda_prefix, "Library", "bin"), # Windows conda layout + os.path.join(conda_prefix, "bin"), # Unix conda layout + ) - binary_path = os.path.join(bin_dir, binary_filename) - if os.path.isfile(binary_path): - return binary_path + for bin_dir in bin_dirs: + if found := _find_file_in_dir(bin_dir, filename_candidates): + return found return None @@ -55,16 +90,10 @@ def _find_binary_in_cuda_home(binary_name: str) -> str | None: if cuda_home is None: return None - if IS_WINDOWS: - binary_filename = f"{binary_name}.exe" - else: - binary_filename = binary_name - + filename_candidates = _get_binary_filename_candidates(binary_name) bin_dir = os.path.join(cuda_home, "bin") - binary_path = os.path.join(bin_dir, binary_filename) - if os.path.isfile(binary_path): - return binary_path - return None + + return _find_file_in_dir(bin_dir, filename_candidates) @functools.cache From 7f7b8b81acb8388901aafdb67480b4cf0532c7c6 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Sat, 10 Jan 2026 13:26:53 -0800 Subject: [PATCH 06/11] wip --- cuda_pathfinder/cuda/pathfinder/__init__.py | 5 + .../_binaries/find_nvidia_binaries.py | 161 ++++++-------- .../_static_libs/find_nvidia_static_libs.py | 205 ++++++++---------- .../pathfinder/_utils/toolchain_tracker.py | 176 +++++++++++++++ cuda_pathfinder/tests/conftest.py | 10 + .../tests/test_toolchain_tracker.py | 152 +++++++++++++ 6 files changed, 499 insertions(+), 210 deletions(-) create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py create mode 100644 cuda_pathfinder/tests/test_toolchain_tracker.py diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index d227c94740..8441b95fc4 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/cuda_pathfinder/cuda/pathfinder/__init__.py @@ -19,6 +19,11 @@ from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( SUPPORTED_STATIC_LIBS as SUPPORTED_NVIDIA_STATIC_LIBS, ) +from cuda.pathfinder._utils.toolchain_tracker import ( + SearchContext as SearchContext, + ToolchainMismatchError as ToolchainMismatchError, + reset_default_context as reset_default_context, +) # Indirections to help Sphinx find the docstrings. #: Mapping from short CUDA Toolkit (CTK) library names to their canonical diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py index 5fefc9103d..74c8ed875b 100644 --- a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py @@ -9,137 +9,102 @@ from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages from cuda.pathfinder._utils.path_utils import _abs_norm +from cuda.pathfinder._utils.toolchain_tracker import ( + SearchContext, + SearchLocation, + ToolchainSource, + get_default_context, +) -def _get_binary_filename_candidates(binary_name: str) -> tuple[str, ...]: - """Generate possible binary filename variations. - - Returns multiple candidates to support fuzzy search across platforms. - The filesystem will naturally filter to what exists. +def _binary_filename_variants(name: str) -> Sequence[str]: + """Generate filename variants for a binary (cross-platform). Args: - binary_name: Base name of the binary (e.g., "nvdisasm"). + name: Base binary name. Returns: - Tuple of possible filenames to try (exact name, with .exe extension). + Tuple of possible filenames (e.g., "nvcc", "nvcc.exe"). """ - # Try exact name first, then with .exe extension - # This works across platforms - non-existent files simply won't be found - return (binary_name, f"{binary_name}.exe") + return (name, f"{name}.exe") -def _find_file_in_dir(directory: str, filename_candidates: Sequence[str]) -> str | None: - """Search for the first existing file from candidates in a directory. +def _get_site_packages_subdirs(binary_name: str) -> Sequence[str]: + """Get site-packages subdirectories for a binary. Args: - directory: Directory to search in. - filename_candidates: Possible filenames to try. + binary_name: Name of the binary. Returns: - Path to first matching file, or None if none found. + List of subdirectories to search, or empty list if binary not in site-packages. """ - for filename in filename_candidates: - file_path = os.path.join(directory, filename) - if os.path.isfile(file_path): - return file_path - return None - - -def _find_binary_under_site_packages(binary_name: str) -> str | None: - """Search for a binary in site-packages directories.""" rel_dirs = SITE_PACKAGES_BINDIRS.get(binary_name) - if rel_dirs is None: - return None - - filename_candidates = _get_binary_filename_candidates(binary_name) + if not rel_dirs: + return [] + # Expand site-packages paths + subdirs = [] for rel_dir in rel_dirs: - for bin_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): - if found := _find_file_in_dir(bin_dir, filename_candidates): - return found - return None - - -def _find_binary_in_conda(binary_name: str) -> str | None: - """Search for a binary in conda prefix. - - Searches common conda bin directory locations across platforms. - """ - conda_prefix = os.environ.get("CONDA_PREFIX") - if not conda_prefix: - return None - - filename_candidates = _get_binary_filename_candidates(binary_name) - - # Try both Windows and Unix bin directory layouts - # The filesystem will naturally tell us what exists - bin_dirs = ( - os.path.join(conda_prefix, "Library", "bin"), # Windows conda layout - os.path.join(conda_prefix, "bin"), # Unix conda layout - ) - - for bin_dir in bin_dirs: - if found := _find_file_in_dir(bin_dir, filename_candidates): - return found - return None + for found_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): + subdirs.append(found_dir) + return subdirs -def _find_binary_in_cuda_home(binary_name: str) -> str | None: - """Search for a binary in CUDA_HOME or CUDA_PATH.""" - cuda_home = get_cuda_home_or_path() - if cuda_home is None: - return None +# Define search locations for binaries +def _create_search_locations(binary_name: str) -> list[SearchLocation]: + """Create search location configurations for a specific binary. - filename_candidates = _get_binary_filename_candidates(binary_name) - bin_dir = os.path.join(cuda_home, "bin") + Args: + binary_name: Name of the binary to search for. - return _find_file_in_dir(bin_dir, filename_candidates) + Returns: + List of SearchLocation objects to try. + """ + return [ + SearchLocation( + source=ToolchainSource.SITE_PACKAGES, + base_dir_func=lambda: None, # Use subdirs for full paths + subdirs=_get_site_packages_subdirs(binary_name), + filename_variants=_binary_filename_variants, + ), + SearchLocation( + source=ToolchainSource.CONDA, + base_dir_func=lambda: os.environ.get("CONDA_PREFIX"), + subdirs=["Library/bin", "bin"], # Windows and Unix layouts + filename_variants=_binary_filename_variants, + ), + SearchLocation( + source=ToolchainSource.CUDA_HOME, + base_dir_func=get_cuda_home_or_path, + subdirs=["bin"], + filename_variants=_binary_filename_variants, + ), + ] @functools.cache -def find_nvidia_binary(binary_name: str) -> str | None: +def find_nvidia_binary(binary_name: str, *, context: SearchContext | None = None) -> str | None: """Locate a CUDA binary executable. Args: - binary_name (str): The name of the binary to find (e.g., ``"nvdisasm"``, - ``"cuobjdump"``). + binary_name: Name of the binary (e.g., "nvdisasm", "cuobjdump"). + context: Optional SearchContext for toolchain consistency tracking. + If None, uses the default module-level context. Returns: - str or None: Absolute path to the discovered binary, or ``None`` if the - binary cannot be found. + Absolute path to the binary, or None if not found. Raises: - RuntimeError: If ``binary_name`` is not in the supported set. - - Search order: - 1. **NVIDIA Python wheels** - - - Scan installed distributions (``site-packages``) for binaries - shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvcc]``). - - 2. **Conda environments** - - - Check Conda-style installation prefixes (``$CONDA_PREFIX/bin`` on - Linux/Mac or ``$CONDA_PREFIX/Library/bin`` on Windows). - - 3. **CUDA Toolkit environment variables** - - - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in the - ``bin`` subdirectory. + RuntimeError: If binary_name is not supported. + ToolchainMismatchError: If binary found in different source than + the context's preferred source. """ if binary_name not in SUPPORTED_BINARIES: raise RuntimeError(f"UNKNOWN {binary_name=}") - # Try site-packages first - if binary_path := _find_binary_under_site_packages(binary_name): - return _abs_norm(binary_path) - - # Try conda prefix - if binary_path := _find_binary_in_conda(binary_name): - return _abs_norm(binary_path) - - # Try CUDA_HOME/CUDA_PATH - if binary_path := _find_binary_in_cuda_home(binary_name): - return _abs_norm(binary_path) + if context is None: + context = get_default_context() - return None + locations = _create_search_locations(binary_name) + path = context.find(binary_name, locations) + return _abs_norm(path) if path else None diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py index a7be94a88f..305eb7f8cf 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py @@ -4,7 +4,7 @@ import functools import glob import os -from typing import Sequence +from typing import Optional, Sequence from cuda.pathfinder._static_libs.artifact_search_config import ARTIFACT_CONFIGS from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( @@ -15,166 +15,147 @@ from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages from cuda.pathfinder._utils.path_utils import _abs_norm from cuda.pathfinder._utils.platform_aware import IS_WINDOWS +from cuda.pathfinder._utils.toolchain_tracker import ( + SearchContext, + SearchLocation, + ToolchainSource, + get_default_context, +) -def _find_artifact_under_site_packages(artifact_name: str) -> str | None: - """Search for an artifact in site-packages directories.""" - rel_dirs = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name) - if rel_dirs is None: - return None - - for rel_dir in rel_dirs: - for lib_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): - artifact_path = os.path.join(lib_dir, artifact_name) - if os.path.isfile(artifact_path): - return artifact_path - return None - - -def _search_paths( - base_dir: str, - subdirs: Sequence[str], - filenames: Sequence[str], -) -> str | None: - """Search for a file in multiple subdirectories of a base directory. +def _static_lib_filename_variants(artifact_name: str) -> Sequence[str]: + """Get filename variants for a static library artifact. Args: - base_dir: The base directory to search in. - subdirs: Subdirectories to check (relative to base_dir). - filenames: Filenames to look for in each subdirectory. + artifact_name: The artifact name to get filenames for. Returns: - First matching file path, or None if not found. + Sequence of filenames to search for. """ - for subdir in subdirs: - dir_path = os.path.join(base_dir, subdir) - for filename in filenames: - file_path = os.path.join(dir_path, filename) - if os.path.isfile(file_path): - return file_path - return None - + config = ARTIFACT_CONFIGS.get(artifact_name) + if not config: + return (artifact_name,) + return config.filenames -def _search_targets_subdirs( - cuda_home: str, - filenames: Sequence[str], -) -> str | None: - """Search in targets/*/lib and targets/*/lib64 subdirectories. - For cross-compilation setups. Returns the first match, preferring - more recently modified targets. +def _get_site_packages_subdirs(artifact_name: str) -> Sequence[str]: + """Get site-packages subdirectories for an artifact. Args: - cuda_home: The CUDA home directory. - filenames: Filenames to search for. + artifact_name: Name of the artifact. Returns: - First matching file path, or None if not found. + List of absolute paths to search in site-packages, or empty if not available. """ - for lib_subdir in ("lib64", "lib"): - pattern = os.path.join(cuda_home, "targets", "*", lib_subdir) - for dir_path in sorted(glob.glob(pattern), reverse=True): - for filename in filenames: - file_path = os.path.join(dir_path, filename) - if os.path.isfile(file_path): - return file_path - return None + rel_dirs = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name) + if not rel_dirs: + return [] + subdirs = [] + for rel_dir in rel_dirs: + for found_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): + subdirs.append(found_dir) + return subdirs -def _find_artifact_in_conda(artifact_name: str) -> str | None: - """Generic conda prefix search for any configured artifact. + +def _get_conda_subdirs(artifact_name: str) -> Sequence[str]: + """Get conda subdirectories for an artifact. Args: - artifact_name: The name of the artifact to find. + artifact_name: Name of the artifact. Returns: - Path to the artifact if found, None otherwise. + List of subdirectories relative to CONDA_PREFIX. """ - conda_prefix = os.environ.get("CONDA_PREFIX") - if not conda_prefix: - return None - config = ARTIFACT_CONFIGS.get(artifact_name) if not config: - return None - - return _search_paths(conda_prefix, config.conda_dirs, config.filenames) + return [] + return config.conda_dirs -def _find_artifact_in_cuda_home(artifact_name: str) -> str | None: - """Generic CUDA_HOME/CUDA_PATH search for any configured artifact. +def _get_cuda_home_subdirs(artifact_name: str) -> Sequence[str]: + """Get CUDA_HOME subdirectories for an artifact, including targets/ search if needed. Args: - artifact_name: The name of the artifact to find. + artifact_name: Name of the artifact. Returns: - Path to the artifact if found, None otherwise. + List of subdirectories to search. """ - cuda_home = get_cuda_home_or_path() - if cuda_home is None: - return None - config = ARTIFACT_CONFIGS.get(artifact_name) if not config: - return None + return [] - # Try standard directories first - if result := _search_paths(cuda_home, config.cuda_home_dirs, config.filenames): - return result + subdirs = list(config.cuda_home_dirs) - # Try targets subdirectories for cross-compilation + # Add targets/* expansion for cross-compilation (Linux only) if config.search_targets_subdirs and not IS_WINDOWS: - if result := _search_targets_subdirs(cuda_home, config.filenames): - return result + cuda_home = get_cuda_home_or_path() + if cuda_home: + for lib_subdir in ("lib64", "lib"): + pattern = os.path.join(cuda_home, "targets", "*", lib_subdir) + for target_dir in sorted(glob.glob(pattern), reverse=True): + # Make relative to cuda_home + rel_path = os.path.relpath(target_dir, cuda_home) + subdirs.append(rel_path) - return None + return subdirs -@functools.cache -def find_nvidia_static_lib(artifact_name: str) -> str | None: - """Locate a CUDA static library or artifact file. +def _create_search_locations(artifact_name: str) -> list[SearchLocation]: + """Create search location configurations for a specific artifact. Args: - artifact_name (str): The name of the artifact to find (e.g., - ``"libdevice.10.bc"``, ``"libcudadevrt.a"``). + artifact_name: Name of the artifact to search for. Returns: - str or None: Absolute path to the discovered artifact, or ``None`` if the - artifact cannot be found. - - Raises: - RuntimeError: If ``artifact_name`` is not in the supported set. - - Search order: - 1. **NVIDIA Python wheels** + List of SearchLocation objects to try. + """ + return [ + SearchLocation( + source=ToolchainSource.SITE_PACKAGES, + base_dir_func=lambda: None, # Use subdirs for full paths + subdirs=_get_site_packages_subdirs(artifact_name), + filename_variants=lambda _: _static_lib_filename_variants(artifact_name), + ), + SearchLocation( + source=ToolchainSource.CONDA, + base_dir_func=lambda: os.environ.get("CONDA_PREFIX"), + subdirs=_get_conda_subdirs(artifact_name), + filename_variants=lambda _: _static_lib_filename_variants(artifact_name), + ), + SearchLocation( + source=ToolchainSource.CUDA_HOME, + base_dir_func=get_cuda_home_or_path, + subdirs=_get_cuda_home_subdirs(artifact_name), + filename_variants=lambda _: _static_lib_filename_variants(artifact_name), + ), + ] - - Scan installed distributions (``site-packages``) for artifacts - shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvvm]``, - ``cuda-toolkit[cudart]``). - 2. **Conda environments** +@functools.cache +def find_nvidia_static_lib(artifact_name: str, *, context: SearchContext | None = None) -> str | None: + """Locate a CUDA static library or artifact file. - - Check Conda-style installation prefixes for the specific artifact - layout. + Args: + artifact_name: Name of the artifact (e.g., "libdevice.10.bc", "libcudadevrt.a"). + context: Optional SearchContext for toolchain consistency tracking. + If None, uses the default module-level context. - 3. **CUDA Toolkit environment variables** + Returns: + Absolute path to the artifact, or None if not found. - - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in - standard CUDA Toolkit directory layouts. + Raises: + RuntimeError: If artifact_name is not supported. + ToolchainMismatchError: If artifact found in different source than + the context's preferred source. """ if artifact_name not in SUPPORTED_STATIC_LIBS: raise RuntimeError(f"UNKNOWN {artifact_name=}") - # Try site-packages first - if artifact_path := _find_artifact_under_site_packages(artifact_name): - return _abs_norm(artifact_path) - - # Try conda prefix (generic, configuration-driven) - if artifact_path := _find_artifact_in_conda(artifact_name): - return _abs_norm(artifact_path) - - # Try CUDA_HOME/CUDA_PATH (generic, configuration-driven) - if artifact_path := _find_artifact_in_cuda_home(artifact_name): - return _abs_norm(artifact_path) + if context is None: + context = get_default_context() - return None + locations = _create_search_locations(artifact_name) + path = context.find(artifact_name, locations) + return _abs_norm(path) if path else None diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py new file mode 100644 index 0000000000..3556296ab2 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Toolchain consistency tracking for CUDA artifact searches.""" + +import os +from dataclasses import dataclass +from enum import Enum, auto +from typing import Callable, Optional, Sequence + + +class ToolchainSource(Enum): + """Source location of a CUDA artifact.""" + + SITE_PACKAGES = auto() + CONDA = auto() + CUDA_HOME = auto() + + +@dataclass(frozen=True) +class SearchLocation: + """Defines where and how to search for artifacts. + + Attributes: + source: Which toolchain source this represents. + base_dir_func: Function that returns the base directory to search, or None if unavailable. + subdirs: Subdirectories to check under the base (e.g., ["bin"], ["Library/bin", "bin"]). + filename_variants: Function that takes artifact name and returns possible filenames. + """ + + source: ToolchainSource + base_dir_func: Callable[[], Optional[str]] + subdirs: Sequence[str] + filename_variants: Callable[[str], Sequence[str]] + + +@dataclass(frozen=True) +class ArtifactRecord: + """Record of a found CUDA artifact.""" + + name: str + path: str + source: ToolchainSource + + +class ToolchainMismatchError(RuntimeError): + """Raised when artifacts from different sources are mixed.""" + + def __init__( + self, + artifact_name: str, + attempted_source: ToolchainSource, + preferred_source: ToolchainSource, + preferred_artifacts: list[ArtifactRecord], + ): + self.artifact_name = artifact_name + self.attempted_source = attempted_source + self.preferred_source = preferred_source + self.preferred_artifacts = preferred_artifacts + + artifact_list = ", ".join(f"'{a.name}'" for a in preferred_artifacts) + message = ( + f"Toolchain mismatch: '{artifact_name}' found in {attempted_source.name}, " + f"but using {preferred_source.name} (previous: {artifact_list})" + ) + super().__init__(message) + + +def search_location(location: SearchLocation, artifact_name: str) -> Optional[str]: + """Search for an artifact in a specific location. + + Args: + location: The search location configuration. + artifact_name: Name of the artifact to find. + + Returns: + Path to artifact if found, None otherwise. + """ + base_dir = location.base_dir_func() + if not base_dir: + return None + + filenames = location.filename_variants(artifact_name) + + for subdir in location.subdirs: + dir_path = os.path.join(base_dir, subdir) + for filename in filenames: + file_path = os.path.join(dir_path, filename) + if os.path.isfile(file_path): + return file_path + + return None + + +class SearchContext: + """Tracks toolchain consistency across artifact searches. + + This context ensures all artifacts come from the same source to prevent + version mismatches. The first artifact found establishes the preferred + source for subsequent searches. + """ + + def __init__(self): + self._artifacts: dict[str, ArtifactRecord] = {} + self._preferred_source: Optional[ToolchainSource] = None + + @property + def preferred_source(self) -> Optional[ToolchainSource]: + """The preferred toolchain source, or None if not yet determined.""" + return self._preferred_source + + def record(self, name: str, path: str, source: ToolchainSource) -> None: + """Record an artifact and enforce consistency. + + Args: + name: Artifact name. + path: Absolute path where found. + source: Source where found. + + Raises: + ToolchainMismatchError: If source conflicts with preferred source. + """ + if self._preferred_source is None: + self._preferred_source = source + elif source != self._preferred_source: + raise ToolchainMismatchError( + artifact_name=name, + attempted_source=source, + preferred_source=self._preferred_source, + preferred_artifacts=list(self._artifacts.values()), + ) + + self._artifacts[name] = ArtifactRecord(name=name, path=path, source=source) + + def find(self, artifact_name: str, locations: Sequence[SearchLocation]) -> Optional[str]: + """Search for artifact respecting toolchain consistency. + + Args: + artifact_name: Name of artifact to find. + locations: Search locations to try. + + Returns: + Path to artifact, or None if not found. + + Raises: + ToolchainMismatchError: If found in different source than preferred. + """ + # Reorder to search preferred source first + if self._preferred_source: + ordered_locations = [loc for loc in locations if loc.source == self._preferred_source] + ordered_locations += [loc for loc in locations if loc.source != self._preferred_source] + else: + ordered_locations = list(locations) + + # Try each location + for location in ordered_locations: + if path := search_location(location, artifact_name): + self.record(artifact_name, path, location.source) + return path + + return None + + +# Module-level default context +_default_context = SearchContext() + + +def get_default_context() -> SearchContext: + """Get the default module-level search context.""" + return _default_context + + +def reset_default_context() -> None: + """Reset the default context to a fresh state.""" + global _default_context + _default_context = SearchContext() diff --git a/cuda_pathfinder/tests/conftest.py b/cuda_pathfinder/tests/conftest.py index 42cff8ac52..8b28416436 100644 --- a/cuda_pathfinder/tests/conftest.py +++ b/cuda_pathfinder/tests/conftest.py @@ -29,3 +29,13 @@ def _append(message): request.config.custom_info.append(f"{request.node.name}: {message}") return _append + + +@pytest.fixture(autouse=True) +def reset_search_context(): + """Reset the default search context between tests.""" + from cuda.pathfinder._utils.toolchain_tracker import reset_default_context + + reset_default_context() + yield + reset_default_context() diff --git a/cuda_pathfinder/tests/test_toolchain_tracker.py b/cuda_pathfinder/tests/test_toolchain_tracker.py new file mode 100644 index 0000000000..a33090630d --- /dev/null +++ b/cuda_pathfinder/tests/test_toolchain_tracker.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for toolchain consistency tracking.""" + +import pytest + +from cuda.pathfinder._utils.toolchain_tracker import ( + ArtifactRecord, + SearchContext, + SearchLocation, + ToolchainMismatchError, + ToolchainSource, + get_default_context, + reset_default_context, + search_location, +) + + +def test_context_initial_state(): + """Test that a new context starts with no preferred source.""" + ctx = SearchContext() + assert ctx.preferred_source is None + + +def test_context_first_record_sets_preference(): + """Test that first recorded artifact sets the preferred source.""" + ctx = SearchContext() + ctx.record("nvcc", "/conda/bin/nvcc", ToolchainSource.CONDA) + assert ctx.preferred_source == ToolchainSource.CONDA + + +def test_context_allows_same_source(): + """Test that context allows multiple artifacts from same source.""" + ctx = SearchContext() + ctx.record("nvcc", "/conda/bin/nvcc", ToolchainSource.CONDA) + ctx.record("nvdisasm", "/conda/bin/nvdisasm", ToolchainSource.CONDA) + assert ctx.preferred_source == ToolchainSource.CONDA + + +def test_context_rejects_different_source(): + """Test that context raises exception for mixed sources.""" + ctx = SearchContext() + ctx.record("nvcc", "/conda/bin/nvcc", ToolchainSource.CONDA) + + with pytest.raises(ToolchainMismatchError) as exc_info: + ctx.record("nvdisasm", "/cuda_home/bin/nvdisasm", ToolchainSource.CUDA_HOME) + + assert exc_info.value.artifact_name == "nvdisasm" + assert exc_info.value.attempted_source == ToolchainSource.CUDA_HOME + assert exc_info.value.preferred_source == ToolchainSource.CONDA + + +def test_find_prefers_established_source(tmp_path): + """Test that find searches preferred source first.""" + ctx = SearchContext() + + # Create test directories + conda_dir = tmp_path / "conda" / "bin" + conda_dir.mkdir(parents=True) + (conda_dir / "nvcc").touch() + (conda_dir / "nvdisasm").touch() # Also add nvdisasm in conda + + cuda_home_dir = tmp_path / "cuda_home" / "bin" + cuda_home_dir.mkdir(parents=True) + (cuda_home_dir / "nvcc").touch() + + site_packages_dir = tmp_path / "site_packages" / "bin" + site_packages_dir.mkdir(parents=True) + (site_packages_dir / "nvdisasm").touch() + + # Define locations + locations = [ + SearchLocation( + source=ToolchainSource.SITE_PACKAGES, + base_dir_func=lambda: str(site_packages_dir.parent), + subdirs=["bin"], + filename_variants=lambda n: (n,), + ), + SearchLocation( + source=ToolchainSource.CONDA, + base_dir_func=lambda: str(conda_dir.parent), + subdirs=["bin"], + filename_variants=lambda n: (n,), + ), + SearchLocation( + source=ToolchainSource.CUDA_HOME, + base_dir_func=lambda: str(cuda_home_dir.parent), + subdirs=["bin"], + filename_variants=lambda n: (n,), + ), + ] + + # First find establishes CONDA preference + result = ctx.find("nvcc", locations) + assert result == str(conda_dir / "nvcc") + assert ctx.preferred_source == ToolchainSource.CONDA + + # Second find should prefer CONDA over SITE_PACKAGES (finds in conda) + result2 = ctx.find("nvdisasm", locations) + assert result2 == str(conda_dir / "nvdisasm") # Should find in conda, not site_packages + + +def test_search_location_basic(tmp_path): + """Test basic search_location functionality.""" + # Create test directory + test_dir = tmp_path / "test" / "bin" + test_dir.mkdir(parents=True) + (test_dir / "nvcc").touch() + + location = SearchLocation( + source=ToolchainSource.CONDA, + base_dir_func=lambda: str(test_dir.parent), + subdirs=["bin"], + filename_variants=lambda n: (n,), + ) + + result = search_location(location, "nvcc") + assert result == str(test_dir / "nvcc") + + +def test_search_location_not_found(tmp_path): + """Test search_location returns None when file not found.""" + location = SearchLocation( + source=ToolchainSource.CONDA, + base_dir_func=lambda: str(tmp_path), + subdirs=["bin"], + filename_variants=lambda n: (n,), + ) + + result = search_location(location, "nvcc") + assert result is None + + +def test_default_context_singleton(): + """Test that get_default_context returns same instance.""" + ctx1 = get_default_context() + ctx2 = get_default_context() + assert ctx1 is ctx2 + + +def test_reset_default_context(): + """Test that reset creates a new default context.""" + ctx1 = get_default_context() + ctx1.record("nvcc", "/conda/bin/nvcc", ToolchainSource.CONDA) + + reset_default_context() + + ctx2 = get_default_context() + assert ctx2.preferred_source is None + # Should be a fresh instance + assert len(ctx2._artifacts) == 0 From 07b31045900496e8a2e836fd7b8efc964369b3e5 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Sat, 10 Jan 2026 13:34:09 -0800 Subject: [PATCH 07/11] renaming --- .../_binaries/find_nvidia_binaries.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py index 74c8ed875b..d59f89a3d3 100644 --- a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py @@ -30,24 +30,24 @@ def _binary_filename_variants(name: str) -> Sequence[str]: def _get_site_packages_subdirs(binary_name: str) -> Sequence[str]: - """Get site-packages subdirectories for a binary. + """Get site-packages sub-directories for a binary. Args: binary_name: Name of the binary. Returns: - List of subdirectories to search, or empty list if binary not in site-packages. + List of sub-directories to search, or empty list if binary not in site-packages. """ - rel_dirs = SITE_PACKAGES_BINDIRS.get(binary_name) - if not rel_dirs: + relative_directories = SITE_PACKAGES_BINDIRS.get(binary_name) + if not relative_directories: return [] # Expand site-packages paths - subdirs = [] - for rel_dir in rel_dirs: - for found_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): - subdirs.append(found_dir) - return subdirs + sub_directories = [] + for relative_directory in relative_directories: + for found_dir in find_sub_dirs_all_sitepackages(tuple(relative_directory.split("/"))): + sub_directories.append(found_dir) + return sub_directories # Define search locations for binaries From 2617324b0e6556053f1ab61ea37b543edb9e220c Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Mon, 12 Jan 2026 09:52:29 -0800 Subject: [PATCH 08/11] wip --- .../_static_libs/artifact_search_config.py | 74 --------- .../_static_libs/find_nvidia_static_libs.py | 147 +++++++++--------- .../supported_nvidia_static_libs.py | 12 +- .../pathfinder/_utils/toolchain_tracker.py | 2 - .../tests/test_find_nvidia_static_libs.py | 6 +- 5 files changed, 83 insertions(+), 158 deletions(-) delete mode 100644 cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py deleted file mode 100644 index 43fefa0e8e..0000000000 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/artifact_search_config.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Configuration for artifact search patterns across platforms and environments.""" - -from dataclasses import dataclass - -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS - - -@dataclass(frozen=True) -class ArtifactSearchConfig: - """Configuration for searching a specific artifact (platform-resolved at init). - - Attributes: - canonical_name: Platform-agnostic identifier for the artifact. - filenames: Filenames to search for on the current platform. - conda_dirs: Directory paths relative to CONDA_PREFIX. - cuda_home_dirs: Directory paths relative to CUDA_HOME. - search_targets_subdirs: Whether to search targets/*/{lib,lib64} subdirs - for cross-compilation setups (Linux only). - """ - - canonical_name: str - filenames: tuple[str, ...] - conda_dirs: tuple[str, ...] - cuda_home_dirs: tuple[str, ...] - search_targets_subdirs: bool = False - - -def _create_config( - canonical_name: str, - linux_filenames: tuple[str, ...], - windows_filenames: tuple[str, ...], - conda_linux_dirs: tuple[str, ...], - conda_windows_dirs: tuple[str, ...], - cuda_home_linux_dirs: tuple[str, ...], - cuda_home_windows_dirs: tuple[str, ...], - search_targets_subdirs: bool = False, -) -> ArtifactSearchConfig: - """Create a platform-specific config by selecting appropriate values at init time.""" - return ArtifactSearchConfig( - canonical_name=canonical_name, - filenames=windows_filenames if IS_WINDOWS else linux_filenames, - conda_dirs=conda_windows_dirs if IS_WINDOWS else conda_linux_dirs, - cuda_home_dirs=cuda_home_windows_dirs if IS_WINDOWS else cuda_home_linux_dirs, - search_targets_subdirs=search_targets_subdirs, - ) - - -# Registry of all supported artifacts with their search configurations -# Platform selection happens once at module import time via _create_config -ARTIFACT_CONFIGS = { - "libcudadevrt.a": _create_config( - canonical_name="cudadevrt", - linux_filenames=("libcudadevrt.a",), - windows_filenames=("cudadevrt.lib",), - conda_linux_dirs=("lib",), - conda_windows_dirs=("Library/lib", "Library/lib/x64"), - cuda_home_linux_dirs=("lib64", "lib"), - cuda_home_windows_dirs=("lib/x64", "lib"), - search_targets_subdirs=True, - ), - "libdevice.10.bc": _create_config( - canonical_name="libdevice", - linux_filenames=("libdevice.10.bc",), - windows_filenames=("libdevice.10.bc",), - conda_linux_dirs=("nvvm/libdevice",), - conda_windows_dirs=("Library/nvvm/libdevice",), - cuda_home_linux_dirs=("nvvm/libdevice",), - cuda_home_windows_dirs=("nvvm/libdevice",), - search_targets_subdirs=False, - ), -} diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py index 305eb7f8cf..a6bd24c1d2 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py @@ -2,11 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 import functools -import glob import os -from typing import Optional, Sequence +from typing import Sequence -from cuda.pathfinder._static_libs.artifact_search_config import ARTIFACT_CONFIGS from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( SITE_PACKAGES_STATIC_LIBDIRS, SUPPORTED_STATIC_LIBS, @@ -23,73 +21,59 @@ ) -def _static_lib_filename_variants(artifact_name: str) -> Sequence[str]: - """Get filename variants for a static library artifact. - - Args: - artifact_name: The artifact name to get filenames for. +# Generic search locations by toolchain source (platform-specific at import time) +if IS_WINDOWS: + CONDA_LIB_SUBDIRS = ("Library/lib", "Library/lib/x64") + CONDA_NVVM_SUBDIRS = ("Library/nvvm/libdevice",) + CUDA_HOME_LIB_SUBDIRS = ("lib/x64", "lib") + CUDA_HOME_NVVM_SUBDIRS = ("nvvm/libdevice",) +else: + CONDA_LIB_SUBDIRS = ("lib",) + CONDA_NVVM_SUBDIRS = ("nvvm/libdevice",) + CUDA_HOME_LIB_SUBDIRS = ("lib64", "lib") + CUDA_HOME_NVVM_SUBDIRS = ("nvvm/libdevice",) - Returns: - Sequence of filenames to search for. - """ - config = ARTIFACT_CONFIGS.get(artifact_name) - if not config: - return (artifact_name,) - return config.filenames - -def _get_site_packages_subdirs(artifact_name: str) -> Sequence[str]: - """Get site-packages subdirectories for an artifact. - - Args: - artifact_name: Name of the artifact. - - Returns: - List of absolute paths to search in site-packages, or empty if not available. - """ - rel_dirs = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name) - if not rel_dirs: - return [] - - subdirs = [] - for rel_dir in rel_dirs: - for found_dir in find_sub_dirs_all_sitepackages(tuple(rel_dir.split("/"))): - subdirs.append(found_dir) - return subdirs - - -def _get_conda_subdirs(artifact_name: str) -> Sequence[str]: - """Get conda subdirectories for an artifact. +def _static_lib_filename_variants(artifact_name: str) -> Sequence[str]: + """Generate platform-appropriate filename variants for an artifact. Args: - artifact_name: Name of the artifact. + artifact_name: Canonical artifact name (e.g., "cudadevrt", "libdevice.10.bc"). Returns: - List of subdirectories relative to CONDA_PREFIX. + Sequence of filenames to search for on this platform. + + Examples: + On Windows: + "cudadevrt" -> ("cudadevrt.lib",) + "libdevice.10.bc" -> ("libdevice.10.bc",) + On Linux: + "cudadevrt" -> ("libcudadevrt.a",) + "libdevice.10.bc" -> ("libdevice.10.bc",) """ - config = ARTIFACT_CONFIGS.get(artifact_name) - if not config: - return [] - return config.conda_dirs + # Files that are the same on all platforms (e.g., .bc bitcode files) + if "." in artifact_name: + return (artifact_name,) + # Platform-specific library naming conventions + if IS_WINDOWS: + return (f"{artifact_name}.lib",) + else: + return (f"lib{artifact_name}.a",) -def _get_cuda_home_subdirs(artifact_name: str) -> Sequence[str]: - """Get CUDA_HOME subdirectories for an artifact, including targets/ search if needed. - Args: - artifact_name: Name of the artifact. +def _get_cuda_home_subdirs_with_targets() -> tuple[str, ...]: + """Get CUDA_HOME subdirectories including expanded targets/* paths. Returns: - List of subdirectories to search. + Tuple of subdirectories to search under CUDA_HOME. """ - config = ARTIFACT_CONFIGS.get(artifact_name) - if not config: - return [] + import glob - subdirs = list(config.cuda_home_dirs) + subdirs = list(CUDA_HOME_LIB_SUBDIRS + CUDA_HOME_NVVM_SUBDIRS) - # Add targets/* expansion for cross-compilation (Linux only) - if config.search_targets_subdirs and not IS_WINDOWS: + # On Linux, also search targets/*/lib64 and targets/*/lib for cross-compilation + if not IS_WINDOWS: cuda_home = get_cuda_home_or_path() if cuda_home: for lib_subdir in ("lib64", "lib"): @@ -99,11 +83,11 @@ def _get_cuda_home_subdirs(artifact_name: str) -> Sequence[str]: rel_path = os.path.relpath(target_dir, cuda_home) subdirs.append(rel_path) - return subdirs + return tuple(subdirs) def _create_search_locations(artifact_name: str) -> list[SearchLocation]: - """Create search location configurations for a specific artifact. + """Create generic search location configurations. Args: artifact_name: Name of the artifact to search for. @@ -111,26 +95,43 @@ def _create_search_locations(artifact_name: str) -> list[SearchLocation]: Returns: List of SearchLocation objects to try. """ - return [ - SearchLocation( - source=ToolchainSource.SITE_PACKAGES, - base_dir_func=lambda: None, # Use subdirs for full paths - subdirs=_get_site_packages_subdirs(artifact_name), - filename_variants=lambda _: _static_lib_filename_variants(artifact_name), - ), + locations = [] + + # Site-packages: Create separate SearchLocation for each found directory + relative_directories = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name) + if relative_directories: + for relative_directory in relative_directories: + for found_dir in find_sub_dirs_all_sitepackages(tuple(relative_directory.split("/"))): + locations.append( + SearchLocation( + source=ToolchainSource.SITE_PACKAGES, + base_dir_func=lambda d=found_dir: d, + subdirs=[""], + filename_variants=_static_lib_filename_variants, + ) + ) + + # Conda: Generic lib and nvvm locations + locations.append( SearchLocation( source=ToolchainSource.CONDA, base_dir_func=lambda: os.environ.get("CONDA_PREFIX"), - subdirs=_get_conda_subdirs(artifact_name), - filename_variants=lambda _: _static_lib_filename_variants(artifact_name), - ), + subdirs=CONDA_LIB_SUBDIRS + CONDA_NVVM_SUBDIRS, + filename_variants=_static_lib_filename_variants, + ) + ) + + # CUDA_HOME: Generic lib and nvvm locations (including targets/* on Linux) + locations.append( SearchLocation( source=ToolchainSource.CUDA_HOME, base_dir_func=get_cuda_home_or_path, - subdirs=_get_cuda_home_subdirs(artifact_name), - filename_variants=lambda _: _static_lib_filename_variants(artifact_name), - ), - ] + subdirs=_get_cuda_home_subdirs_with_targets(), + filename_variants=_static_lib_filename_variants, + ) + ) + + return locations @functools.cache @@ -138,7 +139,9 @@ def find_nvidia_static_lib(artifact_name: str, *, context: SearchContext | None """Locate a CUDA static library or artifact file. Args: - artifact_name: Name of the artifact (e.g., "libdevice.10.bc", "libcudadevrt.a"). + artifact_name: Canonical artifact name (e.g., "libdevice.10.bc", "cudadevrt"). + Platform-specific filenames are resolved automatically: + - "cudadevrt" -> "libcudadevrt.a" on Linux, "cudadevrt.lib" on Windows context: Optional SearchContext for toolchain consistency tracking. If None, uses the default module-level context. diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py index 0ad2e83e58..f0b7a62bfb 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/supported_nvidia_static_libs.py @@ -6,20 +6,18 @@ # SUPPORTED_STATIC_LIBS # SITE_PACKAGES_STATIC_LIBDIRS -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS - -# Supported CUDA static libraries and artifacts that can be found +# Supported CUDA static libraries and artifacts (canonical names) SUPPORTED_STATIC_LIBS_COMMON = ( - "libdevice.10.bc", - "libcudadevrt.a", + "libdevice.10.bc", # Bitcode file (same name on all platforms) + "cudadevrt", # Static device runtime library (libcudadevrt.a on Linux, cudadevrt.lib on Windows) ) SUPPORTED_STATIC_LIBS = SUPPORTED_STATIC_LIBS_COMMON -# Map from artifact name to relative paths under site-packages +# Map from canonical artifact name to relative paths under site-packages SITE_PACKAGES_STATIC_LIBDIRS = { "libdevice.10.bc": ["nvidia/cuda_nvvm/nvvm/libdevice"], - "libcudadevrt.a": [ + "cudadevrt": [ "nvidia/cuda_cudart/lib", # Linux "nvidia/cuda_cudart/lib/x64", # Windows (if present) ], diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py index 3556296ab2..4a9c6f191e 100644 --- a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py +++ b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Toolchain consistency tracking for CUDA artifact searches.""" - import os from dataclasses import dataclass from enum import Enum, auto diff --git a/cuda_pathfinder/tests/test_find_nvidia_static_libs.py b/cuda_pathfinder/tests/test_find_nvidia_static_libs.py index 9d6c14699f..099ce17512 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_static_libs.py +++ b/cuda_pathfinder/tests/test_find_nvidia_static_libs.py @@ -46,9 +46,9 @@ def test_libdevice_specific(info_summary_append): def test_libcudadevrt_specific(info_summary_append): - """Specific test for libcudadevrt.a to ensure it's working.""" - artifact_path = find_nvidia_static_lib("libcudadevrt.a") - info_summary_append(f"libcudadevrt.a path: {artifact_path!r}") + """Specific test for cudadevrt to ensure it's working.""" + artifact_path = find_nvidia_static_lib("cudadevrt") + info_summary_append(f"cudadevrt path: {artifact_path!r}") if artifact_path: assert os.path.isfile(artifact_path) # On Linux it should be .a, on Windows it might be .lib From a688ea093e2fdf66a6c7e099776b92ccc4467c5c Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Mon, 12 Jan 2026 10:12:44 -0800 Subject: [PATCH 09/11] wip --- .../pathfinder/_utils/toolchain_tracker.py | 9 ++++++--- .../tests/test_find_nvidia_binaries.py | 17 ++++------------- .../tests/test_find_nvidia_headers.py | 18 ------------------ .../tests/test_find_nvidia_static_libs.py | 13 ------------- .../tests/test_load_nvidia_dynamic_lib.py | 5 ----- .../tests/test_toolchain_tracker.py | 7 ------- 6 files changed, 10 insertions(+), 59 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py index 4a9c6f191e..d4f91ff1ab 100644 --- a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py +++ b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py @@ -143,10 +143,13 @@ def find(self, artifact_name: str, locations: Sequence[SearchLocation]) -> Optio Raises: ToolchainMismatchError: If found in different source than preferred. """ - # Reorder to search preferred source first + # Reorder to search preferred source first, maintaining original order for ties + # (stable sort: preferred source gets priority 0, others get priority 1) if self._preferred_source: - ordered_locations = [loc for loc in locations if loc.source == self._preferred_source] - ordered_locations += [loc for loc in locations if loc.source != self._preferred_source] + ordered_locations = sorted( + locations, + key=lambda loc: 0 if loc.source == self._preferred_source else 1 + ) else: ordered_locations = list(locations) diff --git a/cuda_pathfinder/tests/test_find_nvidia_binaries.py b/cuda_pathfinder/tests/test_find_nvidia_binaries.py index 15ba94c361..f12ba9ee99 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_binaries.py +++ b/cuda_pathfinder/tests/test_find_nvidia_binaries.py @@ -8,9 +8,6 @@ from cuda.pathfinder import find_nvidia_binary from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES -STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_BINARIES_STRICTNESS", "see_what_works") -assert STRICTNESS in ("see_what_works", "all_must_work") - def test_unknown_binary(): with pytest.raises(RuntimeError, match=r"^UNKNOWN binary_name='unknown-binary'$"): @@ -25,25 +22,19 @@ def test_find_binaries(info_summary_append, binary_name): assert os.path.isfile(binary_path) # Verify the binary name is in the path assert binary_name in os.path.basename(binary_path) - if STRICTNESS == "all_must_work": - assert binary_path is not None def test_nvdisasm_specific(info_summary_append): """Specific test for nvdisasm to ensure it's working.""" binary_path = find_nvidia_binary("nvdisasm") info_summary_append(f"nvdisasm path: {binary_path!r}") - # Only assert if we're in strict mode or if cuda-toolkit is installed - if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): - if binary_path: - assert os.path.isfile(binary_path) + if binary_path: + assert os.path.isfile(binary_path) def test_cuobjdump_specific(info_summary_append): """Specific test for cuobjdump to ensure it's working.""" binary_path = find_nvidia_binary("cuobjdump") info_summary_append(f"cuobjdump path: {binary_path!r}") - # Only assert if we're in strict mode or if cuda-toolkit is installed - if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): - if binary_path: - assert os.path.isfile(binary_path) + if binary_path: + assert os.path.isfile(binary_path) diff --git a/cuda_pathfinder/tests/test_find_nvidia_headers.py b/cuda_pathfinder/tests/test_find_nvidia_headers.py index 494d7c0ae9..2c1f2b51d2 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/cuda_pathfinder/tests/test_find_nvidia_headers.py @@ -29,9 +29,6 @@ SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK, ) -STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS", "see_what_works") -assert STRICTNESS in ("see_what_works", "all_must_work") - NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = { "cusparseLt": r"^nvidia-cusparselt-.*$", "cutensor": r"^cutensor-.*$", @@ -68,19 +65,6 @@ def test_find_non_ctk_headers(info_summary_append, libname): assert hdr_dir is not None hdr_dir_parts = hdr_dir.split(os.path.sep) assert "site-packages" in hdr_dir_parts - elif STRICTNESS == "all_must_work": - assert hdr_dir is not None - if conda_prefix := os.environ.get("CONDA_PREFIX"): - assert hdr_dir.startswith(conda_prefix) - else: - inst_dirs = SUPPORTED_INSTALL_DIRS_NON_CTK.get(libname) - if inst_dirs is not None: - for inst_dir in inst_dirs: - globbed = glob.glob(inst_dir) - if hdr_dir in globbed: - break - else: - raise RuntimeError(f"{hdr_dir=} does not match any {inst_dirs=}") def test_supported_headers_site_packages_ctk_consistency(): @@ -95,5 +79,3 @@ def test_find_ctk_headers(info_summary_append, libname): assert os.path.isdir(hdr_dir) h_filename = SUPPORTED_HEADERS_CTK[libname] assert os.path.isfile(os.path.join(hdr_dir, h_filename)) - if STRICTNESS == "all_must_work": - assert hdr_dir is not None diff --git a/cuda_pathfinder/tests/test_find_nvidia_static_libs.py b/cuda_pathfinder/tests/test_find_nvidia_static_libs.py index 099ce17512..98ba50bdaa 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_static_libs.py +++ b/cuda_pathfinder/tests/test_find_nvidia_static_libs.py @@ -8,9 +8,6 @@ from cuda.pathfinder import find_nvidia_static_lib from cuda.pathfinder._static_libs.supported_nvidia_static_libs import SUPPORTED_STATIC_LIBS -STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_STATIC_LIBS_STRICTNESS", "see_what_works") -assert STRICTNESS in ("see_what_works", "all_must_work") - def test_unknown_artifact(): with pytest.raises(RuntimeError, match=r"^UNKNOWN artifact_name='unknown-artifact'$"): @@ -26,8 +23,6 @@ def test_find_static_libs(info_summary_append, artifact_name): # Verify the artifact name (or its base) is in the path base_name = artifact_name.replace(".10", "") # Handle libdevice.10.bc -> libdevice assert base_name.split(".")[0] in artifact_path.lower() - if STRICTNESS == "all_must_work": - assert artifact_path is not None def test_libdevice_specific(info_summary_append): @@ -39,10 +34,6 @@ def test_libdevice_specific(info_summary_append): assert "libdevice" in artifact_path # Should end with .bc assert artifact_path.endswith(".bc") - # Only assert existence if we're in strict mode or if cuda is installed - if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): - if artifact_path: - assert os.path.isfile(artifact_path) def test_libcudadevrt_specific(info_summary_append): @@ -54,10 +45,6 @@ def test_libcudadevrt_specific(info_summary_append): # On Linux it should be .a, on Windows it might be .lib assert artifact_path.endswith((".a", ".lib")) assert "cudadevrt" in artifact_path.lower() - # Only assert existence if we're in strict mode or if cuda is installed - if STRICTNESS == "all_must_work" or os.environ.get("CUDA_HOME") or os.environ.get("CONDA_PREFIX"): - if artifact_path: - assert os.path.isfile(artifact_path) def test_caching(): diff --git a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py index cff5b74290..bcfe4aaaf1 100644 --- a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py @@ -15,9 +15,6 @@ from cuda.pathfinder._dynamic_libs import supported_nvidia_libs from cuda.pathfinder._utils.platform_aware import IS_WINDOWS, quote_for_shell -STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS", "see_what_works") -assert STRICTNESS in ("see_what_works", "all_must_work") - def test_supported_libnames_linux_sonames_consistency(): assert tuple(sorted(supported_nvidia_libs.SUPPORTED_LIBNAMES_LINUX)) == tuple( @@ -104,8 +101,6 @@ def raise_child_process_failed(): raise_child_process_failed() assert not result.stderr if result.stdout.startswith("CHILD_LOAD_NVIDIA_DYNAMIC_LIB_HELPER_DYNAMIC_LIB_NOT_FOUND_ERROR:"): - if STRICTNESS == "all_must_work" and not _is_expected_load_nvidia_dynamic_lib_failure(libname): - raise_child_process_failed() info_summary_append(f"Not found: {libname=!r}") else: abs_path = json.loads(result.stdout.rstrip()) diff --git a/cuda_pathfinder/tests/test_toolchain_tracker.py b/cuda_pathfinder/tests/test_toolchain_tracker.py index a33090630d..01600beb8c 100644 --- a/cuda_pathfinder/tests/test_toolchain_tracker.py +++ b/cuda_pathfinder/tests/test_toolchain_tracker.py @@ -132,13 +132,6 @@ def test_search_location_not_found(tmp_path): assert result is None -def test_default_context_singleton(): - """Test that get_default_context returns same instance.""" - ctx1 = get_default_context() - ctx2 = get_default_context() - assert ctx1 is ctx2 - - def test_reset_default_context(): """Test that reset creates a new default context.""" ctx1 = get_default_context() From d36cc8b025735a65cee178ff3db6639aa35bab79 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Tue, 13 Jan 2026 17:40:28 -0800 Subject: [PATCH 10/11] pixi.lock updates --- cuda_bindings/pixi.lock | 80 ++++++++++++++++++++--------------------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/cuda_bindings/pixi.lock b/cuda_bindings/pixi.lock index f84d569dff..fb3d0ad393 100644 --- a/cuda_bindings/pixi.lock +++ b/cuda_bindings/pixi.lock @@ -26,7 +26,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-cudart-static-13.1.80-hecca717_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart-static_linux-64-13.1.80-h376f20c_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart_linux-64-13.1.80-h376f20c_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvrtc-13.1.80-hecca717_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvrtc-13.1.115-hecca717_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-13.1.80-h69a702a_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-nvvm-dev_linux-64-13.1.80-ha770c72_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-impl-13.1.80-h4bc722e_0.conda @@ -72,7 +72,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-5_h4a7cf45_openblas.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libcap-2.77-h3ff7636_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-5_h0358290_openblas.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/libcufile-1.16.0.49-hd07211c_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/libcufile-1.16.1.26-hd07211c_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.25-h17f619e_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libdrm-2.4.125-hb03c661_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libegl-1.7.0-ha4b6fd6_2.conda @@ -220,7 +220,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-cudart-static-13.1.80-h8f3c8d4_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart-static_linux-aarch64-13.1.80-h8f3c8d4_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart_linux-aarch64-13.1.80-h8f3c8d4_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvrtc-13.1.80-h8f3c8d4_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvrtc-13.1.115-h8f3c8d4_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvvm-13.1.80-he9431aa_100.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-nvvm-dev_linux-aarch64-13.1.80-h579c4fd_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvvm-impl-13.1.80-h7b14b0b_0.conda @@ -263,7 +263,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libblas-3.11.0-5_haddc8a3_openblas.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcap-2.77-h68e9139_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcblas-3.11.0-5_hd72aa62_openblas.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcufile-1.16.0.49-hbf501ad_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcufile-1.16.1.26-hbf501ad_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libdeflate-1.25-h1af38f5_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libdrm-2.4.125-he30d5cf_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libegl-1.7.0-hd24410f_2.conda @@ -401,7 +401,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-cudart-static-13.1.80-hac47afa_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart-static_win-64-13.1.80-hac47afa_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart_win-64-13.1.80-hac47afa_0.conda - - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvrtc-13.1.80-hac47afa_0.conda + - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvrtc-13.1.115-hac47afa_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvvm-13.1.80-h719f0c7_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-nvvm-dev_win-64-13.1.80-h57928b3_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvvm-impl-13.1.80-h2466b09_0.conda @@ -542,7 +542,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-cudart-static-13.1.80-hecca717_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart-static_linux-64-13.1.80-h376f20c_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart_linux-64-13.1.80-h376f20c_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvrtc-13.1.80-hecca717_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvrtc-13.1.115-hecca717_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-13.1.80-h69a702a_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-nvvm-dev_linux-64-13.1.80-ha770c72_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-impl-13.1.80-h4bc722e_0.conda @@ -588,7 +588,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/libblas-3.11.0-5_h4a7cf45_openblas.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libcap-2.77-h3ff7636_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libcblas-3.11.0-5_h0358290_openblas.conda - - conda: https://conda.anaconda.org/conda-forge/linux-64/libcufile-1.16.0.49-hd07211c_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/libcufile-1.16.1.26-hd07211c_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.25-h17f619e_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libdrm-2.4.125-hb03c661_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libegl-1.7.0-ha4b6fd6_2.conda @@ -736,7 +736,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-cudart-static-13.1.80-h8f3c8d4_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart-static_linux-aarch64-13.1.80-h8f3c8d4_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart_linux-aarch64-13.1.80-h8f3c8d4_0.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvrtc-13.1.80-h8f3c8d4_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvrtc-13.1.115-h8f3c8d4_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvvm-13.1.80-he9431aa_100.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-nvvm-dev_linux-aarch64-13.1.80-h579c4fd_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvvm-impl-13.1.80-h7b14b0b_0.conda @@ -779,7 +779,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libblas-3.11.0-5_haddc8a3_openblas.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcap-2.77-h68e9139_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcblas-3.11.0-5_hd72aa62_openblas.conda - - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcufile-1.16.0.49-hbf501ad_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcufile-1.16.1.26-hbf501ad_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libdeflate-1.25-h1af38f5_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libdrm-2.4.125-he30d5cf_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libegl-1.7.0-hd24410f_2.conda @@ -917,7 +917,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-cudart-static-13.1.80-hac47afa_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart-static_win-64-13.1.80-hac47afa_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-cudart_win-64-13.1.80-hac47afa_0.conda - - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvrtc-13.1.80-hac47afa_0.conda + - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvrtc-13.1.115-hac47afa_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvvm-13.1.80-h719f0c7_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-nvvm-dev_win-64-13.1.80-h57928b3_0.conda - conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvvm-impl-13.1.80-h2466b09_0.conda @@ -1461,10 +1461,10 @@ packages: - cuda-pathfinder >=1.1,<2 - libnvjitlink - cuda-nvrtc - - cuda-nvrtc >=13.1.80,<14.0a0 + - cuda-nvrtc >=13.1.115,<14.0a0 - cuda-nvvm - libcufile - - libcufile >=1.16.0.49,<2.0a0 + - libcufile >=1.16.1.26,<2.0a0 - libgcc >=15 - libgcc >=15 - libstdcxx >=15 @@ -1483,7 +1483,7 @@ packages: - cuda-pathfinder >=1.1,<2 - libnvjitlink - cuda-nvrtc - - cuda-nvrtc >=13.1.80,<14.0a0 + - cuda-nvrtc >=13.1.115,<14.0a0 - cuda-nvvm - vc >=14.1,<15 - vc14_runtime >=14.16.27033 @@ -1502,10 +1502,10 @@ packages: - cuda-pathfinder >=1.1,<2 - libnvjitlink - cuda-nvrtc - - cuda-nvrtc >=13.1.80,<14.0a0 + - cuda-nvrtc >=13.1.115,<14.0a0 - cuda-nvvm - libcufile - - libcufile >=1.16.0.49,<2.0a0 + - libcufile >=1.16.1.26,<2.0a0 - libgcc >=15 - libgcc >=15 - libstdcxx >=15 @@ -1759,39 +1759,39 @@ packages: license: LicenseRef-NVIDIA-End-User-License-Agreement size: 24082 timestamp: 1764883821516 -- conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvrtc-13.1.80-hecca717_0.conda - sha256: d6b326bdbf6fa7bfa0fa617dda547dc585159816b8f130f2535740c4e53fd12c - md5: 7ef874b2dc4ca388ecef3b3893305459 +- conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvrtc-13.1.115-hecca717_0.conda + sha256: 9cc4f9df70c02eea5121cdb0e865207b04cd52591f57ebcac2ba44fada10eb5b + md5: df16c9049d882cdaf4f83a5b90079589 depends: - __glibc >=2.17,<3.0.a0 - cuda-version >=13.1,<13.2.0a0 - libgcc >=14 - libstdcxx >=14 license: LicenseRef-NVIDIA-End-User-License-Agreement - size: 35479197 - timestamp: 1764880529154 -- conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvrtc-13.1.80-h8f3c8d4_0.conda - sha256: 5e10ce4dd84c22c73e58a9f8359fb1e5ef4596afd3a0bc12b9fbde73b388ec0d - md5: 0473ebdb01f2f4024177b024fc19fa72 + size: 35339417 + timestamp: 1768272955912 +- conda: https://conda.anaconda.org/conda-forge/linux-aarch64/cuda-nvrtc-13.1.115-h8f3c8d4_0.conda + sha256: a1ec61512cecb093797e00590ad381ecd5852d2a32440ff22b34f78c743f3d5a + md5: 34da2ff2c64054d65eb8f04d76c40cca depends: - arm-variant * sbsa - cuda-version >=13.1,<13.2.0a0 - libgcc >=14 - libstdcxx >=14 license: LicenseRef-NVIDIA-End-User-License-Agreement - size: 33619044 - timestamp: 1764880672755 -- conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvrtc-13.1.80-hac47afa_0.conda - sha256: 3f67de8a9eb182fa20bbc80bda7185afb676cfe8894f6a0549173bd752a7d2f4 - md5: 7b42337a35cd887ec3eed254b5ed606f + size: 33616576 + timestamp: 1768272976976 +- conda: https://conda.anaconda.org/conda-forge/win-64/cuda-nvrtc-13.1.115-hac47afa_0.conda + sha256: a8869b7d997722f90b9f8a602dc0b1d0d497f2a6f3561dc89383aeb2cd379a66 + md5: 372d3c612a832d5f87d8dd9702d487b2 depends: - cuda-version >=13.1,<13.2.0a0 - ucrt >=10.0.20348.0 - vc >=14.3,<15 - vc14_runtime >=14.44.35208 license: LicenseRef-NVIDIA-End-User-License-Agreement - size: 31012754 - timestamp: 1764880740086 + size: 31006920 + timestamp: 1768273107962 - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-nvvm-13.1.80-h69a702a_0.conda sha256: 84f971ab146e2c822103cfe06f478ece244747a6f2aa565be639a4709d0a1579 md5: 9250c651d8758c8f665dff7519ef21ff @@ -3275,9 +3275,9 @@ packages: license_family: BSD size: 68079 timestamp: 1765819124349 -- conda: https://conda.anaconda.org/conda-forge/linux-64/libcufile-1.16.0.49-hd07211c_0.conda - sha256: 6aabad84132b1f3ee367e5d24291febf8a11d9a7f3967a64fc07e77d9b0b22df - md5: 9cb68a85f8c08f0512931f944f6a75df +- conda: https://conda.anaconda.org/conda-forge/linux-64/libcufile-1.16.1.26-hd07211c_0.conda + sha256: 8c44b5bf947afad827df0df49fe7483cf1b2916694081b2db4fecdfd6a2bacd1 + md5: 48418c48dac04671fa46cb446122b8a5 depends: - __glibc >=2.28,<3.0.a0 - cuda-version >=13.1,<13.2.0a0 @@ -3285,11 +3285,11 @@ packages: - libstdcxx >=14 - rdma-core >=60.0 license: LicenseRef-NVIDIA-End-User-License-Agreement - size: 990030 - timestamp: 1764881892686 -- conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcufile-1.16.0.49-hbf501ad_0.conda - sha256: d03963dc7708ded20340176ade987fc4c3e49da4f7b139a85e69ca7eb413f57a - md5: 315e1b144eaf890519fc63049b6e9228 + size: 990938 + timestamp: 1768273732081 +- conda: https://conda.anaconda.org/conda-forge/linux-aarch64/libcufile-1.16.1.26-hbf501ad_0.conda + sha256: 7451b3e2204e6cad21db501052dfe595c3440213ef3e22c0f9c784012f6a8419 + md5: ee60a24c702ce02de95ae1982c4841d8 depends: - __glibc >=2.28,<3.0.a0 - arm-variant * sbsa @@ -3300,8 +3300,8 @@ packages: constrains: - arm-variant * sbsa license: LicenseRef-NVIDIA-End-User-License-Agreement - size: 887547 - timestamp: 1764881951574 + size: 891752 + timestamp: 1768273724252 - conda: https://conda.anaconda.org/conda-forge/linux-64/libdeflate-1.25-h17f619e_0.conda sha256: aa8e8c4be9a2e81610ddf574e05b64ee131fab5e0e3693210c9d6d2fba32c680 md5: 6c77a605a7a689d17d4819c0f8ac9a00 From 9ee4248be8f3cde22bd0c2f2f26e327875796472 Mon Sep 17 00:00:00 2001 From: Rob Parolin Date: Thu, 15 Jan 2026 00:21:03 -0800 Subject: [PATCH 11/11] wip --- .../_binaries/find_nvidia_binaries.py | 160 ++++++------- .../_static_libs/find_nvidia_static_libs.py | 215 ++++++++---------- .../pathfinder/_utils/filename_resolver.py | 82 +++++++ .../cuda/pathfinder/_utils/platform_paths.py | 57 +++++ .../cuda/pathfinder/_utils/search_factory.py | 172 ++++++++++++++ .../pathfinder/_utils/toolchain_tracker.py | 110 +++++++-- .../cuda/pathfinder/_utils/types.py | 12 + cuda_pathfinder/tests/test_cache_behavior.py | 89 ++++++++ .../tests/test_filename_resolver.py | 54 +++++ cuda_pathfinder/tests/test_platform_paths.py | 56 +++++ .../tests/test_toolchain_tracker.py | 4 +- 11 files changed, 787 insertions(+), 224 deletions(-) create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/filename_resolver.py create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/platform_paths.py create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/search_factory.py create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/types.py create mode 100644 cuda_pathfinder/tests/test_cache_behavior.py create mode 100644 cuda_pathfinder/tests/test_filename_resolver.py create mode 100644 cuda_pathfinder/tests/test_platform_paths.py diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py index d59f89a3d3..4992f3e956 100644 --- a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py +++ b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binaries.py @@ -1,110 +1,116 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Find CUDA binary executables across different installation sources.""" + import functools -import os -from typing import Sequence +from typing import Optional from cuda.pathfinder._binaries.supported_nvidia_binaries import SITE_PACKAGES_BINDIRS, SUPPORTED_BINARIES -from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path -from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.filename_resolver import FilenameResolver from cuda.pathfinder._utils.path_utils import _abs_norm -from cuda.pathfinder._utils.toolchain_tracker import ( - SearchContext, - SearchLocation, - ToolchainSource, - get_default_context, -) - +from cuda.pathfinder._utils.search_factory import create_binary_search_locations +from cuda.pathfinder._utils.toolchain_tracker import SearchContext, get_default_context -def _binary_filename_variants(name: str) -> Sequence[str]: - """Generate filename variants for a binary (cross-platform). +@functools.cache +def _find_nvidia_binary_default(binary_name: str) -> Optional[str]: + """Internal cached version using default context. + Args: - name: Base binary name. - + binary_name: Name of the binary to find. + Returns: - Tuple of possible filenames (e.g., "nvcc", "nvcc.exe"). + Absolute path to the binary, or None if not found. """ - return (name, f"{name}.exe") + return _find_nvidia_binary_impl(binary_name, get_default_context()) -def _get_site_packages_subdirs(binary_name: str) -> Sequence[str]: - """Get site-packages sub-directories for a binary. - +def _find_nvidia_binary_impl(binary_name: str, context: SearchContext) -> Optional[str]: + """Implementation of binary finding logic. + Args: - binary_name: Name of the binary. - + binary_name: Name of the binary to find. + context: SearchContext for toolchain consistency. + Returns: - List of sub-directories to search, or empty list if binary not in site-packages. + Absolute path to the binary, or None if not found. """ - relative_directories = SITE_PACKAGES_BINDIRS.get(binary_name) - if not relative_directories: - return [] - - # Expand site-packages paths - sub_directories = [] - for relative_directory in relative_directories: - for found_dir in find_sub_dirs_all_sitepackages(tuple(relative_directory.split("/"))): - sub_directories.append(found_dir) - return sub_directories - + if binary_name not in SUPPORTED_BINARIES: + raise RuntimeError(f"UNKNOWN {binary_name=}") -# Define search locations for binaries -def _create_search_locations(binary_name: str) -> list[SearchLocation]: - """Create search location configurations for a specific binary. + locations = create_binary_search_locations( + binary_name=binary_name, + site_packages_bindirs=SITE_PACKAGES_BINDIRS, + filename_variants_func=FilenameResolver.for_binary, + ) - Args: - binary_name: Name of the binary to search for. - - Returns: - List of SearchLocation objects to try. - """ - return [ - SearchLocation( - source=ToolchainSource.SITE_PACKAGES, - base_dir_func=lambda: None, # Use subdirs for full paths - subdirs=_get_site_packages_subdirs(binary_name), - filename_variants=_binary_filename_variants, - ), - SearchLocation( - source=ToolchainSource.CONDA, - base_dir_func=lambda: os.environ.get("CONDA_PREFIX"), - subdirs=["Library/bin", "bin"], # Windows and Unix layouts - filename_variants=_binary_filename_variants, - ), - SearchLocation( - source=ToolchainSource.CUDA_HOME, - base_dir_func=get_cuda_home_or_path, - subdirs=["bin"], - filename_variants=_binary_filename_variants, - ), - ] + path = context.find(binary_name, locations) + return _abs_norm(path) if path else None -@functools.cache -def find_nvidia_binary(binary_name: str, *, context: SearchContext | None = None) -> str | None: +def find_nvidia_binary(binary_name: str, *, context: Optional[SearchContext] = None) -> Optional[str]: """Locate a CUDA binary executable. + This function searches for CUDA binaries across multiple installation + sources, ensuring toolchain consistency when multiple artifacts are found. + Args: - binary_name: Name of the binary (e.g., "nvdisasm", "cuobjdump"). + binary_name: The name of the binary to find (e.g., ``"nvdisasm"``, + ``"cuobjdump"``). context: Optional SearchContext for toolchain consistency tracking. - If None, uses the default module-level context. + If None, uses the default module-level context which provides + caching and consistency across multiple calls. Returns: - Absolute path to the binary, or None if not found. + Absolute path to the discovered binary, or ``None`` if the + binary cannot be found. Raises: - RuntimeError: If binary_name is not supported. + RuntimeError: If ``binary_name`` is not in the supported set. ToolchainMismatchError: If binary found in different source than the context's preferred source. - """ - if binary_name not in SUPPORTED_BINARIES: - raise RuntimeError(f"UNKNOWN {binary_name=}") - if context is None: - context = get_default_context() + Search order: + 1. **NVIDIA Python wheels** - locations = _create_search_locations(binary_name) - path = context.find(binary_name, locations) - return _abs_norm(path) if path else None + - Scan installed distributions (``site-packages``) for binaries + shipped in NVIDIA wheels (e.g., ``cuda-nvcc``). + + 2. **Conda environments** + + - Check Conda-style installation prefixes (``$CONDA_PREFIX/bin`` on + Linux/Mac or ``$CONDA_PREFIX/Library/bin`` on Windows). + + 3. **CUDA Toolkit environment variables** + + - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in the + ``bin`` subdirectory. + + Examples: + Basic usage (uses default context with caching): + + >>> from cuda.pathfinder import find_nvidia_binary + >>> path = find_nvidia_binary("nvcc") + >>> if path: + ... print(f"Found nvcc at {path}") + + Using explicit context for isolated search: + + >>> from cuda.pathfinder import SearchContext, find_nvidia_binary + >>> ctx = SearchContext() + >>> nvcc = find_nvidia_binary("nvcc", context=ctx) + >>> nvdisasm = find_nvidia_binary("nvdisasm", context=ctx) + >>> # Both from same source, or ToolchainMismatchError raised + + Note: + When using the default context (context=None), results are cached. + When providing an explicit context, caching is bypassed to allow + for isolated searches with different consistency requirements. + """ + if context is None: + # Use cached version with default context + return _find_nvidia_binary_default(binary_name) + else: + # Bypass cache for explicit context + return _find_nvidia_binary_impl(binary_name, context) diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py index a6bd24c1d2..01de20f0df 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_nvidia_static_libs.py @@ -1,164 +1,133 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Find CUDA static libraries and artifacts across different installation sources.""" + import functools -import os -from typing import Sequence +from typing import Optional from cuda.pathfinder._static_libs.supported_nvidia_static_libs import ( SITE_PACKAGES_STATIC_LIBDIRS, SUPPORTED_STATIC_LIBS, ) -from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path -from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.filename_resolver import FilenameResolver from cuda.pathfinder._utils.path_utils import _abs_norm -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS -from cuda.pathfinder._utils.toolchain_tracker import ( - SearchContext, - SearchLocation, - ToolchainSource, - get_default_context, -) - - -# Generic search locations by toolchain source (platform-specific at import time) -if IS_WINDOWS: - CONDA_LIB_SUBDIRS = ("Library/lib", "Library/lib/x64") - CONDA_NVVM_SUBDIRS = ("Library/nvvm/libdevice",) - CUDA_HOME_LIB_SUBDIRS = ("lib/x64", "lib") - CUDA_HOME_NVVM_SUBDIRS = ("nvvm/libdevice",) -else: - CONDA_LIB_SUBDIRS = ("lib",) - CONDA_NVVM_SUBDIRS = ("nvvm/libdevice",) - CUDA_HOME_LIB_SUBDIRS = ("lib64", "lib") - CUDA_HOME_NVVM_SUBDIRS = ("nvvm/libdevice",) +from cuda.pathfinder._utils.search_factory import create_static_lib_search_locations +from cuda.pathfinder._utils.toolchain_tracker import SearchContext, get_default_context -def _static_lib_filename_variants(artifact_name: str) -> Sequence[str]: - """Generate platform-appropriate filename variants for an artifact. - +@functools.cache +def _find_nvidia_static_lib_default(artifact_name: str) -> Optional[str]: + """Internal cached version using default context. + Args: - artifact_name: Canonical artifact name (e.g., "cudadevrt", "libdevice.10.bc"). - - Returns: - Sequence of filenames to search for on this platform. - - Examples: - On Windows: - "cudadevrt" -> ("cudadevrt.lib",) - "libdevice.10.bc" -> ("libdevice.10.bc",) - On Linux: - "cudadevrt" -> ("libcudadevrt.a",) - "libdevice.10.bc" -> ("libdevice.10.bc",) - """ - # Files that are the same on all platforms (e.g., .bc bitcode files) - if "." in artifact_name: - return (artifact_name,) - - # Platform-specific library naming conventions - if IS_WINDOWS: - return (f"{artifact_name}.lib",) - else: - return (f"lib{artifact_name}.a",) - - -def _get_cuda_home_subdirs_with_targets() -> tuple[str, ...]: - """Get CUDA_HOME subdirectories including expanded targets/* paths. - + artifact_name: Canonical name of the artifact to find. + Returns: - Tuple of subdirectories to search under CUDA_HOME. + Absolute path to the artifact, or None if not found. """ - import glob - - subdirs = list(CUDA_HOME_LIB_SUBDIRS + CUDA_HOME_NVVM_SUBDIRS) - - # On Linux, also search targets/*/lib64 and targets/*/lib for cross-compilation - if not IS_WINDOWS: - cuda_home = get_cuda_home_or_path() - if cuda_home: - for lib_subdir in ("lib64", "lib"): - pattern = os.path.join(cuda_home, "targets", "*", lib_subdir) - for target_dir in sorted(glob.glob(pattern), reverse=True): - # Make relative to cuda_home - rel_path = os.path.relpath(target_dir, cuda_home) - subdirs.append(rel_path) - - return tuple(subdirs) + return _find_nvidia_static_lib_impl(artifact_name, get_default_context()) -def _create_search_locations(artifact_name: str) -> list[SearchLocation]: - """Create generic search location configurations. - +def _find_nvidia_static_lib_impl(artifact_name: str, context: SearchContext) -> Optional[str]: + """Implementation of static library finding logic. + Args: - artifact_name: Name of the artifact to search for. - + artifact_name: Canonical name of the artifact to find. + context: SearchContext for toolchain consistency. + Returns: - List of SearchLocation objects to try. + Absolute path to the artifact, or None if not found. """ - locations = [] - - # Site-packages: Create separate SearchLocation for each found directory - relative_directories = SITE_PACKAGES_STATIC_LIBDIRS.get(artifact_name) - if relative_directories: - for relative_directory in relative_directories: - for found_dir in find_sub_dirs_all_sitepackages(tuple(relative_directory.split("/"))): - locations.append( - SearchLocation( - source=ToolchainSource.SITE_PACKAGES, - base_dir_func=lambda d=found_dir: d, - subdirs=[""], - filename_variants=_static_lib_filename_variants, - ) - ) - - # Conda: Generic lib and nvvm locations - locations.append( - SearchLocation( - source=ToolchainSource.CONDA, - base_dir_func=lambda: os.environ.get("CONDA_PREFIX"), - subdirs=CONDA_LIB_SUBDIRS + CONDA_NVVM_SUBDIRS, - filename_variants=_static_lib_filename_variants, - ) - ) + if artifact_name not in SUPPORTED_STATIC_LIBS: + raise RuntimeError(f"UNKNOWN {artifact_name=}") - # CUDA_HOME: Generic lib and nvvm locations (including targets/* on Linux) - locations.append( - SearchLocation( - source=ToolchainSource.CUDA_HOME, - base_dir_func=get_cuda_home_or_path, - subdirs=_get_cuda_home_subdirs_with_targets(), - filename_variants=_static_lib_filename_variants, - ) + locations = create_static_lib_search_locations( + artifact_name=artifact_name, + site_packages_libdirs=SITE_PACKAGES_STATIC_LIBDIRS, + filename_variants_func=FilenameResolver.for_static_lib, ) - return locations + path = context.find(artifact_name, locations) + return _abs_norm(path) if path else None -@functools.cache -def find_nvidia_static_lib(artifact_name: str, *, context: SearchContext | None = None) -> str | None: +def find_nvidia_static_lib(artifact_name: str, *, context: Optional[SearchContext] = None) -> Optional[str]: """Locate a CUDA static library or artifact file. + This function searches for CUDA static libraries and artifacts (like + bitcode files) across multiple installation sources, ensuring toolchain + consistency when multiple artifacts are found. + Args: artifact_name: Canonical artifact name (e.g., "libdevice.10.bc", "cudadevrt"). Platform-specific filenames are resolved automatically: + - "cudadevrt" -> "libcudadevrt.a" on Linux, "cudadevrt.lib" on Windows + - "libdevice.10.bc" -> "libdevice.10.bc" (same on all platforms) + context: Optional SearchContext for toolchain consistency tracking. - If None, uses the default module-level context. + If None, uses the default module-level context which provides + caching and consistency across multiple calls. Returns: Absolute path to the artifact, or None if not found. Raises: - RuntimeError: If artifact_name is not supported. + RuntimeError: If ``artifact_name`` is not supported. ToolchainMismatchError: If artifact found in different source than the context's preferred source. - """ - if artifact_name not in SUPPORTED_STATIC_LIBS: - raise RuntimeError(f"UNKNOWN {artifact_name=}") - if context is None: - context = get_default_context() + Search order: + 1. **NVIDIA Python wheels** - locations = _create_search_locations(artifact_name) - path = context.find(artifact_name, locations) - return _abs_norm(path) if path else None + - Scan installed distributions (``site-packages``) for libraries + shipped in NVIDIA wheels (e.g., ``cuda-cudart``). + + 2. **Conda environments** + + - Check Conda-style installation prefixes: + + - ``$CONDA_PREFIX/lib`` (Linux/Mac) + - ``$CONDA_PREFIX/Library/lib`` (Windows) + - ``$CONDA_PREFIX/nvvm/libdevice`` (for bitcode files) + + 3. **CUDA Toolkit environment variables** + + - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order) and look in: + + - ``lib64``, ``lib`` subdirectories + - ``nvvm/libdevice`` (for bitcode files) + - ``targets/*/lib64``, ``targets/*/lib`` (Linux cross-compilation) + + Examples: + Basic usage (uses default context with caching): + + >>> from cuda.pathfinder import find_nvidia_static_lib + >>> path = find_nvidia_static_lib("cudadevrt") + >>> if path: + ... print(f"Found cudadevrt at {path}") + + Finding bitcode files: + + >>> libdevice = find_nvidia_static_lib("libdevice.10.bc") + + Using explicit context for isolated search: + + >>> from cuda.pathfinder import SearchContext, find_nvidia_static_lib + >>> ctx = SearchContext() + >>> cudadevrt = find_nvidia_static_lib("cudadevrt", context=ctx) + >>> libdevice = find_nvidia_static_lib("libdevice.10.bc", context=ctx) + >>> # Both from same source, or ToolchainMismatchError raised + + Note: + When using the default context (context=None), results are cached. + When providing an explicit context, caching is bypassed to allow + for isolated searches with different consistency requirements. + """ + if context is None: + # Use cached version with default context + return _find_nvidia_static_lib_default(artifact_name) + else: + # Bypass cache for explicit context + return _find_nvidia_static_lib_impl(artifact_name, context) diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/filename_resolver.py b/cuda_pathfinder/cuda/pathfinder/_utils/filename_resolver.py new file mode 100644 index 0000000000..ee813225a5 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/filename_resolver.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cross-platform filename resolution for CUDA artifacts.""" + +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +class FilenameResolver: + """Resolves artifact names to platform-specific filenames. + + This class handles the platform-specific naming conventions for CUDA + artifacts, allowing the search logic to remain platform-agnostic. + + Examples: + >>> # Binary resolution + >>> FilenameResolver.for_binary("nvcc") + ('nvcc', 'nvcc.exe') + + >>> # Static library resolution on Linux + >>> FilenameResolver.for_static_lib("cudadevrt") + ('libcudadevrt.a',) + + >>> # Static library resolution on Windows + >>> FilenameResolver.for_static_lib("cudadevrt") + ('cudadevrt.lib',) + + >>> # Files with extensions remain unchanged + >>> FilenameResolver.for_static_lib("libdevice.10.bc") + ('libdevice.10.bc',) + """ + + @staticmethod + def for_binary(name: str) -> tuple[str, ...]: + """Generate platform-appropriate binary filename variants. + + Returns both the exact name and Windows .exe variant to support + cross-platform fuzzy search. The filesystem naturally filters to + what exists. + + Args: + name: Base binary name (e.g., "nvcc", "nvdisasm"). + + Returns: + Tuple of possible filenames. Always includes exact name first, + then .exe variant. + """ + # Try exact name first, then with .exe extension + # This works across platforms - non-existent files won't be found + return (name, f"{name}.exe") + + @staticmethod + def for_static_lib(name: str) -> tuple[str, ...]: + """Generate platform-appropriate static library filename variants. + + Handles platform-specific naming conventions for static libraries + and preserves files that already have extensions (like .bc bitcode). + + Args: + name: Canonical artifact name (e.g., "cudadevrt", "libdevice.10.bc"). + + Returns: + Tuple of possible filenames for the current platform. + + Examples: + On Linux: + "cudadevrt" -> ("libcudadevrt.a",) + "libdevice.10.bc" -> ("libdevice.10.bc",) + On Windows: + "cudadevrt" -> ("cudadevrt.lib",) + "libdevice.10.bc" -> ("libdevice.10.bc",) + """ + # Files that already have extensions (e.g., .bc bitcode files) + # are the same on all platforms + if "." in name: + return (name,) + + # Platform-specific library naming conventions + if IS_WINDOWS: + return (f"{name}.lib",) + else: + return (f"lib{name}.a",) diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/platform_paths.py b/cuda_pathfinder/cuda/pathfinder/_utils/platform_paths.py new file mode 100644 index 0000000000..f70eebe8b0 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/platform_paths.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Platform-specific path configurations for CUDA artifacts. + +This module centralizes all platform-specific directory structures for +different CUDA installation sources (conda, CUDA_HOME, etc.). +""" + +from dataclasses import dataclass + +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +@dataclass(frozen=True) +class PlatformPaths: + """Platform-specific directory layouts for CUDA artifacts. + + Attributes: + conda_bin_subdirs: Binary subdirectories under CONDA_PREFIX. + conda_lib_subdirs: Library subdirectories under CONDA_PREFIX. + conda_nvvm_subdirs: NVVM/libdevice subdirectories under CONDA_PREFIX. + cuda_home_bin_subdirs: Binary subdirectories under CUDA_HOME. + cuda_home_lib_subdirs: Library subdirectories under CUDA_HOME. + cuda_home_nvvm_subdirs: NVVM/libdevice subdirectories under CUDA_HOME. + """ + + conda_bin_subdirs: tuple[str, ...] + conda_lib_subdirs: tuple[str, ...] + conda_nvvm_subdirs: tuple[str, ...] + cuda_home_bin_subdirs: tuple[str, ...] + cuda_home_lib_subdirs: tuple[str, ...] + cuda_home_nvvm_subdirs: tuple[str, ...] + + +# Platform-specific paths (determined at import time) +if IS_WINDOWS: + PLATFORM = PlatformPaths( + conda_bin_subdirs=("Library/bin", "bin"), + conda_lib_subdirs=("Library/lib", "Library/lib/x64"), + conda_nvvm_subdirs=("Library/nvvm/libdevice",), + cuda_home_bin_subdirs=("bin",), + cuda_home_lib_subdirs=("lib/x64", "lib"), + cuda_home_nvvm_subdirs=("nvvm/libdevice",), + ) +else: + PLATFORM = PlatformPaths( + conda_bin_subdirs=("bin",), + conda_lib_subdirs=("lib",), + conda_nvvm_subdirs=("nvvm/libdevice",), + cuda_home_bin_subdirs=("bin",), + cuda_home_lib_subdirs=("lib64", "lib"), + cuda_home_nvvm_subdirs=("nvvm/libdevice",), + ) + +# Constants for cross-compilation support (Linux only) +CUDA_TARGETS_LIB_SUBDIRS = ("lib64", "lib") diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/search_factory.py b/cuda_pathfinder/cuda/pathfinder/_utils/search_factory.py new file mode 100644 index 0000000000..b62e8f7033 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/search_factory.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Factory functions for creating standardized search locations. + +This module provides utilities to create SearchLocation objects for common +patterns, reducing code duplication across binary and static library finders. +""" + +import functools +import glob +import os +from typing import Callable, Sequence + +from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path +from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS +from cuda.pathfinder._utils.platform_paths import CUDA_TARGETS_LIB_SUBDIRS, PLATFORM +from cuda.pathfinder._utils.toolchain_tracker import SearchLocation, ToolchainSource +from cuda.pathfinder._utils.types import FilenameVariantFunc + + +@functools.cache +def _discover_cuda_home_lib_subdirs() -> tuple[str, ...]: + """Discover available CUDA_HOME library subdirectories. + + On Linux, this includes cross-compilation targets (targets/*/lib64, etc.). + On Windows, only standard lib directories are included. + + Returns: + Tuple of subdirectories to search under CUDA_HOME for libraries. + + Note: + This function is cached to avoid repeated glob operations. + """ + subdirs = list(PLATFORM.cuda_home_lib_subdirs + PLATFORM.cuda_home_nvvm_subdirs) + + # On Linux, also search targets/*/lib64 and targets/*/lib for cross-compilation + if not IS_WINDOWS: + cuda_home = get_cuda_home_or_path() + if cuda_home: + for lib_subdir in CUDA_TARGETS_LIB_SUBDIRS: + pattern = os.path.join(cuda_home, "targets", "*", lib_subdir) + for target_dir in sorted(glob.glob(pattern), reverse=True): + # Make relative to cuda_home + rel_path = os.path.relpath(target_dir, cuda_home) + subdirs.append(rel_path) + + return tuple(subdirs) + + +def create_standard_search_locations( + artifact_name: str, + site_packages_dirs_map: dict[str, list[str]], + conda_subdirs: Sequence[str], + cuda_home_subdirs: Sequence[str] | Callable[[], Sequence[str]], + filename_variants_func: FilenameVariantFunc, +) -> list[SearchLocation]: + """Create standard SITE_PACKAGES/CONDA/CUDA_HOME search locations. + + This factory function creates a standardized list of search locations + for CUDA artifacts, following the common pattern of checking: + 1. Site-packages (from pip/conda packages) + 2. Conda environment + 3. CUDA_HOME/CUDA_PATH + + Args: + artifact_name: Name of the artifact to search for. + site_packages_dirs_map: Mapping from artifact name to relative + site-packages directories. + conda_subdirs: Subdirectories to search under CONDA_PREFIX. + cuda_home_subdirs: Either a sequence of subdirectories to search + under CUDA_HOME, or a callable that returns such a sequence + (for dynamic discovery). + filename_variants_func: Function to generate platform-specific + filename variants for the artifact. + + Returns: + List of SearchLocation objects in priority order. + """ + locations = [] + + # Site-packages: Create separate SearchLocation for each found directory + site_package_rel_paths = site_packages_dirs_map.get(artifact_name) + if site_package_rel_paths: + for rel_path in site_package_rel_paths: + for absolute_dir in find_sub_dirs_all_sitepackages(tuple(rel_path.split("/"))): + locations.append( + SearchLocation( + source=ToolchainSource.SITE_PACKAGES, + base_dir_func=lambda d=absolute_dir: d, + subdirs=("",), # Already have full path + filename_variants=filename_variants_func, + ) + ) + + # Conda: Generic locations + locations.append( + SearchLocation( + source=ToolchainSource.CONDA, + base_dir_func=lambda: os.environ.get("CONDA_PREFIX"), + subdirs=conda_subdirs, + filename_variants=filename_variants_func, + ) + ) + + # CUDA_HOME: Generic locations (may include dynamic discovery) + if callable(cuda_home_subdirs): + cuda_home_subdirs = cuda_home_subdirs() + + locations.append( + SearchLocation( + source=ToolchainSource.CUDA_HOME, + base_dir_func=get_cuda_home_or_path, + subdirs=cuda_home_subdirs, + filename_variants=filename_variants_func, + ) + ) + + return locations + + +# Convenience functions for specific artifact types + +def create_binary_search_locations( + binary_name: str, + site_packages_bindirs: dict[str, list[str]], + filename_variants_func: FilenameVariantFunc, +) -> list[SearchLocation]: + """Create search locations specifically for binaries. + + Args: + binary_name: Name of the binary to search for. + site_packages_bindirs: Mapping from binary name to site-packages + binary directories. + filename_variants_func: Function to generate binary filename variants. + + Returns: + List of SearchLocation objects for binary search. + """ + return create_standard_search_locations( + artifact_name=binary_name, + site_packages_dirs_map=site_packages_bindirs, + conda_subdirs=PLATFORM.conda_bin_subdirs, + cuda_home_subdirs=PLATFORM.cuda_home_bin_subdirs, + filename_variants_func=filename_variants_func, + ) + + +def create_static_lib_search_locations( + artifact_name: str, + site_packages_libdirs: dict[str, list[str]], + filename_variants_func: FilenameVariantFunc, +) -> list[SearchLocation]: + """Create search locations specifically for static libraries. + + Args: + artifact_name: Name of the static library to search for. + site_packages_libdirs: Mapping from artifact name to site-packages + library directories. + filename_variants_func: Function to generate library filename variants. + + Returns: + List of SearchLocation objects for static library search. + """ + return create_standard_search_locations( + artifact_name=artifact_name, + site_packages_dirs_map=site_packages_libdirs, + conda_subdirs=PLATFORM.conda_lib_subdirs + PLATFORM.conda_nvvm_subdirs, + cuda_home_subdirs=_discover_cuda_home_lib_subdirs, + filename_variants_func=filename_variants_func, + ) diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py index d4f91ff1ab..8f6e4224fe 100644 --- a/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py +++ b/cuda_pathfinder/cuda/pathfinder/_utils/toolchain_tracker.py @@ -1,10 +1,51 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +"""Toolchain consistency tracking for CUDA artifacts. + +This module ensures that all CUDA artifacts (binaries, libraries, headers) +come from the same toolchain source to prevent version mismatches. The first +artifact found establishes a "preferred source" (SITE_PACKAGES, CONDA, or +CUDA_HOME) and subsequent searches prioritize that source. + +The consistency guarantee prevents mixing artifacts from different CUDA +installations, which can lead to ABI incompatibilities, version mismatches, +and runtime errors. + +Example: + Basic usage with default context: + + >>> from cuda.pathfinder import find_nvidia_binary + >>> + >>> # First find establishes preference (e.g., found in CONDA) + >>> nvcc = find_nvidia_binary("nvcc") + >>> + >>> # Subsequent finds prefer CONDA and search it first + >>> nvdisasm = find_nvidia_binary("nvdisasm") + + Using explicit context for isolated searches: + + >>> from cuda.pathfinder import SearchContext, find_nvidia_binary + >>> + >>> ctx = SearchContext() + >>> nvcc = find_nvidia_binary("nvcc", context=ctx) + >>> nvdisasm = find_nvidia_binary("nvdisasm", context=ctx) + >>> # Both from same source, or ToolchainMismatchError raised + + Resetting the default context: + + >>> from cuda.pathfinder import reset_default_context + >>> + >>> reset_default_context() # Clear preference for new search session +""" + import os from dataclasses import dataclass from enum import Enum, auto -from typing import Callable, Optional, Sequence +from itertools import product +from typing import Optional, Sequence + +from cuda.pathfinder._utils.types import BaseDirFunc, FilenameVariantFunc class ToolchainSource(Enum): @@ -21,15 +62,18 @@ class SearchLocation: Attributes: source: Which toolchain source this represents. - base_dir_func: Function that returns the base directory to search, or None if unavailable. - subdirs: Subdirectories to check under the base (e.g., ["bin"], ["Library/bin", "bin"]). - filename_variants: Function that takes artifact name and returns possible filenames. + base_dir_func: Function returning the base directory, or None if unavailable. + This is a function (not a string) to defer environment variable lookups + and allow dynamic discovery. Called at search time, not construction time. + subdirs: Subdirectories to check under the base (e.g., ["bin"]). + filename_variants: Function taking artifact name, returning possible filenames. + Allows platform-specific name resolution (e.g., "nvcc" -> ["nvcc", "nvcc.exe"]). """ source: ToolchainSource - base_dir_func: Callable[[], Optional[str]] + base_dir_func: BaseDirFunc subdirs: Sequence[str] - filename_variants: Callable[[str], Sequence[str]] + filename_variants: FilenameVariantFunc @dataclass(frozen=True) @@ -42,7 +86,18 @@ class ArtifactRecord: class ToolchainMismatchError(RuntimeError): - """Raised when artifacts from different sources are mixed.""" + """Raised when artifacts from different sources are mixed. + + This error indicates that an artifact was found in a different toolchain + source than previously found artifacts, which could lead to version + incompatibilities. + + Attributes: + artifact_name: Name of the artifact that caused the mismatch. + attempted_source: Source where the artifact was found. + preferred_source: Source that was established by previous artifacts. + preferred_artifacts: List of artifacts already found from preferred source. + """ def __init__( self, @@ -59,7 +114,8 @@ def __init__( artifact_list = ", ".join(f"'{a.name}'" for a in preferred_artifacts) message = ( f"Toolchain mismatch: '{artifact_name}' found in {attempted_source.name}, " - f"but using {preferred_source.name} (previous: {artifact_list})" + f"but already using {preferred_source.name} for {artifact_list}. " + f"Call reset_default_context() to clear the preference, or use an explicit SearchContext." ) super().__init__(message) @@ -80,12 +136,10 @@ def search_location(location: SearchLocation, artifact_name: str) -> Optional[st filenames = location.filename_variants(artifact_name) - for subdir in location.subdirs: - dir_path = os.path.join(base_dir, subdir) - for filename in filenames: - file_path = os.path.join(dir_path, filename) - if os.path.isfile(file_path): - return file_path + for subdir, filename in product(location.subdirs, filenames): + file_path = os.path.join(base_dir, subdir, filename) + if os.path.isfile(file_path): + return file_path return None @@ -130,6 +184,23 @@ def record(self, name: str, path: str, source: ToolchainSource) -> None: self._artifacts[name] = ArtifactRecord(name=name, path=path, source=source) + def _reorder_by_preference(self, locations: Sequence[SearchLocation]) -> list[SearchLocation]: + """Reorder locations to search preferred source first. + + Args: + locations: Original search locations. + + Returns: + Reordered list with preferred source first, maintaining original + order for non-preferred sources. + """ + if not self._preferred_source: + return list(locations) + + preferred = [loc for loc in locations if loc.source == self._preferred_source] + others = [loc for loc in locations if loc.source != self._preferred_source] + return preferred + others + def find(self, artifact_name: str, locations: Sequence[SearchLocation]) -> Optional[str]: """Search for artifact respecting toolchain consistency. @@ -143,15 +214,8 @@ def find(self, artifact_name: str, locations: Sequence[SearchLocation]) -> Optio Raises: ToolchainMismatchError: If found in different source than preferred. """ - # Reorder to search preferred source first, maintaining original order for ties - # (stable sort: preferred source gets priority 0, others get priority 1) - if self._preferred_source: - ordered_locations = sorted( - locations, - key=lambda loc: 0 if loc.source == self._preferred_source else 1 - ) - else: - ordered_locations = list(locations) + # Reorder to search preferred source first + ordered_locations = self._reorder_by_preference(locations) # Try each location for location in ordered_locations: diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/types.py b/cuda_pathfinder/cuda/pathfinder/_utils/types.py new file mode 100644 index 0000000000..2202fc1cf1 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/types.py @@ -0,0 +1,12 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Type aliases for pathfinder utilities.""" + +from typing import Callable, Optional, Sequence + +# Type alias for functions that generate filename variants +FilenameVariantFunc = Callable[[str], Sequence[str]] + +# Type alias for functions that return a base directory +BaseDirFunc = Callable[[], Optional[str]] diff --git a/cuda_pathfinder/tests/test_cache_behavior.py b/cuda_pathfinder/tests/test_cache_behavior.py new file mode 100644 index 0000000000..9dda1e6327 --- /dev/null +++ b/cuda_pathfinder/tests/test_cache_behavior.py @@ -0,0 +1,89 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for cache behavior with explicit contexts.""" + +import os + +import pytest + +from cuda.pathfinder import SearchContext, find_nvidia_binary, reset_default_context +from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES + + +def test_default_context_is_cached(): + """Test that default context calls are cached.""" + # Reset to ensure clean state + reset_default_context() + + # Multiple calls with default context should use cache + if SUPPORTED_BINARIES: + binary_name = SUPPORTED_BINARIES[0] + result1 = find_nvidia_binary(binary_name) + result2 = find_nvidia_binary(binary_name) + # Should return the same object (cached) + assert result1 is result2 + + +def test_explicit_context_bypasses_cache(tmp_path): + """Test that explicit context bypasses cache.""" + # Create mock binary directories + dir1 = tmp_path / "install1" / "bin" + dir1.mkdir(parents=True) + binary1 = dir1 / "nvdisasm" + binary1.touch() + + dir2 = tmp_path / "install2" / "bin" + dir2.mkdir(parents=True) + binary2 = dir2 / "nvdisasm" + binary2.touch() + + # Use explicit contexts - they should bypass cache + ctx1 = SearchContext() + ctx2 = SearchContext() + + # Note: In actual use, the search locations would find different paths + # This test verifies the structure allows for independent contexts + assert ctx1 is not ctx2 + assert ctx1.preferred_source is None + assert ctx2.preferred_source is None + + +def test_reset_clears_preference(): + """Test that reset_default_context clears the preference.""" + from cuda.pathfinder._utils.toolchain_tracker import ToolchainSource, get_default_context + + reset_default_context() + ctx = get_default_context() + + # Record something + ctx.record("test", "/path/to/test", ToolchainSource.CONDA) + assert ctx.preferred_source == ToolchainSource.CONDA + + # Reset should clear + reset_default_context() + ctx2 = get_default_context() + assert ctx2.preferred_source is None + + +@pytest.mark.parametrize("binary_name", ["nvdisasm", "cuobjdump"]) +def test_explicit_context_independent(binary_name): + """Test that explicit contexts are independent of default.""" + # Skip if binary not supported + if binary_name not in SUPPORTED_BINARIES: + pytest.skip(f"{binary_name} not in SUPPORTED_BINARIES") + + reset_default_context() + + # Use default context + default_result = find_nvidia_binary(binary_name) + + # Use explicit context - should not affect cache + explicit_ctx = SearchContext() + explicit_result = find_nvidia_binary(binary_name, context=explicit_ctx) + + # Both might find the same file, but search was independent + # If both found something, they should be the same path + if default_result and explicit_result: + # Both found - should be same path (but searched independently) + assert os.path.normpath(default_result) == os.path.normpath(explicit_result) diff --git a/cuda_pathfinder/tests/test_filename_resolver.py b/cuda_pathfinder/tests/test_filename_resolver.py new file mode 100644 index 0000000000..6bc68ccd86 --- /dev/null +++ b/cuda_pathfinder/tests/test_filename_resolver.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for filename resolution utilities.""" + +import pytest + +from cuda.pathfinder._utils.filename_resolver import FilenameResolver + + +def test_binary_resolution(): + """Test binary filename resolution.""" + result = FilenameResolver.for_binary("nvcc") + assert result == ("nvcc", "nvcc.exe") + assert len(result) == 2 + + +def test_binary_resolution_preserves_order(): + """Test that exact name comes first.""" + result = FilenameResolver.for_binary("nvdisasm") + assert result[0] == "nvdisasm" + assert result[1] == "nvdisasm.exe" + + +@pytest.mark.skipif( + not hasattr(__import__("cuda.pathfinder._utils.platform_aware"), "IS_WINDOWS"), + reason="Platform detection not available", +) +def test_static_lib_resolution_platform_specific(): + """Test static library resolution is platform-specific.""" + from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + result = FilenameResolver.for_static_lib("cudadevrt") + + if IS_WINDOWS: + assert result == ("cudadevrt.lib",) + else: + assert result == ("libcudadevrt.a",) + + +def test_static_lib_preserves_extensions(): + """Test that files with extensions are preserved.""" + result = FilenameResolver.for_static_lib("libdevice.10.bc") + assert result == ("libdevice.10.bc",) + + +def test_static_lib_bitcode_files(): + """Test bitcode file resolution.""" + # Bitcode files should be platform-independent + result = FilenameResolver.for_static_lib("libdevice.10.bc") + assert result == ("libdevice.10.bc",) + + result = FilenameResolver.for_static_lib("some.other.bc") + assert result == ("some.other.bc",) diff --git a/cuda_pathfinder/tests/test_platform_paths.py b/cuda_pathfinder/tests/test_platform_paths.py new file mode 100644 index 0000000000..eb7d059fa8 --- /dev/null +++ b/cuda_pathfinder/tests/test_platform_paths.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for platform path configurations.""" + +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS +from cuda.pathfinder._utils.platform_paths import CUDA_TARGETS_LIB_SUBDIRS, PLATFORM + + +def test_platform_paths_structure(): + """Test that PLATFORM has all expected attributes.""" + assert hasattr(PLATFORM, "conda_bin_subdirs") + assert hasattr(PLATFORM, "conda_lib_subdirs") + assert hasattr(PLATFORM, "conda_nvvm_subdirs") + assert hasattr(PLATFORM, "cuda_home_bin_subdirs") + assert hasattr(PLATFORM, "cuda_home_lib_subdirs") + assert hasattr(PLATFORM, "cuda_home_nvvm_subdirs") + + +def test_platform_paths_are_tuples(): + """Test that all paths are tuples (immutable).""" + assert isinstance(PLATFORM.conda_bin_subdirs, tuple) + assert isinstance(PLATFORM.conda_lib_subdirs, tuple) + assert isinstance(PLATFORM.conda_nvvm_subdirs, tuple) + assert isinstance(PLATFORM.cuda_home_bin_subdirs, tuple) + assert isinstance(PLATFORM.cuda_home_lib_subdirs, tuple) + assert isinstance(PLATFORM.cuda_home_nvvm_subdirs, tuple) + + +def test_platform_paths_windows_specific(): + """Test Windows-specific paths.""" + if IS_WINDOWS: + assert "Library/bin" in PLATFORM.conda_bin_subdirs + assert "Library/lib" in PLATFORM.conda_lib_subdirs + assert "lib/x64" in PLATFORM.cuda_home_lib_subdirs + else: + # Unix should have simpler paths + assert "bin" in PLATFORM.conda_bin_subdirs + assert "lib" in PLATFORM.conda_lib_subdirs + assert "lib64" in PLATFORM.cuda_home_lib_subdirs or "lib" in PLATFORM.cuda_home_lib_subdirs + + +def test_cuda_targets_constant(): + """Test CUDA targets constant.""" + assert isinstance(CUDA_TARGETS_LIB_SUBDIRS, tuple) + assert len(CUDA_TARGETS_LIB_SUBDIRS) == 2 + assert "lib64" in CUDA_TARGETS_LIB_SUBDIRS + assert "lib" in CUDA_TARGETS_LIB_SUBDIRS + + +def test_nvvm_subdirs_present(): + """Test that NVVM subdirs are configured.""" + assert len(PLATFORM.conda_nvvm_subdirs) > 0 + assert len(PLATFORM.cuda_home_nvvm_subdirs) > 0 + assert "libdevice" in str(PLATFORM.conda_nvvm_subdirs) + assert "libdevice" in str(PLATFORM.cuda_home_nvvm_subdirs) diff --git a/cuda_pathfinder/tests/test_toolchain_tracker.py b/cuda_pathfinder/tests/test_toolchain_tracker.py index 01600beb8c..8704dc093b 100644 --- a/cuda_pathfinder/tests/test_toolchain_tracker.py +++ b/cuda_pathfinder/tests/test_toolchain_tracker.py @@ -6,7 +6,6 @@ import pytest from cuda.pathfinder._utils.toolchain_tracker import ( - ArtifactRecord, SearchContext, SearchLocation, ToolchainMismatchError, @@ -49,6 +48,9 @@ def test_context_rejects_different_source(): assert exc_info.value.artifact_name == "nvdisasm" assert exc_info.value.attempted_source == ToolchainSource.CUDA_HOME assert exc_info.value.preferred_source == ToolchainSource.CONDA + # Check that the improved error message includes helpful text + assert "reset_default_context()" in str(exc_info.value) + assert "explicit SearchContext" in str(exc_info.value) def test_find_prefers_established_source(tmp_path):