b10-transfer 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

b10_transfer-0.1.2.dist-info/METADATA

@@ -0,0 +1,127 @@
+Metadata-Version: 2.3
+Name: b10-transfer
+Version: 0.1.2
+Summary: Distributed PyTorch file transfer for Baseten - Environment-aware, lock-free file transfer management
+License: MIT
+Keywords: pytorch,file-transfer,cache,machine-learning,inference
+Author: Shounak Ray
+Author-email: shounak.noreply@baseten.co
+Maintainer: Fred Liu
+Maintainer-email: fred.liu.noreply@baseten.co
+Requires-Python: >=3.9,<4.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Dist: torch (>=2.0.0)
+Requires-Dist: triton (>=2.0.0)
+Project-URL: Documentation, https://docs.baseten.co/development/model/b10-transfer
+Project-URL: Homepage, https://docs.baseten.co/development/model/b10-transfer
+Project-URL: Repository, https://pypi.org/project/b10-transfer/
+Description-Content-Type: text/markdown
+
+https://www.notion.so/ml-infra/mega-base-cache-24291d247273805b8e20fe26677b7b0f
+
+# B10 Transfer
+
+PyTorch file transfer for Baseten deployments.
+
+## Usage
+
+```python
+import b10_transfer
+
+# Inside model.load() function
+def load()
+    # Load cache before torch.compile()
+    cache_loaded = b10_transfer.load_compile_cache()
+
+    # ...
+
+    # Your model compilation
+    model = torch.compile(model)
+    # Warm up the model with dummy prompts, and arguments that would be typically used in your requests (e.g resolutions)
+    dummy_input = "What is the capital of France?"
+    model(dummy_input)
+
+    # ...
+
+    # Save cache after compilation
+    if not cache_loaded:
+        b10_transfer.save_compile_cache()
+```
+
+## Configuration
+
+Configure via environment variables:
+
+```bash
+# Cache directories
+export TORCH_CACHE_DIR="/tmp/torchinductor_root"      # Default
+export B10FS_CACHE_DIR="/cache/model/compile_cache"   # Default  
+export LOCAL_WORK_DIR="/app"                          # Default
+
+# Cache limits
+export MAX_CACHE_SIZE_MB="1024"                       # 1GB default
+```
+
+## How It Works
+
+### Environment-Specific Caching
+
+The library automatically creates unique cache keys based on your environment:
+
+```
+torch-2.1.0_cuda-12.1_cc-8.6_triton-2.1.0 → cache_a1b2c3d4e5f6.latest.tar.gz
+torch-2.0.1_cuda-11.8_cc-7.5_triton-2.0.1 → cache_x9y8z7w6v5u4.latest.tar.gz
+torch-2.1.0_cpu_triton-none                → cache_m1n2o3p4q5r6.latest.tar.gz
+```
+
+**Components used:**
+- **PyTorch version** (e.g., `torch-2.1.0`)
+- **CUDA version** (e.g., `cuda-12.1` or `cpu`)
+- **GPU compute capability** (e.g., `cc-8.6` for A100)
+- **Triton version** (e.g., `triton-2.1.0` or `triton-none`)
+
+### Cache Workflow
+
+1. **Load Phase** (startup): Generate environment key, check for matching cache in B10FS, extract to local directory
+2. **Save Phase** (after compilation): Create archive, atomic copy to B10FS with environment-specific filename
+
+### Lock-Free Race Prevention  
+
+Uses journal pattern with atomic filesystem operations for parallel-safe cache saves.
+
+## API Reference
+
+### Functions
+
+- `load_compile_cache() -> bool`: Load cache from B10FS for current environment
+- `save_compile_cache() -> bool`: Save cache to B10FS with environment-specific filename
+- `clear_local_cache() -> bool`: Clear local cache directory
+- `get_cache_info() -> Dict[str, Any]`: Get cache status information for current environment
+- `list_available_caches() -> Dict[str, Any]`: List all cache files with environment details
+
+### Exceptions
+
+- `CacheError`: Base exception for cache operations
+- `CacheValidationError`: Path validation or compatibility check failed
+
+## Performance Impact
+
+### Debugging
+
+Enable debug logging:
+
+```python
+import logging
+logging.getLogger('b10_transfer').setLevel(logging.DEBUG)
+```
+

b10_transfer-0.1.2.dist-info/RECORD

@@ -0,0 +1,12 @@
+b10_transfer/__init__.py,sha256=kZXd7GHMH7PFzr4aXs19MmuFzwEvtwzh0rwyK4jHgHo,641
+b10_transfer/archive.py,sha256=GKb0mi0-YeM7ch4FLAoOLHXw0T6LkRerYad2N2y9TYM,6400
+b10_transfer/cleanup.py,sha256=3RnqWNGMCcko5GQdq1Gr9VPpGzAF5J6x7xjIH9SNZ78,6226
+b10_transfer/constants.py,sha256=EmWCh9AOamCZL3KkSU6YJO6KBkh93OmAIUeEfaZxHL0,4321
+b10_transfer/core.py,sha256=tKA1gWDEpcb_Xfr6njScAGWCxIT3htlyVh9VLg67YMg,15445
+b10_transfer/environment.py,sha256=aC0biEMQrtHk0ke_3epdcq1X9J5fPmPpBVt0fH7XF2Y,5625
+b10_transfer/info.py,sha256=I3iOuImZ5r6DMJTDeBtVvzlSn6IuyPJbLJYUO_OF0ks,6299
+b10_transfer/space_monitor.py,sha256=C_CKDH43bNsWdq60WStSZ3c_nQkWvScQmqU_SYHesew,10531
+b10_transfer/utils.py,sha256=Stee0DFK-8MRRYNIocqaK64cJvfs4jPW3Mpx7zkWV6Y,11932
+b10_transfer-0.1.2.dist-info/METADATA,sha256=P0Yf1VzkqFV4eco8x1yAjgUiahlmjgJGTAee0NNhh6o,4108
+b10_transfer-0.1.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+b10_transfer-0.1.2.dist-info/RECORD,,

b10_transfer/async_torch_cache.py

@@ -1,62 +0,0 @@
-"""Torch-specific async cache operations using the generic async transfer system."""
-
-from pathlib import Path
-from typing import Optional, Callable
-
-from .async_transfers import start_transfer_async
-from .torch_cache import torch_cache_load_callback, torch_cache_save_callback
-from .environment import get_cache_filename
-from .constants import (
-    TORCH_CACHE_DIR,
-    B10FS_CACHE_DIR,
-    MAX_CACHE_SIZE_MB,
-    CACHE_FILE_EXTENSION,
-    CACHE_LATEST_SUFFIX,
-)
-
-
-def load_compile_cache_async(
-    progress_callback: Optional[Callable[[str], None]] = None,
-) -> str:
-    """Start async PyTorch compilation cache load operation."""
-    b10fs_dir = Path(B10FS_CACHE_DIR)
-    torch_dir = Path(TORCH_CACHE_DIR)
-
-    cache_filename = get_cache_filename()
-    cache_file = (
-        b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
-    )
-
-    return start_transfer_async(
-        source=cache_file,
-        dest=torch_dir,
-        callback=torch_cache_load_callback,
-        operation_name="torch_cache_load",
-        progress_callback=progress_callback,
-        monitor_local=True,
-        monitor_b10fs=False,  # No need to monitor b10fs for read operations
-    )
-
-
-def save_compile_cache_async(
-    progress_callback: Optional[Callable[[str], None]] = None,
-) -> str:
-    """Start async PyTorch compilation cache save operation."""
-    b10fs_dir = Path(B10FS_CACHE_DIR)
-    torch_dir = Path(TORCH_CACHE_DIR)
-
-    cache_filename = get_cache_filename()
-    final_file = (
-        b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
-    )
-
-    return start_transfer_async(
-        source=torch_dir,
-        dest=final_file,
-        callback=torch_cache_save_callback,
-        operation_name="torch_cache_save",
-        progress_callback=progress_callback,
-        monitor_local=True,
-        monitor_b10fs=True,
-        max_size_mb=MAX_CACHE_SIZE_MB,
-    )

b10_transfer/async_transfers.py

@@ -1,283 +0,0 @@
-"""Generic async transfer operations with progress tracking."""
-
-import asyncio
-import logging
-import threading
-import time
-from pathlib import Path
-from typing import Optional, Dict, Any, Callable
-from dataclasses import dataclass
-from datetime import datetime, timedelta
-from concurrent.futures import ThreadPoolExecutor
-
-from .core import transfer
-from .constants import AsyncTransferStatus, TransferStatus
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class TransferProgress:
-    operation_id: str
-    operation_name: str
-    status: AsyncTransferStatus
-    started_at: Optional[datetime] = None
-    completed_at: Optional[datetime] = None
-    error_message: Optional[str] = None
-    progress_callback: Optional[Callable[[str], None]] = None
-
-    def update_status(
-        self, status: AsyncTransferStatus, error_message: Optional[str] = None
-    ):
-        self.status = status
-        if error_message:
-            self.error_message = error_message
-        if status == AsyncTransferStatus.IN_PROGRESS and self.started_at is None:
-            self.started_at = datetime.now()
-        elif status in [
-            AsyncTransferStatus.SUCCESS,
-            AsyncTransferStatus.ERROR,
-            AsyncTransferStatus.INTERRUPTED,
-            AsyncTransferStatus.CANCELLED,
-            AsyncTransferStatus.DOES_NOT_EXIST,
-        ]:
-            self.completed_at = datetime.now()
-
-        # Notify callback if provided
-        if self.progress_callback:
-            try:
-                self.progress_callback(self.operation_id)
-            except Exception as e:
-                logger.warning(f"Progress callback failed for {self.operation_id}: {e}")
-
-
-class AsyncTransferManager:
-    def __init__(self):
-        self._transfers: Dict[str, TransferProgress] = {}
-        self._executor = ThreadPoolExecutor(
-            max_workers=2, thread_name_prefix="b10-async-transfer"
-        )
-        self._lock = threading.Lock()
-        self._operation_counter = 0
-
-    def _generate_operation_id(self, operation_name: str) -> str:
-        with self._lock:
-            self._operation_counter += 1
-            return f"{operation_name}_{self._operation_counter}_{int(datetime.now().timestamp())}"
-
-    def start_transfer_async(
-        self,
-        source: Path,
-        dest: Path,
-        callback: Callable,
-        operation_name: str,
-        progress_callback: Optional[Callable[[str], None]] = None,
-        monitor_local: bool = True,
-        monitor_b10fs: bool = True,
-        **callback_kwargs,
-    ) -> str:
-        operation_id = self._generate_operation_id(operation_name)
-
-        progress = TransferProgress(
-            operation_id=operation_id,
-            operation_name=operation_name,
-            status=AsyncTransferStatus.NOT_STARTED,
-            progress_callback=progress_callback,
-        )
-
-        with self._lock:
-            self._transfers[operation_id] = progress
-
-        # Submit the transfer operation to thread pool
-        future = self._executor.submit(
-            self._execute_transfer,
-            operation_id,
-            source,
-            dest,
-            callback,
-            monitor_local,
-            monitor_b10fs,
-            callback_kwargs,
-        )
-
-        logger.info(f"Started async transfer operation: {operation_id}")
-        return operation_id
-
-    def _execute_transfer(
-        self,
-        operation_id: str,
-        source: Path,
-        dest: Path,
-        callback: Callable,
-        monitor_local: bool,
-        monitor_b10fs: bool,
-        callback_kwargs: Dict[str, Any],
-    ) -> None:
-        progress = self._transfers.get(operation_id)
-        if not progress:
-            logger.error(f"Progress tracking lost for operation {operation_id}")
-            return
-
-        try:
-            progress.update_status(AsyncTransferStatus.IN_PROGRESS)
-            logger.info(f"Starting transfer for operation {operation_id}")
-
-            result = transfer(
-                source=source,
-                dest=dest,
-                callback=callback,
-                monitor_local=monitor_local,
-                monitor_b10fs=monitor_b10fs,
-                **callback_kwargs,
-            )
-
-            # Convert TransferStatus to AsyncTransferStatus
-            if result == TransferStatus.SUCCESS:
-                progress.update_status(AsyncTransferStatus.SUCCESS)
-                logger.info(f"Transfer completed successfully: {operation_id}")
-            elif result == TransferStatus.INTERRUPTED:
-                progress.update_status(
-                    AsyncTransferStatus.INTERRUPTED,
-                    "Transfer interrupted due to insufficient disk space",
-                )
-                logger.warning(f"Transfer interrupted: {operation_id}")
-            elif result == TransferStatus.DOES_NOT_EXIST:
-                progress.update_status(
-                    AsyncTransferStatus.DOES_NOT_EXIST,
-                    "Cache file not found",
-                )
-                logger.info(f"Transfer failed - file not found: {operation_id}")
-            else:
-                progress.update_status(
-                    AsyncTransferStatus.ERROR, "Transfer operation failed"
-                )
-                logger.error(f"Transfer failed: {operation_id}")
-
-        except Exception as e:
-            progress.update_status(AsyncTransferStatus.ERROR, str(e))
-            logger.error(
-                f"Transfer operation {operation_id} failed with exception: {e}"
-            )
-
-    def get_transfer_status(self, operation_id: str) -> Optional[TransferProgress]:
-        with self._lock:
-            return self._transfers.get(operation_id)
-
-    def is_transfer_complete(self, operation_id: str) -> bool:
-        progress = self.get_transfer_status(operation_id)
-        if not progress:
-            return False
-
-        return progress.status in [
-            AsyncTransferStatus.SUCCESS,
-            AsyncTransferStatus.ERROR,
-            AsyncTransferStatus.INTERRUPTED,
-            AsyncTransferStatus.CANCELLED,
-            AsyncTransferStatus.DOES_NOT_EXIST,
-        ]
-
-    def wait_for_completion(
-        self, operation_id: str, timeout: Optional[float] = None
-    ) -> bool:
-        start_time = datetime.now()
-
-        while not self.is_transfer_complete(operation_id):
-            if timeout and (datetime.now() - start_time).total_seconds() > timeout:
-                return False
-            time.sleep(0.1)  # Small delay to avoid busy waiting
-
-        return True
-
-    def cancel_transfer(self, operation_id: str) -> bool:
-        progress = self.get_transfer_status(operation_id)
-        if not progress:
-            return False
-
-        if progress.status == AsyncTransferStatus.IN_PROGRESS:
-            progress.update_status(AsyncTransferStatus.CANCELLED)
-            logger.info(f"Marked transfer operation as cancelled: {operation_id}")
-            return True
-
-        return False
-
-    def cleanup_completed_transfers(self, max_age_hours: int = 24) -> int:
-        cutoff_time = datetime.now() - timedelta(hours=max_age_hours)
-        cleaned_count = 0
-
-        with self._lock:
-            to_remove = []
-            for operation_id, progress in self._transfers.items():
-                if (
-                    progress.completed_at
-                    and progress.completed_at < cutoff_time
-                    and self.is_transfer_complete(operation_id)
-                ):
-                    to_remove.append(operation_id)
-
-            for operation_id in to_remove:
-                del self._transfers[operation_id]
-                cleaned_count += 1
-
-        if cleaned_count > 0:
-            logger.info(f"Cleaned up {cleaned_count} completed transfer records")
-
-        return cleaned_count
-
-    def list_active_transfers(self) -> Dict[str, TransferProgress]:
-        with self._lock:
-            return {
-                op_id: progress
-                for op_id, progress in self._transfers.items()
-                if not self.is_transfer_complete(op_id)
-            }
-
-    def shutdown(self) -> None:
-        logger.info("Shutting down async transfer manager...")
-        self._executor.shutdown(wait=True)
-
-
-# Global instance for easy access
-_transfer_manager = AsyncTransferManager()
-
-
-# Generic Public API functions
-def start_transfer_async(
-    source: Path,
-    dest: Path,
-    callback: Callable,
-    operation_name: str,
-    progress_callback: Optional[Callable[[str], None]] = None,
-    monitor_local: bool = True,
-    monitor_b10fs: bool = True,
-    **callback_kwargs,
-) -> str:
-    return _transfer_manager.start_transfer_async(
-        source=source,
-        dest=dest,
-        callback=callback,
-        operation_name=operation_name,
-        progress_callback=progress_callback,
-        monitor_local=monitor_local,
-        monitor_b10fs=monitor_b10fs,
-        **callback_kwargs,
-    )
-
-
-def get_transfer_status(operation_id: str) -> Optional[TransferProgress]:
-    return _transfer_manager.get_transfer_status(operation_id)
-
-
-def is_transfer_complete(operation_id: str) -> bool:
-    return _transfer_manager.is_transfer_complete(operation_id)
-
-
-def wait_for_completion(operation_id: str, timeout: Optional[float] = None) -> bool:
-    return _transfer_manager.wait_for_completion(operation_id, timeout)
-
-
-def cancel_transfer(operation_id: str) -> bool:
-    return _transfer_manager.cancel_transfer(operation_id)
-
-
-def list_active_transfers() -> Dict[str, TransferProgress]:
-    return _transfer_manager.list_active_transfers()