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

1# SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 collections.abc import Callable (empty)

10from dataclasses import dataclass (empty)

11from typing import TYPE_CHECKING (empty)

12 

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

14 _resolve_system_loaded_abs_path_in_subprocess, 

15) 

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

17from cuda.pathfinder._headers.header_descriptor import ( (empty)

18 HEADER_DESCRIPTORS, 

19 platform_include_subdirs, 

20 resolve_conda_anchor, 

21 system_install_dir_patterns, 

22) 

23from cuda.pathfinder._utils.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES (empty)

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

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

26from cuda.pathfinder._utils.path_sort import numeric_aware_path_sort_key (empty)

27 

28if TYPE_CHECKING: (empty)

29 from cuda.pathfinder._headers.header_descriptor import HeaderDescriptor 

30 

31# --------------------------------------------------------------------------- 

32# Data types 

33# --------------------------------------------------------------------------- 

34 

35 

36@dataclass (empty)

37class LocatedHeaderDir: (empty)

38 abs_path: str | None (empty)

39 found_via: str (empty)

40 

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

42 self.abs_path = _abs_norm(self.abs_path) 19 ctx1MNOPQRnzABCDEFGHkIJ

43 

44 

45#: Type alias for a header find step callable. 

46HeaderFindStep = Callable[["HeaderDescriptor"], LocatedHeaderDir | None] (empty)

47 

48# --------------------------------------------------------------------------- 

49# Helpers 

50# --------------------------------------------------------------------------- 

51 

52 

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

54 if path: 19 ctx1MNOPQRnzABCDEFGHkIJ

55 return os.path.normpath(os.path.abspath(path)) 19 ctx1MNOPQRnzABCDEFGHkIJ

56 return None 

57 

58 

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

60 return os.path.isfile(os.path.join(dirpath, basename)) 41 ctx1MNOPQRnazhbAcdefBCDgEFGHkoipqrstuvwxyjmIJ

61 

62 

63def _locate_in_anchor_layout(desc: HeaderDescriptor, anchor_point: str) -> str | None: (empty)

64 """Search for a header under *anchor_point* using the descriptor's layout fields.""" 

65 h_basename = desc.header_basename 35 ctx1nazhbAcdefBCDgEFGHkoipqrstuvwxyjmIJ

66 for rel_dir in desc.anchor_include_rel_dirs: 35 ctx1nazhbAcdefBCDgEFGHkoipqrstuvwxyjmIJ

67 idir = os.path.join(anchor_point, rel_dir) 35 ctx1nazhbAcdefBCDgEFGHkoipqrstuvwxyjmIJ

68 for subdir in platform_include_subdirs(desc): 35 ctx1nazhbAcdefBCDgEFGHkoipqrstuvwxyjmIJ

69 cdir = os.path.join(idir, subdir) 1 ctx1n

70 if _joined_isfile(cdir, h_basename): 1 ctx1n

71 return cdir 1 ctx1n

72 if _joined_isfile(idir, h_basename): 34 ctx1azhbAcdefBCDgEFGHkoipqrstuvwxyjmIJ

73 return idir 12 ctx1zABCDEFGHkIJ

74 return None 22 ctx1ahbcdefgoipqrstuvwxyjm

75 

76 

77# --------------------------------------------------------------------------- 

78# Find steps 

79# --------------------------------------------------------------------------- 

80 

81 

82def find_in_site_packages(desc: HeaderDescriptor) -> LocatedHeaderDir | None: (empty)

83 """Search pip wheel install locations.""" 

84 for sub_dir in desc.site_packages_dirs: 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

85 hdr_dir: str # help mypy 

86 for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))): 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

87 if _joined_isfile(hdr_dir, desc.header_basename): 

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

89 return None 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

90 

91 

92def find_in_conda(desc: HeaderDescriptor) -> LocatedHeaderDir | None: (empty)

93 """Search ``$CONDA_PREFIX``.""" 

94 conda_prefix = os.environ.get("CONDA_PREFIX") 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

95 if not conda_prefix: 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

96 return None 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

97 anchor_point = resolve_conda_anchor(desc, conda_prefix) 

98 if anchor_point is None: 

99 return None 

100 found_header_path = _locate_in_anchor_layout(desc, anchor_point) 

101 if found_header_path: 

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

103 return None 

104 

105 

106def find_in_product_roots(desc: HeaderDescriptor) -> LocatedHeaderDir | None: (empty)

107 """Search roots supplied through product-specific environment variables.""" 

108 for env_var in desc.product_root_env_vars: 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

109 root = os.environ.get(env_var) 4 ctx1ijIJ

110 if not root: 4 ctx1ijIJ

111 continue 2 ctx1ij

112 result = _locate_in_anchor_layout(desc, root) 2 ctx1IJ

113 if result is not None: 2 ctx1IJ

114 return LocatedHeaderDir(abs_path=result, found_via=env_var) 2 ctx1IJ

115 return None 35 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjm

116 

117 

118def find_in_cuda_path(desc: HeaderDescriptor) -> LocatedHeaderDir | None: (empty)

119 """Search ``$CUDA_PATH`` / ``$CUDA_HOME``.""" 

120 cuda_home = get_cuda_path_or_home() 35 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjm

121 if cuda_home is None: 35 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjm

122 return None 3 ctx1KLk

123 result = _locate_in_anchor_layout(desc, cuda_home) 32 ctx1nazhbAcdefBCDgEFGHoipqrstuvwxyjm

124 if result is not None: 32 ctx1nazhbAcdefBCDgEFGHoipqrstuvwxyjm

125 return LocatedHeaderDir(abs_path=result, found_via="CUDA_PATH") 10 ctx1nzABCDEFGH

126 return None 22 ctx1ahbcdefgoipqrstuvwxyjm

127 

128 

129def find_via_ctk_root_canary(desc: HeaderDescriptor) -> LocatedHeaderDir | None: (empty)

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

131 

132 Skips immediately if the descriptor does not opt in (``use_ctk_root_canary``). 

133 Otherwise, system-loads ``cudart`` in a fully isolated Python subprocess, derives 

134 CTK root from the resolved library path, and searches the expected include 

135 layout under that root. 

136 """ 

137 if not desc.use_ctk_root_canary: 25 ctx1ahbcdefgKLkoipqrstuvwxyjm

138 return None 14 ctx1oipqrstuvwxyjm

139 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(CTK_ROOT_CANARY_ANCHOR_LIBNAMES[0]) 11 ctx1ahbcdefgKLk

140 if canary_abs_path is None: 10 ctx1ahbcdefgKk

141 return None 8 ctx1abcdefgK

142 ctk_root = derive_ctk_root(canary_abs_path) 9 ctx1ahbcdefgk

143 if ctk_root is None: 9 ctx1ahbcdefgk

144 return None 

145 result = _locate_in_anchor_layout(desc, ctk_root) 9 ctx1ahbcdefgk

146 if result is not None: 9 ctx1ahbcdefgk

147 return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root") 1 ctx1k

148 return None 8 ctx1ahbcdefg

149 

150 

151def find_in_system_install_dirs(desc: HeaderDescriptor) -> LocatedHeaderDir | None: (empty)

152 """Search system install directories (glob patterns).""" 

153 for pattern in system_install_dir_patterns(desc): 29 ctx1MNOPQRahbcdefgKoipqrstuvwxyjm

154 for hdr_dir in sorted(glob.glob(pattern), key=numeric_aware_path_sort_key, reverse=True): 9 ctx1MNOPQRijm

155 if _joined_isfile(hdr_dir, desc.header_basename): 8 ctx1MNOPQRij

156 return LocatedHeaderDir(abs_path=hdr_dir, found_via="supported_install_dir") 6 ctx1MNOPQR

157 return None 23 ctx1ahbcdefgKoipqrstuvwxyjm

158 

159 

160# --------------------------------------------------------------------------- 

161# Step sequence and cascade runner 

162# --------------------------------------------------------------------------- 

163 

164#: Unified find steps — each step self-gates based on descriptor fields. 

165FIND_STEPS: tuple[HeaderFindStep, ...] = ( (empty)

166 find_in_site_packages, 

167 find_in_conda, 

168 find_in_product_roots, 

169 find_in_cuda_path, 

170 find_via_ctk_root_canary, 

171 find_in_system_install_dirs, 

172) 

173 

174 

175def run_find_steps(desc: HeaderDescriptor, steps: tuple[HeaderFindStep, ...]) -> LocatedHeaderDir | None: (empty)

176 """Run find steps in order, returning the first hit.""" 

177 for step in steps: 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

178 result = step(desc) 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

179 if result is not None: 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

180 return result 13 ctx1nzABCDEFGHkIJ

181 return None 23 ctx1ahbcdefgKoipqrstuvwxyjm

182 

183 

184# --------------------------------------------------------------------------- 

185# Public API 

186# --------------------------------------------------------------------------- 

187 

188 

189@functools.cache (empty)

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

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

192 

193 Args: 

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

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

196 

197 Returns: 

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

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

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

201 

202 Raises: 

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

204 

205 Search order: 

206 1. **NVIDIA Python wheels** — site-packages directories from the descriptor. 

207 2. **Conda environments** — platform-specific conda include layouts. 

208 3. **Product environment variables** — for example, ``CUDNN_PATH``. 

209 4. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. 

210 5. **CTK root canary probe** — subprocess canary (descriptors with 

211 ``use_ctk_root_canary=True`` only). 

212 6. **System install directories** — glob patterns from the descriptor. 

213 """ 

214 desc = HEADER_DESCRIPTORS.get(libname) 38 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJS

215 if desc is None: 38 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJS

216 raise RuntimeError(f"UNKNOWN {libname=}") 1 ctx1S

217 return run_find_steps(desc, FIND_STEPS) 37 ctx1nazhbAcdefBCDgEFGKLHkoipqrstuvwxyjmIJ

218 

219 

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

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

222 

223 Args: 

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

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

226 

227 Returns: 

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

229 if the headers cannot be found. 

230 

231 Raises: 

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

233 

234 Search order: 

235 1. **NVIDIA Python wheels** — site-packages directories from the descriptor. 

236 2. **Conda environments** — platform-specific conda include layouts. 

237 3. **Product environment variables** — for example, ``CUDNN_PATH``. 

238 4. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. 

239 5. **CTK root canary probe** — subprocess canary (descriptors with 

240 ``use_ctk_root_canary=True`` only). 

241 6. **System install directories** — glob patterns from the descriptor. 

242 """ 

243 found = locate_nvidia_header_directory(libname) 34 ctx1nazhbAcdefBCDgEFGKLoipqrstuvwxyjmS

244 return found.abs_path if found else None 32 ctx1nazhbAcdefBCDgEFGKoipqrstuvwxyjm