b10-transfer 0.0.1__tar.gz → 0.0.2__tar.gz

{b10_transfer-0.0.1 → b10_transfer-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: b10-transfer
-Version: 0.0.1
+Version: 0.0.2
 Summary: Distributed PyTorch compilation cache for Baseten - Environment-aware, lock-free compilation cache management
 License: MIT
 Keywords: pytorch,torch.compile,cache,machine-learning,inference

{b10_transfer-0.0.1 → b10_transfer-0.0.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.0.1"
+version = "0.0.2"
 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>"]

{b10_transfer-0.0.1 → b10_transfer-0.0.2}/src/b10_transfer/__init__.py

@@ -21,7 +21,7 @@ from .info import get_cache_info, list_available_caches
 from .constants import SaveStatus, LoadStatus, TransferStatus, AsyncTransferStatus
 
 # Version
-__version__ = "0.0.1"
+__version__ = "0.0.2"
 
 __all__ = [
     "CacheError",

{b10_transfer-0.0.1 → b10_transfer-0.0.2}/src/b10_transfer/torch_cache.py

@@ -40,7 +40,7 @@ logger = logging.getLogger(__name__)
 
 
 def torch_cache_save_callback(
-    source_dir: Path, dest_file: Path, max_size_mb: int
+    source_dir: Path, dest_file: Path, max_size_mb: int = None, *args, **kwargs
 ) -> None:
     """Callback function for saving torch cache: compress then copy to b10fs.
 
@@ -51,8 +51,14 @@ def torch_cache_save_callback(
     Args:
         source_dir: Path to the torch cache directory to compress
         dest_file: Path to the final cache file in b10fs
-        max_size_mb: Maximum allowed archive size in megabytes
+        max_size_mb: Maximum allowed archive size in megabytes (can be passed as kwarg)
+        *args: Additional arguments passed by the transfer system (ignored)
+        **kwargs: Additional keyword arguments passed by the transfer system (may contain max_size_mb)
     """
+    # Handle max_size_mb from kwargs if not provided as positional argument
+    if max_size_mb is None:
+        max_size_mb = kwargs.get("max_size_mb", MAX_CACHE_SIZE_MB)
+
     work_dir = Path(LOCAL_WORK_DIR)
 
     # Create temporary archive in local work directory
@@ -99,7 +105,9 @@ def torch_cache_save_callback(
         raise
 
 
-def torch_cache_load_callback(source_file: Path, dest_dir: Path) -> None:
+def torch_cache_load_callback(
+    source_file: Path, dest_dir: Path, *args, **kwargs
+) -> None:
     """Callback function for loading torch cache: copy from b10fs then extract.
 
     This function handles the torch-specific load logic:
@@ -109,6 +117,8 @@ def torch_cache_load_callback(source_file: Path, dest_dir: Path) -> None:
     Args:
         source_file: Path to the cache file in b10fs
         dest_dir: Path to the torch cache directory where files will be extracted
+        *args: Additional arguments passed by the transfer system (ignored)
+        **kwargs: Additional keyword arguments passed by the transfer system (ignored)
     """
     work_dir = Path(LOCAL_WORK_DIR)