diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py index e39046eec70..93fd9464b2e 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -65,8 +65,12 @@ 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 + 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 requires_rtld_deepbind: bool = False @@ -411,6 +415,27 @@ 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", "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, + ), DescriptorSpec( name="cusolverMp", packaged_with="other", 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..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 @@ -275,12 +275,20 @@ 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 + - First search library-specific roots declared by the descriptor, + 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. **CTK root canary probe (discoverable libs only)** + 5. **Windows Program Files (configured libraries only)** + + - Search descriptor-configured standalone installation roots, such + as versioned x64 cuDNN directories under ``ProgramFiles``, using + the general per-library anchor layout. + + 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 +306,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..6505470973e 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,9 +111,16 @@ def conda_anchor_point(self, conda_prefix: str) -> str: ... 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( self, rel_dirs: tuple[str, ...], + desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], @@ -135,9 +150,19 @@ 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 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 () + def find_in_site_packages( self, rel_dirs: tuple[str, ...], + _desc: LibDescriptor, lib_searched_for: str, error_messages: list[str], attachments: list[str], @@ -192,14 +217,39 @@ 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 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 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: + return () + 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, @@ -211,7 +261,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 5901094fcaa..b8b840c5f88 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 @@ -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 @@ -70,14 +71,26 @@ 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.""" - rel_dirs = platform.anchor_rel_dirs(desc) +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(anchor_point, rel_path))): + for dirname in sorted(glob.glob(os.path.join(root, rel_path))): if os.path.isdir(dirname): - return os.path.normpath(dirname) - return None + 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.""" + 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: + """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: @@ -96,6 +109,25 @@ def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: ) +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. @@ -147,11 +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.""" - 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_anchor_root(ctx, ctk_root, "system-ctk-root") # --------------------------------------------------------------------------- @@ -164,7 +192,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 @@ -176,10 +210,19 @@ 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 _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, rel_dirs, env_var) + if result is not None: + return result return None @@ -197,10 +240,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 _find_under_anchor_root(ctx, cuda_home, "CUDA_PATH") + + +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 + result = _find_under_anchor_root(ctx, os.path.normpath(root), "ProgramFiles") + if result is not None: + return result return None @@ -212,7 +263,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_install_root_env_vars, + find_in_cuda_path, + find_in_program_files_roots, +) # --------------------------------------------------------------------------- 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 b364e224e7e..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). @@ -150,6 +156,21 @@ class HeaderDescriptorSpec: # ----------------------------------------------------------------------- # Third-party / separately packaged headers # ----------------------------------------------------------------------- + HeaderDescriptorSpec( + name="cudnn", + 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, + ), HeaderDescriptorSpec( name="cusolverMp", packaged_with="other", @@ -244,6 +265,18 @@ 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, + 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, + ), HeaderDescriptorSpec( name="nvshmem", packaged_with="other", 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/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/pyproject.toml b/cuda_pathfinder/pyproject.toml index 4fd7a5e3edf..c5840e70ec9 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>=9,<10", "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>=9,<10", "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..5d7fde130c0 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"], @@ -103,6 +112,22 @@ 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") @@ -115,6 +140,27 @@ def test_cusparselt_windows_metadata_matches_wheel_layouts(): ) +@pytest.mark.agent_authored(model="gpt-5") +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", "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 + + @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..34002098d4d 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/cuda_pathfinder/tests/test_find_nvidia_headers.py @@ -21,12 +21,16 @@ 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, ) +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 +47,7 @@ NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES = { "cudensitymat": r"^cudensitymat-.*$", + "cudnn": r"^nvidia-cudnn-cu(?:12|13)$", "cupauliprop": r"^cupauliprop-.*$", "cusolverMp": r"^nvidia-cusolvermp-.*$", "cusparseLt": r"^nvidia-cusparselt-.*$", @@ -53,6 +58,7 @@ "custatevec": r"^custatevec-.*$", "cutlass": r"^nvidia-cutlass$", "mathdx": r"^nvidia-libmathdx-.*$", + "nccl": r"^nvidia-nccl-.*$", "nvshmem": r"^nvidia-nvshmem-.*$", } @@ -67,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", @@ -78,6 +86,35 @@ 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.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 + 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 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 + + @functools.cache def have_distribution_for(libname: str) -> bool: pattern = re.compile(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES[libname]) @@ -86,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() @@ -118,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), @@ -146,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) @@ -154,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: 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_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 fc78e22c708..7d8e32ec2e1 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_install_root_env_vars, + find_in_program_files_roots, find_in_site_packages, run_find_steps, ) @@ -307,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" @@ -645,6 +664,117 @@ 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 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_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=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" + 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_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_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) + (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_roots( + _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_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, + { + "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_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) == () + + # --------------------------------------------------------------------------- # run_find_steps # --------------------------------------------------------------------------- @@ -687,7 +817,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_install_root_env_vars in LATE_FIND_STEPS assert find_in_cuda_path 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_install_root_env_vars, + find_in_cuda_path, + find_in_program_files_roots, + ) == LATE_FIND_STEPS def test_early_and_late_are_disjoint(self): assert not set(EARLY_FIND_STEPS) & set(LATE_FIND_STEPS) 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" )