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

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.1.0"
 
 __all__ = [
     "CacheError",

b10_transfer/async_transfers.py

@@ -39,6 +39,7 @@ class TransferProgress:
             AsyncTransferStatus.ERROR,
             AsyncTransferStatus.INTERRUPTED,
             AsyncTransferStatus.CANCELLED,
+            AsyncTransferStatus.DOES_NOT_EXIST,
         ]:
             self.completed_at = datetime.now()
 
@@ -140,6 +141,12 @@ class AsyncTransferManager:
                     "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"
@@ -166,6 +173,7 @@ class AsyncTransferManager:
             AsyncTransferStatus.ERROR,
             AsyncTransferStatus.INTERRUPTED,
             AsyncTransferStatus.CANCELLED,
+            AsyncTransferStatus.DOES_NOT_EXIST,
         ]
 
     def wait_for_completion(

b10_transfer/constants.py

@@ -113,6 +113,7 @@ class WorkerStatus(Enum):
     SUCCESS = auto()
     ERROR = auto()
     CANCELLED = auto()
+    FILE_NOT_FOUND = auto()
 
 
 class LoadStatus(Enum):
@@ -138,6 +139,7 @@ class TransferStatus(Enum):
     SUCCESS = auto()
     ERROR = auto()
     INTERRUPTED = auto()
+    DOES_NOT_EXIST = auto()
 
 
 class AsyncTransferStatus(Enum):
@@ -147,3 +149,4 @@ class AsyncTransferStatus(Enum):
     ERROR = auto()
     INTERRUPTED = auto()
     CANCELLED = auto()
+    DOES_NOT_EXIST = auto()

b10_transfer/core.py

@@ -11,6 +11,7 @@ from .space_monitor import (
     check_sufficient_disk_space,
     CacheSpaceMonitor,
     CacheOperationInterrupted,
+    CacheFileNotFoundError,
     run_monitored_process,
 )
 from .constants import (
@@ -118,9 +119,13 @@ def transfer(
         if primary_monitor is None:
             # No monitoring requested, execute callback directly
             logger.info(f"Starting transfer (no monitoring): {source} -> {dest}")
-            callback(source, dest, *callback_args, **callback_kwargs)
-            logger.info("Transfer complete")
-            return TransferStatus.SUCCESS
+            try:
+                callback(source, dest, *callback_args, **callback_kwargs)
+                logger.info("Transfer complete")
+                return TransferStatus.SUCCESS
+            except (FileNotFoundError, CacheFileNotFoundError) as e:
+                logger.info(f"Transfer failed - file not found: {e}")
+                return TransferStatus.DOES_NOT_EXIST
 
         # Start the primary space monitor
         primary_monitor.start()
@@ -151,6 +156,10 @@ def transfer(
                 logger.info("Transfer complete (unmonitored)")
                 return TransferStatus.SUCCESS
 
+        except (FileNotFoundError, CacheFileNotFoundError) as e:
+            logger.info(f"Transfer failed - file not found: {e}")
+            return TransferStatus.DOES_NOT_EXIST
+
         except CacheOperationInterrupted as e:
             logger.warning(f"Transfer interrupted: {e}")
             return TransferStatus.INTERRUPTED

b10_transfer/space_monitor.py

@@ -24,6 +24,12 @@ class CacheOperationInterrupted(Exception):
     pass
 
 
+class CacheFileNotFoundError(Exception):
+    """Raised when a cache file is not found during transfer operations."""
+
+    pass
+
+
 def worker_process(cancelled_message: str):
     """Decorator for worker process functions to handle common try/catch/result_queue pattern.
 
@@ -60,6 +66,8 @@ def worker_process(cancelled_message: str):
                 # If we get here, the function completed successfully
                 result_queue.put((WorkerStatus.SUCCESS.value, None))
 
+            except FileNotFoundError as e:
+                result_queue.put((WorkerStatus.FILE_NOT_FOUND.value, str(e)))
             except Exception as e:
                 result_queue.put((WorkerStatus.ERROR.value, str(e)))
 
@@ -286,6 +294,11 @@ def run_monitored_process(
                 if cleanup_func:
                     cleanup_func()
                 raise CacheOperationInterrupted(error_msg)
+            elif status == WorkerStatus.FILE_NOT_FOUND.value:
+                logger.info(
+                    f"{operation_name} worker failed - file not found: {error_msg}"
+                )
+                raise CacheFileNotFoundError(error_msg)
             # status == WorkerStatus.SUCCESS.value - continue normally
 
         logger.debug(f"{operation_name} completed successfully")

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)
 
@@ -126,6 +136,8 @@ def torch_cache_load_callback(source_file: Path, dest_dir: Path) -> None:
             @critical_section_b10fs_file_lock("copy_out")
             def _copy_from_b10fs():
                 logger.info(f"Copying from b10fs: {source_file} -> {temp_archive}")
+                if not source_file.exists():
+                    raise FileNotFoundError(f"Cache file not found: {source_file}")
                 shutil.copy2(source_file, temp_archive)
 
             _copy_from_b10fs()

{b10_transfer-0.0.1.dist-info → b10_transfer-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: b10-transfer
-Version: 0.0.1
+Version: 0.1.0
 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.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+b10_transfer/__init__.py,sha256=Z_p771iwuROcCSNWKjUZ9j-V7ICmbtwr_qet5FCsnkQ,1400
+b10_transfer/archive.py,sha256=GKb0mi0-YeM7ch4FLAoOLHXw0T6LkRerYad2N2y9TYM,6400
+b10_transfer/async_torch_cache.py,sha256=4hMjVR44SLlGes25e_cjgMTywFfIYjH0TnUmg9o-iyI,1903
+b10_transfer/async_transfers.py,sha256=luqdIStT_j4YduImY67HvX5WDurqV9Q5RjEyMI7bh1k,9476
+b10_transfer/cleanup.py,sha256=xjKStmBjaarZPxhPTT1-Ds_pvUR7kdJw5Kp19BLvzzY,6224
+b10_transfer/constants.py,sha256=R2JE_634Ri_9rf8adwiAzcfiej5weAGP1x1ccSZLX8k,4829
+b10_transfer/core.py,sha256=d-aaQwKYqKIafBYBNahNcnOpwcanOSrWLwdzXpjVLBs,6350
+b10_transfer/environment.py,sha256=aC0biEMQrtHk0ke_3epdcq1X9J5fPmPpBVt0fH7XF2Y,5625
+b10_transfer/info.py,sha256=I3iOuImZ5r6DMJTDeBtVvzlSn6IuyPJbLJYUO_OF0ks,6299
+b10_transfer/space_monitor.py,sha256=G_3wLSJa7HTCihSpLoow2oKo2cARJ2PtvY1XOQZl3-s,11028
+b10_transfer/torch_cache.py,sha256=e41mDdnP_h61WNwB7TG5c4a7ecw0-K63ytJiKsX0keY,14907
+b10_transfer/utils.py,sha256=Stee0DFK-8MRRYNIocqaK64cJvfs4jPW3Mpx7zkWV6Y,11932
+b10_transfer-0.1.0.dist-info/METADATA,sha256=wc0a--Bgr-7filvyS4uUAic9fO1JJbKqc5iNp36A-iU,7502
+b10_transfer-0.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+b10_transfer-0.1.0.dist-info/RECORD,,

b10_transfer-0.0.1.dist-info/RECORD

@@ -1,15 +0,0 @@
-b10_transfer/__init__.py,sha256=o1ej-OtAOsfrJbvh5C3PnqxW2qfcO7l8rllVD-07lXE,1400
-b10_transfer/archive.py,sha256=GKb0mi0-YeM7ch4FLAoOLHXw0T6LkRerYad2N2y9TYM,6400
-b10_transfer/async_torch_cache.py,sha256=4hMjVR44SLlGes25e_cjgMTywFfIYjH0TnUmg9o-iyI,1903
-b10_transfer/async_transfers.py,sha256=AAML562qYzF9NyX9AdfiJ0OcQw6vXr985IZWXZSot9Q,9083
-b10_transfer/cleanup.py,sha256=xjKStmBjaarZPxhPTT1-Ds_pvUR7kdJw5Kp19BLvzzY,6224
-b10_transfer/constants.py,sha256=KjSUO6heScDJXQwFlHdeNV4KBBqKz7CKeJzo44-9qMM,4745
-b10_transfer/core.py,sha256=BOnA6FXkZRm74_CtQBMudpx3q7HTEGEORUV26fb6cvQ,5920
-b10_transfer/environment.py,sha256=aC0biEMQrtHk0ke_3epdcq1X9J5fPmPpBVt0fH7XF2Y,5625
-b10_transfer/info.py,sha256=I3iOuImZ5r6DMJTDeBtVvzlSn6IuyPJbLJYUO_OF0ks,6299
-b10_transfer/space_monitor.py,sha256=5pwW643KAHI3mtT61hYf29953UD9LekzWFF1K-QeYbw,10529
-b10_transfer/torch_cache.py,sha256=Oe_OeUPGAlmK9wY-L9w4aPaXOoMnL_kD596hew6ETcw,14192
-b10_transfer/utils.py,sha256=Stee0DFK-8MRRYNIocqaK64cJvfs4jPW3Mpx7zkWV6Y,11932
-b10_transfer-0.0.1.dist-info/METADATA,sha256=hESeWyidAEbtWkIgepBn1Cxlo9--jIj9vcLxM4zP7lY,7502
-b10_transfer-0.0.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-b10_transfer-0.0.1.dist-info/RECORD,,