Coverage for cuda / pathfinder / _dynamic_libs / find_nvidia_dynamic_lib.py: 54%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4import glob 

5import os 

6from collections.abc import Sequence 

7 

8from cuda.pathfinder._dynamic_libs.load_dl_common import DynamicLibNotFoundError 

9from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import ( 

10 SITE_PACKAGES_LIBDIRS_LINUX, 

11 SITE_PACKAGES_LIBDIRS_WINDOWS, 

12 is_suppressed_dll_file, 

13) 

14from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path 

15from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages 

16from cuda.pathfinder._utils.platform_aware import IS_WINDOWS 

17 

18 

19def _no_such_file_in_sub_dirs( 

20 sub_dirs: Sequence[str], file_wild: str, error_messages: list[str], attachments: list[str] 

21) -> None: 

22 error_messages.append(f"No such file: {file_wild}") 

23 for sub_dir in find_sub_dirs_all_sitepackages(sub_dirs): 

24 attachments.append(f' listdir("{sub_dir}"):') 

25 for node in sorted(os.listdir(sub_dir)): 

26 attachments.append(f" {node}") 

27 

28 

29def _find_so_using_nvidia_lib_dirs( 

30 libname: str, so_basename: str, error_messages: list[str], attachments: list[str] 

31) -> str | None: 

32 rel_dirs = SITE_PACKAGES_LIBDIRS_LINUX.get(libname) 

33 if rel_dirs is not None: 

34 sub_dirs_searched = [] 

35 file_wild = so_basename + "*" 

36 for rel_dir in rel_dirs: 

37 sub_dir = tuple(rel_dir.split(os.path.sep)) 

38 for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): 

39 # First look for an exact match 

40 so_name = os.path.join(abs_dir, so_basename) 

41 if os.path.isfile(so_name): 

42 return so_name 

43 # Look for a versioned library 

44 # Using sort here mainly to make the result deterministic. 

45 for so_name in sorted(glob.glob(os.path.join(abs_dir, file_wild))): 

46 if os.path.isfile(so_name): 

47 return so_name 

48 sub_dirs_searched.append(sub_dir) 

49 for sub_dir in sub_dirs_searched: 

50 _no_such_file_in_sub_dirs(sub_dir, file_wild, error_messages, attachments) 

51 return None 

52 

53 

54def _find_dll_under_dir(dirpath: str, file_wild: str) -> str | None: 

55 for path in sorted(glob.glob(os.path.join(dirpath, file_wild))): 

56 if not os.path.isfile(path): 

57 continue 

58 if not is_suppressed_dll_file(os.path.basename(path)): 

59 return path 

60 return None 

61 

62 

63def _find_dll_using_nvidia_bin_dirs( 

64 libname: str, lib_searched_for: str, error_messages: list[str], attachments: list[str] 

65) -> str | None: 

66 rel_dirs = SITE_PACKAGES_LIBDIRS_WINDOWS.get(libname) 

67 if rel_dirs is not None: 

68 sub_dirs_searched = [] 

69 for rel_dir in rel_dirs: 

70 sub_dir = tuple(rel_dir.split(os.path.sep)) 

71 for abs_dir in find_sub_dirs_all_sitepackages(sub_dir): 

72 dll_name = _find_dll_under_dir(abs_dir, lib_searched_for) 

73 if dll_name is not None: 

74 return dll_name 

75 sub_dirs_searched.append(sub_dir) 

76 for sub_dir in sub_dirs_searched: 

77 _no_such_file_in_sub_dirs(sub_dir, lib_searched_for, error_messages, attachments) 

78 return None 

79 

80 

81def _find_lib_dir_using_anchor_point(libname: str, anchor_point: str, linux_lib_dir: str) -> str | None: 

82 # Resolve paths for the four cases: 

83 # Windows/Linux x nvvm yes/no 

84 if IS_WINDOWS: 

85 if libname == "nvvm": # noqa: SIM108 

86 rel_paths = [ 

87 "nvvm/bin/*", # CTK 13 

88 "nvvm/bin", # CTK 12 

89 ] 

90 else: 

91 rel_paths = [ 

92 "bin/x64", # CTK 13 

93 "bin", # CTK 12 

94 ] 

95 else: 

96 if libname == "nvvm": # noqa: SIM108 

97 rel_paths = ["nvvm/lib64"] 

98 else: 

99 rel_paths = [linux_lib_dir] 

100 

101 for rel_path in rel_paths: 

102 for dirname in sorted(glob.glob(os.path.join(anchor_point, rel_path))): 

103 if os.path.isdir(dirname): 

104 return dirname 

105 

106 return None 

107 

108 

109def _find_lib_dir_using_cuda_home(libname: str) -> str | None: 

110 cuda_home = get_cuda_home_or_path() 

111 if cuda_home is None: 

112 return None 

113 return _find_lib_dir_using_anchor_point(libname, anchor_point=cuda_home, linux_lib_dir="lib64") 

114 

115 

116def _find_lib_dir_using_conda_prefix(libname: str) -> str | None: 

117 conda_prefix = os.environ.get("CONDA_PREFIX") 

118 if not conda_prefix: 

119 return None 

120 return _find_lib_dir_using_anchor_point( 

121 libname, anchor_point=os.path.join(conda_prefix, "Library") if IS_WINDOWS else conda_prefix, linux_lib_dir="lib" 

122 ) 

123 

124 

125def _find_so_using_lib_dir( 

126 lib_dir: str, so_basename: str, error_messages: list[str], attachments: list[str] 

127) -> str | None: 

128 so_name = os.path.join(lib_dir, so_basename) 

129 if os.path.isfile(so_name): 

130 return so_name 

131 error_messages.append(f"No such file: {so_name}") 

132 attachments.append(f' listdir("{lib_dir}"):') 

133 if not os.path.isdir(lib_dir): 

134 attachments.append(" DIRECTORY DOES NOT EXIST") 

135 else: 

136 for node in sorted(os.listdir(lib_dir)): 

137 attachments.append(f" {node}") 

138 return None 

139 

140 

141def _find_dll_using_lib_dir( 

142 lib_dir: str, libname: str, error_messages: list[str], attachments: list[str] 

143) -> str | None: 

144 file_wild = libname + "*.dll" 

145 dll_name = _find_dll_under_dir(lib_dir, file_wild) 

146 if dll_name is not None: 

147 return dll_name 

148 error_messages.append(f"No such file: {file_wild}") 

149 attachments.append(f' listdir("{lib_dir}"):') 

150 for node in sorted(os.listdir(lib_dir)): 

151 attachments.append(f" {node}") 

152 return None 

153 

154 

155class _FindNvidiaDynamicLib: 

156 def __init__(self, libname: str): 

157 self.libname = libname 

158 if IS_WINDOWS: 

159 self.lib_searched_for = f"{libname}*.dll" 

160 else: 

161 self.lib_searched_for = f"lib{libname}.so" 

162 self.error_messages: list[str] = [] 

163 self.attachments: list[str] = [] 

164 self.abs_path: str | None = None 

165 

166 def try_site_packages(self) -> str | None: 

167 if IS_WINDOWS: 

168 return _find_dll_using_nvidia_bin_dirs( 

169 self.libname, 

170 self.lib_searched_for, 

171 self.error_messages, 

172 self.attachments, 

173 ) 

174 else: 

175 return _find_so_using_nvidia_lib_dirs( 

176 self.libname, 

177 self.lib_searched_for, 

178 self.error_messages, 

179 self.attachments, 

180 ) 

181 

182 def try_with_conda_prefix(self) -> str | None: 

183 return self._find_using_lib_dir(_find_lib_dir_using_conda_prefix(self.libname)) 

184 

185 def try_with_cuda_home(self) -> str | None: 

186 return self._find_using_lib_dir(_find_lib_dir_using_cuda_home(self.libname)) 

187 

188 def _find_using_lib_dir(self, lib_dir: str | None) -> str | None: 

189 if lib_dir is None: 

190 return None 

191 if IS_WINDOWS: 

192 return _find_dll_using_lib_dir( 

193 lib_dir, 

194 self.libname, 

195 self.error_messages, 

196 self.attachments, 

197 ) 

198 else: 

199 return _find_so_using_lib_dir( 

200 lib_dir, 

201 self.lib_searched_for, 

202 self.error_messages, 

203 self.attachments, 

204 ) 

205 

206 def raise_not_found_error(self) -> None: 

207 err = ", ".join(self.error_messages) 

208 att = "\n".join(self.attachments) 

209 raise DynamicLibNotFoundError(f'Failure finding "{self.lib_searched_for}": {err}\n{att}')