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

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 (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)

32 

33# --------------------------------------------------------------------------- 

34# Data types 

35# --------------------------------------------------------------------------- 

36 

37 

38@dataclass (empty)

39class FindResult: (empty)

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

41 

42 abs_path: str 

43 found_via: str 

44 

45 

46@dataclass (empty)

47class SearchContext: (empty)

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

49 

50 desc: LibDescriptor 

51 platform: SearchPlatform = PLATFORM (empty)

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

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

54 

55 @property (empty)

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

57 return self.desc.name # type: ignore[no-any-return] # mypy can't resolve new sibling module 26 ctx1acbgtqrpiPjdemklnoCDE568QR

58 

59 @property (empty)

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

61 return cast(str, self.platform.lib_searched_for(self.libname)) 25 ctx1acbgtqrpiPjdemklnoCDE56QR

62 

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

64 err = ", ".join(self.error_messages) 4 ctx1tPQR

65 att = "\n".join(self.attachments) 4 ctx1tPQR

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

67 

68 

69#: Type alias for a find step callable. 

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

71 

72 

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

74 """Find the library directory under *anchor_point* using the descriptor's relative paths.""" 

75 rel_dirs = platform.anchor_rel_dirs(desc) 20 ctx1acbfgqruijdeUSTmklno

76 for rel_path in rel_dirs: 20 ctx1acbfgqruijdeUSTmklno

77 for dirname in sorted(glob.glob(os.path.join(anchor_point, rel_path))): 20 ctx1acbfgqruijdeUSTmklno

78 if os.path.isdir(dirname): 17 ctx1acbgqrijdeSTmklno

79 return os.path.normpath(dirname) 17 ctx1acbgqrijdeSTmklno

80 return None 3 ctx1fuU

81 

82 

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

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

85 if lib_dir is None: 17 ctx1acbfgqruijdemklno

86 return None 2 ctx1fu

87 return cast( 15 ctx1acbgqrijdemklno

88 str | None, 

89 ctx.platform.find_in_lib_dir( 

90 lib_dir, 

91 ctx.libname, 

92 ctx.lib_searched_for, 

93 ctx.error_messages, 

94 ctx.attachments, 

95 ), 

96 ) 

97 

98 

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

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

101 

102 Supports: 

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

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

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

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

107 """ 

108 lib_dir = os.path.dirname(resolved_lib_path) 18 ctx1cbsfFGHIJVvwxyzABh

109 basename = os.path.basename(lib_dir) 18 ctx1cbsfFGHIJVvwxyzABh

110 if basename in ("lib64", "lib"): 18 ctx1cbsfFGHIJVvwxyzABh

111 parent = os.path.dirname(lib_dir) 16 ctx1cbfFGHIJvwxyzABh

112 grandparent = os.path.dirname(parent) 16 ctx1cbfFGHIJvwxyzABh

113 if os.path.basename(grandparent) == "targets": 16 ctx1cbfFGHIJvwxyzABh

114 return os.path.dirname(grandparent) 2 ctx1IJ

115 return parent 14 ctx1cbfFGHvwxyzABh

116 return None 6 ctx1cbsfVh

117 

118 

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

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

121 

122 Supports: 

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

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

125 """ 

126 import ntpath 10 ctx1cbsfMKNLOh

127 

128 lib_dir = ntpath.dirname(resolved_lib_path) 10 ctx1cbsfMKNLOh

129 basename = ntpath.basename(lib_dir).lower() 10 ctx1cbsfMKNLOh

130 if basename == "x64": 10 ctx1cbsfMKNLOh

131 parent = ntpath.dirname(lib_dir) 3 ctx1KLh

132 if ntpath.basename(parent).lower() == "bin": 3 ctx1KLh

133 return ntpath.dirname(parent) 3 ctx1KLh

134 elif basename == "bin": 7 ctx1cbsfMNO

135 return ntpath.dirname(lib_dir) 5 ctx1cbfMN

136 return None 2 ctx1sO

137 

138 

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

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

141 ctk_root = _derive_ctk_root_linux(resolved_lib_path) 14 ctx1cbsfZ0vwxyzABh

142 if ctk_root is not None: 14 ctx1cbsfZ0vwxyzABh

143 return ctk_root 12 ctx1cbfZvwxyzABh

144 return _derive_ctk_root_windows(resolved_lib_path) 6 ctx1cbsf0h

145 

146 

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

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

149 lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, ctk_root) 6 ctx1cbfqru

150 abs_path = _find_using_lib_dir(ctx, lib_dir) 6 ctx1cbfqru

151 if abs_path is None: 6 ctx1cbfqru

152 return None 2 ctx1fu

153 return FindResult(abs_path, "system-ctk-root") 4 ctx1cbqr

154 

155 

156# --------------------------------------------------------------------------- 

157# Find steps 

158# --------------------------------------------------------------------------- 

159 

160 

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

162 """Search pip wheel install locations.""" 

163 rel_dirs = ctx.platform.site_packages_rel_dirs(ctx.desc) 8 ctx1apdeCDE1

164 if not rel_dirs: 8 ctx1apdeCDE1

165 return None 1 ctx11

166 abs_path = ctx.platform.find_in_site_packages(rel_dirs, ctx.lib_searched_for, ctx.error_messages, ctx.attachments) 7 ctx1apdeCDE

167 if abs_path is not None: 7 ctx1apdeCDE

168 return FindResult(abs_path, "site-packages") 2 ctx1CD

169 return None 5 ctx1apdeE

170 

171 

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

173 """Search ``$CONDA_PREFIX``.""" 

174 conda_prefix = os.environ.get("CONDA_PREFIX") 10 ctx1apijdekl23

175 if not conda_prefix: 10 ctx1apijdekl23

176 return None 6 ctx1apde23

177 anchor = ctx.platform.conda_anchor_point(conda_prefix) 4 ctx1ijkl

178 lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, anchor) 4 ctx1ijkl

179 abs_path = _find_using_lib_dir(ctx, lib_dir) 4 ctx1ijkl

180 if abs_path is not None: 4 ctx1ijkl

181 return FindResult(abs_path, "conda") 4 ctx1ijkl

182 return None 

183 

184 

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

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

187 

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

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

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

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

192 

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

194 environment variable actually provided the value. 

195 """ 

196 cuda_home = get_cuda_path_or_home() 10 ctx1abgtdemno4

197 if cuda_home is None: 10 ctx1abgtdemno4

198 return None 3 ctx1bt4

199 lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, cuda_home) 7 ctx1agdemno

200 abs_path = _find_using_lib_dir(ctx, lib_dir) 7 ctx1agdemno

201 if abs_path is not None: 7 ctx1agdemno

202 return FindResult(abs_path, "CUDA_PATH") 7 ctx1agdemno

203 return None 

204 

205 

206# --------------------------------------------------------------------------- 

207# Step sequences per strategy 

208# --------------------------------------------------------------------------- 

209 

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

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

212 

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

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

215 

216 

217# --------------------------------------------------------------------------- 

218# Cascade runner 

219# --------------------------------------------------------------------------- 

220 

221 

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

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

224 for step in steps: 11 ctx1abgtpde7WXY

225 result = step(ctx) 10 ctx1abgtpdeWXY

226 if result is not None: 10 ctx1abgtpdeWXY

227 return result 6 ctx1agdeWY

228 return None 8 ctx1abtpde7X