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

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) 

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

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

24 

25if TYPE_CHECKING: (empty)

26 from cuda.pathfinder._headers.header_descriptor import HeaderDescriptor 

27 

28# --------------------------------------------------------------------------- 

29# Data types 

30# --------------------------------------------------------------------------- 

31 

32 

33@dataclass (empty)

34class LocatedHeaderDir: (empty)

35 abs_path: str | None (empty)

36 found_via: str (empty)

37 

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

39 self.abs_path = _abs_norm(self.abs_path) 11 ctx1abucvwxyzAl

40 

41 

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

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

44 

45# --------------------------------------------------------------------------- 

46# Helpers 

47# --------------------------------------------------------------------------- 

48 

49 

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

51 if path: 11 ctx1abucvwxyzAl

52 return os.path.normpath(os.path.abspath(path)) 11 ctx1abucvwxyzAl

53 return None 

54 

55 

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

57 return os.path.isfile(os.path.join(dirpath, basename)) 26 ctx1adbkeufghicvwjxyzAlopqrstn

58 

59 

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

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

62 h_basename = desc.header_basename 26 ctx1adbkeufghicvwjxyzAlopqrstn

63 for rel_dir in desc.anchor_include_rel_dirs: 26 ctx1adbkeufghicvwjxyzAlopqrstn

64 idir = os.path.join(anchor_point, rel_dir) 26 ctx1adbkeufghicvwjxyzAlopqrstn

65 for subdir in platform_include_subdirs(desc): 26 ctx1adbkeufghicvwjxyzAlopqrstn

66 cdir = os.path.join(idir, subdir) 1 ctx1a

67 if _joined_isfile(cdir, h_basename): 1 ctx1a

68 return cdir 1 ctx1a

69 if _joined_isfile(idir, h_basename): 26 ctx1adbkeufghicvwjxyzAlopqrstn

70 return idir 10 ctx1bucvwxyzAl

71 return None 18 ctx1adbkefghicjopqrstn

72 

73 

74# --------------------------------------------------------------------------- 

75# Find steps 

76# --------------------------------------------------------------------------- 

77 

78 

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

80 """Search pip wheel install locations.""" 

81 for sub_dir in desc.site_packages_dirs: 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

82 hdr_dir: str # help mypy 

83 for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))): 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

84 if _joined_isfile(hdr_dir, desc.header_basename): 

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

86 return None 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

87 

88 

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

90 """Search ``$CONDA_PREFIX``.""" 

91 conda_prefix = os.environ.get("CONDA_PREFIX") 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

92 if not conda_prefix: 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

93 return None 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

94 anchor_point = resolve_conda_anchor(desc, conda_prefix) 

95 if anchor_point is None: 

96 return None 

97 found_header_path = _locate_in_anchor_layout(desc, anchor_point) 

98 if found_header_path: 

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

100 return None 

101 

102 

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

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

105 cuda_home = get_cuda_path_or_home() 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

106 if cuda_home is None: 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

107 return None 3 ctx1BCl

108 result = _locate_in_anchor_layout(desc, cuda_home) 25 ctx1adbkeufghicvwjxyzAopqrstn

109 if result is not None: 25 ctx1adbkeufghicvwjxyzAopqrstn

110 return LocatedHeaderDir(abs_path=result, found_via="CUDA_PATH") 10 ctx1abucvwxyzA

111 return None 18 ctx1adbkefghicjopqrstn

112 

113 

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

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

116 

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

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

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

120 layout under that root. 

121 """ 

122 if not desc.use_ctk_root_canary: 21 ctx1adbkefghicjBClopqrstn

123 return None 7 ctx1opqrstn

124 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess("cudart") 14 ctx1adbkefghicjBCl

125 if canary_abs_path is None: 13 ctx1adbkefghicjBl

126 return None 11 ctx1adbefghicjB

127 ctk_root = derive_ctk_root(canary_abs_path) 9 ctx1dkefghijl

128 if ctk_root is None: 9 ctx1dkefghijl

129 return None 

130 result = _locate_in_anchor_layout(desc, ctk_root) 9 ctx1dkefghijl

131 if result is not None: 9 ctx1dkefghijl

132 return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root") 1 ctx1l

133 return None 8 ctx1dkefghij

134 

135 

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

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

138 for pattern in desc.system_install_dirs: 19 ctx1adbkefghicjBopqrstn

139 for hdr_dir in sorted(glob.glob(pattern), reverse=True): 1 ctx1n

140 if _joined_isfile(hdr_dir, desc.header_basename): 

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

142 return None 19 ctx1adbkefghicjBopqrstn

143 

144 

145# --------------------------------------------------------------------------- 

146# Step sequence and cascade runner 

147# --------------------------------------------------------------------------- 

148 

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

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

151 find_in_site_packages, 

152 find_in_conda, 

153 find_in_cuda_path, 

154 find_via_ctk_root_canary, 

155 find_in_system_install_dirs, 

156) 

157 

158 

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

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

161 for step in steps: 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

162 result = step(desc) 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

163 if result is not None: 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

164 return result 11 ctx1abucvwxyzAl

165 return None 19 ctx1adbkefghicjBopqrstn

166 

167 

168# --------------------------------------------------------------------------- 

169# Public API 

170# --------------------------------------------------------------------------- 

171 

172 

173@functools.cache (empty)

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

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

176 

177 Args: 

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

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

180 

181 Returns: 

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

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

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

185 

186 Raises: 

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

188 

189 Search order: 

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

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

192 3. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. 

193 4. **CTK root canary probe** — subprocess canary (descriptors with 

194 ``use_ctk_root_canary=True`` only). 

195 5. **System install directories** — glob patterns from the descriptor. 

196 """ 

197 desc = HEADER_DESCRIPTORS.get(libname) 29 ctx1adbkeufghicvwjxyzBCAlopqrstnD

198 if desc is None: 29 ctx1adbkeufghicvwjxyzBCAlopqrstnD

199 raise RuntimeError(f"UNKNOWN {libname=}") 1 ctx1D

200 return run_find_steps(desc, FIND_STEPS) 28 ctx1adbkeufghicvwjxyzBCAlopqrstn

201 

202 

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

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

205 

206 Args: 

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

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

209 

210 Returns: 

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

212 if the headers cannot be found. 

213 

214 Raises: 

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

216 

217 Search order: 

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

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

220 3. **CUDA Toolkit environment variables** — ``CUDA_PATH`` / ``CUDA_HOME``. 

221 4. **CTK root canary probe** — subprocess canary (descriptors with 

222 ``use_ctk_root_canary=True`` only). 

223 5. **System install directories** — glob patterns from the descriptor. 

224 """ 

225 found = locate_nvidia_header_directory(libname) 27 ctx1adbkeufghicvwjxyzBCopqrstnD

226 return found.abs_path if found else None 25 ctx1adbkeufghicvwjxyzBopqrstn