b10-transfer 0.1.5__tar.gz → 0.1.7__tar.gz

{b10_transfer-0.1.5 → b10_transfer-0.1.7}/PKG-INFO

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

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.1.5"
+version = "0.1.7"
 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.5 → b10_transfer-0.1.7}/src/b10_transfer/__init__.py

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

b10_transfer-0.1.5/src/b10_transfer/core.py → b10_transfer-0.1.7/src/b10_transfer/cache.py

@@ -1,11 +1,13 @@
-import os
+"""Cache operations for PyTorch compilation artifacts.
+
+This module provides functions for loading and saving PyTorch compilation cache
+to/from b10fs shared storage using atomic operations and space monitoring.
+"""
+
 import logging
 import tempfile
-import shutil
 from pathlib import Path
 
-import time
-
 from .environment import get_cache_filename
 from .cleanup import cooperative_cleanup_b10fs
 from .utils import (
@@ -17,7 +19,6 @@ from .utils import (
     safe_unlink,
 )
 from .space_monitor import (
-    check_sufficient_disk_space,
     CacheSpaceMonitor,
     CacheOperationInterrupted,
     run_monitored_process,
@@ -35,10 +36,27 @@ from .constants import (
     CACHE_INCOMPLETE_SUFFIX,
     OperationStatus,
 )
+from .core import transfer
 
 logger = logging.getLogger(__name__)
 
 
+"""
+FIXME(SRAY):
+What about the case in @b10-transfer/ where a single pod finishes an inference request,
+and then the client calls save_compile_cache. And while we are creating the local archive,
+another inference call on the same pod is kicked off, which then modifies the torch cache.
+How would this be handled? Maybe just accept that the cache will be recompiled/overwritten?
+Otherwise you'd need application level coordination to ensure that the cache is not modified
+while we are creating the archive, but this doesn't really seem like a good idea in terms of adoption.
+
+FIXME(SR):
+More things to consider:
+- [possible] What if b10fs dies *during* an op? right now we check for b10fs availability in the beginning of the op... Add some constants instead of just False for load().
+- [possible, and really bad if it happens] potential memory exhaustion during compression if the cache is super super large. very very edge case. higher compression levels also have high memory usage.
+"""
+
+
 def _setup_cache_paths():
     """Common setup for cache operations - returns paths and performs cleanup."""
     # Cooperative cleanup of stale shared resources
@@ -184,22 +202,6 @@ def load_compile_cache() -> OperationStatus:
             return OperationStatus.ERROR
 
 
-"""
-FIXME(SRAY):
-What about the case in @b10-transfer/ where a single pod finishes an inference request,
-and then the client calls save_compile_cache. And while we are creating the local archive,
-another inference call on the same pod is kicked off, which then modifies the torch cache.
-How would this be handled? Maybe just accept that the cache will be recompiled/overwritten?
-Otherwise you'd need application level coordination to ensure that the cache is not modified
-while we are creating the archive, but this doesn't really seem like a good idea in terms of adoption.
-
-FIXME(SR):
-More things to consider:
-- [possible] What if b10fs dies *during* an op? right now we check for b10fs availability in the beginning of the op... Add some constants instead of just False for load().
-- [possible, and really bad if it happens] potential memory exhaustion during compression if the cache is super super large. very very edge case. higher compression levels also have high memory usage.
-"""
-
-
 @timed_fn(logger=logger, name="Saving compile cache")
 @safe_execute("Save failed", False)
 def save_compile_cache() -> OperationStatus:
@@ -281,90 +283,6 @@ def save_compile_cache() -> OperationStatus:
             return OperationStatus.ERROR
 
 
-@timed_fn(logger=logger, name="Transferring file")
-@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
-    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:
-        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 OperationStatus.ERROR).
-        Exception: Any other errors during transfer (caught and returns OperationStatus.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 OperationStatus.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 OperationStatus.SUCCESS
-
-    except CacheOperationInterrupted as e:
-        logger.warning(f"File transfer interrupted: {e}")
-        return OperationStatus.ERROR
-
-    finally:
-        space_monitor.stop()
-
-
 @safe_execute("Clear failed", False)
 def clear_local_cache() -> bool:
     """Clear the local PyTorch compilation cache directory.
@@ -382,6 +300,8 @@ def clear_local_cache() -> bool:
     torch_dir = Path(TORCH_CACHE_DIR)
     if not torch_dir.exists():
         return True
+    import shutil
+
     shutil.rmtree(torch_dir)
     return True
 
@@ -409,27 +329,12 @@ def _cache_compression_worker(
     create_archive(torch_dir, local_temp, max_size_mb)
 
 
-@worker_process("Copy was cancelled before starting")
-def _cache_copy_worker(source_path_str: str, dest_path_str: str) -> None:
-    """Worker process that handles file copy to b10fs.
-
-    This function runs in a separate process to copy the compressed cache file
-    to the b10fs filesystem. It can be terminated externally if disk space becomes insufficient.
-
-    Args:
-        source_path_str: String path to the source file to copy.
-        dest_path_str: String path where the file will be copied.
-    """
-    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:
         if torch_dir.exists():
+            import shutil
+
             shutil.rmtree(torch_dir)
             logger.debug(f"Cleaned up torch directory: {torch_dir}")
     except Exception as e:

b10_transfer-0.1.7/src/b10_transfer/core.py

@@ -0,0 +1,131 @@
+"""Core file transfer operations for b10-transfer.
+
+This module provides generic file transfer functionality with space monitoring
+and error handling for b10fs operations.
+"""
+
+import logging
+import shutil
+from pathlib import Path
+
+from .utils import (
+    timed_fn,
+    safe_execute,
+    safe_unlink,
+)
+from .space_monitor import (
+    check_sufficient_disk_space,
+    CacheSpaceMonitor,
+    CacheOperationInterrupted,
+    run_monitored_process,
+    worker_process,
+)
+from .constants import (
+    B10FS_CACHE_DIR,
+    REQUIRED_B10FS_SPACE_MB,
+    MIN_LOCAL_SPACE_MB,
+    OperationStatus,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@timed_fn(logger=logger, name="Transferring file")
+@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
+    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:
+        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 OperationStatus.ERROR).
+        Exception: Any other errors during transfer (caught and returns OperationStatus.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 OperationStatus.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 OperationStatus.SUCCESS
+
+    except CacheOperationInterrupted as e:
+        logger.warning(f"File transfer interrupted: {e}")
+        return OperationStatus.ERROR
+
+    finally:
+        space_monitor.stop()
+
+
+@worker_process("Copy was cancelled before starting")
+def _cache_copy_worker(source_path_str: str, dest_path_str: str) -> None:
+    """Worker process that handles file copy operations.
+
+    This function runs in a separate process to copy files between locations.
+    It can be terminated externally if disk space becomes insufficient.
+
+    Args:
+        source_path_str: String path to the source file to copy.
+        dest_path_str: String path where the file will be copied.
+    """
+    source_path = Path(source_path_str)
+    dest_path = Path(dest_path_str)
+
+    shutil.copy2(source_path, dest_path)

{b10_transfer-0.1.5 → b10_transfer-0.1.7}/src/b10_transfer/environment.py

@@ -28,15 +28,14 @@ def get_cache_filename() -> str:
     """Get the cache filename prefix for the current environment.
 
     This function generates a cache filename prefix that includes the
-    environment key and hostname to ensure cache files are environment-specific
+    environment key to ensure cache files are environment-specific
     and unique per machine.
 
     Returns:
-        str: Cache filename prefix in format "cache_{environment_key}.{hostname}".
+        str: Cache filename prefix in format "cache_{environment_key}".
     """
     env_key = get_environment_key()
-    hostname = os.uname().nodename or os.getenv("HOSTNAME", UNKNOWN_HOSTNAME)
-    return f"cache_{env_key}.{hostname}"
+    return f"cache_{env_key}"
 
 
 def get_environment_key() -> str: