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
« 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
4import functools
5import os
6from collections.abc import Iterable
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
15class UnsupportedBinaryError(Exception):
16 def __init__(self, utility: str) -> None:
17 super().__init__(utility) 12
18 self.utility = utility 12
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
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
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
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
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
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 )
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
74 ctk_root: str | None = resolve_ctk_root_via_canary(CTK_ROOT_CANARY_ANCHOR_LIBNAMES[0]) 1cdajklebmnifogpqh
75 return ctk_root 1cdajklebmnifogpqh
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
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
108@functools.cache
109def find_nvidia_binary_utility(utility_name: str) -> str | None:
110 """Locate a CUDA binary utility executable.
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.
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).
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.
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.
140 Search order:
141 1. **NVIDIA Python wheels**
143 - Scan installed distributions (``site-packages``) for binary layouts
144 shipped in NVIDIA wheels (e.g., ``cuda-nvcc``).
146 2. **Conda environments**
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).
152 3. **Library-specific standalone installations**
154 - Search the installation paths for the CUDA Toolkit, Nsight Systems,
155 and Nsight Compute.
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.
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.
168 4. **CTK-root canary fallback**
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.
174 Note:
175 Results are cached using ``@functools.cache`` for performance. The cache
176 persists for the lifetime of the process.
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.
182 Lookup is restricted to the trusted directories and the canary-derived
183 CTK root listed above.
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
194 # 1. Search in site-packages (NVIDIA wheels)
195 candidate_dirs = supported_nvidia_binaries.SITE_PACKAGES_BINDIRS.get(utility_name, ()) 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ
196 dirs = [] 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ
198 for sub_dir in candidate_dirs: 1cVRLDNTrdajklebmnifogpqhsMtuvwzABCyOUWxEFGHIJKPQ
199 dirs.extend(find_sub_dirs_all_sitepackages(sub_dir.split(os.sep))) 1cNTrdajklebmnifogpqhstuvwzABCyOPQ
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
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
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
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
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