Coverage for cuda/pathfinder/_static_libs/find_bitcode_lib.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)

6from dataclasses import dataclass (empty)

7from pathlib import Path (empty)

8from typing import NoReturn, TypedDict (empty)

9 

10from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home (empty)

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

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

13 

14 

15class BitcodeLibNotFoundError(RuntimeError): (empty)

16 """Raised when a bitcode library cannot be found.""" 

17 

18 

19@dataclass(frozen=True) (empty)

20class LocatedBitcodeLib: (empty)

21 """Information about a located bitcode library.""" 

22 

23 name: str 

24 abs_path: str 

25 filename: str 

26 found_via: str 

27 

28 

29class _BitcodeLibInfo(TypedDict): (empty)

30 """Bitcode-library metadata with ordered alternatives searched first-to-last.""" 

31 

32 filename: str 

33 rel_path: str 

34 site_packages_dirs: tuple[str, ...] 

35 available_on_windows: bool 

36 

37 

38_SUPPORTED_BITCODE_LIBS_INFO: dict[str, _BitcodeLibInfo] = { (empty)

39 "device": { 

40 "filename": "libdevice.10.bc", 

41 "rel_path": "nvvm/libdevice", 

42 "site_packages_dirs": ( 

43 "nvidia/cu13/nvvm/libdevice", 

44 "nvidia/cuda_nvcc/nvvm/libdevice", 

45 ), 

46 "available_on_windows": True, 

47 }, 

48 "nccl_device": { 

49 "filename": "libnccl_device.bc", 

50 "rel_path": "lib", 

51 "site_packages_dirs": ("nvidia/nccl/lib",), 

52 "available_on_windows": False, 

53 }, 

54 "nvshmem_device": { 

55 "filename": "libnvshmem_device.bc", 

56 "rel_path": "lib", 

57 "site_packages_dirs": ("nvidia/nvshmem/lib",), 

58 "available_on_windows": False, 

59 }, 

60} 

61 

62# Public API: just the supported library names 

63SUPPORTED_BITCODE_LIBS: tuple[str, ...] = tuple( (empty)

64 sorted( 

65 name for name, info in _SUPPORTED_BITCODE_LIBS_INFO.items() if not IS_WINDOWS or info["available_on_windows"] 

66 ) 

67) 

68 

69 

70def _no_such_file_in_dir(directory: Path, filename: str, error_messages: list[str], attachments: list[str]) -> None: (empty)

71 error_messages.append(f"No such file: {directory / filename}") 3 ctx1aef

72 if directory.is_dir(): 3 ctx1aef

73 attachments.append(f' listdir("{directory}"):') 1 ctx1a

74 for node in sorted(node_path.name for node_path in directory.iterdir()): 1 ctx1a

75 attachments.append(f" {node}") 1 ctx1a

76 else: 

77 attachments.append(f' Directory does not exist: "{directory}"') 2 ctx1ef

78 

79 

80class _FindBitcodeLib: (empty)

81 def __init__(self, name: str) -> None: (empty)

82 if name not in _SUPPORTED_BITCODE_LIBS_INFO: # Updated reference 9 ctx1jaghefbcd

83 raise ValueError(f"Unknown bitcode library: '{name}'. Supported: {', '.join(SUPPORTED_BITCODE_LIBS)}") 1 ctx1j

84 self.name: str = name 8 ctx1aghefbcd

85 self.config: _BitcodeLibInfo = _SUPPORTED_BITCODE_LIBS_INFO[name] # Updated reference 8 ctx1aghefbcd

86 self.filename: str = self.config["filename"] 8 ctx1aghefbcd

87 self.rel_path: str = self.config["rel_path"] 8 ctx1aghefbcd

88 self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"] 8 ctx1aghefbcd

89 self.error_messages: list[str] = [] 8 ctx1aghefbcd

90 self.attachments: list[str] = [] 8 ctx1aghefbcd

91 

92 def try_site_packages(self) -> Path | None: (empty)

93 for rel_dir in self.site_packages_dirs: 8 ctx1aghefbcd

94 sub_dir = tuple(rel_dir.split("/")) 8 ctx1aghefbcd

95 for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): 8 ctx1aghefbcd

96 file_path = Path(abs_dir, self.filename) 3 ctx1bcd

97 if file_path.is_file(): 3 ctx1bcd

98 return file_path 3 ctx1bcd

99 return None 8 ctx1aghefbcd

100 

101 def try_with_conda_prefix(self) -> Path | None: (empty)

102 conda_prefix = os.environ.get("CONDA_PREFIX") 8 ctx1aghefbcd

103 if not conda_prefix: 8 ctx1aghefbcd

104 return None 5 ctx1aghef

105 

106 anchor = Path(conda_prefix, "Library") if IS_WINDOWS else Path(conda_prefix) 3 ctx1bcd

107 file_path = anchor / self.rel_path / self.filename 3 ctx1bcd

108 if file_path.is_file(): 3 ctx1bcd

109 return file_path 3 ctx1bcd

110 return None 3 ctx1bcd

111 

112 def try_with_cuda_home(self) -> Path | None: (empty)

113 cuda_home = get_cuda_path_or_home() 8 ctx1aghefbcd

114 if cuda_home is None: 8 ctx1aghefbcd

115 self.error_messages.append("CUDA_HOME/CUDA_PATH not set") 1 ctx1g

116 return None 1 ctx1g

117 

118 anchor = Path(cuda_home) 7 ctx1ahefbcd

119 file_path = anchor / self.rel_path / self.filename 7 ctx1ahefbcd

120 if file_path.is_file(): 7 ctx1ahefbcd

121 return file_path 4 ctx1hbcd

122 

123 _no_such_file_in_dir( 3 ctx1aef

124 anchor / self.rel_path, 

125 self.filename, 

126 self.error_messages, 

127 self.attachments, 

128 ) 

129 return None 3 ctx1aef

130 

131 def raise_not_found_error(self) -> NoReturn: (empty)

132 err = ", ".join(self.error_messages) if self.error_messages else "No search paths available" 4 ctx1agef

133 att = "\n".join(self.attachments) if self.attachments else "" 4 ctx1agef

134 raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}') 4 ctx1agef

135 

136 

137def locate_bitcode_lib(name: str) -> LocatedBitcodeLib: (empty)

138 """Locate a bitcode library by name. 

139 

140 Raises: 

141 ValueError: If ``name`` is not a supported bitcode library. 

142 BitcodeLibNotFoundError: If the bitcode library cannot be found. 

143 """ 

144 finder = _FindBitcodeLib(name) 9 ctx1jaghefbcd

145 

146 abs_path = finder.try_site_packages() 8 ctx1aghefbcd

147 if abs_path is not None: 8 ctx1aghefbcd

148 return LocatedBitcodeLib( 3 ctx1bcd

149 name=name, 

150 abs_path=str(abs_path), 

151 filename=finder.filename, 

152 found_via="site-packages", 

153 ) 

154 

155 abs_path = finder.try_with_conda_prefix() 8 ctx1aghefbcd

156 if abs_path is not None: 8 ctx1aghefbcd

157 return LocatedBitcodeLib( 3 ctx1bcd

158 name=name, 

159 abs_path=str(abs_path), 

160 filename=finder.filename, 

161 found_via="conda", 

162 ) 

163 

164 abs_path = finder.try_with_cuda_home() 8 ctx1aghefbcd

165 if abs_path is not None: 8 ctx1aghefbcd

166 return LocatedBitcodeLib( 4 ctx1hbcd

167 name=name, 

168 abs_path=str(abs_path), 

169 filename=finder.filename, 

170 found_via="CUDA_PATH", 

171 ) 

172 

173 finder.raise_not_found_error() 4 ctx1agef

174 

175 

176@functools.cache (empty)

177def find_bitcode_lib(name: str) -> str: (empty)

178 """Find the absolute path to a bitcode library. 

179 

180 Raises: 

181 ValueError: If ``name`` is not a supported bitcode library. 

182 BitcodeLibNotFoundError: If the bitcode library cannot be found. 

183 """ 

184 return locate_bitcode_lib(name).abs_path 3 ctx1agh