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

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 ctx1azABJf

66 if loaded is not None: 6 ctx1azABJf

67 return loaded 1 ctx1J

68 loaded = LOADER.load_with_system_search(desc) 5 ctx1azABf

69 if loaded is not None: 5 ctx1azABf

70 return loaded 4 ctx1azAf

71 raise DynamicLibNotFoundError( 1 ctx1B

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 ctx1ij

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

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

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 ctx1ij

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

91 else: 

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

93 raise ChildProcessError( 2 ctx1ij

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: 15 ctx1iFGjHqrmstnuvwx

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

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 ctx1j

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

119 

120 if result.returncode != 0: 14 ctx1iFGHqrmstnuvwx

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

122 

123 payload: DynamicLibSubprocessPayload = parse_dynamic_lib_subprocess_payload( 13 ctx1FGHqrmstnuvwx

124 result.stdout, 

125 libname=libname, 

126 error_label="Canary probe child process", 

127 ) 

128 abs_path: str | None = payload.abs_path 10 ctx1qrmstnuvwx

129 if payload.status == STATUS_OK: 10 ctx1qrmstnuvwx

130 return abs_path 5 ctx1qtnvx

131 return None 6 ctx1rmsnuw

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 ctx1mOP

137 

138 

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

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

141 for canary_libname in ctx.desc.ctk_root_canary_anchor_libnames: 7 ctx1obypCDc

142 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname) 7 ctx1obypCDc

143 if canary_abs_path is None: 7 ctx1obypCDc

144 continue 3 ctx1CDc

145 ctk_root = derive_ctk_root(canary_abs_path) 4 ctx1obyp

146 if ctk_root is None: 4 ctx1obyp

147 continue 1 ctx1y

148 find = find_via_ctk_root(ctx, ctk_root) 3 ctx1obp

149 if find is not None: 3 ctx1obp

150 return str(find.abs_path) 2 ctx1ob

151 return None 5 ctx1ypCDc

152 

153 

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

155 desc = LIB_DESCRIPTORS[libname] 13 ctx1abgedKLMhfkcl

156 

157 if libname in _DRIVER_ONLY_LIBNAMES: 13 ctx1abgedKLMhfkcl

158 return _load_driver_lib_no_cache(desc) 5 ctx1aKLMf

159 

160 ctx = SearchContext(desc) 9 ctx1abgedhkcl

161 

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

163 find = run_find_steps(ctx, EARLY_FIND_STEPS) 9 ctx1abgedhkcl

164 

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

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

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

168 loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, find is not None) 9 ctx1abgedhkcl

169 load_dependencies(desc, load_nvidia_dynamic_lib) 9 ctx1abgedhkcl

170 if loaded is not None: 9 ctx1abgedhkcl

171 return loaded 

172 

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

174 if find is not None: 9 ctx1abgedhkcl

175 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 2 ctx1kl

176 

177 loaded = LOADER.load_with_system_search(desc) 7 ctx1abgedhc

178 if loaded is not None: 7 ctx1abgedhc

179 return loaded 3 ctx1adh

180 

181 find = run_find_steps(ctx, LATE_FIND_STEPS) 5 ctx1abgec

182 if find is not None: 5 ctx1abgec

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

184 

185 if desc.ctk_root_canary_anchor_libnames: 3 ctx1bec

186 canary_abs_path = _try_ctk_root_canary(ctx) 2 ctx1bc

187 if canary_abs_path is not None: 2 ctx1bc

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

189 

190 ctx.raise_not_found() 2 ctx1ec

191 

192 

193@functools.cache (empty)

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

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

196 

197 Args: 

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

199 ``"nvvm"``, etc.). 

200 

201 Returns: 

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

203 

204 **Important:** 

205 

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

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

208 

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

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

211 it, leading to crashes or subtle failures. 

212 

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

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

215 

216 Raises: 

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

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

219 supported on this platform. 

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

221 RuntimeError: If Python is not 64-bit. 

222 

223 Search order: 

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

225 

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

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

228 

229 1. **NVIDIA Python wheels** 

230 

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

232 shipped in NVIDIA wheels. 

233 

234 2. **Conda environment** 

235 

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

237 defined automatically in activated conda environments (see 

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

239 

240 3. **OS default mechanisms** 

241 

242 - Fall back to the native loader: 

243 

244 - Linux: ``dlopen()`` 

245 

246 - Windows: ``LoadLibraryExW()`` 

247 

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

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

250 

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

252 ``SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS)``. 

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

254 the system ``PATH``. 

255 

256 4. **Environment variables** 

257 

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

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

260 located. Note that the NVIDIA CTK installer automatically 

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

262 

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

264 

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

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

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

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

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

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

271 

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

273 

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

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

276 libraries the search is simplified to: 

277 

278 0. Already loaded in the current process 

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

280 

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

282 probe) are skipped entirely. 

283 

284 Notes: 

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

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

287 

288 """ 

289 pointer_size_bits = struct.calcsize("P") * 8 6 ctx1adfENI

290 if pointer_size_bits != 64: 6 ctx1adfENI

291 raise RuntimeError( 1 ctx1N

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

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

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

295 ) 

296 if libname not in _ALL_KNOWN_LIBNAMES: 5 ctx1adfEI

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

298 if libname not in _ALL_SUPPORTED_LIBNAMES: 4 ctx1adfE

299 raise DynamicLibNotAvailableError( 1 ctx1E

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

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

302 ) 

303 return _load_lib_no_cache(libname) 3 ctx1adf