Coverage for cuda/pathfinder/_binaries/find_nvidia_binary_utility.py: 98.46%
65 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-29 01:38 +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
7from cuda.pathfinder._binaries import supported_nvidia_binaries
8from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES
9from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home
10from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
11from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
14class UnsupportedBinaryError(Exception):
15 def __init__(self, utility: str) -> None:
16 super().__init__(utility) 1F
17 self.utility = utility 1F
19 def __str__(self) -> str:
20 supported_utilities = ", ".join(supported_nvidia_binaries.SUPPORTED_BINARIES) 1F
21 return f"Binary '{self.utility}' is not supported. Supported utilities are: {supported_utilities}" 1F
24def _normalize_utility_name(utility_name: str) -> str:
25 """Normalize utility name by adding .exe on Windows if needed."""
26 if IS_WINDOWS and not utility_name.lower().endswith((".exe", ".bat", ".cmd")): 1azwvutxqbfghicjklmdneoprs
27 return f"{utility_name}.exe" 1avbfghicjklmdneopr
28 return utility_name 1azwutxqbfghicjklmdneops
31def _is_executable_candidate(path: str) -> bool:
32 if not os.path.isfile(path): 1ABCDabfghicjklmdneop
33 return False 1Aabfghicjklmdneop
34 if IS_WINDOWS: 1ABCDabcde
35 return True 1ABCabcde
36 return os.access(path, os.X_OK) 1ABCDabcde
39def _ctk_bin_subdirs(root: str) -> list[str]:
40 if IS_WINDOWS: 1avutqbfghicjklmdneoprs
41 return [ 1avbfghicjklmdneopr
42 os.path.join(root, "bin", "x64"),
43 os.path.join(root, "bin", "x86_64"),
44 os.path.join(root, "bin"),
45 ]
46 return [os.path.join(root, "bin")] 1autqbfghicjklmdneops
49def _resolve_ctk_root_via_canary() -> str | None:
50 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import resolve_ctk_root_via_canary 1abfghicjklmdneop
52 ctk_root: str | None = resolve_ctk_root_via_canary(CTK_ROOT_CANARY_ANCHOR_LIBNAMES[0]) 1abfghicjklmdneop
53 return ctk_root 1abfghicjklmdneop
56def _resolve_in_trusted_dirs(normalized_name: str, dirs: list[str]) -> str | None:
57 """Resolve ``normalized_name`` against ``dirs`` in order."""
58 seen: set[str] = set() 1ABGCDazwvutxqbfghicjklmdneoprsE
59 for directory in dirs: 1ABGCDazwvutxqbfghicjklmdneoprsE
60 if directory in seen: 1ABGCDawvutqbfghicjklmdneoprsE
61 continue
62 assert directory 1ABGCDawvutqbfghicjklmdneoprsE
63 seen.add(directory) 1ABCDawvutqbfghicjklmdneoprsE
64 candidate = os.path.join(directory, normalized_name) 1ABCDawvutqbfghicjklmdneoprsE
65 if _is_executable_candidate(candidate): 1ABCDawvutqbfghicjklmdneoprsE
66 # Return an absolute path, as the docstring promises (a relative
67 # search dir would otherwise leak a relative result).
68 return os.path.abspath(candidate) 1ABCDawutbcdeE
69 return None 1ABDazvuxqbfghicjklmdneoprs
72@functools.cache
73def find_nvidia_binary_utility(utility_name: str) -> str | None:
74 """Locate a CUDA binary utility executable.
76 Args:
77 utility_name (str): The name of the binary utility to find
78 (e.g., ``"nvdisasm"``, ``"cuobjdump"``). On Windows, the ``.exe``
79 extension will be automatically appended if not present. The function
80 also recognizes ``.bat`` and ``.cmd`` files on Windows.
82 Returns:
83 str or None: Absolute path to the discovered executable, or ``None``
84 if the utility cannot be found. The returned path is normalized
85 (absolute and with resolved separators).
87 Raises:
88 UnsupportedBinaryError: If ``utility_name`` is not in the supported set
89 (see ``SUPPORTED_BINARY_UTILITIES``).
91 Search order:
92 1. **NVIDIA Python wheels**
94 - Scan installed distributions (``site-packages``) for binary layouts
95 shipped in NVIDIA wheels (e.g., ``cuda-nvcc``).
97 2. **Conda environments**
99 - Check Conda-style installation prefixes via ``CONDA_PREFIX``
100 environment variable, which use platform-specific bin directory
101 layouts (``Library/bin`` on Windows, ``bin`` on Linux).
103 3. **CUDA Toolkit environment variables**
105 - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order), searching
106 ``bin/x64``, ``bin/x86_64``, and ``bin`` subdirectories on Windows,
107 or just ``bin`` on Linux.
109 4. **CTK-root canary fallback**
111 - Only when steps 1-3 miss: resolve the ``cudart`` library through the
112 OS dynamic loader, derive the CUDA Toolkit root from it, and search
113 that root's bin layout.
115 Note:
116 Results are cached using ``@functools.cache`` for performance. The cache
117 persists for the lifetime of the process.
119 On Windows, executables are identified by their file extensions
120 (``.exe``, ``.bat``, ``.cmd``). On Unix-like systems, executables
121 are identified by the ``X_OK`` (execute) permission bit.
123 Lookup is restricted to the trusted directories and the canary-derived
124 CTK root listed above.
126 Example:
127 >>> from cuda.pathfinder import find_nvidia_binary_utility
128 >>> nvdisasm = find_nvidia_binary_utility("nvdisasm")
129 >>> if nvdisasm:
130 ... print(f"Found nvdisasm at: {nvdisasm}")
131 """
132 if utility_name not in supported_nvidia_binaries.SUPPORTED_BINARIES: 1azwvutxqbfghicjklmdneoprsF
133 raise UnsupportedBinaryError(utility_name) 1F
135 # 1. Search in site-packages (NVIDIA wheels)
136 candidate_dirs = supported_nvidia_binaries.SITE_PACKAGES_BINDIRS.get(utility_name, ()) 1azwvutxqbfghicjklmdneoprs
137 dirs = [] 1azwvutxqbfghicjklmdneoprs
139 for sub_dir in candidate_dirs: 1azwvutxqbfghicjklmdneoprs
140 dirs.extend(find_sub_dirs_all_sitepackages(sub_dir.split(os.sep))) 1atxqbfghicjklmdneopr
142 # 2. Search in Conda environment
143 if (conda_prefix := os.environ.get("CONDA_PREFIX")) is not None: 1azwvutxqbfghicjklmdneoprs
144 if IS_WINDOWS: 1wtqrs
145 dirs.append(os.path.join(conda_prefix, "Library", "bin")) 1r
146 else:
147 dirs.append(os.path.join(conda_prefix, "bin")) 1wtqs
149 # 3. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH)
150 if (cuda_home := get_cuda_path_or_home()) is not None: 1azwvutxqbfghicjklmdneoprs
151 dirs.extend(_ctk_bin_subdirs(cuda_home)) 1atqbfghicjklmdneoprs
153 normalized_name = _normalize_utility_name(utility_name) 1azwvutxqbfghicjklmdneoprs
154 found = _resolve_in_trusted_dirs(normalized_name, dirs) 1azwvutxqbfghicjklmdneoprs
155 if found is not None: 1azwvutxqbfghicjklmdneoprs
156 return found 1awtbcde
158 # 4. CTK-root canary fallback.
159 ctk_root = _resolve_ctk_root_via_canary() 1azvuxqbfghicjklmdneoprs
160 if ctk_root is not None: 1azvuxqbfghicjklmdneoprs
161 return _resolve_in_trusted_dirs(normalized_name, _ctk_bin_subdirs(ctk_root)) 1avubfghicjklmdneop
162 return None 1azxqfghijklmnoprs