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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: b10-transfer
-Version: 0.1.3
+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.3 → 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.3"
+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.3 → 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.3"
+__version__ = "0.1.4"
 
 __all__ = [
     "CacheError",

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

@@ -90,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
@@ -111,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
@@ -143,9 +131,6 @@ def load_compile_cache() -> LoadStatus:
             logger.warning(f"Cache load interrupted: {e}")
             return LoadStatus.ERROR
 
-        finally:
-            space_monitor.stop()
-
 
 """
 FIXME(SRAY):
@@ -215,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:
@@ -243,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()
 
@@ -283,9 +255,6 @@ 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)
@@ -432,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: