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_home_or_path (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 25 ctx1acbfspqohOidljkmnBCD457PQ

58 

59 @property (empty)

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

61 return cast(str, self.platform.lib_searched_for(self.libname)) 24 ctx1acbfspqohOidljkmnBCD45PQ

62 

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

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

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

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

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) 19 ctx1acbefpqthidTRSljkmn

76 for rel_path in rel_dirs: 19 ctx1acbefpqthidTRSljkmn

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

78 if os.path.isdir(dirname): 16 ctx1acbfpqhidRSljkmn

79 return os.path.normpath(dirname) 16 ctx1acbfpqhidRSljkmn

80 return None 3 ctx1etT

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: 16 ctx1acbefpqthidljkmn

86 return None 2 ctx1et

87 return cast( 14 ctx1acbfpqhidljkmn

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 ctx1cbreEFGHIUuvwxyzAg

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

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

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

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

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

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

115 return parent 14 ctx1cbeEFGuvwxyzAg

116 return None 6 ctx1cbreUg

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 ctx1cbreLJMKNg

127 

128 lib_dir = ntpath.dirname(resolved_lib_path) 10 ctx1cbreLJMKNg

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

130 if basename == "x64": 10 ctx1cbreLJMKNg

131 parent = ntpath.dirname(lib_dir) 3 ctx1JKg

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

133 return ntpath.dirname(parent) 3 ctx1JKg

134 elif basename == "bin": 7 ctx1cbreLMN

135 return ntpath.dirname(lib_dir) 5 ctx1cbeLM

136 return None 2 ctx1rN

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 ctx1cbreYZuvwxyzAg

142 if ctk_root is not None: 14 ctx1cbreYZuvwxyzAg

143 return ctk_root 12 ctx1cbeYuvwxyzAg

144 return _derive_ctk_root_windows(resolved_lib_path) 6 ctx1cbreZg

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 ctx1cbepqt

150 abs_path = _find_using_lib_dir(ctx, lib_dir) 6 ctx1cbepqt

151 if abs_path is None: 6 ctx1cbepqt

152 return None 2 ctx1et

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

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

164 if not rel_dirs: 7 ctx1aodBCD0

165 return None 1 ctx10

166 abs_path = ctx.platform.find_in_site_packages(rel_dirs, ctx.lib_searched_for, ctx.error_messages, ctx.attachments) 6 ctx1aodBCD

167 if abs_path is not None: 6 ctx1aodBCD

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

169 return None 4 ctx1aodD

170 

171 

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

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

174 conda_prefix = os.environ.get("CONDA_PREFIX") 9 ctx1aohidjk12

175 if not conda_prefix: 9 ctx1aohidjk12

176 return None 5 ctx1aod12

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

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

179 abs_path = _find_using_lib_dir(ctx, lib_dir) 4 ctx1hijk

180 if abs_path is not None: 4 ctx1hijk

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

182 return None 

183 

184 

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

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

187 cuda_home = get_cuda_home_or_path() 9 ctx1abfsdlmn3

188 if cuda_home is None: 9 ctx1abfsdlmn3

189 return None 3 ctx1bs3

190 lib_dir = _find_lib_dir_using_anchor(ctx.desc, ctx.platform, cuda_home) 6 ctx1afdlmn

191 abs_path = _find_using_lib_dir(ctx, lib_dir) 6 ctx1afdlmn

192 if abs_path is not None: 6 ctx1afdlmn

193 return FindResult(abs_path, "CUDA_HOME") 6 ctx1afdlmn

194 return None 

195 

196 

197# --------------------------------------------------------------------------- 

198# Step sequences per strategy 

199# --------------------------------------------------------------------------- 

200 

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

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

203 

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

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

206 

207 

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

209# Cascade runner 

210# --------------------------------------------------------------------------- 

211 

212 

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

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

215 for step in steps: 10 ctx1abfsod6VWX

216 result = step(ctx) 9 ctx1abfsodVWX

217 if result is not None: 9 ctx1abfsodVWX

218 return result 5 ctx1afdVX

219 return None 7 ctx1absod6W