From d174f5a98d7481820e11abb413743a950624ddc2 Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Tue, 27 May 2025 16:41:22 -0700 Subject: [PATCH 1/4] find_nvidia_headers.py initial version (untested). --- .../_path_finder/find_nvidia_headers.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py diff --git a/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py b/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py new file mode 100644 index 000000000..bf56097f2 --- /dev/null +++ b/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py @@ -0,0 +1,38 @@ +# Copyright 2025 NVIDIA Corporation. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE + +import functools +import glob +import os + +from cuda.bindings._path_finder.find_sub_dirs import find_sub_dirs_all_sitepackages +from cuda.bindings._path_finder.supported_libs import IS_WINDOWS + + +@functools.cache +def find_nvidia_header_directory(libname: str) -> str: + assert libname == "nvshmem" + assert not IS_WINDOWS + + # Installed from a wheel + nvidia_sub_dirs = ("nvidia", "nvshmem", "include") + 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") + 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_*")): + 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 + + return None From 5ed86c56b0743f4aee8a9d0077c56e307c0b59d2 Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Wed, 28 May 2025 10:09:11 -0700 Subject: [PATCH 2/4] Add tests/test_path_finder_find_headers.py, with hard-coded paths. --- .../tests/test_path_finder_find_headers.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 cuda_bindings/tests/test_path_finder_find_headers.py diff --git a/cuda_bindings/tests/test_path_finder_find_headers.py b/cuda_bindings/tests/test_path_finder_find_headers.py new file mode 100644 index 000000000..5d9b09514 --- /dev/null +++ b/cuda_bindings/tests/test_path_finder_find_headers.py @@ -0,0 +1,24 @@ +# Copyright 2025 NVIDIA Corporation. All rights reserved. +# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE + +import pytest + +from cuda.bindings._path_finder import find_nvidia_headers + + +def test_find_nvidia_header_directory(info_summary_append): + with pytest.raises(AssertionError): + find_nvidia_headers.find_nvidia_header_directory("unknown-libname") + + hdr_dir = find_nvidia_headers.find_nvidia_header_directory("nvshmem") + assert hdr_dir in [ + # pip install nvidia-nvshmem-cu12 + "/home/rgrossekunst/forked/cuda-python/venvs/scratch/lib/python3.12/site-packages/nvidia/nvshmem/include", + # conda create -y -n nvshmem python=3.12 + # conda activate nvshmem + # conda install -y conda-forge::libnvshmem3 conda-forge::libnvshmem-dev + "/home/rgrossekunst/miniforge3/envs/nvshmem/include", + # sudo apt install libnvshmem3-cuda-12 libnvshmem3-dev-cuda-12 + "/usr/include/nvshmem_12", + ] + info_summary_append(f"{hdr_dir=!r}") From a98de70227eda8ac42f9d1e62ab1c5e59ccf9b2e Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Wed, 28 May 2025 10:29:14 -0700 Subject: [PATCH 3/4] Better error message: UNKNOWN libname='unknown-libname' --- .../cuda/bindings/_path_finder/find_nvidia_headers.py | 4 +++- cuda_bindings/tests/test_path_finder_find_headers.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py b/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py index bf56097f2..463235450 100644 --- a/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py +++ b/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py @@ -11,7 +11,9 @@ @functools.cache def find_nvidia_header_directory(libname: str) -> str: - assert libname == "nvshmem" + if libname != "nvshmem": + raise RuntimeError(f"UNKNOWN {libname=}") + assert not IS_WINDOWS # Installed from a wheel diff --git a/cuda_bindings/tests/test_path_finder_find_headers.py b/cuda_bindings/tests/test_path_finder_find_headers.py index 5d9b09514..587cfd0a9 100644 --- a/cuda_bindings/tests/test_path_finder_find_headers.py +++ b/cuda_bindings/tests/test_path_finder_find_headers.py @@ -7,7 +7,7 @@ def test_find_nvidia_header_directory(info_summary_append): - with pytest.raises(AssertionError): + with pytest.raises(RuntimeError, match="^UNKNOWN libname='unknown-libname'$"): find_nvidia_headers.find_nvidia_header_directory("unknown-libname") hdr_dir = find_nvidia_headers.find_nvidia_header_directory("nvshmem") From 6a93fe7b9f3ccb538fa26bba07e5808db97c03d4 Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Wed, 28 May 2025 10:40:04 -0700 Subject: [PATCH 4/4] if libname == "nvshmem" and IS_WINDOWS: return None --- .../cuda/bindings/_path_finder/find_nvidia_headers.py | 4 +++- cuda_bindings/tests/test_path_finder_find_headers.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py b/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py index 463235450..befd4ca81 100644 --- a/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py +++ b/cuda_bindings/cuda/bindings/_path_finder/find_nvidia_headers.py @@ -14,7 +14,9 @@ def find_nvidia_header_directory(libname: str) -> str: if libname != "nvshmem": raise RuntimeError(f"UNKNOWN {libname=}") - assert not IS_WINDOWS + if libname == "nvshmem" and IS_WINDOWS: + # nvshmem has no Windows support. + return None # Installed from a wheel nvidia_sub_dirs = ("nvidia", "nvshmem", "include") diff --git a/cuda_bindings/tests/test_path_finder_find_headers.py b/cuda_bindings/tests/test_path_finder_find_headers.py index 587cfd0a9..d76081789 100644 --- a/cuda_bindings/tests/test_path_finder_find_headers.py +++ b/cuda_bindings/tests/test_path_finder_find_headers.py @@ -11,14 +11,20 @@ def test_find_nvidia_header_directory(info_summary_append): find_nvidia_headers.find_nvidia_header_directory("unknown-libname") hdr_dir = find_nvidia_headers.find_nvidia_header_directory("nvshmem") + # TODO: Find ways to test more meaningfully, and how to avoid HARD-WIRED PATHS in particular. assert hdr_dir in [ # pip install nvidia-nvshmem-cu12 "/home/rgrossekunst/forked/cuda-python/venvs/scratch/lib/python3.12/site-packages/nvidia/nvshmem/include", + # # conda create -y -n nvshmem python=3.12 # conda activate nvshmem # conda install -y conda-forge::libnvshmem3 conda-forge::libnvshmem-dev "/home/rgrossekunst/miniforge3/envs/nvshmem/include", + # # sudo apt install libnvshmem3-cuda-12 libnvshmem3-dev-cuda-12 "/usr/include/nvshmem_12", + # + # nvshmem not available + None, ] info_summary_append(f"{hdr_dir=!r}")