b10-transfer 0.1.8__tar.gz → 0.2.1__tar.gz

{b10_transfer-0.1.8 → b10_transfer-0.2.1}/PKG-INFO

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

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.1.8"
+version = "0.2.1"
 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.8 → b10_transfer-0.2.1}/src/b10_transfer/__init__.py

@@ -9,7 +9,7 @@ from .constants import OperationStatus
 from .logging_utils import get_b10_logger
 
 # Version
-__version__ = "0.1.8"
+__version__ = "0.2.1"
 
 __all__ = [
     "CacheError",

b10_transfer-0.2.1/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.")