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

{b10_transfer-0.1.2 → b10_transfer-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: b10-transfer
-Version: 0.1.2
+Version: 0.1.4
 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.2 → b10_transfer-0.1.4}/pyproject.toml

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

{b10_transfer-0.1.2 → b10_transfer-0.1.4}/src/b10_transfer/__init__.py

@@ -7,7 +7,7 @@ from .info import get_cache_info, list_available_caches
 from .constants import SaveStatus, LoadStatus
 
 # Version
-__version__ = "0.1.2"
+__version__ = "0.1.4"
 
 __all__ = [
     "CacheError",

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

@@ -129,3 +129,10 @@ class SaveStatus(Enum):
     SUCCESS = auto()
     ERROR = auto()
     SKIPPED = auto()
+
+
+class TransferStatus(Enum):
+    """Status values for file transfer operations."""
+
+    SUCCESS = auto()
+    ERROR = auto()

{b10_transfer-0.1.2 → b10_transfer-0.1.4}/src/b10_transfer/core.py

@@ -35,6 +35,7 @@ from .constants import (
     CACHE_INCOMPLETE_SUFFIX,
     LoadStatus,
     SaveStatus,
+    TransferStatus,
 )
 
 logger = logging.getLogger(__name__)
@@ -89,18 +90,6 @@ def load_compile_cache() -> LoadStatus:
             logger.info("Torch cache already loaded, skipping extraction")
             return LoadStatus.SKIPPED
 
-        # Initial disk space check for local operations
-        check_sufficient_disk_space(
-            work_dir, MIN_LOCAL_SPACE_MB, "cache load operations"
-        )
-        logger.debug(
-            f"Initial space check passed: {MIN_LOCAL_SPACE_MB:.1f}MB required on local machine"
-        )
-
-        # Start background space monitoring for local disk
-        space_monitor = CacheSpaceMonitor(MIN_LOCAL_SPACE_MB, work_dir)
-        space_monitor.start()
-
         # Create temp local copy
         with tempfile.NamedTemporaryFile(
             suffix=CACHE_FILE_EXTENSION, dir=work_dir, delete=False
@@ -110,30 +99,30 @@ def load_compile_cache() -> LoadStatus:
 
         try:
             with temp_file_cleanup(temp_path):
-                # Phase 1: Copy from b10fs to local temp file in separate process
+                # Phase 1: Copy from b10fs to local temp file using transfer()
                 @critical_section_b10fs_file_lock("copy_out")
                 def _monitored_copy_from_b10fs():
-                    logger.info(
-                        f"Starting copy from b10fs: {cache_file} -> {temp_path}"
-                    )
-                    run_monitored_process(
-                        _cache_copy_from_b10fs_worker,
-                        (str(cache_file), str(temp_path)),
-                        space_monitor,
-                        "b10fs to local copy",
-                    )
+                    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
-                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),
-                )
+                # 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()
 
             logger.info("Cache load complete")
             return LoadStatus.SUCCESS
@@ -142,9 +131,6 @@ def load_compile_cache() -> LoadStatus:
             logger.warning(f"Cache load interrupted: {e}")
             return LoadStatus.ERROR
 
-        finally:
-            space_monitor.stop()
-
 
 """
 FIXME(SRAY):
@@ -214,26 +200,11 @@ def save_compile_cache() -> SaveStatus:
             logger.info("Cache already exists in b10fs, skipping save")
             return SaveStatus.SKIPPED
 
-        # Initial disk space checks using calculated space requirements
-        check_sufficient_disk_space(
-            work_dir, MAX_CACHE_SIZE_MB, "local temp file creation"
-        )
-        check_sufficient_disk_space(
-            b10fs_dir, REQUIRED_B10FS_SPACE_MB, "cache save to b10fs"
-        )
-        logger.debug(
-            f"Initial space checks passed: {MAX_CACHE_SIZE_MB:.1f}MB local, {REQUIRED_B10FS_SPACE_MB:.1f}MB b10fs"
-        )
-
         temp_file = (
             b10fs_dir
             / f"{cache_filename}{CACHE_INCOMPLETE_SUFFIX}{CACHE_FILE_EXTENSION}"
         )
 
-        # Start background space monitoring
-        space_monitor = CacheSpaceMonitor(REQUIRED_B10FS_SPACE_MB, b10fs_dir)
-        space_monitor.start()
-
         with tempfile.NamedTemporaryFile(
             suffix=CACHE_FILE_EXTENSION, dir=work_dir, delete=False
         ) as f:
@@ -242,30 +213,32 @@ def save_compile_cache() -> SaveStatus:
 
         try:
             with temp_file_cleanup(local_temp):
-                # Phase 1: Compression in separate process
-                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",
-                )
-
-                b10fs_dir.mkdir(parents=True, exist_ok=True)
+                # Phase 1: Compression in separate process with space monitoring
+                space_monitor = CacheSpaceMonitor(REQUIRED_B10FS_SPACE_MB, b10fs_dir)
+                space_monitor.start()
 
-                # Phase 2: Copy to b10fs in separate process
-                @critical_section_b10fs_file_lock("copy_in")
-                def _monitored_copy_to_b10fs():
-                    logger.info(f"Starting copy to b10fs: {local_temp} -> {temp_file}")
+                try:
+                    logger.info(f"Starting compression: {torch_dir} -> {local_temp}")
                     run_monitored_process(
-                        _cache_copy_worker,
-                        (str(local_temp), str(temp_file)),
+                        _cache_compression_worker,
+                        (str(torch_dir), str(local_temp), MAX_CACHE_SIZE_MB),
                         space_monitor,
-                        "b10fs copy",
-                        cleanup_func=lambda: safe_unlink(
-                            temp_file, f"Failed to cleanup interrupted copy {temp_file}"
-                        ),
+                        "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()
 
@@ -282,8 +255,89 @@ def save_compile_cache() -> SaveStatus:
             logger.warning(f"Cache save interrupted: {e}")
             return SaveStatus.ERROR
 
-        finally:
-            space_monitor.stop()
+
+@timed_fn(logger=logger, name="Transferring file")
+@safe_execute("Transfer failed", TransferStatus.ERROR)
+def transfer(source: str, dest: str) -> TransferStatus:
+    """Transfer a file from source to destination with space monitoring.
+
+    This function copies a file from source to destination using the same
+    monitored process approach as the cache operations. It monitors disk space
+    at the destination and can interrupt the transfer if space becomes insufficient.
+
+    Args:
+        source: Path to the source file to copy.
+        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,
+                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).
+    """
+    source_path = Path(source)
+    dest_path = Path(dest)
+
+    # Validate source file exists
+    if not source_path.exists():
+        logger.error(f"Source file does not exist: {source}")
+        return TransferStatus.ERROR
+
+    # Create destination directory if it doesn't exist
+    dest_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Determine appropriate space threshold based on destination directory
+    dest_dir = dest_path.parent
+    if str(dest_dir).startswith(B10FS_CACHE_DIR):
+        # Transferring to b10fs - use b10fs space requirements
+        space_threshold_mb = REQUIRED_B10FS_SPACE_MB
+        logger.debug(
+            f"Transfer to b10fs detected, using {space_threshold_mb:.1f}MB threshold"
+        )
+    else:
+        # Transferring to local directory - use local space requirements
+        space_threshold_mb = MIN_LOCAL_SPACE_MB
+        logger.debug(
+            f"Transfer to local directory detected, using {space_threshold_mb:.1f}MB threshold"
+        )
+
+    # Initial disk space check
+    check_sufficient_disk_space(dest_dir, space_threshold_mb, "file transfer")
+    logger.debug(
+        f"Initial space check passed: {space_threshold_mb:.1f}MB required at destination"
+    )
+
+    # Start background space monitoring for destination directory
+    space_monitor = CacheSpaceMonitor(space_threshold_mb, dest_dir)
+    space_monitor.start()
+
+    try:
+        # Run monitored copy process
+        logger.info(f"Starting transfer: {source} -> {dest}")
+        run_monitored_process(
+            _cache_copy_worker,
+            (str(source_path), str(dest_path)),
+            space_monitor,
+            "file transfer",
+            cleanup_func=lambda: safe_unlink(
+                dest_path, f"Failed to cleanup interrupted transfer {dest_path}"
+            ),
+        )
+
+        logger.info("File transfer complete")
+        return TransferStatus.SUCCESS
+
+    except CacheOperationInterrupted as e:
+        logger.warning(f"File transfer interrupted: {e}")
+        return TransferStatus.ERROR
+
+    finally:
+        space_monitor.stop()
 
 
 @safe_execute("Clear failed", False)
@@ -347,23 +401,6 @@ def _cache_copy_worker(source_path_str: str, dest_path_str: str) -> None:
     shutil.copy2(source_path, dest_path)
 
 
-@worker_process("Copy from b10fs was cancelled before starting")
-def _cache_copy_from_b10fs_worker(source_path_str: str, dest_path_str: str) -> None:
-    """Worker process that handles file copy from b10fs to local machine.
-
-    This function runs in a separate process to copy the cache file from b10fs
-    to the local filesystem. It can be terminated externally if local disk space becomes insufficient.
-
-    Args:
-        source_path_str: String path to the source file in b10fs to copy.
-        dest_path_str: String path where the file will be copied locally.
-    """
-    source_path = Path(source_path_str)
-    dest_path = Path(dest_path_str)
-
-    shutil.copy2(source_path, dest_path)
-
-
 def _cleanup_torch_dir(torch_dir: Path) -> None:
     """Helper function to safely cleanup torch directory during interrupted extraction."""
     try: