Coverage for cuda/pathfinder/_binaries/find_nvidia_binary_utility.py: 97.03%

101 statements  

« prev     ^ index     » next       coverage.py v7.16.0, created at 2026-09-03 02:41 +0000

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4import functools 

5import os 

6from collections.abc import Iterable 

7 

8from cuda.pathfinder._binaries import supported_nvidia_binaries, windows_nsight 

9from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES 

10from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home 

11from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages 

12from cuda.pathfinder._utils.platform_aware import IS_WINDOWS 

13 

14 

15class UnsupportedBinaryError(Exception): 

16 def __init__(self, utility: str) -> None: 

17 super().__init__(utility) 12

18 self.utility = utility 12

19 

20 def __str__(self) -> str: 

21 supported_utilities = ", ".join(supported_nvidia_binaries.SUPPORTED_BINARIES) 12

22 return f"Binary '{self.utility}' is not supported. Supported utilities are: {supported_utilities}" 12

23 

24 

25def _normalize_utility_name(utility_name: str) -> str: 

26 """Normalize utility name by adding .exe on Windows if needed.""" 

27 if IS_WINDOWS and not utility_name.lower().endswith((".exe", ".bat", ".cmd")): 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

28 return f"{utility_name}.exe" 1cLdajklebmnifogpqhsMtuvwzABCyOUWEFGHIJKPQ

29 return utility_name 1cVRDNTrdajklebmnifogpqhx

30 

31 

32def _is_executable_candidate(path: str) -> bool: 

33 if not os.path.isfile(path): 1XYZ0cdajklebmnifogpqh

34 return False 1Xcdajklebmnifogpqh

35 if IS_WINDOWS: 1XYZ0cdefgh

36 return True 1XYZcdefgh

37 return os.access(path, os.X_OK) 1XYZ0cdefgh

38 

39 

40def _ctk_bin_subdirs(root: str) -> list[str]: 

41 if IS_WINDOWS: 1cLDrdajklebmnifogpqhsxJK

42 return [ 1cLdjklemnfogpqhsJK

43 os.path.join(root, "bin", "x64"), 

44 os.path.join(root, "bin", "x86_64"), 

45 os.path.join(root, "bin"), 

46 ] 

47 return [os.path.join(root, "bin")] 1cDrdajklebmnifogpqhx

48 

49 

50def _resolve_candidate_paths(candidates: Iterable[str]) -> str | None: 

51 """Return the first executable candidate, preserving candidate order.""" 

52 seen: set[str] = set() 1abiMtuvwzABCyUWEFGHI

53 for candidate in candidates: 1abiMtuvwzABCyUWEFGHI

54 if candidate in seen: 1aMtuvwzABCyEFGHI

55 continue 

56 seen.add(candidate) 1aMtuvwzABCyEFGHI

57 if _is_executable_candidate(candidate): 1aMtuvwzABCyEFGHI

58 return os.path.abspath(candidate) 1MtuvwzABCyEFGHI

59 return None 1abituvwyUW

60 

61 

62def _find_windows_compute_sanitizer(ctk_root: str) -> str | None: 

63 return _resolve_candidate_paths( 1aEFGHI

64 ( 

65 os.path.join(ctk_root, "bin", "compute-sanitizer.bat"), 

66 os.path.join(ctk_root, "compute-sanitizer", "compute-sanitizer.exe"), 

67 ) 

68 ) 

69 

70 

71def _resolve_ctk_root_via_canary() -> str | None: 

72 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import resolve_ctk_root_via_canary 1cdajklebmnifogpqh

73 

74 ctk_root: str | None = resolve_ctk_root_via_canary(CTK_ROOT_CANARY_ANCHOR_LIBNAMES[0]) 1cdajklebmnifogpqh

75 return ctk_root 1cdajklebmnifogpqh

76 

77 

78def _resolve_in_trusted_dirs(normalized_name: str, dirs: list[str]) -> str | None: 

79 """Resolve ``normalized_name`` against ``dirs`` in order.""" 

80 seen: set[str] = set() 1XY3Z0cVRLDNTrdajklebmnifogpqhszABCOWxJKPQ1

81 for directory in dirs: 1XY3Z0cVRLDNTrdajklebmnifogpqhszABCOWxJKPQ1

82 if directory in seen: 1XY3Z0cRLDNrdajklebmnifogpqhszABCOxJKPQ1

83 continue 

84 assert directory 1XY3Z0cRLDNrdajklebmnifogpqhszABCOxJKPQ1

85 seen.add(directory) 1XYZ0cRLDNrdajklebmnifogpqhszABCOxJKPQ1

86 candidate = os.path.join(directory, normalized_name) 1XYZ0cRLDNrdajklebmnifogpqhszABCOxJKPQ1

87 if _is_executable_candidate(candidate): 1XYZ0cRLDNrdajklebmnifogpqhszABCOxJKPQ1

88 # Return an absolute path, as the docstring promises (a relative 

89 # search dir would otherwise leak a relative result). 

90 return os.path.abspath(candidate) 1XYZ0cRDNdefghOJKPQ1

91 return None 1XY0cVLDTrdajklebmnifogpqhszABCWxJK

92 

93 

94def _resolve_names_in_trusted_dirs(candidate_names: tuple[str, ...], dirs: list[str]) -> str | None: 

95 """Resolve ordered candidate names within each trusted directory.""" 

96 seen: set[str] = set() 1abMtuvwyUEFGHI

97 for directory in dirs: 1abMtuvwyUEFGHI

98 if directory in seen: 1tuvwy

99 continue 

100 assert directory 1tuvwy

101 seen.add(directory) 1tuvwy

102 found = _resolve_candidate_paths(os.path.join(directory, name) for name in candidate_names) 1tuvwy

103 if found is not None: 1tuvwy

104 return found 1y

105 return None 1abMtuvwUEFGHI

106 

107 

108@functools.cache 

109def find_nvidia_binary_utility(utility_name: str) -> str | None: 

110 """Locate a CUDA binary utility executable. 

111 

112 Args: 

113 utility_name (str): The name of the binary utility to find 

114 (e.g., ``"nvdisasm"``, ``"cuobjdump"``). On Windows, the ``.exe`` 

115 extension will be automatically appended if not present. The function 

116 also recognizes ``.bat`` and ``.cmd`` files on Windows. 

117 

118 Returns: 

119 str or None: Absolute path to the discovered executable, or ``None`` 

120 if the utility cannot be found. The returned path is normalized 

121 (absolute and with resolved separators). 

122 

123 Raises: 

124 UnsupportedBinaryError: If ``utility_name`` is not in the supported set 

125 (see ``SUPPORTED_BINARY_UTILITIES``). 

126 RuntimeError: If a native Windows architecture needed for an 

127 architecture-specific utility layout cannot be determined, or an 

128 installed Nsight product has incomplete or invalid registry data. 

129 

130 Windows on ARM (WoA) Note: 

131 Binary utilities execute in separate processes and do not need to match 

132 the Python process architecture. When choosing among architecture-specific 

133 Windows layouts, this API deliberately targets the native machine 

134 architecture rather than the Python interpreter architecture. For 

135 example, standalone ``nsys`` and ``ncu`` discovery under x64 Python on an 

136 Arm64 machine selects the Arm64 target. This differs from 

137 ``load_nvidia_dynamic_lib`` and ``find_static_lib``, which target the 

138 Python interpreter architecture. 

139 

140 Search order: 

141 1. **NVIDIA Python wheels** 

142 

143 - Scan installed distributions (``site-packages``) for binary layouts 

144 shipped in NVIDIA wheels (e.g., ``cuda-nvcc``). 

145 

146 2. **Conda environments** 

147 

148 - Check Conda-style installation prefixes via ``CONDA_PREFIX`` 

149 environment variable, which use platform-specific bin directory 

150 layouts (``Library/bin`` on Windows, ``bin`` on Linux). 

151 

152 3. **Library-specific standalone installations** 

153 

154 - Search the installation paths for the CUDA Toolkit, Nsight Systems, 

155 and Nsight Compute. 

156 

157 3.1. **Nsight installations**: On Windows, locate Nsight Systems and 

158 Nsight Compute from their installer registry entries. Select 

159 architecture-specific binaries using the native machine 

160 architecture, independent of Python. Lookup of the standalone 

161 ``nsys`` and ``ncu`` CLIs is terminal; a miss does not fall 

162 through to CUDA Toolkit locations. 

163 

164 3.2. **CUDA Toolkit installation**: Use ``CUDA_PATH`` or ``CUDA_HOME`` 

165 (in that order), searching ``bin/x64``, ``bin/x86_64``, and 

166 ``bin`` subdirectories on Windows, or just ``bin`` on Linux. 

167 

168 4. **CTK-root canary fallback** 

169 

170 - For utilities that reach this step after the earlier searches miss, 

171 resolve the ``cudart`` library through the OS dynamic loader, derive 

172 the CUDA Toolkit root from it, and search that root's bin layout. 

173 

174 Note: 

175 Results are cached using ``@functools.cache`` for performance. The cache 

176 persists for the lifetime of the process. 

177 

178 On Windows, executables are identified by their file extensions 

179 (``.exe``, ``.bat``, ``.cmd``). On Unix-like systems, executables 

180 are identified by the ``X_OK`` (execute) permission bit. 

181 

182 Lookup is restricted to the trusted directories and the canary-derived 

183 CTK root listed above. 

184 

185 Example: 

186 >>> from cuda.pathfinder import find_nvidia_binary_utility 

187 >>> nvdisasm = find_nvidia_binary_utility("nvdisasm") 

188 >>> if nvdisasm: 

189 ... print(f"Found nvdisasm at: {nvdisasm}") 

190 """ 

191 if utility_name not in supported_nvidia_binaries.SUPPORTED_BINARIES: 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ2

192 raise UnsupportedBinaryError(utility_name) 12

193 

194 # 1. Search in site-packages (NVIDIA wheels) 

195 candidate_dirs = supported_nvidia_binaries.SITE_PACKAGES_BINDIRS.get(utility_name, ()) 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

196 dirs = [] 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

197 

198 for sub_dir in candidate_dirs: 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

199 dirs.extend(find_sub_dirs_all_sitepackages(sub_dir.split(os.sep))) 1cNTrdajklebmnifogpqhstuvwzABCyOPQ

200 

201 # 2. Search in Conda environment 

202 if (conda_prefix := os.environ.get("CONDA_PREFIX")) is not None: 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

203 if IS_WINDOWS: 1RNrstuvwzABCyOxPQ

204 dirs.append(os.path.join(conda_prefix, "Library", "bin")) 1stuvwzABCyOPQ

205 else: 

206 dirs.append(os.path.join(conda_prefix, "bin")) 1RNrx

207 

208 normalized_name = _normalize_utility_name(utility_name) 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

209 if IS_WINDOWS and utility_name in ("compute-sanitizer", "ncu"): 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

210 candidate_names = (f"{utility_name}.bat", normalized_name) 1abMtuvwyUEFGHI

211 found = _resolve_names_in_trusted_dirs(candidate_names, dirs) 1abMtuvwyUEFGHI

212 else: 

213 found = _resolve_in_trusted_dirs(normalized_name, dirs) 1cVRLDNTrdajklebmnifogpqhszABCOWxJKPQ

214 if found is not None: 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ

215 return found 1RNyOPQ

216 

217 # 3. Search library-specific standalone installations. 

218 # 3.1. Standalone Nsight CLI lookup is terminal; CTK does not contain nsys/ncu. 

219 if IS_WINDOWS and utility_name == "nsys": 1cVLDTrdajklebmnifogpqhsMtuvwzABCUWxEFGHIJK

220 return _resolve_candidate_paths(windows_nsight.nsys_candidate_paths()) 1izABCW

221 if IS_WINDOWS and utility_name == "ncu": 1cVLDTrdajklebmnifogpqhsMtuvwUxEFGHIJK

222 return _resolve_candidate_paths(windows_nsight.ncu_candidate_paths()) 1bMtuvwU

223 

224 # 3.2. Search in CUDA Toolkit (CUDA_PATH/CUDA_HOME). 

225 if (cuda_path := get_cuda_path_or_home()) is not None: 1cVLDTrdajklebmnifogpqhsxEFGHIJK

226 if IS_WINDOWS and utility_name == "compute-sanitizer": 1crdajklebmnifogpqhsxEFGHJK

227 found = _find_windows_compute_sanitizer(cuda_path) 1aEFGH

228 else: 

229 found = _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(cuda_path)) 1crdajklebmnifogpqhsxJK

230 if found is not None: 1crdajklebmnifogpqhsxEFGHJK

231 return found 1cdefghEFGHJK

232 

233 # 4. CTK-root canary fallback. 

234 ctk_root = _resolve_ctk_root_via_canary() 1cVLDTrdajklebmnifogpqhsxI

235 if ctk_root is not None: 1cVLDTrdajklebmnifogpqhsxI

236 if IS_WINDOWS and utility_name == "compute-sanitizer": 1cLDdajklebmnifogpqhI

237 return _find_windows_compute_sanitizer(ctk_root) 1I

238 return _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(ctk_root)) 1cLDdajklebmnifogpqh

239 return None 1cVTrajklmnopqsx