Coverage for cuda / pathfinder / _static_libs / find_bitcode_lib.py: 100.00%
77 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-29 01:27 +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 dataclasses import dataclass
7from typing import NoReturn, TypedDict
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 BitcodeLibNotFoundError(RuntimeError):
15 """Raised when a bitcode library cannot be found."""
18@dataclass(frozen=True)
19class LocatedBitcodeLib:
20 """Information about a located bitcode library."""
22 name: str
23 abs_path: str
24 filename: str
25 found_via: str
28class _BitcodeLibInfo(TypedDict):
29 filename: str
30 rel_path: str
31 site_packages_dirs: tuple[str, ...]
32 available_on_windows: bool
35_SUPPORTED_BITCODE_LIBS_INFO: dict[str, _BitcodeLibInfo] = {
36 "device": {
37 "filename": "libdevice.10.bc",
38 "rel_path": os.path.join("nvvm", "libdevice"),
39 "site_packages_dirs": (
40 "nvidia/cu13/nvvm/libdevice",
41 "nvidia/cuda_nvcc/nvvm/libdevice",
42 ),
43 "available_on_windows": True,
44 },
45 "nccl_device": {
46 "filename": "libnccl_device.bc",
47 "rel_path": "lib",
48 "site_packages_dirs": ("nvidia/nccl/lib",),
49 "available_on_windows": False,
50 },
51 "nvshmem_device": {
52 "filename": "libnvshmem_device.bc",
53 "rel_path": "lib",
54 "site_packages_dirs": ("nvidia/nvshmem/lib",),
55 "available_on_windows": False,
56 },
57}
59# Public API: just the supported library names
60SUPPORTED_BITCODE_LIBS: tuple[str, ...] = tuple(
61 sorted(
62 name for name, info in _SUPPORTED_BITCODE_LIBS_INFO.items() if not IS_WINDOWS or info["available_on_windows"]
63 )
64)
67def _no_such_file_in_dir(dir_path: str, filename: str, error_messages: list[str], attachments: list[str]) -> None:
68 error_messages.append(f"No such file: {os.path.join(dir_path, filename)}") 1aef
69 if os.path.isdir(dir_path): 1aef
70 attachments.append(f' listdir("{dir_path}"):') 1a
71 for node in sorted(os.listdir(dir_path)): 1a
72 attachments.append(f" {node}") 1a
73 else:
74 attachments.append(f' Directory does not exist: "{dir_path}"') 1ef
77class _FindBitcodeLib:
78 def __init__(self, name: str) -> None:
79 if name not in _SUPPORTED_BITCODE_LIBS_INFO: # Updated reference 1jaghefbcd
80 raise ValueError(f"Unknown bitcode library: '{name}'. Supported: {', '.join(SUPPORTED_BITCODE_LIBS)}") 1j
81 self.name: str = name 1aghefbcd
82 self.config: _BitcodeLibInfo = _SUPPORTED_BITCODE_LIBS_INFO[name] # Updated reference 1aghefbcd
83 self.filename: str = self.config["filename"] 1aghefbcd
84 self.rel_path: str = self.config["rel_path"] 1aghefbcd
85 self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"] 1aghefbcd
86 self.error_messages: list[str] = [] 1aghefbcd
87 self.attachments: list[str] = [] 1aghefbcd
89 def try_site_packages(self) -> str | None:
90 for rel_dir in self.site_packages_dirs: 1aghefbcd
91 sub_dir = tuple(rel_dir.split("/")) 1aghefbcd
92 for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): 1aghefbcd
93 file_path = os.path.join(abs_dir, self.filename) 1bcd
94 if os.path.isfile(file_path): 1bcd
95 return file_path 1bcd
96 return None 1aghefbcd
98 def try_with_conda_prefix(self) -> str | None:
99 conda_prefix = os.environ.get("CONDA_PREFIX") 1aghefbcd
100 if not conda_prefix: 1aghefbcd
101 return None 1aghef
103 anchor = os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix 1bcd
104 file_path = os.path.join(anchor, self.rel_path, self.filename) 1bcd
105 if os.path.isfile(file_path): 1bcd
106 return file_path 1bcd
107 return None 1bcd
109 def try_with_cuda_home(self) -> str | None:
110 cuda_home = get_cuda_path_or_home() 1aghefbcd
111 if cuda_home is None: 1aghefbcd
112 self.error_messages.append("CUDA_HOME/CUDA_PATH not set") 1g
113 return None 1g
115 file_path = os.path.join(cuda_home, self.rel_path, self.filename) 1ahefbcd
116 if os.path.isfile(file_path): 1ahefbcd
117 return file_path 1hbcd
119 _no_such_file_in_dir( 1aef
120 os.path.join(cuda_home, self.rel_path),
121 self.filename,
122 self.error_messages,
123 self.attachments,
124 )
125 return None 1aef
127 def raise_not_found_error(self) -> NoReturn:
128 err = ", ".join(self.error_messages) if self.error_messages else "No search paths available" 1agef
129 att = "\n".join(self.attachments) if self.attachments else "" 1agef
130 raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}') 1agef
133def locate_bitcode_lib(name: str) -> LocatedBitcodeLib:
134 """Locate a bitcode library by name.
136 Raises:
137 ValueError: If ``name`` is not a supported bitcode library.
138 BitcodeLibNotFoundError: If the bitcode library cannot be found.
139 """
140 finder = _FindBitcodeLib(name) 1jaghefbcd
142 abs_path = finder.try_site_packages() 1aghefbcd
143 if abs_path is not None: 1aghefbcd
144 return LocatedBitcodeLib( 1bcd
145 name=name,
146 abs_path=abs_path,
147 filename=finder.filename,
148 found_via="site-packages",
149 )
151 abs_path = finder.try_with_conda_prefix() 1aghefbcd
152 if abs_path is not None: 1aghefbcd
153 return LocatedBitcodeLib( 1bcd
154 name=name,
155 abs_path=abs_path,
156 filename=finder.filename,
157 found_via="conda",
158 )
160 abs_path = finder.try_with_cuda_home() 1aghefbcd
161 if abs_path is not None: 1aghefbcd
162 return LocatedBitcodeLib( 1hbcd
163 name=name,
164 abs_path=abs_path,
165 filename=finder.filename,
166 found_via="CUDA_PATH",
167 )
169 finder.raise_not_found_error() 1agef
172@functools.cache
173def find_bitcode_lib(name: str) -> str:
174 """Find the absolute path to a bitcode library.
176 Raises:
177 ValueError: If ``name`` is not a supported bitcode library.
178 BitcodeLibNotFoundError: If the bitcode library cannot be found.
179 """
180 return locate_bitcode_lib(name).abs_path 1agh