From 7e9bc6a5d7be3e78bf4fdf042a00d82867a9567f Mon Sep 17 00:00:00 2001 From: Petr Kurapov Date: Mon, 29 Jun 2026 08:01:01 -0500 Subject: [PATCH] Add framework-hinted runtime library discovery to address runtime mismatch on dual-vendor systems. Signed-off-by: Petr Kurapov --- fastsafetensors/common.py | 32 +++++++--- fastsafetensors/copier/dstorage.py | 5 +- fastsafetensors/copier/gds.py | 12 ++-- fastsafetensors/copier/nogds.py | 29 ++++++--- fastsafetensors/copier/unified.py | 2 +- fastsafetensors/threefs_loader.py | 4 +- tests/unit/conftest.py | 4 +- tests/unit/test_runtime_hint.py | 99 ++++++++++++++++++++++++++++++ tests/unit/threefs/conftest.py | 4 +- 9 files changed, 163 insertions(+), 28 deletions(-) create mode 100644 tests/unit/test_runtime_hint.py diff --git a/fastsafetensors/common.py b/fastsafetensors/common.py index bc3148f..f179417 100644 --- a/fastsafetensors/common.py +++ b/fastsafetensors/common.py @@ -69,16 +69,34 @@ def _normalize_windows_dll_path(path: str, source: str) -> str: return normalized -def resolve_cudart_lib_name() -> str: - """Resolve the CUDA runtime library name for the current platform. +def resolve_runtime_lib_name(framework=None) -> str: + """Resolve the GPU runtime library to dlopen for the current platform. - Returns: - On Windows, an absolute DLL path string. On other platforms, "" to use - the compiled-in default. + On Windows, returns an absolute cudart DLL path. On other platforms, maps the framework's declared GPU + vendor to a runtime library so the dlopen'd vendor stays in sync with the + framework's GPU build. Returns "" when there is no usable hint so the caller falls + back to auto-detection. """ - if sys.platform != "win32": - return "" # Non-Windows: use auto-detection (CUDA first, then ROCm) + if sys.platform == "win32": + return _resolve_windows_cudart_lib_name() + if framework is None: + return "" + try: + ver = framework.get_cuda_ver() + except Exception: + return "" + if not ver or "-" not in ver: + return "" + vendor = ver.split("-", 1)[0] + if vendor == "hip": + return "libamdhip64.so" + if vendor == "cuda": + return "libcudart.so" + return "" + +def _resolve_windows_cudart_lib_name() -> str: + """Resolve the absolute cudart DLL path on Windows, "" for the default.""" # Allow explicit override via environment variable override = os.environ.get("FASTSAFETENSORS_CUDART_LIB", "").strip() if override: diff --git a/fastsafetensors/copier/dstorage.py b/fastsafetensors/copier/dstorage.py index 849627c..618474f 100644 --- a/fastsafetensors/copier/dstorage.py +++ b/fastsafetensors/copier/dstorage.py @@ -11,7 +11,7 @@ SafeTensorsMetadata, init_logger, is_gpu_found, - resolve_cudart_lib_name, + resolve_runtime_lib_name, ) from ..frameworks import FrameworkOpBase, TensorBase from ..st_types import Device, DeviceType, DType @@ -114,7 +114,8 @@ def init_dstorage(device_id: int = 0) -> None: load_library_func() if not is_gpu_found(): raise RuntimeError("CUDA runtime not found") - cudart_dll = resolve_cudart_lib_name() + # Windows-only; resolve_runtime_lib_name() returns the cudart DLL path here + cudart_dll = resolve_runtime_lib_name() if not cudart_dll: raise RuntimeError("Could not find CUDA runtime DLL") if _dstorage_dll_dir is None: diff --git a/fastsafetensors/copier/gds.py b/fastsafetensors/copier/gds.py index 583f867..246211b 100644 --- a/fastsafetensors/copier/gds.py +++ b/fastsafetensors/copier/gds.py @@ -165,8 +165,8 @@ def wait_io( _inited_gds = False -def init_gds(): - load_library_func() +def init_gds(framework: Optional[FrameworkOpBase] = None): + load_library_func(framework) global _inited_gds if not _inited_gds: if fstcpp.init_gds() != 0: @@ -181,7 +181,7 @@ def new_gds_file_copier( max_threads: int = 16, **kwargs, ) -> CopierConstructFunc: - init_gds() + init_gds(kwargs.get("framework")) device_is_not_cpu = device.type != DeviceType.CPU if device_is_not_cpu and not is_gpu_found(): raise Exception( @@ -215,8 +215,10 @@ def new_gds_file_copier( from .unified import is_unified_memory_system, new_unified_copier if device_is_not_cpu and is_unified_memory_system(kwargs.get("framework")): - return new_unified_copier(device) - return new_nogds_file_copier(device, bbuf_size_kb, max_threads) + return new_unified_copier(device, framework=kwargs.get("framework")) + return new_nogds_file_copier( + device, bbuf_size_kb, max_threads, framework=kwargs.get("framework") + ) reader = fstcpp.gds_file_reader(max_threads, device_is_not_cpu, device_id) diff --git a/fastsafetensors/copier/nogds.py b/fastsafetensors/copier/nogds.py index 813e177..96d8e3f 100644 --- a/fastsafetensors/copier/nogds.py +++ b/fastsafetensors/copier/nogds.py @@ -5,7 +5,11 @@ from typing import Dict, List from .. import cpp as fstcpp -from ..common import SafeTensorsMetadata, is_gpu_found, resolve_cudart_lib_name +from ..common import ( + SafeTensorsMetadata, + is_gpu_found, + resolve_runtime_lib_name, +) from ..frameworks import FrameworkOpBase, TensorBase from ..st_types import Device, DeviceType, DType from .base import CopierInterface @@ -81,12 +85,23 @@ def wait_io( _loaded_library = False -def load_library_func(): +def load_library_func(framework=None): global _loaded_library - if not _loaded_library: - cudart_lib = resolve_cudart_lib_name() - fstcpp.load_library_functions(cudart_lib) - _loaded_library = True + if _loaded_library: + return + + lib = resolve_runtime_lib_name(framework) + fstcpp.load_library_functions(lib) + if lib and not is_gpu_found(): + # The framework hinted a specific vendor's runtime but loading it found + # no GPU. A GPU-built framework only reports a vendor when it sees a + # device, so this is a real mismatch (wrong/missing runtime for that + # vendor). + raise Exception( + f"[FAIL] framework hinted GPU runtime '{lib}' but no GPU was found " + "after loading it (runtime/devices for that vendor not present)" + ) + _loaded_library = True @register_copier_constructor("nogds") @@ -96,7 +111,7 @@ def new_nogds_file_copier( max_threads: int = 16, **kwargs, ) -> CopierConstructFunc: - load_library_func() + load_library_func(kwargs.get("framework")) device_is_not_cpu = device.type != DeviceType.CPU if device_is_not_cpu and not is_gpu_found(): raise Exception( diff --git a/fastsafetensors/copier/unified.py b/fastsafetensors/copier/unified.py index 5d65e77..c86e712 100644 --- a/fastsafetensors/copier/unified.py +++ b/fastsafetensors/copier/unified.py @@ -122,7 +122,7 @@ def new_unified_copier(device: Device, **kwargs) -> CopierConstructFunc: """ from .nogds import load_library_func - load_library_func() + load_library_func(kwargs.get("framework")) def construct_unified_copier( metadata: SafeTensorsMetadata, diff --git a/fastsafetensors/threefs_loader.py b/fastsafetensors/threefs_loader.py index edb9bde..8a5f32d 100644 --- a/fastsafetensors/threefs_loader.py +++ b/fastsafetensors/threefs_loader.py @@ -3,7 +3,7 @@ from typing import Any, Dict, List, Mapping, Optional from . import cpp as fstcpp -from .common import init_logger, resolve_cudart_lib_name +from .common import init_logger, resolve_runtime_lib_name from .frameworks import get_framework_op from .loader import BaseSafeTensorsFileLoader, loaded_library from .parallel_loader import PipelineParallel @@ -54,7 +54,7 @@ def __init__( global loaded_library if not loaded_library: - fstcpp.load_library_functions(resolve_cudart_lib_name()) + fstcpp.load_library_functions(resolve_runtime_lib_name()) loaded_library = True fstcpp.set_debug_log(debug_log) super().__init__( diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 4bfc87e..0872f34 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -7,7 +7,7 @@ from fastsafetensors import SingleGroup from fastsafetensors import cpp as fstcpp -from fastsafetensors.common import is_gpu_found, resolve_cudart_lib_name +from fastsafetensors.common import is_gpu_found, resolve_runtime_lib_name from fastsafetensors.cpp import load_library_functions from fastsafetensors.frameworks import FrameworkOpBase, get_framework_op from fastsafetensors.st_types import Device @@ -25,7 +25,7 @@ os.makedirs(TMP_DIR, 0o777, True) os.makedirs(GENERATED_DIR, 0o777, True) -load_library_functions(resolve_cudart_lib_name()) +load_library_functions(resolve_runtime_lib_name()) FRAMEWORK = get_framework_op(os.getenv("TEST_FASTSAFETENSORS_FRAMEWORK", "please set")) # Print platform information at test startup diff --git a/tests/unit/test_runtime_hint.py b/tests/unit/test_runtime_hint.py new file mode 100644 index 0000000..5d265fe --- /dev/null +++ b/tests/unit/test_runtime_hint.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the framework-hinted GPU runtime selection.""" + +import sys + +import pytest + +from fastsafetensors import common + + +class _FakeFramework: + def __init__(self, cuda_ver): + self._ver = cuda_ver + + def get_cuda_ver(self): + if isinstance(self._ver, Exception): + raise self._ver + return self._ver + + +@pytest.fixture(autouse=True) +def _force_non_windows(monkeypatch): + # The hint is intentionally a no-op on Windows (cudart resolver owns it). + monkeypatch.setattr(sys, "platform", "linux") + + +def test_none_framework_uses_autodetect(): + assert common.resolve_runtime_lib_name(None) == "" + + +def test_hip_framework_selects_amdhip(): + assert ( + common.resolve_runtime_lib_name(_FakeFramework("hip-7.2.0")) == "libamdhip64.so" + ) + + +def test_cuda_framework_selects_cudart(): + assert ( + common.resolve_runtime_lib_name(_FakeFramework("cuda-12.1")) == "libcudart.so" + ) + + +@pytest.mark.parametrize("ver", ["", "weird", "rocm-7.0"]) +def test_unknown_vendor_uses_autodetect(ver): + assert common.resolve_runtime_lib_name(_FakeFramework(ver)) == "" + + +def test_get_cuda_ver_raises_uses_autodetect(): + assert common.resolve_runtime_lib_name(_FakeFramework(RuntimeError("boom"))) == "" + + +def test_windows_is_noop(monkeypatch): + monkeypatch.setattr(sys, "platform", "win32") + assert common.resolve_runtime_lib_name(_FakeFramework("hip-7.2.0")) == "" + + +def test_load_library_func_hint_with_no_gpu_raises(monkeypatch): + """A hint that finds no GPU is a hard failure""" + from fastsafetensors.copier import nogds + + calls = [] + + def fake_load(lib): + calls.append(lib) + + monkeypatch.setattr(nogds.fstcpp, "load_library_functions", fake_load) + monkeypatch.setattr( + nogds, "resolve_runtime_lib_name", lambda fw=None: "libamdhip64.so" + ) + + monkeypatch.setattr(nogds, "is_gpu_found", lambda: False) + monkeypatch.setattr(nogds, "_loaded_library", False) + + with pytest.raises(Exception, match="libamdhip64.so"): + nogds.load_library_func(_FakeFramework("hip-7.2.0")) + + assert calls == ["libamdhip64.so"] + assert nogds._loaded_library is False + + +def test_load_library_func_hint_succeeds_no_fallback(monkeypatch): + from fastsafetensors.copier import nogds + + calls = [] + monkeypatch.setattr( + nogds.fstcpp, "load_library_functions", lambda lib: calls.append(lib) + ) + monkeypatch.setattr( + nogds, + "resolve_runtime_lib_name", + lambda fw=None: "libamdhip64.so" if fw is not None else "", + ) + monkeypatch.setattr(nogds, "is_gpu_found", lambda: True) + monkeypatch.setattr(nogds, "_loaded_library", False) + + nogds.load_library_func(_FakeFramework("hip-7.2.0")) + + assert calls == ["libamdhip64.so"] diff --git a/tests/unit/threefs/conftest.py b/tests/unit/threefs/conftest.py index 76b371c..a67385e 100644 --- a/tests/unit/threefs/conftest.py +++ b/tests/unit/threefs/conftest.py @@ -14,7 +14,7 @@ from fastsafetensors import SingleGroup from fastsafetensors import cpp as fstcpp -from fastsafetensors.common import is_gpu_found, resolve_cudart_lib_name +from fastsafetensors.common import is_gpu_found, resolve_runtime_lib_name from fastsafetensors.cpp import load_library_functions from fastsafetensors.frameworks import FrameworkOpBase, get_framework_op from fastsafetensors.st_types import Device @@ -34,7 +34,7 @@ def mock_3fs_reader(): yield -load_library_functions(resolve_cudart_lib_name()) +load_library_functions(resolve_runtime_lib_name()) FRAMEWORK = get_framework_op(os.getenv("TEST_FASTSAFETENSORS_FRAMEWORK", "please set"))