speedy-utils 1.1.36__py3-none-any.whl → 1.1.39__py3-none-any.whl

llm_utils/lm/async_lm/async_lm.py

@@ -15,7 +15,7 @@ from pydantic import BaseModel
 
 # from llm_utils.lm.async_lm.async_llm_task import OutputModelType
 from llm_utils.lm.async_lm.async_lm_base import AsyncLMBase
-from speedy_utils import jloads
+from json_repair import jloads
 
 from ._utils import (
     LegacyMsgs,

llm_utils/scripts/fast_vllm.py

@@ -0,0 +1,131 @@
+import os
+import sys
+import shutil
+import time
+import argparse
+import subprocess
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+def get_hf_cache_home():
+    """Locate the Hugging Face cache directory."""
+    if "HF_HOME" in os.environ:
+        return Path(os.environ["HF_HOME"]) / "hub"
+    return Path.home() / ".cache" / "huggingface" / "hub"
+
+def resolve_model_path(model_id, cache_dir):
+    """Find the physical snapshot directory for the given model ID."""
+    dir_name = "models--" + model_id.replace("/", "--")
+    model_root = cache_dir / dir_name
+    if not model_root.exists():
+        raise FileNotFoundError(f"Model folder not found at: {model_root}")
+
+    # 1. Try to find hash via refs/main
+    ref_path = model_root / "refs" / "main"
+    if ref_path.exists():
+        with open(ref_path, "r") as f:
+            commit_hash = f.read().strip()
+        snapshot_path = model_root / "snapshots" / commit_hash
+        if snapshot_path.exists():
+            return snapshot_path
+    
+    # 2. Fallback to the newest snapshot folder
+    snapshots_dir = model_root / "snapshots"
+    if snapshots_dir.exists():
+        subdirs = [x for x in snapshots_dir.iterdir() if x.is_dir()]
+        if subdirs:
+            return sorted(subdirs, key=lambda x: x.stat().st_mtime, reverse=True)[0]
+            
+    raise FileNotFoundError(f"No valid snapshot found in {model_root}")
+
+def copy_worker(src, dst):
+    """Copy a single file, following symlinks to capture actual data."""
+    try:
+        os.makedirs(os.path.dirname(dst), exist_ok=True)
+        # copy2 follows symlinks by default
+        shutil.copy2(src, dst)
+        return os.path.getsize(dst)
+    except Exception as e:
+        return str(e)
+
+def cache_to_ram(model_id, shm_base, workers=64):
+    """Parallel copy from HF cache to the specified RAM directory."""
+    cache_home = get_hf_cache_home()
+    src_path = resolve_model_path(model_id, cache_home)
+    
+    safe_name = model_id.replace("/", "_")
+    dst_path = Path(shm_base) / safe_name
+
+    # Check available space in shm
+    shm_stats = shutil.disk_usage(shm_base)
+    print(f"📦 Source: {src_path}", file=sys.stderr)
+    print(f"🚀 Target RAM: {dst_path} (Available: {shm_stats.free/(1024**3):.1f} GB)", file=sys.stderr)
+
+    files_to_copy = []
+    for root, _, files in os.walk(src_path):
+        for file in files:
+            full_src = Path(root) / file
+            rel_path = full_src.relative_to(src_path)
+            files_to_copy.append((full_src, dst_path / rel_path))
+
+    total_bytes = 0
+    start = time.time()
+    with ThreadPoolExecutor(max_workers=workers) as pool:
+        futures = {pool.submit(copy_worker, s, d): s for s, d in files_to_copy}
+        for i, future in enumerate(as_completed(futures)):
+            res = future.result()
+            if isinstance(res, int):
+                total_bytes += res
+            if i % 100 == 0 or i == len(files_to_copy) - 1:
+                print(f"   Progress: {i+1}/{len(files_to_copy)} files...", end="\r", file=sys.stderr)
+
+    elapsed = time.time() - start
+    print(f"\n✅ Copied {total_bytes/(1024**3):.2f} GB in {elapsed:.2f}s", file=sys.stderr)
+    return dst_path
+
+def main():
+    parser = argparse.ArgumentParser(description="vLLM RAM-cached loader", add_help=False)
+    parser.add_argument("--model", type=str, required=True, help="HuggingFace Model ID")
+    parser.add_argument("--shm-dir", type=str, default="/dev/shm", help="RAM disk mount point")
+    parser.add_argument("--cache-workers", type=int, default=64, help="Threads for copying")
+    parser.add_argument("--keep-cache", action="store_true", help="Do not delete files from RAM on exit")
+    
+    # Capture wrapper args vs vLLM args
+    args, vllm_args = parser.parse_known_args()
+
+    ram_path = None
+    try:
+        # 1. Sync weights to RAM disk
+        ram_path = cache_to_ram(args.model, args.shm_dir, args.cache_workers)
+
+        # 2. Prepare vLLM Command
+        # Point vLLM to the RAM files, but keep the original model ID for the API
+        cmd = [
+            "vllm", "serve", str(ram_path),
+            "--served-model-name", args.model
+        ] + vllm_args
+
+        print(f"\n🔥 Launching vLLM...")
+        print(f"   Command: {' '.join(cmd)}\n", file=sys.stderr)
+        
+        # 3. Run vLLM and wait
+        subprocess.run(cmd, check=True)
+
+    except KeyboardInterrupt:
+        print("\n👋 Process interrupted by user.", file=sys.stderr)
+    except subprocess.CalledProcessError as e:
+        print(f"\n❌ vLLM exited with error: {e}", file=sys.stderr)
+    except Exception as e:
+        print(f"\n❌ Error: {e}", file=sys.stderr)
+    finally:
+        # 4. Cleanup RAM Disk
+        if ram_path and ram_path.exists() and not args.keep_cache:
+            print(f"🧹 Cleaning up RAM cache: {ram_path}", file=sys.stderr)
+            try:
+                shutil.rmtree(ram_path)
+                print("✨ RAM disk cleared.", file=sys.stderr)
+            except Exception as e:
+                print(f"⚠️ Failed to clean {ram_path}: {e}", file=sys.stderr)
+
+if __name__ == "__main__":
+    main()

{speedy_utils-1.1.36.dist-info → speedy_utils-1.1.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 1.1.36
+Version: 1.1.39
 Summary: Fast and easy-to-use package for data science
 Project-URL: Homepage, https://github.com/anhvth/speedy
 Project-URL: Repository, https://github.com/anhvth/speedy

{speedy_utils-1.1.36.dist-info → speedy_utils-1.1.39.dist-info}/RECORD

@@ -16,10 +16,11 @@ llm_utils/lm/utils.py,sha256=dEKFta8S6Mm4LjIctcpFlEGL9RnmLm5DHd2TA70UWuA,12649
 llm_utils/lm/async_lm/__init__.py,sha256=j0xK49ooZ0Dm5GstGGHbmPMrPjd3mOXoJ1H7eAL_Z4g,122
 llm_utils/lm/async_lm/_utils.py,sha256=mB-AueWJJatTx0PXqd_oWc6Kz36cfgDmDTKgiXafCJI,6106
 llm_utils/lm/async_lm/async_llm_task.py,sha256=2PWW4vPW2jYUiGmYFo4-DHrmX5Jm8Iw_1qo6EPL-ytE,18611
-llm_utils/lm/async_lm/async_lm.py,sha256=W8n_S5PKJln9bzO9T525-tqo5DFwBZNXDucz_v-35iU,13761
+llm_utils/lm/async_lm/async_lm.py,sha256=YudjpU3QH8whtnJfIYHFSMWe6lX4Fli7JEiejzrRfYk,13760
 llm_utils/lm/async_lm/async_lm_base.py,sha256=ga5nCzows5Ye3yop41zsUxNYxcj_Vpf02DsfJ1eoE9U,8358
 llm_utils/lm/async_lm/lm_specific.py,sha256=PxP54ltrh9NrLJx7BPib52oYo_aCvDOjf7KzMjp1MYg,3929
 llm_utils/scripts/README.md,sha256=yuOLnLa2od2jp4wVy3rV0rESeiV3o8zol5MNMsZx0DY,999
+llm_utils/scripts/fast_vllm.py,sha256=00UWajLOfTorSMgmgxUOpssdg55oOHneNUY0lhVuRGQ,5128
 llm_utils/scripts/vllm_load_balancer.py,sha256=eQlH07573EDWIBkwc9ef1WvI59anLr4hQqLfZvQk7xk,37133
 llm_utils/scripts/vllm_serve.py,sha256=tPcRB_MbJ01LcqC83RHQ7W9XDS7b1Ntc0fCRdegsNXU,14761
 llm_utils/vector_cache/__init__.py,sha256=oZXpjgBuutI-Pd_pBNYAQMY7-K2C6xv8Qt6a3p88GBQ,879
@@ -50,7 +51,7 @@ vision_utils/README.md,sha256=AIDZZj8jo_QNrEjFyHwd00iOO431s-js-M2dLtVTn3I,5740
 vision_utils/__init__.py,sha256=hF54sT6FAxby8kDVhOvruy4yot8O-Ateey5n96O1pQM,284
 vision_utils/io_utils.py,sha256=pI0Va6miesBysJcllK6NXCay8HpGZsaMWwlsKB2DMgA,26510
 vision_utils/plot.py,sha256=HkNj3osA3moPuupP1VguXfPPOW614dZO5tvC-EFKpKM,12028
-speedy_utils-1.1.36.dist-info/METADATA,sha256=yZYfOkBwR1aiMwAKAxn78sYCbmNnBt-5lmGY6d11hPI,8048
-speedy_utils-1.1.36.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-speedy_utils-1.1.36.dist-info/entry_points.txt,sha256=rwn89AYfBUh9SRJtFbpp-u2JIKiqmZ2sczvqyO6s9cI,289
-speedy_utils-1.1.36.dist-info/RECORD,,
+speedy_utils-1.1.39.dist-info/METADATA,sha256=US1AXLkRjdiPfWya3vNvrZ_zmDwKrKBGYdhIW-d0ypY,8048
+speedy_utils-1.1.39.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+speedy_utils-1.1.39.dist-info/entry_points.txt,sha256=rwn89AYfBUh9SRJtFbpp-u2JIKiqmZ2sczvqyO6s9cI,289
+speedy_utils-1.1.39.dist-info/RECORD,,