Initial version of cuda.pathfinder._find_nvidia_headers for nvshmem

.github/workflows/test-wheel-linux.yml

@@ -298,6 +298,7 @@ jobs:
       - name: Run cuda.pathfinder tests with see_what_works
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: see_what_works
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: see_what_works
         run: run-tests pathfinder
 
       - name: Run cuda.bindings tests
@@ -332,4 +333,5 @@ jobs:
       - name: Run cuda.pathfinder tests with all_must_work
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: all_must_work
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: all_must_work
         run: run-tests pathfinder

.github/workflows/test-wheel-windows.yml

@@ -262,6 +262,7 @@ jobs:
       - name: Run cuda.pathfinder tests with see_what_works
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: see_what_works
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: see_what_works
         shell: bash --noprofile --norc -xeuo pipefail {0}
         run: run-tests pathfinder
 
@@ -299,5 +300,6 @@ jobs:
       - name: Run cuda.pathfinder tests with all_must_work
         env:
           CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS: all_must_work
+          CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS: all_must_work
         shell: bash --noprofile --norc -xeuo pipefail {0}
         run: run-tests pathfinder

ci/tools/run-tests

@@ -31,7 +31,9 @@ popd
 
 if [[ "${test_module}" == "pathfinder" ]]; then
   pushd ./cuda_pathfinder
-  echo "Running pathfinder tests with ${CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS}"
+  echo "Running pathfinder tests with " \
+      "LD:${CUDA_PATHFINDER_TEST_LOAD_NVIDIA_DYNAMIC_LIB_STRICTNESS} " \
+      "FH:${CUDA_PATHFINDER_TEST_FIND_NVIDIA_HEADERS_STRICTNESS}"
   pwd
   pytest -ra -s -v tests/
   popd

cuda_pathfinder/cuda/pathfinder/__init__.py

@@ -7,4 +7,7 @@
 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 (
+    find_nvidia_header_directory as _find_nvidia_header_directory,  # noqa: F401
+)
 from cuda.pathfinder._version import __version__ as __version__

cuda_pathfinder/cuda/pathfinder/_headers/find_nvidia_headers.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import glob
+import os
+from typing import Optional, cast
+
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import IS_WINDOWS
+from cuda.pathfinder._utils.conda_env import get_conda_prefix
+from cuda.pathfinder._utils.env_vars_for_include import iter_env_vars_for_include_dirs
+from cuda.pathfinder._utils.find_sub_dirs import find_sub_dirs_all_sitepackages
+
+
+@functools.cache
+def find_nvidia_header_directory(libname: str) -> Optional[str]:
+    if libname != "nvshmem":
+        raise RuntimeError(f"UNKNOWN {libname=}")
+
+    if libname == "nvshmem" and IS_WINDOWS:
+        # nvshmem has no Windows support.
+        return None
+
+    # Installed from a wheel
+    nvidia_sub_dirs = ("nvidia", "nvshmem", "include")
+    hdr_dir: str  # help mypy
+    for hdr_dir in find_sub_dirs_all_sitepackages(nvidia_sub_dirs):
+        nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+        if os.path.isfile(nvshmem_h_path):
+            return hdr_dir
+
+    conda_prefix = get_conda_prefix()
+    if conda_prefix and os.path.isdir(conda_prefix.path):
+        hdr_dir = os.path.join(conda_prefix.path, "include")
+        if os.path.isdir(hdr_dir):
+            nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+            if os.path.isfile(nvshmem_h_path):
+                return hdr_dir
+
+    for hdr_dir in sorted(glob.glob("/usr/include/nvshmem_*"), reverse=True):
+        if os.path.isdir(hdr_dir):
+            nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+            if os.path.isfile(nvshmem_h_path):
+                return hdr_dir
+
+    for hdr_dir in iter_env_vars_for_include_dirs():
+        if os.path.isdir(hdr_dir):
+            nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+            if os.path.isfile(nvshmem_h_path):
+                return cast(str, hdr_dir)  # help mypy
+
+    return None

cuda_pathfinder/cuda/pathfinder/_utils/conda_env.py

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, Optional
+
+# https://docs.conda.io/projects/conda-build/en/stable/user-guide/environment-variables.html
+
+BUILD_STATES = ("RENDER", "BUILD", "TEST")
+
+
+@dataclass(frozen=True)
+class CondaPrefix:
+    env_state: Literal["RENDER", "BUILD", "TEST", "activated"]
+    path: Path
+
+
+@functools.cache
+def get_conda_prefix() -> Optional[CondaPrefix]:
+    """
+    Return the effective conda prefix.
+    - RENDER, BUILD, TEST: inside conda-build (host prefix at $PREFIX)
+    - activated: user-activated env ($CONDA_PREFIX)
+    - None: neither detected
+    """
+    state = os.getenv("CONDA_BUILD_STATE")
+    if state:
+        if state in BUILD_STATES:
+            p = os.getenv("PREFIX")
+            if p:
+                return CondaPrefix(state, Path(p))  # type: ignore[arg-type]
+        return None
+
+    cp = os.getenv("CONDA_PREFIX")
+    if cp:
+        return CondaPrefix("activated", Path(cp))
+
+    return None

cuda_pathfinder/cuda/pathfinder/_utils/env_vars_for_include.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import sys
+from collections.abc import Iterable
+
+IS_WINDOWS = sys.platform == "win32"
+
+# GCC/Clang-style include vars
+GCC_VNAMES = ("CPATH", "C_INCLUDE_PATH", "CPLUS_INCLUDE_PATH")
+
+# MSVC: INCLUDE is the canonical header search variable
+MSVC_GCC_VNAMES = ("INCLUDE",)
+
+VNAMES: tuple[str, ...] = MSVC_GCC_VNAMES + GCC_VNAMES if IS_WINDOWS else GCC_VNAMES
+
+
+def iter_env_vars_for_include_dirs() -> Iterable[str]:
+    for vname in VNAMES:
+        v = os.getenv(vname)
+        if v:
+            for d in v.split(os.pathsep):
+                if d:
+                    yield d

cuda_pathfinder/tests/test_find_nvidia_headers.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Currently these installations are only manually tested:
+
+# conda create -y -n nvshmem python=3.12
+# conda activate nvshmem
+# conda install -y conda-forge::libnvshmem3 conda-forge::libnvshmem-dev
+
+# wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb
+# sudo dpkg -i cuda-keyring_1.1-1_all.deb
+# sudo apt update
+# sudo apt install libnvshmem3-cuda-12 libnvshmem3-dev-cuda-12
+# sudo apt install libnvshmem3-cuda-13 libnvshmem3-dev-cuda-13
+
+import functools
+import importlib.metadata
+import os
+import re
+
+import pytest
+
+from cuda.pathfinder import _find_nvidia_header_directory as find_nvidia_header_directory
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs 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")
+
+
+@functools.cache
+def have_nvidia_nvshmem_package() -> bool:
+    pattern = re.compile(r"^nvidia-nvshmem-.*$")
+    return any(
+        pattern.match(dist.metadata["Name"]) for dist in importlib.metadata.distributions() if "Name" in dist.metadata
+    )
+
+
+def test_unknown_libname():
+    with pytest.raises(RuntimeError, match=r"^UNKNOWN libname='unknown-libname'$"):
+        find_nvidia_header_directory("unknown-libname")
+
+
+def test_find_libname_nvshmem(info_summary_append):
+    hdr_dir = find_nvidia_header_directory("nvshmem")
+    info_summary_append(f"{hdr_dir=!r}")
+    if IS_WINDOWS:
+        assert hdr_dir is None
+        pytest.skip("nvshmem has no Windows support.")
+    if STRICTNESS == "all_must_work" or have_nvidia_nvshmem_package():
+        assert hdr_dir is not None
+        if have_nvidia_nvshmem_package():
+            hdr_dir_parts = hdr_dir.split(os.path.sep)
+            assert any(
+                sub_dir in hdr_dir_parts
+                for sub_dir in (
+                    "site-packages",  # pip install
+                    "dist-packages",  # apt install
+                )
+            ), hdr_dir

cuda_pathfinder/tests/test_utils_conda_env.py

@@ -0,0 +1,105 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+
+from cuda.pathfinder._utils.conda_env import (
+    BUILD_STATES,
+    CondaPrefix,
+    get_conda_prefix,
+)
+
+# Auto-clean environment & cache before every test -----------------------------
+
+
+@pytest.fixture(autouse=True)
+def _clean_env_and_cache(monkeypatch):
+    # Remove any possibly inherited variables from the test runner environment
+    for k in ("CONDA_BUILD_STATE", "PREFIX", "CONDA_PREFIX"):
+        monkeypatch.delenv(k, raising=False)
+    # Clear the cached result between tests
+    get_conda_prefix.cache_clear()
+    return
+    # (No teardown needed; monkeypatch auto-reverts)
+
+
+# Tests -----------------------------------------------------------------------
+
+
+def test_returns_none_when_no_relevant_env_vars():
+    assert get_conda_prefix() is None
+
+
+@pytest.mark.parametrize("state", BUILD_STATES)
+def test_build_state_returns_prefix_when_present(state, monkeypatch, tmp_path):
+    monkeypatch.setenv("CONDA_BUILD_STATE", state)
+    monkeypatch.setenv("PREFIX", str(tmp_path))
+    res = get_conda_prefix()
+    assert isinstance(res, CondaPrefix)
+    assert res.env_state == state
+    assert res.path == tmp_path
+
+
+@pytest.mark.parametrize("state", BUILD_STATES)
+def test_build_state_requires_prefix_otherwise_none(state, monkeypatch):
+    monkeypatch.setenv("CONDA_BUILD_STATE", state)
+    # No PREFIX set
+    assert get_conda_prefix() is None
+
+
+@pytest.mark.parametrize("state", BUILD_STATES)
+def test_build_state_with_empty_prefix_returns_none(state, monkeypatch):
+    monkeypatch.setenv("CONDA_BUILD_STATE", state)
+    monkeypatch.setenv("PREFIX", "")
+    assert get_conda_prefix() is None
+
+
+def test_activated_env_returns_conda_prefix(monkeypatch, tmp_path):
+    monkeypatch.setenv("CONDA_PREFIX", str(tmp_path))
+    res = get_conda_prefix()
+    assert isinstance(res, CondaPrefix)
+    assert res.env_state == "activated"
+    assert res.path == tmp_path
+
+
+def test_activated_env_ignores_empty_conda_prefix(monkeypatch):
+    monkeypatch.setenv("CONDA_PREFIX", "")
+    assert get_conda_prefix() is None
+
+
+def test_build_state_wins_over_activated_when_valid(monkeypatch, tmp_path):
+    build_p = tmp_path / "host"
+    user_p = tmp_path / "user"
+    monkeypatch.setenv("CONDA_BUILD_STATE", "TEST")
+    monkeypatch.setenv("PREFIX", str(build_p))
+    monkeypatch.setenv("CONDA_PREFIX", str(user_p))
+    res = get_conda_prefix()
+    assert res
+    assert res.env_state == "TEST"
+    assert res.path == build_p
+
+
+def test_unknown_build_state_returns_none_even_if_conda_prefix_set(monkeypatch, tmp_path):
+    # Any non-empty CONDA_BUILD_STATE that is not recognized -> None
+    monkeypatch.setenv("CONDA_BUILD_STATE", "SOMETHING_ELSE")
+    monkeypatch.setenv("CONDA_PREFIX", str(tmp_path))
+    assert get_conda_prefix() is None
+
+
+def test_empty_build_state_treated_as_absent_and_falls_back_to_activated(monkeypatch, tmp_path):
+    # Empty string is falsy -> treated like "not set" -> activated path
+    monkeypatch.setenv("CONDA_BUILD_STATE", "")
+    monkeypatch.setenv("CONDA_PREFIX", str(tmp_path))
+    res = get_conda_prefix()
+    assert res
+    assert res.env_state == "activated"
+    assert res.path == tmp_path
+
+
+def test_have_cache(monkeypatch, tmp_path):
+    monkeypatch.setenv("CONDA_PREFIX", str(tmp_path))
+    res = get_conda_prefix()
+    assert res
+    assert res.path == tmp_path
+    res2 = get_conda_prefix()
+    assert res2 is res