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

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.ctk_root_canary import CTK_ROOT_CANARY_ANCHOR_LIBNAMES (empty)

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

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

25 

26if TYPE_CHECKING: (empty)

27 from cuda.pathfinder._headers.header_descriptor import HeaderDescriptor 

28 

29# --------------------------------------------------------------------------- 

30# Data types 

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

32 

33 

34@dataclass (empty)

35class LocatedHeaderDir: (empty)

36 abs_path: str | None (empty)

37 found_via: str (empty)

38 

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

40 self.abs_path = _abs_norm(self.abs_path) 11 ctx1lxyzABCDEFi

41 

42 

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

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

45 

46# --------------------------------------------------------------------------- 

47# Helpers 

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

49 

50 

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

52 if path: 11 ctx1lxyzABCDEFi

53 return os.path.normpath(os.path.abspath(path)) 11 ctx1lxyzABCDEFi

54 return None 

55 

56 

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

58 return os.path.isfile(os.path.join(dirpath, basename)) 31 ctx1laxhbycdefzABgCDEFimnopqrstuvwk

59 

60 

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

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

63 h_basename = desc.header_basename 31 ctx1laxhbycdefzABgCDEFimnopqrstuvwk

64 for rel_dir in desc.anchor_include_rel_dirs: 31 ctx1laxhbycdefzABgCDEFimnopqrstuvwk

65 idir = os.path.join(anchor_point, rel_dir) 31 ctx1laxhbycdefzABgCDEFimnopqrstuvwk

66 for subdir in platform_include_subdirs(desc): 31 ctx1laxhbycdefzABgCDEFimnopqrstuvwk

67 cdir = os.path.join(idir, subdir) 1 ctx1l

68 if _joined_isfile(cdir, h_basename): 1 ctx1l

69 return cdir 1 ctx1l

70 if _joined_isfile(idir, h_basename): 30 ctx1axhbycdefzABgCDEFimnopqrstuvwk

71 return idir 10 ctx1xyzABCDEFi

72 return None 20 ctx1ahbcdefgmnopqrstuvwk

73 

74 

75# --------------------------------------------------------------------------- 

76# Find steps 

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

78 

79 

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

81 """Search pip wheel install locations.""" 

82 for sub_dir in desc.site_packages_dirs: 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

83 hdr_dir: str # help mypy 

84 for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))): 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

85 if _joined_isfile(hdr_dir, desc.header_basename): 

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

87 return None 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

88 

89 

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

91 """Search ``$CONDA_PREFIX``.""" 

92 conda_prefix = os.environ.get("CONDA_PREFIX") 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

93 if not conda_prefix: 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

94 return None 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

95 anchor_point = resolve_conda_anchor(desc, conda_prefix) 

96 if anchor_point is None: 

97 return None 

98 found_header_path = _locate_in_anchor_layout(desc, anchor_point) 

99 if found_header_path: 

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

101 return None 

102 

103 

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

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

106 cuda_home = get_cuda_path_or_home() 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

107 if cuda_home is None: 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

108 return None 3 ctx1GHi

109 result = _locate_in_anchor_layout(desc, cuda_home) 30 ctx1laxhbycdefzABgCDEFmnopqrstuvwk

110 if result is not None: 30 ctx1laxhbycdefzABgCDEFmnopqrstuvwk

111 return LocatedHeaderDir(abs_path=result, found_via="CUDA_PATH") 10 ctx1lxyzABCDEF

112 return None 20 ctx1ahbcdefgmnopqrstuvwk

113 

114 

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

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

117 

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

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

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

121 layout under that root. 

122 """ 

123 if not desc.use_ctk_root_canary: 23 ctx1ahbcdefgGHimnopqrstuvwk

124 return None 12 ctx1mnopqrstuvwk

125 canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess(CTK_ROOT_CANARY_ANCHOR_LIBNAMES[0]) 11 ctx1ahbcdefgGHi

126 if canary_abs_path is None: 10 ctx1ahbcdefgGi

127 return None 8 ctx1abcdefgG

128 ctk_root = derive_ctk_root(canary_abs_path) 9 ctx1ahbcdefgi

129 if ctk_root is None: 9 ctx1ahbcdefgi

130 return None 

131 result = _locate_in_anchor_layout(desc, ctk_root) 9 ctx1ahbcdefgi

132 if result is not None: 9 ctx1ahbcdefgi

133 return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root") 1 ctx1i

134 return None 8 ctx1ahbcdefg

135 

136 

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

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

139 for pattern in desc.system_install_dirs: 21 ctx1ahbcdefgGmnopqrstuvwk

140 for hdr_dir in sorted(glob.glob(pattern), reverse=True): 1 ctx1k

141 if _joined_isfile(hdr_dir, desc.header_basename): 

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

143 return None 21 ctx1ahbcdefgGmnopqrstuvwk

144 

145 

146# --------------------------------------------------------------------------- 

147# Step sequence and cascade runner 

148# --------------------------------------------------------------------------- 

149 

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

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

152 find_in_site_packages, 

153 find_in_conda, 

154 find_in_cuda_path, 

155 find_via_ctk_root_canary, 

156 find_in_system_install_dirs, 

157) 

158 

159 

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

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

162 for step in steps: 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

163 result = step(desc) 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

164 if result is not None: 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

165 return result 11 ctx1lxyzABCDEFi

166 return None 21 ctx1ahbcdefgGmnopqrstuvwk

167 

168 

169# --------------------------------------------------------------------------- 

170# Public API 

171# --------------------------------------------------------------------------- 

172 

173 

174@functools.cache (empty)

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

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

177 

178 Args: 

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

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

181 

182 Returns: 

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

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

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

186 

187 Raises: 

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

189 

190 Search order: 

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

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

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

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

195 ``use_ctk_root_canary=True`` only). 

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

197 """ 

198 desc = HEADER_DESCRIPTORS.get(libname) 34 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwkI

199 if desc is None: 34 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwkI

200 raise RuntimeError(f"UNKNOWN {libname=}") 1 ctx1I

201 return run_find_steps(desc, FIND_STEPS) 33 ctx1laxhbycdefzABgCDEGHFimnopqrstuvwk

202 

203 

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

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

206 

207 Args: 

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

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

210 

211 Returns: 

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

213 if the headers cannot be found. 

214 

215 Raises: 

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

217 

218 Search order: 

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

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

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

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

223 ``use_ctk_root_canary=True`` only). 

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

225 """ 

226 found = locate_nvidia_header_directory(libname) 32 ctx1laxhbycdefzABgCDEGHmnopqrstuvwkI

227 return found.abs_path if found else None 30 ctx1laxhbycdefzABgCDEGmnopqrstuvwk