Coverage for cuda / pathfinder / _headers / find_nvidia_headers.py: 74.34%

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

2# SPDX-License-Identifier: Apache-2.0 

3 

4from __future__ import annotations (empty)

5 

6import functools (empty)

7import glob (empty)

8import os (empty)

9from dataclasses import dataclass (empty)

10 

11from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import ( (empty)

12 _resolve_system_loaded_abs_path_in_subprocess, 

13) 

14from cuda.pathfinder._dynamic_libs.search_steps import derive_ctk_root (empty)

15from cuda.pathfinder._headers import supported_nvidia_headers (empty)

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

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

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

19 

20 

21@dataclass (empty)

22class LocatedHeaderDir: (empty)

23 abs_path: str | None (empty)

24 found_via: str (empty)

25 

26 def __post_init__(self) -> None: (empty)

27 self.abs_path = _abs_norm(self.abs_path) 10 ctx1aklmnopjrb

28 

29 

30def _abs_norm(path: str | None) -> str | None: (empty)

31 if path: 10 ctx1aklmnopjrb

32 return os.path.normpath(os.path.abspath(path)) 10 ctx1aklmnopjrb

33 return None 

34 

35 

36def _joined_isfile(dirpath: str, basename: str) -> bool: (empty)

37 return os.path.isfile(os.path.join(dirpath, basename)) 17 ctx1ackdlefghmnoipjrb

38 

39 

40def _locate_under_site_packages(sub_dir: str, h_basename: str) -> LocatedHeaderDir | None: (empty)

41 # Installed from a wheel 

42 hdr_dir: str # help mypy 

43 for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))): 22 ctx1ackdlefghmnoipjsurbvwt

44 if _joined_isfile(hdr_dir, h_basename): 

45 return LocatedHeaderDir(abs_path=hdr_dir, found_via="site-packages") 

46 return None 22 ctx1ackdlefghmnoipjsurbvwt

47 

48 

49def _locate_based_on_ctk_layout(libname: str, h_basename: str, anchor_point: str) -> str | None: (empty)

50 parts = [anchor_point] 17 ctx1ackdlefghmnoipjrb

51 if libname == "nvvm": 17 ctx1ackdlefghmnoipjrb

52 parts.append(libname) 1 ctx1j

53 parts.append("include") 17 ctx1ackdlefghmnoipjrb

54 idir = os.path.join(*parts) 17 ctx1ackdlefghmnoipjrb

55 if libname == "cccl": 17 ctx1ackdlefghmnoipjrb

56 if IS_WINDOWS: 1 ctx1a

57 cdir_ctk12 = os.path.join(idir, "targets", "x64") # conda has this anomaly 1 ctx1a

58 cdir_ctk13 = os.path.join(cdir_ctk12, "cccl") 1 ctx1a

59 if _joined_isfile(cdir_ctk13, h_basename): 1 ctx1a

60 return cdir_ctk13 

61 if _joined_isfile(cdir_ctk12, h_basename): 1 ctx1a

62 return cdir_ctk12 

63 cdir = os.path.join(idir, "cccl") # CTK 13 1 ctx1a

64 if _joined_isfile(cdir, h_basename): 1 ctx1a

65 return cdir 1 ctx1a

66 if _joined_isfile(idir, h_basename): 16 ctx1ckdlefghmnoipjrb

67 return idir 9 ctx1klmnopjrb

68 return None 7 ctx1cdefghi

69 

70 

71def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool) -> LocatedHeaderDir | None: (empty)

72 conda_prefix = os.environ.get("CONDA_PREFIX") 22 ctx1ackdlefghmnoipjsurbvwt

73 if not conda_prefix: 22 ctx1ackdlefghmnoipjsurbvwt

74 return None 22 ctx1ackdlefghmnoipjsurbvwt

75 if IS_WINDOWS: 

76 anchor_point = os.path.join(conda_prefix, "Library") 

77 if not os.path.isdir(anchor_point): 

78 return None 

79 else: 

80 if ctk_layout: 

81 targets_include_path = glob.glob(os.path.join(conda_prefix, "targets", "*", "include")) 

82 if not targets_include_path: 

83 return None 

84 if len(targets_include_path) != 1: 

85 # Conda does not support multiple architectures. 

86 # QUESTION(PR#956): Do we want to issue a warning? 

87 return None 

88 include_path = targets_include_path[0] 

89 else: 

90 include_path = os.path.join(conda_prefix, "include") 

91 anchor_point = os.path.dirname(include_path) 

92 found_header_path = _locate_based_on_ctk_layout(libname, h_basename, anchor_point) 

93 if found_header_path: 

94 return LocatedHeaderDir(abs_path=found_header_path, found_via="conda") 

95 return None 

96 

97 

98def _find_ctk_header_directory_via_canary(libname: str, h_basename: str) -> str | None: (empty)

99 """Try CTK header lookup via CTK-root canary probing. 

100 

101 Uses the same canary as dynamic-library CTK-root discovery: system-load 

102 ``cudart`` in a spawned child process, derive CTK root from the resolved 

103 absolute library path, then search the expected CTK include layout under 

104 that root. 

105 """ 

106 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess("cudart") 10 ctx1cdefghisub

107 if canary_abs_path is None: 9 ctx1cdefghisb

108 return None 8 ctx1cdefghis

109 ctk_root = derive_ctk_root(canary_abs_path) 8 ctx1cdefghib

110 if ctk_root is None: 8 ctx1cdefghib

111 return None 

112 return _locate_based_on_ctk_layout(libname, h_basename, ctk_root) 8 ctx1cdefghib

113 

114 

115def _find_ctk_header_directory(libname: str) -> LocatedHeaderDir | None: (empty)

116 h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_CTK[libname] 19 ctx1ackdlefghmnoipjsurb

117 candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK[libname] 19 ctx1ackdlefghmnoipjsurb

118 

119 for cdir in candidate_dirs: 19 ctx1ackdlefghmnoipjsurb

120 if hdr_dir := _locate_under_site_packages(cdir, h_basename): 19 ctx1ackdlefghmnoipjsurb

121 return hdr_dir 

122 

123 if hdr_dir := _find_based_on_conda_layout(libname, h_basename, True): 19 ctx1ackdlefghmnoipjsurb

124 return hdr_dir 

125 

126 cuda_home = get_cuda_home_or_path() 19 ctx1ackdlefghmnoipjsurb

127 if cuda_home and (result := _locate_based_on_ctk_layout(libname, h_basename, cuda_home)): 19 ctx1ackdlefghmnoipjsurb

128 return LocatedHeaderDir(abs_path=result, found_via="CUDA_HOME") 9 ctx1aklmnopjr

129 

130 if result := _find_ctk_header_directory_via_canary(libname, h_basename): 10 ctx1cdefghisub

131 return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root") 1 ctx1b

132 

133 return None 8 ctx1cdefghis

134 

135 

136@functools.cache (empty)

137def locate_nvidia_header_directory(libname: str) -> LocatedHeaderDir | None: (empty)

138 """Locate the header directory for a supported NVIDIA library. 

139 

140 Args: 

141 libname (str): The short name of the library whose headers are needed 

142 (e.g., ``"nvrtc"``, ``"cusolver"``, ``"nvshmem"``). 

143 

144 Returns: 

145 LocatedHeaderDir or None: A LocatedHeaderDir object containing the absolute path 

146 to the discovered header directory and information about where it was found, 

147 or ``None`` if the headers cannot be found. 

148 

149 Raises: 

150 RuntimeError: If ``libname`` is not in the supported set. 

151 

152 Search order: 

153 1. **NVIDIA Python wheels** 

154 

155 - Scan installed distributions (``site-packages``) for header layouts 

156 shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvrtc]``). 

157 

158 2. **Conda environments** 

159 

160 - Check Conda-style installation prefixes, which use platform-specific 

161 include directory layouts. 

162 

163 3. **CUDA Toolkit environment variables** 

164 

165 - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order). 

166 

167 4. **CTK root canary probe** 

168 

169 - Probe a system-loaded ``cudart`` in a spawned child process, 

170 derive the CTK root from the resolved library path, then search 

171 CTK include layout under that root. 

172 """ 

173 

174 if libname in supported_nvidia_headers.SUPPORTED_HEADERS_CTK: 23 ctx1ackdlefghmnoipjsurbvwtx

175 return _find_ctk_header_directory(libname) 19 ctx1ackdlefghmnoipjsurb

176 

177 h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_NON_CTK.get(libname) 4 ctx1vwtx

178 if h_basename is None: 4 ctx1vwtx

179 raise RuntimeError(f"UNKNOWN {libname=}") 1 ctx1x

180 

181 candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_NON_CTK.get(libname, []) 3 ctx1vwt

182 

183 for cdir in candidate_dirs: 3 ctx1vwt

184 if found_hdr := _locate_under_site_packages(cdir, h_basename): 3 ctx1vwt

185 return found_hdr 

186 

187 if found_hdr := _find_based_on_conda_layout(libname, h_basename, False): 3 ctx1vwt

188 return found_hdr 

189 

190 # Fall back to system install directories 

191 candidate_dirs = supported_nvidia_headers.SUPPORTED_INSTALL_DIRS_NON_CTK.get(libname, []) 3 ctx1vwt

192 for cdir in candidate_dirs: 3 ctx1vwt

193 for hdr_dir in sorted(glob.glob(cdir), reverse=True): 1 ctx1t

194 if _joined_isfile(hdr_dir, h_basename): 

195 # For system installs, we don't have a clear found_via, so use "system" 

196 return LocatedHeaderDir(abs_path=hdr_dir, found_via="supported_install_dir") 

197 return None 3 ctx1vwt

198 

199 

200def find_nvidia_header_directory(libname: str) -> str | None: (empty)

201 """Locate the header directory for a supported NVIDIA library. 

202 

203 Args: 

204 libname (str): The short name of the library whose headers are needed 

205 (e.g., ``"nvrtc"``, ``"cusolver"``, ``"nvshmem"``). 

206 

207 Returns: 

208 str or None: Absolute path to the discovered header directory, or ``None`` 

209 if the headers cannot be found. 

210 

211 Raises: 

212 RuntimeError: If ``libname`` is not in the supported set. 

213 

214 Search order: 

215 1. **NVIDIA Python wheels** 

216 

217 - Scan installed distributions (``site-packages``) for header layouts 

218 shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvrtc]``). 

219 

220 2. **Conda environments** 

221 

222 - Check Conda-style installation prefixes, which use platform-specific 

223 include directory layouts. 

224 

225 3. **CUDA Toolkit environment variables** 

226 

227 - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order). 

228 

229 4. **CTK root canary probe** 

230 

231 - Probe a system-loaded ``cudart`` in a spawned child process, 

232 derive the CTK root from the resolved library path, then search 

233 CTK include layout under that root. 

234 """ 

235 found = locate_nvidia_header_directory(libname) 21 ctx1ackdlefghmnoipjsuvwtx

236 return found.abs_path if found else None 19 ctx1ackdlefghmnoipjsvwt