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

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

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

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "b10-transfer"
-version = "0.2.1"
+version = "0.2.3"
 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>"]

{b10_transfer-0.2.1 → b10_transfer-0.2.3}/src/b10_transfer/__init__.py

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

{b10_transfer-0.2.1 → b10_transfer-0.2.3}/src/b10_transfer/cache_cli.py

@@ -1,5 +1,8 @@
 # src/b10_tcache/cli.py
-from __future__ import annotations
+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
@@ -24,9 +27,6 @@ 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),
@@ -78,12 +78,6 @@ def _wait_for_ready(cfg: WaitCfg, logger: logging.Logger) -> bool:
 
 
 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,

{b10_transfer-0.2.1 → b10_transfer-0.2.3}/src/b10_transfer/constants.py

@@ -24,7 +24,9 @@ _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"
+    _torch_cache_dir,
+    ["/tmp/", "/cache/", f"{os.path.expanduser('~')}/.cache"],
+    "TORCHINDUCTOR_CACHE_DIR",
 )
 
 # B10FS cache directory validation