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

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4import functools (empty)

5import os (empty)

6import shutil (empty)

7 

8from cuda.pathfinder._binaries import supported_nvidia_binaries (empty)

9from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path (empty)

10from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages (empty)

11from cuda.pathfinder._utils.platform_aware import IS_WINDOWS (empty)

12 

13 

14class UnsupportedBinaryError(Exception): (empty)

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

16 super().__init__(utility) 1 ctx1w

17 self.utility = utility 1 ctx1w

18 

19 def __str__(self) -> str: (empty)

20 supported_utilities = ", ".join(supported_nvidia_binaries.SUPPORTED_BINARIES) 1 ctx1w

21 return f"Binary '{self.utility}' is not supported. Supported utilities are: {supported_utilities}" 1 ctx1w

22 

23 

24def _normalize_utility_name(utility_name: str) -> str: (empty)

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

26 if IS_WINDOWS and not utility_name.lower().endswith((".exe", ".bat", ".cmd")): 21 ctx1bvurcdefghijklmnopqas

27 return f"{utility_name}.exe" 17 ctx1bcdefghijklmnopqa

28 return utility_name 20 ctx1bvurcdefghijklmnopqs

29 

30 

31@functools.cache (empty)

32def find_nvidia_binary_utility(utility_name: str) -> str | None: (empty)

33 """Locate a CUDA binary utility executable. 

34 

35 Args: 

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

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

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

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

40 

41 Returns: 

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

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

44 (absolute and with resolved separators). 

45 

46 Raises: 

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

48 (see ``SUPPORTED_BINARY_UTILITIES``). 

49 

50 Search order: 

51 1. **NVIDIA Python wheels** 

52 

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

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

55 

56 2. **Conda environments** 

57 

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

59 environment variable, which use platform-specific bin directory 

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

61 

62 3. **CUDA Toolkit environment variables** 

63 

64 - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order), searching 

65 ``bin/x64``, ``bin/x86_64``, and ``bin`` subdirectories on Windows, 

66 or just ``bin`` on Linux. 

67 

68 Note: 

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

70 persists for the lifetime of the process. 

71 

72 On Windows, executables are identified by their file extensions 

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

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

75 

76 Example: 

77 >>> from cuda.pathfinder import find_nvidia_binary_utility 

78 >>> nvdisasm = find_nvidia_binary_utility("nvdisasm") 

79 >>> if nvdisasm: 

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

81 """ 

82 if utility_name not in supported_nvidia_binaries.SUPPORTED_BINARIES: 22 ctx1bvurcdefghijklmnopqasw

83 raise UnsupportedBinaryError(utility_name) 1 ctx1w

84 

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

86 candidate_dirs = supported_nvidia_binaries.SITE_PACKAGES_BINDIRS.get(utility_name, ()) 21 ctx1bvurcdefghijklmnopqas

87 dirs = [] 21 ctx1bvurcdefghijklmnopqas

88 

89 for sub_dir in candidate_dirs: 21 ctx1bvurcdefghijklmnopqas

90 dirs.extend(find_sub_dirs_all_sitepackages(sub_dir.split(os.sep))) 19 ctx1burcdefghijklmnopqa

91 

92 # 2. Search in Conda environment 

93 if (conda_prefix := os.environ.get("CONDA_PREFIX")) is not None: 21 ctx1bvurcdefghijklmnopqas

94 if IS_WINDOWS: 3 ctx1ras

95 dirs.append(os.path.join(conda_prefix, "Library", "bin")) 1 ctx1a

96 else: 

97 dirs.append(os.path.join(conda_prefix, "bin")) 2 ctx1rs

98 

99 # 3. Search in CUDA Toolkit (CUDA_HOME/CUDA_PATH) 

100 if (cuda_home := get_cuda_home_or_path()) is not None: 21 ctx1bvurcdefghijklmnopqas

101 if IS_WINDOWS: 19 ctx1brcdefghijklmnopqas

102 dirs.append(os.path.join(cuda_home, "bin", "x64")) 17 ctx1bcdefghijklmnopqa

103 dirs.append(os.path.join(cuda_home, "bin", "x86_64")) 17 ctx1bcdefghijklmnopqa

104 dirs.append(os.path.join(cuda_home, "bin")) 19 ctx1brcdefghijklmnopqas

105 

106 normalized_name = _normalize_utility_name(utility_name) 21 ctx1bvurcdefghijklmnopqas

107 return shutil.which(normalized_name, path=os.pathsep.join(dirs)) 21 ctx1bvurcdefghijklmnopqas