b10-transfer 0.0.1__py3-none-any.whl

b10_transfer/cleanup.py

@@ -0,0 +1,179 @@
+"""Cooperative cleanup utilities for b10-tcache.
+
+This module provides cooperative cleanup functionality where each pod/replica
+helps maintain the health of shared resources in b10fs by removing stale
+lock files and incomplete cache files.
+"""
+
+import fnmatch
+import time
+import logging
+from pathlib import Path
+from typing import List, Tuple
+
+from .constants import (
+    B10FS_CACHE_DIR,
+    CACHE_INCOMPLETE_SUFFIX,
+    CLEANUP_LOCK_TIMEOUT_SECONDS,
+    CLEANUP_INCOMPLETE_TIMEOUT_SECONDS,
+)
+from .utils import safe_execute, safe_unlink
+
+logger = logging.getLogger(__name__)
+
+
+@safe_execute("Failed to find stale files", [])
+def _find_stale_files(
+    directory: Path, pattern: str, timeout_seconds: int
+) -> List[Path]:
+    """Find files matching pattern that are older than timeout_seconds.
+
+    Args:
+        directory: Directory to search in
+        pattern: Glob pattern to match files (only searches current directory, not subdirs)
+        timeout_seconds: Age threshold in seconds
+
+    Returns:
+        List of Path objects for stale files
+    """
+    if not directory.exists():
+        return []
+
+    current_time = time.time()
+    stale_files = []
+
+    # Use iterdir() + fnmatch for explicit file-only matching in current directory
+
+    for file_path in directory.iterdir():
+        # Skip directories - we only want files
+        if not file_path.is_file():
+            logger.warning(
+                f"Found non-file in b10fs cache directory: {file_path}, skipping consideration for deletion in cleanup phase."
+            )
+            continue
+
+        # Check if filename matches pattern for the type of file we're looking for
+        if not fnmatch.fnmatch(file_path.name, pattern):
+            logger.warning(
+                f"Found non-matching file in b10fs cache directory: {file_path}, skipping consideration for deletion in cleanup phase."
+            )
+            continue
+
+        try:
+            file_age = current_time - file_path.stat().st_mtime
+            if file_age > timeout_seconds:
+                stale_files.append(file_path)
+        except OSError:
+            # File might have been deleted already
+            continue
+
+    return stale_files
+
+
+@safe_execute("Failed to cleanup files", 0)
+def _cleanup_files(files: List[Path], file_type: str) -> int:
+    """Clean up a list of files and return count of successfully cleaned files.
+
+    Args:
+        files: List of file paths to clean up
+        file_type: Description of file type for logging (e.g., "lock", "incomplete")
+
+    Returns:
+        Number of files successfully cleaned up
+    """
+    cleaned_count = 0
+
+    for file_path in files:
+        try:
+            file_age = time.time() - file_path.stat().st_mtime
+            safe_unlink(
+                file_path, f"Failed to clean stale {file_type} file: {file_path}"
+            )
+            cleaned_count += 1
+            logger.debug(
+                f"Cleaned stale {file_type} file: {file_path.name} (age: {file_age:.1f}s)"
+            )
+        except OSError:
+            # File might have been deleted by another pod
+            continue
+
+    return cleaned_count
+
+
+@safe_execute("Cooperative cleanup failed", None)
+def cooperative_cleanup_b10fs() -> None:
+    """Clean up stale shared resources in b10fs cooperatively.
+
+    Each pod/replica calls this function to help maintain system health by
+    removing files that are likely orphaned due to pod crashes or failures.
+
+    Removes:
+    - Lock files older than CLEANUP_LOCK_TIMEOUT_SECONDS (*.lock)
+    - Incomplete cache files older than CLEANUP_INCOMPLETE_TIMEOUT_SECONDS (*.incomplete*)
+
+    Does NOT remove:
+    - Final cache files (*.latest.tar.gz) - these are the actual cached results
+    - Files newer than the configured thresholds (may be from active operations)
+
+    This function is safe to run concurrently from multiple pods as file
+    deletion operations are atomic and missing files are handled gracefully.
+    """
+    b10fs_dir = Path(B10FS_CACHE_DIR)
+    if not b10fs_dir.exists():
+        logger.debug("b10fs cache directory doesn't exist, skipping cleanup")
+        return
+
+    # Find and clean stale lock files
+    stale_locks = _find_stale_files(b10fs_dir, "*.lock", CLEANUP_LOCK_TIMEOUT_SECONDS)
+    cleaned_locks = _cleanup_files(stale_locks, "lock")
+
+    # Find and clean stale incomplete cache files
+    incomplete_pattern = f"*{CACHE_INCOMPLETE_SUFFIX}*"
+    stale_incomplete = _find_stale_files(
+        b10fs_dir, incomplete_pattern, CLEANUP_INCOMPLETE_TIMEOUT_SECONDS
+    )
+    cleaned_incomplete = _cleanup_files(stale_incomplete, "incomplete cache")
+
+    # Log summary
+    total_cleaned = cleaned_locks + cleaned_incomplete
+    if total_cleaned > 0:
+        logger.info(
+            f"Cooperative cleanup: removed {cleaned_locks} stale locks, "
+            f"{cleaned_incomplete} incomplete files"
+        )
+    else:
+        logger.debug("Cooperative cleanup: no stale files found")
+
+
+def get_cleanup_info() -> dict:
+    """Get information about cleanup configuration and current state.
+
+    Returns:
+        dict: Dictionary containing cleanup configuration and statistics:
+            - lock_timeout_seconds: Current lock file cleanup threshold
+            - incomplete_timeout_seconds: Current incomplete file cleanup threshold
+            - b10fs_cache_dir: Path to b10fs cache directory
+            - b10fs_exists: Whether b10fs cache directory exists
+            - stale_locks_count: Number of lock files that would be cleaned
+            - stale_incomplete_count: Number of incomplete files that would be cleaned
+    """
+    b10fs_dir = Path(B10FS_CACHE_DIR)
+
+    info = {
+        "lock_timeout_seconds": CLEANUP_LOCK_TIMEOUT_SECONDS,
+        "incomplete_timeout_seconds": CLEANUP_INCOMPLETE_TIMEOUT_SECONDS,
+        "b10fs_cache_dir": str(b10fs_dir),
+        "b10fs_exists": b10fs_dir.exists(),
+        "stale_locks_count": len(
+            _find_stale_files(b10fs_dir, "*.lock", CLEANUP_LOCK_TIMEOUT_SECONDS)
+        ),
+        "stale_incomplete_count": len(
+            _find_stale_files(
+                b10fs_dir,
+                f"*{CACHE_INCOMPLETE_SUFFIX}*",
+                CLEANUP_INCOMPLETE_TIMEOUT_SECONDS,
+            )
+        ),
+    }
+
+    return info

b10_transfer/constants.py

@@ -0,0 +1,149 @@
+"""Configuration constants for b10-tcache.
+
+This module defines configuration constants for the PyTorch compilation cache system.
+Some values can be overridden by environment variables, but security caps are enforced
+to prevent malicious or accidental misuse in production environments.
+"""
+
+import os
+from enum import Enum, auto
+
+# Import helper functions from utils to avoid duplication
+from .utils import (
+    get_current_username,
+    validate_path_security,
+    validate_boolean_env,
+    apply_cap,
+)
+
+# Cache directories with security validation
+
+# Validate TORCH_CACHE_DIR - allow /tmp and /cache paths
+# TORCHINDUCTOR_CACHE_DIR is what torch uses by default. If it is not set, we use a different value.
+_torch_cache_dir = os.getenv(
+    "TORCHINDUCTOR_CACHE_DIR", f"/tmp/torchinductor_{get_current_username()}"
+)
+TORCH_CACHE_DIR = validate_path_security(
+    _torch_cache_dir, ["/tmp/", "/cache/"], "TORCHINDUCTOR_CACHE_DIR"
+)
+
+# B10FS cache directory validation
+_REQUIRED_TORCH_CACHE_DIR_PREFIX = "/cache/model"
+_b10fs_cache_dir = os.getenv(
+    "B10FS_CACHE_DIR", f"{_REQUIRED_TORCH_CACHE_DIR_PREFIX}/compile_cache"
+)
+B10FS_CACHE_DIR = validate_path_security(
+    _b10fs_cache_dir, [_REQUIRED_TORCH_CACHE_DIR_PREFIX], "B10FS_CACHE_DIR"
+)
+
+# Validate LOCAL_WORK_DIR - allow /app, /tmp, and /cache paths.
+# This is like a "scratch" directory where you can do work (like compression/archival for example)
+_local_work_dir = os.getenv("LOCAL_WORK_DIR", "/app")
+LOCAL_WORK_DIR = validate_path_security(
+    _local_work_dir, ["/app/", "/tmp/", "/cache/"], "LOCAL_WORK_DIR"
+)
+
+# Security caps to prevent resource exhaustion
+_MAX_CACHE_SIZE_CAP_MB = 1 * 1024  # 1GB hard limit per cache archive
+_MAX_CONCURRENT_SAVES_CAP = 100  # Maximum concurrent save operations (only used as estimate for b10fs space requirements/thresholding)
+
+
+# Cache limits (capped for security)
+_user_max_cache_size = int(os.getenv("MAX_CACHE_SIZE_MB", "1024"))
+MAX_CACHE_SIZE_MB = apply_cap(
+    _user_max_cache_size, _MAX_CACHE_SIZE_CAP_MB, "MAX_CACHE_SIZE_MB"
+)
+
+_user_max_concurrent_saves = int(os.getenv("MAX_CONCURRENT_SAVES", "50"))
+MAX_CONCURRENT_SAVES = apply_cap(
+    _user_max_concurrent_saves, _MAX_CONCURRENT_SAVES_CAP, "MAX_CONCURRENT_SAVES"
+)
+
+# Space requirements
+MIN_LOCAL_SPACE_MB = 50 * 1024  # 50GB minimum space on local machine
+REQUIRED_B10FS_SPACE_MB = max(MAX_CONCURRENT_SAVES * MAX_CACHE_SIZE_MB, 100_000)
+
+# B10FS configuration
+# The default is "0" (disabled) to prevent accidental enabling.
+# But this does limit the ability to enable b10fs for debugging purposes.
+# Probably should use B10FS_ENABLED instead for that.
+_baseten_fs_enabled = os.getenv("BASETEN_FS_ENABLED", "0")
+BASETEN_FS_ENABLED = validate_boolean_env(_baseten_fs_enabled, "BASETEN_FS_ENABLED")
+
+# File naming patterns
+CACHE_FILE_EXTENSION = ".tar.gz"
+CACHE_LATEST_SUFFIX = ".latest"
+CACHE_INCOMPLETE_SUFFIX = ".incomplete"
+CACHE_PREFIX = "cache_"
+
+
+# Space monitoring settings
+SPACE_MONITOR_CHECK_INTERVAL_SECONDS = (
+    0.5  # How often to check disk space during operations
+)
+
+# Cooperative cleanup settings
+# Cache operations (load/save) should complete within ~15 seconds under normal conditions
+_LOCK_TIMEOUT_CAP_SECONDS = 3600  # 1 hour hard limit
+_INCOMPLETE_TIMEOUT_CAP_SECONDS = 7200  # 2 hours hard limit
+
+# Lock file cleanup timeout (default: 2x expected operation time)
+_user_lock_timeout = int(
+    os.getenv("CLEANUP_LOCK_TIMEOUT_SECONDS", "30")
+)  # 30 seconds default
+CLEANUP_LOCK_TIMEOUT_SECONDS = apply_cap(
+    _user_lock_timeout, _LOCK_TIMEOUT_CAP_SECONDS, "CLEANUP_LOCK_TIMEOUT_SECONDS"
+)
+
+# Incomplete file cleanup timeout (default: 3x expected operation time)
+_user_incomplete_timeout = int(
+    os.getenv("CLEANUP_INCOMPLETE_TIMEOUT_SECONDS", "60")
+)  # 1 minute default
+CLEANUP_INCOMPLETE_TIMEOUT_SECONDS = apply_cap(
+    _user_incomplete_timeout,
+    _INCOMPLETE_TIMEOUT_CAP_SECONDS,
+    "CLEANUP_INCOMPLETE_TIMEOUT_SECONDS",
+)
+
+
+# Worker process result status enum
+class WorkerStatus(Enum):
+    """Status values for worker process results."""
+
+    SUCCESS = auto()
+    ERROR = auto()
+    CANCELLED = auto()
+
+
+class LoadStatus(Enum):
+    """Status values for cache loading operations."""
+
+    SUCCESS = auto()
+    ERROR = auto()
+    DOES_NOT_EXIST = auto()
+    SKIPPED = auto()
+
+
+class SaveStatus(Enum):
+    """Status values for cache saving operations."""
+
+    SUCCESS = auto()
+    ERROR = auto()
+    SKIPPED = auto()
+
+
+class TransferStatus(Enum):
+    """Status values for generic transfer operations."""
+
+    SUCCESS = auto()
+    ERROR = auto()
+    INTERRUPTED = auto()
+
+
+class AsyncTransferStatus(Enum):
+    NOT_STARTED = auto()
+    IN_PROGRESS = auto()
+    SUCCESS = auto()
+    ERROR = auto()
+    INTERRUPTED = auto()
+    CANCELLED = auto()

b10_transfer/core.py

@@ -0,0 +1,160 @@
+import logging
+from pathlib import Path
+
+from .cleanup import cooperative_cleanup_b10fs
+from .utils import (
+    timed_fn,
+    safe_execute,
+    cache_operation,
+)
+from .space_monitor import (
+    check_sufficient_disk_space,
+    CacheSpaceMonitor,
+    CacheOperationInterrupted,
+    run_monitored_process,
+)
+from .constants import (
+    B10FS_CACHE_DIR,
+    LOCAL_WORK_DIR,
+    REQUIRED_B10FS_SPACE_MB,
+    MIN_LOCAL_SPACE_MB,
+    TransferStatus,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@timed_fn(logger=logger, name="Generic transfer operation")
+@safe_execute("Transfer failed", TransferStatus.ERROR)
+def transfer(
+    source: Path,
+    dest: Path,
+    callback: callable,
+    *callback_args,
+    monitor_local: bool = True,
+    monitor_b10fs: bool = True,
+    **callback_kwargs,
+) -> TransferStatus:
+    """Generic transfer function with space monitoring and atomic operations.
+
+    The actual transfer logic is provided via callback.
+
+    The function handles:
+    - Cooperative cleanup of stale shared resources
+    - Space monitoring during operations (optional for local and b10fs)
+    - Atomic operations using temp files and rename
+    - Automatic cleanup on interruption or failure
+    - Lock management for b10fs operations
+
+    Args:
+        source: Source path for the transfer operation
+        dest: Destination path for the transfer operation
+        callback: Function to perform the actual transfer work
+        *callback_args: Positional arguments to pass to callback
+        monitor_local: Whether to monitor local disk space (default: True)
+        monitor_b10fs: Whether to monitor b10fs disk space (default: True)
+        **callback_kwargs: Keyword arguments to pass to callback
+
+    Returns:
+        TransferStatus:
+              TransferStatus.SUCCESS if transfer completed successfully
+              TransferStatus.ERROR if transfer failed
+              TransferStatus.INTERRUPTED if transfer was interrupted due to insufficient disk space
+
+    Raises:
+        CacheValidationError: If b10fs is not enabled (caught and returns TransferStatus.ERROR).
+        CacheOperationInterrupted: If operations interrupted due to insufficient
+                                  disk space (caught and returns TransferStatus.INTERRUPTED).
+        Exception: Any other errors during transfer (caught and returns TransferStatus.ERROR).
+    """
+    with cache_operation("Transfer"):
+        # Cooperative cleanup of stale shared resources
+        cooperative_cleanup_b10fs()
+
+        b10fs_dir = Path(B10FS_CACHE_DIR)
+        work_dir = Path(LOCAL_WORK_DIR)
+
+        # Determine which paths to monitor based on source/dest
+        local_path = None
+        b10fs_path = None
+
+        if str(source).startswith(str(b10fs_dir)) or str(dest).startswith(
+            str(b10fs_dir)
+        ):
+            b10fs_path = b10fs_dir
+
+        if (
+            str(source).startswith(str(work_dir))
+            or str(dest).startswith(str(work_dir))
+            or not str(source).startswith(str(b10fs_dir))
+            or not str(dest).startswith(str(b10fs_dir))
+        ):
+            local_path = work_dir
+
+        # Initial disk space checks
+        if monitor_local and local_path:
+            check_sufficient_disk_space(
+                local_path, MIN_LOCAL_SPACE_MB, "local transfer operations"
+            )
+            logger.debug(
+                f"Initial local space check passed: {MIN_LOCAL_SPACE_MB:.1f}MB required"
+            )
+
+        if monitor_b10fs and b10fs_path:
+            check_sufficient_disk_space(
+                b10fs_path, REQUIRED_B10FS_SPACE_MB, "b10fs transfer operations"
+            )
+            logger.debug(
+                f"Initial b10fs space check passed: {REQUIRED_B10FS_SPACE_MB:.1f}MB required"
+            )
+
+        # Determine primary space monitor (prioritize b10fs if both are monitored)
+        primary_monitor = None
+        if monitor_b10fs and b10fs_path:
+            primary_monitor = CacheSpaceMonitor(REQUIRED_B10FS_SPACE_MB, b10fs_path)
+        elif monitor_local and local_path:
+            primary_monitor = CacheSpaceMonitor(MIN_LOCAL_SPACE_MB, local_path)
+
+        if primary_monitor is None:
+            # No monitoring requested, execute callback directly
+            logger.info(f"Starting transfer (no monitoring): {source} -> {dest}")
+            callback(source, dest, *callback_args, **callback_kwargs)
+            logger.info("Transfer complete")
+            return TransferStatus.SUCCESS
+
+        # Start the primary space monitor
+        primary_monitor.start()
+
+        try:
+            # Execute the callback using monitored process for continuous space monitoring
+            logger.info(f"Starting monitored transfer: {source} -> {dest}")
+
+            # Try direct callback with run_monitored_process first
+            try:
+                run_monitored_process(
+                    callback,
+                    (source, dest, *callback_args),
+                    primary_monitor,
+                    "transfer callback",
+                )
+                logger.info("Transfer complete (monitored)")
+                return TransferStatus.SUCCESS
+
+            except (TypeError, AttributeError, ImportError, OSError) as e:
+                # Callback not pickleable or other serialization issue
+                logger.warning(
+                    f"Callback not suitable for process isolation, running without monitoring: {e}"
+                )
+
+                # Fallback to direct execution without process isolation
+                callback(source, dest, *callback_args, **callback_kwargs)
+                logger.info("Transfer complete (unmonitored)")
+                return TransferStatus.SUCCESS
+
+        except CacheOperationInterrupted as e:
+            logger.warning(f"Transfer interrupted: {e}")
+            return TransferStatus.INTERRUPTED
+
+        finally:
+            # Stop space monitor
+            primary_monitor.stop()

b10_transfer/environment.py

@@ -0,0 +1,134 @@
+"""Environment detection utilities for GPU cache management.
+
+This module provides functions to generate unique environment keys based on
+GPU hardware and driver information for cache compatibility.
+"""
+
+import hashlib
+import json
+import logging
+import os
+
+# Optional imports - may not be available in all environments
+try:
+    import torch
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    torch = None
+    TORCH_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+KEY_LENGTH = 16
+UNKNOWN_HOSTNAME = "unknown-host"
+
+
+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
+    and unique per machine.
+
+    Returns:
+        str: Cache filename prefix in format "cache_{environment_key}.{hostname}".
+    """
+    env_key = get_environment_key()
+    hostname = os.uname().nodename or os.getenv("HOSTNAME", UNKNOWN_HOSTNAME)
+    return f"cache_{env_key}.{hostname}"
+
+
+def get_environment_key() -> str:
+    """Generate unique environment key based on PyTorch/CUDA/GPU configuration.
+
+    This function creates a deterministic hash key based only on node-specific
+    hardware and driver information to ensure cache compatibility across
+    different environments with identical GPU configurations.
+
+    Returns:
+        str: A 16-character hex hash uniquely identifying the environment.
+
+    Raises:
+        RuntimeError: If PyTorch/CUDA are unavailable or environment key
+                     generation fails for any reason.
+
+    Note:
+        Includes all GPU properties that affect Triton kernel generation.
+        References from PyTorch repository:
+        - Device name: GPU model identification (codecache.py:199)
+        - CUDA version: Driver compatibility (codecache.py:200)
+
+        These next four are bit more embedded in the codebase and not obviously used in the torchinductor_root cache check.
+        They are commented out because:
+        1) They are not explicitly used in the torchinductor_root cache check.
+        2) It's not clear but likely that any violation of these properties will cause local re-compilation when the torch guards activate, not full recompilation.
+        3) We don't want to over-estimate the number of unique environments since that'll cause more cache misses overall.
+        We can add them back if we need to.
+
+        - Compute capability: Available GPU instructions/features (scheduler.py:4286, triton_heuristics.py:480)
+        - Multi-processor count: Affects occupancy and grid sizing (choices.py:210, triton_heuristics.py:539)
+        - Warp size: Thread grouping (triton_heuristics.py:487, triton.py:2763)
+        - Register limits: Affects kernel optimization strategies (triton_heuristics.py:518,536)
+
+        We're also not including the torch and triton versions in the hash, despite the torch compilation cache dependent on these two things.
+        This is because we are saving the cache to the `/cache/model` directory, which is already deployment-specific where the torch/triton versions are constant.
+    """
+    try:
+        _validate_cuda_environment()
+
+        device_properties = torch.cuda.get_device_properties(
+            torch.cuda.current_device()
+        )
+        node_data = _extract_gpu_properties(device_properties, torch.version.cuda)
+
+        node_json = json.dumps(node_data, sort_keys=True)
+        return hashlib.sha256(node_json.encode("utf-8")).hexdigest()[:KEY_LENGTH]
+
+    except (ImportError, RuntimeError, AssertionError) as e:
+        logger.error(f"GPU environment unavailable: {e}")
+        raise RuntimeError(f"Cannot generate environment key: {e}") from e
+    except Exception as e:
+        logger.error(f"Unexpected error during environment key generation: {e}")
+        raise RuntimeError(f"Environment key generation failed: {e}") from e
+
+
+def _validate_cuda_environment() -> None:
+    """Validate that PyTorch and CUDA are available and properly configured.
+
+    Raises:
+        ImportError: If PyTorch is not available
+        RuntimeError: If CUDA is not available or version is missing
+    """
+    if not TORCH_AVAILABLE:
+        raise ImportError("PyTorch not available")
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA must be available - AMD/HIP not supported")
+
+    if torch.version.cuda is None:
+        raise RuntimeError("CUDA version must be available")
+
+
+def _extract_gpu_properties(
+    device_properties: any, cuda_version: str
+) -> dict[str, any]:
+    """Extract relevant GPU properties for environment key generation.
+
+    Args:
+        device_properties: CUDA device properties object
+        cuda_version: CUDA version string
+        SEE docstring of get_environment_key() for more details and why certain properties are excluded.
+
+    Returns:
+        Dict containing GPU properties that affect kernel generation
+    """
+    return {
+        "device_name": device_properties.name,  # GPU model
+        "cuda_version": cuda_version,  # Driver version
+        # "compute_capability": (device_properties.major, device_properties.minor),  # GPU features
+        # "multi_processor_count": device_properties.multi_processor_count,  # SM count for occupancy
+        # "warp_size": device_properties.warp_size,  # Thread grouping size
+        # "regs_per_multiprocessor": getattr(device_properties, "regs_per_multiprocessor", None),  # Register limits
+        # "max_threads_per_multi_processor": getattr(device_properties, "max_threads_per_multi_processor", None),  # Thread limits
+    }