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

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) 5 ctx1ayzAI

66 if loaded is not None: 5 ctx1ayzAI

67 return loaded 1 ctx1I

68 loaded = LOADER.load_with_system_search(desc) 4 ctx1ayzA

69 if loaded is not None: 4 ctx1ayzA

70 return loaded 3 ctx1ayz

71 raise DynamicLibNotFoundError( 1 ctx1A

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 ctx1jk

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

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

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 ctx1jk

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

91 else: 

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

93 raise ChildProcessError( 2 ctx1jk

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(libname: str) -> str | None: (empty)

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

104 try: 13 ctx1jEFkGqrstunvw

105 result = subprocess.run( # noqa: S603 - trusted argv: current interpreter + internal probe module 13 ctx1jEFkGqrstunvw

106 build_dynamic_lib_subprocess_command(MODE_CANARY, libname), 

107 capture_output=True, 

108 text=True, 

109 timeout=_CANARY_PROBE_TIMEOUT_SECONDS, 

110 check=False, 

111 cwd=DYNAMIC_LIB_SUBPROCESS_CWD, 

112 ) 

113 except subprocess.TimeoutExpired as exc: 1 ctx1k

114 _raise_canary_probe_child_process_error(timeout=exc.timeout, stderr=exc.stderr) 1 ctx1k

115 

116 if result.returncode != 0: 12 ctx1jEFGqrstunvw

117 _raise_canary_probe_child_process_error(returncode=result.returncode, stderr=result.stderr) 1 ctx1j

118 

119 payload: DynamicLibSubprocessPayload = parse_dynamic_lib_subprocess_payload( 11 ctx1EFGqrstunvw

120 result.stdout, 

121 libname=libname, 

122 error_label="Canary probe child process", 

123 ) 

124 abs_path: str | None = payload.abs_path 8 ctx1qrstunvw

125 if payload.status == STATUS_OK: 8 ctx1qrstunvw

126 return abs_path 6 ctx1qstunv

127 return None 3 ctx1rnw

128 

129 

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

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

132 for canary_libname in ctx.desc.ctk_root_canary_anchor_libnames: 7 ctx1obxpBCc

133 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(canary_libname) 7 ctx1obxpBCc

134 if canary_abs_path is None: 7 ctx1obxpBCc

135 continue 3 ctx1BCc

136 ctk_root = derive_ctk_root(canary_abs_path) 4 ctx1obxp

137 if ctk_root is None: 4 ctx1obxp

138 continue 1 ctx1x

139 find = find_via_ctk_root(ctx, ctk_root) 3 ctx1obp

140 if find is not None: 3 ctx1obp

141 return str(find.abs_path) 2 ctx1ob

142 return None 5 ctx1xpBCc

143 

144 

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

146 desc = LIB_DESCRIPTORS[libname] 13 ctx1abhgfJKilcmde

147 

148 if libname in _DRIVER_ONLY_LIBNAMES: 13 ctx1abhgfJKilcmde

149 return _load_driver_lib_no_cache(desc) 3 ctx1aJK

150 

151 ctx = SearchContext(desc) 11 ctx1abhgfilcmde

152 

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

154 find = run_find_steps(ctx, EARLY_FIND_STEPS) 11 ctx1abhgfilcmde

155 

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

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

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

159 loaded = LOADER.check_if_already_loaded_from_elsewhere(desc, find is not None) 11 ctx1abhgfilcmde

160 load_dependencies(desc, load_nvidia_dynamic_lib) 11 ctx1abhgfilcmde

161 if loaded is not None: 11 ctx1abhgfilcmde

162 return loaded 

163 

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

165 if find is not None: 11 ctx1abhgfilcmde

166 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 2 ctx1lm

167 

168 loaded = LOADER.load_with_system_search(desc) 9 ctx1abhgficde

169 if loaded is not None: 9 ctx1abhgficde

170 return loaded 3 ctx1afi

171 

172 find = run_find_steps(ctx, LATE_FIND_STEPS) 7 ctx1abhgcde

173 if find is not None: 7 ctx1abhgcde

174 return LOADER.load_with_abs_path(desc, find.abs_path, find.found_via) 4 ctx1ahde

175 

176 if desc.ctk_root_canary_anchor_libnames: 3 ctx1bgc

177 canary_abs_path = _try_ctk_root_canary(ctx) 2 ctx1bc

178 if canary_abs_path is not None: 2 ctx1bc

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

180 

181 ctx.raise_not_found() 2 ctx1gc

182 

183 

184@functools.cache (empty)

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

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

187 

188 Args: 

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

190 ``"nvvm"``, etc.). 

191 

192 Returns: 

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

194 

195 **Important:** 

196 

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

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

199 

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

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

202 it, leading to crashes or subtle failures. 

203 

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

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

206 

207 Raises: 

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

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

210 supported on this platform. 

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

212 RuntimeError: If Python is not 64-bit. 

213 

214 Search order: 

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

216 

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

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

219 

220 1. **NVIDIA Python wheels** 

221 

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

223 shipped in NVIDIA wheels. 

224 

225 2. **Conda environment** 

226 

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

228 defined automatically in activated conda environments (see 

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

230 

231 3. **OS default mechanisms** 

232 

233 - Fall back to the native loader: 

234 

235 - Linux: ``dlopen()`` 

236 

237 - Windows: ``LoadLibraryExW()`` 

238 

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

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

241 

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

243 ``SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_DEFAULT_DIRS)``. 

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

245 the system ``PATH``. 

246 

247 4. **Environment variables** 

248 

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

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

251 located. Note that the NVIDIA CTK installer automatically 

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

253 

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

255 

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

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

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

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

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

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

262 

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

264 

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

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

267 libraries the search is simplified to: 

268 

269 0. Already loaded in the current process 

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

271 

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

273 probe) are skipped entirely. 

274 

275 Notes: 

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

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

278 

279 """ 

280 pointer_size_bits = struct.calcsize("P") * 8 7 ctx1afDLHde

281 if pointer_size_bits != 64: 7 ctx1afDLHde

282 raise RuntimeError( 1 ctx1L

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

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

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

286 ) 

287 if libname not in _ALL_KNOWN_LIBNAMES: 6 ctx1afDHde

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

289 if libname not in _ALL_SUPPORTED_LIBNAMES: 5 ctx1afDde

290 raise DynamicLibNotAvailableError( 1 ctx1D

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

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

293 ) 

294 return _load_lib_no_cache(libname) 4 ctx1afde