diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index 8da4020116..4ff5035b44 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/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 diff --git a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py index 63f8a627fd..637ecabfcb 100644 --- a/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py +++ b/cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py @@ -1,9 +1,12 @@ # 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 @@ -11,6 +14,15 @@ 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,15 +85,18 @@ 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): @@ -89,14 +104,14 @@ def _find_ctk_header_directory(libname: str) -> str | None: 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: @@ -104,8 +119,9 @@ def find_nvidia_header_directory(libname: str) -> str | None: (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 diff --git a/cuda_pathfinder/tests/test_find_nvidia_headers.py b/cuda_pathfinder/tests/test_find_nvidia_headers.py index 494d7c0ae9..f14681546d 100644 --- a/cuda_pathfinder/tests/test_find_nvidia_headers.py +++ b/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))