Coverage for cuda/pathfinder/_dynamic_libs/load_nvidia_dynamic_lib.py: 97.94%

1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations (empty)

5 

6import functools (empty)

7import struct (empty)

8import subprocess (empty)

9import sys (empty)

10from typing import TYPE_CHECKING (empty)

11 

12from cuda.pathfinder._dynamic_libs.lib_descriptor import LIB_DESCRIPTORS (empty)

13from cuda.pathfinder._dynamic_libs.load_dl_common import ( (empty)

14 DynamicLibNotAvailableError, 

15 DynamicLibNotFoundError, 

16 DynamicLibUnknownError, 

17 LoadedDL, 

18 load_dependencies, 

19) 

20from cuda.pathfinder._dynamic_libs.platform_loader import LOADER (empty)

21from cuda.pathfinder._dynamic_libs.search_steps import ( (empty)

22 EARLY_FIND_STEPS, 

23 LATE_FIND_STEPS, 

24 SearchContext, 

25 derive_ctk_root, 

26 find_via_ctk_root, 

27 run_find_steps, 

28) 

29from cuda.pathfinder._dynamic_libs.subprocess_protocol import ( (empty)

30 DYNAMIC_LIB_SUBPROCESS_CWD, 

31 MODE_CANARY, 

32 STATUS_OK, 

33 DynamicLibSubprocessPayload, 

34 build_dynamic_lib_subprocess_command, 

35 parse_dynamic_lib_subprocess_payload, 

36) 

37from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ALL_AVAILABLE_LIBNAMES (empty)

38from cuda.pathfinder._utils.platform_aware import IS_WINDOWS (empty)

39 

40if TYPE_CHECKING: (empty)

41 from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor 

42 

43# All libnames recognized by load_nvidia_dynamic_lib, across all categories 

44# (CTK, third-party, driver). 

45_ALL_KNOWN_LIBNAMES: frozenset[str] = frozenset(LIB_DESCRIPTORS) (empty)

46_PLATFORM_NAME = "Windows" if IS_WINDOWS else "Linux" (empty)

47_CANARY_PROBE_TIMEOUT_SECONDS = 10.0 (empty)

48 

49# Driver libraries: shipped with the NVIDIA display driver, always on the 

50# system linker path. These skip all CTK search steps (site-packages, 

51# conda, CUDA_PATH, canary) and go straight to system search. 

52_DRIVER_ONLY_LIBNAMES = frozenset(name for name, desc in LIB_DESCRIPTORS.items() if desc.packaged_with == "driver") (empty)

53 

54 

55def _load_driver_lib_no_cache(desc: LibDescriptor) -> LoadedDL: (empty)

56 """Load an NVIDIA driver library (system-search only). 

57 

58 Driver libs (libcuda, libnvidia-ml) are part of the display driver, not 

59 the CUDA Toolkit. They are expected to be discoverable via the platform's 

60 native loader mechanisms, so the full CTK search cascade (site-packages, 

61 conda, CUDA_PATH, canary) is unnecessary. 

62 """ 

63 loaded = LOADER.check_if_already_loaded_from_elsewhere(desc) 6 ctx1aCDE1f

64 if loaded is not None: 6 ctx1aCDE1f

65 return loaded 2 ctx1a1

66 loaded = LOADER.load_with_system_search(desc) 5 ctx1aCDEf

67 if loaded is not None: 5 ctx1aCDEf

68 return loaded 4 ctx1aCDf

69 raise DynamicLibNotFoundError( 1 ctx1E

70 f'"{desc.name}" is an NVIDIA driver library and can only be found via' 

71 f" system search. Ensure the NVIDIA display driver is installed." 

72 ) 

73 

74 

75def _coerce_subprocess_output(output: str | bytes | None) -> str: (empty)

76 if isinstance(output, bytes): 2 ctx1no

77 return output.decode(errors="replace") 

78 return "" if output is None else output 2 ctx1no

79 

80 

81def _raise_canary_probe_child_process_error( (empty)

82 *, 

83 returncode: int | None = None, 

84 timeout: float | None = None, 

85 stderr: str | bytes | None = None, 

86) -> None: 

87 if timeout is None: 2 ctx1no

88 error_line = f"Canary probe child process exited with code {returncode}." 1 ctx1n

89 else: 

90 error_line = f"Canary probe child process timed out after {timeout} seconds." 1 ctx1o

91 raise ChildProcessError( 2 ctx1no

92 f"{error_line}\n" 

93 "--- stderr-from-child-process ---\n" 

94 f"{_coerce_subprocess_output(stderr)}" 

95 "<end-of-stderr-from-child-process>\n" 

96 ) 

97 

98 

99@functools.cache (empty)

100def _resolve_system_loaded_abs_path_in_subprocess( (empty)

101 libname: str, 

102 *, 

103 timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS, 

104) -> str | None: 

105 """Resolve a canary library's absolute path in a fresh Python subprocess.""" 

106 try: 15 ctx1nOPoQzAsijtuBvw

107 result = subprocess.run( # noqa: S603 - trusted argv: current interpreter + internal probe module 15 ctx1nOPoQzAsijtuBvw

108 build_dynamic_lib_subprocess_command(MODE_CANARY, libname), 

109 capture_output=True, 

110 text=True, 

111 timeout=timeout, 

112 check=False, 

113 cwd=DYNAMIC_LIB_SUBPROCESS_CWD, 

114 ) 

115 except subprocess.TimeoutExpired as exc: 1 ctx1o

116 _raise_canary_probe_child_process_error(timeout=exc.timeout, stderr=exc.stderr) 1 ctx1o

117 

118 if result.returncode != 0: 14 ctx1nOPQzAsijtuBvw

119 _raise_canary_probe_child_process_error(returncode=result.returncode, stderr=result.stderr) 1 ctx1n

120 

121 payload: DynamicLibSubprocessPayload = parse_dynamic_lib_subprocess_payload( 13 ctx1OPQzAsijtuBvw

122 result.stdout, 

123 libname=libname, 

124 error_label="Canary probe child process", 

125 ) 

126 abs_path: str | None = payload.abs_path 10 ctx1zAsijtuBvw

127 if payload.status == STATUS_OK: 10 ctx1zAsijtuBvw

128 return abs_path 5 ctx1zituv

129 return None 8 ctx1AsjtuBvw

130 

131 

132def _loadable_via_canary_subprocess(libname: str, *, timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS) -> bool: (empty)

133 """Return True if the canary subprocess can resolve ``libname`` via system search.""" 

134 return _resolve_system_loaded_abs_path_in_subprocess(libname, timeout=timeout) is not None 3 ctx1sw7

135 

136 

137def resolve_ctk_root_via_canary(canary_libname: str) -> str | None: (empty)

138 """Resolve the CUDA Toolkit root from a system-loadable canary library. 

139 

140 The canary library's absolute path is resolved by the OS dynamic loader in 

141 an isolated subprocess, which honors ``LD_LIBRARY_PATH`` on Linux and the 

142 native DLL search on Windows. The toolkit root is then derived from that 

143 path. Returns ``None`` if the canary cannot be resolved or no root can be 

144 derived. 

145 """ 

146 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname) 27 ctx1lbpmxy2RSFTGHIJUVijWXKYLMZc

147 if canary_abs_path is None: 27 ctx1lbpmxy2RSFTGHIJUVijWXKYLMZc

148 return None 14 ctx1xy2FGHIJijKLMc

149 ctk_root: str | None = derive_ctk_root(canary_abs_path) 23 ctx1lbpmRSFTGHIJUVijWXKYLMZ

150 return ctk_root 23 ctx1lbpmRSFTGHIJUVijWXKYLMZ

151 

152 

153def _try_ctk_root_canary(ctx: SearchContext) -> str | None: (empty)

154 """Try CTK-root canary fallback for descriptor-configured libraries.""" 

155 for canary_libname in ctx.desc.ctk_root_canary_anchor_libnames: 7 ctx1lbpmxyc

156 ctk_root = resolve_ctk_root_via_canary(canary_libname) 7 ctx1lbpmxyc

157 if ctk_root is None: 7 ctx1lbpmxyc

158 continue 4 ctx1pxyc

159 find = find_via_ctk_root(ctx, ctk_root) 3 ctx1lbm

160 if find is not None: 3 ctx1lbm

161 return str(find.abs_path) 2 ctx1lb

162 return None 5 ctx1pmxyc

163 

164 

165def _load_lib_no_cache(libname: str) -> LoadedDL: (empty)

166 desc = LIB_DESCRIPTORS[libname] 14 ctx1abged345kfhqcr

167 

168 if libname in _DRIVER_ONLY_LIBNAMES: 14 ctx1abged345kfhqcr

169 return _load_driver_lib_no_cache(desc) 5 ctx1a345f

170 

171 ctx = SearchContext(desc) 10 ctx1abgedkhqcr

172 

173 # Phase 1: Try to find the library file on disk (pip wheels, conda). 

174 find = run_find_steps(ctx, EARLY_FIND_STEPS) 10 ctx1abgedkhqcr

175 

176 # Phase 2: Cross-cutting — already-loaded check and dependency loading. 

177 loaded = LOADER.check_if_already_loaded_from_elsewhere(desc) 10 ctx1abgedkhqcr

178 load_dependencies(desc, load_nvidia_dynamic_lib) 10 ctx1abgedkhqcr

179 if loaded is not None: 10 ctx1abgedkhqcr

180 return loaded 

181 

182 # Phase 3: Load from found path, or fall back to system search + late find. 

183 if find is not None: 10 ctx1abgedkhqcr

184 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 2 ctx1qr

185 

186 loaded = LOADER.load_with_system_search(desc) 8 ctx1abgedkhc

187 if loaded is not None: 8 ctx1abgedkhc

188 return loaded 3 ctx1adk

189 

190 find = run_find_steps(ctx, LATE_FIND_STEPS) 6 ctx1abgehc

191 if find is not None: 6 ctx1abgehc

192 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 3 ctx1agh

193 

194 if desc.ctk_root_canary_anchor_libnames: 3 ctx1bec

195 canary_abs_path = _try_ctk_root_canary(ctx) 2 ctx1bc

196 if canary_abs_path is not None: 2 ctx1bc

197 return LOADER.load_with_abs_path(desc, canary_abs_path, "system-ctk-root") 1 ctx1b

198 

199 ctx.raise_not_found() 2 ctx1ec

200 

201 

202@functools.cache (empty)

203def load_nvidia_dynamic_lib(libname: str) -> LoadedDL: (empty)

204 """Load an NVIDIA dynamic library by name. 

205 

206 Args: 

207 libname (str): The short name of the library to load (e.g., ``"cudart"``, 

208 ``"nvvm"``, etc.). 

209 

210 Returns: 

211 LoadedDL: Object containing the OS library handle and absolute path. 

212 

213 **Important:** 

214 

215 **Never close the returned handle.** Do **not** call ``dlclose`` (Linux) or 

216 ``FreeLibrary`` (Windows) on the ``LoadedDL._handle_uint``. 

217 

218 **Why:** the return value is cached (``functools.cache``) and shared across the 

219 process. Closing the handle can unload the module while other code still uses 

220 it, leading to crashes or subtle failures. 

221 

222 This applies to Linux and Windows. For context, see issue #1011: 

223 https://github.com/NVIDIA/cuda-python/issues/1011 

224 

225 Raises: 

226 DynamicLibUnknownError: If ``libname`` is not a recognized library name. 

227 DynamicLibNotAvailableError: If ``libname`` is recognized but not 

228 supported on this platform. 

229 DynamicLibNotFoundError: If the library cannot be found or loaded. 

230 RuntimeError: If Python is not 64-bit. 

231 

232 Windows on ARM (WoA) Note: 

233 On Windows, this API aims to load a dynamic library whose architecture 

234 matches the Python interpreter architecture. For example, x64 Python 

235 running on an Arm64 machine targets an x64 DLL, while native Arm64 Python 

236 targets an Arm64 DLL. A library loaded into the Python process must be 

237 compatible with that process. This differs from 

238 ``find_nvidia_binary_utility``, which targets the native machine 

239 architecture when selecting architecture-specific executables. 

240 

241 Search order: 

242 0. **Already loaded in the current process** 

243 

244 - If a matching library is already loaded by some other component, 

245 return its absolute path and handle and skip the rest of the search. 

246 

247 1. **NVIDIA Python wheels** 

248 

249 - Scan installed distributions (``site-packages``) to find libraries 

250 shipped in NVIDIA wheels. 

251 

252 2. **Conda environment** 

253 

254 - Conda installations are discovered via ``CONDA_PREFIX``, which is 

255 defined automatically in activated conda environments (see 

256 https://docs.conda.io/projects/conda-build/en/stable/user-guide/environment-variables.html). 

257 

258 3. **OS default mechanisms** 

259 

260 - Fall back to the native loader: 

261 

262 - Linux: ``dlopen()`` 

263 

264 - Windows: ``LoadLibraryExW()`` 

265 

266 On Linux, CUDA Toolkit (CTK) system installs with system config updates are 

267 usually discovered via ``/etc/ld.so.conf.d/*cuda*.conf``. 

268 

269 On Windows, under Python 3.8+, CPython configures the process with 

270 ``SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS)``. 

271 As a result, the native DLL search used here does **not** include 

272 the system ``PATH``. 

273 

274 4. **Environment variables** 

275 

276 - First search library-specific roots declared by the descriptor, 

277 such as ``CUDNN_PATH`` and ``NCCL_HOME``, using their 

278 platform-specific product layouts. Then use ``CUDA_PATH`` or 

279 ``CUDA_HOME`` (in that order). 

280 On Windows, ``CUDA_PATH`` is the typical way system-installed CTK 

281 DLLs are located. Note that the NVIDIA CTK installer automatically 

282 adds ``CUDA_PATH`` to the system-wide environment. 

283 

284 5. **Windows Program Files (configured libraries only)** 

285 

286 - Search descriptor-configured standalone installation roots, such 

287 as versioned x64 cuDNN directories under ``ProgramFiles``, using 

288 the general per-library anchor layout. 

289 

290 6. **CTK root canary probe (discoverable libs only)** 

291 

292 - For selected libraries whose shared object doesn't reside on the 

293 standard linker path (currently ``nvvm``), attempt to derive CTK 

294 root by system-loading a well-known CTK canary library in a 

295 subprocess and then searching relative to that root. On Windows, 

296 the canary uses the same native ``LoadLibraryExW`` semantics as 

297 step 3, so there is also no ``PATH``-based discovery. 

298 

299 **Driver libraries** (``"cuda"``, ``"nvml"``): 

300 

301 These are part of the NVIDIA display driver (not the CUDA Toolkit) and 

302 are expected to be reachable via the native OS loader path. For these 

303 libraries the search is simplified to: 

304 

305 0. Already loaded in the current process 

306 1. OS default mechanisms (``dlopen`` / ``LoadLibraryExW``) 

307 

308 The non-driver steps (site-packages, conda, environment roots, 

309 ``ProgramFiles``, and canary probe) are skipped entirely. 

310 

311 Notes: 

312 The search is performed **per library**. There is currently no mechanism to 

313 guarantee that multiple libraries are all resolved from the same location. 

314 

315 """ 

316 pointer_size_bits = struct.calcsize("P") * 8 6 ctx1adfN60

317 if pointer_size_bits != 64: 6 ctx1adfN60

318 raise RuntimeError( 1 ctx16

319 f"cuda.pathfinder.load_nvidia_dynamic_lib() requires 64-bit Python." 

320 f" Currently running: {pointer_size_bits}-bit Python" 

321 f" {sys.version_info.major}.{sys.version_info.minor}" 

322 ) 

323 if libname not in _ALL_KNOWN_LIBNAMES: 5 ctx1adfN0

324 raise DynamicLibUnknownError(f"Unknown library name: {libname!r}. Known names: {sorted(_ALL_KNOWN_LIBNAMES)}") 1 ctx10

325 if libname not in ALL_AVAILABLE_LIBNAMES: 4 ctx1adfN

326 raise DynamicLibNotAvailableError( 1 ctx1N

327 f"Library name {libname!r} is known but not available on {_PLATFORM_NAME}. " 

328 f"Supported names on {_PLATFORM_NAME}: {sorted(ALL_AVAILABLE_LIBNAMES)}" 

329 ) 

330 return _load_lib_no_cache(libname) 3 ctx1adf