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

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

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

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

{b10_transfer-0.1.2 → b10_transfer-0.1.3}/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.3}/src/b10_transfer/core.py

@@ -35,6 +35,7 @@ from .constants import (
     CACHE_INCOMPLETE_SUFFIX,
     LoadStatus,
     SaveStatus,
+    TransferStatus,
 )
 
 logger = logging.getLogger(__name__)
@@ -286,6 +287,90 @@ def save_compile_cache() -> SaveStatus:
             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)
 def clear_local_cache() -> bool:
     """Clear the local PyTorch compilation cache directory.