b10-transfer 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

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/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/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.

{b10_transfer-0.1.2.dist-info → b10_transfer-0.1.3.dist-info}/METADATA

@@ -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.dist-info → b10_transfer-0.1.3.dist-info}/RECORD

@@ -1,12 +1,12 @@
-b10_transfer/__init__.py,sha256=kZXd7GHMH7PFzr4aXs19MmuFzwEvtwzh0rwyK4jHgHo,641
+b10_transfer/__init__.py,sha256=qaqigpp3gmzR4-Cr9aH1ilaYQgp9UB_DzzpsmwX_D-o,641
 b10_transfer/archive.py,sha256=GKb0mi0-YeM7ch4FLAoOLHXw0T6LkRerYad2N2y9TYM,6400
 b10_transfer/cleanup.py,sha256=3RnqWNGMCcko5GQdq1Gr9VPpGzAF5J6x7xjIH9SNZ78,6226
-b10_transfer/constants.py,sha256=EmWCh9AOamCZL3KkSU6YJO6KBkh93OmAIUeEfaZxHL0,4321
-b10_transfer/core.py,sha256=tKA1gWDEpcb_Xfr6njScAGWCxIT3htlyVh9VLg67YMg,15445
+b10_transfer/constants.py,sha256=qCViKTyfHTLpiFVF2SwsbHp2IMz3kg3syxJfgRAq2dc,4446
+b10_transfer/core.py,sha256=UVjzcqCqDMqpXqiXzdC6c8nJ_2tM35zlDl-Jp1Gvn20,18657
 b10_transfer/environment.py,sha256=aC0biEMQrtHk0ke_3epdcq1X9J5fPmPpBVt0fH7XF2Y,5625
 b10_transfer/info.py,sha256=I3iOuImZ5r6DMJTDeBtVvzlSn6IuyPJbLJYUO_OF0ks,6299
 b10_transfer/space_monitor.py,sha256=C_CKDH43bNsWdq60WStSZ3c_nQkWvScQmqU_SYHesew,10531
 b10_transfer/utils.py,sha256=Stee0DFK-8MRRYNIocqaK64cJvfs4jPW3Mpx7zkWV6Y,11932
-b10_transfer-0.1.2.dist-info/METADATA,sha256=P0Yf1VzkqFV4eco8x1yAjgUiahlmjgJGTAee0NNhh6o,4108
-b10_transfer-0.1.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-b10_transfer-0.1.2.dist-info/RECORD,,
+b10_transfer-0.1.3.dist-info/METADATA,sha256=dNTL7He36x0j9nZpbf5Rhaa4FYHeuNaim2zyXCKq_lE,4108
+b10_transfer-0.1.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+b10_transfer-0.1.3.dist-info/RECORD,,