File size: 4,296 Bytes
76f9669 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import functools
import os
import shutil
from cuda.pathfinder._binaries import supported_nvidia_binaries
from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
class UnsupportedBinaryError(Exception):
def __init__(self, utility: str) -> None:
super().__init__(utility)
self.utility = utility
def __str__(self) -> str:
supported_utilities = ", ".join(supported_nvidia_binaries.SUPPORTED_BINARIES)
return f"Binary '{self.utility}' is not supported. Supported utilities are: {supported_utilities}"
def _normalize_utility_name(utility_name: str) -> str:
"""Normalize utility name by adding .exe on Windows if needed."""
if IS_WINDOWS and not utility_name.lower().endswith((".exe", ".bat", ".cmd")):
return f"{utility_name}.exe"
return utility_name
@functools.cache
def find_nvidia_binary_utility(utility_name: str) -> str | None:
"""Locate a CUDA binary utility executable.
Args:
utility_name (str): The name of the binary utility to find
(e.g., ``"nvdisasm"``, ``"cuobjdump"``). On Windows, the ``.exe``
extension will be automatically appended if not present. The function
also recognizes ``.bat`` and ``.cmd`` files on Windows.
Returns:
str or None: Absolute path to the discovered executable, or ``None``
if the utility cannot be found. The returned path is normalized
(absolute and with resolved separators).
Raises:
UnsupportedBinaryError: If ``utility_name`` is not in the supported set
(see ``SUPPORTED_BINARY_UTILITIES``).
Search order:
1. **NVIDIA Python wheels**
- Scan installed distributions (``site-packages``) for binary layouts
shipped in NVIDIA wheels (e.g., ``cuda-nvcc``).
2. **Conda environments**
- Check Conda-style installation prefixes via ``CONDA_PREFIX``
environment variable, which use platform-specific bin directory
layouts (``Library/bin`` on Windows, ``bin`` on Linux).
3. **CUDA Toolkit environment variables**
- Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order), searching
``bin/x64``, ``bin/x86_64``, and ``bin`` subdirectories on Windows,
or just ``bin`` on Linux.
Note:
Results are cached using ``@functools.cache`` for performance. The cache
persists for the lifetime of the process.
On Windows, executables are identified by their file extensions
(``.exe``, ``.bat``, ``.cmd``). On Unix-like systems, executables
are identified by the ``X_OK`` (execute) permission bit.
Example:
>>> from cuda.pathfinder import find_nvidia_binary_utility
>>> nvdisasm = find_nvidia_binary_utility("nvdisasm")
>>> if nvdisasm:
... print(f"Found nvdisasm at: {nvdisasm}")
"""
if utility_name not in supported_nvidia_binaries.SUPPORTED_BINARIES:
raise UnsupportedBinaryError(utility_name)
# 1. Search in site-packages (NVIDIA wheels)
candidate_dirs = supported_nvidia_binaries.SITE_PACKAGES_BINDIRS.get(utility_name, ())
dirs = []
for sub_dir in candidate_dirs:
dirs.extend(find_sub_dirs_all_sitepackages(sub_dir.split(os.sep)))
# 2. Search in Conda environment
if (conda_prefix := os.environ.get("CONDA_PREFIX")) is not None:
if IS_WINDOWS:
dirs.append(os.path.join(conda_prefix, "Library", "bin"))
else:
dirs.append(os.path.join(conda_prefix, "bin"))
# 3. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH)
if (cuda_home := get_cuda_home_or_path()) is not None:
if IS_WINDOWS:
dirs.append(os.path.join(cuda_home, "bin", "x64"))
dirs.append(os.path.join(cuda_home, "bin", "x86_64"))
dirs.append(os.path.join(cuda_home, "bin"))
normalized_name = _normalize_utility_name(utility_name)
return shutil.which(normalized_name, path=os.pathsep.join(dirs))
|