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_")