Initial version of cuda.pathfinder._find_nvidia_headers for nvshmem

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

@@ -255,6 +255,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
@@ -289,4 +290,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

@@ -261,6 +261,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
 
@@ -298,5 +299,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,42 @@
+# 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
+
+from cuda.pathfinder._dynamic_libs.supported_nvidia_libs import IS_WINDOWS
+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 = os.environ.get("CONDA_PREFIX")
+    if conda_prefix and os.path.isdir(conda_prefix):
+        hdr_dir = os.path.join(conda_prefix, "include")
+        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):
+        nvshmem_h_path = os.path.join(hdr_dir, "nvshmem.h")
+        if os.path.isfile(nvshmem_h_path):
+            return hdr_dir
+
+    return None

cuda_pathfinder/cuda/pathfinder/_version.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "1.2.2a0"
+__version__ = "1.2.2"

cuda_pathfinder/docs/nv-versions.json

@@ -3,6 +3,10 @@
         "version": "latest",
         "url": "https://nvidia.github.io/cuda-python/cuda-pathfinder/latest/"
     },
+    {
+        "version": "1.2.2",
+        "url": "https://nvidia.github.io/cuda-python/cuda-pathfinder/1.2.2/"
+    },
     {
         "version": "1.2.1",
         "url": "https://nvidia.github.io/cuda-python/cuda-pathfinder/1.2.1/"

cuda_pathfinder/docs/source/release.rst

@@ -7,6 +7,7 @@ Release Notes
 .. toctree::
    :maxdepth: 3
 
+   1.2.2 <release/1.2.2-notes>
    1.2.1 <release/1.2.1-notes>
    1.2.0 <release/1.2.0-notes>
    1.1.0 <release/1.1.0-notes>

cuda_pathfinder/docs/source/release/1.2.2-notes.rst

@@ -0,0 +1,19 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-License-Identifier: Apache-2.0
+
+.. module:: cuda.pathfinder
+
+``cuda-pathfinder`` 1.2.2 Release notes
+=======================================
+
+Released on Sep 8, 2025
+
+
+Highlights
+----------
+
+* Support nccl library (`PR #945 <https://github.com/NVIDIA/cuda-python/pull/945>`_)
+
+* Add experimental ``cuda.pathfinder._find_nvidia_headers`` API,
+  currently limited to supporting ``nvshmem``
+  (`PR #661 <https://github.com/NVIDIA/cuda-python/pull/661>`_)

cuda_pathfinder/tests/test_find_nvidia_headers.py

@@ -0,0 +1,60 @@
+# 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 hdr_dir:
+        assert os.path.isdir(hdr_dir)
+        assert os.path.isfile(os.path.join(hdr_dir, "nvshmem.h"))
+    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 "site-packages" in hdr_dir_parts
+        elif conda_prefix := os.getenv("CONDA_PREFIX"):
+            assert hdr_dir.startswith(conda_prefix)
+        else:
+            assert hdr_dir.startswith("/usr/include/nvshmem_")