Skip to content

Commit 3c96fcc

Browse files
committed
refactor(pathfinder): introduce header descriptor catalog
Mirror the dynamic-library descriptor-catalog pattern for headers. All per-header metadata (basename, site-packages dirs, anchor layout, platform availability, system install dirs) is now authored once in HeaderDescriptorSpec entries in header_descriptor_catalog.py. - header_descriptor.py provides a name-keyed HEADER_DESCRIPTORS registry - supported_nvidia_headers.py derives its legacy tables from the catalog - find_nvidia_headers.py uses descriptors directly, eliminating name-based if/else branching (nvvm, cccl special cases) in favor of descriptor fields (anchor_include_rel_dirs, include_subdirs, etc.) Public API (LocatedHeaderDir, find_nvidia_header_directory, locate_nvidia_header_directory, SUPPORTED_HEADERS_CTK) is unchanged. Made-with: Cursor
1 parent a772555 commit 3c96fcc

4 files changed

Lines changed: 293 additions & 118 deletions

File tree

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

Lines changed: 67 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
# SPDX-License-Identifier: Apache-2.0
33

44
from __future__ import annotations
@@ -7,16 +7,20 @@
77
import glob
88
import os
99
from dataclasses import dataclass
10+
from typing import TYPE_CHECKING
1011

1112
from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
1213
_resolve_system_loaded_abs_path_in_subprocess,
1314
)
1415
from cuda.pathfinder._dynamic_libs.search_steps import derive_ctk_root
15-
from cuda.pathfinder._headers import supported_nvidia_headers
16+
from cuda.pathfinder._headers.header_descriptor import HEADER_DESCRIPTORS
1617
from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
1718
from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
1819
from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
1920

21+
if TYPE_CHECKING:
22+
from cuda.pathfinder._headers.header_descriptor import HeaderDescriptor
23+
2024

2125
@dataclass
2226
class LocatedHeaderDir:
@@ -38,37 +42,33 @@ def _joined_isfile(dirpath: str, basename: str) -> bool:
3842

3943

4044
def _locate_under_site_packages(sub_dir: str, h_basename: str) -> LocatedHeaderDir | None:
41-
# Installed from a wheel
4245
hdr_dir: str # help mypy
4346
for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))):
4447
if _joined_isfile(hdr_dir, h_basename):
4548
return LocatedHeaderDir(abs_path=hdr_dir, found_via="site-packages")
4649
return None
4750

4851

49-
def _locate_based_on_ctk_layout(libname: str, h_basename: str, anchor_point: str) -> str | None:
50-
parts = [anchor_point]
51-
if libname == "nvvm":
52-
parts.append(libname)
53-
parts.append("include")
54-
idir = os.path.join(*parts)
55-
if libname == "cccl":
52+
def _locate_in_anchor_layout(desc: HeaderDescriptor, anchor_point: str) -> str | None:
53+
"""Search for a header under *anchor_point* using the descriptor's layout fields."""
54+
h_basename = desc.header_basename
55+
for rel_dir in desc.anchor_include_rel_dirs:
56+
idir = os.path.join(anchor_point, rel_dir)
5657
if IS_WINDOWS:
57-
cdir_ctk12 = os.path.join(idir, "targets", "x64") # conda has this anomaly
58-
cdir_ctk13 = os.path.join(cdir_ctk12, "cccl")
59-
if _joined_isfile(cdir_ctk13, h_basename):
60-
return cdir_ctk13
61-
if _joined_isfile(cdir_ctk12, h_basename):
62-
return cdir_ctk12
63-
cdir = os.path.join(idir, "cccl") # CTK 13
64-
if _joined_isfile(cdir, h_basename):
65-
return cdir
66-
if _joined_isfile(idir, h_basename):
67-
return idir
58+
for subdir in desc.include_subdirs_windows:
59+
cdir = os.path.join(idir, subdir)
60+
if _joined_isfile(cdir, h_basename):
61+
return cdir
62+
for subdir in desc.include_subdirs:
63+
cdir = os.path.join(idir, subdir)
64+
if _joined_isfile(cdir, h_basename):
65+
return cdir
66+
if _joined_isfile(idir, h_basename):
67+
return idir
6868
return None
6969

7070

71-
def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool) -> LocatedHeaderDir | None:
71+
def _find_in_conda(desc: HeaderDescriptor) -> LocatedHeaderDir | None:
7272
conda_prefix = os.environ.get("CONDA_PREFIX")
7373
if not conda_prefix:
7474
return None
@@ -77,25 +77,23 @@ def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool)
7777
if not os.path.isdir(anchor_point):
7878
return None
7979
else:
80-
if ctk_layout:
80+
if desc.packaged_with == "ctk":
8181
targets_include_path = glob.glob(os.path.join(conda_prefix, "targets", "*", "include"))
8282
if not targets_include_path:
8383
return None
8484
if len(targets_include_path) != 1:
85-
# Conda does not support multiple architectures.
86-
# QUESTION(PR#956): Do we want to issue a warning?
8785
return None
8886
include_path = targets_include_path[0]
8987
else:
9088
include_path = os.path.join(conda_prefix, "include")
9189
anchor_point = os.path.dirname(include_path)
92-
found_header_path = _locate_based_on_ctk_layout(libname, h_basename, anchor_point)
90+
found_header_path = _locate_in_anchor_layout(desc, anchor_point)
9391
if found_header_path:
9492
return LocatedHeaderDir(abs_path=found_header_path, found_via="conda")
9593
return None
9694

9795

98-
def _find_ctk_header_directory_via_canary(libname: str, h_basename: str) -> str | None:
96+
def _find_via_ctk_root_canary(desc: HeaderDescriptor) -> str | None:
9997
"""Try CTK header lookup via CTK-root canary probing.
10098
10199
Uses the same canary as dynamic-library CTK-root discovery: system-load
@@ -109,26 +107,43 @@ def _find_ctk_header_directory_via_canary(libname: str, h_basename: str) -> str
109107
ctk_root = derive_ctk_root(canary_abs_path)
110108
if ctk_root is None:
111109
return None
112-
return _locate_based_on_ctk_layout(libname, h_basename, ctk_root)
110+
return _locate_in_anchor_layout(desc, ctk_root)
111+
112+
113+
def _find_in_system_install_dirs(desc: HeaderDescriptor) -> LocatedHeaderDir | None:
114+
"""Search system install directories (glob patterns) for non-CTK headers."""
115+
for pattern in desc.system_install_dirs:
116+
for hdr_dir in sorted(glob.glob(pattern), reverse=True):
117+
if _joined_isfile(hdr_dir, desc.header_basename):
118+
return LocatedHeaderDir(abs_path=hdr_dir, found_via="supported_install_dir")
119+
return None
113120

114121

115-
def _find_ctk_header_directory(libname: str) -> LocatedHeaderDir | None:
116-
h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_CTK[libname]
117-
candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK[libname]
122+
def _locate_header_directory(desc: HeaderDescriptor) -> LocatedHeaderDir | None:
123+
"""Unified header search using a descriptor."""
124+
h_basename = desc.header_basename
118125

119-
for cdir in candidate_dirs:
126+
# 1. Site-packages (pip wheels)
127+
for cdir in desc.site_packages_dirs:
120128
if hdr_dir := _locate_under_site_packages(cdir, h_basename):
121129
return hdr_dir
122130

123-
if hdr_dir := _find_based_on_conda_layout(libname, h_basename, True):
131+
# 2. Conda
132+
if hdr_dir := _find_in_conda(desc):
124133
return hdr_dir
125134

126-
cuda_home = get_cuda_home_or_path()
127-
if cuda_home and (result := _locate_based_on_ctk_layout(libname, h_basename, cuda_home)):
128-
return LocatedHeaderDir(abs_path=result, found_via="CUDA_HOME")
135+
# CTK-only steps: CUDA_HOME and canary probe
136+
if desc.packaged_with == "ctk":
137+
cuda_home = get_cuda_home_or_path()
138+
if cuda_home and (result := _locate_in_anchor_layout(desc, cuda_home)):
139+
return LocatedHeaderDir(abs_path=result, found_via="CUDA_HOME")
129140

130-
if result := _find_ctk_header_directory_via_canary(libname, h_basename):
131-
return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root")
141+
if result := _find_via_ctk_root_canary(desc):
142+
return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root")
143+
144+
# Non-CTK fallback: system install directories
145+
if desc.system_install_dirs and (hdr_dir := _find_in_system_install_dirs(desc)):
146+
return hdr_dir
132147

133148
return None
134149

@@ -160,41 +175,24 @@ def locate_nvidia_header_directory(libname: str) -> LocatedHeaderDir | None:
160175
- Check Conda-style installation prefixes, which use platform-specific
161176
include directory layouts.
162177
163-
3. **CUDA Toolkit environment variables**
178+
3. **CUDA Toolkit environment variables** *(CTK headers only)*
164179
165180
- Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
166181
167-
4. **CTK root canary probe**
182+
4. **CTK root canary probe** *(CTK headers only)*
168183
169184
- Probe a system-loaded ``cudart`` in a spawned child process,
170185
derive the CTK root from the resolved library path, then search
171186
CTK include layout under that root.
172-
"""
173187
174-
if libname in supported_nvidia_headers.SUPPORTED_HEADERS_CTK:
175-
return _find_ctk_header_directory(libname)
188+
5. **System install directories** *(non-CTK headers only)*
176189
177-
h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_NON_CTK.get(libname)
178-
if h_basename is None:
190+
- Search well-known system install paths (e.g., ``/usr/include/nvshmem_*``).
191+
"""
192+
desc = HEADER_DESCRIPTORS.get(libname)
193+
if desc is None:
179194
raise RuntimeError(f"UNKNOWN {libname=}")
180-
181-
candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_NON_CTK.get(libname, [])
182-
183-
for cdir in candidate_dirs:
184-
if found_hdr := _locate_under_site_packages(cdir, h_basename):
185-
return found_hdr
186-
187-
if found_hdr := _find_based_on_conda_layout(libname, h_basename, False):
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, [])
192-
for cdir in candidate_dirs:
193-
for hdr_dir in sorted(glob.glob(cdir), reverse=True):
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
195+
return _locate_header_directory(desc)
198196

199197

200198
def find_nvidia_header_directory(libname: str) -> str | None:
@@ -222,15 +220,19 @@ def find_nvidia_header_directory(libname: str) -> str | None:
222220
- Check Conda-style installation prefixes, which use platform-specific
223221
include directory layouts.
224222
225-
3. **CUDA Toolkit environment variables**
223+
3. **CUDA Toolkit environment variables** *(CTK headers only)*
226224
227225
- Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
228226
229-
4. **CTK root canary probe**
227+
4. **CTK root canary probe** *(CTK headers only)*
230228
231229
- Probe a system-loaded ``cudart`` in a spawned child process,
232230
derive the CTK root from the resolved library path, then search
233231
CTK include layout under that root.
232+
233+
5. **System install directories** *(non-CTK headers only)*
234+
235+
- Search well-known system install paths (e.g., ``/usr/include/nvshmem_*``).
234236
"""
235237
found = locate_nvidia_header_directory(libname)
236238
return found.abs_path if found else None
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
"""Per-header descriptor and registry.
5+
6+
The canonical authored data lives in :mod:`header_descriptor_catalog`. This
7+
module provides a name-keyed registry consumed by the runtime search path.
8+
"""
9+
10+
from __future__ import annotations
11+
12+
from typing import TypeAlias
13+
14+
from cuda.pathfinder._headers.header_descriptor_catalog import (
15+
HEADER_DESCRIPTOR_CATALOG,
16+
HeaderDescriptorSpec,
17+
)
18+
19+
HeaderDescriptor: TypeAlias = HeaderDescriptorSpec
20+
21+
#: Canonical registry of all known header libraries.
22+
HEADER_DESCRIPTORS: dict[str, HeaderDescriptor] = {desc.name: desc for desc in HEADER_DESCRIPTOR_CATALOG}

0 commit comments

Comments
 (0)