File size: 7,292 Bytes
76f9669 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 | # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import functools
import struct
import sys
from cuda.pathfinder._dynamic_libs.find_nvidia_dynamic_lib import _FindNvidiaDynamicLib
from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError, LoadedDL, load_dependencies
from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
SUPPORTED_LINUX_SONAMES,
SUPPORTED_WINDOWS_DLLS,
)
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
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,
)
# 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()
)
# 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"))
def _load_driver_lib_no_cache(libname: str) -> LoadedDL:
"""Load an NVIDIA driver library (system-search only).
Driver libs (libcuda, libnvidia-ml) are part of the display driver, not
the CUDA Toolkit. They are always on the system linker path, so the
full CTK search cascade (site-packages, conda, CUDA_HOME, canary) is
unnecessary.
"""
loaded = check_if_already_loaded_from_elsewhere(libname, False)
if loaded is not None:
return loaded
loaded = load_with_system_search(libname)
if loaded is not None:
return loaded
raise DynamicLibNotFoundError(
f'"{libname}" is an NVIDIA driver library and can only be found via'
f" system search. Ensure the NVIDIA display driver is installed."
)
def _load_lib_no_cache(libname: str) -> LoadedDL:
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)
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
abs_path = finder.try_with_cuda_home()
if abs_path is None:
finder.raise_not_found_error()
else:
found_via = "CUDA_HOME"
return load_with_abs_path(libname, abs_path, found_via)
@functools.cache
def load_nvidia_dynamic_lib(libname: str) -> LoadedDL:
"""Load an NVIDIA dynamic library by name.
Args:
libname (str): The short name of the library to load (e.g., ``"cudart"``,
``"nvvm"``, etc.).
Returns:
LoadedDL: Object containing the OS library handle and absolute path.
**Important:**
**Never close the returned handle.** Do **not** call ``dlclose`` (Linux) or
``FreeLibrary`` (Windows) on the ``LoadedDL._handle_uint``.
**Why:** the return value is cached (``functools.cache``) and shared across the
process. Closing the handle can unload the module while other code still uses
it, leading to crashes or subtle failures.
This applies to Linux and Windows. For context, see issue #1011:
https://github.com/NVIDIA/cuda-python/issues/1011
Raises:
ValueError: If ``libname`` is not a recognized library name.
DynamicLibNotFoundError: If the library cannot be found or loaded.
RuntimeError: If Python is not 64-bit.
Search order:
0. **Already loaded in the current process**
- If a matching library is already loaded by some other component,
return its absolute path and handle and skip the rest of the search.
1. **NVIDIA Python wheels**
- Scan installed distributions (``site-packages``) to find libraries
shipped in NVIDIA wheels.
2. **Conda environment**
- Conda installations are discovered via ``CONDA_PREFIX``, which is
defined automatically in activated conda environments (see
https://docs.conda.io/projects/conda-build/en/stable/user-guide/environment-variables.html).
3. **OS default mechanisms**
- Fall back to the native loader:
- Linux: ``dlopen()``
- Windows: ``LoadLibraryW()``
- CUDA Toolkit (CTK) system installs with system config updates are often
discovered via:
- Linux: ``/etc/ld.so.conf.d/*cuda*.conf``
- Windows: ``C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\vX.Y\\bin``
on the system ``PATH``.
4. **Environment variables**
- If set, use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
**Driver libraries** (``"cuda"``, ``"nvml"``):
These are part of the NVIDIA display driver (not the CUDA Toolkit) and
are always on the system linker path. For these libraries the search
is simplified to:
0. Already loaded in the current process
1. OS default mechanisms (``dlopen`` / ``LoadLibraryW``)
The CTK-specific steps (site-packages, conda, ``CUDA_HOME``, canary
probe) are skipped entirely.
Notes:
The search is performed **per library**. There is currently no mechanism to
guarantee that multiple libraries are all resolved from the same location.
"""
pointer_size_bits = struct.calcsize("P") * 8
if pointer_size_bits != 64:
raise RuntimeError(
f"cuda.pathfinder.load_nvidia_dynamic_lib() requires 64-bit Python."
f" Currently running: {pointer_size_bits}-bit Python"
f" {sys.version_info.major}.{sys.version_info.minor}"
)
if libname not in _ALL_SUPPORTED_LIBNAMES:
raise ValueError(f"Unsupported library name: {libname!r}. Supported names: {sorted(_ALL_SUPPORTED_LIBNAMES)}")
return _load_lib_no_cache(libname)
|