Introduce LocatedHeaderDir to report the header discovery method (#1536)

* initial

* resolve a few tests

* precommit

* address reviews

* test updates

* address reviews

* address reviews

Author: brandon-b-miller

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -9,7 +9,11 @@
 from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import (
     SUPPORTED_LIBNAMES as SUPPORTED_NVIDIA_LIBNAMES,  # noqa: F401
 )
+from cuda.pathfinder._headers.find_nvidia_headers import LocatedHeaderDir as LocatedHeaderDir
 from cuda.pathfinder._headers.find_nvidia_headers import find_nvidia_header_directory as find_nvidia_header_directory
+from cuda.pathfinder._headers.find_nvidia_headers import (
+    locate_nvidia_header_directory as locate_nvidia_header_directory,
+)
 from cuda.pathfinder._headers.supported_nvidia_headers import SUPPORTED_HEADERS_CTK as _SUPPORTED_HEADERS_CTK
 
 from cuda.pathfinder._version import __version__  # isort: skip  # noqa: F401

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -1,16 +1,28 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import functools
 import glob
 import os
+from dataclasses import dataclass
 
 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
 from cuda.pathfinder._utils.platform_aware import IS_WINDOWS
 
 
+@dataclass
+class LocatedHeaderDir:
+    abs_path: str | None
+    found_via: str
+
+    def __post_init__(self) -> None:
+        self.abs_path = _abs_norm(self.abs_path)
+
+
 def _abs_norm(path: str | None) -> str | None:
     if path:
         return os.path.normpath(os.path.abspath(path))
@@ -21,16 +33,16 @@ def _joined_isfile(dirpath: str, basename: str) -> bool:
     return os.path.isfile(os.path.join(dirpath, basename))
 
 
-def _find_under_site_packages(sub_dir: str, h_basename: str) -> str | None:
+def _locate_under_site_packages(sub_dir: str, h_basename: str) -> LocatedHeaderDir | None:
     # Installed from a wheel
     hdr_dir: str  # help mypy
     for hdr_dir in find_sub_dirs_all_sitepackages(tuple(sub_dir.split("/"))):
         if _joined_isfile(hdr_dir, h_basename):
-            return hdr_dir
+            return LocatedHeaderDir(abs_path=hdr_dir, found_via="site-packages")
     return None
 
 
-def _find_based_on_ctk_layout(libname: str, h_basename: str, anchor_point: str) -> str | None:
+def _locate_based_on_ctk_layout(libname: str, h_basename: str, anchor_point: str) -> str | None:
     parts = [anchor_point]
     if libname == "nvvm":
         parts.append(libname)
@@ -52,7 +64,7 @@ def _find_based_on_ctk_layout(libname: str, h_basename: str, anchor_point: str)
     return None
 
 
-def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool) -> str | None:
+def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool) -> LocatedHeaderDir | None:
     conda_prefix = os.environ.get("CONDA_PREFIX")
     if not conda_prefix:
         return None
@@ -73,39 +85,43 @@ def _find_based_on_conda_layout(libname: str, h_basename: str, ctk_layout: bool)
         else:
             include_path = os.path.join(conda_prefix, "include")
         anchor_point = os.path.dirname(include_path)
-    return _find_based_on_ctk_layout(libname, h_basename, anchor_point)
+    found_header_path = _locate_based_on_ctk_layout(libname, h_basename, anchor_point)
+    if found_header_path:
+        return LocatedHeaderDir(abs_path=found_header_path, found_via="conda")
+    return None
 
 
-def _find_ctk_header_directory(libname: str) -> str | None:
+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]
 
     for cdir in candidate_dirs:
-        if hdr_dir := _find_under_site_packages(cdir, h_basename):
+        if hdr_dir := _locate_under_site_packages(cdir, h_basename):
             return hdr_dir
 
     if hdr_dir := _find_based_on_conda_layout(libname, h_basename, True):
         return hdr_dir
 
     cuda_home = get_cuda_home_or_path()
     if cuda_home:  # noqa: SIM102
-        if result := _find_based_on_ctk_layout(libname, h_basename, cuda_home):
-            return result
+        if result := _locate_based_on_ctk_layout(libname, h_basename, cuda_home):
+            return LocatedHeaderDir(abs_path=result, found_via="CUDA_HOME")
 
     return None
 
 
 @functools.cache
-def find_nvidia_header_directory(libname: str) -> str | None:
+def locate_nvidia_header_directory(libname: str) -> LocatedHeaderDir | None:
     """Locate the header directory for a supported NVIDIA library.
 
     Args:
         libname (str): The short name of the library whose headers are needed
             (e.g., ``"nvrtc"``, ``"cusolver"``, ``"nvshmem"``).
 
     Returns:
-        str or None: Absolute path to the discovered header directory, or ``None``
-        if the headers cannot be found.
+        LocatedHeaderDir or None: A LocatedHeaderDir object containing the absolute path
+        to the discovered header directory and information about where it was found,
+        or ``None`` if the headers cannot be found.
 
     Raises:
         RuntimeError: If ``libname`` is not in the supported set.
@@ -127,25 +143,59 @@ def find_nvidia_header_directory(libname: str) -> str | None:
     """
 
     if libname in supported_nvidia_headers.SUPPORTED_HEADERS_CTK:
-        return _abs_norm(_find_ctk_header_directory(libname))
+        return _find_ctk_header_directory(libname)
 
     h_basename = supported_nvidia_headers.SUPPORTED_HEADERS_NON_CTK.get(libname)
     if h_basename is None:
         raise RuntimeError(f"UNKNOWN {libname=}")
 
     candidate_dirs = supported_nvidia_headers.SUPPORTED_SITE_PACKAGE_HEADER_DIRS_NON_CTK.get(libname, [])
-    hdr_dir: str | None  # help mypy
+
     for cdir in candidate_dirs:
-        if hdr_dir := _find_under_site_packages(cdir, h_basename):
-            return _abs_norm(hdr_dir)
+        if found_hdr := _locate_under_site_packages(cdir, h_basename):
+            return found_hdr
 
-    if hdr_dir := _find_based_on_conda_layout(libname, h_basename, False):
-        return _abs_norm(hdr_dir)
+    if found_hdr := _find_based_on_conda_layout(libname, h_basename, False):
+        return found_hdr
 
+    # Fall back to system install directories
     candidate_dirs = supported_nvidia_headers.SUPPORTED_INSTALL_DIRS_NON_CTK.get(libname, [])
     for cdir in candidate_dirs:
         for hdr_dir in sorted(glob.glob(cdir), reverse=True):
             if _joined_isfile(hdr_dir, h_basename):
-                return _abs_norm(hdr_dir)
-
+                # For system installs, we don't have a clear found_via, so use "system"
+                return LocatedHeaderDir(abs_path=hdr_dir, found_via="supported_install_dir")
     return None
+
+
+def find_nvidia_header_directory(libname: str) -> str | None:
+    """Locate the header directory for a supported NVIDIA library.
+
+    Args:
+        libname (str): The short name of the library whose headers are needed
+            (e.g., ``"nvrtc"``, ``"cusolver"``, ``"nvshmem"``).
+
+    Returns:
+        str or None: Absolute path to the discovered header directory, or ``None``
+        if the headers cannot be found.
+
+    Raises:
+        RuntimeError: If ``libname`` is not in the supported set.
+
+    Search order:
+        1. **NVIDIA Python wheels**
+
+           - Scan installed distributions (``site-packages``) for header layouts
+             shipped in NVIDIA wheels (e.g., ``cuda-toolkit[nvrtc]``).
+
+        2. **Conda environments**
+
+           - Check Conda-style installation prefixes, which use platform-specific
+             include directory layouts.
+
+        3. **CUDA Toolkit environment variables**
+
+           - Use ``CUDA_HOME`` or ``CUDA_PATH`` (in that order).
+    """
+    found = locate_nvidia_header_directory(libname)
+    return found.abs_path if found else None

cuda_pathfinder/tests/test_find_nvidia_headers.py

@@ -19,7 +19,7 @@
 
 import pytest
 
-from cuda.pathfinder import find_nvidia_header_directory
+from cuda.pathfinder import LocatedHeaderDir, find_nvidia_header_directory, locate_nvidia_header_directory
 from cuda.pathfinder._headers.supported_nvidia_headers import (
     SUPPORTED_HEADERS_CTK,
     SUPPORTED_HEADERS_CTK_ALL,
@@ -44,6 +44,11 @@ def test_unknown_libname():
         find_nvidia_header_directory("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")
+
+
 def test_non_ctk_importlib_metadata_distributions_names():
     # Ensure the dict keys above stay in sync with supported_nvidia_headers
     assert sorted(NON_CTK_IMPORTLIB_METADATA_DISTRIBUTIONS_NAMES) == sorted(SUPPORTED_HEADERS_NON_CTK_ALL)
@@ -58,10 +63,14 @@ def have_distribution_for(libname: str) -> bool:
 
 
 @pytest.mark.parametrize("libname", SUPPORTED_HEADERS_NON_CTK.keys())
-def test_find_non_ctk_headers(info_summary_append, libname):
+def test_locate_non_ctk_headers(info_summary_append, libname):
     hdr_dir = find_nvidia_header_directory(libname)
+    located_hdr_dir = locate_nvidia_header_directory(libname)
+    assert hdr_dir is None if not located_hdr_dir else hdr_dir == located_hdr_dir.abs_path
+
     info_summary_append(f"{hdr_dir=!r}")
     if hdr_dir:
+        _located_hdr_dir_asserts(located_hdr_dir)
         assert os.path.isdir(hdr_dir)
         assert os.path.isfile(os.path.join(hdr_dir, SUPPORTED_HEADERS_NON_CTK[libname]))
     if have_distribution_for(libname):
@@ -88,10 +97,14 @@ def test_supported_headers_site_packages_ctk_consistency():
 
 
 @pytest.mark.parametrize("libname", SUPPORTED_HEADERS_CTK.keys())
-def test_find_ctk_headers(info_summary_append, libname):
+def test_locate_ctk_headers(info_summary_append, libname):
     hdr_dir = find_nvidia_header_directory(libname)
+    located_hdr_dir = locate_nvidia_header_directory(libname)
+    assert hdr_dir is None if not located_hdr_dir else hdr_dir == located_hdr_dir.abs_path
+
     info_summary_append(f"{hdr_dir=!r}")
     if hdr_dir:
+        _located_hdr_dir_asserts(located_hdr_dir)
         assert os.path.isdir(hdr_dir)
         h_filename = SUPPORTED_HEADERS_CTK[libname]
         assert os.path.isfile(os.path.join(hdr_dir, h_filename))