diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index a64b5ef64c7..383d88de889 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/cuda_pathfinder/cuda/pathfinder/__init__.py @@ -11,6 +11,11 @@ find_nvidia_binary_utility as find_nvidia_binary_utility, ) from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES as _SUPPORTED_BINARIES +from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import LocatedDynamicLib as LocatedDynamicLib +from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import find_nvidia_dynamic_lib as find_nvidia_dynamic_lib +from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import ( + locate_nvidia_dynamic_lib as locate_nvidia_dynamic_lib, +) from cuda.pathfinder._dynamic_libs.load_dl_common import ( DynamicLibNotAvailableError as DynamicLibNotAvailableError, ) @@ -76,8 +81,8 @@ SUPPORTED_BINARY_UTILITIES = _SUPPORTED_BINARIES #: Tuple of CUDA Toolkit dynamic library names supported by -#: :func:`load_nvidia_dynamic_lib` for the current operating system and -#: interpreter architecture. +#: :func:`load_nvidia_dynamic_lib` and :func:`locate_nvidia_dynamic_lib` for +#: the current operating system and interpreter architecture. SUPPORTED_NVIDIA_LIBNAMES = _SUPPORTED_NVIDIA_LIBNAMES #: Tuple of supported bitcode library names that can be resolved 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 new file mode 100644 index 00000000000..86e71660381 --- /dev/null +++ b/cuda_pathfinder/cuda/pathfinder/_dynamic_libs/find_nvidia_dynamic_lib.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import functools +import os +from dataclasses import dataclass + +from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS +from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( + _DRIVER_ONLY_LIBNAMES, + _resolve_system_loaded_abs_path_in_subprocess, + _try_ctk_root_canary, + check_libname, + raise_driver_lib_not_found, +) +from cuda.pathfinder._dynamic_libs.search_steps import ( + EARLY_FIND_STEPS, + LATE_FIND_STEPS, + FindResult, + SearchContext, + run_find_steps, +) + + +@dataclass(frozen=True) +class LocatedDynamicLib: + """Information about a located dynamic library.""" + + name: str + abs_path: str + filename: str + found_via: str + + +def _find_with_system_search(libname: str) -> FindResult | None: + # The native loader search can only be observed by loading, so it runs in + # a fresh subprocess and the caller's process stays untouched. + abs_path = _resolve_system_loaded_abs_path_in_subprocess(libname) + if abs_path is None: + return None + return FindResult(abs_path, "system-search") + + +def _locate_lib(libname: str) -> FindResult: + desc = LIB_DESCRIPTORS[libname] + + if libname in _DRIVER_ONLY_LIBNAMES: + find = _find_with_system_search(libname) + if find is None: + raise_driver_lib_not_found(desc) + return find + + ctx = SearchContext(desc) + find = run_find_steps(ctx, EARLY_FIND_STEPS) + + # load_nvidia_dynamic_lib loads the required dependencies at this point and + # fails if one of them cannot be found. + for dep in desc.dependencies: + locate_nvidia_dynamic_lib(dep) + + if find is not None: + return find + + find = _find_with_system_search(libname) + if find is not None: + return find + + find = run_find_steps(ctx, LATE_FIND_STEPS) + if find is not None: + return find + + if desc.ctk_root_canary_anchor_libnames: + canary_abs_path = _try_ctk_root_canary(ctx) + if canary_abs_path is not None: + return FindResult(canary_abs_path, "system-ctk-root") + + ctx.raise_not_found() + + +def locate_nvidia_dynamic_lib(libname: str) -> LocatedDynamicLib: + """Locate an NVIDIA dynamic library by name without loading it. + + Returns the file that :func:`load_nvidia_dynamic_lib` would load in a fresh + process. Use it where a path is needed but the library must not be loaded + into the current process, for example to pass a library to a build system. + + Like :func:`load_nvidia_dynamic_lib`, it fails if a library that ``libname`` + requires cannot be found. Optional dependencies are not checked. + + Args: + libname (str): The short name of the library to locate (e.g., + ``"cudart"``, ``"nvvm"``, etc.). + + Returns: + LocatedDynamicLib: The absolute path of the library file and how it + was found. + + Raises: + DynamicLibUnknownError: If ``libname`` is not a recognized library name. + DynamicLibNotAvailableError: If ``libname`` is recognized but not + supported on this platform. + DynamicLibNotFoundError: If the library cannot be found. + + Search order: + The steps of :func:`load_nvidia_dynamic_lib`, in the same order, except + that a library already loaded in the current process is not taken into + account: + + 1. NVIDIA Python wheels (``site-packages``) + 2. Conda environment (``CONDA_PREFIX``) + 3. OS default mechanisms (``dlopen()`` / ``LoadLibraryExW()``), run in + a fresh Python subprocess + 4. Environment variables (library-specific roots, then ``CUDA_PATH`` / + ``CUDA_HOME``) + 5. Windows Program Files (configured libraries only) + 6. CTK root canary probe (discoverable libs only) + + Driver libraries (``"cuda"``, ``"nvml"``) use step 3 only. + + Notes: + The result can differ from what :func:`load_nvidia_dynamic_lib` returns + in the current process when another component has already loaded the + library, or influences the native loader search (for example through + its ``RPATH``). A file that is found but fails to load is still + returned, since nothing is loaded. + + The result of the native loader search (step 3) is cached for the life + of the process and shared with the CTK root canary probe, because each + search starts a Python subprocess. + """ + check_libname(libname) + find = _locate_lib(libname) + return LocatedDynamicLib( + name=libname, + abs_path=find.abs_path, + filename=os.path.basename(find.abs_path), + found_via=find.found_via, + ) + + +@functools.cache +def find_nvidia_dynamic_lib(libname: str) -> str: + """Find the absolute path of an NVIDIA dynamic library without loading it. + + See :func:`locate_nvidia_dynamic_lib` for the search order. + + Raises: + DynamicLibUnknownError: If ``libname`` is not a recognized library name. + DynamicLibNotAvailableError: If ``libname`` is recognized but not + supported on this platform. + DynamicLibNotFoundError: If the library cannot be found. + """ + return locate_nvidia_dynamic_lib(libname).abs_path 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 2b0761be1a6..1616536d632 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 @@ -7,7 +7,7 @@ import struct import subprocess import sys -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, NoReturn from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS from cuda.pathfinder._dynamic_libs.load_dl_common import ( @@ -66,12 +66,27 @@ def _load_driver_lib_no_cache(desc: LibDescriptor) -> LoadedDL: loaded = LOADER.load_with_system_search(desc) if loaded is not None: return loaded + raise_driver_lib_not_found(desc) + + +def raise_driver_lib_not_found(desc: LibDescriptor) -> NoReturn: raise DynamicLibNotFoundError( f'"{desc.name}" is an NVIDIA driver library and can only be found via' f" system search. Ensure the NVIDIA display driver is installed." ) +def check_libname(libname: str) -> None: + """Raise if ``libname`` is unknown or not available on this platform.""" + if libname not in _ALL_KNOWN_LIBNAMES: + raise DynamicLibUnknownError(f"Unknown library name: {libname!r}. Known names: {sorted(_ALL_KNOWN_LIBNAMES)}") + if libname not in ALL_AVAILABLE_LIBNAMES: + raise DynamicLibNotAvailableError( + f"Library name {libname!r} is known but not available on {_PLATFORM_NAME}. " + f"Supported names on {_PLATFORM_NAME}: {sorted(ALL_AVAILABLE_LIBNAMES)}" + ) + + def _coerce_subprocess_output(output: str | bytes | None) -> str: if isinstance(output, bytes): return output.decode(errors="replace") @@ -320,11 +335,5 @@ def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: f" Currently running: {pointer_size_bits}-bit Python" f" {sys.version_info.major}.{sys.version_info.minor}" ) - if libname not in _ALL_KNOWN_LIBNAMES: - raise DynamicLibUnknownError(f"Unknown library name: {libname!r}. Known names: {sorted(_ALL_KNOWN_LIBNAMES)}") - if libname not in ALL_AVAILABLE_LIBNAMES: - raise DynamicLibNotAvailableError( - f"Library name {libname!r} is known but not available on {_PLATFORM_NAME}. " - f"Supported names on {_PLATFORM_NAME}: {sorted(ALL_AVAILABLE_LIBNAMES)}" - ) + check_libname(libname) return _load_lib_no_cache(libname) diff --git a/cuda_pathfinder/docs/source/api.rst b/cuda_pathfinder/docs/source/api.rst index f65014923f9..c02e8846c88 100644 --- a/cuda_pathfinder/docs/source/api.rst +++ b/cuda_pathfinder/docs/source/api.rst @@ -6,8 +6,8 @@ ``cuda.pathfinder`` API Reference ================================= -The ``cuda.pathfinder`` module provides utilities for loading NVIDIA dynamic libraries, -locating NVIDIA C/C++ header directories, finding CUDA binary utilities, and locating +The ``cuda.pathfinder`` module provides utilities for loading and locating NVIDIA dynamic +libraries, locating NVIDIA C/C++ header directories, finding CUDA binary utilities, and locating CUDA bitcode and static libraries. .. NOTE: The source of truth for public APIs is cuda_pathfinder/cuda/pathfinder/__init__.py. @@ -21,6 +21,9 @@ CUDA bitcode and static libraries. SUPPORTED_NVIDIA_LIBNAMES load_nvidia_dynamic_lib LoadedDL + find_nvidia_dynamic_lib + locate_nvidia_dynamic_lib + LocatedDynamicLib DynamicLibNotFoundError DynamicLibUnknownError DynamicLibNotAvailableError diff --git a/cuda_pathfinder/tests/test_find_nvidia_dynamic_lib.py b/cuda_pathfinder/tests/test_find_nvidia_dynamic_lib.py new file mode 100644 index 00000000000..ba6920c0853 --- /dev/null +++ b/cuda_pathfinder/tests/test_find_nvidia_dynamic_lib.py @@ -0,0 +1,259 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest +from child_load_nvidia_dynamic_lib_helper import ( + build_child_process_failed_for_libname_message, + run_load_nvidia_dynamic_lib_in_subprocess, +) + +from cuda.pathfinder import ( + DynamicLibNotAvailableError, + DynamicLibNotFoundError, + DynamicLibUnknownError, + LocatedDynamicLib, + find_nvidia_dynamic_lib, + locate_nvidia_dynamic_lib, +) +from cuda.pathfinder._dynamic_libs import find_nvidia_dynamic_lib as find_mod +from cuda.pathfinder._dynamic_libs import load_nvidia_dynamic_lib as load_mod +from cuda.pathfinder._dynamic_libs import supported_nvidia_libs +from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import _resolve_system_loaded_abs_path_in_subprocess +from cuda.pathfinder._dynamic_libs.search_steps import EARLY_FIND_STEPS, LATE_FIND_STEPS, FindResult +from cuda.pathfinder._dynamic_libs.subprocess_protocol import STATUS_NOT_FOUND, parse_dynamic_lib_subprocess_payload +from cuda.pathfinder._utils.platform_aware import IS_WINDOWS, quote_for_shell + +_MODULE = "cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib" + + +@pytest.fixture(autouse=True) +def _clear_caches(): + find_nvidia_dynamic_lib.cache_clear() + _resolve_system_loaded_abs_path_in_subprocess.cache_clear() + yield + find_nvidia_dynamic_lib.cache_clear() + _resolve_system_loaded_abs_path_in_subprocess.cache_clear() + + +@pytest.fixture +def no_in_process_loading(mocker): + for name in ("check_if_already_loaded_from_elsewhere", "load_with_system_search", "load_with_abs_path"): + mocker.patch.object(load_mod.LOADER, name, side_effect=AssertionError(f"LOADER.{name} was called")) + + +def _patch_find_steps(mocker, *, early=None, late=None): + def run_find_steps(_ctx, steps): + if steps is EARLY_FIND_STEPS: + return early + assert steps is LATE_FIND_STEPS + return late + + return mocker.patch(f"{_MODULE}.run_find_steps", side_effect=run_find_steps) + + +@pytest.mark.parametrize( + "libname", + supported_nvidia_libs.SUPPORTED_WINDOWS_DLLS if IS_WINDOWS else supported_nvidia_libs.SUPPORTED_LINUX_SONAMES, +) +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_locate_matches_load_in_fresh_process(info_summary_append, libname): + timeout = 120 if IS_WINDOWS else 30 + result = run_load_nvidia_dynamic_lib_in_subprocess(libname, timeout=timeout) + if result.returncode != 0: + raise RuntimeError(build_child_process_failed_for_libname_message(libname, result)) + payload = parse_dynamic_lib_subprocess_payload( + result.stdout, + libname=libname, + error_label="Load subprocess child process", + ) + + if payload.status == STATUS_NOT_FOUND: + with pytest.raises(DynamicLibNotFoundError): + locate_nvidia_dynamic_lib(libname) + info_summary_append(f"Not found: {libname=!r}") + return + + located = locate_nvidia_dynamic_lib(libname) + info_summary_append(f"abs_path={quote_for_shell(located.abs_path)} found_via={located.found_via}") + assert located.name == libname + assert located.filename == os.path.basename(located.abs_path) + assert os.path.samefile(located.abs_path, payload.abs_path) + assert find_nvidia_dynamic_lib(libname) == located.abs_path + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_early_find_skips_system_search(mocker, tmp_path): + lib = tmp_path / "libcudart.so.13" + _patch_find_steps(mocker, early=FindResult(str(lib), "site-packages")) + probe = mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess") + + assert locate_nvidia_dynamic_lib("cudart") == LocatedDynamicLib( + name="cudart", + abs_path=str(lib), + filename="libcudart.so.13", + found_via="site-packages", + ) + probe.assert_not_called() + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_system_search_runs_before_late_find_steps(mocker, tmp_path): + lib = tmp_path / "libcudart.so.13" + run_find_steps = _patch_find_steps(mocker, late=FindResult("unused", "CUDA_PATH")) + probe = mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=str(lib)) + + located = locate_nvidia_dynamic_lib("cudart") + + assert located.abs_path == str(lib) + assert located.found_via == "system-search" + probe.assert_called_once_with("cudart") + assert [call.args[1] for call in run_find_steps.call_args_list] == [EARLY_FIND_STEPS] + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_late_find_steps_after_system_search(mocker, tmp_path): + lib = tmp_path / "libcudart.so.13" + _patch_find_steps(mocker, late=FindResult(str(lib), "CUDA_PATH")) + mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) + + located = locate_nvidia_dynamic_lib("cudart") + + assert located.abs_path == str(lib) + assert located.found_via == "CUDA_PATH" + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_ctk_root_canary_is_last(mocker, tmp_path): + lib = tmp_path / "libnvvm.so.4" + _patch_find_steps(mocker) + mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) + canary = mocker.patch(f"{_MODULE}._try_ctk_root_canary", return_value=str(lib)) + + located = locate_nvidia_dynamic_lib("nvvm") + + assert located.abs_path == str(lib) + assert located.found_via == "system-ctk-root" + canary.assert_called_once() + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_not_found_raises(mocker): + _patch_find_steps(mocker) + mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) + mocker.patch(f"{_MODULE}._try_ctk_root_canary", return_value=None) + + with pytest.raises(DynamicLibNotFoundError, match="Failure finding"): + locate_nvidia_dynamic_lib("nvvm") + with pytest.raises(DynamicLibNotFoundError, match="Failure finding"): + find_nvidia_dynamic_lib("nvvm") + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_driver_lib_uses_system_search_only(mocker, tmp_path): + lib = tmp_path / "libcuda.so.1" + run_find_steps = mocker.patch(f"{_MODULE}.run_find_steps") + mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=str(lib)) + + located = locate_nvidia_dynamic_lib("cuda") + + assert located.abs_path == str(lib) + assert located.found_via == "system-search" + run_find_steps.assert_not_called() + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_driver_lib_not_found_raises(mocker): + mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) + + with pytest.raises(DynamicLibNotFoundError, match="NVIDIA driver library"): + locate_nvidia_dynamic_lib("cuda") + + +def _patch_early_find_by_libname(mocker, tmp_path, found): + """Early find steps hit for the libnames in ``found``; everything else misses.""" + + def run_find_steps(ctx, steps): + if steps is EARLY_FIND_STEPS and ctx.libname in found: + return FindResult(str(tmp_path / f"lib{ctx.libname}.so"), "site-packages") + return None + + mocker.patch(f"{_MODULE}.run_find_steps", side_effect=run_find_steps) + mocker.patch(f"{_MODULE}._resolve_system_loaded_abs_path_in_subprocess", return_value=None) + mocker.patch(f"{_MODULE}._try_ctk_root_canary", return_value=None) + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_missing_required_dependency_raises(mocker, tmp_path): + _patch_early_find_by_libname(mocker, tmp_path, {"cublas"}) + + with pytest.raises(DynamicLibNotFoundError, match="cublasLt"): + locate_nvidia_dynamic_lib("cublas") + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_required_dependencies_are_checked_transitively(mocker, tmp_path): + # cusolver requires cublas, which requires cublasLt. + _patch_early_find_by_libname(mocker, tmp_path, {"cusolver", "nvJitLink", "cusparse", "cublas"}) + + with pytest.raises(DynamicLibNotFoundError, match="cublasLt"): + locate_nvidia_dynamic_lib("cusolver") + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_present_dependencies_do_not_change_the_result(mocker, tmp_path): + _patch_early_find_by_libname(mocker, tmp_path, {"cublas", "cublasLt"}) + + located = locate_nvidia_dynamic_lib("cublas") + + assert located.abs_path == str(tmp_path / "libcublas.so") + assert located.found_via == "site-packages" + + +@pytest.mark.usefixtures("no_in_process_loading") +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_optional_dependencies_are_not_required(mocker, tmp_path): + desc = load_mod.LIB_DESCRIPTORS["cudss"] + assert desc.optional_dependencies + _patch_early_find_by_libname(mocker, tmp_path, {"cudss", *desc.dependencies}) + + assert locate_nvidia_dynamic_lib("cudss").abs_path == str(tmp_path / "libcudss.so") + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_find_is_cached(mocker, tmp_path): + lib = tmp_path / "libcudart.so.13" + locate = mocker.patch( + f"{_MODULE}.locate_nvidia_dynamic_lib", + return_value=LocatedDynamicLib("cudart", str(lib), lib.name, "site-packages"), + ) + + assert find_mod.find_nvidia_dynamic_lib("cudart") == str(lib) + assert find_mod.find_nvidia_dynamic_lib("cudart") == str(lib) + locate.assert_called_once_with("cudart") + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_unknown_libname_raises_dynamic_lib_unknown_error(): + with pytest.raises(DynamicLibUnknownError, match=r"Unknown library name: 'not_a_real_lib'"): + locate_nvidia_dynamic_lib("not_a_real_lib") + + +@pytest.mark.agent_authored(model="claude-opus-5-5") +def test_known_but_platform_unavailable_libname_raises_dynamic_lib_not_available_error(monkeypatch): + monkeypatch.setattr(load_mod, "_ALL_KNOWN_LIBNAMES", frozenset(("known_but_unavailable",))) + monkeypatch.setattr(load_mod, "ALL_AVAILABLE_LIBNAMES", frozenset()) + monkeypatch.setattr(load_mod, "_PLATFORM_NAME", "TestOS") + with pytest.raises(DynamicLibNotAvailableError, match=r"known_but_unavailable.*not available on TestOS"): + locate_nvidia_dynamic_lib("known_but_unavailable")