biondizzle/vllm-with-lmcache
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix loop

Browse Source
This commit is contained in:
biondizzle
2026-04-23 00:36:30 +00:00
parent 66e37bf3fb
commit 7899827a03
2 changed files with 67 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Dockerfile
View File

@@ -51,3 +51,4 @@ RUN mkdir -p /opt/vllm-shim/vllm/entrypoints/openai \
/opt/vllm-shim/vllm/entrypoints/cli/__init__.py
ENV PYTHONPATH="/opt/vllm-shim:${PYTHONPATH}"
ENV PYTHONUNBUFFERED=1

90
vllm_shim_module.py
View File

@@ -3,9 +3,9 @@
vLLM shim with custom weights download.
Intercepts `python -m vllm.entrypoints.openai.api_server` so that
if --model points to a URL (http/https), we download + extract it
to a local cache dir, then replace --model with that local path
before handing off to the real vLLM server.
if --model or the positional model arg (after "serve") points to a URL,
we download + extract it to a local cache dir, then replace it with
the local path before handing off to the real vLLM server.
Supported archive formats (detected from URL extension):
.tar, .tar.gz, .tgz, .tar.bz2, .tar.xz, .zip
@@ -15,7 +15,7 @@ import sys
import subprocess
import datetime
import shutil
import tempfile
import time
import urllib.parse
import urllib.request
@@ -23,13 +23,17 @@ import urllib.request
# Production stack mounts the PVC at /data — use a subdir so it persists across pod restarts
CACHE_DIR = os.environ.get("VLLM_WEIGHTS_CACHE", "/data/weights")
# The shim dir that shadows the vllm package — must be stripped from PYTHONPATH
# before exec'ing the real vLLM, otherwise we loop forever.
SHIM_DIR = "/opt/vllm-shim"
def log(msg: str):
"""Write to both stdout and the shim log file."""
log_path = os.environ.get("VLLM_SHIM_LOG", "/tmp/vllm-shim.log")
ts = datetime.datetime.now().isoformat()
line = f"[{ts}] {msg}"
print(line)
print(line, flush=True)
try:
with open(log_path, "a") as f:
f.write(line + "\n")
@@ -37,18 +41,19 @@ def log(msg: str):
pass
def is_url(value: str) -> bool:
return value.startswith("http://") or value.startswith("https://")
def detect_archive_type(url: str) -> str:
"""
Detect archive type from URL path extension.
Returns one of: 'tar', 'tar.gz', 'tar.bz2', 'tar.xz', 'zip', or '' (unknown).
"""
# Strip query string and fragment
path = urllib.parse.urlparse(url).path
# Check multi-part extensions first
for ext in (".tar.gz", ".tar.bz2", ".tar.xz"):
if path.endswith(ext):
return ext.lstrip(".")
# Single-part
_, ext = os.path.splitext(path)
mapping = {
".tar": "tar",
@@ -89,7 +94,6 @@ def _download_progress(block_num, block_size, total_size):
return
downloaded = block_num * block_size
pct = min(downloaded * 100 // total_size, 100)
# Only print every 10% to avoid spam
if pct % 10 == 0 and pct > 0:
mb_down = downloaded / (1024 * 1024)
mb_total = total_size / (1024 * 1024)
@@ -105,8 +109,6 @@ def extract_archive(archive_path: str, dest_dir: str, archive_type: str):
elif archive_type == "tar.bz2":
shutil.unpack_archive(archive_path, dest_dir, "bztar")
elif archive_type == "tar.xz":
# shutil.unpack_archive doesn't support xztar in all Pythons,
# use subprocess for reliability
subprocess.run(
["tar", "-xJf", archive_path, "-C", dest_dir],
check=True,
@@ -125,13 +127,11 @@ def find_model_dir(extract_dir: str) -> str:
After extraction, find the directory containing the actual model weights.
Walks the tree looking for .safetensors files and returns the directory
that contains one. This handles archives with extra parent dirs,
nested structures, or flat extractions — doesn't matter how people
compressed it, we find the safetensors.
nested structures, or flat extractions.
"""
for root, dirs, files in os.walk(extract_dir):
if any(f.endswith(".safetensors") for f in files):
return root
# Fallback: if no safetensors found, use the old heuristic
log("WARNING: No .safetensors files found in extracted archive, falling back to single-dir heuristic")
entries = [e for e in os.listdir(extract_dir)
if not e.startswith(".") and e != "__MACOSX"]
@@ -146,10 +146,9 @@ def download_and_extract_model(url: str) -> str:
Uses a cache keyed by URL filename to avoid re-downloading.
"""
url_filename = os.path.basename(urllib.parse.urlparse(url).path)
cache_key = os.path.splitext(url_filename)[0] # e.g. "model-v1.0"
cache_key = os.path.splitext(url_filename)[0]
local_dir = os.path.join(CACHE_DIR, cache_key)
# Already extracted?
if os.path.isdir(local_dir) and os.listdir(local_dir):
model_path = find_model_dir(local_dir)
log(f"Using cached weights: {model_path}")
@@ -163,13 +162,11 @@ def download_and_extract_model(url: str) -> str:
f"Supported extensions: .tar, .tar.gz, .tgz, .tar.bz2, .tar.xz, .zip"
)
# Download to a temp file in the cache dir
tmp_archive = os.path.join(CACHE_DIR, url_filename + ".tmp")
try:
download_file(url, tmp_archive)
extract_archive(tmp_archive, local_dir, archive_type)
finally:
# Clean up the archive file to save space
if os.path.exists(tmp_archive):
os.remove(tmp_archive)
@@ -178,13 +175,19 @@ def download_and_extract_model(url: str) -> str:
def parse_args(args):
"""
Parse argv, intercepting --model.
If --model is a URL, download+extract and replace with local path.
Parse argv, intercepting --model and positional model args.
Production stack invokes: python -m vllm.entrypoints.openai.api_server serve <model-url> ...
The model can appear as:
- --model <url>
- --model=<url>
- A positional arg after "serve" subcommand
If the value is a URL, download+extract and replace with local path.
Returns the modified argv list.
"""
result = []
i = 0
model_replaced = False
saw_serve = False
while i < len(args):
arg = args[i]
@@ -192,7 +195,7 @@ def parse_args(args):
# --model=<value>
if arg.startswith("--model="):
value = arg.split("=", 1)[1]
if value.startswith("http://") or value.startswith("https://"):
if is_url(value):
local_path = download_and_extract_model(value)
result.append(f"--model={local_path}")
model_replaced = True
@@ -207,7 +210,7 @@ def parse_args(args):
i += 1
if i < len(args):
value = args[i]
if value.startswith("http://") or value.startswith("https://"):
if is_url(value):
local_path = download_and_extract_model(value)
result.append(local_path)
model_replaced = True
@@ -216,15 +219,51 @@ def parse_args(args):
i += 1
continue
# "serve" subcommand — next positional is the model
if arg == "serve":
result.append(arg)
saw_serve = True
i += 1
# The next non-flag argument is the model
if i < len(args) and not args[i].startswith("-") and is_url(args[i]):
local_path = download_and_extract_model(args[i])
result.append(local_path)
model_replaced = True
i += 1
continue
# Positional model arg when there's no "serve" subcommand
# (first non-flag arg if no serve seen)
if not arg.startswith("-") and not saw_serve and not model_replaced:
if is_url(arg):
local_path = download_and_extract_model(arg)
result.append(local_path)
model_replaced = True
i += 1
continue
result.append(arg)
i += 1
if model_replaced:
log("--model URL was replaced with local path")
log("Model URL was replaced with local path")
return result
def strip_shim_from_pythonpath():
"""
Remove the shim directory from PYTHONPATH so that when we exec the
real vLLM, Python doesn't find our shadow package again (infinite loop).
"""
pp = os.environ.get("PYTHONPATH", "")
parts = [p for p in pp.split(":") if p != SHIM_DIR]
new_pp = ":".join(parts)
if new_pp != pp:
os.environ["PYTHONPATH"] = new_pp
log(f"Stripped {SHIM_DIR} from PYTHONPATH (was: {pp!r}, now: {new_pp!r})")
def main():
args = sys.argv[1:]
@@ -233,9 +272,12 @@ def main():
log(f" Invoked as: python -m {__name__} {' '.join(args)}")
log("=" * 50)
# Intercept --model if it's a URL
# Intercept --model / positional model if it's a URL
modified_args = parse_args(args)
# Strip our shim from PYTHONPATH so the real vLLM resolves correctly
strip_shim_from_pythonpath()
# Build the real vLLM command
vllm_cmd = [sys.executable, "-m", "vllm.entrypoints.openai.api_server"] + modified_args