b10-transfer 0.0.1__tar.gz

b10_transfer-0.0.1/PKG-INFO

@@ -0,0 +1,219 @@
+Metadata-Version: 2.3
+Name: b10-transfer
+Version: 0.0.1
+Summary: Distributed PyTorch compilation cache for Baseten - Environment-aware, lock-free compilation cache management
+License: MIT
+Keywords: pytorch,torch.compile,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 compilation cache for Baseten deployments.
+
+## Usage
+
+### Synchronous Operations (Blocking)
+
+```python
+import b10_transfer
+
+# Inside model.load() function
+def load():
+    # Load cache before torch.compile()
+    status = 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 status != b10_transfer.LoadStatus.SUCCESS:
+        b10_transfer.save_compile_cache()
+```
+
+### Asynchronous Operations (Non-blocking)
+
+```python
+import b10_transfer
+
+def load_with_async_cache():
+    # Start async cache load (returns immediately with operation ID)
+    operation_id = b10_transfer.load_compile_cache_async()
+    
+    # Check status periodically
+    while not b10_transfer.is_transfer_complete(operation_id):
+        status = b10_transfer.get_transfer_status(operation_id)
+        print(f"Cache load status: {status.status}")
+        time.sleep(1)
+    
+    # Get final status
+    final_status = b10_transfer.get_transfer_status(operation_id)
+    if final_status.status == b10_transfer.AsyncTransferStatus.SUCCESS:
+        print("Cache loaded successfully!")
+    
+    # Your model compilation...
+    model = torch.compile(model)
+    
+    # Async save
+    save_op_id = b10_transfer.save_compile_cache_async()
+    
+    # You can continue with other work while save happens in background
+    # Or wait for completion if needed
+    b10_transfer.wait_for_completion(save_op_id, timeout=300)  # 5 minute timeout
+
+# With progress callback
+def on_progress(operation_id: str):
+    status = b10_transfer.get_transfer_status(operation_id)
+    print(f"Transfer {operation_id}: {status.status}")
+
+operation_id = b10_transfer.load_compile_cache_async(progress_callback=on_progress)
+```
+
+### Generic Async Operations
+
+You can also use the generic async system for custom transfer operations:
+
+```python
+import b10_transfer
+from pathlib import Path
+
+def my_custom_callback(source: Path, dest: Path):
+    # Your custom transfer logic here
+    # This could be any file operation, compression, etc.
+    shutil.copy2(source, dest)
+
+# Start a generic async transfer
+operation_id = b10_transfer.start_transfer_async(
+    source=Path("/source/file.txt"),
+    dest=Path("/dest/file.txt"),
+    callback=my_custom_callback,
+    operation_name="custom_file_copy",
+    monitor_local=True,
+    monitor_b10fs=False
+)
+
+# Use the same progress tracking as torch cache operations
+b10_transfer.wait_for_completion(operation_id)
+```
+
+## 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
+
+### Synchronous Functions
+
+- `load_compile_cache() -> LoadStatus`: Load cache from B10FS for current environment
+- `save_compile_cache() -> SaveStatus`: 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
+
+### Generic Asynchronous Functions
+
+- `start_transfer_async(source, dest, callback, operation_name, **kwargs) -> str`: Start any async transfer operation
+- `get_transfer_status(operation_id: str) -> TransferProgress`: Get current status of async operation
+- `is_transfer_complete(operation_id: str) -> bool`: Check if async operation has completed
+- `wait_for_completion(operation_id: str, timeout=None) -> bool`: Wait for async operation to complete
+- `cancel_transfer(operation_id: str) -> bool`: Attempt to cancel running operation
+- `list_active_transfers() -> Dict[str, TransferProgress]`: Get all active transfer operations
+
+### Torch Cache Async Functions
+
+- `load_compile_cache_async(progress_callback=None) -> str`: Start async cache load, returns operation ID
+- `save_compile_cache_async(progress_callback=None) -> str`: Start async cache save, returns operation ID
+
+### Status Enums
+
+- `LoadStatus`: SUCCESS, ERROR, DOES_NOT_EXIST, SKIPPED
+- `SaveStatus`: SUCCESS, ERROR, SKIPPED  
+- `AsyncTransferStatus`: NOT_STARTED, IN_PROGRESS, SUCCESS, ERROR, INTERRUPTED, CANCELLED
+
+### Data Classes
+
+- `TransferProgress`: Contains operation_id, status, started_at, completed_at, error_message
+
+### Exceptions
+
+- `CacheError`: Base exception for cache operations
+- `CacheValidationError`: Path validation or compatibility check failed
+- `CacheOperationInterrupted`: Operation interrupted due to insufficient disk space
+
+## Performance Impact
+
+### Debugging
+
+Enable debug logging:
+
+```python
+import logging
+logging.getLogger('b10_tcache').setLevel(logging.DEBUG)
+```
+

b10_transfer-0.0.1/README.md

@@ -0,0 +1,189 @@
+https://www.notion.so/ml-infra/mega-base-cache-24291d247273805b8e20fe26677b7b0f
+
+# B10 Transfer
+
+PyTorch compilation cache for Baseten deployments.
+
+## Usage
+
+### Synchronous Operations (Blocking)
+
+```python
+import b10_transfer
+
+# Inside model.load() function
+def load():
+    # Load cache before torch.compile()
+    status = 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 status != b10_transfer.LoadStatus.SUCCESS:
+        b10_transfer.save_compile_cache()
+```
+
+### Asynchronous Operations (Non-blocking)
+
+```python
+import b10_transfer
+
+def load_with_async_cache():
+    # Start async cache load (returns immediately with operation ID)
+    operation_id = b10_transfer.load_compile_cache_async()
+    
+    # Check status periodically
+    while not b10_transfer.is_transfer_complete(operation_id):
+        status = b10_transfer.get_transfer_status(operation_id)
+        print(f"Cache load status: {status.status}")
+        time.sleep(1)
+    
+    # Get final status
+    final_status = b10_transfer.get_transfer_status(operation_id)
+    if final_status.status == b10_transfer.AsyncTransferStatus.SUCCESS:
+        print("Cache loaded successfully!")
+    
+    # Your model compilation...
+    model = torch.compile(model)
+    
+    # Async save
+    save_op_id = b10_transfer.save_compile_cache_async()
+    
+    # You can continue with other work while save happens in background
+    # Or wait for completion if needed
+    b10_transfer.wait_for_completion(save_op_id, timeout=300)  # 5 minute timeout
+
+# With progress callback
+def on_progress(operation_id: str):
+    status = b10_transfer.get_transfer_status(operation_id)
+    print(f"Transfer {operation_id}: {status.status}")
+
+operation_id = b10_transfer.load_compile_cache_async(progress_callback=on_progress)
+```
+
+### Generic Async Operations
+
+You can also use the generic async system for custom transfer operations:
+
+```python
+import b10_transfer
+from pathlib import Path
+
+def my_custom_callback(source: Path, dest: Path):
+    # Your custom transfer logic here
+    # This could be any file operation, compression, etc.
+    shutil.copy2(source, dest)
+
+# Start a generic async transfer
+operation_id = b10_transfer.start_transfer_async(
+    source=Path("/source/file.txt"),
+    dest=Path("/dest/file.txt"),
+    callback=my_custom_callback,
+    operation_name="custom_file_copy",
+    monitor_local=True,
+    monitor_b10fs=False
+)
+
+# Use the same progress tracking as torch cache operations
+b10_transfer.wait_for_completion(operation_id)
+```
+
+## 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
+
+### Synchronous Functions
+
+- `load_compile_cache() -> LoadStatus`: Load cache from B10FS for current environment
+- `save_compile_cache() -> SaveStatus`: 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
+
+### Generic Asynchronous Functions
+
+- `start_transfer_async(source, dest, callback, operation_name, **kwargs) -> str`: Start any async transfer operation
+- `get_transfer_status(operation_id: str) -> TransferProgress`: Get current status of async operation
+- `is_transfer_complete(operation_id: str) -> bool`: Check if async operation has completed
+- `wait_for_completion(operation_id: str, timeout=None) -> bool`: Wait for async operation to complete
+- `cancel_transfer(operation_id: str) -> bool`: Attempt to cancel running operation
+- `list_active_transfers() -> Dict[str, TransferProgress]`: Get all active transfer operations
+
+### Torch Cache Async Functions
+
+- `load_compile_cache_async(progress_callback=None) -> str`: Start async cache load, returns operation ID
+- `save_compile_cache_async(progress_callback=None) -> str`: Start async cache save, returns operation ID
+
+### Status Enums
+
+- `LoadStatus`: SUCCESS, ERROR, DOES_NOT_EXIST, SKIPPED
+- `SaveStatus`: SUCCESS, ERROR, SKIPPED  
+- `AsyncTransferStatus`: NOT_STARTED, IN_PROGRESS, SUCCESS, ERROR, INTERRUPTED, CANCELLED
+
+### Data Classes
+
+- `TransferProgress`: Contains operation_id, status, started_at, completed_at, error_message
+
+### Exceptions
+
+- `CacheError`: Base exception for cache operations
+- `CacheValidationError`: Path validation or compatibility check failed
+- `CacheOperationInterrupted`: Operation interrupted due to insufficient disk space
+
+## Performance Impact
+
+### Debugging
+
+Enable debug logging:
+
+```python
+import logging
+logging.getLogger('b10_tcache').setLevel(logging.DEBUG)
+```

b10_transfer-0.0.1/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry]
+name = "b10-transfer"
+version = "0.0.1"
+description = "Distributed PyTorch compilation cache for Baseten - Environment-aware, lock-free compilation cache management"
+authors = ["Shounak Ray <shounak.noreply@baseten.co>", "Fred Liu <fred.liu.noreply@baseten.co>"]
+maintainers = ["Fred Liu <fred.liu.noreply@baseten.co>", "Shounak Ray <shounak.noreply@baseten.co>"]
+readme = "README.md"
+homepage = "https://docs.baseten.co/development/model/b10-transfer"
+documentation = "https://docs.baseten.co/development/model/b10-transfer"
+repository = "https://pypi.org/project/b10-transfer/"
+license = "MIT"
+keywords = ["pytorch", "torch.compile", "cache", "machine-learning", "inference"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+packages = [{include = "b10_transfer", from = "src"}]
+
+[tool.poetry.dependencies]
+python = "^3.9"
+torch = ">=2.0.0"
+triton = ">=2.0.0"
+

b10_transfer-0.0.1/src/b10_transfer/__init__.py

@@ -0,0 +1,51 @@
+"""B10 Transfer - Lock-free PyTorch compilation cache for Baseten."""
+
+from .core import transfer
+from .torch_cache import load_compile_cache, save_compile_cache, clear_local_cache
+from .async_transfers import (
+    start_transfer_async,
+    get_transfer_status,
+    is_transfer_complete,
+    wait_for_completion,
+    cancel_transfer,
+    list_active_transfers,
+    TransferProgress,
+)
+from .async_torch_cache import (
+    load_compile_cache_async,
+    save_compile_cache_async,
+)
+from .utils import CacheError, CacheValidationError
+from .space_monitor import CacheOperationInterrupted
+from .info import get_cache_info, list_available_caches
+from .constants import SaveStatus, LoadStatus, TransferStatus, AsyncTransferStatus
+
+# Version
+__version__ = "0.0.1"
+
+__all__ = [
+    "CacheError",
+    "CacheValidationError",
+    "CacheOperationInterrupted",
+    "SaveStatus",
+    "LoadStatus",
+    "TransferStatus",
+    "AsyncTransferStatus",
+    "transfer",
+    "load_compile_cache",
+    "save_compile_cache",
+    "clear_local_cache",
+    "get_cache_info",
+    "list_available_caches",
+    # Generic async operations
+    "start_transfer_async",
+    "get_transfer_status",
+    "is_transfer_complete",
+    "wait_for_completion",
+    "cancel_transfer",
+    "list_active_transfers",
+    "TransferProgress",
+    # Torch-specific async operations
+    "load_compile_cache_async",
+    "save_compile_cache_async",
+]

b10_transfer-0.0.1/src/b10_transfer/archive.py

@@ -0,0 +1,175 @@
+import os
+import logging
+import subprocess
+from pathlib import Path
+
+from .utils import timed_fn, safe_unlink, CacheValidationError, validate_path_security
+from .constants import MAX_CACHE_SIZE_MB
+
+logger = logging.getLogger(__name__)
+
+
+class ArchiveError(Exception):
+    """Archive operation failed."""
+
+    pass
+
+
+def get_file_size_mb(file_path: Path) -> float:
+    """Get the size of a file in megabytes.
+
+    Args:
+        file_path: Path to the file to measure.
+
+    Returns:
+        float: File size in megabytes, or 0.0 if file doesn't exist or
+               can't be accessed.
+
+    Raises:
+        No exceptions are raised; OSError is caught and returns 0.0.
+    """
+    try:
+        return file_path.stat().st_size / (1024 * 1024)
+    except OSError:
+        return 0.0
+
+
+def _compress_directory_to_tar(source_dir: Path, target_file: Path) -> None:
+    """Compress directory contents to a gzipped tar archive using system tar.
+
+    This function recursively compresses all files in the source directory
+    into a gzipped tar archive using the system tar command for better performance.
+
+    Args:
+        source_dir: Path to the directory to compress.
+        target_file: Path where the compressed archive will be created.
+
+    Raises:
+        subprocess.CalledProcessError: If tar command fails.
+        OSError: If source directory can't be read or target file can't be written.
+    """
+    # Use system tar command for better performance
+    # -czf: create, gzip, file
+    # -C: change to directory before archiving
+    cmd = ["tar", "-czf", str(target_file), "-C", str(source_dir), "."]
+
+    try:
+        subprocess.run(cmd, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        raise OSError(f"tar compression failed: {e.stderr}") from e
+
+
+@timed_fn(logger=logger, name="Creating archive")
+def create_archive(
+    source_dir: Path, target_file: Path, max_size_mb: int = MAX_CACHE_SIZE_MB
+) -> None:
+    """Create a compressed archive with path validation and size limits.
+
+    This function safely creates a gzipped tar archive from a source directory
+    with security validation and size constraints. It validates paths to prevent
+    directory traversal attacks and enforces maximum archive size limits.
+
+    Args:
+        source_dir: Path to the directory to archive. Must exist and be within
+                   allowed directories (/tmp/ or its parent).
+        target_file: Path where the archive will be created. Must be within
+                    allowed directories (/app or /cache).
+        max_size_mb: Maximum allowed archive size in megabytes. Defaults to MAX_CACHE_SIZE_MB.
+
+    Raises:
+        CacheValidationError: If paths are outside allowed directories.
+        ArchiveError: If source directory doesn't exist, archive creation fails,
+                     or archive exceeds size limit.
+    """
+    # Validate paths
+    validate_path_security(
+        str(source_dir),
+        ["/tmp/", str(source_dir.parent)],
+        f"Source directory {source_dir}",
+        CacheValidationError,
+    )
+    validate_path_security(
+        str(target_file),
+        ["/app", "/cache"],
+        f"Target file {target_file}",
+        CacheValidationError,
+    )
+
+    if not source_dir.exists():
+        raise ArchiveError(f"Source directory missing: {source_dir}")
+
+    target_file.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        _compress_directory_to_tar(source_dir, target_file)
+        size_mb = get_file_size_mb(target_file)
+
+        if size_mb > max_size_mb:
+            safe_unlink(
+                target_file, f"Failed to delete oversized archive {target_file}"
+            )
+            raise ArchiveError(f"Archive too large: {size_mb:.1f}MB > {max_size_mb}MB")
+
+    except Exception as e:
+        safe_unlink(target_file, f"Failed to cleanup failed archive {target_file}")
+        raise ArchiveError(f"Archive creation failed: {e}") from e
+
+
+@timed_fn(logger=logger, name="Extracting archive")
+def extract_archive(archive_file: Path, target_dir: Path) -> None:
+    """Extract a compressed archive with security validation.
+
+    This function safely extracts a gzipped tar archive to a target directory
+    with security checks to prevent directory traversal attacks. It validates
+    both the archive and target paths, and inspects archive contents for
+    malicious paths before extraction.
+
+    Args:
+        archive_file: Path to the archive file to extract. Must exist and be
+                     within allowed directories (/app or /cache).
+        target_dir: Path to the directory where files will be extracted. Must
+                   be within allowed directories (/tmp/ or its parent).
+
+    Raises:
+        CacheValidationError: If paths are outside allowed directories or if
+                            archive contains unsafe paths (absolute paths or
+                            paths with '..' components).
+        ArchiveError: If archive file doesn't exist or extraction fails.
+    """
+    # Validate paths
+    validate_path_security(
+        str(archive_file),
+        ["/app", "/cache"],
+        f"Archive file {archive_file}",
+        CacheValidationError,
+    )
+    validate_path_security(
+        str(target_dir),
+        ["/tmp/", str(target_dir.parent)],
+        f"Target directory {target_dir}",
+        CacheValidationError,
+    )
+
+    if not archive_file.exists():
+        raise ArchiveError(f"Archive missing: {archive_file}")
+
+    try:
+        target_dir.mkdir(parents=True, exist_ok=True)
+
+        # First, perform security check by listing archive contents
+        list_cmd = ["tar", "-tzf", str(archive_file)]
+        result = subprocess.run(list_cmd, check=True, capture_output=True, text=True)
+
+        # Security check on all paths in the archive
+        for path in result.stdout.strip().split("\n"):
+            if path and (os.path.isabs(path) or ".." in path):
+                raise CacheValidationError(f"Unsafe path in archive: {path}")
+
+        # Extract using system tar command for better performance
+        extract_cmd = ["tar", "-xzf", str(archive_file), "-C", str(target_dir)]
+        subprocess.run(extract_cmd, check=True, capture_output=True, text=True)
+
+    except subprocess.CalledProcessError as e:
+        raise ArchiveError(f"tar extraction failed: {e.stderr}") from e
+    except Exception as e:
+        raise ArchiveError(f"Extraction failed: {e}") from e