b10-transfer 0.2.4__tar.gz → 0.3.1__tar.gz

{b10_transfer-0.2.4 → b10_transfer-0.3.1}/PKG-INFO

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

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.2.4"
+version = "0.3.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>"]
@@ -29,7 +29,7 @@ classifiers = [
 packages = [{include = "b10_transfer", from = "src"}]
 
 [tool.poetry.scripts]
-b10-transfer = "b10_transfer.cache_cli:main"
+b10-compile-cache = "b10_transfer.cache_cli:main"
 
 [tool.poetry.dependencies]
 python = "^3.9"

{b10_transfer-0.2.4 → b10_transfer-0.3.1}/src/b10_transfer/__init__.py

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

{b10_transfer-0.2.4 → b10_transfer-0.3.1}/src/b10_transfer/cache_cli.py

@@ -1,18 +1,15 @@
 # src/b10_tcache/cli.py
-import os
-
-VLLM_CACHE_DIR = os.getenv("VLLM_CACHE_ROOT", "~/.cache/vllm")
-os.environ["TORCHINDUCTOR_CACHE_DIR"] = VLLM_CACHE_DIR
-
-from .cache import load_compile_cache, save_compile_cache
-
 import logging
+import os
 import sys
 import time
 import urllib.error
 import urllib.request
 from dataclasses import dataclass
 
+from .cache import load_compile_cache, save_compile_cache
+from .constants import OperationStatus
+
 
 @dataclass(frozen=True)
 class WaitCfg:
@@ -90,9 +87,9 @@ def main() -> None:
 
     # 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.")
+        if load_compile_cache() == OperationStatus.SUCCESS:
+            logger.info("Compile cache loaded successfully.")
+            exit(0)
     except Exception as e:
         logger.exception("load_compile_cache() failed: %s", e)
 
@@ -101,19 +98,17 @@ def main() -> None:
         ready = _wait_for_ready(cfg, logger)
     except Exception as e:
         logger.exception("Readiness wait crashed: %s", e)
-        sys.exit(3)
+        sys.exit(1)
 
     if not ready:
         # Loop timed out. Safe exit.
-        sys.exit(4)
+        sys.exit(2)
 
     # 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)
+        sys.exit(3)
 
     logger.info("vLLM automatic torch compile cache done.")