diff --git a/Dockerfile b/Dockerfile index 91e2c85..0d00763 100644 --- a/Dockerfile +++ b/Dockerfile @@ -33,3 +33,21 @@ COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_pa # Copy over minimax parsers with kwargs fixes COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py COPY minimax_m2_parser.py /usr/local/lib/python3.12/dist-packages/vllm/parser/minimax_m2_parser.py + +# Copy vLLM shim that intercepts --model to download custom weights from URLs +COPY vllm_shim_module.py /opt/vllm-shim/vllm_shim_module.py + +# Shadow `python -m vllm.*` invocations via PYTHONPATH +# The shim masquerades as the vllm package so python -m vllm/entrypoints/openai/api_server +# hits our interceptor first, which downloads weights then execs the real vLLM +RUN mkdir -p /opt/vllm-shim/vllm/entrypoints/openai \ + /opt/vllm-shim/vllm/entrypoints/cli && \ + cp /opt/vllm-shim/vllm_shim_module.py /opt/vllm-shim/vllm/__main__.py && \ + cp /opt/vllm-shim/vllm_shim_module.py /opt/vllm-shim/vllm/entrypoints/openai/api_server.py && \ + cp /opt/vllm-shim/vllm_shim_module.py /opt/vllm-shim/vllm/entrypoints/cli/main.py && \ + touch /opt/vllm-shim/vllm/__init__.py \ + /opt/vllm-shim/vllm/entrypoints/__init__.py \ + /opt/vllm-shim/vllm/entrypoints/openai/__init__.py \ + /opt/vllm-shim/vllm/entrypoints/cli/__init__.py + +ENV PYTHONPATH="/opt/vllm-shim:${PYTHONPATH}" diff --git a/vllm_shim_module.py b/vllm_shim_module.py new file mode 100644 index 0000000..9044ea5 --- /dev/null +++ b/vllm_shim_module.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +""" +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. + +Supported archive formats (detected from URL extension): + .tar, .tar.gz, .tgz, .tar.bz2, .tar.xz, .zip +""" +import os +import sys +import subprocess +import datetime +import shutil +import tempfile +import urllib.parse +import urllib.request + +# Where to cache downloaded+extracted weights +# 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") + + +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) + try: + with open(log_path, "a") as f: + f.write(line + "\n") + except Exception: + pass + + +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", + ".tgz": "tar.gz", + ".zip": "zip", + } + return mapping.get(ext.lower(), "") + + +MAX_DOWNLOAD_RETRIES = int(os.environ.get("VLLM_SHIM_MAX_RETRIES", "5")) +RETRY_DELAY_SECONDS = 5 + + +def download_file(url: str, dest: str): + """Download url to dest with retries and a progress indicator.""" + for attempt in range(1, MAX_DOWNLOAD_RETRIES + 1): + try: + log(f"Downloading {url} -> {dest} (attempt {attempt}/{MAX_DOWNLOAD_RETRIES})") + urllib.request.urlretrieve(url, dest, reporthook=_download_progress) + log(f"Download complete: {dest}") + return + except Exception as e: + log(f"Download attempt {attempt} failed: {e}") + if os.path.exists(dest): + os.remove(dest) + if attempt < MAX_DOWNLOAD_RETRIES: + wait = RETRY_DELAY_SECONDS * attempt + log(f"Retrying in {wait}s...") + time.sleep(wait) + else: + log(f"All {MAX_DOWNLOAD_RETRIES} download attempts failed") + raise + + +def _download_progress(block_num, block_size, total_size): + """Simple download progress callback.""" + if total_size <= 0: + 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) + sys.stdout.write(f"\r {pct}% ({mb_down:.0f}/{mb_total:.0f} MB)") + sys.stdout.flush() + + +def extract_archive(archive_path: str, dest_dir: str, archive_type: str): + """Extract archive to dest_dir based on archive_type.""" + log(f"Extracting {archive_path} ({archive_type}) -> {dest_dir}") + if archive_type == "tar.gz" or archive_type == "tgz": + shutil.unpack_archive(archive_path, dest_dir, "gztar") + 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, + ) + elif archive_type == "tar": + shutil.unpack_archive(archive_path, dest_dir, "tar") + elif archive_type == "zip": + shutil.unpack_archive(archive_path, dest_dir, "zip") + else: + raise ValueError(f"Unsupported archive type: {archive_type}") + log(f"Extraction complete: {dest_dir}") + + +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. + """ + 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"] + if len(entries) == 1 and os.path.isdir(os.path.join(extract_dir, entries[0])): + return os.path.join(extract_dir, entries[0]) + return extract_dir + + +def download_and_extract_model(url: str) -> str: + """ + Download a model from URL, extract it, and return the local path. + 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" + 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}") + return model_path + + os.makedirs(local_dir, exist_ok=True) + archive_type = detect_archive_type(url) + if not archive_type: + raise ValueError( + f"Cannot determine archive type from URL: {url}\n" + 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) + + return find_model_dir(local_dir) + + +def parse_args(args): + """ + Parse argv, intercepting --model. + If --model is a URL, download+extract and replace with local path. + Returns the modified argv list. + """ + result = [] + i = 0 + model_replaced = False + + while i < len(args): + arg = args[i] + + # --model= + if arg.startswith("--model="): + value = arg.split("=", 1)[1] + if value.startswith("http://") or value.startswith("https://"): + local_path = download_and_extract_model(value) + result.append(f"--model={local_path}") + model_replaced = True + else: + result.append(arg) + i += 1 + continue + + # --model + if arg == "--model": + result.append(arg) + i += 1 + if i < len(args): + value = args[i] + if value.startswith("http://") or value.startswith("https://"): + local_path = download_and_extract_model(value) + result.append(local_path) + model_replaced = True + else: + result.append(value) + i += 1 + continue + + result.append(arg) + i += 1 + + if model_replaced: + log("--model URL was replaced with local path") + + return result + + +def main(): + args = sys.argv[1:] + + log("=" * 50) + log("vLLM Custom Weights Shim") + log(f" Invoked as: python -m {__name__} {' '.join(args)}") + log("=" * 50) + + # Intercept --model if it's a URL + modified_args = parse_args(args) + + # Build the real vLLM command + vllm_cmd = [sys.executable, "-m", "vllm.entrypoints.openai.api_server"] + modified_args + + log(f"Launching vLLM: {' '.join(vllm_cmd)}") + + # Exec into vLLM — replace this process so signals flow through cleanly + os.execvp(vllm_cmd[0], vllm_cmd) + + +if __name__ == "__main__": + main() + +# Also run if imported as a module (some invocation paths just import the file) +main()