fix(pathfinder): use CTK canary fallback for header discovery

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -8,6 +8,10 @@
 import os
 from dataclasses import dataclass
 
+from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
+    _resolve_system_loaded_abs_path_in_subprocess,
+)
+from cuda.pathfinder._dynamic_libs.search_steps import derive_ctk_root
 from cuda.pathfinder._headers import supported_nvidia_headers
 from cuda.pathfinder._utils.env_vars import get_cuda_home_or_path
 from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
@@ -91,6 +95,23 @@ def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool)
     return None
 
 
+def _find_ctk_header_directory_via_canary(libname: str, h_basename: str) -> str | None:
+    """Try CTK header lookup via CTK-root canary probing.
+
+    Uses the same canary as dynamic-library CTK-root discovery: system-load
+    ``cudart`` in a spawned child process, derive CTK root from the resolved
+    absolute library path, then search the expected CTK include layout under
+    that root.
+    """
+    canary_abs_path = _resolve_system_loaded_abs_path_in_subprocess("cudart")
+    if canary_abs_path is None:
+        return None
+    ctk_root = derive_ctk_root(canary_abs_path)
+    if ctk_root is None:
+        return None
+    return _locate_based_on_ctk_layout(libname, h_basename, ctk_root)
+
+
 def _find_ctk_header_directory(libname: str) -> LocatedHeaderDir | None:
     h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_CTK[libname]
     candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK[libname]
@@ -106,6 +127,9 @@ def _find_ctk_header_directory(libname: str) -> LocatedHeaderDir | None:
     if cuda_home and (result := _locate_based_on_ctk_layout(libname, h_basename, cuda_home)):
         return LocatedHeaderDir(abs_path=result, found_via="CUDA_HOME")
 
+    if result := _find_ctk_header_directory_via_canary(libname, h_basename):
+        return LocatedHeaderDir(abs_path=result, found_via="system-ctk-root")
+
     return None
 
 
@@ -139,6 +163,12 @@ def locate_nvidia_header_directory(libname: str) -> LocatedHeaderDir | None:
         3. **CUDA Toolkit environment variables**
 
            - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
+
+        4. **CTK root canary probe**
+
+           - Probe a system-loaded ``cudart`` in a spawned child process,
+             derive the CTK root from the resolved library path, then search
+             CTK include layout under that root.
     """
 
     if libname in supported_nvidia_headers.SUPPORTED_HEADERS_CTK:
@@ -195,6 +225,12 @@ def find_nvidia_header_directory(libname: str) -> str | None:
         3. **CUDA Toolkit environment variables**
 
            - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
+
+        4. **CTK root canary probe**
+
+           - Probe a system-loaded ``cudart`` in a spawned child process,
+             derive the CTK root from the resolved library path, then search
+             CTK include layout under that root.
     """
     found = locate_nvidia_header_directory(libname)
     return found.abs_path if found else None

cuda_pathfinder/tests/test_find_nvidia_headers.py

@@ -16,10 +16,15 @@
 import importlib.metadata
 import os
 import re
+from pathlib import Path
 
 import pytest
 
+import cuda.pathfinder._headers.find_nvidia_headers as find_nvidia_headers_module
 from cuda.pathfinder import LocatedHeaderDir, find_nvidia_header_directory, locate_nvidia_header_directory
+from cuda.pathfinder._dynamic_libs.load_nvidia_dynamic_lib import (
+    _resolve_system_loaded_abs_path_in_subprocess,
+)
 from cuda.pathfinder._headers.supported_nvidia_headers import (
     SUPPORTED_HEADERS_CTK,
     SUPPORTED_HEADERS_CTK_ALL,
@@ -28,6 +33,7 @@
     SUPPORTED_INSTALL_DIRS_NON_CTK,
     SUPPORTED_SITE_PACKAGE_HEADER_DIRS_CTK,
 )
+from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS", "see_what_works")
 assert STRICTNESS in ("see_what_works", "all_must_work")
@@ -46,7 +52,13 @@ def test_unknown_libname():
 
 def _located_hdr_dir_asserts(located_hdr_dir):
     assert isinstance(located_hdr_dir, LocatedHeaderDir)
-    assert located_hdr_dir.found_via in ("site-packages", "conda", "CUDA_HOME", "supported_install_dir")
+    assert located_hdr_dir.found_via in (
+        "site-packages",
+        "conda",
+        "CUDA_HOME",
+        "system-ctk-root",
+        "supported_install_dir",
+    )
 
 
 def test_non_ctk_importlib_metadata_distributions_names():
@@ -62,6 +74,36 @@ def have_distribution_for(libname: str) -> bool:
     )
 
 
+@pytest.fixture
+def clear_locate_nvidia_header_cache():
+    locate_nvidia_header_directory.cache_clear()
+    _resolve_system_loaded_abs_path_in_subprocess.cache_clear()
+    yield
+    locate_nvidia_header_directory.cache_clear()
+    _resolve_system_loaded_abs_path_in_subprocess.cache_clear()
+
+
+def _create_ctk_header(ctk_root: Path, libname: str) -> str:
+    """Create a fake CTK header file and return its directory."""
+    header_basename = SUPPORTED_HEADERS_CTK[libname]
+    if libname == "nvvm":
+        header_dir = ctk_root / "nvvm" / "include"
+    elif libname == "cccl":
+        header_dir = ctk_root / "include" / "cccl"
+    else:
+        header_dir = ctk_root / "include"
+    header_path = header_dir / header_basename
+    header_path.parent.mkdir(parents=True, exist_ok=True)
+    header_path.touch()
+    return str(header_dir)
+
+
+def _fake_cudart_canary_abs_path(ctk_root: Path) -> str:
+    if IS_WINDOWS:
+        return str(ctk_root / "bin" / "x64" / "cudart64_13.dll")
+    return str(ctk_root / "lib64" / "libcudart.so.13")
+
+
 @pytest.mark.parametrize("libname", SUPPORTED_HEADERS_NON_CTK.keys())
 def test_locate_non_ctk_headers(info_summary_append, libname):
     hdr_dir = find_nvidia_header_directory(libname)
@@ -110,3 +152,85 @@ def test_locate_ctk_headers(info_summary_append, libname):
         assert os.path.isfile(os.path.join(hdr_dir, h_filename))
     if STRICTNESS == "all_must_work":
         assert hdr_dir is not None
+
+
+@pytest.mark.usefixtures("clear_locate_nvidia_header_cache")
+def test_locate_ctk_headers_uses_canary_fallback_when_cuda_home_unset(tmp_path, monkeypatch, mocker):
+    ctk_root = tmp_path / "cuda-system"
+    expected_hdr_dir = _create_ctk_header(ctk_root, "cudart")
+
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    monkeypatch.delenv("CUDA_HOME", raising=False)
+    monkeypatch.delenv("CUDA_PATH", raising=False)
+    mocker.patch.object(find_nvidia_headers_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    probe = mocker.patch.object(
+        find_nvidia_headers_module,
+        "_resolve_system_loaded_abs_path_in_subprocess",
+        return_value=_fake_cudart_canary_abs_path(ctk_root),
+    )
+
+    located_hdr_dir = locate_nvidia_header_directory("cudart")
+
+    assert located_hdr_dir is not None
+    assert located_hdr_dir.abs_path == expected_hdr_dir
+    assert located_hdr_dir.found_via == "system-ctk-root"
+    probe.assert_called_once_with("cudart")
+
+
+@pytest.mark.usefixtures("clear_locate_nvidia_header_cache")
+def test_locate_ctk_headers_cuda_home_takes_priority_over_canary(tmp_path, monkeypatch, mocker):
+    cuda_home = tmp_path / "cuda-home"
+    expected_hdr_dir = _create_ctk_header(cuda_home, "cudart")
+    canary_root = tmp_path / "cuda-system"
+    _create_ctk_header(canary_root, "cudart")
+
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    monkeypatch.setenv("CUDA_HOME", str(cuda_home))
+    monkeypatch.delenv("CUDA_PATH", raising=False)
+    mocker.patch.object(find_nvidia_headers_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    probe = mocker.patch.object(
+        find_nvidia_headers_module,
+        "_resolve_system_loaded_abs_path_in_subprocess",
+        return_value=_fake_cudart_canary_abs_path(canary_root),
+    )
+
+    located_hdr_dir = locate_nvidia_header_directory("cudart")
+
+    assert located_hdr_dir is not None
+    assert located_hdr_dir.abs_path == expected_hdr_dir
+    assert located_hdr_dir.found_via == "CUDA_HOME"
+    probe.assert_not_called()
+
+
+@pytest.mark.usefixtures("clear_locate_nvidia_header_cache")
+def test_locate_ctk_headers_canary_miss_paths_are_non_fatal(monkeypatch, mocker):
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    monkeypatch.delenv("CUDA_HOME", raising=False)
+    monkeypatch.delenv("CUDA_PATH", raising=False)
+    mocker.patch.object(find_nvidia_headers_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    mocker.patch.object(
+        find_nvidia_headers_module,
+        "_resolve_system_loaded_abs_path_in_subprocess",
+        return_value=None,
+    )
+
+    assert locate_nvidia_header_directory("cudart") is None
+    assert find_nvidia_header_directory("cudart") is None
+
+
+@pytest.mark.usefixtures("clear_locate_nvidia_header_cache")
+def test_locate_ctk_headers_canary_probe_errors_are_not_masked(monkeypatch, mocker):
+    monkeypatch.delenv("CONDA_PREFIX", raising=False)
+    monkeypatch.delenv("CUDA_HOME", raising=False)
+    monkeypatch.delenv("CUDA_PATH", raising=False)
+    mocker.patch.object(find_nvidia_headers_module, "find_sub_dirs_all_sitepackages", return_value=[])
+    mocker.patch.object(
+        find_nvidia_headers_module,
+        "_resolve_system_loaded_abs_path_in_subprocess",
+        side_effect=RuntimeError("canary probe failed"),
+    )
+
+    with pytest.raises(RuntimeError, match="canary probe failed"):
+        locate_nvidia_header_directory("cudart")
+    with pytest.raises(RuntimeError, match="canary probe failed"):
+        find_nvidia_header_directory("cudart")