unifiedefficientloader 0.3.2__cp39-abi3-win_amd64.whl

unifiedefficientloader/__init__.py

@@ -0,0 +1,66 @@
+import importlib.util
+
+def check_dependencies(*packages):
+    """
+    Check if required packages are installed.
+    Throws a descriptive error if not.
+    """
+    missing = []
+    for pkg in packages:
+        if importlib.util.find_spec(pkg) is None:
+            missing.append(pkg)
+
+    if missing:
+        missing_str = ", ".join(missing)
+        raise ImportError(
+            f"Missing required packages for unifiedefficientloader: {missing_str}. "
+            f"Please install them using: pip install {missing_str}"
+        )
+
+# Pre-check torch as it is the foundation of most of these tools
+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",
+    "MemoryEfficientSafeOpen",
+    "dict_to_tensor",
+    "tensor_to_dict",
+    "transfer_to_gpu_pinned",
+    "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/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/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)