From 3d1eab90ea77e63ae58667d26bd44b9ebef1afae Mon Sep 17 00:00:00 2001 From: Jonathan Dekhtiar Date: Fri, 21 Aug 2026 00:01:28 -0500 Subject: [PATCH 1/5] Add cuDNN discovery and NCCL header support - cuDNN: added dynamic-library loading and header discovery. - NCCL: dynamic-library loading already existed; added the missing header discovery. --- .../_dynamic_libs/descriptor_catalog.py | 11 ++++++++++ .../_headers/header_descriptor_catalog.py | 17 ++++++++++++++ cuda_pathfinder/pyproject.toml | 2 ++ .../tests/test_descriptor_catalog.py | 14 ++++++++++++ .../tests/test_find_nvidia_headers.py | 22 +++++++++++++++++++ .../conda_create_for_pathfinder_testing.ps1 | 1 + .../conda_create_for_pathfinder_testing.sh | 2 ++ 7 files changed, 69 insertions(+) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index e39046eec70..8bf5535ea92 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -411,6 +411,17 @@ class DescriptorSpec: dependencies=("nvshmem_host",), requires_rtld_deepbind=True, ), + DescriptorSpec( + name="cudnn", + packaged_with="other", + linux_sonames=("libcudnn.so.9",), + windows_dlls=("cudnn64_9.dll",), + supported_windows_arch=("x64",), + site_packages_linux=("nvidia/cudnn/lib",), + site_packages_windows=WindowsSearchDirs.x64_only("nvidia/cudnn/bin"), + anchor_rel_dirs_windows=WindowsSearchDirs.x64_only("bin/x64", "bin"), + requires_add_dll_directory=True, + ), DescriptorSpec( name="cusolverMp", packaged_with="other", diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py index b364e224e7e..36e182a3c8c 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py @@ -150,6 +150,14 @@ class HeaderDescriptorSpec: # ----------------------------------------------------------------------- # Third-party / separately packaged headers # ----------------------------------------------------------------------- + HeaderDescriptorSpec( + name="cudnn", + packaged_with="other", + header_basename="cudnn.h", + site_packages_dirs=("nvidia/cudnn/include",), + conda_targets_layout=False, + use_ctk_root_canary=False, + ), HeaderDescriptorSpec( name="cusolverMp", packaged_with="other", @@ -244,6 +252,15 @@ class HeaderDescriptorSpec: conda_targets_layout=False, use_ctk_root_canary=False, ), + HeaderDescriptorSpec( + name="nccl", + packaged_with="other", + header_basename="nccl.h", + site_packages_dirs=("nvidia/nccl/include",), + available_on_windows=False, + conda_targets_layout=False, + use_ctk_root_canary=False, + ), HeaderDescriptorSpec( name="nvshmem", packaged_with="other", diff --git a/cuda_pathfinder/pyproject.toml b/cuda_pathfinder/pyproject.toml index 4fd7a5e3edf..022d868568c 100644 --- a/cuda_pathfinder/pyproject.toml +++ b/cuda_pathfinder/pyproject.toml @@ -24,6 +24,7 @@ cu12 = [ "cuquantum-cu12; sys_platform != 'win32'", "cutensor-cu12", "nvidia-cublasmp-cu12; sys_platform != 'win32'", + "nvidia-cudnn-cu12", "nvidia-cudss-cu12", "nvidia-cufftmp-cu12; sys_platform != 'win32'", "nvidia-cusolvermp-cu12; sys_platform != 'win32'", @@ -39,6 +40,7 @@ cu13 = [ "cutensor-cu13", "nvidia-cublasmp-cu13; sys_platform != 'win32'", "nvidia-cudla; platform_system == 'Linux' and platform_machine == 'aarch64'", + "nvidia-cudnn-cu13", "nvidia-cudss-cu13", "nvidia-cufftmp-cu13; sys_platform != 'win32'", "nvidia-cusolvermp-cu13; sys_platform != 'win32'", diff --git a/cuda_pathfinder/tests/test_descriptor_catalog.py b/cuda_pathfinder/tests/test_descriptor_catalog.py index 3b643aa2e77..02990267ba7 100644 --- a/cuda_pathfinder/tests/test_descriptor_catalog.py +++ b/cuda_pathfinder/tests/test_descriptor_catalog.py @@ -115,6 +115,20 @@ def test_cusparselt_windows_metadata_matches_wheel_layouts(): ) +@pytest.mark.agent_authored(model="gpt-5") +def test_cudnn_metadata_matches_wheel_layouts(): + spec = _CATALOG_BY_NAME["cudnn"] + assert spec.packaged_with == "other" + assert spec.linux_sonames == ("libcudnn.so.9",) + assert spec.windows_dlls == ("cudnn64_9.dll",) + assert spec.supported_windows_arch == ("x64",) + assert spec.site_packages_linux == ("nvidia/cudnn/lib",) + assert spec.site_packages_windows == WindowsSearchDirs.x64_only("nvidia/cudnn/bin") + assert spec.anchor_rel_dirs_windows == WindowsSearchDirs.x64_only("bin/x64", "bin") + assert spec.dependencies == () + assert spec.requires_add_dll_directory + + @pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) def test_ctk_root_canary_anchors_reference_known_ctk_libs(spec: DescriptorSpec): for anchor in spec.ctk_root_canary_anchor_libnames: diff --git a/cuda_pathfinder/tests/test_find_nvidia_headers.py b/cuda_pathfinder/tests/test_find_nvidia_headers.py index 3e045dae265..9e36e8acc56 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/cuda_pathfinder/tests/test_find_nvidia_headers.py @@ -27,6 +27,7 @@ from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( _resolve_system_loaded_abs_path_in_subprocess, ) +from cuda.pathfinder._headers.header_descriptor import HEADER_DESCRIPTORS from cuda.pathfinder._headers.supported_nvidia_headers import ( SUPPORTED_HEADERS_CTK, SUPPORTED_HEADERS_CTK_ALL, @@ -43,6 +44,7 @@ NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = { "cudensitymat": r"^cudensitymat-.*$", + "cudnn": r"^nvidia-cudnn-.*$", "cupauliprop": r"^cupauliprop-.*$", "cusolverMp": r"^nvidia-cusolvermp-.*$", "cusparseLt": r"^nvidia-cusparselt-.*$", @@ -53,6 +55,7 @@ "custatevec": r"^custatevec-.*$", "cutlass": r"^nvidia-cutlass$", "mathdx": r"^nvidia-libmathdx-.*$", + "nccl": r"^nvidia-nccl-.*$", "nvshmem": r"^nvidia-nvshmem-.*$", } @@ -78,6 +81,25 @@ def test_non_ctk_importlib_metadata_distributions_names(): assert sorted(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES) == sorted(SUPPORTED_HEADERS_NON_CTK_ALL) +@pytest.mark.agent_authored(model="gpt-5") +def test_cudnn_and_nccl_header_metadata_matches_wheel_layouts(): + cudnn = HEADER_DESCRIPTORS["cudnn"] + assert cudnn.header_basename == "cudnn.h" + assert cudnn.site_packages_dirs == ("nvidia/cudnn/include",) + assert cudnn.available_on_linux + assert cudnn.available_on_windows + assert not cudnn.conda_targets_layout + assert not cudnn.use_ctk_root_canary + + nccl = HEADER_DESCRIPTORS["nccl"] + assert nccl.header_basename == "nccl.h" + assert nccl.site_packages_dirs == ("nvidia/nccl/include",) + assert nccl.available_on_linux + assert not nccl.available_on_windows + assert not nccl.conda_targets_layout + assert not nccl.use_ctk_root_canary + + @functools.cache def have_distribution_for(libname: str) -> bool: pattern = re.compile(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES[libname]) diff --git a/toolshed/conda_create_for_pathfinder_testing.ps1 b/toolshed/conda_create_for_pathfinder_testing.ps1 index 0f93c3ab026..94a917d6ff0 100644 --- a/toolshed/conda_create_for_pathfinder_testing.ps1 +++ b/toolshed/conda_create_for_pathfinder_testing.ps1 @@ -19,6 +19,7 @@ conda activate "pathfinder_testing_cu$CudaMajorMinorPatch" # Keep this list aligned with the Windows-installable subset of # cuda_pathfinder/pyproject.toml. $cpkgs = @( + "cudnn", "cusparselt-dev", "cutensor", "cutlass", diff --git a/toolshed/conda_create_for_pathfinder_testing.sh b/toolshed/conda_create_for_pathfinder_testing.sh index 4e38c5dbe88..3512f80ec43 100755 --- a/toolshed/conda_create_for_pathfinder_testing.sh +++ b/toolshed/conda_create_for_pathfinder_testing.sh @@ -26,6 +26,7 @@ set -u # cuda_pathfinder/pyproject.toml. cpkgs=( "cuquantum" + "cudnn" "cusparselt-dev" "cutensor" "cutlass" @@ -34,6 +35,7 @@ cpkgs=( "libcufftmp-dev" "libcusolvermp-dev" "libmathdx-dev" + "nccl" "libnvshmem3" "libnvshmem-dev" ) From 1b30583b79b3610d2190d18e78c43b62a26dcf6b Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Fri, 21 Aug 2026 08:50:36 -0700 Subject: [PATCH 2/5] Fix cuDNN dynamic loading on Windows --- .../_dynamic_libs/descriptor_catalog.py | 10 +++ .../_dynamic_libs/load_dl_common.py | 15 ++++ .../_dynamic_libs/load_nvidia_dynamic_lib.py | 21 +++-- .../_dynamic_libs/search_platform.py | 20 +++++ .../pathfinder/_dynamic_libs/search_steps.py | 32 ++++++- .../cuda/pathfinder/_utils/path_sort.py | 23 +++++ .../tests/test_descriptor_catalog.py | 17 +++- cuda_pathfinder/tests/test_load_dl_common.py | 88 +++++++++++++++++++ .../tests/test_load_nvidia_dynamic_lib.py | 2 +- cuda_pathfinder/tests/test_search_steps.py | 78 ++++++++++++++++ 10 files changed, 296 insertions(+), 10 deletions(-) create mode 100644 cuda_pathfinder/cuda/pathfinder/_utils/path_sort.py create mode 100644 cuda_pathfinder/tests/test_load_dl_common.py diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index 8bf5535ea92..961dfb98758 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -65,8 +65,11 @@ class DescriptorSpec: site_packages_linux: tuple[str, ...] = () site_packages_windows: WindowsSearchDirs = WindowsSearchDirs() dependencies: tuple[str, ...] = () + optional_dependencies: tuple[str, ...] = () anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib") anchor_rel_dirs_windows: WindowsSearchDirs = DEFAULT_WINDOWS_CTK_ANCHOR_DIRS + windows_root_env_vars: tuple[str, ...] = () + program_files_rel_dirs_windows: WindowsSearchDirs = WindowsSearchDirs() ctk_root_canary_anchor_libnames: tuple[str, ...] = () requires_add_dll_directory: bool = False requires_rtld_deepbind: bool = False @@ -419,7 +422,14 @@ class DescriptorSpec: supported_windows_arch=("x64",), site_packages_linux=("nvidia/cudnn/lib",), site_packages_windows=WindowsSearchDirs.x64_only("nvidia/cudnn/bin"), + dependencies=("cublasLt",), + optional_dependencies=("nvrtc",), anchor_rel_dirs_windows=WindowsSearchDirs.x64_only("bin/x64", "bin"), + windows_root_env_vars=("CUDNN_PATH",), + program_files_rel_dirs_windows=WindowsSearchDirs.x64_only( + "NVIDIA/CUDNN/v9.*/bin/x64", + "NVIDIA/CUDNN/v9.*/bin", + ), requires_add_dll_directory=True, ), DescriptorSpec( diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py index 8d7987d00e2..c3917073032 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py @@ -32,5 +32,20 @@ class LoadedDL: def load_dependencies(desc: LibDescriptor, load_func: Callable[[str], LoadedDL]) -> None: + """Load required dependencies, then best-effort runtime dependencies. + + A plain ``DynamicLibNotFoundError`` from an optional dependency is + suppressed. More specific contract errors and failures while loading a + dependency that was found remain errors. + """ for dep in desc.dependencies: load_func(dep) + for dep in desc.optional_dependencies: + try: + load_func(dep) + except DynamicLibNotFoundError as exc: + # Both public contract errors inherit DynamicLibNotFoundError, but + # neither an unknown descriptor nor platform incompatibility means + # that an optional runtime component is simply absent. + if type(exc) is not DynamicLibNotFoundError: + raise diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py index 61bf31720d3..93073c21a7e 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py @@ -275,12 +275,19 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: 4. **Environment variables** - - If set, use ``CUDA_PATH`` or ``CUDA_HOME`` (in that order). - On Windows, this is the typical way system-installed CTK DLLs are - located. Note that the NVIDIA CTK installer automatically - adds ``CUDA_PATH`` to the system-wide environment. + - First search library-specific roots declared by the descriptor, + such as ``CUDNN_PATH``, then use ``CUDA_PATH`` or ``CUDA_HOME`` + (in that order). On Windows, ``CUDA_PATH`` is the typical way + system-installed CTK DLLs are located. Note that the NVIDIA CTK + installer automatically adds ``CUDA_PATH`` to the system-wide + environment. - 5. **CTK root canary probe (discoverable libs only)** + 5. **Windows Program Files (configured libraries only)** + + - Search descriptor-configured standalone installation layouts, + such as versioned cuDNN directories under ``ProgramFiles``. + + 6. **CTK root canary probe (discoverable libs only)** - For selected libraries whose shared object doesn't reside on the standard linker path (currently ``nvvm``), attempt to derive CTK @@ -298,8 +305,8 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: 0. Already loaded in the current process 1. OS default mechanisms (``dlopen`` / ``LoadLibraryExW``) - The CTK-specific steps (site-packages, conda, ``CUDA_PATH``, canary - probe) are skipped entirely. + The non-driver steps (site-packages, conda, environment roots, + ``ProgramFiles``, and canary probe) are skipped entirely. Notes: The search is performed **per library**. There is currently no mechanism to diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index 2d6a5f016a7..e71a2c3b652 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py @@ -103,6 +103,10 @@ def conda_anchor_point(self, conda_prefix: str) -> str: ... def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ... + def root_environment_variables(self, desc: LibDescriptor) -> tuple[str, ...]: ... + + def program_files_dir_globs(self, desc: LibDescriptor) -> tuple[str, ...]: ... + def find_in_site_packages( self, rel_dirs: tuple[str, ...], @@ -135,6 +139,12 @@ def conda_anchor_point(self, conda_prefix: str) -> str: def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: return cast(tuple[str, ...], desc.anchor_rel_dirs_linux) + def root_environment_variables(self, _desc: LibDescriptor) -> tuple[str, ...]: + return () + + def program_files_dir_globs(self, _desc: LibDescriptor) -> tuple[str, ...]: + return () + def find_in_site_packages( self, rel_dirs: tuple[str, ...], @@ -192,6 +202,16 @@ def conda_anchor_point(self, conda_prefix: str) -> str: def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: return cast(tuple[str, ...], desc.anchor_rel_dirs_windows.for_arch(self.target_arch)) + def root_environment_variables(self, desc: LibDescriptor) -> tuple[str, ...]: + return cast(tuple[str, ...], desc.windows_root_env_vars) + + def program_files_dir_globs(self, desc: LibDescriptor) -> tuple[str, ...]: + program_files = os.environ.get("PROGRAMW6432") or os.environ.get("PROGRAMFILES") + if not program_files: + return () + rel_dirs = desc.program_files_rel_dirs_windows.for_arch(self.target_arch) + return tuple(os.path.join(program_files, rel_dir) for rel_dir in rel_dirs) + def find_in_site_packages( self, rel_dirs: tuple[str, ...], diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py index 5901094fcaa..82dcdee4003 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py @@ -29,6 +29,7 @@ from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError from cuda.pathfinder._dynamic_libs.search_platform import PLATFORM, SearchPlatform from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home +from cuda.pathfinder._utils.path_sort import natural_path_sort_key # --------------------------------------------------------------------------- # Data types @@ -183,6 +184,19 @@ def find_in_conda(ctx: SearchContext) -> FindResult | None: return None +def find_in_descriptor_environment_roots(ctx: SearchContext) -> FindResult | None: + """Search roots named by descriptor-specific environment variables.""" + for env_var in ctx.platform.root_environment_variables(ctx.desc): + root = os.environ.get(env_var) + if not root: + continue + lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, root) + abs_path = _find_using_lib_dir(ctx, lib_dir) + if abs_path is not None: + return FindResult(abs_path, env_var) + return None + + def find_in_cuda_path(ctx: SearchContext) -> FindResult | None: """Search ``$CUDA_PATH`` / ``$CUDA_HOME``. @@ -204,6 +218,18 @@ def find_in_cuda_path(ctx: SearchContext) -> FindResult | None: return None +def find_in_program_files(ctx: SearchContext) -> FindResult | None: + """Search descriptor-configured installation directories under Program Files.""" + for dir_glob in ctx.platform.program_files_dir_globs(ctx.desc): + for lib_dir in sorted(glob.glob(dir_glob), key=natural_path_sort_key, reverse=True): + if not os.path.isdir(lib_dir): + continue + abs_path = _find_using_lib_dir(ctx, os.path.normpath(lib_dir)) + if abs_path is not None: + return FindResult(abs_path, "ProgramFiles") + return None + + # --------------------------------------------------------------------------- # Step sequences per strategy # --------------------------------------------------------------------------- @@ -212,7 +238,11 @@ def find_in_cuda_path(ctx: SearchContext) -> FindResult | None: EARLY_FIND_STEPS: tuple[FindStep, ...] = (find_in_site_packages, find_in_conda) #: Find steps that run after system search fails. -LATE_FIND_STEPS: tuple[FindStep, ...] = (find_in_cuda_path,) +LATE_FIND_STEPS: tuple[FindStep, ...] = ( + find_in_descriptor_environment_roots, + find_in_cuda_path, + find_in_program_files, +) # --------------------------------------------------------------------------- diff --git a/cuda_pathfinder/cuda/pathfinder/_utils/path_sort.py b/cuda_pathfinder/cuda/pathfinder/_utils/path_sort.py new file mode 100644 index 00000000000..7c162f8c142 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_utils/path_sort.py @@ -0,0 +1,23 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic path ordering helpers.""" + +from __future__ import annotations + +import os +import re + +_DIGIT_RUN = re.compile(r"(\d+)") + + +def natural_path_sort_key(path: str) -> tuple[tuple[int, int, str], ...]: + """Return a key that compares numeric path components by value.""" + key: list[tuple[int, int, str]] = [] + for part in _DIGIT_RUN.split(os.path.normcase(path)): + if part.isdigit(): + normalized = part.lstrip("0") or "0" + key.append((1, len(normalized), normalized)) + else: + key.append((0, 0, part)) + return tuple(key) diff --git a/cuda_pathfinder/tests/test_descriptor_catalog.py b/cuda_pathfinder/tests/test_descriptor_catalog.py index 02990267ba7..9bf59e0a276 100644 --- a/cuda_pathfinder/tests/test_descriptor_catalog.py +++ b/cuda_pathfinder/tests/test_descriptor_catalog.py @@ -52,6 +52,15 @@ def test_no_self_dependency(spec: DescriptorSpec): assert spec.name not in spec.dependencies, f"{spec.name} lists itself as a dependency" +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +@pytest.mark.agent_authored(model="gpt-5") +def test_optional_dependencies_reference_existing_entries(spec: DescriptorSpec): + for dep in spec.optional_dependencies: + assert dep in _CATALOG_BY_NAME, f"{spec.name} optionally depends on unknown library {dep!r}" + assert dep != spec.name, f"{spec.name} lists itself as an optional dependency" + assert dep not in spec.dependencies, f"{spec.name} lists {dep!r} as both required and optional" + + @pytest.mark.parametrize( "spec", [s for s in DESCRIPTOR_CATALOG if s.packaged_with == "driver"], @@ -125,7 +134,13 @@ def test_cudnn_metadata_matches_wheel_layouts(): assert spec.site_packages_linux == ("nvidia/cudnn/lib",) assert spec.site_packages_windows == WindowsSearchDirs.x64_only("nvidia/cudnn/bin") assert spec.anchor_rel_dirs_windows == WindowsSearchDirs.x64_only("bin/x64", "bin") - assert spec.dependencies == () + assert spec.dependencies == ("cublasLt",) + assert spec.optional_dependencies == ("nvrtc",) + assert spec.windows_root_env_vars == ("CUDNN_PATH",) + assert spec.program_files_rel_dirs_windows == WindowsSearchDirs.x64_only( + "NVIDIA/CUDNN/v9.*/bin/x64", + "NVIDIA/CUDNN/v9.*/bin", + ) assert spec.requires_add_dll_directory diff --git a/cuda_pathfinder/tests/test_load_dl_common.py b/cuda_pathfinder/tests/test_load_dl_common.py new file mode 100644 index 00000000000..61a204c41ee --- /dev/null +++ b/cuda_pathfinder/tests/test_load_dl_common.py @@ -0,0 +1,88 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from cuda.pathfinder._dynamic_libs.descriptor_catalog import DescriptorSpec +from cuda.pathfinder._dynamic_libs.load_dl_common import ( + DynamicLibNotAvailableError, + DynamicLibNotFoundError, + DynamicLibUnknownError, + LoadedDL, + load_dependencies, +) + + +def _loaded(name: str) -> LoadedDL: + return LoadedDL(f"/{name}", False, 1, "test") + + +@pytest.mark.agent_authored(model="gpt-5") +def test_load_dependencies_loads_required_then_optional_dependencies(): + desc = DescriptorSpec( + name="subject", + packaged_with="other", + dependencies=("required",), + optional_dependencies=("optional",), + ) + calls = [] + + def load_func(name): + calls.append(name) + return _loaded(name) + + load_dependencies(desc, load_func) + + assert calls == ["required", "optional"] + + +@pytest.mark.agent_authored(model="gpt-5") +def test_load_dependencies_continues_after_optional_dependency_is_absent(): + desc = DescriptorSpec( + name="subject", + packaged_with="other", + optional_dependencies=("absent", "available"), + ) + calls = [] + + def load_func(name): + calls.append(name) + if name == "absent": + raise DynamicLibNotFoundError(name) + return _loaded(name) + + load_dependencies(desc, load_func) + + assert calls == ["absent", "available"] + + +@pytest.mark.parametrize("error_type", (DynamicLibUnknownError, DynamicLibNotAvailableError, RuntimeError)) +@pytest.mark.agent_authored(model="gpt-5") +def test_load_dependencies_propagates_malformed_or_unloadable_optional_dependency(error_type): + desc = DescriptorSpec(name="subject", packaged_with="other", optional_dependencies=("broken",)) + + def load_func(name): + raise error_type(name) + + with pytest.raises(error_type): + load_dependencies(desc, load_func) + + +@pytest.mark.agent_authored(model="gpt-5") +def test_load_dependencies_keeps_required_dependencies_fail_fast(): + desc = DescriptorSpec( + name="subject", + packaged_with="other", + dependencies=("required",), + optional_dependencies=("optional",), + ) + calls = [] + + def load_func(name): + calls.append(name) + raise DynamicLibNotFoundError(name) + + with pytest.raises(DynamicLibNotFoundError): + load_dependencies(desc, load_func) + + assert calls == ["required"] diff --git a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py index 66ede86c6ad..f08d100a1ae 100644 --- a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib.py @@ -86,7 +86,7 @@ def test_libname_dict_values_are_unique(dict_name): def test_supported_libnames_windows_libnames_requiring_os_add_dll_directory_consistency(): assert not ( set(supported_nvidia_libs.LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY) - - set(supported_nvidia_libs.SUPPORTED_LIBNAMES_WINDOWS) + - set(supported_nvidia_libs.SUPPORTED_WINDOWS_DLLS) ) diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index fc78e22c708..6708b37cbf7 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -25,6 +25,8 @@ _find_lib_dir_using_anchor, find_in_conda, find_in_cuda_path, + find_in_descriptor_environment_roots, + find_in_program_files, find_in_site_packages, run_find_steps, ) @@ -645,6 +647,72 @@ def test_nvvm_windows_checks_binary_arch(self, mocker, tmp_path, target_arch, ma assert any(f"No {target_arch}-compatible PE file" in message for message in ctx.error_messages) +# --------------------------------------------------------------------------- +# Descriptor-specific Windows install roots +# --------------------------------------------------------------------------- + + +class TestDescriptorSpecificWindowsRoots: + @pytest.mark.agent_authored(model="gpt-5") + def test_cudnn_path_finds_archive_layout(self, mocker, tmp_path): + bin_dir = tmp_path / "bin" / "x64" + bin_dir.mkdir(parents=True) + dll = bin_dir / "cudnn64_9.dll" + dll.touch() + mocker.patch.dict(os.environ, {"CUDNN_PATH": str(tmp_path)}) + + result = find_in_descriptor_environment_roots( + _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) + ) + + assert result == FindResult(str(dll), "CUDNN_PATH") + + @pytest.mark.agent_authored(model="gpt-5") + def test_program_files_finds_versioned_cudnn_install(self, mocker, tmp_path): + bin_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.24" / "bin" + bin_dir.mkdir(parents=True) + dll = bin_dir / "cudnn64_9.dll" + dll.touch() + mocker.patch.dict(os.environ, {"PROGRAMFILES": str(tmp_path), "PROGRAMW6432": ""}) + + result = find_in_program_files( + _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) + ) + + assert result == FindResult(str(dll), "ProgramFiles") + + @pytest.mark.agent_authored(model="gpt-5") + def test_program_files_prefers_newest_numeric_cudnn_version(self, mocker, tmp_path): + older_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.9" / "bin" + newer_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.10" / "bin" + older_dir.mkdir(parents=True) + newer_dir.mkdir(parents=True) + (older_dir / "cudnn64_9.dll").touch() + newer_dll = newer_dir / "cudnn64_9.dll" + newer_dll.touch() + mocker.patch.dict(os.environ, {"PROGRAMFILES": str(tmp_path), "PROGRAMW6432": ""}) + + result = find_in_program_files( + _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) + ) + + assert result == FindResult(str(newer_dll), "ProgramFiles") + + @pytest.mark.agent_authored(model="gpt-5") + def test_cudnn_windows_roots_are_not_searched_for_arm64(self, mocker, tmp_path): + bin_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.24" / "bin" + bin_dir.mkdir(parents=True) + (bin_dir / "cudnn64_9.dll").touch() + mocker.patch.dict( + os.environ, + {"CUDNN_PATH": str(tmp_path), "PROGRAMFILES": str(tmp_path), "PROGRAMW6432": ""}, + ) + ctx = _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="arm64")) + + assert find_in_descriptor_environment_roots(ctx) is None + assert find_in_program_files(ctx) is None + + # --------------------------------------------------------------------------- # run_find_steps # --------------------------------------------------------------------------- @@ -687,7 +755,17 @@ def test_early_find_steps_contains_expected(self): assert find_in_conda in EARLY_FIND_STEPS def test_late_find_steps_contains_expected(self): + assert find_in_descriptor_environment_roots in LATE_FIND_STEPS assert find_in_cuda_path in LATE_FIND_STEPS + assert find_in_program_files in LATE_FIND_STEPS + + @pytest.mark.agent_authored(model="gpt-5") + def test_late_find_steps_use_specific_roots_before_generic_roots(self): + assert ( + find_in_descriptor_environment_roots, + find_in_cuda_path, + find_in_program_files, + ) == LATE_FIND_STEPS def test_early_and_late_are_disjoint(self): assert not set(EARLY_FIND_STEPS) & set(LATE_FIND_STEPS) From 25214c0480cfb83e351c762109e6b1e72e03141a Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Fri, 21 Aug 2026 08:50:56 -0700 Subject: [PATCH 3/5] Complete cuDNN and NCCL header discovery --- .../_headers/find_nvidia_headers.py | 33 +++- .../pathfinder/_headers/header_descriptor.py | 15 ++ .../_headers/header_descriptor_catalog.py | 18 +- .../_headers/supported_nvidia_headers.py | 3 +- cuda_pathfinder/pyproject.toml | 4 +- .../tests/test_find_nvidia_headers.py | 164 +++++++++++++++++- 6 files changed, 222 insertions(+), 15 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py index f5f56817141..0a7ee16df8a 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py @@ -18,10 +18,12 @@ HEADER_DESCRIPTORS, platform_include_subdirs, resolve_conda_anchor, + system_install_dir_patterns, ) from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.path_sort import natural_path_sort_key if TYPE_CHECKING: from cuda.pathfinder._headers.header_descriptor import HeaderDescriptor @@ -101,6 +103,18 @@ def find_in_conda(desc: HeaderDescriptor) -> LocatedHeaderDir | None: return None +def find_in_product_roots(desc: HeaderDescriptor) -> LocatedHeaderDir | None: + """Search roots supplied through product-specific environment variables.""" + for env_var in desc.product_root_env_vars: + root = os.environ.get(env_var) + if not root: + continue + result = _locate_in_anchor_layout(desc, root) + if result is not None: + return LocatedHeaderDir(abs_path=result, found_via=env_var) + return None + + def find_in_cuda_path(desc: HeaderDescriptor) -> LocatedHeaderDir | None: """Search ``$CUDA_PATH`` / ``$CUDA_HOME``.""" cuda_home = get_cuda_path_or_home() @@ -136,8 +150,8 @@ def find_via_ctk_root_canary(desc: HeaderDescriptor) -> LocatedHeaderDir | None: def find_in_system_install_dirs(desc: HeaderDescriptor) -> LocatedHeaderDir | None: """Search system install directories (glob patterns).""" - for pattern in desc.system_install_dirs: - for hdr_dir in sorted(glob.glob(pattern), reverse=True): + for pattern in system_install_dir_patterns(desc): + for hdr_dir in sorted(glob.glob(pattern), key=natural_path_sort_key, reverse=True): if _joined_isfile(hdr_dir, desc.header_basename): return LocatedHeaderDir(abs_path=hdr_dir, found_via="supported_install_dir") return None @@ -151,6 +165,7 @@ def find_in_system_install_dirs(desc: HeaderDescriptor) -> LocatedHeaderDir | No FIND_STEPS: tuple[HeaderFindStep, ...] = ( find_in_site_packages, find_in_conda, + find_in_product_roots, find_in_cuda_path, find_via_ctk_root_canary, find_in_system_install_dirs, @@ -190,10 +205,11 @@ def locate_nvidia_header_directory(libname: str) -> LocatedHeaderDir | None: Search order: 1. **NVIDIA Python wheels** — site-packages directories from the descriptor. 2. **Conda environments** — platform-specific conda include layouts. - 3. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. - 4. **CTK root canary probe** — subprocess canary (descriptors with + 3. **Product environment variables** — for example, ``CUDNN_PATH``. + 4. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. + 5. **CTK root canary probe** — subprocess canary (descriptors with ``use_ctk_root_canary=True`` only). - 5. **System install directories** — glob patterns from the descriptor. + 6. **System install directories** — glob patterns from the descriptor. """ desc = HEADER_DESCRIPTORS.get(libname) if desc is None: @@ -218,10 +234,11 @@ def find_nvidia_header_directory(libname: str) -> str | None: Search order: 1. **NVIDIA Python wheels** — site-packages directories from the descriptor. 2. **Conda environments** — platform-specific conda include layouts. - 3. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. - 4. **CTK root canary probe** — subprocess canary (descriptors with + 3. **Product environment variables** — for example, ``CUDNN_PATH``. + 4. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. + 5. **CTK root canary probe** — subprocess canary (descriptors with ``use_ctk_root_canary=True`` only). - 5. **System install directories** — glob patterns from the descriptor. + 6. **System install directories** — glob patterns from the descriptor. """ found = locate_nvidia_header_directory(libname) return found.abs_path if found else None diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor.py b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor.py index 609dcab7184..3b108f618a5 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor.py @@ -12,6 +12,7 @@ import glob import os +import sysconfig from typing import TypeAlias, cast from cuda.pathfinder._headers.header_descriptor_catalog import ( @@ -37,6 +38,20 @@ def platform_include_subdirs(desc: HeaderDescriptor) -> tuple[str, ...]: return cast(tuple[str, ...], desc.include_subdirs) +def system_install_dir_patterns(desc: HeaderDescriptor) -> tuple[str, ...]: + """Return platform-aware, expanded system include-directory patterns.""" + patterns: list[str] = [] + if IS_WINDOWS: + patterns.extend(os.path.expandvars(pattern) for pattern in desc.system_install_dirs_windows) + return tuple(patterns) + if desc.use_linux_multiarch_include_dir: + multiarch = sysconfig.get_config_var("MULTIARCH") + if isinstance(multiarch, str) and multiarch: + patterns.append(os.path.join("/usr/include", multiarch)) + patterns.extend(os.path.expandvars(pattern) for pattern in desc.system_install_dirs) + return tuple(patterns) + + def resolve_conda_anchor(desc: HeaderDescriptor, conda_prefix: str) -> str | None: """Resolve the conda anchor point for header search on the current platform. diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py index 36e182a3c8c..14e4a3a28e6 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/header_descriptor_catalog.py @@ -21,12 +21,18 @@ class HeaderDescriptorSpec: available_on_windows: bool = True # Relative path(s) from anchor point to the include directory. anchor_include_rel_dirs: tuple[str, ...] = ("include",) + # Product-specific environment variables whose values are anchor points. + product_root_env_vars: tuple[str, ...] = () # Subdirectories within the include dir to check before the include dir itself. include_subdirs: tuple[str, ...] = () # Windows-only additional subdirectories within the include dir. include_subdirs_windows: tuple[str, ...] = () - # System install directories (glob patterns). + # Linux system install directories (glob patterns; environment variables are expanded). system_install_dirs: tuple[str, ...] = () + # Windows system install directories (glob patterns; environment variables are expanded). + system_install_dirs_windows: tuple[str, ...] = () + # Whether to search /usr/include/ before system_install_dirs. + use_linux_multiarch_include_dir: bool = False # Whether to use targets//include layout for conda on Linux. conda_targets_layout: bool = True # Whether to attempt CTK-root canary probing (spawns a subprocess). @@ -155,6 +161,13 @@ class HeaderDescriptorSpec: packaged_with="other", header_basename="cudnn.h", site_packages_dirs=("nvidia/cudnn/include",), + product_root_env_vars=("CUDNN_PATH",), + system_install_dirs=( + "/usr/include", + "/usr/local/include", + ), + system_install_dirs_windows=("${ProgramFiles}/NVIDIA/CUDNN/v9.*/include",), + use_linux_multiarch_include_dir=True, conda_targets_layout=False, use_ctk_root_canary=False, ), @@ -258,6 +271,9 @@ class HeaderDescriptorSpec: header_basename="nccl.h", site_packages_dirs=("nvidia/nccl/include",), available_on_windows=False, + anchor_include_rel_dirs=("include", "build/include"), + product_root_env_vars=("NCCL_HOME",), + system_install_dirs=("/usr/include", "/usr/local/include"), conda_targets_layout=False, use_ctk_root_canary=False, ), diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/supported_nvidia_headers.py b/cuda_pathfinder/cuda/pathfinder/_headers/supported_nvidia_headers.py index c7b40834e67..34736baad8d 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/supported_nvidia_headers.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/supported_nvidia_headers.py @@ -12,6 +12,7 @@ from typing import Final +from cuda.pathfinder._headers.header_descriptor import system_install_dir_patterns from cuda.pathfinder._headers.header_descriptor_catalog import HEADER_DESCRIPTOR_CATALOG from cuda.pathfinder._utils.platform_aware import IS_WINDOWS @@ -77,5 +78,5 @@ } SUPPORTED_INSTALL_DIRS_NON_CTK: Final[dict[str, tuple[str, ...]]] = { - desc.name: desc.system_install_dirs for desc in _NON_CTK_DESCRIPTORS if desc.system_install_dirs + desc.name: patterns for desc in _NON_CTK_DESCRIPTORS if (patterns := system_install_dir_patterns(desc)) } diff --git a/cuda_pathfinder/pyproject.toml b/cuda_pathfinder/pyproject.toml index 022d868568c..c5840e70ec9 100644 --- a/cuda_pathfinder/pyproject.toml +++ b/cuda_pathfinder/pyproject.toml @@ -24,7 +24,7 @@ cu12 = [ "cuquantum-cu12; sys_platform != 'win32'", "cutensor-cu12", "nvidia-cublasmp-cu12; sys_platform != 'win32'", - "nvidia-cudnn-cu12", + "nvidia-cudnn-cu12>=9,<10", "nvidia-cudss-cu12", "nvidia-cufftmp-cu12; sys_platform != 'win32'", "nvidia-cusolvermp-cu12; sys_platform != 'win32'", @@ -40,7 +40,7 @@ cu13 = [ "cutensor-cu13", "nvidia-cublasmp-cu13; sys_platform != 'win32'", "nvidia-cudla; platform_system == 'Linux' and platform_machine == 'aarch64'", - "nvidia-cudnn-cu13", + "nvidia-cudnn-cu13>=9,<10", "nvidia-cudss-cu13", "nvidia-cufftmp-cu13; sys_platform != 'win32'", "nvidia-cusolvermp-cu13; sys_platform != 'win32'", diff --git a/cuda_pathfinder/tests/test_find_nvidia_headers.py b/cuda_pathfinder/tests/test_find_nvidia_headers.py index 9e36e8acc56..34002098d4d 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/cuda_pathfinder/tests/test_find_nvidia_headers.py @@ -21,8 +21,11 @@ import pytest from conftest import skip_if_missing_libnvcudla_so +from packaging.requirements import Requirement +from packaging.version import Version import cuda.pathfinder._headers.find_nvidia_headers as find_nvidia_headers_module +import cuda.pathfinder._headers.header_descriptor as header_descriptor_module from cuda.pathfinder import LocatedHeaderDir, find_nvidia_header_directory, locate_nvidia_header_directory from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( _resolve_system_loaded_abs_path_in_subprocess, @@ -44,7 +47,7 @@ NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = { "cudensitymat": r"^cudensitymat-.*$", - "cudnn": r"^nvidia-cudnn-.*$", + "cudnn": r"^nvidia-cudnn-cu(?:12|13)$", "cupauliprop": r"^cupauliprop-.*$", "cusolverMp": r"^nvidia-cusolvermp-.*$", "cusparseLt": r"^nvidia-cusparselt-.*$", @@ -70,6 +73,8 @@ def _located_hdr_dir_asserts(located_hdr_dir): assert located_hdr_dir.found_via in ( "site-packages", "conda", + "CUDNN_PATH", + "NCCL_HOME", "CUDA_PATH", "system-ctk-root", "supported_install_dir", @@ -86,6 +91,13 @@ def test_cudnn_and_nccl_header_metadata_matches_wheel_layouts(): cudnn = HEADER_DESCRIPTORS["cudnn"] assert cudnn.header_basename == "cudnn.h" assert cudnn.site_packages_dirs == ("nvidia/cudnn/include",) + assert cudnn.product_root_env_vars == ("CUDNN_PATH",) + assert cudnn.system_install_dirs == ( + "/usr/include", + "/usr/local/include", + ) + assert cudnn.system_install_dirs_windows == ("${ProgramFiles}/NVIDIA/CUDNN/v9.*/include",) + assert cudnn.use_linux_multiarch_include_dir assert cudnn.available_on_linux assert cudnn.available_on_windows assert not cudnn.conda_targets_layout @@ -96,6 +108,9 @@ def test_cudnn_and_nccl_header_metadata_matches_wheel_layouts(): assert nccl.site_packages_dirs == ("nvidia/nccl/include",) assert nccl.available_on_linux assert not nccl.available_on_windows + assert nccl.anchor_include_rel_dirs == ("include", "build/include") + assert nccl.product_root_env_vars == ("NCCL_HOME",) + assert nccl.system_install_dirs == ("/usr/include", "/usr/local/include") assert not nccl.conda_targets_layout assert not nccl.use_ctk_root_canary @@ -108,6 +123,39 @@ def have_distribution_for(libname: str) -> bool: ) +@pytest.mark.parametrize( + ("distribution_name", "expected"), + [ + ("nvidia-cudnn-cu12", True), + ("nvidia-cudnn-cu13", True), + ("nvidia-cudnn-frontend", False), + ("nvidia-cudnn-jit-cu12", False), + ("nvidia-cudnn-jit-cu13", False), + ], +) +@pytest.mark.agent_authored(model="gpt-5") +def test_cudnn_distribution_pattern_only_matches_backend_wheels(distribution_name, expected): + pattern = re.compile(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES["cudnn"]) + + assert bool(pattern.match(distribution_name)) is expected + + +@pytest.mark.parametrize( + "requirement", + ["nvidia-cudnn-cu12>=9,<10", "nvidia-cudnn-cu13>=9,<10"], +) +@pytest.mark.agent_authored(model="gpt-5") +def test_cudnn_test_dependencies_are_bounded_to_major_nine(requirement): + pyproject_text = (Path(__file__).parents[1] / "pyproject.toml").read_text(encoding="utf-8") + parsed = Requirement(requirement) + + assert f'"{requirement}",' in pyproject_text + assert Version("8.9") not in parsed.specifier + assert Version("9.0") in parsed.specifier + assert Version("9.24.0.43") in parsed.specifier + assert Version("10.0") not in parsed.specifier + + @pytest.fixture def clear_locate_nvidia_header_cache(): locate_nvidia_header_directory.cache_clear() @@ -140,6 +188,114 @@ def _fake_cudart_canary_abs_path(ctk_root: Path) -> str: return str(ctk_root / "lib64" / "libcudart.so.13") +@pytest.mark.parametrize( + ("libname", "env_var", "include_rel_dir"), + [ + ("cudnn", "CUDNN_PATH", "include"), + ("nccl", "NCCL_HOME", "build/include"), + ], +) +@pytest.mark.usefixtures("clear_locate_nvidia_header_cache") +@pytest.mark.agent_authored(model="gpt-5") +def test_locate_non_ctk_headers_uses_product_root(tmp_path, monkeypatch, mocker, libname, env_var, include_rel_dir): + product_root = tmp_path / libname + header_dir = product_root / include_rel_dir + header_dir.mkdir(parents=True) + (header_dir / HEADER_DESCRIPTORS[libname].header_basename).touch() + + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.delenv("CUDNN_PATH", raising=False) + monkeypatch.delenv("NCCL_HOME", raising=False) + monkeypatch.setenv(env_var, str(product_root)) + monkeypatch.setenv("CUDA_HOME", str(tmp_path / "unused-cuda-home")) + monkeypatch.delenv("CUDA_PATH", raising=False) + mocker.patch.object(find_nvidia_headers_module, "find_sub_dirs_all_sitepackages", return_value=[]) + + located_hdr_dir = locate_nvidia_header_directory(libname) + + assert located_hdr_dir is not None + assert located_hdr_dir.abs_path == str(header_dir) + assert located_hdr_dir.found_via == env_var + + +@pytest.mark.parametrize( + ("libname", "system_pattern"), + [ + ("cudnn", "/usr/include"), + ("cudnn", "/usr/local/include"), + ("nccl", "/usr/include"), + ("nccl", "/usr/local/include"), + ], +) +@pytest.mark.agent_authored(model="gpt-5") +def test_find_in_system_install_dirs_uses_native_linux_layouts(tmp_path, mocker, libname, system_pattern): + header_dir = tmp_path / libname + header_dir.mkdir() + (header_dir / HEADER_DESCRIPTORS[libname].header_basename).touch() + mocker.patch.object(header_descriptor_module, "IS_WINDOWS", False) + glob_mock = mocker.patch.object( + find_nvidia_headers_module.glob, + "glob", + side_effect=lambda pattern: [str(header_dir)] if pattern == system_pattern else [], + ) + + located_hdr_dir = find_nvidia_headers_module.find_in_system_install_dirs(HEADER_DESCRIPTORS[libname]) + + assert located_hdr_dir is not None + assert located_hdr_dir.abs_path == str(header_dir) + assert located_hdr_dir.found_via == "supported_install_dir" + assert any(call.args == (system_pattern,) for call in glob_mock.call_args_list) + + +@pytest.mark.agent_authored(model="gpt-5") +def test_find_in_system_install_dirs_uses_running_linux_multiarch(tmp_path, mocker): + header_dir = tmp_path / "multiarch" + header_dir.mkdir() + (header_dir / "cudnn.h").touch() + mocker.patch.object(header_descriptor_module, "IS_WINDOWS", False) + mocker.patch.object(header_descriptor_module.sysconfig, "get_config_var", return_value="x86_64-linux-gnu") + expected_pattern = "/usr/include/x86_64-linux-gnu" + glob_mock = mocker.patch.object( + find_nvidia_headers_module.glob, + "glob", + side_effect=lambda pattern: [str(header_dir)] if pattern == expected_pattern else [], + ) + + located_hdr_dir = find_nvidia_headers_module.find_in_system_install_dirs(HEADER_DESCRIPTORS["cudnn"]) + + assert located_hdr_dir is not None + assert located_hdr_dir.abs_path == str(header_dir) + assert located_hdr_dir.found_via == "supported_install_dir" + assert glob_mock.call_args_list[0].args == (expected_pattern,) + + +@pytest.mark.agent_authored(model="gpt-5") +def test_find_in_system_install_dirs_expands_program_files_and_prefers_newest_cudnn(tmp_path, monkeypatch, mocker): + program_files = tmp_path / "Program Files" + older_header_dir = program_files / "NVIDIA" / "CUDNN" / "v9.9" / "include" + newer_header_dir = program_files / "NVIDIA" / "CUDNN" / "v9.10" / "include" + for header_dir in (older_header_dir, newer_header_dir): + header_dir.mkdir(parents=True) + (header_dir / "cudnn.h").touch() + mocker.patch.object(header_descriptor_module, "IS_WINDOWS", True) + monkeypatch.setenv("ProgramFiles", str(program_files)) + expected_pattern = str(program_files / "NVIDIA" / "CUDNN" / "v9.*" / "include") + glob_mock = mocker.patch.object( + find_nvidia_headers_module.glob, + "glob", + side_effect=lambda pattern: ( + [str(older_header_dir), str(newer_header_dir)] if pattern == expected_pattern else [] + ), + ) + + located_hdr_dir = find_nvidia_headers_module.find_in_system_install_dirs(HEADER_DESCRIPTORS["cudnn"]) + + assert located_hdr_dir is not None + assert located_hdr_dir.abs_path == str(newer_header_dir) + assert located_hdr_dir.found_via == "supported_install_dir" + glob_mock.assert_called_once_with(expected_pattern) + + # TODO: remove the Python 3.15 guard once 3.15 is officially supported _CUTLASS_SKIP = pytest.mark.skipif( sys.version_info >= (3, 15), @@ -168,7 +324,9 @@ def test_locate_non_ctk_headers(info_summary_append, libname): assert "site-packages" in Path(hdr_dir).parts elif STRICTNESS == "all_must_work": assert hdr_dir is not None - if conda_prefix := os.environ.get("CONDA_PREFIX"): + if located_hdr_dir.found_via in HEADER_DESCRIPTORS[libname].product_root_env_vars: + assert hdr_dir.startswith(os.environ[located_hdr_dir.found_via]) + elif conda_prefix := os.environ.get("CONDA_PREFIX"): assert hdr_dir.startswith(conda_prefix) else: inst_dirs = SUPPORTED_INSTALL_DIRS_NON_CTK.get(libname) @@ -176,7 +334,7 @@ def test_locate_non_ctk_headers(info_summary_append, libname): for inst_dir in inst_dirs: # Absolute glob pattern: Path.glob needs a separate base dir, # and the wildcard is not pinned to the last component. - globbed = glob.glob(inst_dir) + globbed = glob.glob(os.path.expandvars(inst_dir)) if hdr_dir in globbed: break else: From 955b5e7404d41bcba04c0008ceec4d81eeec527e Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Fri, 21 Aug 2026 10:01:25 -0700 Subject: [PATCH 4/5] Unify dynamic library installation root discovery --- .../_dynamic_libs/descriptor_catalog.py | 11 +-- .../_dynamic_libs/load_nvidia_dynamic_lib.py | 5 +- .../_dynamic_libs/search_platform.py | 48 ++++++++--- .../pathfinder/_dynamic_libs/search_steps.py | 86 +++++++++++-------- .../tests/test_descriptor_catalog.py | 8 +- cuda_pathfinder/tests/test_search_steps.py | 49 +++++++---- 6 files changed, 126 insertions(+), 81 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index 961dfb98758..75122d7cfcc 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -68,8 +68,8 @@ class DescriptorSpec: optional_dependencies: tuple[str, ...] = () anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib") anchor_rel_dirs_windows: WindowsSearchDirs = DEFAULT_WINDOWS_CTK_ANCHOR_DIRS - windows_root_env_vars: tuple[str, ...] = () - program_files_rel_dirs_windows: WindowsSearchDirs = WindowsSearchDirs() + install_root_env_vars_windows: tuple[str, ...] = () + program_files_root_globs_windows: WindowsSearchDirs = WindowsSearchDirs() ctk_root_canary_anchor_libnames: tuple[str, ...] = () requires_add_dll_directory: bool = False requires_rtld_deepbind: bool = False @@ -425,11 +425,8 @@ class DescriptorSpec: dependencies=("cublasLt",), optional_dependencies=("nvrtc",), anchor_rel_dirs_windows=WindowsSearchDirs.x64_only("bin/x64", "bin"), - windows_root_env_vars=("CUDNN_PATH",), - program_files_rel_dirs_windows=WindowsSearchDirs.x64_only( - "NVIDIA/CUDNN/v9.*/bin/x64", - "NVIDIA/CUDNN/v9.*/bin", - ), + install_root_env_vars_windows=("CUDNN_PATH",), + program_files_root_globs_windows=WindowsSearchDirs.x64_only("NVIDIA/CUDNN/v9.*"), requires_add_dll_directory=True, ), DescriptorSpec( diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py index 93073c21a7e..7f231021ea0 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py @@ -284,8 +284,9 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: 5. **Windows Program Files (configured libraries only)** - - Search descriptor-configured standalone installation layouts, - such as versioned cuDNN directories under ``ProgramFiles``. + - Search descriptor-configured standalone installation roots, such + as versioned cuDNN directories under ``ProgramFiles``, using the + same per-library relative layout as other root-based searches. 6. **CTK root canary probe (discoverable libs only)** diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index e71a2c3b652..25c034b7edb 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py @@ -63,8 +63,13 @@ def _find_so_in_rel_dirs( return None -def _find_dll_under_dir(dirpath: str, file_wild: str, target_arch: str | None = None) -> str | None: - for path in sorted(glob.glob(os.path.join(dirpath, file_wild))): +def _find_descriptor_dll_under_dir( + dirpath: str, + desc: LibDescriptor, + target_arch: str | None = None, +) -> str | None: + for dll_basename in reversed(cast(tuple[str, ...], desc.windows_dlls)): + path = os.path.join(dirpath, dll_basename) if not os.path.isfile(path): continue if is_suppressed_dll_file(os.path.basename(path)): @@ -77,15 +82,18 @@ def _find_dll_under_dir(dirpath: str, file_wild: str, target_arch: str | None = def _find_dll_in_rel_dirs( rel_dirs: tuple[str, ...], + desc: LibDescriptor, + target_arch: str, lib_searched_for: str, error_messages: list[str], attachments: list[str], ) -> str | None: sub_dirs_searched: list[tuple[str, ...]] = [] + checked_arch = target_arch if desc.requires_windows_binary_arch_check else None for rel_dir in rel_dirs: sub_dir = PurePath(rel_dir).parts for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): - dll_name = _find_dll_under_dir(abs_dir, lib_searched_for) + dll_name = _find_descriptor_dll_under_dir(abs_dir, desc, checked_arch) if dll_name is not None: return dll_name sub_dirs_searched.append(sub_dir) @@ -103,13 +111,14 @@ def conda_anchor_point(self, conda_prefix: str) -> str: ... def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ... - def root_environment_variables(self, desc: LibDescriptor) -> tuple[str, ...]: ... + def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]: ... - def program_files_dir_globs(self, desc: LibDescriptor) -> tuple[str, ...]: ... + def program_files_root_globs(self, desc: LibDescriptor) -> tuple[str, ...]: ... def find_in_site_packages( self, rel_dirs: tuple[str, ...], + desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], @@ -139,15 +148,16 @@ def conda_anchor_point(self, conda_prefix: str) -> str: def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: return cast(tuple[str, ...], desc.anchor_rel_dirs_linux) - def root_environment_variables(self, _desc: LibDescriptor) -> tuple[str, ...]: + def install_root_env_vars(self, _desc: LibDescriptor) -> tuple[str, ...]: return () - def program_files_dir_globs(self, _desc: LibDescriptor) -> tuple[str, ...]: + def program_files_root_globs(self, _desc: LibDescriptor) -> tuple[str, ...]: return () def find_in_site_packages( self, rel_dirs: tuple[str, ...], + _desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], @@ -202,24 +212,34 @@ def conda_anchor_point(self, conda_prefix: str) -> str: def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: return cast(tuple[str, ...], desc.anchor_rel_dirs_windows.for_arch(self.target_arch)) - def root_environment_variables(self, desc: LibDescriptor) -> tuple[str, ...]: - return cast(tuple[str, ...], desc.windows_root_env_vars) + def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]: + if self.target_arch not in desc.supported_windows_arch: + return () + return cast(tuple[str, ...], desc.install_root_env_vars_windows) - def program_files_dir_globs(self, desc: LibDescriptor) -> tuple[str, ...]: + def program_files_root_globs(self, desc: LibDescriptor) -> tuple[str, ...]: program_files = os.environ.get("PROGRAMW6432") or os.environ.get("PROGRAMFILES") if not program_files: return () - rel_dirs = desc.program_files_rel_dirs_windows.for_arch(self.target_arch) - return tuple(os.path.join(program_files, rel_dir) for rel_dir in rel_dirs) + rel_globs = desc.program_files_root_globs_windows.for_arch(self.target_arch) + return tuple(os.path.join(program_files, rel_glob) for rel_glob in rel_globs) def find_in_site_packages( self, rel_dirs: tuple[str, ...], + desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], ) -> str | None: - return _find_dll_in_rel_dirs(rel_dirs, lib_searched_for, error_messages, attachments) + return _find_dll_in_rel_dirs( + rel_dirs, + desc, + self.target_arch, + lib_searched_for, + error_messages, + attachments, + ) def find_in_lib_dir( self, @@ -231,7 +251,7 @@ def find_in_lib_dir( ) -> str | None: file_wild = desc.name + "*.dll" target_arch = self.target_arch if desc.requires_windows_binary_arch_check else None - dll_name = _find_dll_under_dir(lib_dir, file_wild, target_arch) + dll_name = _find_descriptor_dll_under_dir(lib_dir, desc, target_arch) if dll_name is not None: return dll_name if target_arch is None: diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py index 82dcdee4003..891468eed4b 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py @@ -21,7 +21,7 @@ import glob import os -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from typing import NoReturn, cast @@ -71,14 +71,22 @@ def raise_not_found(self) -> NoReturn: FindStep = Callable[[SearchContext], FindResult | None] -def _find_lib_dir_using_anchor(desc: LibDescriptor, platform: SearchPlatform, anchor_point: str) -> str | None: - """Find the library directory under *anchor_point* using the descriptor's relative paths.""" +def _iter_lib_dirs_using_anchor( + desc: LibDescriptor, + platform: SearchPlatform, + anchor_point: str, +) -> Iterator[str]: + """Yield existing library directories under *anchor_point* in descriptor order.""" rel_dirs = platform.anchor_rel_dirs(desc) for rel_path in rel_dirs: for dirname in sorted(glob.glob(os.path.join(anchor_point, rel_path))): if os.path.isdir(dirname): - return os.path.normpath(dirname) - return None + yield os.path.normpath(dirname) + + +def _find_lib_dir_using_anchor(desc: LibDescriptor, platform: SearchPlatform, anchor_point: str) -> str | None: + """Find the first library directory under *anchor_point*.""" + return next(_iter_lib_dirs_using_anchor(desc, platform, anchor_point), None) def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: @@ -97,6 +105,15 @@ def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: ) +def _find_under_root(ctx: SearchContext, root: str, found_via: str) -> FindResult | None: + """Resolve descriptor anchors under *root*, then find the requested library.""" + for lib_dir in _iter_lib_dirs_using_anchor(ctx.desc, ctx.platform, root): + abs_path = _find_using_lib_dir(ctx, lib_dir) + if abs_path is not None: + return FindResult(abs_path, found_via) + return None + + def _derive_ctk_root_linux(resolved_lib_path: str) -> str | None: """Derive CTK root from Linux canary path. @@ -148,11 +165,7 @@ def derive_ctk_root(resolved_lib_path: str) -> str | None: def find_via_ctk_root(ctx: SearchContext, ctk_root: str) -> FindResult | None: """Find a library under a previously derived CTK root.""" - lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, ctk_root) - abs_path = _find_using_lib_dir(ctx, lib_dir) - if abs_path is None: - return None - return FindResult(abs_path, "system-ctk-root") + return _find_under_root(ctx, ctk_root, "system-ctk-root") # --------------------------------------------------------------------------- @@ -165,7 +178,13 @@ def find_in_site_packages(ctx: SearchContext) -> FindResult | None: rel_dirs = ctx.platform.site_packages_rel_dirs(ctx.desc) if not rel_dirs: return None - abs_path = ctx.platform.find_in_site_packages(rel_dirs, ctx.lib_searched_for, ctx.error_messages, ctx.attachments) + abs_path = ctx.platform.find_in_site_packages( + rel_dirs, + ctx.desc, + ctx.lib_searched_for, + ctx.error_messages, + ctx.attachments, + ) if abs_path is not None: return FindResult(abs_path, "site-packages") return None @@ -177,23 +196,18 @@ def find_in_conda(ctx: SearchContext) -> FindResult | None: if not conda_prefix: return None anchor = ctx.platform.conda_anchor_point(conda_prefix) - lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, anchor) - abs_path = _find_using_lib_dir(ctx, lib_dir) - if abs_path is not None: - return FindResult(abs_path, "conda") - return None + return _find_under_root(ctx, anchor, "conda") -def find_in_descriptor_environment_roots(ctx: SearchContext) -> FindResult | None: - """Search roots named by descriptor-specific environment variables.""" - for env_var in ctx.platform.root_environment_variables(ctx.desc): +def find_in_install_root_env_vars(ctx: SearchContext) -> FindResult | None: + """Search installation roots named by descriptor-specific environment variables.""" + for env_var in ctx.platform.install_root_env_vars(ctx.desc): root = os.environ.get(env_var) if not root: continue - lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, root) - abs_path = _find_using_lib_dir(ctx, lib_dir) - if abs_path is not None: - return FindResult(abs_path, env_var) + result = _find_under_root(ctx, root, env_var) + if result is not None: + return result return None @@ -211,22 +225,18 @@ def find_in_cuda_path(ctx: SearchContext) -> FindResult | None: cuda_home = get_cuda_path_or_home() if cuda_home is None: return None - lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, cuda_home) - abs_path = _find_using_lib_dir(ctx, lib_dir) - if abs_path is not None: - return FindResult(abs_path, "CUDA_PATH") - return None + return _find_under_root(ctx, cuda_home, "CUDA_PATH") -def find_in_program_files(ctx: SearchContext) -> FindResult | None: - """Search descriptor-configured installation directories under Program Files.""" - for dir_glob in ctx.platform.program_files_dir_globs(ctx.desc): - for lib_dir in sorted(glob.glob(dir_glob), key=natural_path_sort_key, reverse=True): - if not os.path.isdir(lib_dir): +def find_in_program_files_roots(ctx: SearchContext) -> FindResult | None: + """Search descriptor-configured installation roots under Program Files.""" + for root_glob in ctx.platform.program_files_root_globs(ctx.desc): + for root in sorted(glob.glob(root_glob), key=natural_path_sort_key, reverse=True): + if not os.path.isdir(root): continue - abs_path = _find_using_lib_dir(ctx, os.path.normpath(lib_dir)) - if abs_path is not None: - return FindResult(abs_path, "ProgramFiles") + result = _find_under_root(ctx, os.path.normpath(root), "ProgramFiles") + if result is not None: + return result return None @@ -239,9 +249,9 @@ def find_in_program_files(ctx: SearchContext) -> FindResult | None: #: Find steps that run after system search fails. LATE_FIND_STEPS: tuple[FindStep, ...] = ( - find_in_descriptor_environment_roots, + find_in_install_root_env_vars, find_in_cuda_path, - find_in_program_files, + find_in_program_files_roots, ) diff --git a/cuda_pathfinder/tests/test_descriptor_catalog.py b/cuda_pathfinder/tests/test_descriptor_catalog.py index 9bf59e0a276..d4b0b65748e 100644 --- a/cuda_pathfinder/tests/test_descriptor_catalog.py +++ b/cuda_pathfinder/tests/test_descriptor_catalog.py @@ -112,6 +112,7 @@ def test_windows_search_dirs_do_not_include_unsupported_arches(spec: DescriptorS if arch not in spec.supported_windows_arch: assert not spec.site_packages_windows.for_arch(arch) assert not spec.anchor_rel_dirs_windows.for_arch(arch) + assert not spec.program_files_root_globs_windows.for_arch(arch) @pytest.mark.agent_authored(model="gpt-5") @@ -136,11 +137,8 @@ def test_cudnn_metadata_matches_wheel_layouts(): assert spec.anchor_rel_dirs_windows == WindowsSearchDirs.x64_only("bin/x64", "bin") assert spec.dependencies == ("cublasLt",) assert spec.optional_dependencies == ("nvrtc",) - assert spec.windows_root_env_vars == ("CUDNN_PATH",) - assert spec.program_files_rel_dirs_windows == WindowsSearchDirs.x64_only( - "NVIDIA/CUDNN/v9.*/bin/x64", - "NVIDIA/CUDNN/v9.*/bin", - ) + assert spec.install_root_env_vars_windows == ("CUDNN_PATH",) + assert spec.program_files_root_globs_windows == WindowsSearchDirs.x64_only("NVIDIA/CUDNN/v9.*") assert spec.requires_add_dll_directory diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index 6708b37cbf7..abb67869b8a 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -25,8 +25,8 @@ _find_lib_dir_using_anchor, find_in_conda, find_in_cuda_path, - find_in_descriptor_environment_roots, - find_in_program_files, + find_in_install_root_env_vars, + find_in_program_files_roots, find_in_site_packages, run_find_steps, ) @@ -309,6 +309,23 @@ def test_found_windows(self, mocker, tmp_path): assert result.abs_path == str(dll) assert result.found_via == "site-packages" + @pytest.mark.agent_authored(model="gpt-5") + def test_windows_requires_an_exact_descriptor_dll_name(self, mocker, tmp_path): + bin_dir = tmp_path / "nvidia" / "cudnn" / "bin" + bin_dir.mkdir(parents=True) + (bin_dir / "cudnn_adv64_9.dll").touch() + + mocker.patch( + f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages", + return_value=[str(bin_dir)], + ) + + result = find_in_site_packages( + _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) + ) + + assert result is None + @pytest.mark.agent_authored(model="gpt-5") def test_found_windows_arm64_prefers_cuda13_arch_dir_to_cuda12(self, mocker, tmp_path): x86_64_dir = tmp_path / "nvidia" / "cu13" / "bin" / "x86_64" @@ -652,7 +669,7 @@ def test_nvvm_windows_checks_binary_arch(self, mocker, tmp_path, target_arch, ma # --------------------------------------------------------------------------- -class TestDescriptorSpecificWindowsRoots: +class TestWindowsInstallRoots: @pytest.mark.agent_authored(model="gpt-5") def test_cudnn_path_finds_archive_layout(self, mocker, tmp_path): bin_dir = tmp_path / "bin" / "x64" @@ -661,7 +678,7 @@ def test_cudnn_path_finds_archive_layout(self, mocker, tmp_path): dll.touch() mocker.patch.dict(os.environ, {"CUDNN_PATH": str(tmp_path)}) - result = find_in_descriptor_environment_roots( + result = find_in_install_root_env_vars( _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) ) @@ -670,20 +687,22 @@ def test_cudnn_path_finds_archive_layout(self, mocker, tmp_path): @pytest.mark.agent_authored(model="gpt-5") def test_program_files_finds_versioned_cudnn_install(self, mocker, tmp_path): bin_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.24" / "bin" - bin_dir.mkdir(parents=True) + x64_bin_dir = bin_dir / "x64" + x64_bin_dir.mkdir(parents=True) + (x64_bin_dir / "cudnn_adv64_9.dll").touch() dll = bin_dir / "cudnn64_9.dll" dll.touch() mocker.patch.dict(os.environ, {"PROGRAMFILES": str(tmp_path), "PROGRAMW6432": ""}) - result = find_in_program_files( + result = find_in_program_files_roots( _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) ) assert result == FindResult(str(dll), "ProgramFiles") @pytest.mark.agent_authored(model="gpt-5") - def test_program_files_prefers_newest_numeric_cudnn_version(self, mocker, tmp_path): - older_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.9" / "bin" + def test_program_files_prefers_newest_numeric_cudnn_version_across_layouts(self, mocker, tmp_path): + older_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.9" / "bin" / "x64" newer_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.10" / "bin" older_dir.mkdir(parents=True) newer_dir.mkdir(parents=True) @@ -692,7 +711,7 @@ def test_program_files_prefers_newest_numeric_cudnn_version(self, mocker, tmp_pa newer_dll.touch() mocker.patch.dict(os.environ, {"PROGRAMFILES": str(tmp_path), "PROGRAMW6432": ""}) - result = find_in_program_files( + result = find_in_program_files_roots( _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) ) @@ -709,8 +728,8 @@ def test_cudnn_windows_roots_are_not_searched_for_arm64(self, mocker, tmp_path): ) ctx = _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="arm64")) - assert find_in_descriptor_environment_roots(ctx) is None - assert find_in_program_files(ctx) is None + assert find_in_install_root_env_vars(ctx) is None + assert find_in_program_files_roots(ctx) is None # --------------------------------------------------------------------------- @@ -755,16 +774,16 @@ def test_early_find_steps_contains_expected(self): assert find_in_conda in EARLY_FIND_STEPS def test_late_find_steps_contains_expected(self): - assert find_in_descriptor_environment_roots in LATE_FIND_STEPS + assert find_in_install_root_env_vars in LATE_FIND_STEPS assert find_in_cuda_path in LATE_FIND_STEPS - assert find_in_program_files in LATE_FIND_STEPS + assert find_in_program_files_roots in LATE_FIND_STEPS @pytest.mark.agent_authored(model="gpt-5") def test_late_find_steps_use_specific_roots_before_generic_roots(self): assert ( - find_in_descriptor_environment_roots, + find_in_install_root_env_vars, find_in_cuda_path, - find_in_program_files, + find_in_program_files_roots, ) == LATE_FIND_STEPS def test_early_and_late_are_disjoint(self): From ddaa78fc01138188526a1099f53e4c7e8fc33380 Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Fri, 21 Aug 2026 10:51:18 -0700 Subject: [PATCH 5/5] Add cuDNN Windows ARM64 archive discovery --- .../_dynamic_libs/descriptor_catalog.py | 9 ++- .../_dynamic_libs/load_nvidia_dynamic_lib.py | 14 ++--- .../_dynamic_libs/search_platform.py | 10 +++ .../pathfinder/_dynamic_libs/search_steps.py | 41 +++++++++---- .../tests/test_descriptor_catalog.py | 23 ++++++- ...st_load_nvidia_dynamic_lib_using_mocker.py | 42 ++++++++++++- cuda_pathfinder/tests/test_search_steps.py | 61 ++++++++++++++++--- 7 files changed, 167 insertions(+), 33 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index 75122d7cfcc..93fd9464b2e 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -69,6 +69,7 @@ class DescriptorSpec: anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib") anchor_rel_dirs_windows: WindowsSearchDirs = DEFAULT_WINDOWS_CTK_ANCHOR_DIRS install_root_env_vars_windows: tuple[str, ...] = () + install_root_env_rel_dirs_windows: WindowsSearchDirs = WindowsSearchDirs() program_files_root_globs_windows: WindowsSearchDirs = WindowsSearchDirs() ctk_root_canary_anchor_libnames: tuple[str, ...] = () requires_add_dll_directory: bool = False @@ -419,13 +420,19 @@ class DescriptorSpec: packaged_with="other", linux_sonames=("libcudnn.so.9",), windows_dlls=("cudnn64_9.dll",), - supported_windows_arch=("x64",), + supported_windows_arch=("x64", "arm64"), site_packages_linux=("nvidia/cudnn/lib",), site_packages_windows=WindowsSearchDirs.x64_only("nvidia/cudnn/bin"), dependencies=("cublasLt",), optional_dependencies=("nvrtc",), + # The ARM64 layout is verified only for the standalone archive rooted + # at CUDNN_PATH, not for conda, CUDA_PATH, or Program Files installs. anchor_rel_dirs_windows=WindowsSearchDirs.x64_only("bin/x64", "bin"), install_root_env_vars_windows=("CUDNN_PATH",), + install_root_env_rel_dirs_windows=WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64",), + ), program_files_root_globs_windows=WindowsSearchDirs.x64_only("NVIDIA/CUDNN/v9.*"), requires_add_dll_directory=True, ), diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py index 7f231021ea0..a6982d5a553 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py @@ -276,17 +276,17 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: 4. **Environment variables** - First search library-specific roots declared by the descriptor, - such as ``CUDNN_PATH``, then use ``CUDA_PATH`` or ``CUDA_HOME`` - (in that order). On Windows, ``CUDA_PATH`` is the typical way - system-installed CTK DLLs are located. Note that the NVIDIA CTK - installer automatically adds ``CUDA_PATH`` to the system-wide - environment. + such as ``CUDNN_PATH``, using their architecture-specific product + layouts. Then use ``CUDA_PATH`` or ``CUDA_HOME`` (in that order). + On Windows, ``CUDA_PATH`` is the typical way system-installed CTK + DLLs are located. Note that the NVIDIA CTK installer automatically + adds ``CUDA_PATH`` to the system-wide environment. 5. **Windows Program Files (configured libraries only)** - Search descriptor-configured standalone installation roots, such - as versioned cuDNN directories under ``ProgramFiles``, using the - same per-library relative layout as other root-based searches. + as versioned x64 cuDNN directories under ``ProgramFiles``, using + the general per-library anchor layout. 6. **CTK root canary probe (discoverable libs only)** diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py index 25c034b7edb..6505470973e 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py @@ -113,6 +113,8 @@ def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ... def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]: ... + def install_root_env_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ... + def program_files_root_globs(self, desc: LibDescriptor) -> tuple[str, ...]: ... def find_in_site_packages( @@ -151,6 +153,9 @@ def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: def install_root_env_vars(self, _desc: LibDescriptor) -> tuple[str, ...]: return () + def install_root_env_rel_dirs(self, _desc: LibDescriptor) -> tuple[str, ...]: + return () + def program_files_root_globs(self, _desc: LibDescriptor) -> tuple[str, ...]: return () @@ -217,6 +222,11 @@ def install_root_env_vars(self, desc: LibDescriptor) -> tuple[str, ...]: return () return cast(tuple[str, ...], desc.install_root_env_vars_windows) + def install_root_env_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: + if self.target_arch not in desc.supported_windows_arch: + return () + return cast(tuple[str, ...], desc.install_root_env_rel_dirs_windows.for_arch(self.target_arch)) + def program_files_root_globs(self, desc: LibDescriptor) -> tuple[str, ...]: program_files = os.environ.get("PROGRAMW6432") or os.environ.get("PROGRAMFILES") if not program_files: diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py index 891468eed4b..b8b840c5f88 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py @@ -71,17 +71,21 @@ def raise_not_found(self) -> NoReturn: FindStep = Callable[[SearchContext], FindResult | None] +def _iter_lib_dirs(root: str, rel_dirs: tuple[str, ...]) -> Iterator[str]: + """Yield existing library directories under *root* in descriptor order.""" + for rel_path in rel_dirs: + for dirname in sorted(glob.glob(os.path.join(root, rel_path))): + if os.path.isdir(dirname): + yield os.path.normpath(dirname) + + def _iter_lib_dirs_using_anchor( desc: LibDescriptor, platform: SearchPlatform, anchor_point: str, ) -> Iterator[str]: """Yield existing library directories under *anchor_point* in descriptor order.""" - rel_dirs = platform.anchor_rel_dirs(desc) - for rel_path in rel_dirs: - for dirname in sorted(glob.glob(os.path.join(anchor_point, rel_path))): - if os.path.isdir(dirname): - yield os.path.normpath(dirname) + yield from _iter_lib_dirs(anchor_point, platform.anchor_rel_dirs(desc)) def _find_lib_dir_using_anchor(desc: LibDescriptor, platform: SearchPlatform, anchor_point: str) -> str | None: @@ -105,15 +109,25 @@ def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: ) -def _find_under_root(ctx: SearchContext, root: str, found_via: str) -> FindResult | None: - """Resolve descriptor anchors under *root*, then find the requested library.""" - for lib_dir in _iter_lib_dirs_using_anchor(ctx.desc, ctx.platform, root): +def _find_under_root( + ctx: SearchContext, + root: str, + rel_dirs: tuple[str, ...], + found_via: str, +) -> FindResult | None: + """Resolve *rel_dirs* under *root*, then find the requested library.""" + for lib_dir in _iter_lib_dirs(root, rel_dirs): abs_path = _find_using_lib_dir(ctx, lib_dir) if abs_path is not None: return FindResult(abs_path, found_via) return None +def _find_under_anchor_root(ctx: SearchContext, root: str, found_via: str) -> FindResult | None: + """Resolve the descriptor's general anchors under *root*.""" + return _find_under_root(ctx, root, ctx.platform.anchor_rel_dirs(ctx.desc), found_via) + + def _derive_ctk_root_linux(resolved_lib_path: str) -> str | None: """Derive CTK root from Linux canary path. @@ -165,7 +179,7 @@ def derive_ctk_root(resolved_lib_path: str) -> str | None: def find_via_ctk_root(ctx: SearchContext, ctk_root: str) -> FindResult | None: """Find a library under a previously derived CTK root.""" - return _find_under_root(ctx, ctk_root, "system-ctk-root") + return _find_under_anchor_root(ctx, ctk_root, "system-ctk-root") # --------------------------------------------------------------------------- @@ -196,16 +210,17 @@ def find_in_conda(ctx: SearchContext) -> FindResult | None: if not conda_prefix: return None anchor = ctx.platform.conda_anchor_point(conda_prefix) - return _find_under_root(ctx, anchor, "conda") + return _find_under_anchor_root(ctx, anchor, "conda") def find_in_install_root_env_vars(ctx: SearchContext) -> FindResult | None: """Search installation roots named by descriptor-specific environment variables.""" + rel_dirs = ctx.platform.install_root_env_rel_dirs(ctx.desc) for env_var in ctx.platform.install_root_env_vars(ctx.desc): root = os.environ.get(env_var) if not root: continue - result = _find_under_root(ctx, root, env_var) + result = _find_under_root(ctx, root, rel_dirs, env_var) if result is not None: return result return None @@ -225,7 +240,7 @@ def find_in_cuda_path(ctx: SearchContext) -> FindResult | None: cuda_home = get_cuda_path_or_home() if cuda_home is None: return None - return _find_under_root(ctx, cuda_home, "CUDA_PATH") + return _find_under_anchor_root(ctx, cuda_home, "CUDA_PATH") def find_in_program_files_roots(ctx: SearchContext) -> FindResult | None: @@ -234,7 +249,7 @@ def find_in_program_files_roots(ctx: SearchContext) -> FindResult | None: for root in sorted(glob.glob(root_glob), key=natural_path_sort_key, reverse=True): if not os.path.isdir(root): continue - result = _find_under_root(ctx, os.path.normpath(root), "ProgramFiles") + result = _find_under_anchor_root(ctx, os.path.normpath(root), "ProgramFiles") if result is not None: return result return None diff --git a/cuda_pathfinder/tests/test_descriptor_catalog.py b/cuda_pathfinder/tests/test_descriptor_catalog.py index d4b0b65748e..5d7fde130c0 100644 --- a/cuda_pathfinder/tests/test_descriptor_catalog.py +++ b/cuda_pathfinder/tests/test_descriptor_catalog.py @@ -112,9 +112,24 @@ def test_windows_search_dirs_do_not_include_unsupported_arches(spec: DescriptorS if arch not in spec.supported_windows_arch: assert not spec.site_packages_windows.for_arch(arch) assert not spec.anchor_rel_dirs_windows.for_arch(arch) + assert not spec.install_root_env_rel_dirs_windows.for_arch(arch) assert not spec.program_files_root_globs_windows.for_arch(arch) +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +@pytest.mark.agent_authored(model="gpt-5") +def test_install_root_env_metadata_is_complete(spec: DescriptorSpec): + has_env_vars = bool(spec.install_root_env_vars_windows) + has_rel_dirs = any(spec.install_root_env_rel_dirs_windows.for_arch(arch) for arch in _VALID_WINDOWS_ARCHES) + + assert has_env_vars == has_rel_dirs, f"{spec.name} must define both installation-root env vars and relative dirs" + if has_env_vars: + for arch in spec.supported_windows_arch: + assert spec.install_root_env_rel_dirs_windows.for_arch(arch), ( + f"{spec.name} exposes installation-root env vars without {arch} relative dirs" + ) + + @pytest.mark.agent_authored(model="gpt-5") def test_cusparselt_windows_metadata_matches_wheel_layouts(): spec = _CATALOG_BY_NAME["cusparseLt"] @@ -126,18 +141,22 @@ def test_cusparselt_windows_metadata_matches_wheel_layouts(): @pytest.mark.agent_authored(model="gpt-5") -def test_cudnn_metadata_matches_wheel_layouts(): +def test_cudnn_metadata_matches_supported_layouts(): spec = _CATALOG_BY_NAME["cudnn"] assert spec.packaged_with == "other" assert spec.linux_sonames == ("libcudnn.so.9",) assert spec.windows_dlls == ("cudnn64_9.dll",) - assert spec.supported_windows_arch == ("x64",) + assert spec.supported_windows_arch == ("x64", "arm64") assert spec.site_packages_linux == ("nvidia/cudnn/lib",) assert spec.site_packages_windows == WindowsSearchDirs.x64_only("nvidia/cudnn/bin") assert spec.anchor_rel_dirs_windows == WindowsSearchDirs.x64_only("bin/x64", "bin") assert spec.dependencies == ("cublasLt",) assert spec.optional_dependencies == ("nvrtc",) assert spec.install_root_env_vars_windows == ("CUDNN_PATH",) + assert spec.install_root_env_rel_dirs_windows == WindowsSearchDirs( + x64=("bin/x64", "bin"), + arm64=("bin/arm64",), + ) assert spec.program_files_root_globs_windows == WindowsSearchDirs.x64_only("NVIDIA/CUDNN/v9.*") assert spec.requires_add_dll_directory diff --git a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib_using_mocker.py b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib_using_mocker.py index f46ad43356b..09c8bd3ea61 100644 --- a/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib_using_mocker.py +++ b/cuda_pathfinder/tests/test_load_nvidia_dynamic_lib_using_mocker.py @@ -10,7 +10,8 @@ _load_lib_no_cache, _resolve_system_loaded_abs_path_in_subprocess, ) -from cuda.pathfinder._dynamic_libs.search_steps import EARLY_FIND_STEPS +from cuda.pathfinder._dynamic_libs.search_platform import WindowsSearchPlatform +from cuda.pathfinder._dynamic_libs.search_steps import EARLY_FIND_STEPS, SearchContext from cuda.pathfinder._utils.platform_aware import IS_WINDOWS _MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib" @@ -45,6 +46,45 @@ def _create_cupti_in_ctk(ctk_root): return cupti_lib +# --------------------------------------------------------------------------- +# cuDNN Windows ARM64 archive +# --------------------------------------------------------------------------- + + +@pytest.mark.agent_authored(model="gpt-5") +def test_cudnn_arm64_archive_layout_reaches_loader(tmp_path, mocker, monkeypatch): + bin_dir = tmp_path / "bin" / "arm64" + bin_dir.mkdir(parents=True) + dll = bin_dir / "cudnn64_9.dll" + dll.touch() + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.setenv("CUDNN_PATH", str(tmp_path)) + + desc = load_mod.LIB_DESCRIPTORS["cudnn"] + ctx = SearchContext(desc, platform=WindowsSearchPlatform(target_arch="arm64")) + mocker.patch(f"{_MODULE}.SearchContext", return_value=ctx) + mocker.patch.object(load_mod.LOADER, "check_if_already_loaded_from_elsewhere", return_value=None) + + def _load_dependency(name): + if name == "nvrtc": + raise DynamicLibNotFoundError(name) + return _make_loaded_dl(name, "dependency") + + load_dependency = mocker.patch(f"{_MODULE}.load_nvidia_dynamic_lib", side_effect=_load_dependency) + mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None) + load_with_abs_path = mocker.patch.object( + load_mod.LOADER, + "load_with_abs_path", + side_effect=lambda _desc, path, via: _make_loaded_dl(path, via), + ) + + result = _load_lib_no_cache("cudnn") + + assert result == _make_loaded_dl(str(dll), "CUDNN_PATH") + assert [call.args for call in load_dependency.call_args_list] == [("cublasLt",), ("nvrtc",)] + load_with_abs_path.assert_called_once_with(desc, str(dll), "CUDNN_PATH") + + # --------------------------------------------------------------------------- # Conda tests # Note: Site-packages and CTK are covered by real CI tests. diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py index abb67869b8a..7d8e32ec2e1 100644 --- a/cuda_pathfinder/tests/test_search_steps.py +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -670,20 +670,49 @@ def test_nvvm_windows_checks_binary_arch(self, mocker, tmp_path, target_arch, ma class TestWindowsInstallRoots: + @pytest.mark.parametrize( + ("target_arch", "archive_bin_rel_dir"), + [ + ("x64", "bin/x64"), + ("x64", "bin"), + ("arm64", "bin/arm64"), + ], + ) @pytest.mark.agent_authored(model="gpt-5") - def test_cudnn_path_finds_archive_layout(self, mocker, tmp_path): - bin_dir = tmp_path / "bin" / "x64" + def test_cudnn_path_finds_supported_archive_layouts(self, mocker, tmp_path, target_arch, archive_bin_rel_dir): + bin_dir = tmp_path / archive_bin_rel_dir bin_dir.mkdir(parents=True) dll = bin_dir / "cudnn64_9.dll" dll.touch() mocker.patch.dict(os.environ, {"CUDNN_PATH": str(tmp_path)}) result = find_in_install_root_env_vars( - _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="x64")) + _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch=target_arch)) ) assert result == FindResult(str(dll), "CUDNN_PATH") + @pytest.mark.parametrize( + ("target_arch", "other_arch_bin_rel_dir"), + [ + ("x64", "bin/arm64"), + ("arm64", "bin/x64"), + ("arm64", "bin"), + ], + ) + @pytest.mark.agent_authored(model="gpt-5") + def test_cudnn_path_does_not_cross_architectures(self, mocker, tmp_path, target_arch, other_arch_bin_rel_dir): + bin_dir = tmp_path / other_arch_bin_rel_dir + bin_dir.mkdir(parents=True) + (bin_dir / "cudnn64_9.dll").touch() + mocker.patch.dict(os.environ, {"CUDNN_PATH": str(tmp_path)}) + + result = find_in_install_root_env_vars( + _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch=target_arch)) + ) + + assert result is None + @pytest.mark.agent_authored(model="gpt-5") def test_program_files_finds_versioned_cudnn_install(self, mocker, tmp_path): bin_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.24" / "bin" @@ -718,18 +747,32 @@ def test_program_files_prefers_newest_numeric_cudnn_version_across_layouts(self, assert result == FindResult(str(newer_dll), "ProgramFiles") @pytest.mark.agent_authored(model="gpt-5") - def test_cudnn_windows_roots_are_not_searched_for_arm64(self, mocker, tmp_path): - bin_dir = tmp_path / "NVIDIA" / "CUDNN" / "v9.24" / "bin" - bin_dir.mkdir(parents=True) - (bin_dir / "cudnn64_9.dll").touch() + def test_cudnn_arm64_archive_layout_is_not_assumed_for_other_roots(self, mocker, tmp_path): + conda_root = tmp_path / "conda" + cuda_root = tmp_path / "cuda" + program_files_root = tmp_path / "Program Files" + for bin_dir in ( + conda_root / "Library" / "bin" / "arm64", + cuda_root / "bin" / "arm64", + program_files_root / "NVIDIA" / "CUDNN" / "v9.25" / "bin" / "arm64", + ): + bin_dir.mkdir(parents=True) + (bin_dir / "cudnn64_9.dll").touch() mocker.patch.dict( os.environ, - {"CUDNN_PATH": str(tmp_path), "PROGRAMFILES": str(tmp_path), "PROGRAMW6432": ""}, + { + "CONDA_PREFIX": str(conda_root), + "PROGRAMFILES": str(program_files_root), + "PROGRAMW6432": "", + }, ) + mocker.patch(f"{_STEPS_MOD}.get_cuda_path_or_home", return_value=str(cuda_root)) ctx = _ctx(LIB_DESCRIPTORS["cudnn"], platform=WindowsSearchPlatform(target_arch="arm64")) - assert find_in_install_root_env_vars(ctx) is None + assert find_in_conda(ctx) is None + assert find_in_cuda_path(ctx) is None assert find_in_program_files_roots(ctx) is None + assert ctx.platform.site_packages_rel_dirs(ctx.desc) == () # ---------------------------------------------------------------------------