unifiedefficientloader 0.2.1__tar.gz → 0.2.3__tar.gz

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unifiedefficientloader
-Version: 0.2.1
+Version: 0.2.3
 Summary: A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts.
 Author: silveroxides
 License: MIT
@@ -111,6 +111,33 @@ with UnifiedSafetensorsLoader("model.safetensors", low_memory=True) as loader:
             loader.mark_processed(key)
 ```
 
+### Direct-to-GPU Streaming (Zero-Copy)
+
+For the absolute fastest loading times on CUDA devices, use the `direct_gpu=True` flag. This creates a pipeline that pre-allocates pinned memory pools and GPU memory slabs. Tensors are loaded from disk directly into pinned buffers, and immediately asynchronously copied to the GPU using CUDA streams, hiding the PCIe transfer latency completely behind the disk I/O.
+
+```python
+from unifiedefficientloader import UnifiedSafetensorsLoader
+
+with UnifiedSafetensorsLoader("model.safetensors", low_memory=True, direct_gpu=True) as loader:
+    keys_to_load = loader.keys()
+    
+    # async_stream will automatically coordinate disk -> pinned buffer -> GPU slab -> tensor header
+    stream = loader.async_stream(
+        keys_to_load,
+        batch_size=8,
+        prefetch_batches=2,
+        direct_gpu=True # optional here since we passed it in __init__
+    )
+    
+    for batch in stream:
+        for key, gpu_tensor in batch:
+            # gpu_tensor is already on the GPU!
+            assert gpu_tensor.device.type == "cuda"
+            
+            # ... process gpu_tensor ...
+            loader.mark_processed(key)
+```
+
 ### Tensor/Dict Conversion
 
 ```python

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/README.md

@@ -85,6 +85,33 @@ with UnifiedSafetensorsLoader("model.safetensors", low_memory=True) as loader:
             loader.mark_processed(key)
 ```
 
+### Direct-to-GPU Streaming (Zero-Copy)
+
+For the absolute fastest loading times on CUDA devices, use the `direct_gpu=True` flag. This creates a pipeline that pre-allocates pinned memory pools and GPU memory slabs. Tensors are loaded from disk directly into pinned buffers, and immediately asynchronously copied to the GPU using CUDA streams, hiding the PCIe transfer latency completely behind the disk I/O.
+
+```python
+from unifiedefficientloader import UnifiedSafetensorsLoader
+
+with UnifiedSafetensorsLoader("model.safetensors", low_memory=True, direct_gpu=True) as loader:
+    keys_to_load = loader.keys()
+    
+    # async_stream will automatically coordinate disk -> pinned buffer -> GPU slab -> tensor header
+    stream = loader.async_stream(
+        keys_to_load,
+        batch_size=8,
+        prefetch_batches=2,
+        direct_gpu=True # optional here since we passed it in __init__
+    )
+    
+    for batch in stream:
+        for key, gpu_tensor in batch:
+            # gpu_tensor is already on the GPU!
+            assert gpu_tensor.device.type == "cuda"
+            
+            # ... process gpu_tensor ...
+            loader.mark_processed(key)
+```
+
 ### Tensor/Dict Conversion
 
 ```python

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "unifiedefficientloader"
-version = "0.2.1"
+version = "0.2.3"
 description = "A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts."
 readme = "README.md"
 authors = [
@@ -34,4 +34,9 @@ log_cli_format = "%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(linen
 log_cli_date_format = "%Y-%m-%d %H:%M:%S"
 filterwarnings = [
     "ignore:.*argument 'device' of Tensor.*:DeprecationWarning"
-]
+]
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["unifiedefficientloader*"]
+exclude = ["reference"]

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/setup.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 """Minimal setup.py for backward compatibility with legacy pip install workflows."""
 from setuptools import setup
 

unifiedefficientloader-0.2.3/tests/test_direct_gpu.py

@@ -0,0 +1,95 @@
+import os
+import tempfile
+import pytest
+
+try:
+    import torch
+    from safetensors.torch import save_file
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+from unifiedefficientloader import MemoryEfficientSafeOpen
+
+@pytest.fixture
+def sample_safetensors():
+    if not HAS_TORCH:
+        pytest.skip("Requires torch and safetensors")
+        
+    with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as f:
+        path = f.name
+        
+    tensors = {
+        "weight1": torch.randn(10, 10),
+        "weight2": torch.randn(20, 20),
+        "bias": torch.zeros(10),
+    }
+    save_file(tensors, path)
+    
+    yield path, tensors
+    
+    if os.path.exists(path):
+        os.remove(path)
+
+@pytest.mark.skipif(not HAS_TORCH or not torch.cuda.is_available(), reason="Requires CUDA")
+def test_direct_gpu_streaming(sample_safetensors):
+    path, original_tensors = sample_safetensors
+    
+    loader = MemoryEfficientSafeOpen(path, low_memory=True, direct_gpu=True)
+    
+    # Test load_all which uses async_stream under the hood
+    loaded_tensors = loader.load_all()
+    
+    for key, orig_tensor in original_tensors.items():
+        assert key in loaded_tensors
+        loaded_tensor = loaded_tensors[key]
+        
+        # Verify it's on GPU
+        assert loaded_tensor.device.type == "cuda"
+        
+        # Verify data matches
+        torch.testing.assert_close(loaded_tensor.cpu(), orig_tensor)
+        
+    loader.close()
+
+@pytest.mark.skipif(not HAS_TORCH or not torch.cuda.is_available(), reason="Requires CUDA")
+def test_direct_gpu_async_stream(sample_safetensors):
+    path, original_tensors = sample_safetensors
+    
+    loader = MemoryEfficientSafeOpen(path, low_memory=True, direct_gpu=True)
+    
+    stream = loader.async_stream(
+        keys=list(original_tensors.keys()),
+        batch_size=2,
+        prefetch_batches=1,
+        direct_gpu=True
+    )
+    
+    loaded_count = 0
+    for batch in stream:
+        for key, tensor in batch:
+            assert tensor.device.type == "cuda"
+            torch.testing.assert_close(tensor.cpu(), original_tensors[key])
+            loaded_count += 1
+            
+    assert loaded_count == len(original_tensors)
+    loader.close()
+
+@pytest.mark.skipif(not HAS_TORCH, reason="Requires torch")
+def test_direct_gpu_fallback_no_cuda(sample_safetensors, monkeypatch):
+    # Force cuda to be unavailable
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    
+    path, original_tensors = sample_safetensors
+    
+    # Should fallback to CPU silently
+    loader = MemoryEfficientSafeOpen(path, low_memory=True, direct_gpu=True)
+    
+    loaded_tensors = loader.load_all()
+    
+    for key, orig_tensor in original_tensors.items():
+        loaded_tensor = loaded_tensors[key]
+        assert loaded_tensor.device.type == "cpu"
+        torch.testing.assert_close(loaded_tensor, orig_tensor)
+        
+    loader.close()

unifiedefficientloader-0.2.3/tests/test_logging.py

@@ -0,0 +1,51 @@
+import os
+import torch
+from unifiedefficientloader import (
+    UnifiedSafetensorsLoader,
+    setup_logging,
+    VERBOSE_LEVEL,
+    DEBUG_LEVEL
+)
+from safetensors.torch import save_file
+
+def test_logging():
+    # 1. Create a dummy safetensors file
+    dummy_file = "test_logging.safetensors"
+    tensors = {
+        "weight1": torch.randn(10, 10),
+        "weight2": torch.randn(5, 5)
+    }
+    save_file(tensors, dummy_file)
+
+    try:
+        # 2. Test NORMAL logging (default)
+        print("\n--- Testing NORMAL Logging (Default) ---")
+        setup_logging("NORMAL")
+        with UnifiedSafetensorsLoader(dummy_file, low_memory=False) as loader:
+            _ = loader.get_tensor("weight1")
+
+        # 3. Test VERBOSE logging
+        print("\n--- Testing VERBOSE Logging ---")
+        setup_logging("VERBOSE")
+        with UnifiedSafetensorsLoader(dummy_file, low_memory=True) as loader:
+            _ = loader.get_tensor("weight1")
+            _ = loader.get_tensor("weight2")
+
+        # 4. Test DEBUG logging (includes function traces)
+        print("\n--- Testing DEBUG Logging ---")
+        setup_logging("DEBUG")
+        with UnifiedSafetensorsLoader(dummy_file, low_memory=True) as loader:
+            _ = loader.get_tensor("weight1")
+
+        # 5. Test MINIMAL logging
+        print("\n--- Testing MINIMAL Logging ---")
+        setup_logging("MINIMAL")
+        with UnifiedSafetensorsLoader(dummy_file, low_memory=False) as loader:
+            _ = loader.get_tensor("weight1")
+
+    finally:
+        if os.path.exists(dummy_file):
+            os.remove(dummy_file)
+
+if __name__ == "__main__":
+    test_logging()

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/unifiedefficientloader/__init__.py

@@ -9,7 +9,7 @@ def check_dependencies(*packages):
     for pkg in packages:
         if importlib.util.find_spec(pkg) is None:
             missing.append(pkg)
-    
+
     if missing:
         missing_str = ", ".join(missing)
         raise ImportError(
@@ -23,6 +23,22 @@ check_dependencies("torch")
 from .memory_efficient_loader import UnifiedSafetensorsLoader, MemoryEfficientSafeOpen
 from .tensor_utils import dict_to_tensor, tensor_to_dict
 from .pinned_transfer import transfer_to_gpu_pinned, set_verbose, get_pinned_transfer_stats, reset_pinned_transfer_stats
+from .gpu_buffer_pool import GpuBufferPool
+from .pinned_buffer_pool import PinnedBufferPool
+from .logging_utils import (
+    setup_logging,
+    MINIMAL_LEVEL,
+    NORMAL_LEVEL,
+    VERBOSE_LEVEL,
+    DEBUG_LEVEL,
+    debug,
+    verbose,
+    normal,
+    info,
+    minimal,
+    warning,
+    error
+)
 
 __all__ = [
     "UnifiedSafetensorsLoader",
@@ -33,4 +49,18 @@ __all__ = [
     "set_verbose",
     "get_pinned_transfer_stats",
     "reset_pinned_transfer_stats",
-]
+    "GpuBufferPool",
+    "PinnedBufferPool",
+    "setup_logging",
+    "MINIMAL_LEVEL",
+    "NORMAL_LEVEL",
+    "VERBOSE_LEVEL",
+    "DEBUG_LEVEL",
+    "debug",
+    "verbose",
+    "normal",
+    "info",
+    "minimal",
+    "warning",
+    "error",
+]

unifiedefficientloader-0.2.3/unifiedefficientloader/gpu_buffer_pool.py

@@ -0,0 +1,40 @@
+"""
+GPU memory buffer pool for direct-to-GPU streaming.
+
+Maintains a pool of pre-allocated GPU tensors to avoid allocation overhead
+and ensure strictly bounded VRAM usage during streaming.
+"""
+from typing import Tuple, Optional
+import torch
+from . import logging_utils
+
+logger = logging_utils.get_logger(__name__)
+
+class GpuBufferPool:
+    """Manages a pool of fixed-size GPU memory buffers."""
+    def __init__(self, size_bytes: int, num_buffers: int, device: str = "cuda"):
+        import torch
+        import queue
+        self.device = device
+        self.size_bytes = size_bytes
+        self.num_buffers = num_buffers
+
+        logging_utils.verbose(f"Initializing GpuBufferPool: {num_buffers} buffers of {size_bytes / (1024**2):.2f} MB each on {device}.")
+
+        self.buffers = []
+        for _ in range(num_buffers):
+            buf = torch.empty(size_bytes, dtype=torch.uint8, device=device)
+            self.buffers.append(buf)
+
+        self.free_queue = queue.Queue()
+        for i in range(num_buffers):
+            self.free_queue.put(i)
+
+    def acquire(self) -> Tuple[int, 'torch.Tensor']:
+        """Acquire a free buffer. Blocks if empty."""
+        idx = self.free_queue.get()
+        return idx, self.buffers[idx]
+
+    def release(self, idx: int):
+        """Release buffer back to pool."""
+        self.free_queue.put(idx)

unifiedefficientloader-0.2.3/unifiedefficientloader/logging_utils.py

@@ -0,0 +1,117 @@
+import logging
+import sys
+import functools
+
+# Custom Levels
+# MINIMAL (30): WARNING+ (Reduced)
+# NORMAL (20): INFO+ (Default)
+# VERBOSE (15): Custom+ (Increased)
+# DEBUG (10): DEBUG+ (Every function call)
+
+MINIMAL_LEVEL = 30 # Use logging.WARNING
+NORMAL_LEVEL = 20  # Use logging.INFO
+VERBOSE_LEVEL = 15 # Custom level between INFO and DEBUG
+DEBUG_LEVEL = 10   # logging.DEBUG
+
+logging.addLevelName(VERBOSE_LEVEL, "VERBOSE")
+logging.addLevelName(MINIMAL_LEVEL, "MINIMAL")
+
+class CustomFormatter(logging.Formatter):
+    def format(self, record):
+        # Save original format to restore it later
+        orig_fmt = self._style._fmt
+
+        if record.levelno <= DEBUG_LEVEL:
+             # Debug: Full trace info
+            self._style._fmt = "[%(levelname)s] %(name)s:%(lineno)d - %(message)s"
+        elif record.levelno <= VERBOSE_LEVEL:
+            # Verbose: Detail
+            self._style._fmt = "[%(levelname)s] %(message)s"
+        elif record.levelno <= NORMAL_LEVEL:
+            # Normal: Standard output
+            self._style._fmt = "%(message)s"
+        else:
+            # Minimal/Warning
+            self._style._fmt = "[%(levelname)s] %(message)s"
+
+        result = super().format(record)
+
+        # Restore original format
+        self._style._fmt = orig_fmt
+        return result
+
+def setup_logging(verbose_arg: str = "NORMAL"):
+    """
+    Setup logging based on verbosity name.
+    """
+    level_map = {
+        "DEBUG": DEBUG_LEVEL,
+        "VERBOSE": VERBOSE_LEVEL,
+        "NORMAL": NORMAL_LEVEL,
+        "MINIMAL": MINIMAL_LEVEL
+    }
+
+    level = level_map.get(verbose_arg.upper(), NORMAL_LEVEL)
+
+    logger = logging.getLogger("unifiedefficientloader")
+    logger.setLevel(level)
+
+    # Clear existing handlers to prevent duplicates
+    if logger.handlers:
+        logger.handlers.clear()
+
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(CustomFormatter())
+    logger.addHandler(handler)
+
+    return logger
+
+def get_logger(name=None):
+    if name:
+        if not name.startswith("unifiedefficientloader"):
+            name = f"unifiedefficientloader.{name}"
+        return logging.getLogger(name)
+    return logging.getLogger("unifiedefficientloader")
+
+# Decorator for DEBUG level tracing
+def log_debug(func):
+    """Decorator to log function entry/exit with args (DEBUG level only)."""
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # We only want to construct the string if debug is enabled to save perf
+        logger = get_logger(func.__module__.split('.')[-1])
+        if logger.isEnabledFor(DEBUG_LEVEL):
+            arg_str = ", ".join([repr(a) for a in args])
+            kw_str = ", ".join([f"{k}={v!r}" for k, v in kwargs.items()])
+            all_args = ", ".join(filter(None, [arg_str, kw_str]))
+            logger.log(DEBUG_LEVEL, f"CALL {func.__name__}({all_args})")
+
+        result = func(*args, **kwargs)
+
+        if logger.isEnabledFor(DEBUG_LEVEL):
+             logger.log(DEBUG_LEVEL, f"RET {func.__name__} -> {type(result)}")
+        return result
+    return wrapper
+
+# Convenience wrappers
+def debug(msg, *args, **kwargs):
+    get_logger().log(DEBUG_LEVEL, msg, *args, **kwargs)
+
+def verbose(msg, *args, **kwargs):
+    get_logger().log(VERBOSE_LEVEL, msg, *args, **kwargs)
+
+def normal(msg, *args, **kwargs):
+    get_logger().log(NORMAL_LEVEL, msg, *args, **kwargs)
+
+def info(msg, *args, **kwargs):
+    """Alias for normal/INFO level logging."""
+    normal(msg, *args, **kwargs)
+
+def minimal(msg, *args, **kwargs):
+    get_logger().log(MINIMAL_LEVEL, msg, *args, **kwargs)
+
+def warning(msg, *args, **kwargs):
+    get_logger().warning(msg, *args, **kwargs)
+
+def error(msg, *args, **kwargs):
+    get_logger().error(msg, *args, **kwargs)

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/unifiedefficientloader/memory_efficient_loader.py

@@ -7,10 +7,11 @@ Requires `torch`, `safetensors`, and optionally `tqdm`.
 import gc
 import json
 import struct
-import logging
 from typing import Dict, Optional, Tuple
 
-logger = logging.getLogger(__name__)
+from . import logging_utils
+
+logger = logging_utils.get_logger(__name__)
 
 def _ensure_torch():
     try:
@@ -51,46 +52,57 @@ class UnifiedSafetensorsLoader:
                 loader.mark_processed(key)  # Frees memory in low_memory mode
     """
 
-    def __init__(self, filename: str, low_memory: bool = False):
+    @logging_utils.log_debug
+    def __init__(self, filename: str, low_memory: bool = False, direct_gpu: bool = False):
         """Initialize the loader.
 
         Args:
             filename: Path to safetensors file
             low_memory: If True, use streaming mode; if False, preload all tensors
+            direct_gpu: If True, stream directly to GPU pinned/slab memory (requires low_memory=True)
         """
         torch = _ensure_torch()
         safe_open = _ensure_safetensors()
-        
+
         self.filename = filename
         self.low_memory = low_memory
+        self.direct_gpu = direct_gpu
+
+        if self.direct_gpu and not self.low_memory:
+            logging_utils.warning("direct_gpu=True requires low_memory=True. Forcing low_memory=True.")
+            self.low_memory = True
+
         self._tensors: Dict[str, 'torch.Tensor'] = {}
+        self._gpu_buffer_indices: Dict[str, int] = {}
+        self._gpu_pool = None
+
         self._all_keys = []
         self._file = None
         self._header = None
         self._header_size = None
         self._metadata: Dict[str, str] = {}
 
-        if low_memory:
+        if self.low_memory:
             # Streaming mode: read header only
             self._header, self._header_size = self._read_header()
             self._file = None # Opened lazily to support multiprocessing DataLoader
             self._all_keys = [k for k in self._header.keys() if k != "__metadata__"]
             # Extract metadata from header (safetensors stores it under __metadata__ key)
             self._metadata = self._header.get("__metadata__", {})
-            logger.debug(f"Initialized Low-memory mode: parsed header of size {self._header_size} bytes.")
-            logger.debug(f"Found {len(self._all_keys)} tensors (streaming mode)")
+            logging_utils.verbose(f"Initialized Low-memory mode: parsed header of size {self._header_size} bytes.")
+            logging_utils.verbose(f"Found {len(self._all_keys)} tensors (streaming mode)")
         else:
             # Standard mode: preload all tensors
-            with safe_open(filename, framework="pt", device="cpu") as f:
+            with safe_open(self.filename, framework="pt", device="cpu") as f:
                 self._metadata = f.metadata() or {}
                 self._all_keys = list(f.keys())
-                print(f"Loading {len(self._all_keys)} tensors from source file...")
+                logging_utils.normal(f"Loading {len(self._all_keys)} tensors from source file...")
                 try:
                     from tqdm import tqdm
-                    iterator = tqdm(self._all_keys, desc="Loading tensors")
+                    iterator = tqdm(self._all_keys, desc="Loading tensors", disable=not logger.isEnabledFor(logging_utils.NORMAL_LEVEL))
                 except ImportError:
                     iterator = self._all_keys
-                
+
                 for key in iterator:
                     self._tensors[key] = f.get_tensor(key)
 
@@ -141,6 +153,7 @@ class UnifiedSafetensorsLoader:
         """Get tensor ndim without loading tensor data."""
         return len(self.get_shape(key))
 
+    @logging_utils.log_debug
     def get_tensor(self, key: str) -> 'torch.Tensor':
         """Get a tensor by key.
 
@@ -162,7 +175,7 @@ class UnifiedSafetensorsLoader:
         offset_start, offset_end = metadata["data_offsets"]
 
         if offset_start != offset_end:
-            logger.debug(f"Loading tensor '{key}' from offset {offset_start} to {offset_end} ({(offset_end - offset_start)} bytes)")
+            logging_utils.debug(f"Loading tensor '{key}' from offset {offset_start} to {offset_end} ({(offset_end - offset_start)} bytes)")
             self._file.seek(self._header_size + 8 + offset_start)
             # Use bytearray to create a writable buffer, avoiding PyTorch warning
             # about non-writable tensors from read-only bytes.
@@ -177,12 +190,17 @@ class UnifiedSafetensorsLoader:
         """Mark a tensor as processed, freeing memory if in low-memory mode.
 
         In standard mode, optionally deletes from cache.
-        In low-memory mode, this is a no-op (tensor was never cached).
+        In low-memory mode, frees GPU buffer back to pool if direct_gpu.
         """
         if not self.low_memory and key in self._tensors:
             del self._tensors[key]
             gc.collect()
 
+        if self.direct_gpu and key in self._gpu_buffer_indices:
+            idx = self._gpu_buffer_indices.pop(key)
+            if self._gpu_pool:
+                self._gpu_pool.release(idx)
+
     def _read_header(self):
         """Read and parse the safetensors header."""
         with open(self.filename, "rb") as f:
@@ -270,61 +288,152 @@ class UnifiedSafetensorsLoader:
 
     def async_stream(self, keys: list, batch_size: int = 1, prefetch_batches: int = 2, pin_memory: bool = False):
         """Asynchronously stream tensors from disk.
-        
+
         Args:
             keys: List of tensor keys to load
             batch_size: Number of tensors to yield in each batch
             prefetch_batches: Number of batches to pre-fetch in background
             pin_memory: If True, tensors will be pinned in CPU memory (sequentially in main thread)
-            
+            direct_gpu: Stream via pinned buffer directly to GPU
+
         Yields:
             List of (key, tensor) tuples
         """
         import threading
         import queue
         from concurrent.futures import ThreadPoolExecutor
+        import os
 
         torch = _ensure_torch()
         thread_local = threading.local()
 
+        # Initialize GPU slab and Pinned Buffer Pool if direct_gpu
+        pinned_pool = None
+        cuda_stream = None
+
+        if self.direct_gpu and torch.cuda.is_available():
+            try:
+                from .gpu_buffer_pool import GpuBufferPool
+                from .pinned_buffer_pool import PinnedBufferPool
+
+                # Pre-calculate required slab size
+                max_tensor_bytes = 0
+                for k in keys:
+                    meta = self._header[k]
+                    start, end = meta["data_offsets"]
+                    sz = end - start
+                    max_tensor_bytes = max(max_tensor_bytes, sz)
+
+                # Initialize pools (size of largest tensor)
+                # We need a larger pool to allow the GPU to lag behind the CPU without stalling
+                max_workers = min(16, max(4, batch_size))
+                max_in_flight = max(max_workers, prefetch_batches * batch_size)
+
+                # Double the buffers for a smooth pipeline
+                num_buffers = (max_in_flight + max_workers) * 2 + 2
+
+                # Assign pool to instance to survive the generator lifetime
+                if not getattr(self, '_gpu_pool', None):
+                    self._gpu_pool = GpuBufferPool(max_tensor_bytes, num_buffers)
+
+                pinned_pool = PinnedBufferPool(max_tensor_bytes, num_buffers)
+                cuda_stream = torch.cuda.Stream()
+
+                logging_utils.normal(f"Direct GPU pipeline initialized: {num_buffers} buffers, max {max_tensor_bytes / (1024**2):.1f}MB each (Total VRAM: {(num_buffers*max_tensor_bytes)/(1024**2):.1f}MB)")
+
+            except Exception as e:
+                logging_utils.warning(f"Failed to initialize direct GPU pipeline: {e}. Falling back.")
+                self.direct_gpu = False
+                pinned_pool = None
+        elif self.direct_gpu:
+            logging_utils.warning("direct_gpu=True requested but CUDA is not available. Falling back to CPU.")
+            self.direct_gpu = False
+
         def get_file_handle():
             if not hasattr(thread_local, 'file'):
                 thread_local.file = open(self.filename, "rb")
             return thread_local.file
 
         def _worker_load(key):
+            buf_idx = None
+            gpu_idx = None
             try:
-                # Direct thread-safe read
                 metadata = self._header[key]
                 offset_start, offset_end = metadata["data_offsets"]
-                if offset_start != offset_end:
-                    f = get_file_handle()
-                    f.seek(self._header_size + 8 + offset_start)
-                    tensor_bytes = bytearray(offset_end - offset_start)
-                    f.readinto(tensor_bytes)
-                else:
-                    tensor_bytes = None
+                sz = offset_end - offset_start
+
+                if self.direct_gpu and sz > 0:
+                    # Direct GPU Pipeline Path
+                    buf_idx, pinned_buf = pinned_pool.acquire()
 
-                tensor = self._deserialize_tensor(tensor_bytes, metadata)
-                return key, tensor, None
+                    try:
+                        # Schedule GPU transfer
+                        gpu_idx, gpu_buf = self._gpu_pool.acquire()
+
+                        try:
+                            # Read into pinned memory directly (Zero-Copy CPU path)
+                            import ctypes
+                            view = pinned_buf[:sz]
+
+                            # Create a ctypes c_uint8 array spanning the pinned buffer memory
+                            # This allows f.readinto() to write bytes directly to the torch tensor memory
+                            c_uint8_array = (ctypes.c_uint8 * sz).from_address(view.data_ptr())
+
+                            f = get_file_handle()
+                            f.seek(self._header_size + 8 + offset_start)
+                            f.readinto(c_uint8_array)
+
+                            gpu_view = gpu_buf[:sz]
+
+                            with torch.cuda.stream(cuda_stream):
+                                gpu_view.copy_(view, non_blocking=True)
+
+                                # Create event to track when copy finishes
+                                event = torch.cuda.Event()
+                                event.record()
+
+                            # Critical: wait for stream before allowing worker to finish
+                            # If worker finishes, buffer might be overwritten by next worker
+                            # if pool sizing is tight.
+                            # In direct_gpu, the tensor is the gpu_view.
+                            return key, gpu_view, metadata, buf_idx, gpu_idx, event
+
+                        except Exception as e:
+                            # If reading or copying fails, release GPU buffer
+                            self._gpu_pool.release(gpu_idx)
+                            raise e
+
+                    except Exception as e:
+                        # If acquiring GPU buffer fails, release pinned buffer
+                        pinned_pool.release(buf_idx)
+                        raise e
+                else:
+                    # Standard CPU Path
+                    if offset_start != offset_end:
+                        f = get_file_handle()
+                        f.seek(self._header_size + 8 + offset_start)
+                        tensor_bytes = bytearray(offset_end - offset_start)
+                        f.readinto(tensor_bytes)
+                    else:
+                        tensor_bytes = None
+
+                    tensor = self._deserialize_tensor(tensor_bytes, metadata)
+                    return key, tensor, None, None, None, None
             except Exception as e:
-                # Fallback info for main thread
-                return key, None, e
+                return key, None, e, None, None, None
+
+        max_workers = min(16, max(4, batch_size))
+        max_in_flight = max(max_workers, prefetch_batches * batch_size)
 
         # Queue for individual (key, tensor) pairs
-        # Size it to hold enough for prefetch_batches
-        q = queue.Queue(maxsize=prefetch_batches * batch_size)
-        
+        # Size it to hold enough for prefetch_batches PLUS max_workers to prevent stalling
+        q = queue.Queue(maxsize=max_in_flight + max_workers)
+
         def _producer():
-            # Use a reasonable number of workers for I/O bound tasks
-            max_workers = min(16, max(4, batch_size))
-            # Limit task submission to maintain backpressure on memory
-            max_in_flight = max(max_workers, prefetch_batches * batch_size)
-            
             with ThreadPoolExecutor(max_workers=max_workers) as executor:
                 futures = []
                 key_iter = iter(keys)
-                
+
                 # Fill the pipeline
                 for _ in range(max_in_flight):
                     try:
@@ -332,50 +441,81 @@ class UnifiedSafetensorsLoader:
                         futures.append(executor.submit(_worker_load, k))
                     except StopIteration:
                         break
-                
+
                 while futures:
                     # Maintain order by taking the first future
                     f = futures.pop(0)
                     result = f.result() # Blocks until this specific tensor is loaded
                     q.put(result)       # Blocks if the consumption queue is full
-                    
+
                     # Submit next task if available
                     try:
                         k = next(key_iter)
                         futures.append(executor.submit(_worker_load, k))
                     except StopIteration:
                         pass
-                        
+
             q.put(None) # Sentinel
 
+        producer_thread = threading.local()
         producer_thread = threading.Thread(target=_producer, daemon=True)
         producer_thread.start()
 
         batch = []
+        pending_pinned = [] # Track (event, buf_idx) to release later
+
         while True:
             res = q.get()
             if res is None:
+                # Synchronize and cleanup any remaining buffers on exit
+                for ev, idx in pending_pinned:
+                    ev.synchronize()
+                    pinned_pool.release(idx)
                 if batch:
                     yield batch
                 break
-            
-            k, t, err = res
-            if err is not None:
-                logger.warning(f"Async load failed for {k}, falling back to sync: {err}")
+
+            k, t, err, buf_idx, gpu_idx, event = res
+            if err is not None and not isinstance(err, dict):
+                logging_utils.warning(f"Async load failed for {k}, falling back to sync: {err}")
                 # Fallback synchronous load
                 try:
                     t = self.get_tensor(k)
                 except Exception as sync_err:
-                    logger.error(f"Sync fallback also failed for {k}: {sync_err}")
+                    logging_utils.error(f"Sync fallback also failed for {k}: {sync_err}")
                     raise sync_err
-            
+
+            if buf_idx is not None and event is not None:
+                # Don't block here! Yield the tensor with its event.
+                # Only release the PREVIOUS batch's buffers.
+                # This creates a sliding window of safety.
+                while len(pending_pinned) >= (max_in_flight + 1):
+                    ev, idx = pending_pinned.pop(0)
+                    ev.synchronize() # Wait only if we MUST reuse a buffer
+                    pinned_pool.release(idx)
+
+                pending_pinned.append((event, buf_idx))
+
+                # Register GPU index for cleanup
+                self._gpu_buffer_indices[k] = gpu_idx
+
+                # Reshape GPU view to tensor
+                meta = err # we reused err for metadata in direct_gpu path
+                dtype = self._get_torch_dtype(meta["dtype"])
+                shape = meta["shape"]
+
+                if meta["dtype"] in ["F8_E5M2", "F8_E4M3"]:
+                    t = self._convert_float8(t, meta["dtype"], shape)
+                else:
+                    t = t.view(dtype).reshape(shape)
+
             # Pin memory sequentially in the main thread to avoid OS-level lock contention
-            if pin_memory and t.device.type == 'cpu':
+            elif pin_memory and t.device.type == 'cpu':
                 try:
                     t = t.pin_memory()
                 except Exception as e:
-                    logger.warning(f"Failed to pin memory for {k}: {e}")
-            
+                    logging_utils.warning(f"Failed to pin memory for {k}: {e}")
+
             batch.append((k, t))
             if len(batch) == batch_size:
                 yield batch

unifiedefficientloader-0.2.3/unifiedefficientloader/pinned_buffer_pool.py

@@ -0,0 +1,39 @@
+"""
+Pinned memory pool manager.
+
+Provides pre-allocated pinned buffers to avoid per-tensor allocation overhead.
+"""
+from typing import Optional
+import torch
+from . import logging_utils
+
+logger = logging_utils.get_logger(__name__)
+
+class PinnedBufferPool:
+    """Manages a pool of pinned memory buffers for fast disk-to-GPU transfer."""
+    def __init__(self, size_bytes: int, num_buffers: int):
+        import torch
+        self.size_bytes = size_bytes
+        self.num_buffers = num_buffers
+
+        logging_utils.verbose(f"Initializing PinnedBufferPool: {num_buffers} buffers of {size_bytes / (1024**2):.2f} MB each.")
+
+        self.buffers = []
+        for _ in range(num_buffers):
+            buf = torch.empty(size_bytes, dtype=torch.uint8, pin_memory=True)
+            self.buffers.append(buf)
+
+        import queue
+        self.free_queue = queue.Queue()
+        for i in range(num_buffers):
+            self.free_queue.put(i)
+
+    def acquire(self) -> tuple[int, 'torch.Tensor']:
+        """Acquire a free buffer. Blocks if empty."""
+        idx = self.free_queue.get()
+        return idx, self.buffers[idx]
+
+    def release(self, idx: int):
+        """Release buffer back to pool."""
+        self.free_queue.put(idx)
+

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/unifiedefficientloader/pinned_transfer.py

@@ -4,10 +4,11 @@ Pinned memory utilities for faster CPU→GPU tensor transfers.
 Pinned (page-locked) memory enables faster DMA transfers to GPU.
 Uses PyTorch's native pin_memory() with non_blocking transfers.
 """
-import logging
 from typing import Optional
 
-logger = logging.getLogger(__name__)
+from . import logging_utils
+
+logger = logging_utils.get_logger(__name__)
 
 def _ensure_torch():
     try:
@@ -21,9 +22,14 @@ _verbose = False
 _pinned_transfer_stats = {"pinned": 0, "fallback": 0}
 
 def set_verbose(enabled: bool):
-    """Enable/disable verbose output for pinned transfers."""
+    """
+    Enable/disable verbose output for pinned transfers.
+    Also adjusts logging level to VERBOSE if enabled.
+    """
     global _verbose
     _verbose = enabled
+    if enabled:
+        logging_utils.setup_logging("VERBOSE")
 
 def get_pinned_transfer_stats():
     """Return pinned transfer statistics for verification."""
@@ -34,10 +40,12 @@ def reset_pinned_transfer_stats():
     global _pinned_transfer_stats
     _pinned_transfer_stats = {"pinned": 0, "fallback": 0}
 
+@logging_utils.log_debug
 def transfer_to_gpu_pinned(
     tensor,
     device: str = 'cuda',
-    dtype = None
+    dtype = None,
+    non_blocking: bool = True
 ):
     """Transfer tensor to GPU using pinned memory for faster transfer."""
     torch = _ensure_torch()
@@ -46,44 +54,47 @@ def transfer_to_gpu_pinned(
     # Skip if not a CPU tensor or CUDA unavailable
     if tensor.device.type != 'cpu' or not torch.cuda.is_available():
         if dtype is not None:
-            return tensor.to(device=device, dtype=dtype)
-        return tensor.to(device=device)
+            return tensor.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        return tensor.to(device=device, non_blocking=non_blocking)
 
     # Skip if target is not CUDA
     if not str(device).startswith('cuda'):
         if dtype is not None:
-            return tensor.to(device=device, dtype=dtype)
-        return tensor.to(device=device)
+            return tensor.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        return tensor.to(device=device, non_blocking=non_blocking)
 
     try:
         pinned = tensor.pin_memory()
 
         if dtype is not None:
-            result = pinned.to(device=device, dtype=dtype, non_blocking=True)
+            result = pinned.to(device=device, dtype=dtype, non_blocking=non_blocking)
         else:
-            result = pinned.to(device=device, non_blocking=True)
+            result = pinned.to(device=device, non_blocking=non_blocking)
 
         torch.cuda.current_stream().synchronize()
 
         # One-time confirmation on first success
         if _pinned_transfer_stats["pinned"] == 0:
-            logger.debug("[pinned_transfer] Pinned memory active - faster GPU transfers enabled")
+            logging_utils.verbose("[pinned_transfer] Pinned memory active - faster GPU transfers enabled")
 
         _pinned_transfer_stats["pinned"] += 1
+
+        msg = f"[pinned_transfer] Pinned: {tensor.shape} ({tensor.numel() * tensor.element_size() / 1024:.1f} KB)"
         if _verbose:
-            logger.debug(f"[pinned_transfer] Pinned: {tensor.shape} ({tensor.numel() * tensor.element_size() / 1024:.1f} KB)")
+            logging_utils.normal(msg)
         else:
-            logger.debug(f"[pinned_transfer] Transferred tensor {tensor.shape} to {device} via pinned memory")
+            logging_utils.verbose(msg)
 
         return result
 
     except Exception as e:
         _pinned_transfer_stats["fallback"] += 1
+        msg = f"[pinned_transfer] Fallback transfer to {device} due to error: {e}"
         if _verbose:
-            logger.debug(f"[pinned_transfer] Fallback: {e}")
+            logging_utils.warning(msg)
         else:
-            logger.debug(f"[pinned_transfer] Fallback transfer to {device} due to error: {e}")
+            logging_utils.verbose(msg)
 
         if dtype is not None:
-            return tensor.to(device=device, dtype=dtype)
-        return tensor.to(device=device)
+            return tensor.to(device=device, dtype=dtype, non_blocking=non_blocking)
+        return tensor.to(device=device, non_blocking=non_blocking)

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/unifiedefficientloader/tensor_utils.py

@@ -5,10 +5,11 @@ Provides serialization helpers for dictionary/tensor conversion.
 Requires `torch`.
 """
 import json
-import logging
 from typing import Dict, Tuple
 
-logger = logging.getLogger(__name__)
+from . import logging_utils
+
+logger = logging_utils.get_logger(__name__)
 
 def _ensure_torch():
     try:
@@ -18,6 +19,7 @@ def _ensure_torch():
         raise ImportError("The 'torch' package is required but not installed. Please install it.")
 
 
+@logging_utils.log_debug
 def dict_to_tensor(data_dict: dict):
     """
     Convert a dictionary to a torch.uint8 tensor containing JSON bytes.
@@ -32,9 +34,10 @@ def dict_to_tensor(data_dict: dict):
     json_str = json.dumps(data_dict)
     byte_data = json_str.encode("utf-8")
     tensor_data = torch.tensor(list(byte_data), dtype=torch.uint8)
-    logger.debug(f"dict_to_tensor: serialized dict to uint8 tensor of shape {tensor_data.shape}")
+    logging_utils.debug(f"dict_to_tensor: serialized dict to uint8 tensor of shape {tensor_data.shape}")
     return tensor_data
 
+@logging_utils.log_debug
 def tensor_to_dict(tensor_data) -> dict:
     """
     Convert a torch.uint8 tensor containing JSON bytes to a dictionary.
@@ -50,5 +53,5 @@ def tensor_to_dict(tensor_data) -> dict:
     byte_data = bytes(tensor_data.tolist())
     json_str = byte_data.decode("utf-8")
     data_dict = json.loads(json_str)
-    logger.debug(f"tensor_to_dict: deserialized tensor of shape {tensor_data.shape} to dict with keys: {list(data_dict.keys())}")
-    return data_dict
+    logging_utils.debug(f"tensor_to_dict: deserialized tensor of shape {tensor_data.shape} to dict with keys: {list(data_dict.keys())}")
+    return data_dict

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/unifiedefficientloader.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unifiedefficientloader
-Version: 0.2.1
+Version: 0.2.3
 Summary: A unified interface for memory efficient per tensor loading of safetensors files as raw bytes from offset, handling CPU/GPU pinned transfers, and converting between tensors and dicts.
 Author: silveroxides
 License: MIT
@@ -111,6 +111,33 @@ with UnifiedSafetensorsLoader("model.safetensors", low_memory=True) as loader:
             loader.mark_processed(key)
 ```
 
+### Direct-to-GPU Streaming (Zero-Copy)
+
+For the absolute fastest loading times on CUDA devices, use the `direct_gpu=True` flag. This creates a pipeline that pre-allocates pinned memory pools and GPU memory slabs. Tensors are loaded from disk directly into pinned buffers, and immediately asynchronously copied to the GPU using CUDA streams, hiding the PCIe transfer latency completely behind the disk I/O.
+
+```python
+from unifiedefficientloader import UnifiedSafetensorsLoader
+
+with UnifiedSafetensorsLoader("model.safetensors", low_memory=True, direct_gpu=True) as loader:
+    keys_to_load = loader.keys()
+    
+    # async_stream will automatically coordinate disk -> pinned buffer -> GPU slab -> tensor header
+    stream = loader.async_stream(
+        keys_to_load,
+        batch_size=8,
+        prefetch_batches=2,
+        direct_gpu=True # optional here since we passed it in __init__
+    )
+    
+    for batch in stream:
+        for key, gpu_tensor in batch:
+            # gpu_tensor is already on the GPU!
+            assert gpu_tensor.device.type == "cuda"
+            
+            # ... process gpu_tensor ...
+            loader.mark_processed(key)
+```
+
 ### Tensor/Dict Conversion
 
 ```python

{unifiedefficientloader-0.2.1 → unifiedefficientloader-0.2.3}/unifiedefficientloader.egg-info/SOURCES.txt

@@ -2,9 +2,14 @@ LICENSE
 README.md
 pyproject.toml
 setup.py
+tests/test_direct_gpu.py
+tests/test_logging.py
 tests/test_utils.py
 unifiedefficientloader/__init__.py
+unifiedefficientloader/gpu_buffer_pool.py
+unifiedefficientloader/logging_utils.py
 unifiedefficientloader/memory_efficient_loader.py
+unifiedefficientloader/pinned_buffer_pool.py
 unifiedefficientloader/pinned_transfer.py
 unifiedefficientloader/tensor_utils.py
 unifiedefficientloader.egg-info/PKG-INFO