diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..06fd7da3ed --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,77 @@ +# cuda_pathfinder agent instructions + +You are working on `cuda_pathfinder`, a Python sub-package of the +[cuda-python](https://github.com/NVIDIA/cuda-python) monorepo. It finds and +loads NVIDIA dynamic libraries (CTK, third-party, and driver) across Linux and +Windows. + +## Workspace + +The workspace root is `cuda_pathfinder/` inside the monorepo. Use the +`working_directory` parameter on the Shell tool when you need the monorepo root +(one level up). + +## Conventions + +- **Python**: all source is pure Python (no Cython in this sub-package). +- **Testing**: `pytest` with `pytest-mock` (`mocker` fixture). Use + `spawned_process_runner` for real-loading tests that need process isolation + (dynamic linker state leaks across tests otherwise). Use the + `info_summary_append` fixture to emit `INFO` lines visible in CI/QA logs. +- **STRICTNESS env var**: `CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS` + controls whether missing libs are tolerated (`see_what_works`, default) or + fatal (`all_must_work`). +- **Formatting/linting**: rely on pre-commit (runs automatically on commit). Do + not run formatters manually. +- **Imports**: use `from cuda.pathfinder._dynamic_libs...` for internal imports + in tests; public API is `from cuda.pathfinder import load_nvidia_dynamic_lib`. + +## Testing guidelines + +- **Real tests over mocks**: mocks are fine for hard-to-reach branches (e.g. + 24-bit Python), but every loading path must also have a real-loading test that + runs in a spawned child process. Track results with `INFO` lines so CI logs + show what actually loaded. +- **No real lib names in negative tests**: when parametrizing unsupported / + invalid libnames, use obviously fake names (`"bogus"`, `"not_a_real_lib"`) + to avoid confusion when searching the codebase. +- **`functools.cache` awareness**: `load_nvidia_dynamic_lib` is cached. Tests + that patch internals it depends on must call + `load_nvidia_dynamic_lib.cache_clear()` first, or use a child process for + isolation. + +## Key modules + +- `cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py` -- main entry + point and dispatch logic (CTK vs driver). +- `cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py` -- canonical + registry of sonames, DLLs, site-packages paths, and dependencies. +- `cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py` -- CTK search + cascade (site-packages, conda, CUDA_HOME). +- `tests/child_load_nvidia_dynamic_lib_helper.py` -- lightweight helper + imported by spawned child processes (avoids re-importing the full test + module). + +### Fix all code review findings from lib-descriptor-refactor review + +**Request:** Fix all 8 findings from the external code review. + +**Actions (in worktree `cuda_pathfinder_refactor`):** +1. `search_steps.py`: Restored `os.path.normpath(dirname)` in + `_find_lib_dir_using_anchor` (regression from pre-refactor fix). Added + `NoReturn` annotation to `raise_not_found`. +2. `search_platform.py`: Guarded `os.listdir(lib_dir)` in + `WindowsSearchPlatform.find_in_lib_dir` with `os.path.isdir` check to + prevent crash on missing directory. +3. `test_descriptor_catalog.py`: Rewrote tautological tests as structural + invariant tests (uniqueness, valid names, strategy values, dep graph, + soname/dll format, driver lib constraints). 237 new parametrized cases. +4. `platform_loader.py`: Eliminated `WindowsLoader`/`LinuxLoader` boilerplate + classes — assign the platform module directly as `LOADER`. Removed stale + `type: ignore`. +5. `descriptor_catalog.py`: Trimmed default-valued fields from all entries, + added `# ---` section comments (CTK / third-party / driver). +6. `load_nvidia_dynamic_lib.py`: Fixed import layout — `TYPE_CHECKING` block + now properly separated after unconditional imports. + +All 742 tests pass, all pre-commit hooks green. diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/canary_probe_subprocess.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/canary_probe_subprocess.py index 3f1525dd48..95ed61a0ea 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/canary_probe_subprocess.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/canary_probe_subprocess.py @@ -4,18 +4,17 @@ import json +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS - -if IS_WINDOWS: - from cuda.pathfinder._dynamic_libs.load_dl_windows import load_with_system_search -else: - from cuda.pathfinder._dynamic_libs.load_dl_linux import load_with_system_search +from cuda.pathfinder._dynamic_libs.platform_loader import LOADER def _probe_canary_abs_path(libname: str) -> str | None: + desc = LIB_DESCRIPTORS.get(libname) + if desc is None: + raise ValueError(f"Unsupported canary library name: {libname!r}") try: - loaded: LoadedDL | None = load_with_system_search(libname) + loaded: LoadedDL | None = LOADER.load_with_system_search(desc) except DynamicLibNotFoundError: return None if loaded is None: diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py new file mode 100644 index 0000000000..7bc2477401 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/descriptor_catalog.py @@ -0,0 +1,364 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canonical authored descriptor catalog for dynamic libraries.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +Strategy = Literal["ctk", "other", "driver"] + + +@dataclass(frozen=True, slots=True) +class DescriptorSpec: + name: str + strategy: Strategy + linux_sonames: tuple[str, ...] = () + windows_dlls: tuple[str, ...] = () + site_packages_linux: tuple[str, ...] = () + site_packages_windows: tuple[str, ...] = () + dependencies: tuple[str, ...] = () + anchor_rel_dirs_linux: tuple[str, ...] = ("lib64", "lib") + anchor_rel_dirs_windows: tuple[str, ...] = ("bin/x64", "bin") + ctk_root_canary_anchor_libnames: tuple[str, ...] = () + requires_add_dll_directory: bool = False + requires_rtld_deepbind: bool = False + + +DESCRIPTOR_CATALOG: tuple[DescriptorSpec, ...] = ( + # ----------------------------------------------------------------------- + # CTK (CUDA Toolkit) libraries + # ----------------------------------------------------------------------- + DescriptorSpec( + name="cudart", + strategy="ctk", + linux_sonames=("libcudart.so.12", "libcudart.so.13"), + windows_dlls=("cudart64_12.dll", "cudart64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_runtime/bin"), + ), + DescriptorSpec( + name="nvfatbin", + strategy="ctk", + linux_sonames=("libnvfatbin.so.12", "libnvfatbin.so.13"), + windows_dlls=("nvfatbin_120_0.dll", "nvfatbin_130_0.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/nvfatbin/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/nvfatbin/bin"), + ), + DescriptorSpec( + name="nvJitLink", + strategy="ctk", + linux_sonames=("libnvJitLink.so.12", "libnvJitLink.so.13"), + windows_dlls=("nvJitLink_120_0.dll", "nvJitLink_130_0.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/nvjitlink/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/nvjitlink/bin"), + ), + DescriptorSpec( + name="nvrtc", + strategy="ctk", + linux_sonames=("libnvrtc.so.12", "libnvrtc.so.13"), + windows_dlls=("nvrtc64_120_0.dll", "nvrtc64_130_0.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_nvrtc/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvrtc/bin"), + requires_add_dll_directory=True, + ), + DescriptorSpec( + name="nvvm", + strategy="ctk", + linux_sonames=("libnvvm.so.4",), + windows_dlls=("nvvm64.dll", "nvvm64_40_0.dll", "nvvm70.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cuda_nvcc/nvvm/lib64"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvcc/nvvm/bin"), + anchor_rel_dirs_linux=("nvvm/lib64",), + anchor_rel_dirs_windows=("nvvm/bin/*", "nvvm/bin"), + ctk_root_canary_anchor_libnames=("cudart",), + ), + DescriptorSpec( + name="cublas", + strategy="ctk", + linux_sonames=("libcublas.so.12", "libcublas.so.13"), + windows_dlls=("cublas64_12.dll", "cublas64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cublas/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), + dependencies=("cublasLt",), + ), + DescriptorSpec( + name="cublasLt", + strategy="ctk", + linux_sonames=("libcublasLt.so.12", "libcublasLt.so.13"), + windows_dlls=("cublasLt64_12.dll", "cublasLt64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cublas/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), + ), + DescriptorSpec( + name="cufft", + strategy="ctk", + linux_sonames=("libcufft.so.11", "libcufft.so.12"), + windows_dlls=("cufft64_11.dll", "cufft64_12.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cufft/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"), + requires_add_dll_directory=True, + ), + DescriptorSpec( + name="cufftw", + strategy="ctk", + linux_sonames=("libcufftw.so.11", "libcufftw.so.12"), + windows_dlls=("cufftw64_11.dll", "cufftw64_12.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cufft/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"), + dependencies=("cufft",), + ), + DescriptorSpec( + name="curand", + strategy="ctk", + linux_sonames=("libcurand.so.10",), + windows_dlls=("curand64_10.dll",), + site_packages_linux=("nvidia/cu13/lib", "nvidia/curand/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/curand/bin"), + ), + DescriptorSpec( + name="cusolver", + strategy="ctk", + linux_sonames=("libcusolver.so.11", "libcusolver.so.12"), + windows_dlls=("cusolver64_11.dll", "cusolver64_12.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cusolver/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"), + dependencies=("nvJitLink", "cusparse", "cublasLt", "cublas"), + ), + DescriptorSpec( + name="cusolverMg", + strategy="ctk", + linux_sonames=("libcusolverMg.so.11", "libcusolverMg.so.12"), + windows_dlls=("cusolverMg64_11.dll", "cusolverMg64_12.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cusolver/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"), + dependencies=("nvJitLink", "cublasLt", "cublas"), + ), + DescriptorSpec( + name="cusparse", + strategy="ctk", + linux_sonames=("libcusparse.so.12",), + windows_dlls=("cusparse64_12.dll",), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cusparse/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cusparse/bin"), + dependencies=("nvJitLink",), + ), + DescriptorSpec( + name="nppc", + strategy="ctk", + linux_sonames=("libnppc.so.12", "libnppc.so.13"), + windows_dlls=("nppc64_12.dll", "nppc64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + ), + DescriptorSpec( + name="nppial", + strategy="ctk", + linux_sonames=("libnppial.so.12", "libnppial.so.13"), + windows_dlls=("nppial64_12.dll", "nppial64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppicc", + strategy="ctk", + linux_sonames=("libnppicc.so.12", "libnppicc.so.13"), + windows_dlls=("nppicc64_12.dll", "nppicc64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppidei", + strategy="ctk", + linux_sonames=("libnppidei.so.12", "libnppidei.so.13"), + windows_dlls=("nppidei64_12.dll", "nppidei64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppif", + strategy="ctk", + linux_sonames=("libnppif.so.12", "libnppif.so.13"), + windows_dlls=("nppif64_12.dll", "nppif64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppig", + strategy="ctk", + linux_sonames=("libnppig.so.12", "libnppig.so.13"), + windows_dlls=("nppig64_12.dll", "nppig64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppim", + strategy="ctk", + linux_sonames=("libnppim.so.12", "libnppim.so.13"), + windows_dlls=("nppim64_12.dll", "nppim64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppist", + strategy="ctk", + linux_sonames=("libnppist.so.12", "libnppist.so.13"), + windows_dlls=("nppist64_12.dll", "nppist64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppisu", + strategy="ctk", + linux_sonames=("libnppisu.so.12", "libnppisu.so.13"), + windows_dlls=("nppisu64_12.dll", "nppisu64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nppitc", + strategy="ctk", + linux_sonames=("libnppitc.so.12", "libnppitc.so.13"), + windows_dlls=("nppitc64_12.dll", "nppitc64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="npps", + strategy="ctk", + linux_sonames=("libnpps.so.12", "libnpps.so.13"), + windows_dlls=("npps64_12.dll", "npps64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/npp/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), + dependencies=("nppc",), + ), + DescriptorSpec( + name="nvblas", + strategy="ctk", + linux_sonames=("libnvblas.so.12", "libnvblas.so.13"), + windows_dlls=("nvblas64_12.dll", "nvblas64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cublas/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), + dependencies=("cublas", "cublasLt"), + ), + DescriptorSpec( + name="nvjpeg", + strategy="ctk", + linux_sonames=("libnvjpeg.so.12", "libnvjpeg.so.13"), + windows_dlls=("nvjpeg64_12.dll", "nvjpeg64_13.dll"), + site_packages_linux=("nvidia/cu13/lib", "nvidia/nvjpeg/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/nvjpeg/bin"), + ), + DescriptorSpec( + name="cufile", + strategy="ctk", + linux_sonames=("libcufile.so.0",), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cufile/lib"), + ), + # ----------------------------------------------------------------------- + # Third-party / separately packaged libraries + # ----------------------------------------------------------------------- + DescriptorSpec( + name="cublasmp", + strategy="other", + linux_sonames=("libcublasmp.so.0",), + site_packages_linux=("nvidia/cublasmp/cu13/lib", "nvidia/cublasmp/cu12/lib"), + dependencies=("cublas", "cublasLt", "nvshmem_host"), + ), + DescriptorSpec( + name="cufftMp", + strategy="other", + linux_sonames=("libcufftMp.so.12", "libcufftMp.so.11"), + site_packages_linux=("nvidia/cufftmp/cu13/lib", "nvidia/cufftmp/cu12/lib"), + dependencies=("nvshmem_host",), + requires_rtld_deepbind=True, + ), + DescriptorSpec( + name="mathdx", + strategy="other", + linux_sonames=("libmathdx.so.0",), + windows_dlls=("mathdx64_0.dll",), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cu12/lib"), + site_packages_windows=("nvidia/cu13/bin/x86_64", "nvidia/cu12/bin"), + dependencies=("nvrtc",), + ), + DescriptorSpec( + name="cudss", + strategy="other", + linux_sonames=("libcudss.so.0",), + windows_dlls=("cudss64_0.dll",), + site_packages_linux=("nvidia/cu13/lib", "nvidia/cu12/lib"), + site_packages_windows=("nvidia/cu13/bin", "nvidia/cu12/bin"), + dependencies=("cublas", "cublasLt"), + ), + DescriptorSpec( + name="cusparseLt", + strategy="other", + linux_sonames=("libcusparseLt.so.0",), + windows_dlls=("cusparseLt.dll",), + site_packages_linux=("nvidia/cusparselt/lib",), + site_packages_windows=("nvidia/cusparselt/bin",), + ), + DescriptorSpec( + name="cutensor", + strategy="other", + linux_sonames=("libcutensor.so.2",), + windows_dlls=("cutensor.dll",), + site_packages_linux=("cutensor/lib",), + site_packages_windows=("cutensor/bin",), + dependencies=("cublasLt",), + ), + DescriptorSpec( + name="cutensorMg", + strategy="other", + linux_sonames=("libcutensorMg.so.2",), + windows_dlls=("cutensorMg.dll",), + site_packages_linux=("cutensor/lib",), + site_packages_windows=("cutensor/bin",), + dependencies=("cutensor", "cublasLt"), + ), + DescriptorSpec( + name="nccl", + strategy="other", + linux_sonames=("libnccl.so.2",), + site_packages_linux=("nvidia/nccl/lib",), + ), + DescriptorSpec( + name="nvpl_fftw", + strategy="other", + linux_sonames=("libnvpl_fftw.so.0",), + site_packages_linux=("nvpl/lib",), + ), + DescriptorSpec( + name="nvshmem_host", + strategy="other", + linux_sonames=("libnvshmem_host.so.3",), + site_packages_linux=("nvidia/nvshmem/lib",), + ), + # ----------------------------------------------------------------------- + # Driver libraries (system-search only, no CTK cascade) + # ----------------------------------------------------------------------- + DescriptorSpec( + name="cuda", + strategy="driver", + linux_sonames=("libcuda.so.1",), + windows_dlls=("nvcuda.dll",), + ), + DescriptorSpec( + name="nvml", + strategy="driver", + linux_sonames=("libnvidia-ml.so.1",), + windows_dlls=("nvml.dll",), + ), +) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py deleted file mode 100644 index b93523f36a..0000000000 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py +++ /dev/null @@ -1,277 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import glob -import os -from collections.abc import Sequence - -from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError -from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( - SITE_PACKAGES_LIBDIRS_LINUX, - SITE_PACKAGES_LIBDIRS_WINDOWS, - is_suppressed_dll_file, -) -from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path -from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS - - -def _no_such_file_in_sub_dirs( - sub_dirs: Sequence[str], file_wild: str, error_messages: list[str], attachments: list[str] -) -> None: - error_messages.append(f"No such file: {file_wild}") - for sub_dir in find_sub_dirs_all_sitepackages(sub_dirs): - attachments.append(f' listdir("{sub_dir}"):') - for node in sorted(os.listdir(sub_dir)): - attachments.append(f" {node}") - - -def _find_so_using_nvidia_lib_dirs( - libname: str, so_basename: str, error_messages: list[str], attachments: list[str] -) -> str | None: - rel_dirs = SITE_PACKAGES_LIBDIRS_LINUX.get(libname) - if rel_dirs is not None: - sub_dirs_searched = [] - file_wild = so_basename + "*" - for rel_dir in rel_dirs: - sub_dir = tuple(rel_dir.split(os.path.sep)) - for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): - # First look for an exact match - so_name = os.path.join(abs_dir, so_basename) - if os.path.isfile(so_name): - return so_name - # Look for a versioned library - # Using sort here mainly to make the result deterministic. - for so_name in sorted(glob.glob(os.path.join(abs_dir, file_wild))): - if os.path.isfile(so_name): - return so_name - sub_dirs_searched.append(sub_dir) - for sub_dir in sub_dirs_searched: - _no_such_file_in_sub_dirs(sub_dir, file_wild, error_messages, attachments) - return None - - -def _find_dll_under_dir(dirpath: str, file_wild: str) -> str | None: - for path in sorted(glob.glob(os.path.join(dirpath, file_wild))): - if not os.path.isfile(path): - continue - if not is_suppressed_dll_file(os.path.basename(path)): - return path - return None - - -def _find_dll_using_nvidia_bin_dirs( - libname: str, lib_searched_for: str, error_messages: list[str], attachments: list[str] -) -> str | None: - rel_dirs = SITE_PACKAGES_LIBDIRS_WINDOWS.get(libname) - if rel_dirs is not None: - sub_dirs_searched = [] - for rel_dir in rel_dirs: - sub_dir = tuple(rel_dir.split(os.path.sep)) - for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): - dll_name = _find_dll_under_dir(abs_dir, lib_searched_for) - if dll_name is not None: - return dll_name - sub_dirs_searched.append(sub_dir) - for sub_dir in sub_dirs_searched: - _no_such_file_in_sub_dirs(sub_dir, lib_searched_for, error_messages, attachments) - return None - - -def _find_lib_dir_using_anchor_point(libname: str, anchor_point: str, linux_lib_dir: str) -> str | None: - # Resolve paths for the four cases: - # Windows/Linux x nvvm yes/no - if IS_WINDOWS: - if libname == "nvvm": # noqa: SIM108 - rel_paths = [ - "nvvm/bin/*", # CTK 13 - "nvvm/bin", # CTK 12 - ] - else: - rel_paths = [ - "bin/x64", # CTK 13 - "bin", # CTK 12 - ] - else: - if libname == "nvvm": # noqa: SIM108 - rel_paths = ["nvvm/lib64"] - else: - rel_paths = [linux_lib_dir] - - for rel_path in rel_paths: - 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 - - -def _find_lib_dir_using_cuda_home(libname: str) -> str | None: - cuda_home = get_cuda_home_or_path() - if cuda_home is None: - return None - return _find_lib_dir_using_anchor_point(libname, anchor_point=cuda_home, linux_lib_dir="lib64") - - -def _find_lib_dir_using_conda_prefix(libname: str) -> str | None: - conda_prefix = os.environ.get("CONDA_PREFIX") - if not conda_prefix: - return None - return _find_lib_dir_using_anchor_point( - libname, anchor_point=os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix, linux_lib_dir="lib" - ) - - -def _find_so_using_lib_dir( - lib_dir: str, so_basename: str, error_messages: list[str], attachments: list[str] -) -> str | None: - so_name = os.path.join(lib_dir, so_basename) - if os.path.isfile(so_name): - return so_name - error_messages.append(f"No such file: {so_name}") - attachments.append(f' listdir("{lib_dir}"):') - if not os.path.isdir(lib_dir): - attachments.append(" DIRECTORY DOES NOT EXIST") - else: - for node in sorted(os.listdir(lib_dir)): - attachments.append(f" {node}") - return None - - -def _find_dll_using_lib_dir( - lib_dir: str, libname: str, error_messages: list[str], attachments: list[str] -) -> str | None: - file_wild = libname + "*.dll" - dll_name = _find_dll_under_dir(lib_dir, file_wild) - if dll_name is not None: - return dll_name - error_messages.append(f"No such file: {file_wild}") - attachments.append(f' listdir("{lib_dir}"):') - for node in sorted(os.listdir(lib_dir)): - attachments.append(f" {node}") - return None - - -def _derive_ctk_root_linux(resolved_lib_path: str) -> str | None: - """Derive the CTK installation root from a resolved library path on Linux. - - Standard system CTK layout: ``$CTK_ROOT/lib64/libfoo.so.XX`` - (some installs use ``lib`` instead of ``lib64``). - Also handles target-specific layouts: - ``$CTK_ROOT/targets//lib64/libfoo.so.XX`` (or ``lib``). - - Returns None if the path doesn't match a recognized layout. - """ - lib_dir = os.path.dirname(resolved_lib_path) - basename = os.path.basename(lib_dir) - if basename in ("lib64", "lib"): - parent = os.path.dirname(lib_dir) - grandparent = os.path.dirname(parent) - if os.path.basename(grandparent) == "targets": - # This corresponds to /.../targets//lib{,64} - return os.path.dirname(grandparent) - return parent - return None - - -def _derive_ctk_root_windows(resolved_lib_path: str) -> str | None: - """Derive the CTK installation root from a resolved library path on Windows. - - Handles two CTK layouts: - - CTK 13: ``$CTK_ROOT/bin/x64/foo.dll`` - - CTK 12: ``$CTK_ROOT/bin/foo.dll`` - - Returns None if the path doesn't match a recognized layout. - - Uses ``ntpath`` explicitly so the function is testable on any platform. - """ - import ntpath - - lib_dir = ntpath.dirname(resolved_lib_path) - basename = ntpath.basename(lib_dir).lower() - if basename == "x64": - parent = ntpath.dirname(lib_dir) - if ntpath.basename(parent).lower() == "bin": - return ntpath.dirname(parent) - elif basename == "bin": - return ntpath.dirname(lib_dir) - return None - - -def derive_ctk_root(resolved_lib_path: str) -> str | None: - """Derive the CTK installation root from a resolved library path. - - Given the absolute path of a loaded CTK shared library, walk up the - directory tree to find the CTK root. Returns None if the path doesn't - match any recognized CTK directory layout. - """ - if IS_WINDOWS: - return _derive_ctk_root_windows(resolved_lib_path) - return _derive_ctk_root_linux(resolved_lib_path) - - -class _FindNvidiaDynamicLib: - def __init__(self, libname: str): - self.libname = libname - if IS_WINDOWS: - self.lib_searched_for = f"{libname}*.dll" - else: - self.lib_searched_for = f"lib{libname}.so" - self.error_messages: list[str] = [] - self.attachments: list[str] = [] - self.abs_path: str | None = None - - def try_site_packages(self) -> str | None: - if IS_WINDOWS: - return _find_dll_using_nvidia_bin_dirs( - self.libname, - self.lib_searched_for, - self.error_messages, - self.attachments, - ) - else: - return _find_so_using_nvidia_lib_dirs( - self.libname, - self.lib_searched_for, - self.error_messages, - self.attachments, - ) - - def try_with_conda_prefix(self) -> str | None: - return self._find_using_lib_dir(_find_lib_dir_using_conda_prefix(self.libname)) - - def try_with_cuda_home(self) -> str | None: - return self._find_using_lib_dir(_find_lib_dir_using_cuda_home(self.libname)) - - def try_via_ctk_root(self, ctk_root: str) -> str | None: - """Find the library under a derived CTK root directory. - - Uses :func:`_find_lib_dir_using_anchor_point` which already knows - about non-standard sub-paths (e.g. ``nvvm/lib64`` for nvvm). - """ - return self._find_using_lib_dir( - _find_lib_dir_using_anchor_point(self.libname, anchor_point=ctk_root, linux_lib_dir="lib64") - ) - - def _find_using_lib_dir(self, lib_dir: str | None) -> str | None: - if lib_dir is None: - return None - if IS_WINDOWS: - return _find_dll_using_lib_dir( - lib_dir, - self.libname, - self.error_messages, - self.attachments, - ) - else: - return _find_so_using_lib_dir( - lib_dir, - self.lib_searched_for, - self.error_messages, - self.attachments, - ) - - def raise_not_found_error(self) -> None: - err = ", ".join(self.error_messages) - att = "\n".join(self.attachments) - raise DynamicLibNotFoundError(f'Failure finding "{self.lib_searched_for}": {err}\n{att}') diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/lib_descriptor.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/lib_descriptor.py new file mode 100644 index 0000000000..d3ab749bd5 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/lib_descriptor.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Per-library descriptor and registry. + +The canonical authored data lives in :mod:`descriptor_catalog`. This module +provides a name-keyed registry consumed by the runtime search/load path. +""" + +from __future__ import annotations + +from typing import TypeAlias + +from cuda.pathfinder._dynamic_libs.descriptor_catalog import ( + DESCRIPTOR_CATALOG, + DescriptorSpec, +) +from cuda.pathfinder._dynamic_libs.descriptor_catalog import ( + Strategy as Strategy, +) + +# Keep the historical type name for downstream imports. +LibDescriptor: TypeAlias = DescriptorSpec + + +#: Canonical registry of all known libraries. +LIB_DESCRIPTORS: dict[str, LibDescriptor] = {desc.name: desc for desc in DESCRIPTOR_CATALOG} 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 91e6284a00..68717a415a 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_common.py @@ -1,10 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING -from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import DIRECT_DEPENDENCIES +if TYPE_CHECKING: + from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor class DynamicLibNotFoundError(RuntimeError): @@ -19,6 +23,6 @@ class LoadedDL: found_via: str -def load_dependencies(libname: str, load_func: Callable[[str], LoadedDL]) -> None: - for dep in DIRECT_DEPENDENCIES.get(libname, ()): +def load_dependencies(desc: LibDescriptor, load_func: Callable[[str], LoadedDL]) -> None: + for dep in desc.dependencies: load_func(dep) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py index 4d2bae5b90..e4f2d0d817 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_linux.py @@ -1,17 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + import contextlib import ctypes import ctypes.util import os -from typing import cast +from typing import TYPE_CHECKING, cast from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL -from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( - LIBNAMES_REQUIRING_RTLD_DEEPBIND, - SUPPORTED_LINUX_SONAMES, -) + +if TYPE_CHECKING: + from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor CDLL_MODE = os.RTLD_NOW | os.RTLD_GLOBAL @@ -124,34 +125,37 @@ def abs_path_for_dynamic_library(libname: str, handle: ctypes.CDLL) -> str: return os.path.join(l_origin, os.path.basename(l_name)) -def get_candidate_sonames(libname: str) -> list[str]: - # Reverse tabulated names to achieve new → old search order. - candidate_sonames = list(reversed(SUPPORTED_LINUX_SONAMES.get(libname, ()))) - candidate_sonames.append(f"lib{libname}.so") - return candidate_sonames +def _candidate_sonames(desc: LibDescriptor) -> list[str]: + # Reverse tabulated names to achieve new -> old search order. + candidates = list(reversed(desc.linux_sonames)) + candidates.append(f"lib{desc.name}.so") + return candidates -def check_if_already_loaded_from_elsewhere(libname: str, _have_abs_path: bool) -> LoadedDL | None: - for soname in get_candidate_sonames(libname): +def check_if_already_loaded_from_elsewhere(desc: LibDescriptor, _have_abs_path: bool) -> LoadedDL | None: + for soname in _candidate_sonames(desc): try: handle = ctypes.CDLL(soname, mode=os.RTLD_NOLOAD) except OSError: continue else: return LoadedDL( - abs_path_for_dynamic_library(libname, handle), True, handle._handle, "was-already-loaded-from-elsewhere" + abs_path_for_dynamic_library(desc.name, handle), + True, + handle._handle, + "was-already-loaded-from-elsewhere", ) return None -def _load_lib(libname: str, filename: str) -> ctypes.CDLL: +def _load_lib(desc: LibDescriptor, filename: str) -> ctypes.CDLL: cdll_mode = CDLL_MODE - if libname in LIBNAMES_REQUIRING_RTLD_DEEPBIND: + if desc.requires_rtld_deepbind: cdll_mode |= os.RTLD_DEEPBIND return ctypes.CDLL(filename, cdll_mode) -def load_with_system_search(libname: str) -> LoadedDL | None: +def load_with_system_search(desc: LibDescriptor) -> LoadedDL | None: """Try to load a library using system search paths. Args: @@ -163,15 +167,15 @@ def load_with_system_search(libname: str) -> LoadedDL | None: Raises: RuntimeError: If the library is loaded but no expected symbol is found """ - for soname in get_candidate_sonames(libname): + for soname in _candidate_sonames(desc): try: - handle = _load_lib(libname, soname) + handle = _load_lib(desc, soname) except OSError: pass else: - abs_path = abs_path_for_dynamic_library(libname, handle) + abs_path = abs_path_for_dynamic_library(desc.name, handle) if abs_path is None: - raise RuntimeError(f"No expected symbol for {libname=!r}") + raise RuntimeError(f"No expected symbol for libname={desc.name!r}") return LoadedDL(abs_path, False, handle._handle, "system-search") return None @@ -195,22 +199,23 @@ def _work_around_known_bugs(libname: str, found_path: str) -> None: ctypes.CDLL(dep_path, CDLL_MODE) -def load_with_abs_path(libname: str, found_path: str, found_via: str | None = None) -> LoadedDL: +def load_with_abs_path(desc: LibDescriptor, found_path: str, found_via: str | None = None) -> LoadedDL: """Load a dynamic library from the given path. Args: - libname: The name of the library to load - found_path: The absolute path to the library file + desc: Descriptor for the library to load. + found_path: The absolute path to the library file. + found_via: Label indicating how the path was discovered. Returns: - A LoadedDL object representing the loaded library + A LoadedDL object representing the loaded library. Raises: - RuntimeError: If the library cannot be loaded + RuntimeError: If the library cannot be loaded. """ - _work_around_known_bugs(libname, found_path) + _work_around_known_bugs(desc.name, found_path) try: - handle = _load_lib(libname, found_path) + handle = _load_lib(desc, found_path) except OSError as e: raise RuntimeError(f"Failed to dlopen {found_path}: {e}") from e return LoadedDL(found_path, False, handle._handle, found_via) diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py index b9f15ea50b..cf4e32d0d8 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/load_dl_windows.py @@ -1,16 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + import ctypes import ctypes.wintypes import os import struct +from typing import TYPE_CHECKING from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL -from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( - LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY, - SUPPORTED_WINDOWS_DLLS, -) + +if TYPE_CHECKING: + from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor # Mirrors WinBase.h (unfortunately not defined already elsewhere) WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100 @@ -99,12 +101,12 @@ def abs_path_for_dynamic_library(libname: str, handle: ctypes.wintypes.HMODULE) return buffer.value -def check_if_already_loaded_from_elsewhere(libname: str, have_abs_path: bool) -> LoadedDL | None: - for dll_name in SUPPORTED_WINDOWS_DLLS.get(libname, ()): +def check_if_already_loaded_from_elsewhere(desc: LibDescriptor, have_abs_path: bool) -> LoadedDL | None: + for dll_name in desc.windows_dlls: handle = kernel32.GetModuleHandleW(dll_name) if handle: - abs_path = abs_path_for_dynamic_library(libname, handle) - if have_abs_path and libname in LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY: + abs_path = abs_path_for_dynamic_library(desc.name, handle) + if have_abs_path and desc.requires_add_dll_directory: # This is a side-effect if the pathfinder loads the library via # load_with_abs_path(). To make the side-effect more deterministic, # activate it even if the library was already loaded from elsewhere. @@ -113,7 +115,7 @@ def check_if_already_loaded_from_elsewhere(libname: str, have_abs_path: bool) -> return None -def load_with_system_search(libname: str) -> LoadedDL | None: +def load_with_system_search(desc: LibDescriptor) -> LoadedDL | None: """Try to load a DLL using system search paths. Args: @@ -122,30 +124,31 @@ def load_with_system_search(libname: str) -> LoadedDL | None: Returns: A LoadedDL object if successful, None if the library cannot be loaded """ - # Reverse tabulated names to achieve new → old search order. - for dll_name in reversed(SUPPORTED_WINDOWS_DLLS.get(libname, ())): + # Reverse tabulated names to achieve new -> old search order. + for dll_name in reversed(desc.windows_dlls): handle = kernel32.LoadLibraryExW(dll_name, None, 0) if handle: - abs_path = abs_path_for_dynamic_library(libname, handle) + abs_path = abs_path_for_dynamic_library(desc.name, handle) return LoadedDL(abs_path, False, ctypes_handle_to_unsigned_int(handle), "system-search") return None -def load_with_abs_path(libname: str, found_path: str, found_via: str | None = None) -> LoadedDL: +def load_with_abs_path(desc: LibDescriptor, found_path: str, found_via: str | None = None) -> LoadedDL: """Load a dynamic library from the given path. Args: - libname: The name of the library to load - found_path: The absolute path to the DLL file + desc: Descriptor for the library to load. + found_path: The absolute path to the DLL file. + found_via: Label indicating how the path was discovered. Returns: - A LoadedDL object representing the loaded library + A LoadedDL object representing the loaded library. Raises: - RuntimeError: If the DLL cannot be loaded + RuntimeError: If the DLL cannot be loaded. """ - if libname in LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY: + if desc.requires_add_dll_directory: add_dll_directory(found_path) flags = WINBASE_LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | WINBASE_LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR 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 3e22d62b30..87f09fd79d 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 @@ -1,54 +1,42 @@ # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + import functools import json import struct import sys +from typing import TYPE_CHECKING from cuda.pathfinder._dynamic_libs.canary_probe_subprocess import probe_canary_abs_path_and_print_json -from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import ( - _FindNvidiaDynamicLib, - derive_ctk_root, -) +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL, load_dependencies -from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( - _CTK_ROOT_CANARY_ANCHOR_LIBNAMES, - _CTK_ROOT_CANARY_DISCOVERABLE_LIBNAMES, - SUPPORTED_LINUX_SONAMES, - SUPPORTED_WINDOWS_DLLS, +from cuda.pathfinder._dynamic_libs.platform_loader import LOADER +from cuda.pathfinder._dynamic_libs.search_steps import ( + EARLY_FIND_STEPS, + LATE_FIND_STEPS, + SearchContext, + derive_ctk_root, + find_via_ctk_root, + run_find_steps, ) -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS from cuda.pathfinder._utils.spawned_process_runner import run_in_spawned_child_process -if IS_WINDOWS: - from cuda.pathfinder._dynamic_libs.load_dl_windows import ( - check_if_already_loaded_from_elsewhere, - load_with_abs_path, - load_with_system_search, - ) -else: - from cuda.pathfinder._dynamic_libs.load_dl_linux import ( - check_if_already_loaded_from_elsewhere, - load_with_abs_path, - load_with_system_search, - ) +if TYPE_CHECKING: + from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor # All libnames recognized by load_nvidia_dynamic_lib, across all categories -# (CTK, third-party, driver). Built from the platform-appropriate soname/DLL -# registry so that platform-specific libs (e.g. cufile on Linux) are included -# only where they apply. -_ALL_SUPPORTED_LIBNAMES: frozenset[str] = frozenset( - (SUPPORTED_WINDOWS_DLLS if IS_WINDOWS else SUPPORTED_LINUX_SONAMES).keys() -) +# (CTK, third-party, driver). +_ALL_SUPPORTED_LIBNAMES: frozenset[str] = frozenset(LIB_DESCRIPTORS) # Driver libraries: shipped with the NVIDIA display driver, always on the # system linker path. These skip all CTK search steps (site-packages, # conda, CUDA_HOME, canary) and go straight to system search. -_DRIVER_ONLY_LIBNAMES = frozenset(("cuda", "nvml")) +_DRIVER_ONLY_LIBNAMES = frozenset(name for name, desc in LIB_DESCRIPTORS.items() if desc.strategy == "driver") -def _load_driver_lib_no_cache(libname: str) -> LoadedDL: +def _load_driver_lib_no_cache(desc: LibDescriptor) -> LoadedDL: """Load an NVIDIA driver library (system-search only). Driver libs (libcuda, libnvidia-ml) are part of the display driver, not @@ -56,31 +44,21 @@ def _load_driver_lib_no_cache(libname: str) -> LoadedDL: full CTK search cascade (site-packages, conda, CUDA_HOME, canary) is unnecessary. """ - loaded = check_if_already_loaded_from_elsewhere(libname, False) + loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, False) if loaded is not None: return loaded - loaded = load_with_system_search(libname) + loaded = LOADER.load_with_system_search(desc) if loaded is not None: return loaded raise DynamicLibNotFoundError( - f'"{libname}" is an NVIDIA driver library and can only be found via' + f'"{desc.name}" is an NVIDIA driver library and can only be found via' f" system search. Ensure the NVIDIA display driver is installed." ) @functools.cache def _resolve_system_loaded_abs_path_in_subprocess(libname: str) -> str | None: - """Resolve a library's system-search absolute path in a child process. - - This runs in a spawned (not forked) child process. Spawning is important - because it starts from a fresh interpreter state, so the child does not - inherit already-loaded CUDA dynamic libraries from the parent process - (especially the well-known canary probe library). - - That keeps any side-effects of loading the canary library scoped to the - child process instead of polluting the current process, and ensures the - canary probe is an independent system-search attempt. - """ + """Resolve a canary library's absolute path in a spawned child process.""" result = run_in_spawned_child_process( probe_canary_abs_path_and_print_json, args=(libname,), @@ -88,7 +66,7 @@ def _resolve_system_loaded_abs_path_in_subprocess(libname: str) -> str | None: rethrow=True, ) - # Read the final non-empty stdout line in case earlier lines are emitted. + # Use the final non-empty line in case earlier output lines are emitted. lines = [line for line in result.stdout.splitlines() if line.strip()] if not lines: raise RuntimeError(f"Canary probe child process produced no stdout payload for {libname!r}") @@ -105,80 +83,58 @@ def _resolve_system_loaded_abs_path_in_subprocess(libname: str) -> str | None: raise RuntimeError(f"Canary probe child process emitted unexpected payload for {libname!r}: {payload!r}") -def _try_ctk_root_canary(finder: _FindNvidiaDynamicLib) -> str | None: - """Derive the CTK root from a system-installed canary lib. - - For discoverable libs (currently nvvm) whose shared object doesn't reside - on the standard linker path, we locate a well-known CTK lib that IS on - the linker path via system search, derive the CTK installation root from - its resolved path, and then look for the target lib relative to that root. - - The canary load is performed in a subprocess to avoid introducing loader - state into the current process. - """ - for canary_libname in _CTK_ROOT_CANARY_ANCHOR_LIBNAMES: +def _try_ctk_root_canary(ctx: SearchContext) -> str | None: + """Try CTK-root canary fallback for descriptor-configured libraries.""" + for canary_libname in ctx.desc.ctk_root_canary_anchor_libnames: canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname) if canary_abs_path is None: continue ctk_root = derive_ctk_root(canary_abs_path) if ctk_root is None: continue - abs_path: str | None = finder.try_via_ctk_root(ctk_root) - if abs_path is not None: - return abs_path + find = find_via_ctk_root(ctx, ctk_root) + if find is not None: + return str(find.abs_path) return None def _load_lib_no_cache(libname: str) -> LoadedDL: + desc = LIB_DESCRIPTORS[libname] + if libname in _DRIVER_ONLY_LIBNAMES: - return _load_driver_lib_no_cache(libname) - - finder = _FindNvidiaDynamicLib(libname) - abs_path = finder.try_site_packages() - if abs_path is not None: - found_via = "site-packages" - else: - abs_path = finder.try_with_conda_prefix() - if abs_path is not None: - found_via = "conda" - - # If the library was already loaded by someone else, reproduce any OS-specific - # side-effects we would have applied on a direct absolute-path load (e.g., - # AddDllDirectory on Windows for libs that require it). - loaded = check_if_already_loaded_from_elsewhere(libname, abs_path is not None) - - # Load dependencies regardless of who loaded the primary lib first. - # Doing this *after* the side-effect ensures dependencies resolve consistently - # relative to the actually loaded location. - load_dependencies(libname, load_nvidia_dynamic_lib) + return _load_driver_lib_no_cache(desc) + + ctx = SearchContext(desc) + # Phase 1: Try to find the library file on disk (pip wheels, conda). + find = run_find_steps(ctx, EARLY_FIND_STEPS) + + # Phase 2: Cross-cutting — already-loaded check and dependency loading. + # The already-loaded check on Windows uses the "have we found a path?" + # flag to decide whether to apply AddDllDirectory side-effects. + loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, find is not None) + load_dependencies(desc, load_nvidia_dynamic_lib) if loaded is not None: return loaded - if abs_path is None: - loaded = load_with_system_search(libname) - if loaded is not None: - return loaded + # Phase 3: Load from found path, or fall back to system search + late find. + if find is not None: + return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) + + loaded = LOADER.load_with_system_search(desc) + if loaded is not None: + return loaded - abs_path = finder.try_with_cuda_home() - if abs_path is not None: - found_via = "CUDA_HOME" - else: - if libname not in _CTK_ROOT_CANARY_DISCOVERABLE_LIBNAMES: - finder.raise_not_found_error() + find = run_find_steps(ctx, LATE_FIND_STEPS) + if find is not None: + return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) - # Canary probe (discoverable libs only): if the direct system - # search and CUDA_HOME both failed (e.g. nvvm isn't on the linker - # path and CUDA_HOME is unset), try to discover the CTK root by - # loading a well-known CTK lib in a subprocess, then look for the - # target lib relative to that root. - abs_path = _try_ctk_root_canary(finder) - if abs_path is not None: - found_via = "system-ctk-root" - else: - finder.raise_not_found_error() + if desc.ctk_root_canary_anchor_libnames: + canary_abs_path = _try_ctk_root_canary(ctx) + if canary_abs_path is not None: + return LOADER.load_with_abs_path(desc, canary_abs_path, "system-ctk-root") - return load_with_abs_path(libname, abs_path, found_via) + ctx.raise_not_found() @functools.cache @@ -249,10 +205,9 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: 5. **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 discover the CTK installation root by system-loading a - well-known CTK library (``cudart``) in a subprocess, then derive - the root from its resolved absolute path. + standard linker path (currently ``nvvm``), attempt to derive CTK + root by system-loading a well-known CTK canary library in a + subprocess and then searching relative to that root. **Driver libraries** (``"cuda"``, ``"nvml"``): diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/platform_loader.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/platform_loader.py new file mode 100644 index 0000000000..9b108a57ac --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/platform_loader.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Platform loader seam for OS-specific dynamic linking. + +This module provides a small interface that hides the Linux vs Windows +implementation details of: + +- already-loaded checks +- system-search loading +- absolute-path loading + +The orchestration logic in :mod:`load_nvidia_dynamic_lib` should not need to +branch on platform; it calls through the loader instance exported here. +""" + +from __future__ import annotations + +from typing import Protocol + +from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor +from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +class PlatformLoader(Protocol): + def check_if_already_loaded_from_elsewhere(self, desc: LibDescriptor, have_abs_path: bool) -> LoadedDL | None: ... + + def load_with_system_search(self, desc: LibDescriptor) -> LoadedDL | None: ... + + def load_with_abs_path(self, desc: LibDescriptor, found_path: str, found_via: str | None = None) -> LoadedDL: ... + + +if IS_WINDOWS: + from cuda.pathfinder._dynamic_libs import load_dl_windows as _impl +else: + from cuda.pathfinder._dynamic_libs import load_dl_linux as _impl + +# The platform modules already expose functions matching the PlatformLoader +# protocol. Wrap in a simple namespace so callers use LOADER.method() syntax. +LOADER: PlatformLoader = _impl diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py new file mode 100644 index 0000000000..817ac0b65f --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_platform.py @@ -0,0 +1,202 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Platform abstraction for filesystem search steps. + +The goal is to keep :mod:`search_steps` platform-agnostic: it should not branch +on OS flags like ``IS_WINDOWS``. Instead, it calls through the single +``PLATFORM`` instance exported here. +""" + +from __future__ import annotations + +import glob +import os +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Protocol, cast + +from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor +from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import is_suppressed_dll_file +from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS + + +def _no_such_file_in_sub_dirs( + sub_dirs: Sequence[str], file_wild: str, error_messages: list[str], attachments: list[str] +) -> None: + error_messages.append(f"No such file: {file_wild}") + for sub_dir in find_sub_dirs_all_sitepackages(sub_dirs): + attachments.append(f' listdir("{sub_dir}"):') + for node in sorted(os.listdir(sub_dir)): + attachments.append(f" {node}") + + +def _find_so_in_rel_dirs( + rel_dirs: tuple[str, ...], + so_basename: str, + error_messages: list[str], + attachments: list[str], +) -> str | None: + sub_dirs_searched: list[tuple[str, ...]] = [] + file_wild = so_basename + "*" + for rel_dir in rel_dirs: + sub_dir = tuple(rel_dir.split(os.path.sep)) + for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): + so_name = os.path.join(abs_dir, so_basename) + if os.path.isfile(so_name): + return so_name + for so_name in sorted(glob.glob(os.path.join(abs_dir, file_wild))): + if os.path.isfile(so_name): + return so_name + sub_dirs_searched.append(sub_dir) + for sub_dir in sub_dirs_searched: + _no_such_file_in_sub_dirs(sub_dir, file_wild, error_messages, attachments) + return None + + +def _find_dll_under_dir(dirpath: str, file_wild: str) -> str | None: + for path in sorted(glob.glob(os.path.join(dirpath, file_wild))): + if not os.path.isfile(path): + continue + if not is_suppressed_dll_file(os.path.basename(path)): + return path + return None + + +def _find_dll_in_rel_dirs( + rel_dirs: tuple[str, ...], + lib_searched_for: str, + error_messages: list[str], + attachments: list[str], +) -> str | None: + sub_dirs_searched: list[tuple[str, ...]] = [] + for rel_dir in rel_dirs: + sub_dir = tuple(rel_dir.split(os.path.sep)) + for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): + dll_name = _find_dll_under_dir(abs_dir, lib_searched_for) + if dll_name is not None: + return dll_name + sub_dirs_searched.append(sub_dir) + for sub_dir in sub_dirs_searched: + _no_such_file_in_sub_dirs(sub_dir, lib_searched_for, error_messages, attachments) + return None + + +class SearchPlatform(Protocol): + def lib_searched_for(self, libname: str) -> str: ... + + def site_packages_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ... + + def conda_anchor_point(self, conda_prefix: str) -> str: ... + + def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: ... + + def find_in_site_packages( + self, + rel_dirs: tuple[str, ...], + lib_searched_for: str, + error_messages: list[str], + attachments: list[str], + ) -> str | None: ... + + def find_in_lib_dir( + self, + lib_dir: str, + libname: str, + lib_searched_for: str, + error_messages: list[str], + attachments: list[str], + ) -> str | None: ... + + +@dataclass(frozen=True, slots=True) +class LinuxSearchPlatform: + def lib_searched_for(self, libname: str) -> str: + return f"lib{libname}.so" + + def site_packages_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: + return cast(tuple[str, ...], desc.site_packages_linux) + + def conda_anchor_point(self, conda_prefix: str) -> str: + return conda_prefix + + def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: + return cast(tuple[str, ...], desc.anchor_rel_dirs_linux) + + def find_in_site_packages( + self, + rel_dirs: tuple[str, ...], + lib_searched_for: str, + error_messages: list[str], + attachments: list[str], + ) -> str | None: + return _find_so_in_rel_dirs(rel_dirs, lib_searched_for, error_messages, attachments) + + def find_in_lib_dir( + self, + lib_dir: str, + _libname: str, + lib_searched_for: str, + error_messages: list[str], + attachments: list[str], + ) -> str | None: + so_name = os.path.join(lib_dir, lib_searched_for) + if os.path.isfile(so_name): + return so_name + error_messages.append(f"No such file: {so_name}") + attachments.append(f' listdir("{lib_dir}"):') + if not os.path.isdir(lib_dir): + attachments.append(" DIRECTORY DOES NOT EXIST") + else: + for node in sorted(os.listdir(lib_dir)): + attachments.append(f" {node}") + return None + + +@dataclass(frozen=True, slots=True) +class WindowsSearchPlatform: + def lib_searched_for(self, libname: str) -> str: + return f"{libname}*.dll" + + def site_packages_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: + return cast(tuple[str, ...], desc.site_packages_windows) + + def conda_anchor_point(self, conda_prefix: str) -> str: + return os.path.join(conda_prefix, "Library") + + def anchor_rel_dirs(self, desc: LibDescriptor) -> tuple[str, ...]: + return cast(tuple[str, ...], desc.anchor_rel_dirs_windows) + + def find_in_site_packages( + self, + rel_dirs: tuple[str, ...], + 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) + + def find_in_lib_dir( + self, + lib_dir: str, + libname: str, + _lib_searched_for: str, + error_messages: list[str], + attachments: list[str], + ) -> str | None: + file_wild = libname + "*.dll" + dll_name = _find_dll_under_dir(lib_dir, file_wild) + if dll_name is not None: + return dll_name + error_messages.append(f"No such file: {file_wild}") + attachments.append(f' listdir("{lib_dir}"):') + if not os.path.isdir(lib_dir): + attachments.append(" DIRECTORY DOES NOT EXIST") + else: + for node in sorted(os.listdir(lib_dir)): + attachments.append(f" {node}") + return None + + +PLATFORM: SearchPlatform = WindowsSearchPlatform() if IS_WINDOWS else LinuxSearchPlatform() diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py new file mode 100644 index 0000000000..7c7e36f70d --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/search_steps.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Composable search steps for locating NVIDIA libraries. + +Each find step is a callable with signature:: + + (SearchContext) -> FindResult | None + +Find steps locate a library file on disk without loading it. The +orchestrator in :mod:`load_nvidia_dynamic_lib` handles loading, the +already-loaded check, and dependency resolution. + +Step sequences are defined per search strategy so that adding a new +step or strategy only requires adding a function and a tuple entry. + +This module is intentionally platform-agnostic: it does not branch on the +current operating system. Platform differences are routed through the +:data:`~cuda.pathfinder._dynamic_libs.search_platform.PLATFORM` instance. +""" + +import glob +import os +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import NoReturn, cast + +from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor +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_home_or_path + +# --------------------------------------------------------------------------- +# Data types +# --------------------------------------------------------------------------- + + +@dataclass +class FindResult: + """A library file located on disk (not yet loaded).""" + + abs_path: str + found_via: str + + +@dataclass +class SearchContext: + """Mutable state accumulated during the search cascade.""" + + desc: LibDescriptor + platform: SearchPlatform = PLATFORM + error_messages: list[str] = field(default_factory=list) + attachments: list[str] = field(default_factory=list) + + @property + def libname(self) -> str: + return self.desc.name # type: ignore[no-any-return] # mypy can't resolve new sibling module + + @property + def lib_searched_for(self) -> str: + return cast(str, self.platform.lib_searched_for(self.libname)) + + def raise_not_found(self) -> NoReturn: + err = ", ".join(self.error_messages) + att = "\n".join(self.attachments) + raise DynamicLibNotFoundError(f'Failure finding "{self.lib_searched_for}": {err}\n{att}') + + +#: Type alias for a find step callable. +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) + 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 + + +def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: + """Find a library file in a resolved lib directory.""" + if lib_dir is None: + return None + return cast( + str | None, + ctx.platform.find_in_lib_dir( + lib_dir, + ctx.libname, + ctx.lib_searched_for, + ctx.error_messages, + ctx.attachments, + ), + ) + + +def _derive_ctk_root_linux(resolved_lib_path: str) -> str | None: + """Derive CTK root from Linux canary path. + + Supports: + - ``$CTK_ROOT/lib64/libfoo.so.*`` + - ``$CTK_ROOT/lib/libfoo.so.*`` + - ``$CTK_ROOT/targets//lib64/libfoo.so.*`` + - ``$CTK_ROOT/targets//lib/libfoo.so.*`` + """ + lib_dir = os.path.dirname(resolved_lib_path) + basename = os.path.basename(lib_dir) + if basename in ("lib64", "lib"): + parent = os.path.dirname(lib_dir) + grandparent = os.path.dirname(parent) + if os.path.basename(grandparent) == "targets": + return os.path.dirname(grandparent) + return parent + return None + + +def _derive_ctk_root_windows(resolved_lib_path: str) -> str | None: + """Derive CTK root from Windows canary path. + + Supports: + - ``$CTK_ROOT/bin/x64/foo.dll`` (CTK 13 style) + - ``$CTK_ROOT/bin/foo.dll`` (CTK 12 style) + """ + import ntpath + + lib_dir = ntpath.dirname(resolved_lib_path) + basename = ntpath.basename(lib_dir).lower() + if basename == "x64": + parent = ntpath.dirname(lib_dir) + if ntpath.basename(parent).lower() == "bin": + return ntpath.dirname(parent) + elif basename == "bin": + return ntpath.dirname(lib_dir) + return None + + +def derive_ctk_root(resolved_lib_path: str) -> str | None: + """Derive CTK root from a resolved canary library path.""" + ctk_root = _derive_ctk_root_linux(resolved_lib_path) + if ctk_root is not None: + return ctk_root + return _derive_ctk_root_windows(resolved_lib_path) + + +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") + + +# --------------------------------------------------------------------------- +# Find steps +# --------------------------------------------------------------------------- + + +def find_in_site_packages(ctx: SearchContext) -> FindResult | None: + """Search pip wheel install locations.""" + 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) + if abs_path is not None: + return FindResult(abs_path, "site-packages") + return None + + +def find_in_conda(ctx: SearchContext) -> FindResult | None: + """Search ``$CONDA_PREFIX``.""" + conda_prefix = os.environ.get("CONDA_PREFIX") + 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 + + +def find_in_cuda_home(ctx: SearchContext) -> FindResult | None: + """Search ``$CUDA_HOME`` / ``$CUDA_PATH``.""" + cuda_home = get_cuda_home_or_path() + 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_HOME") + return None + + +# --------------------------------------------------------------------------- +# Step sequences per strategy +# --------------------------------------------------------------------------- + +#: Find steps that run before the already-loaded check and system search. +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_home,) + + +# --------------------------------------------------------------------------- +# Cascade runner +# --------------------------------------------------------------------------- + + +def run_find_steps(ctx: SearchContext, steps: tuple[FindStep, ...]) -> FindResult | None: + """Run find steps in order, returning the first hit.""" + for step in steps: + result = step(ctx) + if result is not None: + return result + return None diff --git a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py index d4225233c2..d3f1a9a848 100644 --- a/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/supported_nvidia_libs.py @@ -1,450 +1,70 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -# THIS FILE NEEDS TO BE REVIEWED/UPDATED FOR EACH CTK RELEASE -# Likely candidates for updates are: -# SUPPORTED_LIBNAMES -# SUPPORTED_WINDOWS_DLLS -# SUPPORTED_LINUX_SONAMES +"""Legacy table exports derived from the authored descriptor catalog. -from cuda.pathfinder._utils.platform_aware import IS_WINDOWS +The canonical data entry point is :mod:`descriptor_catalog`. This module keeps +historical constant names for backward compatibility by deriving them from the +catalog. +""" -SUPPORTED_LIBNAMES_COMMON = ( - # Core CUDA Runtime and Compiler - "cudart", - "nvfatbin", - "nvJitLink", - "nvrtc", - "nvvm", - # Math Libraries - "cublas", - "cublasLt", - "cufft", - "cufftw", - "curand", - "cusolver", - "cusolverMg", - "cusparse", - "nppc", - "nppial", - "nppicc", - "nppidei", - "nppif", - "nppig", - "nppim", - "nppist", - "nppisu", - "nppitc", - "npps", - "nvblas", - # Other - "nvjpeg", -) +from __future__ import annotations -# Note: The `cufile_rdma` information is intentionally retained (commented out) -# despite not being actively used in the current build. It took a nontrivial -# amount of effort to determine the SONAME, dependencies, and expected symbols -# for this special-case library, especially given its RDMA/MLX5 dependencies -# and limited availability. Keeping this as a reference avoids having to -# reconstruct the information from scratch in the future. +from cuda.pathfinder._dynamic_libs.descriptor_catalog import DESCRIPTOR_CATALOG +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS -SUPPORTED_LIBNAMES_LINUX_ONLY = ( - "cufile", - # "cufile_rdma", # Requires libmlx5.so +_CTK_DESCRIPTORS = tuple(desc for desc in DESCRIPTOR_CATALOG if desc.strategy == "ctk") +_OTHER_DESCRIPTORS = tuple(desc for desc in DESCRIPTOR_CATALOG if desc.strategy == "other") +_DRIVER_DESCRIPTORS = tuple(desc for desc in DESCRIPTOR_CATALOG if desc.strategy == "driver") +_NON_CTK_DESCRIPTORS = _OTHER_DESCRIPTORS + _DRIVER_DESCRIPTORS + +SUPPORTED_LIBNAMES_COMMON = tuple(desc.name for desc in _CTK_DESCRIPTORS if desc.linux_sonames and desc.windows_dlls) +SUPPORTED_LIBNAMES_LINUX_ONLY = tuple( + desc.name for desc in _CTK_DESCRIPTORS if desc.linux_sonames and not desc.windows_dlls +) +SUPPORTED_LIBNAMES_WINDOWS_ONLY = tuple( + desc.name for desc in _CTK_DESCRIPTORS if desc.windows_dlls and not desc.linux_sonames ) -SUPPORTED_LIBNAMES_LINUX = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY -SUPPORTED_LIBNAMES_WINDOWS_ONLY = () +SUPPORTED_LIBNAMES_LINUX = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY SUPPORTED_LIBNAMES_WINDOWS = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_WINDOWS_ONLY - SUPPORTED_LIBNAMES_ALL = SUPPORTED_LIBNAMES_COMMON + SUPPORTED_LIBNAMES_LINUX_ONLY + SUPPORTED_LIBNAMES_WINDOWS_ONLY SUPPORTED_LIBNAMES = SUPPORTED_LIBNAMES_WINDOWS if IS_WINDOWS else SUPPORTED_LIBNAMES_LINUX -# Based on ldd output for Linux x86_64 nvidia-*-cu12 wheels (12.8.1) -DIRECT_DEPENDENCIES_CTK = { - "cublas": ("cublasLt",), - "cufftw": ("cufft",), - # "cufile_rdma": ("cufile",), - "cusolver": ("nvJitLink", "cusparse", "cublasLt", "cublas"), - "cusolverMg": ("nvJitLink", "cublasLt", "cublas"), - "cusparse": ("nvJitLink",), - "nppial": ("nppc",), - "nppicc": ("nppc",), - "nppidei": ("nppc",), - "nppif": ("nppc",), - "nppig": ("nppc",), - "nppim": ("nppc",), - "nppist": ("nppc",), - "nppisu": ("nppc",), - "nppitc": ("nppc",), - "npps": ("nppc",), - "nvblas": ("cublas", "cublasLt"), -} -DIRECT_DEPENDENCIES = DIRECT_DEPENDENCIES_CTK | { - "mathdx": ("nvrtc",), - "cublasmp": ("cublas", "cublasLt", "nvshmem_host"), - "cufftMp": ("nvshmem_host",), - "cudss": ("cublas", "cublasLt"), - "cutensor": ("cublasLt",), - "cutensorMg": ("cutensor", "cublasLt"), -} +DIRECT_DEPENDENCIES_CTK = {desc.name: desc.dependencies for desc in _CTK_DESCRIPTORS if desc.dependencies} +DIRECT_DEPENDENCIES = {desc.name: desc.dependencies for desc in DESCRIPTOR_CATALOG if desc.dependencies} -# Based on these files: -# cuda_12.0.1_525.85.12_linux.run -# cuda_12.1.1_530.30.02_linux.run -# cuda_12.2.2_535.104.05_linux.run -# cuda_12.3.2_545.23.08_linux.run -# cuda_12.4.1_550.54.15_linux.run -# cuda_12.5.1_555.42.06_linux.run -# cuda_12.6.3_560.35.05_linux.run -# cuda_12.8.1_570.124.06_linux.run -# cuda_12.9.1_575.57.08_linux.run -# cuda_13.0.2_580.95.05_linux.run -# cuda_13.1.0_590.44.01_linux.run -# Generated with toolshed/build_pathfinder_sonames.py -# Please keep in old → new sort order. -SUPPORTED_LINUX_SONAMES_CTK = { - "cublas": ( - "libcublas.so.12", - "libcublas.so.13", - ), - "cublasLt": ( - "libcublasLt.so.12", - "libcublasLt.so.13", - ), - "cudart": ( - "libcudart.so.12", - "libcudart.so.13", - ), - "cufft": ( - "libcufft.so.11", - "libcufft.so.12", - ), - "cufftw": ( - "libcufftw.so.11", - "libcufftw.so.12", - ), - "cufile": ("libcufile.so.0",), - # "cufile_rdma": ("libcufile_rdma.so.1",), - "curand": ("libcurand.so.10",), - "cusolver": ( - "libcusolver.so.11", - "libcusolver.so.12", - ), - "cusolverMg": ( - "libcusolverMg.so.11", - "libcusolverMg.so.12", - ), - "cusparse": ("libcusparse.so.12",), - "nppc": ( - "libnppc.so.12", - "libnppc.so.13", - ), - "nppial": ( - "libnppial.so.12", - "libnppial.so.13", - ), - "nppicc": ( - "libnppicc.so.12", - "libnppicc.so.13", - ), - "nppidei": ( - "libnppidei.so.12", - "libnppidei.so.13", - ), - "nppif": ( - "libnppif.so.12", - "libnppif.so.13", - ), - "nppig": ( - "libnppig.so.12", - "libnppig.so.13", - ), - "nppim": ( - "libnppim.so.12", - "libnppim.so.13", - ), - "nppist": ( - "libnppist.so.12", - "libnppist.so.13", - ), - "nppisu": ( - "libnppisu.so.12", - "libnppisu.so.13", - ), - "nppitc": ( - "libnppitc.so.12", - "libnppitc.so.13", - ), - "npps": ( - "libnpps.so.12", - "libnpps.so.13", - ), - "nvJitLink": ( - "libnvJitLink.so.12", - "libnvJitLink.so.13", - ), - "nvblas": ( - "libnvblas.so.12", - "libnvblas.so.13", - ), - "nvfatbin": ( - "libnvfatbin.so.12", - "libnvfatbin.so.13", - ), - "nvjpeg": ( - "libnvjpeg.so.12", - "libnvjpeg.so.13", - ), - "nvrtc": ( - "libnvrtc.so.12", - "libnvrtc.so.13", - ), - "nvvm": ("libnvvm.so.4",), -} -SUPPORTED_LINUX_SONAMES_OTHER = { - "cublasmp": ("libcublasmp.so.0",), - "cufftMp": ("libcufftMp.so.12", "libcufftMp.so.11"), - "mathdx": ("libmathdx.so.0",), - "cudss": ("libcudss.so.0",), - "cusparseLt": ("libcusparseLt.so.0",), - "cutensor": ("libcutensor.so.2",), - "cutensorMg": ("libcutensorMg.so.2",), - "nccl": ("libnccl.so.2",), - "nvpl_fftw": ("libnvpl_fftw.so.0",), - "nvshmem_host": ("libnvshmem_host.so.3",), -} -# Driver libraries: shipped with the NVIDIA driver, always on the system -# linker path. Only system search is needed (no site-packages / conda / -# CUDA_HOME). -SUPPORTED_LINUX_SONAMES_DRIVER = { - "cuda": ("libcuda.so.1",), - "nvml": ("libnvidia-ml.so.1",), -} +SUPPORTED_LINUX_SONAMES_CTK = {desc.name: desc.linux_sonames for desc in _CTK_DESCRIPTORS if desc.linux_sonames} +SUPPORTED_LINUX_SONAMES_OTHER = {desc.name: desc.linux_sonames for desc in _OTHER_DESCRIPTORS if desc.linux_sonames} +SUPPORTED_LINUX_SONAMES_DRIVER = {desc.name: desc.linux_sonames for desc in _DRIVER_DESCRIPTORS if desc.linux_sonames} SUPPORTED_LINUX_SONAMES = SUPPORTED_LINUX_SONAMES_CTK | SUPPORTED_LINUX_SONAMES_OTHER | SUPPORTED_LINUX_SONAMES_DRIVER -# Based on these files: -# cuda_12.0.1_528.33_windows.exe -# cuda_12.1.1_531.14_windows.exe -# cuda_12.2.2_537.13_windows.exe -# cuda_12.3.2_546.12_windows.exe -# cuda_12.4.1_551.78_windows.exe -# cuda_12.5.1_555.85_windows.exe -# cuda_12.6.3_561.17_windows.exe -# cuda_12.8.1_572.61_windows.exe -# cuda_12.9.1_576.57_windows.exe -# cuda_13.0.2_windows.exe -# cuda_13.1.0_windows.exe -# Generated with toolshed/build_pathfinder_dlls.py -# Please keep in old → new sort order. -SUPPORTED_WINDOWS_DLLS_CTK = { - "cublas": ( - "cublas64_12.dll", - "cublas64_13.dll", - ), - "cublasLt": ( - "cublasLt64_12.dll", - "cublasLt64_13.dll", - ), - "cudart": ( - "cudart64_12.dll", - "cudart64_13.dll", - ), - "cufft": ( - "cufft64_11.dll", - "cufft64_12.dll", - ), - "cufftw": ( - "cufftw64_11.dll", - "cufftw64_12.dll", - ), - "curand": ("curand64_10.dll",), - "cusolver": ( - "cusolver64_11.dll", - "cusolver64_12.dll", - ), - "cusolverMg": ( - "cusolverMg64_11.dll", - "cusolverMg64_12.dll", - ), - "cusparse": ("cusparse64_12.dll",), - "nppc": ( - "nppc64_12.dll", - "nppc64_13.dll", - ), - "nppial": ( - "nppial64_12.dll", - "nppial64_13.dll", - ), - "nppicc": ( - "nppicc64_12.dll", - "nppicc64_13.dll", - ), - "nppidei": ( - "nppidei64_12.dll", - "nppidei64_13.dll", - ), - "nppif": ( - "nppif64_12.dll", - "nppif64_13.dll", - ), - "nppig": ( - "nppig64_12.dll", - "nppig64_13.dll", - ), - "nppim": ( - "nppim64_12.dll", - "nppim64_13.dll", - ), - "nppist": ( - "nppist64_12.dll", - "nppist64_13.dll", - ), - "nppisu": ( - "nppisu64_12.dll", - "nppisu64_13.dll", - ), - "nppitc": ( - "nppitc64_12.dll", - "nppitc64_13.dll", - ), - "npps": ( - "npps64_12.dll", - "npps64_13.dll", - ), - "nvJitLink": ( - "nvJitLink_120_0.dll", - "nvJitLink_130_0.dll", - ), - "nvblas": ( - "nvblas64_12.dll", - "nvblas64_13.dll", - ), - "nvfatbin": ( - "nvfatbin_120_0.dll", - "nvfatbin_130_0.dll", - ), - "nvjpeg": ( - "nvjpeg64_12.dll", - "nvjpeg64_13.dll", - ), - "nvrtc": ( - "nvrtc64_120_0.dll", - "nvrtc64_130_0.dll", - ), - "nvvm": ( - "nvvm64.dll", - "nvvm64_40_0.dll", - "nvvm70.dll", - ), -} -SUPPORTED_WINDOWS_DLLS_OTHER = { - "mathdx": ("mathdx64_0.dll",), - "cudss": ("cudss64_0.dll",), - "cusparseLt": ("cusparseLt.dll",), - "cutensor": ("cutensor.dll",), - "cutensorMg": ("cutensorMg.dll",), -} -SUPPORTED_WINDOWS_DLLS_DRIVER = { - "cuda": ("nvcuda.dll",), - "nvml": ("nvml.dll",), -} +SUPPORTED_WINDOWS_DLLS_CTK = {desc.name: desc.windows_dlls for desc in _CTK_DESCRIPTORS if desc.windows_dlls} +SUPPORTED_WINDOWS_DLLS_OTHER = {desc.name: desc.windows_dlls for desc in _OTHER_DESCRIPTORS if desc.windows_dlls} +SUPPORTED_WINDOWS_DLLS_DRIVER = {desc.name: desc.windows_dlls for desc in _DRIVER_DESCRIPTORS if desc.windows_dlls} SUPPORTED_WINDOWS_DLLS = SUPPORTED_WINDOWS_DLLS_CTK | SUPPORTED_WINDOWS_DLLS_OTHER | SUPPORTED_WINDOWS_DLLS_DRIVER -LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY = ( - "cufft", - "nvrtc", +LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY = tuple( + desc.name for desc in DESCRIPTOR_CATALOG if desc.requires_add_dll_directory and desc.windows_dlls +) +LIBNAMES_REQUIRING_RTLD_DEEPBIND = tuple( + desc.name for desc in DESCRIPTOR_CATALOG if desc.requires_rtld_deepbind and desc.linux_sonames ) - -LIBNAMES_REQUIRING_RTLD_DEEPBIND = ("cufftMp",) - -# CTK root canary probe config: -# - anchor libs: expected on the standard system loader path and used to derive -# CTK root in an isolated child process. -# - discoverable libs: libs that are allowed to use the CTK-root canary fallback. -_CTK_ROOT_CANARY_ANCHOR_LIBNAMES = ("cudart",) -_CTK_ROOT_CANARY_DISCOVERABLE_LIBNAMES = ("nvvm",) # Based on output of toolshed/make_site_packages_libdirs_linux.py SITE_PACKAGES_LIBDIRS_LINUX_CTK = { - "cublas": ("nvidia/cu13/lib", "nvidia/cublas/lib"), - "cublasLt": ("nvidia/cu13/lib", "nvidia/cublas/lib"), - "cudart": ("nvidia/cu13/lib", "nvidia/cuda_runtime/lib"), - "cufft": ("nvidia/cu13/lib", "nvidia/cufft/lib"), - "cufftw": ("nvidia/cu13/lib", "nvidia/cufft/lib"), - "cufile": ("nvidia/cu13/lib", "nvidia/cufile/lib"), - # "cufile_rdma": ("nvidia/cu13/lib", "nvidia/cufile/lib"), - "curand": ("nvidia/cu13/lib", "nvidia/curand/lib"), - "cusolver": ("nvidia/cu13/lib", "nvidia/cusolver/lib"), - "cusolverMg": ("nvidia/cu13/lib", "nvidia/cusolver/lib"), - "cusparse": ("nvidia/cu13/lib", "nvidia/cusparse/lib"), - "nppc": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppial": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppicc": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppidei": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppif": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppig": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppim": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppist": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppisu": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nppitc": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "npps": ("nvidia/cu13/lib", "nvidia/npp/lib"), - "nvJitLink": ("nvidia/cu13/lib", "nvidia/nvjitlink/lib"), - "nvblas": ("nvidia/cu13/lib", "nvidia/cublas/lib"), - "nvfatbin": ("nvidia/cu13/lib", "nvidia/nvfatbin/lib"), - "nvjpeg": ("nvidia/cu13/lib", "nvidia/nvjpeg/lib"), - "nvrtc": ("nvidia/cu13/lib", "nvidia/cuda_nvrtc/lib"), - "nvvm": ("nvidia/cu13/lib", "nvidia/cuda_nvcc/nvvm/lib64"), + desc.name: desc.site_packages_linux for desc in _CTK_DESCRIPTORS if desc.site_packages_linux } SITE_PACKAGES_LIBDIRS_LINUX_OTHER = { - "cublasmp": ("nvidia/cublasmp/cu13/lib", "nvidia/cublasmp/cu12/lib"), - "cudss": ("nvidia/cu13/lib", "nvidia/cu12/lib"), - "cufftMp": ("nvidia/cufftmp/cu13/lib", "nvidia/cufftmp/cu12/lib"), - "cusparseLt": ("nvidia/cusparselt/lib",), - "cutensor": ("cutensor/lib",), - "cutensorMg": ("cutensor/lib",), - "mathdx": ("nvidia/cu13/lib", "nvidia/cu12/lib"), - "nccl": ("nvidia/nccl/lib",), - "nvpl_fftw": ("nvpl/lib",), - "nvshmem_host": ("nvidia/nvshmem/lib",), + desc.name: desc.site_packages_linux for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_linux } SITE_PACKAGES_LIBDIRS_LINUX = SITE_PACKAGES_LIBDIRS_LINUX_CTK | SITE_PACKAGES_LIBDIRS_LINUX_OTHER -# Based on output of toolshed/make_site_packages_libdirs_windows.py SITE_PACKAGES_LIBDIRS_WINDOWS_CTK = { - "cublas": ("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), - "cublasLt": ("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), - "cudart": ("nvidia/cu13/bin/x86_64", "nvidia/cuda_runtime/bin"), - "cufft": ("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"), - "cufftw": ("nvidia/cu13/bin/x86_64", "nvidia/cufft/bin"), - "curand": ("nvidia/cu13/bin/x86_64", "nvidia/curand/bin"), - "cusolver": ("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"), - "cusolverMg": ("nvidia/cu13/bin/x86_64", "nvidia/cusolver/bin"), - "cusparse": ("nvidia/cu13/bin/x86_64", "nvidia/cusparse/bin"), - "nppc": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppial": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppicc": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppidei": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppif": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppig": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppim": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppist": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppisu": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nppitc": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "npps": ("nvidia/cu13/bin/x86_64", "nvidia/npp/bin"), - "nvJitLink": ("nvidia/cu13/bin/x86_64", "nvidia/nvjitlink/bin"), - "nvblas": ("nvidia/cu13/bin/x86_64", "nvidia/cublas/bin"), - "nvfatbin": ("nvidia/cu13/bin/x86_64", "nvidia/nvfatbin/bin"), - "nvjpeg": ("nvidia/cu13/bin/x86_64", "nvidia/nvjpeg/bin"), - "nvrtc": ("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvrtc/bin"), - "nvvm": ("nvidia/cu13/bin/x86_64", "nvidia/cuda_nvcc/nvvm/bin"), + desc.name: desc.site_packages_windows for desc in _CTK_DESCRIPTORS if desc.site_packages_windows } SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER = { - "cudss": ("nvidia/cu13/bin", "nvidia/cu12/bin"), - "mathdx": ("nvidia/cu13/bin/x86_64", "nvidia/cu12/bin"), - "cusparseLt": ("nvidia/cusparselt/bin",), - "cutensor": ("cutensor/bin",), - "cutensorMg": ("cutensor/bin",), + desc.name: desc.site_packages_windows for desc in _NON_CTK_DESCRIPTORS if desc.site_packages_windows } SITE_PACKAGES_LIBDIRS_WINDOWS = SITE_PACKAGES_LIBDIRS_WINDOWS_CTK | SITE_PACKAGES_LIBDIRS_WINDOWS_OTHER diff --git a/cuda_pathfinder/tests/test_ctk_root_discovery.py b/cuda_pathfinder/tests/test_ctk_root_discovery.py index be20da57af..bcfce30c25 100644 --- a/cuda_pathfinder/tests/test_ctk_root_discovery.py +++ b/cuda_pathfinder/tests/test_ctk_root_discovery.py @@ -4,22 +4,30 @@ import pytest -from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import ( - _derive_ctk_root_linux, - _derive_ctk_root_windows, - _FindNvidiaDynamicLib, - derive_ctk_root, -) +from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as load_mod +from cuda.pathfinder._dynamic_libs import search_steps as steps_mod +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( _load_lib_no_cache, _resolve_system_loaded_abs_path_in_subprocess, _try_ctk_root_canary, ) +from cuda.pathfinder._dynamic_libs.search_steps import ( + SearchContext, + _derive_ctk_root_linux, + _derive_ctk_root_windows, + derive_ctk_root, + find_via_ctk_root, +) from cuda.pathfinder._utils.platform_aware import IS_WINDOWS _MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib" -_FIND_MODULE = "cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib" +_STEPS_MODULE = "cuda.pathfinder._dynamic_libs.search_steps" + + +def _ctx(libname: str = "nvvm") -> SearchContext: + return SearchContext(LIB_DESCRIPTORS[libname]) @pytest.fixture(autouse=True) @@ -124,17 +132,22 @@ def test_derive_ctk_root_windows_case_insensitive_x64(): def test_derive_ctk_root_dispatches_to_linux(mocker): - mocker.patch(f"{_FIND_MODULE}.IS_WINDOWS", False) + linux_derive = mocker.patch(f"{_STEPS_MODULE}._derive_ctk_root_linux", return_value="/usr/local/cuda") + windows_derive = mocker.patch(f"{_STEPS_MODULE}._derive_ctk_root_windows") assert derive_ctk_root("/usr/local/cuda/lib64/libcudart.so.13") == "/usr/local/cuda" + linux_derive.assert_called_once_with("/usr/local/cuda/lib64/libcudart.so.13") + windows_derive.assert_not_called() def test_derive_ctk_root_dispatches_to_windows(mocker): - mocker.patch(f"{_FIND_MODULE}.IS_WINDOWS", True) + mocker.patch(f"{_STEPS_MODULE}._derive_ctk_root_linux", return_value=None) + windows_derive = mocker.patch(f"{_STEPS_MODULE}._derive_ctk_root_windows", return_value=r"C:\CUDA\v13") assert derive_ctk_root(r"C:\CUDA\v13\bin\cudart64_13.dll") == r"C:\CUDA\v13" + windows_derive.assert_called_once_with(r"C:\CUDA\v13\bin\cudart64_13.dll") # --------------------------------------------------------------------------- -# _FindNvidiaDynamicLib.try_via_ctk_root +# find_via_ctk_root # --------------------------------------------------------------------------- @@ -142,21 +155,27 @@ def test_try_via_ctk_root_finds_nvvm(tmp_path): ctk_root = tmp_path / "cuda-13" nvvm_lib = _create_nvvm_in_ctk(ctk_root) - assert _FindNvidiaDynamicLib("nvvm").try_via_ctk_root(str(ctk_root)) == str(nvvm_lib) + result = find_via_ctk_root(_ctx("nvvm"), str(ctk_root)) + assert result is not None + assert result.abs_path == str(nvvm_lib) + assert result.found_via == "system-ctk-root" def test_try_via_ctk_root_returns_none_when_dir_missing(tmp_path): ctk_root = tmp_path / "cuda-13" ctk_root.mkdir() - assert _FindNvidiaDynamicLib("nvvm").try_via_ctk_root(str(ctk_root)) is None + assert find_via_ctk_root(_ctx("nvvm"), str(ctk_root)) is None def test_try_via_ctk_root_regular_lib(tmp_path): ctk_root = tmp_path / "cuda-13" cudart_lib = _create_cudart_in_ctk(ctk_root) - assert _FindNvidiaDynamicLib("cudart").try_via_ctk_root(str(ctk_root)) == str(cudart_lib) + result = find_via_ctk_root(_ctx("cudart"), str(ctk_root)) + assert result is not None + assert result.abs_path == str(cudart_lib) + assert result.found_via == "system-ctk-root" # --------------------------------------------------------------------------- @@ -230,16 +249,16 @@ def test_canary_finds_nvvm(tmp_path, mocker): f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=_fake_canary_path(ctk_root), ) - parent_system_loader = mocker.patch(f"{_MODULE}.load_with_system_search") + parent_system_loader = mocker.patch.object(load_mod.LOADER, "load_with_system_search") - assert _try_ctk_root_canary(_FindNvidiaDynamicLib("nvvm")) == str(nvvm_lib) + assert _try_ctk_root_canary(_ctx("nvvm")) == str(nvvm_lib) probe.assert_called_once_with("cudart") parent_system_loader.assert_not_called() def test_canary_returns_none_when_subprocess_probe_fails(mocker): mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) - assert _try_ctk_root_canary(_FindNvidiaDynamicLib("nvvm")) is None + assert _try_ctk_root_canary(_ctx("nvvm")) is None def test_canary_returns_none_when_ctk_root_unrecognized(mocker): @@ -247,7 +266,7 @@ def test_canary_returns_none_when_ctk_root_unrecognized(mocker): f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value="/weird/path/libcudart.so.13", ) - assert _try_ctk_root_canary(_FindNvidiaDynamicLib("nvvm")) is None + assert _try_ctk_root_canary(_ctx("nvvm")) is None def test_canary_returns_none_when_nvvm_not_in_ctk_root(tmp_path, mocker): @@ -259,12 +278,12 @@ def test_canary_returns_none_when_nvvm_not_in_ctk_root(tmp_path, mocker): f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=_fake_canary_path(ctk_root), ) - assert _try_ctk_root_canary(_FindNvidiaDynamicLib("nvvm")) is None + assert _try_ctk_root_canary(_ctx("nvvm")) is None def test_canary_skips_when_abs_path_none(mocker): mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) - assert _try_ctk_root_canary(_FindNvidiaDynamicLib("nvvm")) is None + assert _try_ctk_root_canary(_ctx("nvvm")) is None # --------------------------------------------------------------------------- @@ -279,12 +298,17 @@ def _isolate_load_cascade(mocker): This lets the ordering tests focus on system-search, CUDA_HOME, and the canary probe without needing a real site-packages or conda environment. """ - # No wheels installed - mocker.patch.object(_FindNvidiaDynamicLib, "try_site_packages", return_value=None) - # No conda env - mocker.patch.object(_FindNvidiaDynamicLib, "try_with_conda_prefix", return_value=None) + + # Skip EARLY_FIND_STEPS (site-packages + conda) so tests can focus on + # system-search, CUDA_HOME and canary behavior. + def _run_find_steps_with_early_disabled(ctx, steps): + if steps is load_mod.EARLY_FIND_STEPS: + return None + return steps_mod.run_find_steps(ctx, steps) + + mocker.patch(f"{_MODULE}.run_find_steps", side_effect=_run_find_steps_with_early_disabled) # Lib not already loaded by another component - mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) + mocker.patch.object(load_mod.LOADER, "check_if_already_loaded_from_elsewhere", return_value=None) # Skip transitive dependency loading mocker.patch(f"{_MODULE}.load_dependencies") @@ -302,15 +326,16 @@ def test_cuda_home_takes_priority_over_canary(tmp_path, mocker): canary_mock = mocker.MagicMock(return_value=_fake_canary_path(canary_root)) # System search finds nothing for nvvm. - mocker.patch(f"{_MODULE}.load_with_system_search", return_value=None) + mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None) # Canary subprocess probe would find cudart if consulted. mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", side_effect=canary_mock) # CUDA_HOME points to a separate root that also has nvvm - mocker.patch(f"{_FIND_MODULE}.get_cuda_home_or_path", return_value=str(cuda_home_root)) + mocker.patch(f"{_STEPS_MODULE}.get_cuda_home_or_path", return_value=str(cuda_home_root)) # Capture the final load call - mocker.patch( - f"{_MODULE}.load_with_abs_path", - side_effect=lambda _libname, path, via: _make_loaded_dl(path, via), + 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("nvvm") @@ -328,18 +353,19 @@ def test_canary_fires_only_after_all_earlier_steps_fail(tmp_path, mocker): nvvm_lib = _create_nvvm_in_ctk(canary_root) # System search: nvvm not on linker path. - mocker.patch(f"{_MODULE}.load_with_system_search", return_value=None) + mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None) # Canary subprocess probe finds cudart under a system CTK root. mocker.patch( f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=_fake_canary_path(canary_root), ) # No CUDA_HOME set - mocker.patch(f"{_FIND_MODULE}.get_cuda_home_or_path", return_value=None) + mocker.patch(f"{_STEPS_MODULE}.get_cuda_home_or_path", return_value=None) # Capture the final load call - mocker.patch( - f"{_MODULE}.load_with_abs_path", - side_effect=lambda _libname, path, via: _make_loaded_dl(path, via), + 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("nvvm") @@ -351,8 +377,8 @@ def test_canary_fires_only_after_all_earlier_steps_fail(tmp_path, mocker): @pytest.mark.usefixtures("_isolate_load_cascade") def test_non_discoverable_lib_skips_canary_probe(mocker): # Force fallback path for a lib that is not canary-discoverable. - mocker.patch(f"{_MODULE}.load_with_system_search", return_value=None) - mocker.patch(f"{_FIND_MODULE}.get_cuda_home_or_path", return_value=None) + mocker.patch.object(load_mod.LOADER, "load_with_system_search", return_value=None) + mocker.patch(f"{_STEPS_MODULE}.get_cuda_home_or_path", return_value=None) canary_probe = mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess") with pytest.raises(DynamicLibNotFoundError): diff --git a/cuda_pathfinder/tests/test_descriptor_catalog.py b/cuda_pathfinder/tests/test_descriptor_catalog.py new file mode 100644 index 0000000000..a299f80c42 --- /dev/null +++ b/cuda_pathfinder/tests/test_descriptor_catalog.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Structural invariant tests for the authored descriptor catalog. + +These verify properties that should always hold for any valid catalog +entry, rather than comparing the catalog against itself. +""" + +from __future__ import annotations + +import re + +import pytest + +from cuda.pathfinder._dynamic_libs.descriptor_catalog import DESCRIPTOR_CATALOG, DescriptorSpec + +_VALID_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_VALID_STRATEGIES = {"ctk", "other", "driver"} +_CATALOG_BY_NAME = {spec.name: spec for spec in DESCRIPTOR_CATALOG} + + +def test_catalog_names_are_unique(): + names = [spec.name for spec in DESCRIPTOR_CATALOG] + assert len(names) == len(set(names)) + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_name_is_valid_identifier(spec: DescriptorSpec): + assert _VALID_NAME_RE.match(spec.name), f"{spec.name!r} is not a valid Python identifier" + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_strategy_is_valid(spec: DescriptorSpec): + assert spec.strategy in _VALID_STRATEGIES + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_has_at_least_one_soname_or_dll(spec: DescriptorSpec): + assert spec.linux_sonames or spec.windows_dlls, f"{spec.name} has no sonames or DLLs" + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_dependencies_reference_existing_entries(spec: DescriptorSpec): + for dep in spec.dependencies: + assert dep in _CATALOG_BY_NAME, f"{spec.name} depends on unknown library {dep!r}" + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +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", + [s for s in DESCRIPTOR_CATALOG if s.strategy == "driver"], + ids=lambda s: s.name, +) +def test_driver_libs_have_no_site_packages(spec: DescriptorSpec): + """Driver libs are system-search-only; site-packages paths would be unused.""" + assert not spec.site_packages_linux, f"driver lib {spec.name} has site_packages_linux" + assert not spec.site_packages_windows, f"driver lib {spec.name} has site_packages_windows" + + +@pytest.mark.parametrize( + "spec", + [s for s in DESCRIPTOR_CATALOG if s.strategy == "driver"], + ids=lambda s: s.name, +) +def test_driver_libs_have_no_dependencies(spec: DescriptorSpec): + """Driver libs skip the full cascade and shouldn't declare deps.""" + assert not spec.dependencies, f"driver lib {spec.name} has dependencies" + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_linux_sonames_look_like_sonames(spec: DescriptorSpec): + for soname in spec.linux_sonames: + assert soname.startswith("lib"), f"Unexpected Linux soname format: {soname}" + assert ".so" in soname, f"Unexpected Linux soname format: {soname}" + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_windows_dlls_look_like_dlls(spec: DescriptorSpec): + for dll in spec.windows_dlls: + assert dll.endswith(".dll"), f"Unexpected Windows DLL format: {dll}" + + +@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: + assert anchor in _CATALOG_BY_NAME, f"{spec.name} has unknown canary anchor {anchor!r}" + assert _CATALOG_BY_NAME[anchor].strategy == "ctk", f"{spec.name} has non-CTK canary anchor {anchor!r}" + + +@pytest.mark.parametrize("spec", DESCRIPTOR_CATALOG, ids=lambda s: s.name) +def test_only_ctk_libs_define_ctk_root_canary_anchors(spec: DescriptorSpec): + if spec.ctk_root_canary_anchor_libnames: + assert spec.strategy == "ctk", f"{spec.name} defines canary anchors but is not a CTK lib" diff --git a/cuda_pathfinder/tests/test_driver_lib_loading.py b/cuda_pathfinder/tests/test_driver_lib_loading.py index ed36833c62..d8d463599a 100644 --- a/cuda_pathfinder/tests/test_driver_lib_loading.py +++ b/cuda_pathfinder/tests/test_driver_lib_loading.py @@ -14,6 +14,7 @@ import pytest from child_load_nvidia_dynamic_lib_helper import build_child_process_failed_for_libname_message, child_process_func +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( _DRIVER_ONLY_LIBNAMES, @@ -27,6 +28,10 @@ assert STRICTNESS in ("see_what_works", "all_must_work") _MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib" +_LOADER_MODULE = "cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib.LOADER" + +_CUDA_DESC = LIB_DESCRIPTORS["cuda"] +_NVML_DESC = LIB_DESCRIPTORS["nvml"] def _make_loaded_dl(path, found_via): @@ -40,47 +45,44 @@ def _make_loaded_dl(path, found_via): def test_driver_lib_returns_already_loaded(mocker): already = LoadedDL("/usr/lib/libcuda.so.1", True, 0xBEEF, "was-already-loaded-from-elsewhere") - mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=already) - mocker.patch(f"{_MODULE}.load_with_system_search") + mocker.patch(f"{_LOADER_MODULE}.check_if_already_loaded_from_elsewhere", return_value=already) + mocker.patch(f"{_LOADER_MODULE}.load_with_system_search") - result = _load_driver_lib_no_cache("cuda") + result = _load_driver_lib_no_cache(_CUDA_DESC) assert result is already - # system search should not have been called - from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as mod + from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import LOADER - mod.load_with_system_search.assert_not_called() + LOADER.load_with_system_search.assert_not_called() def test_driver_lib_falls_through_to_system_search(mocker): loaded = _make_loaded_dl("/usr/lib/libcuda.so.1", "system-search") - mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) - mocker.patch(f"{_MODULE}.load_with_system_search", return_value=loaded) + mocker.patch(f"{_LOADER_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) + mocker.patch(f"{_LOADER_MODULE}.load_with_system_search", return_value=loaded) - result = _load_driver_lib_no_cache("cuda") + result = _load_driver_lib_no_cache(_CUDA_DESC) assert result is loaded assert result.found_via == "system-search" def test_driver_lib_raises_when_not_found(mocker): - mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) - mocker.patch(f"{_MODULE}.load_with_system_search", return_value=None) + mocker.patch(f"{_LOADER_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) + mocker.patch(f"{_LOADER_MODULE}.load_with_system_search", return_value=None) with pytest.raises(DynamicLibNotFoundError, match="NVIDIA driver library"): - _load_driver_lib_no_cache("nvml") + _load_driver_lib_no_cache(_NVML_DESC) def test_driver_lib_does_not_search_site_packages(mocker): """Driver libs must not go through the CTK search cascade.""" loaded = _make_loaded_dl("/usr/lib/libcuda.so.1", "system-search") - mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) - mocker.patch(f"{_MODULE}.load_with_system_search", return_value=loaded) - - from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib + mocker.patch(f"{_LOADER_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) + mocker.patch(f"{_LOADER_MODULE}.load_with_system_search", return_value=loaded) - spy = mocker.spy(_FindNvidiaDynamicLib, "try_site_packages") - _load_driver_lib_no_cache("cuda") + spy = mocker.patch(f"{_MODULE}.run_find_steps") + _load_driver_lib_no_cache(_CUDA_DESC) spy.assert_not_called() @@ -97,22 +99,17 @@ def test_load_lib_no_cache_dispatches_to_driver_path(libname, mocker): result = _load_lib_no_cache(libname) assert result is loaded - mock_driver.assert_called_once_with(libname) + mock_driver.assert_called_once_with(LIB_DESCRIPTORS[libname]) def test_load_lib_no_cache_does_not_dispatch_ctk_lib_to_driver_path(mocker): """Ensure regular CTK libs don't take the driver shortcut.""" mock_driver = mocker.patch(f"{_MODULE}._load_driver_lib_no_cache") - # Let the normal path run far enough to prove the driver path wasn't used. - # We'll make it fail quickly at check_if_already_loaded_from_elsewhere. - from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib - - mocker.patch.object(_FindNvidiaDynamicLib, "try_site_packages", return_value=None) - mocker.patch.object(_FindNvidiaDynamicLib, "try_with_conda_prefix", return_value=None) - mocker.patch(f"{_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) + mocker.patch(f"{_MODULE}.run_find_steps", return_value=None) + mocker.patch(f"{_LOADER_MODULE}.check_if_already_loaded_from_elsewhere", return_value=None) mocker.patch(f"{_MODULE}.load_dependencies") mocker.patch( - f"{_MODULE}.load_with_system_search", + f"{_LOADER_MODULE}.load_with_system_search", return_value=_make_loaded_dl("/usr/lib/libcudart.so.13", "system-search"), ) diff --git a/cuda_pathfinder/tests/test_lib_descriptor.py b/cuda_pathfinder/tests/test_lib_descriptor.py new file mode 100644 index 0000000000..22f07ea159 --- /dev/null +++ b/cuda_pathfinder/tests/test_lib_descriptor.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests verifying that the LibDescriptor registry faithfully represents +the existing data tables in supported_nvidia_libs.py.""" + +import pytest + +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS +from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( + DIRECT_DEPENDENCIES, + LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY, + LIBNAMES_REQUIRING_RTLD_DEEPBIND, + SITE_PACKAGES_LIBDIRS_LINUX, + SITE_PACKAGES_LIBDIRS_WINDOWS, + SUPPORTED_LIBNAMES, + SUPPORTED_LINUX_SONAMES, + SUPPORTED_WINDOWS_DLLS, +) + +# --------------------------------------------------------------------------- +# Registry completeness +# --------------------------------------------------------------------------- + + +def test_registry_covers_all_linux_sonames(): + assert set(SUPPORTED_LINUX_SONAMES) <= set(LIB_DESCRIPTORS) + + +def test_registry_covers_all_windows_dlls(): + assert set(SUPPORTED_WINDOWS_DLLS) <= set(LIB_DESCRIPTORS) + + +def test_registry_has_no_extra_entries(): + expected = set(SUPPORTED_LINUX_SONAMES) | set(SUPPORTED_WINDOWS_DLLS) + assert set(LIB_DESCRIPTORS) == expected + + +# --------------------------------------------------------------------------- +# Per-field consistency with source dicts +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_linux_sonames_match(name): + assert LIB_DESCRIPTORS[name].linux_sonames == SUPPORTED_LINUX_SONAMES.get(name, ()) + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_windows_dlls_match(name): + assert LIB_DESCRIPTORS[name].windows_dlls == SUPPORTED_WINDOWS_DLLS.get(name, ()) + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_site_packages_linux_match(name): + assert LIB_DESCRIPTORS[name].site_packages_linux == SITE_PACKAGES_LIBDIRS_LINUX.get(name, ()) + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_site_packages_windows_match(name): + assert LIB_DESCRIPTORS[name].site_packages_windows == SITE_PACKAGES_LIBDIRS_WINDOWS.get(name, ()) + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_dependencies_match(name): + assert LIB_DESCRIPTORS[name].dependencies == DIRECT_DEPENDENCIES.get(name, ()) + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_requires_add_dll_directory_match(name): + assert LIB_DESCRIPTORS[name].requires_add_dll_directory == (name in LIBNAMES_REQUIRING_OS_ADD_DLL_DIRECTORY) + + +@pytest.mark.parametrize("name", sorted(LIB_DESCRIPTORS)) +def test_requires_rtld_deepbind_match(name): + assert LIB_DESCRIPTORS[name].requires_rtld_deepbind == (name in LIBNAMES_REQUIRING_RTLD_DEEPBIND) + + +# --------------------------------------------------------------------------- +# Strategy classification +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("name", sorted(SUPPORTED_LIBNAMES)) +def test_ctk_libs_have_ctk_strategy(name): + assert LIB_DESCRIPTORS[name].strategy == "ctk" + + +def test_other_libs_have_other_strategy(): + # Spot-check a few known "other" libs + for name in ("nccl", "cutensor", "cusparseLt"): + if name in LIB_DESCRIPTORS: + assert LIB_DESCRIPTORS[name].strategy == "other", name + + +# --------------------------------------------------------------------------- +# Descriptor properties +# --------------------------------------------------------------------------- + + +def test_descriptor_is_frozen(): + desc = LIB_DESCRIPTORS["cudart"] + with pytest.raises(AttributeError): + desc.name = "bogus" # type: ignore[misc] diff --git a/cuda_pathfinder/tests/test_search_steps.py b/cuda_pathfinder/tests/test_search_steps.py new file mode 100644 index 0000000000..ef14e29abe --- /dev/null +++ b/cuda_pathfinder/tests/test_search_steps.py @@ -0,0 +1,337 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for the composable search steps and cascade runner.""" + +from __future__ import annotations + +import os + +import pytest + +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS, LibDescriptor +from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError +from cuda.pathfinder._dynamic_libs.search_platform import LinuxSearchPlatform, WindowsSearchPlatform +from cuda.pathfinder._dynamic_libs.search_steps import ( + EARLY_FIND_STEPS, + LATE_FIND_STEPS, + FindResult, + SearchContext, + _find_lib_dir_using_anchor, + find_in_conda, + find_in_cuda_home, + find_in_site_packages, + run_find_steps, +) + +_STEPS_MOD = "cuda.pathfinder._dynamic_libs.search_steps" +_PLAT_MOD = "cuda.pathfinder._dynamic_libs.search_platform" + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_desc(name: str = "cudart", **overrides) -> LibDescriptor: + defaults = { + "name": name, + "strategy": "ctk", + "linux_sonames": ("libcudart.so",), + "windows_dlls": ("cudart64_12.dll",), + "site_packages_linux": (os.path.join("nvidia", "cuda_runtime", "lib"),), + "site_packages_windows": (os.path.join("nvidia", "cuda_runtime", "bin"),), + } + defaults.update(overrides) + return LibDescriptor(**defaults) + + +def _ctx(desc: LibDescriptor | None = None, *, platform=None) -> SearchContext: + if platform is None: + platform = LinuxSearchPlatform() + return SearchContext(desc or _make_desc(), platform=platform) + + +# --------------------------------------------------------------------------- +# SearchContext +# --------------------------------------------------------------------------- + + +class TestSearchContext: + def test_libname_delegates_to_descriptor(self): + ctx = _ctx(_make_desc(name="nvrtc")) + assert ctx.libname == "nvrtc" + + def test_lib_searched_for_linux(self): + ctx = SearchContext(_make_desc(name="cublas"), platform=LinuxSearchPlatform()) + assert ctx.lib_searched_for == "libcublas.so" + + def test_lib_searched_for_windows(self): + ctx = SearchContext(_make_desc(name="cublas"), platform=WindowsSearchPlatform()) + assert ctx.lib_searched_for == "cublas*.dll" + + def test_raise_not_found_includes_messages(self): + ctx = _ctx() + ctx.error_messages.append("No such file: libcudart.so*") + ctx.attachments.append(' listdir("/some/dir"):') + with pytest.raises(DynamicLibNotFoundError, match="No such file"): + ctx.raise_not_found() + + def test_raise_not_found_empty_messages(self): + ctx = _ctx() + with pytest.raises(DynamicLibNotFoundError): + ctx.raise_not_found() + + +# --------------------------------------------------------------------------- +# find_in_site_packages +# --------------------------------------------------------------------------- + + +class TestFindInSitePackages: + def test_returns_none_when_no_rel_dirs(self): + desc = _make_desc(site_packages_linux=(), site_packages_windows=()) + result = find_in_site_packages(_ctx(desc)) + assert result is None + + def test_found_linux(self, mocker, tmp_path): + lib_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib" + lib_dir.mkdir(parents=True) + so_file = lib_dir / "libcudart.so" + so_file.touch() + + mocker.patch( + f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages", + return_value=[str(lib_dir)], + ) + + desc = _make_desc( + site_packages_linux=(os.path.join("nvidia", "cuda_runtime", "lib"),), + ) + result = find_in_site_packages(_ctx(desc, platform=LinuxSearchPlatform())) + assert result is not None + assert result.abs_path == str(so_file) + assert result.found_via == "site-packages" + + def test_found_windows(self, mocker, tmp_path): + bin_dir = tmp_path / "nvidia" / "cuda_runtime" / "bin" + bin_dir.mkdir(parents=True) + dll = bin_dir / "cudart64_12.dll" + dll.touch() + + mocker.patch( + f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages", + return_value=[str(bin_dir)], + ) + mocker.patch(f"{_PLAT_MOD}.is_suppressed_dll_file", return_value=False) + + desc = _make_desc( + name="cudart", + site_packages_windows=(os.path.join("nvidia", "cuda_runtime", "bin"),), + ) + result = find_in_site_packages(_ctx(desc, platform=WindowsSearchPlatform())) + assert result is not None + assert result.abs_path == str(dll) + assert result.found_via == "site-packages" + + def test_not_found_appends_error(self, mocker, tmp_path): + empty_dir = tmp_path / "nvidia" / "cuda_runtime" / "lib" + empty_dir.mkdir(parents=True) + + mocker.patch( + f"{_PLAT_MOD}.find_sub_dirs_all_sitepackages", + return_value=[str(empty_dir)], + ) + + ctx = _ctx(platform=LinuxSearchPlatform()) + result = find_in_site_packages(ctx) + assert result is None + assert any("No such file" in m for m in ctx.error_messages) + + +# --------------------------------------------------------------------------- +# find_in_conda +# --------------------------------------------------------------------------- + + +class TestFindInConda: + def test_returns_none_without_conda_prefix(self, mocker): + mocker.patch.dict(os.environ, {}, clear=True) + assert find_in_conda(_ctx()) is None + + def test_returns_none_with_empty_conda_prefix(self, mocker): + mocker.patch.dict(os.environ, {"CONDA_PREFIX": ""}) + assert find_in_conda(_ctx()) is None + + def test_found_linux(self, mocker, tmp_path): + lib_dir = tmp_path / "lib" + lib_dir.mkdir() + so_file = lib_dir / "libcudart.so" + so_file.touch() + + mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)}) + + result = find_in_conda(_ctx(platform=LinuxSearchPlatform())) + assert result is not None + assert result.abs_path == str(so_file) + assert result.found_via == "conda" + + def test_found_windows(self, mocker, tmp_path): + bin_dir = tmp_path / "Library" / "bin" + bin_dir.mkdir(parents=True) + dll = bin_dir / "cudart64_12.dll" + dll.touch() + + mocker.patch.dict(os.environ, {"CONDA_PREFIX": str(tmp_path)}) + + result = find_in_conda(_ctx(platform=WindowsSearchPlatform())) + assert result is not None + assert result.abs_path == str(dll) + assert result.found_via == "conda" + + +# --------------------------------------------------------------------------- +# find_in_cuda_home +# --------------------------------------------------------------------------- + + +class TestFindInCudaHome: + def test_returns_none_without_env_var(self, mocker): + mocker.patch(f"{_STEPS_MOD}.get_cuda_home_or_path", return_value=None) + assert find_in_cuda_home(_ctx(platform=LinuxSearchPlatform())) is None + + def test_found_linux(self, mocker, tmp_path): + lib_dir = tmp_path / "lib64" + lib_dir.mkdir() + so_file = lib_dir / "libcudart.so" + so_file.touch() + + mocker.patch(f"{_STEPS_MOD}.get_cuda_home_or_path", return_value=str(tmp_path)) + + result = find_in_cuda_home(_ctx(platform=LinuxSearchPlatform())) + assert result is not None + assert result.abs_path == str(so_file) + assert result.found_via == "CUDA_HOME" + + def test_found_windows(self, mocker, tmp_path): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + dll = bin_dir / "cudart64_12.dll" + dll.touch() + + mocker.patch(f"{_STEPS_MOD}.get_cuda_home_or_path", return_value=str(tmp_path)) + + result = find_in_cuda_home(_ctx(platform=WindowsSearchPlatform())) + assert result is not None + assert result.abs_path == str(dll) + assert result.found_via == "CUDA_HOME" + + +# --------------------------------------------------------------------------- +# run_find_steps +# --------------------------------------------------------------------------- + + +class TestRunFindSteps: + def test_returns_first_hit(self): + hit = FindResult("/path/to/lib.so", "step-a") + + def step_a(_ctx): + return hit + + def step_b(_ctx): + raise AssertionError("step_b should not be called") + + result = run_find_steps(_ctx(), (step_a, step_b)) + assert result is hit + + def test_returns_none_when_all_miss(self): + result = run_find_steps(_ctx(), (lambda _: None, lambda _: None)) + assert result is None + + def test_empty_steps(self): + assert run_find_steps(_ctx(), ()) is None + + def test_skips_nones_returns_later_hit(self): + hit = FindResult("/later/lib.so", "step-c") + result = run_find_steps(_ctx(), (lambda _: None, lambda _: hit)) + assert result is hit + + +# --------------------------------------------------------------------------- +# Step tuple sanity checks +# --------------------------------------------------------------------------- + + +class TestStepTuples: + def test_early_find_steps_contains_expected(self): + assert find_in_site_packages in EARLY_FIND_STEPS + assert find_in_conda in EARLY_FIND_STEPS + + def test_late_find_steps_contains_expected(self): + assert find_in_cuda_home in LATE_FIND_STEPS + + def test_early_and_late_are_disjoint(self): + assert not set(EARLY_FIND_STEPS) & set(LATE_FIND_STEPS) + + +# --------------------------------------------------------------------------- +# Data-driven anchor paths +# --------------------------------------------------------------------------- + + +class TestAnchorRelDirs: + """Verify that descriptor anchor paths drive directory resolution.""" + + def test_nvvm_has_custom_linux_paths(self): + desc = LIB_DESCRIPTORS["nvvm"] + assert desc.anchor_rel_dirs_linux == ("nvvm/lib64",) + + def test_nvvm_has_custom_windows_paths(self): + desc = LIB_DESCRIPTORS["nvvm"] + assert desc.anchor_rel_dirs_windows == ("nvvm/bin/*", "nvvm/bin") + + @pytest.mark.parametrize("libname", ["cudart", "cublas", "nvrtc"]) + def test_regular_ctk_libs_use_defaults(self, libname): + desc = LIB_DESCRIPTORS[libname] + assert desc.anchor_rel_dirs_linux == ("lib64", "lib") + assert desc.anchor_rel_dirs_windows == ("bin/x64", "bin") + + def test_find_lib_dir_uses_descriptor_linux(self, tmp_path): + (tmp_path / "nvvm" / "lib64").mkdir(parents=True) + + desc = _make_desc(name="nvvm", anchor_rel_dirs_linux=("nvvm/lib64",)) + result = _find_lib_dir_using_anchor(desc, LinuxSearchPlatform(), str(tmp_path)) + assert result is not None + assert result.endswith(os.path.join("nvvm", "lib64")) + + def test_find_lib_dir_uses_descriptor_windows(self, tmp_path): + (tmp_path / "nvvm" / "bin").mkdir(parents=True) + + desc = _make_desc(name="nvvm", anchor_rel_dirs_windows=("nvvm/bin/*", "nvvm/bin")) + result = _find_lib_dir_using_anchor(desc, WindowsSearchPlatform(), str(tmp_path)) + assert result is not None + assert result.endswith(os.path.join("nvvm", "bin")) + + def test_find_lib_dir_returns_none_when_no_match(self, tmp_path): + desc = _make_desc(anchor_rel_dirs_linux=("nonexistent",)) + assert _find_lib_dir_using_anchor(desc, LinuxSearchPlatform(), str(tmp_path)) is None + + def test_nvvm_cuda_home_linux(self, mocker, tmp_path): + """End-to-end: find_in_cuda_home resolves nvvm under its custom subdir.""" + mocker.patch(f"{_STEPS_MOD}.get_cuda_home_or_path", return_value=str(tmp_path)) + + nvvm_dir = tmp_path / "nvvm" / "lib64" + nvvm_dir.mkdir(parents=True) + so_file = nvvm_dir / "libnvvm.so" + so_file.touch() + + desc = _make_desc( + name="nvvm", + linux_sonames=("libnvvm.so",), + anchor_rel_dirs_linux=("nvvm/lib64",), + ) + result = find_in_cuda_home(_ctx(desc, platform=LinuxSearchPlatform())) + assert result is not None + assert result.abs_path == str(so_file) + assert result.found_via == "CUDA_HOME"