b10-transfer 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

b10_transfer/__init__.py

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

b10_transfer/constants.py

@@ -114,25 +114,10 @@ class WorkerStatus(Enum):
     CANCELLED = auto()
 
 
-class LoadStatus(Enum):
-    """Status values for cache loading operations."""
-
-    SUCCESS = auto()
-    ERROR = auto()
-    DOES_NOT_EXIST = auto()
-    SKIPPED = auto()
-
-
-class SaveStatus(Enum):
-    """Status values for cache saving operations."""
-
-    SUCCESS = auto()
-    ERROR = auto()
-    SKIPPED = auto()
-
-
-class TransferStatus(Enum):
-    """Status values for file transfer operations."""
+class OperationStatus(Enum):
+    """Status values for all b10-transfer operations (load, save, transfer)."""
 
     SUCCESS = auto()
     ERROR = auto()
+    DOES_NOT_EXIST = auto()  # Used by load operations when cache file not found
+    SKIPPED = auto()  # Used by load/save operations when operation not needed

b10_transfer/core.py

@@ -33,17 +33,83 @@ from .constants import (
     CACHE_FILE_EXTENSION,
     CACHE_LATEST_SUFFIX,
     CACHE_INCOMPLETE_SUFFIX,
-    LoadStatus,
-    SaveStatus,
-    TransferStatus,
+    OperationStatus,
 )
 
 logger = logging.getLogger(__name__)
 
 
+def _setup_cache_paths():
+    """Common setup for cache operations - returns paths and performs cleanup."""
+    # Cooperative cleanup of stale shared resources
+    cooperative_cleanup_b10fs()
+
+    b10fs_dir = Path(B10FS_CACHE_DIR)
+    torch_dir = Path(TORCH_CACHE_DIR)
+    work_dir = Path(LOCAL_WORK_DIR)
+
+    return b10fs_dir, torch_dir, work_dir
+
+
+def _get_cache_file_paths(cache_filename: str, b10fs_dir: Path):
+    """Generate cache file paths for a given cache filename."""
+    final_file = (
+        b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
+    )
+    temp_file = (
+        b10fs_dir / f"{cache_filename}{CACHE_INCOMPLETE_SUFFIX}{CACHE_FILE_EXTENSION}"
+    )
+    return final_file, temp_file
+
+
+def _run_with_space_monitoring(
+    space_threshold_mb: float,
+    monitor_dir: Path,
+    operation_name: str,
+    worker_func,
+    worker_args: tuple,
+    cleanup_func=None,
+):
+    """Helper to run an operation with space monitoring."""
+    space_monitor = CacheSpaceMonitor(space_threshold_mb, monitor_dir)
+    space_monitor.start()
+
+    try:
+        logger.info(
+            f"Starting {operation_name}: {' -> '.join(str(arg) for arg in worker_args[:2])}"
+        )
+        run_monitored_process(
+            worker_func,
+            worker_args,
+            space_monitor,
+            operation_name,
+            cleanup_func=cleanup_func,
+        )
+    finally:
+        space_monitor.stop()
+
+
+def _transfer_with_b10fs_lock(
+    source: str, dest: str, lock_type: str, cleanup_on_failure=True
+):
+    """Transfer a file with b10fs file locking and error handling."""
+
+    @critical_section_b10fs_file_lock(lock_type)
+    def _locked_transfer():
+        result = transfer(source, dest)
+        if result != OperationStatus.SUCCESS:
+            if cleanup_on_failure:
+                safe_unlink(
+                    Path(dest), f"Failed to cleanup after failed transfer {dest}"
+                )
+            raise Exception(f"Failed to transfer {source} -> {dest}")
+
+    _locked_transfer()
+
+
 @timed_fn(logger=logger, name="Loading compile cache")
 @safe_execute("Load failed", False)
-def load_compile_cache() -> LoadStatus:
+def load_compile_cache() -> OperationStatus:
     """Load PyTorch compilation cache from b10fs to local torch cache directory.
 
     This function implements a lock-free pattern to safely load cached PyTorch
@@ -55,40 +121,33 @@ def load_compile_cache() -> LoadStatus:
     extraction phases, interrupting operations if space falls below MIN_LOCAL_SPACE_MB.
 
     Returns:
-        LoadStatus:
-              LoadStatus.SUCCESS if cache was successfully loaded
-              LoadStatus.SKIPPED if already exists
-              LoadStatus.ERROR if b10fs is unavailable, local disk space is insufficient, or loading failed.
-              LoadStatus.DOES_NOT_EXIST if no cache file was found.
+        OperationStatus:
+              OperationStatus.SUCCESS if cache was successfully loaded
+              OperationStatus.SKIPPED if already exists
+              OperationStatus.ERROR if b10fs is unavailable, local disk space is insufficient, or loading failed.
+              OperationStatus.DOES_NOT_EXIST if no cache file was found.
 
     Raises:
-        CacheValidationError: If b10fs is not enabled (caught and returns LoadStatus.ERROR).
+        CacheValidationError: If b10fs is not enabled (caught and returns OperationStatus.ERROR).
         CacheOperationInterrupted: If operations interrupted due to insufficient
-                                  local disk space (caught and returns LoadStatus.ERROR).
-        Exception: Any other errors during loading (caught and returns LoadStatus.ERROR).
+                                  local disk space (caught and returns OperationStatus.ERROR).
+        Exception: Any other errors during loading (caught and returns OperationStatus.ERROR).
     """
     with cache_operation("Load"):
-        # Cooperative cleanup of stale shared resources
-        cooperative_cleanup_b10fs()
-
-        b10fs_dir = Path(B10FS_CACHE_DIR)
-        torch_dir = Path(TORCH_CACHE_DIR)
-        work_dir = Path(LOCAL_WORK_DIR)
+        b10fs_dir, torch_dir, work_dir = _setup_cache_paths()
 
         cache_filename = get_cache_filename()
-        cache_file = (
-            b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
-        )
-        logger.debug(f"Looking for cache file: {cache_file}")
+        final_file, _ = _get_cache_file_paths(cache_filename, b10fs_dir)
+        logger.debug(f"Looking for cache file: {final_file}")
 
-        if not cache_file.exists():
+        if not final_file.exists():
             logger.info("No cache file found in b10fs")
-            return LoadStatus.DOES_NOT_EXIST
+            return OperationStatus.DOES_NOT_EXIST
 
         # Skip if already loaded
         if torch_dir.exists() and any(torch_dir.iterdir()):
             logger.info("Torch cache already loaded, skipping extraction")
-            return LoadStatus.SKIPPED
+            return OperationStatus.SKIPPED
 
         # Create temp local copy
         with tempfile.NamedTemporaryFile(
@@ -99,37 +158,30 @@ def load_compile_cache() -> LoadStatus:
 
         try:
             with temp_file_cleanup(temp_path):
-                # Phase 1: Copy from b10fs to local temp file using transfer()
-                @critical_section_b10fs_file_lock("copy_out")
-                def _monitored_copy_from_b10fs():
-                    result = transfer(str(cache_file), str(temp_path))
-                    if result != TransferStatus.SUCCESS:
-                        raise Exception("Failed to copy cache file from b10fs")
-
-                _monitored_copy_from_b10fs()
-
-                # Phase 2: Extract archive in separate process with space monitoring
-                space_monitor = CacheSpaceMonitor(MIN_LOCAL_SPACE_MB, work_dir)
-                space_monitor.start()
-
-                try:
-                    logger.info(f"Starting extraction: {temp_path} -> {torch_dir}")
-                    run_monitored_process(
-                        _cache_extract_worker,
-                        (str(temp_path), str(torch_dir)),
-                        space_monitor,
-                        "archive extraction",
-                        cleanup_func=lambda: _cleanup_torch_dir(torch_dir),
-                    )
-                finally:
-                    space_monitor.stop()
+                # Phase 1: Copy from b10fs to local temp file
+                _transfer_with_b10fs_lock(
+                    str(final_file),
+                    str(temp_path),
+                    "copy_out",
+                    cleanup_on_failure=False,
+                )
+
+                # Phase 2: Extract archive with space monitoring
+                _run_with_space_monitoring(
+                    MIN_LOCAL_SPACE_MB,
+                    work_dir,
+                    "archive extraction",
+                    _cache_extract_worker,
+                    (str(temp_path), str(torch_dir)),
+                    cleanup_func=lambda: _cleanup_torch_dir(torch_dir),
+                )
 
             logger.info("Cache load complete")
-            return LoadStatus.SUCCESS
+            return OperationStatus.SUCCESS
 
         except CacheOperationInterrupted as e:
             logger.warning(f"Cache load interrupted: {e}")
-            return LoadStatus.ERROR
+            return OperationStatus.ERROR
 
 
 """
@@ -150,7 +202,7 @@ More things to consider:
 
 @timed_fn(logger=logger, name="Saving compile cache")
 @safe_execute("Save failed", False)
-def save_compile_cache() -> SaveStatus:
+def save_compile_cache() -> OperationStatus:
     """Save local PyTorch compilation cache to b10fs using atomic journal pattern.
 
     This function creates an archive of the local torch cache directory and
@@ -164,46 +216,34 @@ def save_compile_cache() -> SaveStatus:
     space becomes insufficient, finally performing an atomic rename to the final cache file.
 
     Returns:
-        SaveStatus:
-              SaveStatus.SUCCESS if cache was successfully saved or already exists
-              SaveStatus.ERROR if b10fs is unavailable, insufficient disk space caused interruption,
+        OperationStatus:
+              OperationStatus.SUCCESS if cache was successfully saved or already exists
+              OperationStatus.ERROR if b10fs is unavailable, insufficient disk space caused interruption,
                 no cache exists to save, or saving failed.
-              SaveStatus.SKIPPED if no cache exists to save or cache already exists in b10fs
+              OperationStatus.SKIPPED if no cache exists to save or cache already exists in b10fs
 
     Raises:
-        CacheValidationError: If b10fs is not enabled (caught and returns SaveStatus.ERROR).
+        CacheValidationError: If b10fs is not enabled (caught and returns OperationStatus.ERROR).
         CacheOperationInterrupted: If operations interrupted due to insufficient
-                                  disk space (caught and returns SaveStatus.ERROR).
-        ArchiveError: If archive creation fails (caught and returns SaveStatus.ERROR).
-        Exception: Any other errors during saving (caught and returns SaveStatus.ERROR).
+                                  disk space (caught and returns OperationStatus.ERROR).
+        ArchiveError: If archive creation fails (caught and returns OperationStatus.ERROR).
+        Exception: Any other errors during saving (caught and returns OperationStatus.ERROR).
     """
     with cache_operation("Save"):
-        # Cooperative cleanup of stale shared resources
-        cooperative_cleanup_b10fs()
-
-        b10fs_dir = Path(B10FS_CACHE_DIR)
-        torch_dir = Path(TORCH_CACHE_DIR)
-        work_dir = Path(LOCAL_WORK_DIR)
+        b10fs_dir, torch_dir, work_dir = _setup_cache_paths()
 
         # Check if anything to save
         if not torch_dir.exists() or not any(torch_dir.iterdir()):
             logger.info("No torch cache to save")
-            return SaveStatus.SKIPPED
+            return OperationStatus.SKIPPED
 
         cache_filename = get_cache_filename()
-        final_file = (
-            b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
-        )
+        final_file, temp_file = _get_cache_file_paths(cache_filename, b10fs_dir)
 
         # Check for existing cache first (early exit)
         if final_file.exists():
             logger.info("Cache already exists in b10fs, skipping save")
-            return SaveStatus.SKIPPED
-
-        temp_file = (
-            b10fs_dir
-            / f"{cache_filename}{CACHE_INCOMPLETE_SUFFIX}{CACHE_FILE_EXTENSION}"
-        )
+            return OperationStatus.SKIPPED
 
         with tempfile.NamedTemporaryFile(
             suffix=CACHE_FILE_EXTENSION, dir=work_dir, delete=False
@@ -213,34 +253,19 @@ def save_compile_cache() -> SaveStatus:
 
         try:
             with temp_file_cleanup(local_temp):
-                # Phase 1: Compression in separate process with space monitoring
-                space_monitor = CacheSpaceMonitor(REQUIRED_B10FS_SPACE_MB, b10fs_dir)
-                space_monitor.start()
-
-                try:
-                    logger.info(f"Starting compression: {torch_dir} -> {local_temp}")
-                    run_monitored_process(
-                        _cache_compression_worker,
-                        (str(torch_dir), str(local_temp), MAX_CACHE_SIZE_MB),
-                        space_monitor,
-                        "compression",
-                    )
-                finally:
-                    space_monitor.stop()
-
-                # Phase 2: Copy to b10fs using transfer()
-                @critical_section_b10fs_file_lock("copy_in")
-                def _monitored_copy_to_b10fs():
-                    result = transfer(str(local_temp), str(temp_file))
-                    if result != TransferStatus.SUCCESS:
-                        # Clean up the temp file if transfer failed
-                        safe_unlink(
-                            temp_file,
-                            f"Failed to cleanup after failed copy {temp_file}",
-                        )
-                        raise Exception("Failed to copy cache file to b10fs")
-
-                _monitored_copy_to_b10fs()
+                # Phase 1: Compression with space monitoring
+                _run_with_space_monitoring(
+                    REQUIRED_B10FS_SPACE_MB,
+                    b10fs_dir,
+                    "compression",
+                    _cache_compression_worker,
+                    (str(torch_dir), str(local_temp), MAX_CACHE_SIZE_MB),
+                )
+
+                # Phase 2: Copy to b10fs with locking
+                _transfer_with_b10fs_lock(
+                    str(local_temp), str(temp_file), "copy_in", cleanup_on_failure=True
+                )
 
                 # Phase 3: Atomic rename (fast, don't interrupt)
                 logger.info(
@@ -249,16 +274,16 @@ def save_compile_cache() -> SaveStatus:
                 temp_file.rename(final_file)
 
             logger.info("Cache save complete")
-            return SaveStatus.SUCCESS
+            return OperationStatus.SUCCESS
 
         except CacheOperationInterrupted as e:
             logger.warning(f"Cache save interrupted: {e}")
-            return SaveStatus.ERROR
+            return OperationStatus.ERROR
 
 
 @timed_fn(logger=logger, name="Transferring file")
-@safe_execute("Transfer failed", TransferStatus.ERROR)
-def transfer(source: str, dest: str) -> TransferStatus:
+@safe_execute("Transfer failed", OperationStatus.ERROR)
+def transfer(source: str, dest: str) -> OperationStatus:
     """Transfer a file from source to destination with space monitoring.
 
     This function copies a file from source to destination using the same
@@ -270,15 +295,15 @@ def transfer(source: str, dest: str) -> TransferStatus:
         dest: Path to the destination where the file will be copied.
 
     Returns:
-        TransferStatus:
-              TransferStatus.SUCCESS if transfer was successful
-              TransferStatus.ERROR if transfer failed due to insufficient disk space,
+        OperationStatus:
+              OperationStatus.SUCCESS if transfer was successful
+              OperationStatus.ERROR if transfer failed due to insufficient disk space,
                 file not found, or other errors.
 
     Raises:
         CacheOperationInterrupted: If transfer interrupted due to insufficient
-                                  disk space (caught and returns TransferStatus.ERROR).
-        Exception: Any other errors during transfer (caught and returns TransferStatus.ERROR).
+                                  disk space (caught and returns OperationStatus.ERROR).
+        Exception: Any other errors during transfer (caught and returns OperationStatus.ERROR).
     """
     source_path = Path(source)
     dest_path = Path(dest)
@@ -286,7 +311,7 @@ def transfer(source: str, dest: str) -> TransferStatus:
     # Validate source file exists
     if not source_path.exists():
         logger.error(f"Source file does not exist: {source}")
-        return TransferStatus.ERROR
+        return OperationStatus.ERROR
 
     # Create destination directory if it doesn't exist
     dest_path.parent.mkdir(parents=True, exist_ok=True)
@@ -330,11 +355,11 @@ def transfer(source: str, dest: str) -> TransferStatus:
         )
 
         logger.info("File transfer complete")
-        return TransferStatus.SUCCESS
+        return OperationStatus.SUCCESS
 
     except CacheOperationInterrupted as e:
         logger.warning(f"File transfer interrupted: {e}")
-        return TransferStatus.ERROR
+        return OperationStatus.ERROR
 
     finally:
         space_monitor.stop()

{b10_transfer-0.1.4.dist-info → b10_transfer-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: b10-transfer
-Version: 0.1.4
+Version: 0.1.5
 Summary: Distributed PyTorch file transfer for Baseten - Environment-aware, lock-free file transfer management
 License: MIT
 Keywords: pytorch,file-transfer,cache,machine-learning,inference

{b10_transfer-0.1.4.dist-info → b10_transfer-0.1.5.dist-info}/RECORD

@@ -1,12 +1,12 @@
-b10_transfer/__init__.py,sha256=LKMroIusY1itfMVrJT07xLS1XVehwr54Wk5dhEl8MzY,641
+b10_transfer/__init__.py,sha256=S7FJhF-BMPbXBqjQbXIj_Dl0r0-kSQD91l-wSD7COTI,647
 b10_transfer/archive.py,sha256=GKb0mi0-YeM7ch4FLAoOLHXw0T6LkRerYad2N2y9TYM,6400
 b10_transfer/cleanup.py,sha256=3RnqWNGMCcko5GQdq1Gr9VPpGzAF5J6x7xjIH9SNZ78,6226
-b10_transfer/constants.py,sha256=qCViKTyfHTLpiFVF2SwsbHp2IMz3kg3syxJfgRAq2dc,4446
-b10_transfer/core.py,sha256=XWLuwjHXuhh-6abZMAl2yuLB7R2deyUc6gGPn6-Yfkc,17006
+b10_transfer/constants.py,sha256=iuLShDW6hInhyz2YTQ8CzBanqW4chCkQOAzPZkCtOoA,4322
+b10_transfer/core.py,sha256=ScD-O9ot-ciVhX121fCchDu8_vM7izSFMxPtenU5RS0,17381
 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.4.dist-info/METADATA,sha256=69s3ACBUFzGB7J97eVt4aCGSXrIpld1oV0Wj8Z0HLZ8,4108
-b10_transfer-0.1.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-b10_transfer-0.1.4.dist-info/RECORD,,
+b10_transfer-0.1.5.dist-info/METADATA,sha256=lPyE9VvN1srmMZxn7tCOPEEkPHRZE8JsdAVIOeW2SPs,4108
+b10_transfer-0.1.5.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+b10_transfer-0.1.5.dist-info/RECORD,,