feat(binaries): find binaries using pathfinder

rparolin · cpcloud · commit aa6cecec8651 · 2026-02-03T08:24:40.000-05:00
diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py
@@ -3,6 +3,10 @@
 
 """cuda.pathfinder public APIs"""
 
+from cuda.pathfinder._binaries.find_nvidia_binary_utility import (
+    find_nvidia_binary_utility as find_nvidia_binary_utility,
+)
+from cuda.pathfinder._binaries.supported_nvidia_binaries import SUPPORTED_BINARIES as _SUPPORTED_BINARIES
 from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError as DynamicLibNotFoundError
 from cuda.pathfinder._dynamic_libs.load_dl_common import LoadedDL as LoadedDL
 from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import load_nvidia_dynamic_lib as load_nvidia_dynamic_lib
@@ -25,6 +29,12 @@
 #: (e.g., ``"cufile"`` may be Linux-only).
 SUPPORTED_HEADERS_CTK = _SUPPORTED_HEADERS_CTK
 
+#: Tuple of supported CUDA binary utility names that can be located
+#: via ``find_nvidia_binary_utility()``. Platform-aware (e.g., some
+#: utilities may be available only on Linux or Windows).
+#: Example utilities: ``"nvdisasm"``, ``"cuobjdump"``, ``"nvcc"``.
+SUPPORTED_BINARY_UTILITIES = _SUPPORTED_BINARIES
+
 # Backward compatibility: _find_nvidia_header_directory was added in release 1.2.2.
 # It will be removed in release 1.2.4.
 _find_nvidia_header_directory = find_nvidia_header_directory
diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py b/cuda_pathfinder/cuda/pathfinder/_binaries/find_nvidia_binary_utility.py
@@ -0,0 +1,107 @@
+# 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:
+        if bin_dir := find_sub_dirs_all_sitepackages(sub_dir.split(os.sep)):
+            dirs.extend(bin_dir)
+
+    # 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"))
+        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))
diff --git a/cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py b/cuda_pathfinder/cuda/pathfinder/_binaries/supported_nvidia_binaries.py
@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+
+# Common CUDA binary utilities available on both Linux and Windows
+SUPPORTED_BINARIES_ALL = (
+    # Core compilation tools
+    "nvcc",
+    "nvdisasm",
+    "cuobjdump",
+    "nvprune",
+    "fatbinary",
+    "bin2c",
+    "nvlink",
+    # Runtime/debugging tools
+    "cuda-gdb",
+    "cuda-gdbserver",
+    "compute-sanitizer",
+    # Profiling tools
+    "nvprof",
+    "nsys",
+    "nsight-sys",
+    "ncu",
+    "nsight-compute",
+)
+
+SUPPORTED_BINARIES = SUPPORTED_BINARIES_ALL
+
+# Site-packages bin directories where binaries might be found
+# Based on NVIDIA wheel layouts (same for Linux and Windows)
+SITE_PACKAGES_BINDIRS = {
+    "nvcc": ("nvidia/cuda_nvcc/bin",),
+    "nvdisasm": ("nvidia/cuda_nvcc/bin",),
+    "cuobjdump": ("nvidia/cuda_nvcc/bin",),
+    "nvprune": ("nvidia/cuda_nvcc/bin",),
+    "fatbinary": ("nvidia/cuda_nvcc/bin",),
+    "bin2c": ("nvidia/cuda_nvcc/bin",),
+    "nvlink": ("nvidia/cuda_nvcc/bin",),
+    "cuda-gdb": ("nvidia/cuda_nvcc/bin",),
+    "cuda-gdbserver": ("nvidia/cuda_nvcc/bin",),
+    "compute-sanitizer": ("nvidia/cuda_nvcc/bin",),
+    "nvprof": ("nvidia/cuda_nvcc/bin",),
+    "nsys": ("nvidia/nsight_systems/bin",),
+    "nsight-sys": ("nvidia/nsight_systems/bin",),
+    "ncu": ("nvidia/nsight_compute/bin",),
+    "nsight-compute": ("nvidia/nsight_compute/bin",),
+}
diff --git a/cuda_pathfinder/docs/source/api.rst b/cuda_pathfinder/docs/source/api.rst
@@ -7,7 +7,7 @@
 =================================
 
 The ``cuda.pathfinder`` module provides utilities for loading NVIDIA dynamic libraries,
-and experimental APIs for locating NVIDIA C/C++ header directories.
+locating NVIDIA C/C++ header directories, and finding CUDA binary utilities.
 
 .. autosummary::
    :toctree: generated/
@@ -20,3 +20,6 @@ and experimental APIs for locating NVIDIA C/C++ header directories.
    SUPPORTED_HEADERS_CTK
    SUPPORTED_HEADERS_NON_CTK
    find_nvidia_header_directory
+
+   SUPPORTED_BINARY_UTILITIES
+   find_nvidia_binary_utility
diff --git a/cuda_pathfinder/tests/test_find_nvidia_binaries.py b/cuda_pathfinder/tests/test_find_nvidia_binaries.py
@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import pytest
+
+from cuda.pathfinder import find_nvidia_binary_utility
+from cuda.pathfinder._binaries.find_nvidia_binary_utility import UnsupportedBinaryError
+from cuda.pathfinder._binaries.supported_nvidia_binaries import (
+    SITE_PACKAGES_BINDIRS,
+    SUPPORTED_BINARIES,
+    SUPPORTED_BINARIES_ALL,
+)
+
+
+def test_unknown_utility_name():
+    with pytest.raises(UnsupportedBinaryError, match=r"'unknown-utility' is not supported"):
+        find_nvidia_binary_utility("unknown-utility")
+
+
+@pytest.mark.parametrize("utility_name", SUPPORTED_BINARIES)
+def test_find_binary_utilities(info_summary_append, utility_name):
+    bin_path = find_nvidia_binary_utility(utility_name)
+    info_summary_append(f"{bin_path=!r}")
+
+    if bin_path:
+        assert os.path.isfile(bin_path), f"Path exists but is not a file: {bin_path}"
+        # Note: We verify the file exists but don't check executability here because
+        # permissions may vary in test environments (e.g., mounted filesystems, CI
+        # containers). The _is_executable() check is tested separately in unit tests.
+
+
+def test_supported_binaries_consistency():
+    # Ensure SUPPORTED_BINARIES is a subset of SUPPORTED_BINARIES_ALL
+    assert set(SUPPORTED_BINARIES).issubset(SUPPORTED_BINARIES_ALL)
+
+
+def test_site_packages_bindirs_consistency():
+    """Verify SITE_PACKAGES_BINDIRS keys are in SUPPORTED_BINARIES_ALL."""
+
+    assert set(SITE_PACKAGES_BINDIRS).issubset(SUPPORTED_BINARIES_ALL)
+
+
+def test_caching_per_utility():
+    """Verify that different utilities have independent cache entries."""
+    nvdisasm1 = find_nvidia_binary_utility("nvdisasm")
+    nvcc1 = find_nvidia_binary_utility("nvcc")
+    nvdisasm2 = find_nvidia_binary_utility("nvdisasm")
+    nvcc2 = find_nvidia_binary_utility("nvcc")
+
+    # Same utility should return cached result
+    assert nvdisasm1 is nvdisasm2
+    assert nvcc1 is nvcc2
+
+    # Different utilities should have different results (unless at least one of
+    # them is None)
+    if nvdisasm1 is not None and nvcc1 is not None:
+        assert nvdisasm1 != nvcc1
