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._utils.platform_aware import IS_WINDOWS (empty)

38 

39if TYPE_CHECKING: (empty)

40 from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor 

41 

42# All libnames recognized by load_nvidia_dynamic_lib, across all categories 

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

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

45_ALL_SUPPORTED_LIBNAMES: frozenset[str] = frozenset( (empty)

46 name for name, desc in LIB_DESCRIPTORS.items() if (desc.windows_dlls if IS_WINDOWS else desc.linux_sonames) 

47) 

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

49_CANARY_PROBE_TIMEOUT_SECONDS = 10.0 (empty)

50 

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

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

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

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

55 

56 

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

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

59 

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

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

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

63 conda, CUDA_PATH, canary) is unnecessary. 

64 """ 

65 loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, False) 6 ctx1aDEF1f

66 if loaded is not None: 6 ctx1aDEF1f

67 return loaded 2 ctx1a1

68 loaded = LOADER.load_with_system_search(desc) 5 ctx1aDEFf

69 if loaded is not None: 5 ctx1aDEFf

70 return loaded 4 ctx1aDEf

71 raise DynamicLibNotFoundError( 1 ctx1F

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

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

74 ) 

75 

76 

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

78 if isinstance(output, bytes): 2 ctx1mn

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

80 return "" if output is None else output 2 ctx1mn

81 

82 

83def _raise_canary_probe_child_process_error( (empty)

84 *, 

85 returncode: int | None = None, 

86 timeout: float | None = None, 

87 stderr: str | bytes | None = None, 

88) -> None: 

89 if timeout is None: 2 ctx1mn

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

91 else: 

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

93 raise ChildProcessError( 2 ctx1mn

94 f"{error_line}\n" 

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

96 f"{_coerce_subprocess_output(stderr)}" 

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

98 ) 

99 

100 

101@functools.cache (empty)

102def _resolve_system_loaded_abs_path_in_subprocess( (empty)

103 libname: str, 

104 *, 

105 timeout: float = _CANARY_PROBE_TIMEOUT_SECONDS, 

106) -> str | None: 

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

108 try: 17 ctx1mRSnTwxrhiyzABsCt

109 result = subprocess.run( # noqa: S603 - trusted argv: current interpreter + internal probe module 17 ctx1mRSnTwxrhiyzABsCt

110 build_dynamic_lib_subprocess_command(MODE_CANARY, libname), 

111 capture_output=True, 

112 text=True, 

113 timeout=timeout, 

114 check=False, 

115 cwd=DYNAMIC_LIB_SUBPROCESS_CWD, 

116 ) 

117 except subprocess.TimeoutExpired as exc: 1 ctx1n

118 _raise_canary_probe_child_process_error(timeout=exc.timeout, stderr=exc.stderr) 1 ctx1n

119 

120 if result.returncode != 0: 16 ctx1mRSTwxrhiyzABsCt

121 _raise_canary_probe_child_process_error(returncode=result.returncode, stderr=result.stderr) 1 ctx1m

122 

123 payload: DynamicLibSubprocessPayload = parse_dynamic_lib_subprocess_payload( 15 ctx1RSTwxrhiyzABsCt

124 result.stdout, 

125 libname=libname, 

126 error_label="Canary probe child process", 

127 ) 

128 abs_path: str | None = payload.abs_path 12 ctx1wxrhiyzABsCt

129 if payload.status == STATUS_OK: 12 ctx1wxrhiyzABsCt

130 return abs_path 5 ctx1wizAs

131 return None 8 ctx1xrhyBsCt

132 

133 

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

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

136 return _resolve_system_loaded_abs_path_in_subprocess(libname, timeout=timeout) is not None 3 ctx1rt7

137 

138 

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

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

141 

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

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

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

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

146 derived. 

147 """ 

148 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname) 26 ctx1kboluv2UVGWHhIJXKLiMYNZOPc

149 if canary_abs_path is None: 26 ctx1kboluv2UVGWHhIJXKLiMYNZOPc

150 return None 16 ctx1uv2GHhIJKLiMNOPc

151 ctk_root: str | None = derive_ctk_root(canary_abs_path) 22 ctx1kbolUVGWHhIJXKLiMYNZOP

152 return ctk_root 22 ctx1kbolUVGWHhIJXKLiMYNZOP

153 

154 

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

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

157 for canary_libname in ctx.desc.ctk_root_canary_anchor_libnames: 7 ctx1kboluvc

158 ctk_root = resolve_ctk_root_via_canary(canary_libname) 7 ctx1kboluvc

159 if ctk_root is None: 7 ctx1kboluvc

160 continue 4 ctx1ouvc

161 find = find_via_ctk_root(ctx, ctk_root) 3 ctx1kbl

162 if find is not None: 3 ctx1kbl

163 return str(find.abs_path) 2 ctx1kb

164 return None 5 ctx1oluvc

165 

166 

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

168 desc = LIB_DESCRIPTORS[libname] 13 ctx1abged345jfpcq

169 

170 if libname in _DRIVER_ONLY_LIBNAMES: 13 ctx1abged345jfpcq

171 return _load_driver_lib_no_cache(desc) 5 ctx1a345f

172 

173 ctx = SearchContext(desc) 9 ctx1abgedjpcq

174 

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

176 find = run_find_steps(ctx, EARLY_FIND_STEPS) 9 ctx1abgedjpcq

177 

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

179 # The already-loaded check on Windows uses the "have we found a path?" 

180 # flag to decide whether to apply AddDllDirectory side-effects. 

181 loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, find is not None) 9 ctx1abgedjpcq

182 load_dependencies(desc, load_nvidia_dynamic_lib) 9 ctx1abgedjpcq

183 if loaded is not None: 9 ctx1abgedjpcq

184 return loaded 

185 

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

187 if find is not None: 9 ctx1abgedjpcq

188 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 2 ctx1pq

189 

190 loaded = LOADER.load_with_system_search(desc) 7 ctx1abgedjc

191 if loaded is not None: 7 ctx1abgedjc

192 return loaded 3 ctx1adj

193 

194 find = run_find_steps(ctx, LATE_FIND_STEPS) 5 ctx1abgec

195 if find is not None: 5 ctx1abgec

196 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 2 ctx1ag

197 

198 if desc.ctk_root_canary_anchor_libnames: 3 ctx1bec

199 canary_abs_path = _try_ctk_root_canary(ctx) 2 ctx1bc

200 if canary_abs_path is not None: 2 ctx1bc

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

202 

203 ctx.raise_not_found() 2 ctx1ec

204 

205 

206@functools.cache (empty)

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

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

209 

210 Args: 

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

212 ``"nvvm"``, etc.). 

213 

214 Returns: 

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

216 

217 **Important:** 

218 

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

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

221 

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

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

224 it, leading to crashes or subtle failures. 

225 

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

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

228 

229 Raises: 

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

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

232 supported on this platform. 

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

234 RuntimeError: If Python is not 64-bit. 

235 

236 Search order: 

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

238 

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

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

241 

242 1. **NVIDIA Python wheels** 

243 

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

245 shipped in NVIDIA wheels. 

246 

247 2. **Conda environment** 

248 

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

250 defined automatically in activated conda environments (see 

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

252 

253 3. **OS default mechanisms** 

254 

255 - Fall back to the native loader: 

256 

257 - Linux: ``dlopen()`` 

258 

259 - Windows: ``LoadLibraryExW()`` 

260 

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

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

263 

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

265 ``SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS)``. 

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

267 the system ``PATH``. 

268 

269 4. **Environment variables** 

270 

271 - If set, use ``CUDA_PATH`` or ``CUDA_HOME`` (in that order). 

272 On Windows, this is the typical way system-installed CTK DLLs are 

273 located. Note that the NVIDIA CTK installer automatically 

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

275 

276 5. **CTK root canary probe (discoverable libs only)** 

277 

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

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

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

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

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

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

284 

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

286 

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

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

289 libraries the search is simplified to: 

290 

291 0. Already loaded in the current process 

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

293 

294 The CTK-specific steps (site-packages, conda, ``CUDA_PATH``, canary 

295 probe) are skipped entirely. 

296 

297 Notes: 

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

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

300 

301 """ 

302 pointer_size_bits = struct.calcsize("P") * 8 6 ctx1adfQ60

303 if pointer_size_bits != 64: 6 ctx1adfQ60

304 raise RuntimeError( 1 ctx16

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

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

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

308 ) 

309 if libname not in _ALL_KNOWN_LIBNAMES: 5 ctx1adfQ0

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

311 if libname not in _ALL_SUPPORTED_LIBNAMES: 4 ctx1adfQ

312 raise DynamicLibNotAvailableError( 1 ctx1Q

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

314 f"Supported names on {_PLATFORM_NAME}: {sorted(_ALL_SUPPORTED_LIBNAMES)}" 

315 ) 

316 return _load_lib_no_cache(libname) 3 ctx1adf