b10-transfer 0.1.1__tar.gz → 0.1.2__tar.gz

b10_transfer-0.1.2/PKG-INFO

@@ -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/README.md

@@ -0,0 +1,97 @@
+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.1 → b10_transfer-0.1.2}/pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.1.1"
-description = "Distributed PyTorch compilation cache for Baseten - Environment-aware, lock-free compilation cache management"
+version = "0.1.2"
+description = "Distributed PyTorch file transfer for Baseten - Environment-aware, lock-free file transfer 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"
@@ -13,7 +13,7 @@ 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"]
+keywords = ["pytorch", "file-transfer", "cache", "machine-learning", "inference"]
 classifiers = [
     "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",

b10_transfer-0.1.2/src/b10_transfer/__init__.py

@@ -0,0 +1,23 @@
+"""B10 Transfer - Lock-free PyTorch file transfer for Baseten."""
+
+from .core import load_compile_cache, save_compile_cache, clear_local_cache
+from .utils import CacheError, CacheValidationError
+from .space_monitor import CacheOperationInterrupted
+from .info import get_cache_info, list_available_caches
+from .constants import SaveStatus, LoadStatus
+
+# Version
+__version__ = "0.1.2"
+
+__all__ = [
+    "CacheError",
+    "CacheValidationError",
+    "CacheOperationInterrupted",
+    "SaveStatus",
+    "LoadStatus",
+    "load_compile_cache",
+    "save_compile_cache",
+    "clear_local_cache",
+    "get_cache_info",
+    "list_available_caches",
+]

{b10_transfer-0.1.1 → b10_transfer-0.1.2}/src/b10_transfer/constants.py

@@ -36,8 +36,7 @@ B10FS_CACHE_DIR = validate_path_security(
     _b10fs_cache_dir, [_REQUIRED_TORCH_CACHE_DIR_PREFIX], "B10FS_CACHE_DIR"
 )
 
-# Validate LOCAL_WORK_DIR - allow /app, /tmp, and /cache paths.
-# This is like a "scratch" directory where you can do work (like compression/archival for example)
+# Validate LOCAL_WORK_DIR - allow /app, /tmp, and /cache paths
 _local_work_dir = os.getenv("LOCAL_WORK_DIR", "/app")
 LOCAL_WORK_DIR = validate_path_security(
     _local_work_dir, ["/app/", "/tmp/", "/cache/"], "LOCAL_WORK_DIR"
@@ -113,7 +112,6 @@ class WorkerStatus(Enum):
     SUCCESS = auto()
     ERROR = auto()
     CANCELLED = auto()
-    FILE_NOT_FOUND = auto()
 
 
 class LoadStatus(Enum):
@@ -131,22 +129,3 @@ class SaveStatus(Enum):
     SUCCESS = auto()
     ERROR = auto()
     SKIPPED = auto()
-
-
-class TransferStatus(Enum):
-    """Status values for generic transfer operations."""
-
-    SUCCESS = auto()
-    ERROR = auto()
-    INTERRUPTED = auto()
-    DOES_NOT_EXIST = auto()
-
-
-class AsyncTransferStatus(Enum):
-    NOT_STARTED = auto()
-    IN_PROGRESS = auto()
-    SUCCESS = auto()
-    ERROR = auto()
-    INTERRUPTED = auto()
-    CANCELLED = auto()
-    DOES_NOT_EXIST = auto()