b10-transfer 0.1.8__py3-none-any.whl → 0.2.0__py3-none-any.whl

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.0"
 
 __all__ = [
     "CacheError",

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.8.dist-info → b10_transfer-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: b10-transfer
-Version: 0.1.8
+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.8.dist-info → b10_transfer-0.2.0.dist-info}/RECORD

@@ -1,6 +1,7 @@
-b10_transfer/__init__.py,sha256=ik92ay-2iZ3qcRWOP9yNuWiM6TVKzFZBA_m8UTWiKCQ,729
+b10_transfer/__init__.py,sha256=K4E2V4gYvauE65ZCgS0Artx6aVJpEFsGAUDNjHE0Sj4,729
 b10_transfer/archive.py,sha256=gunAZ6oTUz0SxodUCho1uP-MNLuQyuhMKFusx7E0xDg,6439
 b10_transfer/cache.py,sha256=k759Cs4IIUdAJ90ctK1m8ws3Fceje8ckug09L8X6W5I,17518
+b10_transfer/cache_cli.py,sha256=I_MOUgYqMdtxjX0caP__jmp957koQy86VsSDwn4CmtU,3602
 b10_transfer/cleanup.py,sha256=IAjRlpzCcXVrokTXdOsm_EJo_-U5kiy8KJoip8V4vvk,6345
 b10_transfer/constants.py,sha256=iuLShDW6hInhyz2YTQ8CzBanqW4chCkQOAzPZkCtOoA,4322
 b10_transfer/core.py,sha256=Ny4lViiPup4Zb_OCx_Z2IcbcWq42Ymj_zF7xyVroHxE,4668
@@ -9,6 +10,7 @@ b10_transfer/info.py,sha256=WlpSQNEKi93d3EYo2HnbbzrLazOhK9dpX5w7Kf6DVfA,6339
 b10_transfer/logging_utils.py,sha256=vnjnuVsVO5bibHSHiF5sYhHqfprIezU0fM3YaCugMC0,3488
 b10_transfer/space_monitor.py,sha256=-hFc9f29K-7gF1vDZRrKDa-FrHkmD6OBLsV8l8knrMM,10761
 b10_transfer/utils.py,sha256=XG4dGLAlQQyUPeIK48OzuDkz1SOpdEs2_GrXKd8fr5M,12061
-b10_transfer-0.1.8.dist-info/METADATA,sha256=bV0Qk7Dwsq3ZWroigvDCKqvzERNAdCDiHaMv6pqCauA,4108
-b10_transfer-0.1.8.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-b10_transfer-0.1.8.dist-info/RECORD,,
+b10_transfer-0.2.0.dist-info/METADATA,sha256=iP567IJPB9rTq3L4FSFz8QQVF89vnak4iJvrBbgzkRA,4108
+b10_transfer-0.2.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+b10_transfer-0.2.0.dist-info/entry_points.txt,sha256=6vj8KpPhV5qkaix1Vs7c--RL8nmuaGx_jIvkp4mRXoY,60
+b10_transfer-0.2.0.dist-info/RECORD,,

b10_transfer-0.2.0.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+b10-transfer=b10_transfer.cache_cli:main
+