From 1a16604298a88c1f5a10a41ddd8256a528caed58 Mon Sep 17 00:00:00 2001 From: benjaming Date: Mon, 6 Jul 2026 16:54:26 +0000 Subject: [PATCH 1/3] Add exact-name bitcode library lookup --- cuda_pathfinder/cuda/pathfinder/__init__.py | 6 +- .../_static_libs/find_bitcode_lib.py | 53 ++++++-- cuda_pathfinder/docs/source/api.rst | 1 + .../tests/test_find_bitcode_lib.py | 127 +++++++++++++++++- 4 files changed, 171 insertions(+), 16 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index dc818dfd08f..5c155017797 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/cuda_pathfinder/cuda/pathfinder/__init__.py @@ -41,6 +41,9 @@ from cuda.pathfinder._static_libs.find_bitcode_lib import ( find_bitcode_lib as find_bitcode_lib, ) +from cuda.pathfinder._static_libs.find_bitcode_lib import ( + find_bitcode_lib_by_name as find_bitcode_lib_by_name, +) from cuda.pathfinder._static_libs.find_bitcode_lib import ( locate_bitcode_lib as locate_bitcode_lib, ) @@ -77,7 +80,8 @@ SUPPORTED_BINARY_UTILITIES = _SUPPORTED_BINARIES #: Tuple of supported bitcode library names that can be resolved -#: via ``locate_bitcode_lib()`` and ``find_bitcode_lib()``. +#: via ``locate_bitcode_lib()``, ``find_bitcode_lib()``, and +#: ``find_bitcode_lib_by_name()``. #: Example value: ``"device"``. SUPPORTED_BITCODE_LIBS = _SUPPORTED_BITCODE_LIBS diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py index ac038aadfe7..fbfb10befe1 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import functools +import ntpath import os from dataclasses import dataclass from typing import NoReturn, TypedDict @@ -74,13 +75,21 @@ def _no_such_file_in_dir(dir_path: str, filename: str, error_messages: list[str] attachments.append(f' Directory does not exist: "{dir_path}"') +def _validate_filename(filename: str) -> str: + if not isinstance(filename, str): + raise TypeError(f"filename must be a string, got {type(filename).__name__}") + if filename in ("", ".", "..") or "\0" in filename or ntpath.basename(filename) != filename: + raise ValueError(f"filename must be a file name without a directory: {filename!r}") + return filename + + class _FindBitcodeLib: - def __init__(self, name: str) -> None: + def __init__(self, name: str, *, filename: str | None = None) -> None: if name not in _SUPPORTED_BITCODE_LIBS_INFO: # Updated reference raise ValueError(f"Unknown bitcode library: '{name}'. Supported: {', '.join(SUPPORTED_BITCODE_LIBS)}") self.name: str = name self.config: _BitcodeLibInfo = _SUPPORTED_BITCODE_LIBS_INFO[name] # Updated reference - self.filename: str = self.config["filename"] + self.filename: str = self.config["filename"] if filename is None else _validate_filename(filename) self.rel_path: str = self.config["rel_path"] self.site_packages_dirs: tuple[str, ...] = self.config["site_packages_dirs"] self.error_messages: list[str] = [] @@ -130,14 +139,8 @@ def raise_not_found_error(self) -> NoReturn: raise BitcodeLibNotFoundError(f'Failure finding "{self.filename}": {err}\n{att}') -def locate_bitcode_lib(name: str) -> LocatedBitcodeLib: - """Locate a bitcode library by name. - - Raises: - ValueError: If ``name`` is not a supported bitcode library. - BitcodeLibNotFoundError: If the bitcode library cannot be found. - """ - finder = _FindBitcodeLib(name) +def _locate_bitcode_lib(name: str, *, filename: str | None = None) -> LocatedBitcodeLib: + finder = _FindBitcodeLib(name, filename=filename) abs_path = finder.try_site_packages() if abs_path is not None: @@ -169,6 +172,16 @@ def locate_bitcode_lib(name: str) -> LocatedBitcodeLib: finder.raise_not_found_error() +def locate_bitcode_lib(name: str) -> LocatedBitcodeLib: + """Locate a bitcode library by name. + + Raises: + ValueError: If ``name`` is not a supported bitcode library. + BitcodeLibNotFoundError: If the bitcode library cannot be found. + """ + return _locate_bitcode_lib(name) + + @functools.cache def find_bitcode_lib(name: str) -> str: """Find the absolute path to a bitcode library. @@ -178,3 +191,23 @@ def find_bitcode_lib(name: str) -> str: BitcodeLibNotFoundError: If the bitcode library cannot be found. """ return locate_bitcode_lib(name).abs_path + + +@functools.cache +def find_bitcode_lib_by_name(name: str, filename: str) -> str: + """Find an exact bitcode filename using a supported library's search paths. + + This lets a library define its own filename convention, including any + architecture or other attributes encoded in the filename. + + Args: + name: Supported bitcode library whose configured directories to search. + filename: Exact file name to find. Directory components are not allowed. + + Raises: + TypeError: If ``filename`` is not a string. + ValueError: If ``name`` is unsupported or ``filename`` is not a file name. + BitcodeLibNotFoundError: If ``filename`` cannot be found. + """ + _validate_filename(filename) + return _locate_bitcode_lib(name, filename=filename).abs_path diff --git a/cuda_pathfinder/docs/source/api.rst b/cuda_pathfinder/docs/source/api.rst index e49478c09ec..c5cd45a9e2b 100644 --- a/cuda_pathfinder/docs/source/api.rst +++ b/cuda_pathfinder/docs/source/api.rst @@ -35,6 +35,7 @@ CUDA bitcode and static libraries. SUPPORTED_BITCODE_LIBS find_bitcode_lib + find_bitcode_lib_by_name locate_bitcode_lib LocatedBitcodeLib BitcodeLibNotFoundError diff --git a/cuda_pathfinder/tests/test_find_bitcode_lib.py b/cuda_pathfinder/tests/test_find_bitcode_lib.py index 659b068f0ff..c04a4a35ddc 100644 --- a/cuda_pathfinder/tests/test_find_bitcode_lib.py +++ b/cuda_pathfinder/tests/test_find_bitcode_lib.py @@ -11,6 +11,7 @@ SUPPORTED_BITCODE_LIBS, BitcodeLibNotFoundError, find_bitcode_lib, + find_bitcode_lib_by_name, locate_bitcode_lib, ) from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home @@ -30,15 +31,17 @@ def _bitcode_lib_filename(libname: str) -> str: @pytest.fixture def clear_find_bitcode_lib_cache(): find_bitcode_lib_module.find_bitcode_lib.cache_clear() + find_bitcode_lib_module.find_bitcode_lib_by_name.cache_clear() get_cuda_path_or_home.cache_clear() yield find_bitcode_lib_module.find_bitcode_lib.cache_clear() + find_bitcode_lib_module.find_bitcode_lib_by_name.cache_clear() get_cuda_path_or_home.cache_clear() -def _make_bitcode_lib_file(dir_path: Path, libname: str) -> str: +def _make_bitcode_lib_file(dir_path: Path, filename: str) -> str: dir_path.mkdir(parents=True, exist_ok=True) - file_path = dir_path / _bitcode_lib_filename(libname) + file_path = dir_path / filename file_path.touch() return str(file_path) @@ -92,14 +95,15 @@ def test_locate_bitcode_lib(info_summary_append, libname): @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") @pytest.mark.parametrize("libname", SUPPORTED_BITCODE_LIBS) def test_locate_bitcode_lib_search_order(monkeypatch, tmp_path, libname): + filename = _bitcode_lib_filename(libname) site_packages_lib_dir = _site_packages_bitcode_lib_dir_under(tmp_path / "site-packages", libname) - site_packages_path = _make_bitcode_lib_file(site_packages_lib_dir, libname) + site_packages_path = _make_bitcode_lib_file(site_packages_lib_dir, filename) conda_prefix = tmp_path / "conda-prefix" - conda_path = _make_bitcode_lib_file(_bitcode_lib_dir_under(_conda_anchor(conda_prefix), libname), libname) + conda_path = _make_bitcode_lib_file(_bitcode_lib_dir_under(_conda_anchor(conda_prefix), libname), filename) cuda_home = tmp_path / "cuda-home" - cuda_home_path = _make_bitcode_lib_file(_bitcode_lib_dir_under(cuda_home, libname), libname) + cuda_home_path = _make_bitcode_lib_file(_bitcode_lib_dir_under(cuda_home, libname), filename) site_packages_sub_dirs = tuple( tuple(rel_dir.split("/")) for rel_dir in _bitcode_lib_info(libname)["site_packages_dirs"] @@ -135,6 +139,70 @@ def find_expected_sub_dir(sub_dir): assert located_lib.found_via == "CUDA_PATH" +@pytest.mark.usefixtures("clear_find_bitcode_lib_cache") +@pytest.mark.agent_authored(model="gpt-5") +def test_find_bitcode_lib_by_name_search_order(monkeypatch, tmp_path): + libname = "device" + filename = "libdevice_sm_90.bc" + site_packages_lib_dir = _site_packages_bitcode_lib_dir_under(tmp_path / "site-packages", libname) + site_packages_path = _make_bitcode_lib_file(site_packages_lib_dir, filename) + + conda_prefix = tmp_path / "conda-prefix" + conda_path = _make_bitcode_lib_file(_bitcode_lib_dir_under(_conda_anchor(conda_prefix), libname), filename) + + cuda_home = tmp_path / "cuda-home" + cuda_home_path = _make_bitcode_lib_file(_bitcode_lib_dir_under(cuda_home, libname), filename) + + site_packages_sub_dirs = tuple( + tuple(rel_dir.split("/")) for rel_dir in _bitcode_lib_info(libname)["site_packages_dirs"] + ) + + def find_expected_sub_dir(sub_dir): + assert sub_dir in site_packages_sub_dirs + if sub_dir == site_packages_sub_dirs[0]: + return [str(site_packages_lib_dir)] + return [] + + monkeypatch.setattr( + find_bitcode_lib_module, + "find_sub_dirs_all_sitepackages", + find_expected_sub_dir, + ) + monkeypatch.setenv("CONDA_PREFIX", str(conda_prefix)) + monkeypatch.setenv("CUDA_HOME", str(cuda_home)) + monkeypatch.delenv("CUDA_PATH", raising=False) + + assert find_bitcode_lib_by_name(libname, filename) == site_packages_path + os.remove(site_packages_path) + find_bitcode_lib_by_name.cache_clear() + + assert find_bitcode_lib_by_name(libname, filename) == conda_path + os.remove(conda_path) + find_bitcode_lib_by_name.cache_clear() + + assert find_bitcode_lib_by_name(libname, filename) == cuda_home_path + + +@pytest.mark.usefixtures("clear_find_bitcode_lib_cache") +@pytest.mark.agent_authored(model="gpt-5") +def test_find_bitcode_lib_by_name_cache_keeps_filenames_separate(monkeypatch, tmp_path): + lib_dir = _site_packages_bitcode_lib_dir_under(tmp_path / "site-packages", "device") + sm80_path = _make_bitcode_lib_file(lib_dir, "libdevice_sm_80.bc") + sm90_path = _make_bitcode_lib_file(lib_dir, "libdevice_sm_90.bc") + + monkeypatch.setattr( + find_bitcode_lib_module, + "find_sub_dirs_all_sitepackages", + lambda _sub_dir: [str(lib_dir)], + ) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.delenv("CUDA_HOME", raising=False) + monkeypatch.delenv("CUDA_PATH", raising=False) + + assert find_bitcode_lib_by_name("device", "libdevice_sm_80.bc") == sm80_path + assert find_bitcode_lib_by_name("device", "libdevice_sm_90.bc") == sm90_path + + @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") def test_find_bitcode_lib_not_found_error_includes_cuda_home_directory_listing(monkeypatch, tmp_path): cuda_home = tmp_path / "cuda-home" @@ -162,6 +230,32 @@ def test_find_bitcode_lib_not_found_error_includes_cuda_home_directory_listing(m assert "README.txt" in message +@pytest.mark.usefixtures("clear_find_bitcode_lib_cache") +@pytest.mark.agent_authored(model="gpt-5") +def test_find_bitcode_lib_by_name_not_found_error_uses_requested_filename(monkeypatch, tmp_path): + filename = "libdevice_sm_90.bc" + cuda_home = tmp_path / "cuda-home" + lib_dir = _bitcode_lib_dir_under(cuda_home, "device") + lib_dir.mkdir(parents=True) + (lib_dir / "libdevice.10.bc").touch() + + monkeypatch.setattr( + find_bitcode_lib_module, + "find_sub_dirs_all_sitepackages", + lambda _sub_dir: [], + ) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.setenv("CUDA_HOME", str(cuda_home)) + monkeypatch.delenv("CUDA_PATH", raising=False) + + with pytest.raises(BitcodeLibNotFoundError, match=rf'Failure finding "{filename}"') as exc_info: + find_bitcode_lib_by_name("device", filename) + + message = str(exc_info.value) + assert f"No such file: {lib_dir / filename}" in message + assert "libdevice.10.bc" in message + + @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") def test_find_bitcode_lib_not_found_error_without_cuda_home(monkeypatch): monkeypatch.setattr( @@ -183,3 +277,26 @@ def test_find_bitcode_lib_not_found_error_without_cuda_home(monkeypatch): def test_find_bitcode_lib_invalid_name(): with pytest.raises(ValueError, match="Unknown bitcode library"): find_bitcode_lib_module.locate_bitcode_lib("invalid") + + +@pytest.mark.parametrize( + "filename", + ("", ".", "..", "../file.bc", "subdir/file.bc", r"subdir\file.bc", r"C:\lib\file.bc", "bad\0name.bc"), +) +@pytest.mark.agent_authored(model="gpt-5") +def test_find_bitcode_lib_by_name_rejects_paths(filename): + with pytest.raises(ValueError, match="without a directory"): + find_bitcode_lib_by_name("device", filename) + + +@pytest.mark.parametrize("filename", (None, 90)) +@pytest.mark.agent_authored(model="gpt-5") +def test_find_bitcode_lib_by_name_requires_string(filename): + with pytest.raises(TypeError, match="filename must be a string"): + find_bitcode_lib_by_name("device", filename) + + +@pytest.mark.agent_authored(model="gpt-5") +def test_find_bitcode_lib_by_name_invalid_project(): + with pytest.raises(ValueError, match="Unknown bitcode library"): + find_bitcode_lib_by_name("not_a_real_lib", "libdevice_sm_90.bc") From 820ffba9003efff99f826743cc027ed7229948f3 Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Wed, 8 Jul 2026 14:10:09 -0700 Subject: [PATCH 2/3] Add exact-name bitcode locate API --- cuda_pathfinder/cuda/pathfinder/__init__.py | 7 +++-- .../_static_libs/find_bitcode_lib.py | 22 ++++++++++++-- cuda_pathfinder/docs/source/api.rst | 1 + .../tests/test_find_bitcode_lib.py | 30 ++++++++++++++----- 4 files changed, 49 insertions(+), 11 deletions(-) diff --git a/cuda_pathfinder/cuda/pathfinder/__init__.py b/cuda_pathfinder/cuda/pathfinder/__init__.py index 5c155017797..a0ffb622a2f 100644 --- a/cuda_pathfinder/cuda/pathfinder/__init__.py +++ b/cuda_pathfinder/cuda/pathfinder/__init__.py @@ -47,6 +47,9 @@ from cuda.pathfinder._static_libs.find_bitcode_lib import ( locate_bitcode_lib as locate_bitcode_lib, ) +from cuda.pathfinder._static_libs.find_bitcode_lib import ( + locate_bitcode_lib_by_name as locate_bitcode_lib_by_name, +) from cuda.pathfinder._static_libs.find_static_lib import ( SUPPORTED_STATIC_LIBS as _SUPPORTED_STATIC_LIBS, ) @@ -80,8 +83,8 @@ SUPPORTED_BINARY_UTILITIES = _SUPPORTED_BINARIES #: Tuple of supported bitcode library names that can be resolved -#: via ``locate_bitcode_lib()``, ``find_bitcode_lib()``, and -#: ``find_bitcode_lib_by_name()``. +#: via ``locate_bitcode_lib()``, ``locate_bitcode_lib_by_name()``, +#: ``find_bitcode_lib()``, and ``find_bitcode_lib_by_name()``. #: Example value: ``"device"``. SUPPORTED_BITCODE_LIBS = _SUPPORTED_BITCODE_LIBS diff --git a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py index fbfb10befe1..8a671247673 100644 --- a/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py +++ b/cuda_pathfinder/cuda/pathfinder/_static_libs/find_bitcode_lib.py @@ -182,6 +182,25 @@ def locate_bitcode_lib(name: str) -> LocatedBitcodeLib: return _locate_bitcode_lib(name) +def locate_bitcode_lib_by_name(name: str, filename: str) -> LocatedBitcodeLib: + """Locate an exact bitcode filename using a supported library's search paths. + + This lets a library define its own filename convention, including any + architecture or other attributes encoded in the filename. + + Args: + name: Supported bitcode library whose configured directories to search. + filename: Exact file name to locate. Directory components are not allowed. + + Raises: + TypeError: If ``filename`` is not a string. + ValueError: If ``name`` is unsupported or ``filename`` is not a file name. + BitcodeLibNotFoundError: If ``filename`` cannot be found. + """ + _validate_filename(filename) + return _locate_bitcode_lib(name, filename=filename) + + @functools.cache def find_bitcode_lib(name: str) -> str: """Find the absolute path to a bitcode library. @@ -209,5 +228,4 @@ def find_bitcode_lib_by_name(name: str, filename: str) -> str: ValueError: If ``name`` is unsupported or ``filename`` is not a file name. BitcodeLibNotFoundError: If ``filename`` cannot be found. """ - _validate_filename(filename) - return _locate_bitcode_lib(name, filename=filename).abs_path + return locate_bitcode_lib_by_name(name, filename).abs_path diff --git a/cuda_pathfinder/docs/source/api.rst b/cuda_pathfinder/docs/source/api.rst index c5cd45a9e2b..2cc6641b13a 100644 --- a/cuda_pathfinder/docs/source/api.rst +++ b/cuda_pathfinder/docs/source/api.rst @@ -37,6 +37,7 @@ CUDA bitcode and static libraries. find_bitcode_lib find_bitcode_lib_by_name locate_bitcode_lib + locate_bitcode_lib_by_name LocatedBitcodeLib BitcodeLibNotFoundError diff --git a/cuda_pathfinder/tests/test_find_bitcode_lib.py b/cuda_pathfinder/tests/test_find_bitcode_lib.py index c04a4a35ddc..bceddda02cf 100644 --- a/cuda_pathfinder/tests/test_find_bitcode_lib.py +++ b/cuda_pathfinder/tests/test_find_bitcode_lib.py @@ -13,6 +13,7 @@ find_bitcode_lib, find_bitcode_lib_by_name, locate_bitcode_lib, + locate_bitcode_lib_by_name, ) from cuda.pathfinder._utils.env_vars import get_cuda_path_or_home @@ -141,7 +142,7 @@ def find_expected_sub_dir(sub_dir): @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") @pytest.mark.agent_authored(model="gpt-5") -def test_find_bitcode_lib_by_name_search_order(monkeypatch, tmp_path): +def test_bitcode_lib_by_name_search_order(monkeypatch, tmp_path): libname = "device" filename = "libdevice_sm_90.bc" site_packages_lib_dir = _site_packages_bitcode_lib_dir_under(tmp_path / "site-packages", libname) @@ -172,14 +173,26 @@ def find_expected_sub_dir(sub_dir): monkeypatch.setenv("CUDA_HOME", str(cuda_home)) monkeypatch.delenv("CUDA_PATH", raising=False) + located_lib = locate_bitcode_lib_by_name(libname, filename) + assert located_lib.abs_path == site_packages_path + assert located_lib.filename == filename + assert located_lib.found_via == "site-packages" assert find_bitcode_lib_by_name(libname, filename) == site_packages_path os.remove(site_packages_path) find_bitcode_lib_by_name.cache_clear() + located_lib = locate_bitcode_lib_by_name(libname, filename) + assert located_lib.abs_path == conda_path + assert located_lib.filename == filename + assert located_lib.found_via == "conda" assert find_bitcode_lib_by_name(libname, filename) == conda_path os.remove(conda_path) find_bitcode_lib_by_name.cache_clear() + located_lib = locate_bitcode_lib_by_name(libname, filename) + assert located_lib.abs_path == cuda_home_path + assert located_lib.filename == filename + assert located_lib.found_via == "CUDA_PATH" assert find_bitcode_lib_by_name(libname, filename) == cuda_home_path @@ -283,20 +296,23 @@ def test_find_bitcode_lib_invalid_name(): "filename", ("", ".", "..", "../file.bc", "subdir/file.bc", r"subdir\file.bc", r"C:\lib\file.bc", "bad\0name.bc"), ) +@pytest.mark.parametrize("lookup", (find_bitcode_lib_by_name, locate_bitcode_lib_by_name)) @pytest.mark.agent_authored(model="gpt-5") -def test_find_bitcode_lib_by_name_rejects_paths(filename): +def test_bitcode_lib_by_name_rejects_paths(lookup, filename): with pytest.raises(ValueError, match="without a directory"): - find_bitcode_lib_by_name("device", filename) + lookup("device", filename) @pytest.mark.parametrize("filename", (None, 90)) +@pytest.mark.parametrize("lookup", (find_bitcode_lib_by_name, locate_bitcode_lib_by_name)) @pytest.mark.agent_authored(model="gpt-5") -def test_find_bitcode_lib_by_name_requires_string(filename): +def test_bitcode_lib_by_name_requires_string(lookup, filename): with pytest.raises(TypeError, match="filename must be a string"): - find_bitcode_lib_by_name("device", filename) + lookup("device", filename) +@pytest.mark.parametrize("lookup", (find_bitcode_lib_by_name, locate_bitcode_lib_by_name)) @pytest.mark.agent_authored(model="gpt-5") -def test_find_bitcode_lib_by_name_invalid_project(): +def test_bitcode_lib_by_name_invalid_project(lookup): with pytest.raises(ValueError, match="Unknown bitcode library"): - find_bitcode_lib_by_name("not_a_real_lib", "libdevice_sm_90.bc") + lookup("not_a_real_lib", "libdevice_sm_90.bc") From 526ce5388e5caa23113a1ff3ac863c234b9a1371 Mon Sep 17 00:00:00 2001 From: "Ralf W. Grosse-Kunstleve" Date: Wed, 8 Jul 2026 14:13:43 -0700 Subject: [PATCH 3/3] Test unpinned NVSHMEM bitcode layouts --- cuda_pathfinder/pyproject.toml | 4 +- .../tests/test_find_bitcode_lib.py | 98 ++++++++++++++++++- 2 files changed, 99 insertions(+), 3 deletions(-) diff --git a/cuda_pathfinder/pyproject.toml b/cuda_pathfinder/pyproject.toml index 86f33b726e9..667bbf7983a 100644 --- a/cuda_pathfinder/pyproject.toml +++ b/cuda_pathfinder/pyproject.toml @@ -29,7 +29,7 @@ cu12 = [ "nvidia-cusparselt-cu12", "nvidia-libmathdx-cu12", "nvidia-nccl-cu12; sys_platform != 'win32'", - "nvidia-nvshmem-cu12<3.7; sys_platform != 'win32'", + "nvidia-nvshmem-cu12; sys_platform != 'win32'", ] cu13 = [ "cuda-toolkit[nvcc,cublas,nvrtc,cudart,cufft,curand,cusolver,cusparse,npp,nvfatbin,nvjitlink,nvjpeg,cccl,cupti,profiler,nvvm]==13.*", @@ -43,7 +43,7 @@ cu13 = [ "nvidia-cusparselt-cu13", "nvidia-libmathdx-cu13", "nvidia-nccl-cu13; sys_platform != 'win32'", - "nvidia-nvshmem-cu13<3.7; sys_platform != 'win32'", + "nvidia-nvshmem-cu13; sys_platform != 'win32'", ] host = [ # TODO: remove the Python 3.15 guard once 3.15 is officially supported diff --git a/cuda_pathfinder/tests/test_find_bitcode_lib.py b/cuda_pathfinder/tests/test_find_bitcode_lib.py index bceddda02cf..005211e61a8 100644 --- a/cuda_pathfinder/tests/test_find_bitcode_lib.py +++ b/cuda_pathfinder/tests/test_find_bitcode_lib.py @@ -1,7 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import importlib.metadata import os +import re from pathlib import Path import pytest @@ -20,6 +22,10 @@ STRICTNESS = os.environ.get("CUDA_PATHFINDER_TEST_FIND_NVIDIA_BITCODE_LIB_STRICTNESS", "see_what_works") assert STRICTNESS in ("see_what_works", "all_must_work") +_NVSHMEM_DISTRIBUTION_PATTERN = re.compile(r"^nvidia-nvshmem-cu(?:12|13)$", re.IGNORECASE) +_NVSHMEM_LIB_DIR_PARTS = ("nvidia", "nvshmem", "lib") +_NVSHMEM_LEGACY_FILENAME = "libnvshmem_device.bc" + def _bitcode_lib_info(libname: str): return find_bitcode_lib_module._SUPPORTED_BITCODE_LIBS_INFO[libname] @@ -62,6 +68,26 @@ def _conda_anchor(conda_prefix: Path) -> Path: return conda_prefix +def _installed_nvshmem_distributions(): + return tuple( + dist + for dist in importlib.metadata.distributions() + if "Name" in dist.metadata and _NVSHMEM_DISTRIBUTION_PATTERN.fullmatch(dist.metadata["Name"]) + ) + + +def _installed_nvshmem_bitcode_paths(distributions): + paths = { + Path(dist.locate_file(file)) + for dist in distributions + for file in (dist.files or ()) + if file.parts[-4:-1] == _NVSHMEM_LIB_DIR_PARTS + and file.name.startswith("libnvshmem_device") + and file.suffix == ".bc" + } + return tuple(sorted(paths)) + + def _located_bitcode_lib_asserts(located_bitcode_lib): """Common assertions for a located bitcode library.""" assert located_bitcode_lib is not None @@ -74,7 +100,7 @@ def _located_bitcode_lib_asserts(located_bitcode_lib): @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") -@pytest.mark.parametrize("libname", SUPPORTED_BITCODE_LIBS) +@pytest.mark.parametrize("libname", tuple(name for name in SUPPORTED_BITCODE_LIBS if name != "nvshmem_device")) def test_locate_bitcode_lib(info_summary_append, libname): try: located_lib = locate_bitcode_lib(libname) @@ -93,6 +119,40 @@ def test_locate_bitcode_lib(info_summary_append, libname): assert os.path.basename(lib_path) == expected_filename +@pytest.mark.skipif( + find_bitcode_lib_module.IS_WINDOWS, + reason="NVSHMEM wheel test dependencies are not installed on Windows", +) +@pytest.mark.usefixtures("clear_find_bitcode_lib_cache") +@pytest.mark.agent_authored(model="gpt-5") +def test_locate_installed_nvshmem_bitcode(info_summary_append): + distributions = _installed_nvshmem_distributions() + if not distributions: + if STRICTNESS == "all_must_work": + pytest.fail("No nvidia-nvshmem-cu12/13 distribution is installed") + info_summary_append("NVSHMEM distribution not installed") + return + + bitcode_paths = _installed_nvshmem_bitcode_paths(distributions) + versions = ", ".join(f"{dist.metadata['Name']}=={dist.version}" for dist in distributions) + info_summary_append(f"{versions}: {[path.name for path in bitcode_paths]}") + assert bitcode_paths, f"No NVSHMEM device bitcode files found in {versions}" + + for expected_path in bitcode_paths: + located_lib = locate_bitcode_lib_by_name("nvshmem_device", expected_path.name) + assert Path(located_lib.abs_path).samefile(expected_path) + assert located_lib.filename == expected_path.name + assert located_lib.found_via == "site-packages" + assert Path(find_bitcode_lib_by_name("nvshmem_device", expected_path.name)).samefile(expected_path) + + legacy_paths = [path for path in bitcode_paths if path.name == _NVSHMEM_LEGACY_FILENAME] + if legacy_paths: + assert len(legacy_paths) == 1 + located_lib = locate_bitcode_lib("nvshmem_device") + assert Path(located_lib.abs_path).samefile(legacy_paths[0]) + assert Path(find_bitcode_lib("nvshmem_device")).samefile(legacy_paths[0]) + + @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") @pytest.mark.parametrize("libname", SUPPORTED_BITCODE_LIBS) def test_locate_bitcode_lib_search_order(monkeypatch, tmp_path, libname): @@ -140,6 +200,42 @@ def find_expected_sub_dir(sub_dir): assert located_lib.found_via == "CUDA_PATH" +@pytest.mark.skipif( + find_bitcode_lib_module.IS_WINDOWS, + reason="NVSHMEM wheel test dependencies are not installed on Windows", +) +@pytest.mark.usefixtures("clear_find_bitcode_lib_cache") +@pytest.mark.agent_authored(model="gpt-5") +def test_nvshmem_legacy_unversioned_bitcode_filename(monkeypatch, tmp_path): + libname = "nvshmem_device" + assert _bitcode_lib_filename(libname) == _NVSHMEM_LEGACY_FILENAME + + lib_dir = _site_packages_bitcode_lib_dir_under(tmp_path / "site-packages", libname) + expected_path = _make_bitcode_lib_file(lib_dir, _NVSHMEM_LEGACY_FILENAME) + expected_sub_dir = tuple(_bitcode_lib_info(libname)["site_packages_dirs"][0].split("/")) + + def find_expected_sub_dir(sub_dir): + assert sub_dir == expected_sub_dir + return [str(lib_dir)] + + monkeypatch.setattr(find_bitcode_lib_module, "find_sub_dirs_all_sitepackages", find_expected_sub_dir) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.delenv("CUDA_HOME", raising=False) + monkeypatch.delenv("CUDA_PATH", raising=False) + + located_lib = locate_bitcode_lib(libname) + assert located_lib.abs_path == expected_path + assert located_lib.filename == _NVSHMEM_LEGACY_FILENAME + assert located_lib.found_via == "site-packages" + assert find_bitcode_lib(libname) == expected_path + + located_by_name = locate_bitcode_lib_by_name(libname, _NVSHMEM_LEGACY_FILENAME) + assert located_by_name.abs_path == expected_path + assert located_by_name.filename == _NVSHMEM_LEGACY_FILENAME + assert located_by_name.found_via == "site-packages" + assert find_bitcode_lib_by_name(libname, _NVSHMEM_LEGACY_FILENAME) == expected_path + + @pytest.mark.usefixtures("clear_find_bitcode_lib_cache") @pytest.mark.agent_authored(model="gpt-5") def test_bitcode_lib_by_name_search_order(monkeypatch, tmp_path):