Coverage for cuda/pathfinder/_dynamic_libs/search_steps.py: 97.56%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4"""Composable search steps for locating NVIDIA libraries. 

5 

6Each find step is a callable with signature:: 

7 

8 (SearchContext) -> FindResult | None 

9 

10Find steps locate a library file on disk without loading it. The 

11orchestrator in :mod:`load_nvidia_dynamic_lib` handles loading, the 

12already-loaded check, and dependency resolution. 

13 

14Step sequences are defined per search strategy so that adding a new 

15step or strategy only requires adding a function and a tuple entry. 

16 

17This module is intentionally platform-agnostic: it does not branch on the 

18current operating system. Platform differences are routed through the 

19:data:`~cuda.pathfinder._dynamic_libs.search_platform.PLATFORM` instance. 

20""" 

21 

22import glob (empty)

23import os (empty)

24from collections.abc import Callable, Iterator (empty)

25from dataclasses import dataclass, field (empty)

26from typing import NoReturn, cast (empty)

27 

28from cuda.pathfinder._dynamic_libs.lib_descriptor import LibDescriptor (empty)

29from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError (empty)

30from cuda.pathfinder._dynamic_libs.search_platform import PLATFORM, SearchPlatform (empty)

31from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home (empty)

32from cuda.pathfinder._utils.path_sort import numeric_aware_path_sort_key (empty)

33 

34# --------------------------------------------------------------------------- 

35# Data types 

36# --------------------------------------------------------------------------- 

37 

38 

39@dataclass (empty)

40class FindResult: (empty)

41 """A library file located on disk (not yet loaded).""" 

42 

43 abs_path: str 

44 found_via: str 

45 

46 

47@dataclass (empty)

48class SearchContext: (empty)

49 """Mutable state accumulated during the search cascade.""" 

50 

51 desc: LibDescriptor 

52 platform: SearchPlatform = PLATFORM (empty)

53 error_messages: list[str] = field(default_factory=list) (empty)

54 attachments: list[str] = field(default_factory=list) (empty)

55 

56 @property (empty)

57 def libname(self) -> str: (empty)

58 return self.desc.name # type: ignore[no-any-return] # mypy can't resolve new sibling module 1 ctx2Eb

59 

60 @property (empty)

61 def lib_searched_for(self) -> str: (empty)

62 return cast(str, self.platform.lib_searched_for(self.desc)) 6 ctx2u pbFbGbqbrb

63 

64 def raise_not_found(self) -> NoReturn: (empty)

65 err = ", ".join(self.error_messages) 4 ctx2u pbqbrb

66 att = "\n".join(self.attachments) 4 ctx2u pbqbrb

67 raise DynamicLibNotFoundError(f'Failure finding "{self.lib_searched_for}": {err}\n{att}') 4 ctx2u pbqbrb

68 

69 

70#: Type alias for a find step callable. 

71FindStep = Callable[[SearchContext], FindResult | None] (empty)

72 

73 

74def _iter_lib_dirs(root: str, rel_dirs: tuple[str, ...]) -> Iterator[str]: (empty)

75 """Yield existing library directories under *root* in descriptor order.""" 

76 for rel_path in rel_dirs: 47 ctx2a c b e f Q R ] S d v w sb^ _ ` F x y z A B C D E G H I J K L M i j k l m n N U V o p q r s t

77 for dirname in sorted(glob.glob(os.path.join(root, rel_path))): 46 ctx2a c b e f Q R ] S d v w sb^ _ ` F x y z A B C D E G H I J K L M i j k l m n U V o p q r s t

78 if os.path.isdir(dirname): 41 ctx1acbfQRSdvw^_`FxyzABCDEGHIJKLMijklmnopqrst

79 yield os.path.normpath(dirname) 41 ctx1acbfQRSdvw^_`FxyzABCDEGHIJKLMijklmnopqrst

80 

81 

82def _iter_lib_dirs_using_anchor( (empty)

83 desc: LibDescriptor, 

84 platform: SearchPlatform, 

85 anchor_point: str, 

86) -> Iterator[str]: 

87 """Yield existing library directories under *anchor_point* in descriptor order.""" 

88 yield from _iter_lib_dirs(anchor_point, platform.anchor_rel_dirs(desc)) 4 ctx2sb^ _ `

89 

90 

91def _find_lib_dir_using_anchor(desc: LibDescriptor, platform: SearchPlatform, anchor_point: str) -> str | None: (empty)

92 """Find the first library directory under *anchor_point*.""" 

93 return next(_iter_lib_dirs_using_anchor(desc, platform, anchor_point), None) 4 ctx2sb^ _ `

94 

95 

96def _find_using_lib_dir(ctx: SearchContext, lib_dir: str | None) -> str | None: (empty)

97 """Find a library file in a resolved lib directory.""" 

98 if lib_dir is None: 38 ctx1acbfQRSdvwFxyzABCDEGHIJKLMijklmnopqrst

99 return None 

100 return cast( 38 ctx1acbfQRSdvwFxyzABCDEGHIJKLMijklmnopqrst

101 str | None, 

102 ctx.platform.find_in_lib_dir( 

103 lib_dir, 

104 ctx.desc, 

105 ctx.error_messages, 

106 ctx.attachments, 

107 ), 

108 ) 

109 

110 

111def _find_under_root( (empty)

112 ctx: SearchContext, 

113 root: str, 

114 rel_dirs: tuple[str, ...], 

115 found_via: str, 

116) -> FindResult | None: 

117 """Resolve *rel_dirs* under *root*, then find the requested library.""" 

118 for lib_dir in _iter_lib_dirs(root, rel_dirs): 43 ctx1acbefQR]SdvwFxyzABCDEGHIJKLMijklmnNUVopqrst

119 abs_path = _find_using_lib_dir(ctx, lib_dir) 38 ctx1acbfQRSdvwFxyzABCDEGHIJKLMijklmnopqrst

120 if abs_path is not None: 38 ctx1acbfQRSdvwFxyzABCDEGHIJKLMijklmnopqrst

121 return FindResult(abs_path, found_via) 31 ctx1acbfQRSdvwFxyzAGHIKLijklmnpqrst

122 return None 12 ctx1e]BCDEJMNUVo

123 

124 

125def _find_under_anchor_root(ctx: SearchContext, root: str, found_via: str) -> FindResult | None: (empty)

126 """Resolve the descriptor's general anchors under *root*.""" 

127 return _find_under_root(ctx, root, ctx.platform.anchor_rel_dirs(ctx.desc), found_via) 30 ctx1acbefQR]SvwFxyzABCDEGHIJKLMNst

128 

129 

130def _derive_ctk_root_linux(resolved_lib_path: str) -> str | None: (empty)

131 """Derive CTK root from Linux canary path. 

132 

133 Supports: 

134 - ``$CTK_ROOT/lib64/libfoo.so.*`` 

135 - ``$CTK_ROOT/lib/libfoo.so.*`` 

136 - ``$CTK_ROOT/targets/<triple>/lib64/libfoo.so.*`` 

137 - ``$CTK_ROOT/targets/<triple>/lib/libfoo.so.*`` 

138 """ 

139 lib_dir = os.path.dirname(resolved_lib_path) 38 ctx2c b O e . / : ; = tbP h W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - g

140 basename = os.path.basename(lib_dir) 38 ctx2c b O e . / : ; = tbP h W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - g

141 if basename in ("lib64", "lib"): 38 ctx2c b O e . / : ; = tbP h W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - g

142 parent = os.path.dirname(lib_dir) 35 ctx1cbe./:;=hWXYZ0123456789!#$%'()*+,-g

143 grandparent = os.path.dirname(parent) 35 ctx1cbe./:;=hWXYZ0123456789!#$%'()*+,-g

144 if os.path.basename(grandparent) == "targets": 35 ctx1cbe./:;=hWXYZ0123456789!#$%'()*+,-g

145 return os.path.dirname(grandparent) 2 ctx1;=

146 return parent 33 ctx1cbe./:hWXYZ0123456789!#$%'()*+,-g

147 return None 8 ctx2c b O e tbP h g

148 

149 

150def _derive_ctk_root_windows(resolved_lib_path: str) -> str | None: (empty)

151 """Derive CTK root from Windows canary path. 

152 

153 Supports: 

154 - ``$CTK_ROOT/bin/x64/foo.dll`` (CTK 13 style) 

155 - ``$CTK_ROOT/bin/arm64/foo.dll`` (Windows on Arm CTK 13 style) 

156 - ``$CTK_ROOT/bin/foo.dll`` (CTK 12 style) 

157 """ 

158 import ntpath 13 ctx1cbOe{?|@[}Phg

159 

160 lib_dir = ntpath.dirname(resolved_lib_path) 13 ctx1cbOe{?|@[}Phg

161 basename = ntpath.basename(lib_dir).lower() 13 ctx1cbOe{?|@[}Phg

162 if basename in ("x64", "arm64"): 13 ctx1cbOe{?|@[}Phg

163 parent = ntpath.dirname(lib_dir) 4 ctx1?@[g

164 if ntpath.basename(parent).lower() == "bin": 4 ctx1?@[g

165 return ntpath.dirname(parent) 4 ctx1?@[g

166 elif basename == "bin": 9 ctx1cbOe{|}Ph

167 return ntpath.dirname(lib_dir) 6 ctx1cbe{|h

168 return None 3 ctx1O}P

169 

170 

171def derive_ctk_root(resolved_lib_path: str) -> str | None: (empty)

172 """Derive CTK root from a resolved canary library path.""" 

173 ctk_root = _derive_ctk_root_linux(resolved_lib_path) 34 ctx2c b O e xbybP h W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - g

174 if ctk_root is not None: 34 ctx2c b O e xbybP h W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - g

175 return ctk_root 31 ctx2c b e xbh W X Y Z 0 1 2 3 4 5 6 7 8 9 ! # $ % ' ( ) * + , - g

176 return _derive_ctk_root_windows(resolved_lib_path) 8 ctx2c b O e ybP h g

177 

178 

179def find_via_ctk_root(ctx: SearchContext, ctk_root: str) -> FindResult | None: (empty)

180 """Find a library under a previously derived CTK root.""" 

181 return _find_under_anchor_root(ctx, ctk_root, "system-ctk-root") 7 ctx1cbeQR]S

182 

183 

184# --------------------------------------------------------------------------- 

185# Find steps 

186# --------------------------------------------------------------------------- 

187 

188 

189def find_in_site_packages(ctx: SearchContext) -> FindResult | None: (empty)

190 """Search pip wheel install locations.""" 

191 rel_dirs = ctx.platform.site_packages_rel_dirs(ctx.desc) 20 ctx2a T d ~ abbbcbdbebfbgbhbibjbkblbmbzbnbob

192 if not rel_dirs: 20 ctx2a T d ~ abbbcbdbebfbgbhbibjbkblbmbzbnbob

193 return None 2 ctx2d zb

194 abs_path = ctx.platform.find_in_site_packages( 18 ctx2a T ~ abbbcbdbebfbgbhbibjbkblbmbnbob

195 rel_dirs, 

196 ctx.desc, 

197 ctx.error_messages, 

198 ctx.attachments, 

199 ) 

200 if abs_path is not None: 18 ctx2a T ~ abbbcbdbebfbgbhbibjbkblbmbnbob

201 return FindResult(abs_path, "site-packages") 8 ctx2~ abcbdbebfbgbmb

202 return None 10 ctx2a T bbhbibjbkblbnbob

203 

204 

205def find_in_conda(ctx: SearchContext) -> FindResult | None: (empty)

206 """Search ``$CONDA_PREFIX``.""" 

207 conda_prefix = os.environ.get("CONDA_PREFIX") 16 ctx2a T d v w x y z A B C D E AbBbN

208 if not conda_prefix: 16 ctx2a T d v w x y z A B C D E AbBbN

209 return None 5 ctx2a T d AbBb

210 anchor = ctx.platform.conda_anchor_point(conda_prefix) 11 ctx1vwxyzABCDEN

211 return _find_under_anchor_root(ctx, anchor, "conda") 11 ctx1vwxyzABCDEN

212 

213 

214def find_in_install_root_env_vars(ctx: SearchContext) -> FindResult | None: (empty)

215 """Search installation roots named by descriptor-specific environment variables.""" 

216 rel_dirs = ctx.platform.install_root_env_rel_dirs(ctx.desc) 17 ctx1abfudijklmnUVopqr

217 for env_var in ctx.platform.install_root_env_vars(ctx.desc): 17 ctx1abfudijklmnUVopqr

218 root = os.environ.get(env_var) 13 ctx1dijklmnUVopqr

219 if not root: 13 ctx1dijklmnUVopqr

220 continue 

221 result = _find_under_root(ctx, root, rel_dirs, env_var) 13 ctx1dijklmnUVopqr

222 if result is not None: 13 ctx1dijklmnUVopqr

223 return result 10 ctx1dijklmnpqr

224 return None 7 ctx1abfuUVo

225 

226 

227def find_in_cuda_path(ctx: SearchContext) -> FindResult | None: (empty)

228 """Search ``$CUDA_PATH`` / ``$CUDA_HOME``. 

229 

230 On Windows, this is the normal fallback for system-installed CTK DLLs when 

231 they are not already discoverable via the native ``LoadLibraryExW(..., 0)`` 

232 path used by :func:`cuda.pathfinder._dynamic_libs.load_dl_windows.load_with_system_search`. 

233 Python 3.8+ does not include ``PATH`` in that native DLL search. 

234 

235 The returned ``found_via`` is always ``"CUDA_PATH"`` regardless of which 

236 environment variable actually provided the value. 

237 """ 

238 cuda_home = get_cuda_path_or_home() 14 ctx2a b f u F G H I J K L M CbN

239 if cuda_home is None: 14 ctx2a b f u F G H I J K L M CbN

240 return None 3 ctx2b u Cb

241 return _find_under_anchor_root(ctx, cuda_home, "CUDA_PATH") 11 ctx1afFGHIJKLMN

242 

243 

244def find_in_program_files_roots(ctx: SearchContext) -> FindResult | None: (empty)

245 """Search descriptor-configured installation roots under Program Files.""" 

246 for root_glob in ctx.platform.program_files_root_globs(ctx.desc): 5 ctx1buNst

247 for root in sorted(glob.glob(root_glob), key=numeric_aware_path_sort_key, reverse=True): 2 ctx1st

248 if not os.path.isdir(root): 2 ctx1st

249 continue 

250 result = _find_under_anchor_root(ctx, os.path.normpath(root), "ProgramFiles") 2 ctx1st

251 if result is not None: 2 ctx1st

252 return result 2 ctx1st

253 return None 3 ctx1buN

254 

255 

256# --------------------------------------------------------------------------- 

257# Step sequences per strategy 

258# --------------------------------------------------------------------------- 

259 

260#: Find steps that run before the already-loaded check and system search. 

261EARLY_FIND_STEPS: tuple[FindStep, ...] = (find_in_site_packages, find_in_conda) (empty)

262 

263#: Find steps that run after system search fails. 

264LATE_FIND_STEPS: tuple[FindStep, ...] = ( (empty)

265 find_in_install_root_env_vars, 

266 find_in_cuda_path, 

267 find_in_program_files_roots, 

268) 

269 

270 

271# --------------------------------------------------------------------------- 

272# Cascade runner 

273# --------------------------------------------------------------------------- 

274 

275 

276def run_find_steps(ctx: SearchContext, steps: tuple[FindStep, ...]) -> FindResult | None: (empty)

277 """Run find steps in order, returning the first hit.""" 

278 for step in steps: 10 ctx2a b f u T d Dbubvbwb

279 result = step(ctx) 9 ctx2a b f u T d ubvbwb

280 if result is not None: 9 ctx2a b f u T d ubvbwb

281 return result 5 ctx2a f d ubwb

282 return None 7 ctx2a b u T d Dbvb