b10-transfer 0.0.1__py3-none-any.whl

b10_transfer/__init__.py

@@ -0,0 +1,51 @@
+"""B10 Transfer - Lock-free PyTorch compilation cache for Baseten."""
+
+from .core import transfer
+from .torch_cache import load_compile_cache, save_compile_cache, clear_local_cache
+from .async_transfers import (
+    start_transfer_async,
+    get_transfer_status,
+    is_transfer_complete,
+    wait_for_completion,
+    cancel_transfer,
+    list_active_transfers,
+    TransferProgress,
+)
+from .async_torch_cache import (
+    load_compile_cache_async,
+    save_compile_cache_async,
+)
+from .utils import CacheError, CacheValidationError
+from .space_monitor import CacheOperationInterrupted
+from .info import get_cache_info, list_available_caches
+from .constants import SaveStatus, LoadStatus, TransferStatus, AsyncTransferStatus
+
+# Version
+__version__ = "0.0.1"
+
+__all__ = [
+    "CacheError",
+    "CacheValidationError",
+    "CacheOperationInterrupted",
+    "SaveStatus",
+    "LoadStatus",
+    "TransferStatus",
+    "AsyncTransferStatus",
+    "transfer",
+    "load_compile_cache",
+    "save_compile_cache",
+    "clear_local_cache",
+    "get_cache_info",
+    "list_available_caches",
+    # Generic async operations
+    "start_transfer_async",
+    "get_transfer_status",
+    "is_transfer_complete",
+    "wait_for_completion",
+    "cancel_transfer",
+    "list_active_transfers",
+    "TransferProgress",
+    # Torch-specific async operations
+    "load_compile_cache_async",
+    "save_compile_cache_async",
+]

b10_transfer/archive.py

@@ -0,0 +1,175 @@
+import os
+import logging
+import subprocess
+from pathlib import Path
+
+from .utils import timed_fn, safe_unlink, CacheValidationError, validate_path_security
+from .constants import MAX_CACHE_SIZE_MB
+
+logger = logging.getLogger(__name__)
+
+
+class ArchiveError(Exception):
+    """Archive operation failed."""
+
+    pass
+
+
+def get_file_size_mb(file_path: Path) -> float:
+    """Get the size of a file in megabytes.
+
+    Args:
+        file_path: Path to the file to measure.
+
+    Returns:
+        float: File size in megabytes, or 0.0 if file doesn't exist or
+               can't be accessed.
+
+    Raises:
+        No exceptions are raised; OSError is caught and returns 0.0.
+    """
+    try:
+        return file_path.stat().st_size / (1024 * 1024)
+    except OSError:
+        return 0.0
+
+
+def _compress_directory_to_tar(source_dir: Path, target_file: Path) -> None:
+    """Compress directory contents to a gzipped tar archive using system tar.
+
+    This function recursively compresses all files in the source directory
+    into a gzipped tar archive using the system tar command for better performance.
+
+    Args:
+        source_dir: Path to the directory to compress.
+        target_file: Path where the compressed archive will be created.
+
+    Raises:
+        subprocess.CalledProcessError: If tar command fails.
+        OSError: If source directory can't be read or target file can't be written.
+    """
+    # Use system tar command for better performance
+    # -czf: create, gzip, file
+    # -C: change to directory before archiving
+    cmd = ["tar", "-czf", str(target_file), "-C", str(source_dir), "."]
+
+    try:
+        subprocess.run(cmd, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        raise OSError(f"tar compression failed: {e.stderr}") from e
+
+
+@timed_fn(logger=logger, name="Creating archive")
+def create_archive(
+    source_dir: Path, target_file: Path, max_size_mb: int = MAX_CACHE_SIZE_MB
+) -> None:
+    """Create a compressed archive with path validation and size limits.
+
+    This function safely creates a gzipped tar archive from a source directory
+    with security validation and size constraints. It validates paths to prevent
+    directory traversal attacks and enforces maximum archive size limits.
+
+    Args:
+        source_dir: Path to the directory to archive. Must exist and be within
+                   allowed directories (/tmp/ or its parent).
+        target_file: Path where the archive will be created. Must be within
+                    allowed directories (/app or /cache).
+        max_size_mb: Maximum allowed archive size in megabytes. Defaults to MAX_CACHE_SIZE_MB.
+
+    Raises:
+        CacheValidationError: If paths are outside allowed directories.
+        ArchiveError: If source directory doesn't exist, archive creation fails,
+                     or archive exceeds size limit.
+    """
+    # Validate paths
+    validate_path_security(
+        str(source_dir),
+        ["/tmp/", str(source_dir.parent)],
+        f"Source directory {source_dir}",
+        CacheValidationError,
+    )
+    validate_path_security(
+        str(target_file),
+        ["/app", "/cache"],
+        f"Target file {target_file}",
+        CacheValidationError,
+    )
+
+    if not source_dir.exists():
+        raise ArchiveError(f"Source directory missing: {source_dir}")
+
+    target_file.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        _compress_directory_to_tar(source_dir, target_file)
+        size_mb = get_file_size_mb(target_file)
+
+        if size_mb > max_size_mb:
+            safe_unlink(
+                target_file, f"Failed to delete oversized archive {target_file}"
+            )
+            raise ArchiveError(f"Archive too large: {size_mb:.1f}MB > {max_size_mb}MB")
+
+    except Exception as e:
+        safe_unlink(target_file, f"Failed to cleanup failed archive {target_file}")
+        raise ArchiveError(f"Archive creation failed: {e}") from e
+
+
+@timed_fn(logger=logger, name="Extracting archive")
+def extract_archive(archive_file: Path, target_dir: Path) -> None:
+    """Extract a compressed archive with security validation.
+
+    This function safely extracts a gzipped tar archive to a target directory
+    with security checks to prevent directory traversal attacks. It validates
+    both the archive and target paths, and inspects archive contents for
+    malicious paths before extraction.
+
+    Args:
+        archive_file: Path to the archive file to extract. Must exist and be
+                     within allowed directories (/app or /cache).
+        target_dir: Path to the directory where files will be extracted. Must
+                   be within allowed directories (/tmp/ or its parent).
+
+    Raises:
+        CacheValidationError: If paths are outside allowed directories or if
+                            archive contains unsafe paths (absolute paths or
+                            paths with '..' components).
+        ArchiveError: If archive file doesn't exist or extraction fails.
+    """
+    # Validate paths
+    validate_path_security(
+        str(archive_file),
+        ["/app", "/cache"],
+        f"Archive file {archive_file}",
+        CacheValidationError,
+    )
+    validate_path_security(
+        str(target_dir),
+        ["/tmp/", str(target_dir.parent)],
+        f"Target directory {target_dir}",
+        CacheValidationError,
+    )
+
+    if not archive_file.exists():
+        raise ArchiveError(f"Archive missing: {archive_file}")
+
+    try:
+        target_dir.mkdir(parents=True, exist_ok=True)
+
+        # First, perform security check by listing archive contents
+        list_cmd = ["tar", "-tzf", str(archive_file)]
+        result = subprocess.run(list_cmd, check=True, capture_output=True, text=True)
+
+        # Security check on all paths in the archive
+        for path in result.stdout.strip().split("\n"):
+            if path and (os.path.isabs(path) or ".." in path):
+                raise CacheValidationError(f"Unsafe path in archive: {path}")
+
+        # Extract using system tar command for better performance
+        extract_cmd = ["tar", "-xzf", str(archive_file), "-C", str(target_dir)]
+        subprocess.run(extract_cmd, check=True, capture_output=True, text=True)
+
+    except subprocess.CalledProcessError as e:
+        raise ArchiveError(f"tar extraction failed: {e.stderr}") from e
+    except Exception as e:
+        raise ArchiveError(f"Extraction failed: {e}") from e

b10_transfer/async_torch_cache.py

@@ -0,0 +1,62 @@
+"""Torch-specific async cache operations using the generic async transfer system."""
+
+from pathlib import Path
+from typing import Optional, Callable
+
+from .async_transfers import start_transfer_async
+from .torch_cache import torch_cache_load_callback, torch_cache_save_callback
+from .environment import get_cache_filename
+from .constants import (
+    TORCH_CACHE_DIR,
+    B10FS_CACHE_DIR,
+    MAX_CACHE_SIZE_MB,
+    CACHE_FILE_EXTENSION,
+    CACHE_LATEST_SUFFIX,
+)
+
+
+def load_compile_cache_async(
+    progress_callback: Optional[Callable[[str], None]] = None,
+) -> str:
+    """Start async PyTorch compilation cache load operation."""
+    b10fs_dir = Path(B10FS_CACHE_DIR)
+    torch_dir = Path(TORCH_CACHE_DIR)
+
+    cache_filename = get_cache_filename()
+    cache_file = (
+        b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
+    )
+
+    return start_transfer_async(
+        source=cache_file,
+        dest=torch_dir,
+        callback=torch_cache_load_callback,
+        operation_name="torch_cache_load",
+        progress_callback=progress_callback,
+        monitor_local=True,
+        monitor_b10fs=False,  # No need to monitor b10fs for read operations
+    )
+
+
+def save_compile_cache_async(
+    progress_callback: Optional[Callable[[str], None]] = None,
+) -> str:
+    """Start async PyTorch compilation cache save operation."""
+    b10fs_dir = Path(B10FS_CACHE_DIR)
+    torch_dir = Path(TORCH_CACHE_DIR)
+
+    cache_filename = get_cache_filename()
+    final_file = (
+        b10fs_dir / f"{cache_filename}{CACHE_LATEST_SUFFIX}{CACHE_FILE_EXTENSION}"
+    )
+
+    return start_transfer_async(
+        source=torch_dir,
+        dest=final_file,
+        callback=torch_cache_save_callback,
+        operation_name="torch_cache_save",
+        progress_callback=progress_callback,
+        monitor_local=True,
+        monitor_b10fs=True,
+        max_size_mb=MAX_CACHE_SIZE_MB,
+    )

b10_transfer/async_transfers.py

@@ -0,0 +1,275 @@
+"""Generic async transfer operations with progress tracking."""
+
+import asyncio
+import logging
+import threading
+import time
+from pathlib import Path
+from typing import Optional, Dict, Any, Callable
+from dataclasses import dataclass
+from datetime import datetime, timedelta
+from concurrent.futures import ThreadPoolExecutor
+
+from .core import transfer
+from .constants import AsyncTransferStatus, TransferStatus
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TransferProgress:
+    operation_id: str
+    operation_name: str
+    status: AsyncTransferStatus
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    error_message: Optional[str] = None
+    progress_callback: Optional[Callable[[str], None]] = None
+
+    def update_status(
+        self, status: AsyncTransferStatus, error_message: Optional[str] = None
+    ):
+        self.status = status
+        if error_message:
+            self.error_message = error_message
+        if status == AsyncTransferStatus.IN_PROGRESS and self.started_at is None:
+            self.started_at = datetime.now()
+        elif status in [
+            AsyncTransferStatus.SUCCESS,
+            AsyncTransferStatus.ERROR,
+            AsyncTransferStatus.INTERRUPTED,
+            AsyncTransferStatus.CANCELLED,
+        ]:
+            self.completed_at = datetime.now()
+
+        # Notify callback if provided
+        if self.progress_callback:
+            try:
+                self.progress_callback(self.operation_id)
+            except Exception as e:
+                logger.warning(f"Progress callback failed for {self.operation_id}: {e}")
+
+
+class AsyncTransferManager:
+    def __init__(self):
+        self._transfers: Dict[str, TransferProgress] = {}
+        self._executor = ThreadPoolExecutor(
+            max_workers=2, thread_name_prefix="b10-async-transfer"
+        )
+        self._lock = threading.Lock()
+        self._operation_counter = 0
+
+    def _generate_operation_id(self, operation_name: str) -> str:
+        with self._lock:
+            self._operation_counter += 1
+            return f"{operation_name}_{self._operation_counter}_{int(datetime.now().timestamp())}"
+
+    def start_transfer_async(
+        self,
+        source: Path,
+        dest: Path,
+        callback: Callable,
+        operation_name: str,
+        progress_callback: Optional[Callable[[str], None]] = None,
+        monitor_local: bool = True,
+        monitor_b10fs: bool = True,
+        **callback_kwargs,
+    ) -> str:
+        operation_id = self._generate_operation_id(operation_name)
+
+        progress = TransferProgress(
+            operation_id=operation_id,
+            operation_name=operation_name,
+            status=AsyncTransferStatus.NOT_STARTED,
+            progress_callback=progress_callback,
+        )
+
+        with self._lock:
+            self._transfers[operation_id] = progress
+
+        # Submit the transfer operation to thread pool
+        future = self._executor.submit(
+            self._execute_transfer,
+            operation_id,
+            source,
+            dest,
+            callback,
+            monitor_local,
+            monitor_b10fs,
+            callback_kwargs,
+        )
+
+        logger.info(f"Started async transfer operation: {operation_id}")
+        return operation_id
+
+    def _execute_transfer(
+        self,
+        operation_id: str,
+        source: Path,
+        dest: Path,
+        callback: Callable,
+        monitor_local: bool,
+        monitor_b10fs: bool,
+        callback_kwargs: Dict[str, Any],
+    ) -> None:
+        progress = self._transfers.get(operation_id)
+        if not progress:
+            logger.error(f"Progress tracking lost for operation {operation_id}")
+            return
+
+        try:
+            progress.update_status(AsyncTransferStatus.IN_PROGRESS)
+            logger.info(f"Starting transfer for operation {operation_id}")
+
+            result = transfer(
+                source=source,
+                dest=dest,
+                callback=callback,
+                monitor_local=monitor_local,
+                monitor_b10fs=monitor_b10fs,
+                **callback_kwargs,
+            )
+
+            # Convert TransferStatus to AsyncTransferStatus
+            if result == TransferStatus.SUCCESS:
+                progress.update_status(AsyncTransferStatus.SUCCESS)
+                logger.info(f"Transfer completed successfully: {operation_id}")
+            elif result == TransferStatus.INTERRUPTED:
+                progress.update_status(
+                    AsyncTransferStatus.INTERRUPTED,
+                    "Transfer interrupted due to insufficient disk space",
+                )
+                logger.warning(f"Transfer interrupted: {operation_id}")
+            else:
+                progress.update_status(
+                    AsyncTransferStatus.ERROR, "Transfer operation failed"
+                )
+                logger.error(f"Transfer failed: {operation_id}")
+
+        except Exception as e:
+            progress.update_status(AsyncTransferStatus.ERROR, str(e))
+            logger.error(
+                f"Transfer operation {operation_id} failed with exception: {e}"
+            )
+
+    def get_transfer_status(self, operation_id: str) -> Optional[TransferProgress]:
+        with self._lock:
+            return self._transfers.get(operation_id)
+
+    def is_transfer_complete(self, operation_id: str) -> bool:
+        progress = self.get_transfer_status(operation_id)
+        if not progress:
+            return False
+
+        return progress.status in [
+            AsyncTransferStatus.SUCCESS,
+            AsyncTransferStatus.ERROR,
+            AsyncTransferStatus.INTERRUPTED,
+            AsyncTransferStatus.CANCELLED,
+        ]
+
+    def wait_for_completion(
+        self, operation_id: str, timeout: Optional[float] = None
+    ) -> bool:
+        start_time = datetime.now()
+
+        while not self.is_transfer_complete(operation_id):
+            if timeout and (datetime.now() - start_time).total_seconds() > timeout:
+                return False
+            time.sleep(0.1)  # Small delay to avoid busy waiting
+
+        return True
+
+    def cancel_transfer(self, operation_id: str) -> bool:
+        progress = self.get_transfer_status(operation_id)
+        if not progress:
+            return False
+
+        if progress.status == AsyncTransferStatus.IN_PROGRESS:
+            progress.update_status(AsyncTransferStatus.CANCELLED)
+            logger.info(f"Marked transfer operation as cancelled: {operation_id}")
+            return True
+
+        return False
+
+    def cleanup_completed_transfers(self, max_age_hours: int = 24) -> int:
+        cutoff_time = datetime.now() - timedelta(hours=max_age_hours)
+        cleaned_count = 0
+
+        with self._lock:
+            to_remove = []
+            for operation_id, progress in self._transfers.items():
+                if (
+                    progress.completed_at
+                    and progress.completed_at < cutoff_time
+                    and self.is_transfer_complete(operation_id)
+                ):
+                    to_remove.append(operation_id)
+
+            for operation_id in to_remove:
+                del self._transfers[operation_id]
+                cleaned_count += 1
+
+        if cleaned_count > 0:
+            logger.info(f"Cleaned up {cleaned_count} completed transfer records")
+
+        return cleaned_count
+
+    def list_active_transfers(self) -> Dict[str, TransferProgress]:
+        with self._lock:
+            return {
+                op_id: progress
+                for op_id, progress in self._transfers.items()
+                if not self.is_transfer_complete(op_id)
+            }
+
+    def shutdown(self) -> None:
+        logger.info("Shutting down async transfer manager...")
+        self._executor.shutdown(wait=True)
+
+
+# Global instance for easy access
+_transfer_manager = AsyncTransferManager()
+
+
+# Generic Public API functions
+def start_transfer_async(
+    source: Path,
+    dest: Path,
+    callback: Callable,
+    operation_name: str,
+    progress_callback: Optional[Callable[[str], None]] = None,
+    monitor_local: bool = True,
+    monitor_b10fs: bool = True,
+    **callback_kwargs,
+) -> str:
+    return _transfer_manager.start_transfer_async(
+        source=source,
+        dest=dest,
+        callback=callback,
+        operation_name=operation_name,
+        progress_callback=progress_callback,
+        monitor_local=monitor_local,
+        monitor_b10fs=monitor_b10fs,
+        **callback_kwargs,
+    )
+
+
+def get_transfer_status(operation_id: str) -> Optional[TransferProgress]:
+    return _transfer_manager.get_transfer_status(operation_id)
+
+
+def is_transfer_complete(operation_id: str) -> bool:
+    return _transfer_manager.is_transfer_complete(operation_id)
+
+
+def wait_for_completion(operation_id: str, timeout: Optional[float] = None) -> bool:
+    return _transfer_manager.wait_for_completion(operation_id, timeout)
+
+
+def cancel_transfer(operation_id: str) -> bool:
+    return _transfer_manager.cancel_transfer(operation_id)
+
+
+def list_active_transfers() -> Dict[str, TransferProgress]:
+    return _transfer_manager.list_active_transfers()