Add framework-hinted runtime library discovery to address runtime mis…

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:

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:

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)
 

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(

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,

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__(

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

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"]

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