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

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

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

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.1.7"
+version = "0.2.0"
 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>"]
@@ -28,6 +28,9 @@ classifiers = [
 ]
 packages = [{include = "b10_transfer", from = "src"}]
 
+[tool.poetry.scripts]
+b10-transfer = "b10_transfer.cache_cli:main"
+
 [tool.poetry.dependencies]
 python = "^3.9"
 torch = ">=2.0.0"

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/__init__.py

@@ -6,9 +6,10 @@ from .utils import CacheError, CacheValidationError
 from .space_monitor import CacheOperationInterrupted
 from .info import get_cache_info, list_available_caches
 from .constants import OperationStatus
+from .logging_utils import get_b10_logger
 
 # Version
-__version__ = "0.1.7"
+__version__ = "0.2.0"
 
 __all__ = [
     "CacheError",
@@ -21,4 +22,5 @@ __all__ = [
     "transfer",
     "get_cache_info",
     "list_available_caches",
+    "get_b10_logger",
 ]

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/archive.py

@@ -5,8 +5,9 @@ from pathlib import Path
 
 from .utils import timed_fn, safe_unlink, CacheValidationError, validate_path_security
 from .constants import MAX_CACHE_SIZE_MB
+from .logging_utils import get_b10_logger
 
-logger = logging.getLogger(__name__)
+logger = get_b10_logger(__name__)
 
 
 class ArchiveError(Exception):

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/cache.py

@@ -8,6 +8,8 @@ import logging
 import tempfile
 from pathlib import Path
 
+from .logging_utils import get_b10_logger
+
 from .environment import get_cache_filename
 from .cleanup import cooperative_cleanup_b10fs
 from .utils import (
@@ -38,7 +40,7 @@ from .constants import (
 )
 from .core import transfer
 
-logger = logging.getLogger(__name__)
+logger = get_b10_logger(__name__)
 
 
 """
@@ -94,7 +96,7 @@ def _run_with_space_monitoring(
 
     try:
         logger.info(
-            f"Starting {operation_name}: {' -> '.join(str(arg) for arg in worker_args[:2])}"
+            f"[MONITORING] Starting {operation_name} with space monitoring: {' -> '.join(str(arg) for arg in worker_args[:2])}"
         )
         run_monitored_process(
             worker_func,
@@ -114,14 +116,29 @@ def _transfer_with_b10fs_lock(
 
     @critical_section_b10fs_file_lock(lock_type)
     def _locked_transfer():
+        # Get file size for logging
+        source_path = Path(source)
+        source_size_mb = (
+            source_path.stat().st_size / (1024 * 1024) if source_path.exists() else 0
+        )
+        logger.info(
+            f"[TRANSFER] Starting locked transfer: {source} -> {dest} (size: {source_size_mb:.2f} MB, lock: {lock_type})"
+        )
+
         result = transfer(source, dest)
         if result != OperationStatus.SUCCESS:
+            logger.error(f"[TRANSFER] Transfer failed with status: {result}")
             if cleanup_on_failure:
+                logger.info(
+                    f"[TRANSFER] Cleaning up failed transfer destination: {dest}"
+                )
                 safe_unlink(
                     Path(dest), f"Failed to cleanup after failed transfer {dest}"
                 )
             raise Exception(f"Failed to transfer {source} -> {dest}")
 
+        logger.info(f"[TRANSFER] Transfer completed successfully: {source} -> {dest}")
+
     _locked_transfer()
 
 
@@ -156,15 +173,20 @@ def load_compile_cache() -> OperationStatus:
 
         cache_filename = get_cache_filename()
         final_file, _ = _get_cache_file_paths(cache_filename, b10fs_dir)
-        logger.debug(f"Looking for cache file: {final_file}")
+        logger.info(f"[LOADING] Searching for cache file: {final_file}")
 
         if not final_file.exists():
-            logger.info("No cache file found in b10fs")
+            logger.info(f"[LOADING] No cache file found in b10fs at: {final_file}")
             return OperationStatus.DOES_NOT_EXIST
 
         # Skip if already loaded
         if torch_dir.exists() and any(torch_dir.iterdir()):
-            logger.info("Torch cache already loaded, skipping extraction")
+            size_mb = sum(
+                f.stat().st_size for f in torch_dir.rglob("*") if f.is_file()
+            ) / (1024 * 1024)
+            logger.info(
+                f"[LOADING] Torch cache already exists at {torch_dir}, skipping extraction (size: {size_mb:.2f} MB)"
+            )
             return OperationStatus.SKIPPED
 
         # Create temp local copy
@@ -172,11 +194,14 @@ def load_compile_cache() -> OperationStatus:
             suffix=CACHE_FILE_EXTENSION, dir=work_dir, delete=False
         ) as f:
             temp_path = Path(f.name)
-        logger.debug(f"Created temporary file for cache: {temp_path}")
+        logger.info(f"[LOADING] Created temporary file for cache download: {temp_path}")
 
         try:
             with temp_file_cleanup(temp_path):
                 # Phase 1: Copy from b10fs to local temp file
+                logger.info(
+                    f"[LOADING] Phase 1: Copying cache from b10fs to local temp file ({final_file} -> {temp_path})"
+                )
                 _transfer_with_b10fs_lock(
                     str(final_file),
                     str(temp_path),
@@ -185,6 +210,9 @@ def load_compile_cache() -> OperationStatus:
                 )
 
                 # Phase 2: Extract archive with space monitoring
+                logger.info(
+                    f"[LOADING] Phase 2: Extracting cache archive to torch directory ({temp_path} -> {torch_dir})"
+                )
                 _run_with_space_monitoring(
                     MIN_LOCAL_SPACE_MB,
                     work_dir,
@@ -194,11 +222,22 @@ def load_compile_cache() -> OperationStatus:
                     cleanup_func=lambda: _cleanup_torch_dir(torch_dir),
                 )
 
-            logger.info("Cache load complete")
+            # Calculate final cache size for logging
+            final_size_mb = (
+                sum(f.stat().st_size for f in torch_dir.rglob("*") if f.is_file())
+                / (1024 * 1024)
+                if torch_dir.exists()
+                else 0
+            )
+            logger.info(
+                f"[LOADING] Cache load completed successfully (final size: {final_size_mb:.2f} MB)"
+            )
             return OperationStatus.SUCCESS
 
         except CacheOperationInterrupted as e:
-            logger.warning(f"Cache load interrupted: {e}")
+            logger.warning(
+                f"[LOADING] Cache load interrupted due to insufficient disk space: {e}"
+            )
             return OperationStatus.ERROR
 
 
@@ -236,7 +275,7 @@ def save_compile_cache() -> OperationStatus:
 
         # Check if anything to save
         if not torch_dir.exists() or not any(torch_dir.iterdir()):
-            logger.info("No torch cache to save")
+            logger.info(f"[SAVING] No torch cache found at {torch_dir} to save")
             return OperationStatus.SKIPPED
 
         cache_filename = get_cache_filename()
@@ -244,18 +283,30 @@ def save_compile_cache() -> OperationStatus:
 
         # Check for existing cache first (early exit)
         if final_file.exists():
-            logger.info("Cache already exists in b10fs, skipping save")
+            file_size_mb = final_file.stat().st_size / (1024 * 1024)
+            logger.info(
+                f"[SAVING] Cache already exists in b10fs at {final_file} (size: {file_size_mb:.2f} MB), skipping save"
+            )
             return OperationStatus.SKIPPED
 
         with tempfile.NamedTemporaryFile(
             suffix=CACHE_FILE_EXTENSION, dir=work_dir, delete=False
         ) as f:
             local_temp = Path(f.name)
-        logger.debug(f"Created local temp file for archive: {local_temp}")
+        # Calculate source cache size for logging
+        source_size_mb = sum(
+            f.stat().st_size for f in torch_dir.rglob("*") if f.is_file()
+        ) / (1024 * 1024)
+        logger.info(
+            f"[SAVING] Created local temp file for archive: {local_temp} (source cache size: {source_size_mb:.2f} MB)"
+        )
 
         try:
             with temp_file_cleanup(local_temp):
                 # Phase 1: Compression with space monitoring
+                logger.info(
+                    f"[SAVING] Phase 1: Compressing torch cache directory ({torch_dir} -> {local_temp}, max size: {MAX_CACHE_SIZE_MB} MB)"
+                )
                 _run_with_space_monitoring(
                     REQUIRED_B10FS_SPACE_MB,
                     b10fs_dir,
@@ -265,21 +316,30 @@ def save_compile_cache() -> OperationStatus:
                 )
 
                 # Phase 2: Copy to b10fs with locking
+                compressed_size_mb = local_temp.stat().st_size / (1024 * 1024)
+                logger.info(
+                    f"[SAVING] Phase 2: Copying compressed archive to b10fs ({local_temp} -> {temp_file}, size: {compressed_size_mb:.2f} MB)"
+                )
                 _transfer_with_b10fs_lock(
                     str(local_temp), str(temp_file), "copy_in", cleanup_on_failure=True
                 )
 
                 # Phase 3: Atomic rename (fast, don't interrupt)
                 logger.info(
-                    f"Renaming temp file to final cache file: {temp_file} -> {final_file}"
+                    f"[SAVING] Phase 3: Atomically renaming temp file to final cache file: {temp_file} -> {final_file}"
                 )
                 temp_file.rename(final_file)
 
-            logger.info("Cache save complete")
+            final_file_size_mb = final_file.stat().st_size / (1024 * 1024)
+            logger.info(
+                f"[SAVING] Cache save completed successfully (final file: {final_file}, size: {final_file_size_mb:.2f} MB)"
+            )
             return OperationStatus.SUCCESS
 
         except CacheOperationInterrupted as e:
-            logger.warning(f"Cache save interrupted: {e}")
+            logger.warning(
+                f"[SAVING] Cache save interrupted due to insufficient disk space: {e}"
+            )
             return OperationStatus.ERROR
 
 
@@ -299,10 +359,23 @@ def clear_local_cache() -> bool:
     """
     torch_dir = Path(TORCH_CACHE_DIR)
     if not torch_dir.exists():
+        logger.info(
+            f"[CLEARING] No torch cache directory found at {torch_dir}, nothing to clear"
+        )
         return True
+
+    # Calculate size before clearing for logging
+    size_mb = sum(f.stat().st_size for f in torch_dir.rglob("*") if f.is_file()) / (
+        1024 * 1024
+    )
+    logger.info(
+        f"[CLEARING] Removing torch cache directory: {torch_dir} (size: {size_mb:.2f} MB)"
+    )
+
     import shutil
 
     shutil.rmtree(torch_dir)
+    logger.info(f"[CLEARING] Successfully cleared torch cache directory: {torch_dir}")
     return True
 
 
@@ -326,6 +399,8 @@ def _cache_compression_worker(
     # Import here to avoid issues with multiprocessing
     from .archive import create_archive
 
+    # Note: We can't use the main logger here due to multiprocessing
+    # The create_archive function should handle its own logging
     create_archive(torch_dir, local_temp, max_size_mb)
 
 
@@ -336,9 +411,11 @@ def _cleanup_torch_dir(torch_dir: Path) -> None:
             import shutil
 
             shutil.rmtree(torch_dir)
-            logger.debug(f"Cleaned up torch directory: {torch_dir}")
+            logger.info(
+                f"[CLEANUP] Successfully cleaned up torch directory: {torch_dir}"
+            )
     except Exception as e:
-        logger.error(f"Failed to cleanup torch directory {torch_dir}: {e}")
+        logger.error(f"[CLEANUP] Failed to cleanup torch directory {torch_dir}: {e}")
 
 
 @worker_process("Extraction was cancelled before starting")
@@ -358,4 +435,6 @@ def _cache_extract_worker(archive_path_str: str, dest_dir_str: str) -> None:
     # Import here to avoid issues with multiprocessing
     from .archive import extract_archive
 
+    # Note: We can't use the main logger here due to multiprocessing
+    # The extract_archive function should handle its own logging
     extract_archive(archive_path, dest_dir)

b10_transfer-0.2.0/src/b10_transfer/cache_cli.py

@@ -0,0 +1,124 @@
+# src/b10_tcache/cli.py
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+import urllib.error
+import urllib.request
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class WaitCfg:
+    url: str
+    timeout_s: float
+    interval_s: float
+    loglevel: str
+
+
+DEFAULT_URL = os.getenv("B10_TRANSFER_VLLM_URL", "http://127.0.0.1:8000/v1/models")
+DEFAULT_TIMEOUT_S = float(os.getenv("B10_TRANSFER_TIMEOUT_S", "1800"))  # 30m default
+DEFAULT_INTERVAL_S = float(os.getenv("B10_TRANSFER_INTERVAL_S", "2"))
+DEFAULT_LOGLEVEL = os.getenv("B10_TRANSFER_CLI_LOGLEVEL", "INFO").upper()
+
+
+VLLM_CACHE_DIR = os.getenv("VLLM_CACHE_ROOT", "~/.cache/vllm")
+
+
+def _setup_logging(level: str) -> logging.Logger:
+    logging.basicConfig(
+        level=getattr(logging, level, logging.INFO),
+        format="%(asctime)s | %(levelname)s | %(message)s",
+    )
+    return logging.getLogger("b10_transfer.cli")
+
+
+def _http_ok(url: str, logger: logging.Logger) -> bool:
+    """
+    Return True if vLLM readiness looks good.
+
+    We consider it 'ready' if GET <url> returns 200.
+    """
+    try:
+        req = urllib.request.Request(url, method="GET")
+        with urllib.request.urlopen(req, timeout=5) as resp:
+            if resp.status != 200:
+                return False
+            return True
+    except (urllib.error.URLError, urllib.error.HTTPError) as e:
+        logger.debug("Readiness probe failed: %s", e)
+        return False
+    except Exception as e:
+        logger.debug("Unexpected readiness error: %s", e)
+        return False
+
+
+def _wait_for_ready(cfg: WaitCfg, logger: logging.Logger) -> bool:
+    t0 = time.monotonic()
+    logger.info(
+        "Waiting for vLLM readiness at %s (timeout=%.0fs, interval=%.1fs)",
+        cfg.url,
+        cfg.timeout_s,
+        cfg.interval_s,
+    )
+
+    while True:
+        if _http_ok(cfg.url, logger):
+            logger.info("vLLM reported ready at %s", cfg.url)
+            return True
+        if time.monotonic() - t0 > cfg.timeout_s:
+            logger.error(
+                "Timed out after %.0fs waiting for vLLM readiness.", cfg.timeout_s
+            )
+            return False
+
+        time.sleep(cfg.interval_s)
+
+
+def main() -> None:
+    # Configure torch compile cache location
+    os.environ["TORCHINDUCTOR_CACHE_DIR"] = VLLM_CACHE_DIR
+
+    # Import here to allow environment variables to be set before the imported script uses them
+    from cache import load_compile_cache, save_compile_cache
+
+    cfg = WaitCfg(
+        url=DEFAULT_URL,
+        timeout_s=DEFAULT_TIMEOUT_S,
+        interval_s=DEFAULT_INTERVAL_S,
+        loglevel=DEFAULT_LOGLEVEL,
+    )
+
+    logger = _setup_logging(cfg.loglevel)
+
+    # 1) Preload any existing cache (non-fatal on error)
+    try:
+        logger.info("Calling load_compile_cache() …")
+        load_compile_cache()
+        logger.info("load_compile_cache() returned.")
+    except Exception as e:
+        logger.exception("load_compile_cache() failed: %s", e)
+
+    # 2) Wait for vLLM HTTP to be ready
+    try:
+        ready = _wait_for_ready(cfg, logger)
+    except Exception as e:
+        logger.exception("Readiness wait crashed: %s", e)
+        sys.exit(3)
+
+    if not ready:
+        # Loop timed out. Safe exit.
+        sys.exit(4)
+
+    # 3) Save compile cache
+    try:
+        logger.info("Calling save_compile_cache() …")
+        save_compile_cache()
+        logger.info("save_compile_cache() completed.")
+    except Exception as e:
+        logger.exception("save_compile_cache() failed: %s", e)
+        sys.exit(5)
+
+    logger.info("vLLM automatic torch compile cache done.")

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/cleanup.py

@@ -18,8 +18,9 @@ from .constants import (
     CLEANUP_INCOMPLETE_TIMEOUT_SECONDS,
 )
 from .utils import safe_execute, safe_unlink
+from .logging_utils import get_b10_logger
 
-logger = logging.getLogger(__name__)
+logger = get_b10_logger(__name__)
 
 
 @safe_execute("Failed to find stale files", [])
@@ -48,14 +49,14 @@ def _find_stale_files(
         # 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."
+                f"[CLEANUP] 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."
+                f"[CLEANUP] Found non-matching file in b10fs cache directory: {file_path}, skipping consideration for deletion in cleanup phase."
             )
             continue
 
@@ -91,7 +92,7 @@ def _cleanup_files(files: List[Path], file_type: str) -> int:
             )
             cleaned_count += 1
             logger.debug(
-                f"Cleaned stale {file_type} file: {file_path.name} (age: {file_age:.1f}s)"
+                f"[CLEANUP] Cleaned stale {file_type} file: {file_path.name} (age: {file_age:.1f}s)"
             )
         except OSError:
             # File might have been deleted by another pod
@@ -120,7 +121,7 @@ def cooperative_cleanup_b10fs() -> None:
     """
     b10fs_dir = Path(B10FS_CACHE_DIR)
     if not b10fs_dir.exists():
-        logger.debug("b10fs cache directory doesn't exist, skipping cleanup")
+        logger.debug("[CLEANUP] b10fs cache directory doesn't exist, skipping cleanup")
         return
 
     # Find and clean stale lock files
@@ -138,11 +139,11 @@ def cooperative_cleanup_b10fs() -> None:
     total_cleaned = cleaned_locks + cleaned_incomplete
     if total_cleaned > 0:
         logger.info(
-            f"Cooperative cleanup: removed {cleaned_locks} stale locks, "
+            f"[CLEANUP] Cooperative cleanup completed: removed {cleaned_locks} stale locks, "
             f"{cleaned_incomplete} incomplete files"
         )
     else:
-        logger.debug("Cooperative cleanup: no stale files found")
+        logger.debug("[CLEANUP] Cooperative cleanup completed: no stale files found")
 
 
 def get_cleanup_info() -> dict:

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/core.py

@@ -26,8 +26,9 @@ from .constants import (
     MIN_LOCAL_SPACE_MB,
     OperationStatus,
 )
+from .logging_utils import get_b10_logger
 
-logger = logging.getLogger(__name__)
+logger = get_b10_logger(__name__)
 
 
 @timed_fn(logger=logger, name="Transferring file")
@@ -59,7 +60,7 @@ def transfer(source: str, dest: str) -> OperationStatus:
 
     # Validate source file exists
     if not source_path.exists():
-        logger.error(f"Source file does not exist: {source}")
+        logger.error(f"[TRANSFER] Source file does not exist: {source}")
         return OperationStatus.ERROR
 
     # Create destination directory if it doesn't exist
@@ -71,19 +72,19 @@ def transfer(source: str, dest: str) -> OperationStatus:
         # 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"
+            f"[TRANSFER] 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"
+            f"[TRANSFER] 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"
+        f"[TRANSFER] Initial space check passed: {space_threshold_mb:.1f}MB required at destination"
     )
 
     # Start background space monitoring for destination directory
@@ -92,7 +93,7 @@ def transfer(source: str, dest: str) -> OperationStatus:
 
     try:
         # Run monitored copy process
-        logger.info(f"Starting transfer: {source} -> {dest}")
+        logger.info(f"[TRANSFER] Starting file transfer: {source} -> {dest}")
         run_monitored_process(
             _cache_copy_worker,
             (str(source_path), str(dest_path)),
@@ -103,11 +104,15 @@ def transfer(source: str, dest: str) -> OperationStatus:
             ),
         )
 
-        logger.info("File transfer complete")
+        logger.info(
+            f"[TRANSFER] File transfer completed successfully: {source} -> {dest}"
+        )
         return OperationStatus.SUCCESS
 
     except CacheOperationInterrupted as e:
-        logger.warning(f"File transfer interrupted: {e}")
+        logger.warning(
+            f"[TRANSFER] File transfer interrupted due to insufficient disk space: {e}"
+        )
         return OperationStatus.ERROR
 
     finally:

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

@@ -18,7 +18,9 @@ except ImportError:
     torch = None
     TORCH_AVAILABLE = False
 
-logger = logging.getLogger(__name__)
+from .logging_utils import get_b10_logger
+
+logger = get_b10_logger(__name__)
 
 KEY_LENGTH = 16
 UNKNOWN_HOSTNAME = "unknown-host"
@@ -85,10 +87,12 @@ def get_environment_key() -> str:
         return hashlib.sha256(node_json.encode("utf-8")).hexdigest()[:KEY_LENGTH]
 
     except (ImportError, RuntimeError, AssertionError) as e:
-        logger.error(f"GPU environment unavailable: {e}")
+        logger.error(f"[ENVIRONMENT] 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}")
+        logger.error(
+            f"[ENVIRONMENT] Unexpected error during environment key generation: {e}"
+        )
         raise RuntimeError(f"Environment key generation failed: {e}") from e
 
 

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/info.py

@@ -13,7 +13,9 @@ from .constants import (
 )
 from .utils import safe_execute, _is_b10fs_enabled
 
-logger = logging.getLogger(__name__)
+from .logging_utils import get_b10_logger
+
+logger = get_b10_logger(__name__)
 
 
 @safe_execute("Failed to calculate local cache size", None)

b10_transfer-0.2.0/src/b10_transfer/logging_utils.py

@@ -0,0 +1,117 @@
+"""Centralized logging utilities for b10-transfer package with colored output."""
+
+import logging
+from typing import Optional
+
+
+class ColoredFormatter(logging.Formatter):
+    """Custom formatter that adds colors and b10-transfer prefix to log messages."""
+
+    # ANSI color codes
+    COLORS = {
+        "cyan": "\033[96m",
+        "green": "\033[92m",
+        "red": "\033[91m",
+        "yellow": "\033[93m",
+        "reset": "\033[0m",
+    }
+
+    def format(self, record):
+        # Add the b10-transfer prefix to the message
+        original_msg = record.getMessage()
+
+        # Determine color based on log level and message content
+        color = self._get_message_color(record, original_msg)
+
+        # Format the message with color and prefix
+        colored_msg = f"{self.COLORS[color]}[b10-transfer log] {original_msg}{self.COLORS['reset']}"
+
+        # Temporarily replace the message for formatting
+        record.msg = colored_msg
+        record.args = ()
+
+        # Use the parent formatter
+        formatted = super().format(record)
+
+        return formatted
+
+    def _get_message_color(self, record, message: str) -> str:
+        """Determine the appropriate color for the log message."""
+        # Red for errors and failures
+        if record.levelno >= logging.ERROR:
+            return "red"
+
+        # Red for warning messages that indicate failures
+        if record.levelno == logging.WARNING and any(
+            keyword in message.lower()
+            for keyword in ["failed", "error", "interrupted", "cancelled", "abort"]
+        ):
+            return "red"
+
+        # Green for success messages
+        if any(
+            keyword in message.lower()
+            for keyword in [
+                "completed successfully",
+                "success",
+                "complete",
+                "finished",
+                "saved",
+                "loaded",
+                "extracted",
+                "compressed",
+                "transferred",
+                "cleared successfully",
+            ]
+        ):
+            return "green"
+
+        # Default to cyan
+        return "cyan"
+
+
+def get_b10_logger(name: str) -> logging.Logger:
+    """Get a logger configured with b10-transfer colored formatting.
+
+    Args:
+        name: The logger name (typically __name__)
+
+    Returns:
+        Logger configured with colored b10-transfer formatting
+    """
+    logger = logging.getLogger(name)
+
+    # Only add handler if it doesn't already exist
+    if not any(
+        isinstance(h, logging.StreamHandler)
+        and isinstance(h.formatter, ColoredFormatter)
+        for h in logger.handlers
+    ):
+        # Create handler with colored formatter
+        handler = logging.StreamHandler()
+        formatter = ColoredFormatter("%(levelname)s - %(message)s")
+        handler.setFormatter(formatter)
+
+        # Add handler to logger
+        logger.addHandler(handler)
+        logger.setLevel(logging.INFO)
+
+        # Prevent duplicate messages from parent loggers
+        logger.propagate = False
+
+    return logger
+
+
+def log_success(logger: logging.Logger, message: str):
+    """Log a success message that will be colored green."""
+    logger.info(message)
+
+
+def log_failure(logger: logging.Logger, message: str, level: int = logging.ERROR):
+    """Log a failure message that will be colored red."""
+    logger.log(level, message)
+
+
+def log_info(logger: logging.Logger, message: str):
+    """Log an info message that will be colored cyan."""
+    logger.info(message)

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/space_monitor.py

@@ -14,8 +14,9 @@ from multiprocessing import Process, Queue
 from functools import wraps
 
 from .constants import WorkerStatus, SPACE_MONITOR_CHECK_INTERVAL_SECONDS
+from .logging_utils import get_b10_logger
 
-logger = logging.getLogger(__name__)
+logger = get_b10_logger(__name__)
 
 
 class CacheOperationInterrupted(Exception):
@@ -161,7 +162,7 @@ class CacheSpaceMonitor:
         self.thread = threading.Thread(target=self._monitor, daemon=True)
         self.thread.start()
         logger.debug(
-            f"Started space monitor for {self.path} (required: {self.required_space_mb:.1f}MB)"
+            f"[MONITORING] Started space monitor for {self.path} (required: {self.required_space_mb:.1f}MB)"
         )
 
     def _monitor(self) -> None:
@@ -170,18 +171,18 @@ class CacheSpaceMonitor:
             try:
                 available_mb = get_available_disk_space_mb(self.path)
                 logger.debug(
-                    f"[SpaceMonitor] Available space: {available_mb:.1f}MB (required: {self.required_space_mb:.1f}MB)"
+                    f"[MONITORING] Available space: {available_mb:.1f}MB (required: {self.required_space_mb:.1f}MB)"
                 )
 
                 if available_mb < self.required_space_mb:
                     logger.error(
-                        f"CRITICAL: Space ({available_mb:.1f}MB) below required {self.required_space_mb:.1f}MB. Signaling stop!"
+                        f"[MONITORING] CRITICAL: Space ({available_mb:.1f}MB) below required {self.required_space_mb:.1f}MB. Signaling stop!"
                     )
                     self.stop_operation.set()
                     break
 
             except Exception as e:
-                logger.warning(f"Space monitor error: {e}")
+                logger.warning(f"[MONITORING] Space monitor error: {e}")
 
             time.sleep(self.check_interval)
 
@@ -197,7 +198,7 @@ class CacheSpaceMonitor:
         """Stop the background monitoring thread."""
         self.stop_operation.set()
         if self.thread is not None:
-            logger.debug("Stopped space monitor")
+            logger.debug("[MONITORING] Stopped space monitor")
 
 
 def cleanup_process(
@@ -217,7 +218,7 @@ def cleanup_process(
         process.terminate()
         process.join(timeout=timeout)
         if process.is_alive():
-            logger.warning(f"Force killing {operation_name} process")
+            logger.warning(f"[MONITORING] Force killing {operation_name} process")
             process.kill()
             process.join()
 
@@ -260,7 +261,9 @@ def run_monitored_process(
         # Monitor the process
         while process.is_alive():
             if space_monitor.should_stop():
-                logger.warning(f"Low disk space detected, cancelling {operation_name}")
+                logger.warning(
+                    f"[MONITORING] Low disk space detected, cancelling {operation_name}"
+                )
                 stop_event.set()
                 cleanup_process(process, operation_name)
 
@@ -280,7 +283,9 @@ def run_monitored_process(
         if not result_queue.empty():
             status, error_msg = result_queue.get()
             if status == WorkerStatus.ERROR.value:
-                logger.error(f"{operation_name} worker failed: {error_msg}")
+                logger.error(
+                    f"[MONITORING] {operation_name} worker failed: {error_msg}"
+                )
                 raise Exception(error_msg)
             elif status == WorkerStatus.CANCELLED.value:
                 if cleanup_func:
@@ -288,12 +293,12 @@ def run_monitored_process(
                 raise CacheOperationInterrupted(error_msg)
             # status == WorkerStatus.SUCCESS.value - continue normally
 
-        logger.debug(f"{operation_name} completed successfully")
+        logger.debug(f"[MONITORING] {operation_name} completed successfully")
 
     except Exception as e:
         # Ensure process is cleaned up
         cleanup_process(process, operation_name)
 
         if not isinstance(e, CacheOperationInterrupted):
-            logger.error(f"{operation_name} failed: {e}")
+            logger.error(f"[MONITORING] {operation_name} failed: {e}")
         raise

{b10_transfer-0.1.7 → b10_transfer-0.2.0}/src/b10_transfer/utils.py

@@ -6,7 +6,9 @@ from pathlib import Path
 from contextlib import contextmanager
 from typing import Generator, Any
 
-logger = logging.getLogger(__name__)
+from .logging_utils import get_b10_logger
+
+logger = get_b10_logger(__name__)
 
 # Lock file settings
 LOCK_WAIT_SLEEP_SECONDS = 1.0  # How long to wait between lock file checks
@@ -137,7 +139,7 @@ def apply_cap(value: int, cap: int, name: str) -> int:
     """
     if value > cap:
         logger.warning(
-            f"{name} capped at {cap} (requested {value}) for security/stability"
+            f"[UTILS] {name} capped at {cap} (requested {value}) for security/stability"
         )
         return cap
     return value
@@ -159,11 +161,11 @@ def timed_fn(logger=logger, name=None):
 
     def decorator(fn):
         def wrapper(*args, **kwargs):
-            logger.info(f"{name or fn.__name__} started")
+            logger.info(f"[TIMING] {name or fn.__name__} started")
             start = time.perf_counter()
             result = fn(*args, **kwargs)
             logger.info(
-                f"{name or fn.__name__} finished in {time.perf_counter() - start:.2f}s"
+                f"[TIMING] {name or fn.__name__} finished in {time.perf_counter() - start:.2f}s"
             )
             return result
 
@@ -193,7 +195,7 @@ def safe_execute(error_message: str, default_return: Any = None):
             try:
                 return func(*args, **kwargs)
             except Exception as e:
-                logger.error(f"{error_message}: {e}")
+                logger.error(f"[ERROR] {error_message}: {e}")
                 return default_return
 
         return wrapper
@@ -233,7 +235,7 @@ def critical_section_b10fs_file_lock(name):
 
             lock_file = lock_dir / f"{name}.lock"
             while lock_file.exists():
-                logger.debug("Waiting for lock file to be released...")
+                logger.debug("[LOCKING] Waiting for lock file to be released...")
                 time.sleep(LOCK_WAIT_SLEEP_SECONDS)
 
             try:
@@ -267,9 +269,9 @@ def safe_unlink(
     try:
         file_path.unlink(missing_ok=True)
         if success_message:
-            logger.debug(success_message)
+            logger.debug(f"[UTILS] {success_message}")
     except Exception as e:
-        logger.error(f"{error_message}: {e}")
+        logger.error(f"[UTILS] {error_message}: {e}")
 
 
 @contextmanager
@@ -348,8 +350,8 @@ def cache_operation(operation_name: str) -> Generator[None, None, None]:
         _validate_b10fs_available()
         yield
     except CacheValidationError as e:
-        logger.debug(f"{operation_name} failed: {e}")
+        logger.debug(f"[OPERATION] {operation_name} failed: {e}")
         raise
     except Exception as e:
-        logger.debug(f"{operation_name} failed: {e}")
+        logger.debug(f"[OPERATION] {operation_name} failed: {e}")
         raise