fastsafetensor-3fs-reader 0.3.3__py3-none-any.whl

fastsafetensor_3fs_reader/__init__.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""fastsafetensor_3fs_reader -- 3FS USRBIO file reader for fastsafetensors.
+
+Quick start::
+
+    from fastsafetensor_3fs_reader import ThreeFSFileReader, is_available
+
+    if is_available():
+        reader = ThreeFSFileReader(mount_point="/mnt/3fs")
+        headers = reader.read_headers_batch(["/mnt/3fs/model.safetensors"])
+        reader.close()
+
+Backend auto-selection (override via ``FASTSAFETENSORS_BACKEND``)::
+
+    cpp -> python -> mock
+"""
+
+from ._lib_preload import get_hf3fs_lib_path, preload_hf3fs_library
+
+preload_hf3fs_library()  # must run before any backend import
+
+from ._mount_utils import extract_mount_point
+from .interface import FileReaderInterface
+from .mock import MockFileReader
+
+from ._backend import (  # noqa: E402
+    create_reader,
+    get_backend,
+    init_backend,
+    is_available,
+)
+
+# init_backend() must run BEFORE importing ThreeFSFileReader: Python's
+# ``from mod import name`` captures the value at import time.
+init_backend()
+
+from ._backend import ThreeFSFileReader  # noqa: E402
+
+__all__ = [
+    "FileReaderInterface",
+    "ThreeFSFileReader",
+    "MockFileReader",
+    "extract_mount_point",
+    "get_hf3fs_lib_path",
+    "is_available",
+    "get_backend",
+    "create_reader",
+]

fastsafetensor_3fs_reader/_backend.py

@@ -0,0 +1,121 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Backend selection: ``FASTSAFETENSORS_BACKEND`` → cpp / python / mock."""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+from .interface import FileReaderInterface
+from .mock import MockFileReader
+
+logger = logging.getLogger(__name__)
+
+_VALID_BACKENDS = ("cpp", "python", "mock", "auto")
+
+_BACKEND: str = "mock"
+ThreeFSFileReader: type | None = None
+
+
+def _load_backend(name: str) -> None:
+    global ThreeFSFileReader, _BACKEND
+    if name == "cpp":
+        from .reader_cpp import ThreeFSFileReaderCpp
+
+        ThreeFSFileReader = ThreeFSFileReaderCpp
+        _BACKEND = "cpp"
+    elif name == "python":
+        from .reader_py import ThreeFSFileReaderPy
+
+        ThreeFSFileReader = ThreeFSFileReaderPy
+        _BACKEND = "python"
+    elif name == "mock":
+        ThreeFSFileReader = MockFileReader
+        _BACKEND = "mock"
+    else:
+        raise ValueError(f"Unknown backend: {name!r}")
+
+
+def init_backend() -> None:
+    """Auto-select backend (cpp → python → mock).
+
+    Override with ``FASTSAFETENSORS_BACKEND=cpp|python|mock``.
+    """
+    forced = os.environ.get("FASTSAFETENSORS_BACKEND", "").lower().strip()
+    if forced and forced not in _VALID_BACKENDS:
+        raise ValueError(
+            f"FASTSAFETENSORS_BACKEND={forced!r} is invalid. "
+            f"Valid values: {', '.join(_VALID_BACKENDS)} (or unset)"
+        )
+
+    if forced and forced != "auto":
+        _load_backend(forced)
+        logger.info(
+            "using backend=%r (forced via FASTSAFETENSORS_BACKEND)",
+            _BACKEND,
+        )
+    else:
+        for candidate in ("cpp", "python"):
+            try:
+                _load_backend(candidate)
+                logger.info(
+                    "using backend=%r (auto-selected)", _BACKEND
+                )
+                break
+            except ImportError as exc:
+                logger.debug(
+                    "backend=%r not available (%s), trying next",
+                    candidate,
+                    exc,
+                )
+
+        if ThreeFSFileReader is None:
+            _load_backend("mock")
+            logger.warning(
+                "no real 3FS backend available "
+                "(cpp/python both failed), falling back to mock backend"
+            )
+
+
+def is_available() -> bool:
+    return _BACKEND in ("cpp", "python")
+
+
+def get_backend() -> str:
+    return _BACKEND
+
+
+def create_reader(backend: str = "auto", **kwargs: Any) -> FileReaderInterface:
+    """Create a reader instance, optionally forcing a specific backend.
+
+    ``**kwargs`` are forwarded to the reader constructor.
+    """
+    if backend == "auto":
+        if ThreeFSFileReader is None:
+            raise RuntimeError("No backend is available")
+        return ThreeFSFileReader(**kwargs)
+    elif backend == "cpp":
+        from .reader_cpp import ThreeFSFileReaderCpp
+
+        return ThreeFSFileReaderCpp(**kwargs)
+    elif backend == "python":
+        from .reader_py import ThreeFSFileReaderPy
+
+        return ThreeFSFileReaderPy(**kwargs)
+    elif backend == "mock":
+        return MockFileReader(**kwargs)
+    else:
+        raise ValueError(
+            f"backend={backend!r} is invalid. Valid values: {', '.join(_VALID_BACKENDS)}"
+        )
+
+
+__all__ = [
+    "ThreeFSFileReader",
+    "init_backend",
+    "is_available",
+    "get_backend",
+    "create_reader",
+]

fastsafetensor_3fs_reader/_cuda_utils.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""CUDA Runtime helpers for safe host-to-target memory copies via ctypes."""
+
+import ctypes
+import ctypes.util
+
+_cudart_lib = None
+
+class _CudaPointerAttributes(ctypes.Structure):
+    _fields_ = [
+        ("type", ctypes.c_int),
+        ("device", ctypes.c_int),
+        ("devicePointer", ctypes.c_void_p),
+        ("hostPointer", ctypes.c_void_p),
+    ]
+
+def _load_cudart():
+    global _cudart_lib
+    if _cudart_lib is not None:
+        return _cudart_lib
+
+    for name in ("libcudart.so", "libcudart.so.12", "libcudart.so.11"):
+        try:
+            _cudart_lib = ctypes.CDLL(name)
+            return _cudart_lib
+        except OSError:
+            continue
+
+    path = ctypes.util.find_library("cudart")
+    if path:
+        try:
+            _cudart_lib = ctypes.CDLL(path)
+            return _cudart_lib
+        except OSError:
+            pass
+
+    return None
+
+def _get_cuda_ptr_type(ptr: int) -> str:
+    """Return the CUDA memory type of *ptr*.
+
+    Returns:
+        ``'device'``  -- confirmed CUDA device or managed memory.
+        ``'host'``    -- confirmed host (pageable or pinned) memory.
+        ``'unknown'`` -- ``cudaPointerGetAttributes`` failed.
+    """
+    lib = _load_cudart()
+    if lib is None:
+        return "unknown"
+    try:
+        attrs = _CudaPointerAttributes()
+        err = lib.cudaPointerGetAttributes(ctypes.byref(attrs), ctypes.c_void_p(ptr))
+        if err == 0:
+            # cudaMemoryTypeDevice == 2, cudaMemoryTypeManaged == 3
+            if attrs.type in (2, 3):
+                return "device"
+            else:
+                return "host"
+        else:
+            # Clear the sticky CUDA error so subsequent calls are not affected.
+            if hasattr(lib, "cudaGetLastError"):
+                lib.cudaGetLastError()
+    except Exception:
+        pass
+    return "unknown"
+
+def _is_cuda_ptr(ptr: int) -> bool:
+    """Kept for backward compatibility; prefer ``_get_cuda_ptr_type``."""
+    return _get_cuda_ptr_type(ptr) == "device"
+
+def _cuda_memcpy(dst: int, src: int, nbytes: int, kind: int = 1) -> None:
+    """Call ``cudaMemcpy`` via ctypes.
+
+    Args:
+        kind:   ``cudaMemcpyKind`` value.
+                ``1`` = ``cudaMemcpyHostToDevice`` (default).
+                ``4`` = ``cudaMemcpyDefault`` -- CUDA runtime infers
+                src/dst memory types at call time.
+
+    Raises:
+        RuntimeError: If ``cudaMemcpy`` returns a non-zero error code.
+        OSError: If ``libcudart.so`` cannot be loaded.
+    """
+    lib = _load_cudart()
+    if lib is None:
+        raise OSError("libcudart.so not found; cannot perform cudaMemcpy")
+    err = lib.cudaMemcpy(
+        ctypes.c_void_p(dst),
+        ctypes.c_void_p(src),
+        ctypes.c_size_t(nbytes),
+        ctypes.c_int(kind),
+    )
+    if err != 0:
+        raise RuntimeError(f"cudaMemcpy failed with CUDA error code {err}")
+
+_cudaMemcpy_fn = None
+_cudaHostRegister_fn = None
+_cudaHostUnregister_fn = None
+
+def _get_fast_cudaMemcpy():
+    """Return a cached ctypes function pointer for ``cudaMemcpy``."""
+    global _cudaMemcpy_fn
+    if _cudaMemcpy_fn is None:
+        lib = _load_cudart()
+        if lib is not None:
+            fn = lib.cudaMemcpy
+            fn.restype = ctypes.c_int
+            fn.argtypes = [
+                ctypes.c_void_p,
+                ctypes.c_void_p,
+                ctypes.c_size_t,
+                ctypes.c_int,
+            ]
+            _cudaMemcpy_fn = fn
+    return _cudaMemcpy_fn
+
+def _fast_cuda_memcpy(dst: int, src: int, nbytes: int, kind: int = 4) -> None:
+    """Call ``cudaMemcpy`` using a cached function pointer.
+
+    Args:
+        kind:   ``cudaMemcpyKind``.  Default ``4`` = ``cudaMemcpyDefault``.
+
+    Raises:
+        OSError:      If ``libcudart.so`` cannot be loaded.
+        RuntimeError: If ``cudaMemcpy`` returns a non-zero error code.
+    """
+    fn = _get_fast_cudaMemcpy()
+    if fn is None:
+        raise OSError("libcudart.so not found; cannot perform cudaMemcpy")
+    err = fn(dst, src, nbytes, kind)
+    if err != 0:
+        raise RuntimeError(f"cudaMemcpy failed with CUDA error code {err}")
+
+def _cuda_host_register(ptr: int, size: int) -> bool:
+    """Register a host memory region as CUDA pinned memory.
+
+    Returns:
+        ``True`` if registration succeeded, ``False`` otherwise.
+    """
+    lib = _load_cudart()
+    if lib is None:
+        return False
+    try:
+        fn = lib.cudaHostRegister
+        fn.restype = ctypes.c_int
+        fn.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint]
+        err = fn(ctypes.c_void_p(ptr), ctypes.c_size_t(size), ctypes.c_uint(0))
+        if err != 0:
+            lib.cudaGetLastError()
+            return False
+        return True
+    except Exception:
+        return False
+
+def _cuda_host_unregister(ptr: int) -> None:
+    """Unregister a previously pinned host memory region.
+
+    Failure is silently ignored: the CUDA context may already be destroyed
+    at process exit, in which case CUDA cleans up automatically.
+    """
+    lib = _load_cudart()
+    if lib is None:
+        return
+    try:
+        fn = lib.cudaHostUnregister
+        fn.restype = ctypes.c_int
+        fn.argtypes = [ctypes.c_void_p]
+        fn(ctypes.c_void_p(ptr))
+        lib.cudaGetLastError()
+    except Exception:
+        pass
+
+def _copy_host_to_target(
+    staging_buf: bytearray,
+    staging_ptr: int,
+    target_ptr: int,
+    nbytes: int,
+) -> None:
+    """Copy *nbytes* from *staging_buf* to *target_ptr*.
+
+    * No-op if *target_ptr* is 0.
+    * Uses ``cudaMemcpyHostToDevice`` for confirmed GPU pointers.
+    * Uses ``cudaMemcpyDefault`` when pointer type is unknown.
+    * Falls back to ``ctypes.memmove`` if CUDA is unavailable.
+    """
+    if target_ptr == 0 or nbytes == 0:
+        return
+
+    ptr_type = _get_cuda_ptr_type(target_ptr)
+
+    if ptr_type == "device":
+        _cuda_memcpy(target_ptr, staging_ptr, nbytes, kind=1)
+    else:
+        try:
+            _cuda_memcpy(target_ptr, staging_ptr, nbytes, kind=4)
+            return
+        except (RuntimeError, OSError):
+            pass
+        ctypes.memmove(target_ptr, staging_ptr, nbytes)
+
+def _copy_target_to_host(
+    src_ptr: int,
+    dst_buf: bytearray,
+    nbytes: int,
+) -> None:
+    """Inverse of ``_copy_host_to_target``: copy from *src_ptr* to *dst_buf*."""
+    if src_ptr == 0 or nbytes == 0:
+        return
+
+    if len(dst_buf) < nbytes:
+        raise ValueError(f"dst_buf too small: {len(dst_buf)} < {nbytes}")
+
+    dst_ptr = ctypes.addressof((ctypes.c_char * len(dst_buf)).from_buffer(dst_buf))
+
+    ptr_type = _get_cuda_ptr_type(src_ptr)
+
+    if ptr_type == "device":
+        _cuda_memcpy(dst_ptr, src_ptr, nbytes, kind=2)
+    else:
+        try:
+            _cuda_memcpy(dst_ptr, src_ptr, nbytes, kind=4)
+            return
+        except (RuntimeError, OSError):
+            pass
+        ctypes.memmove(dst_buf, src_ptr, nbytes)

fastsafetensor_3fs_reader/_lib_preload.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Auto-discover and preload ``libhf3fs_api_shared.so`` with ``RTLD_GLOBAL``.
+
+Pre-loads the library so that the C++ extension (``_core_v2``) and the
+Python backend (``hf3fs_fuse.io``) can resolve symbols at import time
+without requiring the user to set ``LD_LIBRARY_PATH``.
+
+Discovery priority (first match wins):
+    1. ``HF3FS_LIB_DIR`` environment variable.
+    2. ``LD_LIBRARY_PATH`` directories.
+    3. ``hf3fs_py_usrbio`` pip install path.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import importlib.util
+import logging
+import os
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+_LIB_NAME = "libhf3fs_api_shared.so"
+_preloaded_path: str | None = None
+
+
+def _glob_lib(directory: Path) -> str | None:
+    """Return the first matching ``libhf3fs_api_shared.so*`` in *directory*."""
+    if not directory.is_dir():
+        return None
+    for entry in directory.glob(f"{_LIB_NAME}*"):
+        if entry.is_file():
+            return str(entry)
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Discovery helpers
+# ---------------------------------------------------------------------------
+
+def _find_in_env_var() -> str | None:
+    lib_dir = os.environ.get("HF3FS_LIB_DIR")
+    if not lib_dir:
+        return None
+    candidate = os.path.join(lib_dir, _LIB_NAME)
+    if os.path.isfile(candidate):
+        return candidate
+    return _glob_lib(Path(lib_dir))
+
+
+def _find_in_ld_library_path() -> str | None:
+    ld_path = os.environ.get("LD_LIBRARY_PATH", "")
+    if not ld_path:
+        return None
+    for directory in ld_path.split(os.pathsep):
+        directory = directory.strip()
+        if not directory:
+            continue
+        candidate = os.path.join(directory, _LIB_NAME)
+        if os.path.isfile(candidate):
+            return candidate
+        found = _glob_lib(Path(directory))
+        if found:
+            return found
+    return None
+
+
+def _find_in_pip_packages() -> str | None:
+    """Locate the library relative to ``hf3fs_py_usrbio``'s install path.
+
+    Handles two install layouts:
+    * Package directory: ``site-packages/hf3fs_py_usrbio/__init__.py``
+    * Single-file C extension: ``site-packages/hf3fs_py_usrbio.cpython-*.so``
+    """
+    try:
+        spec = importlib.util.find_spec("hf3fs_py_usrbio")
+    except (ModuleNotFoundError, ValueError):
+        return None
+
+    if spec is None or spec.origin is None:
+        return None
+
+    origin = Path(spec.origin)
+
+    if origin.name == "__init__.py" or origin.suffix == ".py":
+        pkg_dir = origin.parent
+        site_packages = pkg_dir.parent
+    else:
+        # Single-file .so sitting directly in site-packages
+        pkg_dir = None
+        site_packages = origin.parent
+
+    candidates = [
+        site_packages / "hf3fs_py_usrbio.libs",   # auditwheel convention
+        site_packages,
+    ]
+    if pkg_dir is not None:
+        candidates.extend([pkg_dir / "lib", pkg_dir])
+
+    for directory in candidates:
+        found = _glob_lib(directory)
+        if found:
+            return found
+
+    return None
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+def preload_hf3fs_library() -> bool:
+    """Discover and pre-load ``libhf3fs_api_shared.so`` with ``RTLD_GLOBAL``.
+
+    Returns ``True`` if the library was loaded (or already loaded), ``False``
+    otherwise.
+    """
+    global _preloaded_path
+
+    if _preloaded_path is not None:
+        return True
+
+    for finder, source in (
+        (_find_in_env_var, "HF3FS_LIB_DIR"),
+        (_find_in_ld_library_path, "LD_LIBRARY_PATH"),
+        (_find_in_pip_packages, "hf3fs_py_usrbio pip package"),
+    ):
+        path = finder()
+        if path:
+            return _do_preload(path, source=source)
+
+    logger.debug(
+        "preload_hf3fs_library: %s not found in any search path", _LIB_NAME
+    )
+    return False
+
+
+def get_hf3fs_lib_path() -> str | None:
+    return _preloaded_path
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+def _do_preload(path: str, *, source: str) -> bool:
+    global _preloaded_path
+    try:
+        ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+        _preloaded_path = path
+        logger.info(
+            "preload_hf3fs_library: loaded %s from %s (source: %s)",
+            _LIB_NAME,
+            path,
+            source,
+        )
+        return True
+    except OSError as exc:
+        logger.warning(
+            "preload_hf3fs_library: failed to load %s from %s (source: %s): %s",
+            _LIB_NAME,
+            path,
+            source,
+            exc,
+        )
+        return False
+
+
+__all__ = ["preload_hf3fs_library", "get_hf3fs_lib_path"]

fastsafetensor_3fs_reader/_mount_utils.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for resolving the 3FS mount point from file paths."""
+
+from __future__ import annotations
+
+try:
+    from hf3fs_py_usrbio import extract_mount_point
+
+except ImportError:
+    # No-op stub when hf3fs_py_usrbio is not installed.
+    def extract_mount_point(path: str) -> str:  # type: ignore[misc]
+        return ""
+
+
+__all__ = ["extract_mount_point"]