diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py index 1727cca607..13f47fc2b5 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py +++ b/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 diff --git a/cuda_pathfinder/tests/test_find_nvidia_headers.py b/cuda_pathfinder/tests/test_find_nvidia_headers.py index f14681546d..2732de216b 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/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")