From 7724b030262716aebc6cffb811e30dd3e4ea3f24 Mon Sep 17 00:00:00 2001 From: danstotts-ops Date: Tue, 29 Sep 2026 22:31:45 -0700 Subject: [PATCH 1/2] feat: add vLLM v2.27.0 worker source for Holo4 Hub --- .gitignore | 7 + Dockerfile | 65 +++++++ LICENSE | 21 +++ UPSTREAM.md | 7 + builder/requirements.txt | 10 + handler.py | 164 +++++++++++++++++ src/__init__.py | 0 src/args_builder.py | 343 +++++++++++++++++++++++++++++++++++ src/download_model.py | 47 +++++ src/handler.py | 164 +++++++++++++++++ src/main.py | 193 ++++++++++++++++++++ src/startup_errors.py | 118 ++++++++++++ tests/requirements.txt | 4 + tests/test_args_builder.py | 161 ++++++++++++++++ tests/test_handler.py | 103 +++++++++++ tests/test_startup_errors.py | 128 +++++++++++++ 16 files changed, 1535 insertions(+) create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 LICENSE create mode 100644 UPSTREAM.md create mode 100644 builder/requirements.txt create mode 100644 handler.py create mode 100644 src/__init__.py create mode 100644 src/args_builder.py create mode 100644 src/download_model.py create mode 100644 src/handler.py create mode 100644 src/main.py create mode 100644 src/startup_errors.py create mode 100644 tests/requirements.txt create mode 100644 tests/test_args_builder.py create mode 100644 tests/test_handler.py create mode 100644 tests/test_startup_errors.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4458cf8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ + +runpod.toml +*.pyc +.env +test/* +vllm-base/vllm-* +.DS_Store \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..b0253ba --- /dev/null +++ b/Dockerfile @@ -0,0 +1,65 @@ +# Worker image = official vLLM OpenAI server image + RunPod serverless wrapper. +# vLLM upgrades are now a single build ARG: +# docker buildx build --build-arg VLLM_VERSION=v0.23.0 ... +ARG VLLM_VERSION=v0.29.0 +FROM vllm/vllm-openai:${VLLM_VERSION} +# Re-declare so the stage can reference it in RUN steps below. +ARG VLLM_VERSION + +# RunPod serverless SDK + HTTP proxy deps (vLLM itself comes from the base image). +COPY builder/requirements.txt /requirements.txt +RUN python3 -m ensurepip --upgrade 2>/dev/null || true \ + && python3 -m pip install --no-cache-dir -r /requirements.txt + +# bitsandbytes moved out of vLLM core into the out-of-tree vllm-bnb-plugin in +# v0.28.0 (vllm-project/vllm#43529). The --no-deps is load-bearing: the plugin +# declares an UNPINNED `vllm` dependency, and letting pip resolve it replaces +# the base image's vLLM — it once backtracked to vllm 0.19.1, producing an +# image tagged v0.28.0 that actually contained 0.19.1 (broken deep_ep/NCCL +# mix). Its only real runtime need is the `bitsandbytes` package, pinned in +# builder/requirements.txt. +RUN python3 -m pip install --no-cache-dir --no-deps "vllm-bnb-plugin>=0.0.3" + +# Guard rail: the base image's vLLM must survive our pip installs — fail the +# build loudly if the installed version no longer matches the requested tag. +RUN test "$(python3 -c 'import vllm; print(vllm.__version__)')" = "${VLLM_VERSION#v}" + +# Setup for Option 2: building the image with the model baked in. +ARG MODEL_NAME="" +ARG MODEL_REVISION="" +ARG TOKENIZER_NAME="" +ARG TOKENIZER_REVISION="" +ARG QUANTIZATION="" +ARG BASE_PATH="/runpod-volume" + +ENV MODEL_NAME=$MODEL_NAME \ + MODEL_REVISION=$MODEL_REVISION \ + TOKENIZER_NAME=$TOKENIZER_NAME \ + TOKENIZER_REVISION=$TOKENIZER_REVISION \ + QUANTIZATION=$QUANTIZATION \ + BASE_PATH=$BASE_PATH \ + # The RunPod network volume mounts at $BASE_PATH; keep the HF cache there so + # model downloads persist across worker cold starts. + HF_HOME="${BASE_PATH}/huggingface-cache/hub" \ + HUGGINGFACE_HUB_CACHE="${BASE_PATH}/huggingface-cache/hub" \ + HF_DATASETS_CACHE="${BASE_PATH}/huggingface-cache/datasets" \ + HF_HUB_ENABLE_HF_TRANSFER=0 \ + TOKENIZERS_PARALLELISM=false + +COPY src /src + +# Optionally bake the model into the image at build time. Pass the +# HF_TOKEN build secret if the repo is gated: +# docker buildx build --secret id=HF_TOKEN ... --build-arg MODEL_NAME=... +RUN --mount=type=secret,id=HF_TOKEN,required=false \ + if [ -n "$MODEL_NAME" ]; then \ + if [ -f /run/secrets/HF_TOKEN ]; then \ + export HF_TOKEN=$(cat /run/secrets/HF_TOKEN); \ + fi && \ + python3 /src/download_model.py; \ + fi + +# main.py spawns `vllm serve` with args built from the environment, waits for +# /health, then starts the RunPod serverless loop. Explicit ENTRYPOINT so the +# base image's own entrypoint can never swallow our command. +ENTRYPOINT ["python3", "/src/main.py"] diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..9055f99 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Runpod + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/UPSTREAM.md b/UPSTREAM.md new file mode 100644 index 0000000..c291428 --- /dev/null +++ b/UPSTREAM.md @@ -0,0 +1,7 @@ +# Worker source + +The worker implementation (`Dockerfile`, `builder/`, `src/`, and `tests/`) is copied from [`runpod-workers/worker-vllm` tag `v2.27.0`](https://github.com/runpod-workers/worker-vllm/tree/v2.27.0), commit `76054c22c79c515f07065f523598d8efb2f9b682`. Its MIT license is in `LICENSE`. + +The root `handler.py` is a byte-for-byte copy of `src/handler.py` to meet the Runpod Hub publishing guide's root-file requirement. The Dockerfile runs the `src/` copy. Keep the copies identical when changing the handler. + +The three Holo4 bundle files in `.runpod/` and `README.md` were copied unchanged from the newer bundle comment on [DR-1532](https://linear.app/runpod/issue/DR-1532/review-serverless-worker-and-hub-listing-tests-green-marked-official). diff --git a/builder/requirements.txt b/builder/requirements.txt new file mode 100644 index 0000000..a0cef11 --- /dev/null +++ b/builder/requirements.txt @@ -0,0 +1,10 @@ +# vLLM itself comes from the vllm/vllm-openai base image; this is only what +# the RunPod wrapper adds on top. +runpod~=1.12.0 +aiohttp>=3.9 +huggingface-hub>=0.24 +# bitsandbytes moved out of vLLM core into an out-of-tree plugin in v0.28.0 +# (vllm-project/vllm#43529); without it LOAD_FORMAT/QUANTIZATION=bitsandbytes +# configs die at model load. This is its runtime library — the plugin itself +# is installed in the Dockerfile with --no-deps (see the comment there). +bitsandbytes>=0.50.2 diff --git a/handler.py b/handler.py new file mode 100644 index 0000000..c359654 --- /dev/null +++ b/handler.py @@ -0,0 +1,164 @@ +"""RunPod Serverless handler that proxies jobs to the local vLLM OpenAI server. + +`main.py` starts `vllm serve` on 127.0.0.1:VLLM_PORT and only then starts the +RunPod serverless loop, so by the time a job arrives the vLLM HTTP server is up +and fully backwards/forwards compatible — we never import vLLM. The one +exception: when vLLM could not start for a reason a restart cannot fix, +`main.py` sets `startup_error` and every job is answered with that message. + +Accepted job input shapes (all under job["input"]): + +1. RunPod OpenAI passthrough (what the platform sends on /openai/v1/...): + {"openai_route": "/v1/chat/completions", "openai_input": {...}} +2. Generic proxy to any vLLM route: + {"route": "/v1/completions", "body": {...}, "method": "POST"} + ("method" is optional — bodies POST, body-less requests GET) +3. Legacy shorthand: + {"prompt": "...", "sampling_params": {...}, "stream": false} + {"messages": [...], "sampling_params": {...}, "stream": true} + +When the request body has "stream": true, raw SSE chunks are yielded as they +arrive; otherwise the parsed JSON response is yielded once. +""" + +import logging +import os +from typing import Any, AsyncGenerator, Optional, Tuple + +import aiohttp + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + +VLLM_PORT = os.getenv("VLLM_PORT", "8000") +VLLM_BASE_URL = os.getenv("VLLM_BASE_URL", f"http://127.0.0.1:{VLLM_PORT}") +REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", "3600")) + +DEFAULT_CHAT_ROUTE = "/v1/chat/completions" +DEFAULT_COMPLETION_ROUTE = "/v1/completions" + +# Set by main.py once the vLLM subprocess is running, so we can fail fast +# instead of hanging on a dead server. +vllm_process = None + +# Set by main.py when vLLM failed to start for a reason a restart cannot fix +# (see startup_errors.py). Every job then answers with the cause instead of +# dialling a server that is not there. +startup_error: Optional[str] = None + +_default_model_cache: Optional[str] = None + + +def _is_vllm_alive() -> bool: + return vllm_process is None or vllm_process.poll() is None + + +async def _default_model(session: aiohttp.ClientSession) -> Optional[str]: + """Best-effort model id for legacy shorthand requests that omit "model".""" + global _default_model_cache + if _default_model_cache is not None: + return _default_model_cache + + served = os.getenv("SERVED_MODEL_NAME") or os.getenv("OPENAI_SERVED_MODEL_NAME_OVERRIDE") + if served: + _default_model_cache = served.split(",")[0].strip() + return _default_model_cache + + try: + async with session.get(f"{VLLM_BASE_URL}/v1/models") as resp: + data = await resp.json(content_type=None) + model_id = (data.get("data") or [{}])[0].get("id") + if model_id: + _default_model_cache = model_id + return model_id + except Exception as e: + logging.warning("Could not resolve default model from /v1/models: %s", e) + + return os.getenv("MODEL_NAME") + + +def _normalize_job_input(job_input: dict) -> Tuple[str, str, Optional[dict]]: + """Return (route, method, body) for any accepted job input shape.""" + if job_input.get("openai_input"): + return job_input.get("openai_route") or DEFAULT_CHAT_ROUTE, "POST", job_input["openai_input"] + + if job_input.get("openai_route"): + # Bare route without a body (e.g. /v1/models) is proxied as a GET. + return job_input["openai_route"], "GET", None + + if job_input.get("route"): + body = job_input.get("body") + # The verb defaults to the payload: bodies POST, body-less requests GET + # so read-only routes (/v1/models, /health) don't need "method" spelled out. + method = (job_input.get("method") or ("POST" if body else "GET")).upper() + return job_input["route"], method, body + + # Legacy shorthand: prompt/messages + sampling_params at the top level. + messages = job_input.get("messages") + prompt = job_input.get("prompt") + if messages is None and prompt is None: + raise ValueError( + "Job input must contain one of: openai_input (+openai_route), " + "route (+body), or prompt/messages." + ) + + sampling_params = dict(job_input.get("sampling_params") or {}) + body = {**sampling_params, "stream": bool(job_input.get("stream", False))} + if messages is not None: + body["messages"] = messages + return DEFAULT_CHAT_ROUTE, "POST", body + body["prompt"] = prompt + return DEFAULT_COMPLETION_ROUTE, "POST", body + + +def _error(message: str, error_type: str = "worker_error") -> dict: + return {"error": {"message": message, "type": error_type, "code": None}} + + +async def handler(job: dict) -> AsyncGenerator[Any, None]: + if startup_error: + yield _error(startup_error, error_type="startup_error") + return + + job_input = job.get("input") or {} + + try: + route, method, body = _normalize_job_input(job_input) + except ValueError as e: + yield _error(str(e)) + return + + if not _is_vllm_alive(): + yield _error("vLLM server process is not running; worker is unhealthy") + return + + headers = {"Content-Type": "application/json"} + + if method != "GET" and body is not None and "model" not in body: + timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) + try: + async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session: + model = await _default_model(session) + if model: + body = {**body, "model": model} + except Exception: + pass + + timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) + try: + async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session: + async with session.request(method, f"{VLLM_BASE_URL}{route}", json=body) as resp: + if resp.status >= 400: + detail = await resp.text() + logging.error("vLLM %s %s returned HTTP %s: %s", method, route, resp.status, detail) + yield _error(f"vLLM returned HTTP {resp.status}: {detail}") + return + + wants_stream = isinstance(body, dict) and body.get("stream") is True + if wants_stream: + async for chunk in resp.content.iter_any(): + yield chunk.decode("utf-8", errors="replace") + else: + yield await resp.json(content_type=None) + except aiohttp.ClientError as e: + logging.exception("Request to vLLM failed") + yield _error(f"Request to vLLM failed: {e}") diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/args_builder.py b/src/args_builder.py new file mode 100644 index 0000000..0228b14 --- /dev/null +++ b/src/args_builder.py @@ -0,0 +1,343 @@ +"""Translate environment variables into `vllm serve` CLI arguments. + +Rules (evaluated in order, later wins): + +1. Generic scan: every env var whose name matches a known vLLM flag in + UPPER_SNAKE_CASE form is emitted as a CLI flag. ``MAX_MODEL_LEN=4096`` turns + into ``--max-model-len 4096``. Anything the allowlist does not cover can be + passed verbatim through ``VLLM_EXTRA_ARGS``. + + Flag categories (regenerated against the current vLLM CLI — stable docs and + ``main`` sources — in 2026-07; earlier versions of this list came from + ``vllm serve --help`` of vLLM 0.11): + + - ``VALUE_FLAGS`` take a value, passed through verbatim and parsed by vLLM. + Exception: a ``0`` for a flag in ``ZERO_MEANS_UNSET`` (currently just + ``--max-num-batched-tokens``) drops the flag so vLLM picks its own default. + - ``PAIRED_BOOL_FLAGS`` are argparse ``BooleanOptionalAction`` flags + (``--flag`` / ``--no-flag``) that take NO value. A truthy env value emits + the bare ``--flag``; a falsy value emits the ``--no-flag`` form. Older + vLLM releases also accepted ``--flag false``; current vLLM rejects that + with ``unrecognized arguments``, so never emit a value after these. + - ``BOOL_FLAGS`` are plain store-true flags: truthy emits the bare flag, + falsy omits it entirely. + - ``REMOVED_FLAGS`` are gone from current vLLM. Setting their env var logs + a warning naming the replacement and emits nothing, so a stale template + default cannot crash startup. + +2. Aliases: a handful of historical env names that do not match a flag + (``MODEL_NAME`` -> ``--model``, ...). + +3. ``VLLM_EXTRA_ARGS``: raw argument string appended last, so it overrides + anything above when argparse sees the same flag twice. + Example: ``VLLM_EXTRA_ARGS="--override-generation-config '{"max_new_tokens": 512}'"``. +""" + +import logging +import os +import shlex +from typing import Mapping, Optional + +TRUE_VALUES = {"true", "1", "yes", "on"} +FALSE_VALUES = {"false", "0", "no", "off"} + +# Historical env names that do not match the UPPER_SNAKE of their flag. +ENV_ALIASES = { + "MODEL_NAME": "--model", + "MODEL_REVISION": "--revision", + "TOKENIZER_NAME": "--tokenizer", + "OPENAI_SERVED_MODEL_NAME_OVERRIDE": "--served-model-name", + "CUSTOM_CHAT_TEMPLATE": "--chat-template", + "VLLM_CONFIG_FILE": "--config", +} + +# Flags owned by the wrapper itself (or deliberately not exposed generically): +# --host/--port are pinned by main.py, --config comes from the VLLM_CONFIG_FILE +# alias, --model has the MODEL_NAME alias. --api-key stays reserved because the +# worker intentionally does not support it: the internal server binds to +# loopback only and RunPod authenticates callers at the platform. +RESERVED_ENV_VARS = frozenset({"HOST", "PORT", "API_KEY", "MODEL", "CONFIG", "HELP"}) + +# Flags that take a value. vLLM parses int/float/JSON values itself. +VALUE_FLAGS = frozenset({ + "--additional-config", + "--allowed-headers", + "--allowed-local-media-path", + "--allowed-media-domains", + "--allowed-methods", + "--allowed-origins", + "--api-server-count", + "--attention-backend", + "--block-size", + "--chat-template", + "--chat-template-content-format", + "--code-revision", + "--collect-detailed-traces", + "--compilation-config", + "--config-format", + "--convert", + "--cpu-offload-gb", + "--data-parallel-address", + "--data-parallel-backend", + "--data-parallel-rank", + "--data-parallel-rpc-port", + "--data-parallel-size", + "--data-parallel-size-local", + "--data-parallel-start-rank", + "--dbo-decode-token-threshold", + "--dbo-prefill-token-threshold", + "--decode-context-parallel-size", + "--default-mm-loras", + "--distributed-executor-backend", + "--download-dir", + "--dtype", + "--eplb-config", + "--expert-placement-strategy", + "--generation-config", + "--gpu-memory-utilization", + "--h11-max-header-count", + "--h11-max-incomplete-event-size", + "--hf-config-path", + "--hf-overrides", + "--hf-token", + "--ignore-patterns", + "--io-processor-plugin", + "--kv-cache-dtype", + "--kv-cache-memory-bytes", + "--kv-events-config", + "--kv-offloading-backend", + "--kv-offloading-size", + "--kv-transfer-config", + "--limit-mm-per-prompt", + "--load-format", + "--log-config-file", + "--logits-processors", + "--logprobs-mode", + "--long-prefill-token-threshold", + "--lora-dtype", + "--lora-modules", + "--mamba-cache-dtype", + "--mamba-ssm-cache-dtype", + "--max-cpu-loras", + "--max-log-len", + "--max-logprobs", + "--max-lora-rank", + "--max-loras", + "--max-model-len", + "--max-num-batched-tokens", + "--max-num-seqs", + "--max-parallel-loading-workers", + "--media-io-kwargs", + "--middleware", + "--mm-encoder-tp-mode", + "--mm-processor-cache-gb", + "--mm-processor-cache-type", + "--mm-processor-kwargs", + "--mm-shm-cache-max-object-size-mb", + "--model", + "--model-impl", + "--model-loader-extra-config", + "--num-gpu-blocks-override", + "--otlp-traces-endpoint", + "--override-generation-config", + "--pipeline-parallel-size", + "--pooler-config", + "--prefix-caching-hash-algo", + "--pt-load-map-location", + "--quantization", + "--reasoning-parser", + "--response-role", + "--revision", + "--root-path", + "--runner", + "--safetensors-load-strategy", + "--scheduler-cls", + "--scheduling-policy", + "--seed", + "--served-model-name", + "--show-hidden-metrics-for-version", + "--speculative-config", + "--ssl-ca-certs", + "--ssl-cert-reqs", + "--ssl-certfile", + "--ssl-keyfile", + "--structured-outputs-config", + "--tensor-parallel-size", + "--tokenizer", + "--tokenizer-mode", + "--tokenizer-revision", + "--tool-call-parser", + "--tool-parser-plugin", + "--tool-server", + "--uds", + "--uvicorn-log-level", + "--video-pruning-rate", + "--worker-cls", + "--worker-extension-cls", +}) + +# argparse BooleanOptionalAction flags: `--flag` / `--no-flag`, no value. +# Truthy env -> bare flag, falsy env -> `--no-flag`. +PAIRED_BOOL_FLAGS = frozenset({ + "--allow-credentials", + "--async-scheduling", + "--data-parallel-hybrid-lb", + "--disable-cascade-attn", + "--disable-chunked-mm-input", + "--disable-custom-all-reduce", + "--disable-fastapi-docs", + "--disable-hybrid-kv-cache-manager", + "--disable-sliding-window", + "--disable-uvicorn-access-log", + "--enable-auto-tool-choice", + "--enable-chunked-prefill", + "--enable-dbo", + "--enable-eplb", + "--enable-expert-parallel", + "--enable-force-include-usage", + "--enable-log-outputs", + "--enable-log-requests", + "--enable-lora", + "--enable-prefix-caching", + "--enable-prompt-embeds", + "--enable-prompt-tokens-details", + "--enable-request-id-headers", + "--enable-server-load-tracking", + "--enable-sleep-mode", + "--enable-ssl-refresh", + "--enable-tokenizer-info-endpoint", + "--enforce-eager", + "--exclude-tools-when-tool-choice-none", + "--fully-sharded-loras", + "--interleave-mm-strings", + "--kv-sharing-fast-prefill", + "--log-error-stack", + "--ray-workers-use-nsight", + "--return-tokens-as-token-ids", + "--skip-mm-profiling", + "--skip-tokenizer-init", + "--trust-remote-code", + "--trust-request-chat-template", + "--use-tqdm-on-load", +}) + +# argparse store-true flags: truthy env value -> bare flag, falsy -> omitted. +BOOL_FLAGS = frozenset({ + "--disable-log-stats", +}) + +# Value flags where "0" means "leave it to vLLM", because 0 is rejected by +# vLLM's config validation (SchedulerConfig requires max_num_batched_tokens >= 1) +# while endpoint templates commonly carry a stray MAX_NUM_BATCHED_TOKENS=0. +# Omitting the flag lets vLLM size the batch for the actual GPU. Note this is NOT +# generalized: plenty of flags take a legitimate 0 (--seed 0, --cpu-offload-gb 0). +ZERO_MEANS_UNSET = frozenset({ + "--max-num-batched-tokens", +}) + +# Flags current vLLM no longer accepts. Setting one logs a warning and emits +# nothing instead of dying with `unrecognized arguments` at container start. +# The note names the replacement (if any) and appears in the warning. +REMOVED_FLAGS = { + "--calculate-kv-scales": "removed in v0.28.0 (#49389); no replacement", + "--cuda-graph-sizes": "moved into COMPILATION_CONFIG as cudagraph_capture_sizes", + "--disable-frontend-multiprocessing": "no replacement", + "--disable-log-requests": "inverted upstream; use ENABLE_LOG_REQUESTS=false", + "--disable-mm-preprocessor-cache": "no replacement; tune via MM_PROCESSOR_CACHE_GB", + "--enable-lora-bias": "no replacement", + "--enable-multimodal-encoder-data-parallel": "no replacement", + "--eplb-log-balancedness": "now a field of EPLB_CONFIG (--eplb-config)", + "--eplb-step-interval": "now a field of EPLB_CONFIG (--eplb-config)", + "--eplb-window-size": "now a field of EPLB_CONFIG (--eplb-config)", + "--guided-decoding-backend": "use STRUCTURED_OUTPUTS_CONFIG (--structured-outputs-config)", + "--guided-decoding-disable-additional-properties": "use STRUCTURED_OUTPUTS_CONFIG", + "--guided-decoding-disable-any-whitespace": "use STRUCTURED_OUTPUTS_CONFIG", + "--guided-decoding-disable-fallback": "use STRUCTURED_OUTPUTS_CONFIG", + "--logits-processor-pattern": "use LOGITS_PROCESSORS (--logits-processors)", + "--lora-extra-vocab-size": "no replacement", + "--max-long-partial-prefills": "removed in v0.27.0; no replacement", + "--max-num-partial-prefills": "removed in v0.27.0; no replacement", + "--num-lookahead-slots": "speculative decoding is configured via SPECULATIVE_CONFIG", + "--num-redundant-experts": "now a field of EPLB_CONFIG (--eplb-config)", + "--override-attention-dtype": "removed in v0.28.0 (#48684); no replacement", + "--override-pooler-config": "use POOLER_CONFIG (--pooler-config)", + "--rope-scaling": "pass rope_scaling via HF_OVERRIDES (--hf-overrides)", + "--rope-theta": "pass rope_theta via HF_OVERRIDES (--hf-overrides)", + "--swap-space": "V1 has no swap space; no replacement", + "--task": "use RUNNER (--runner) / CONVERT (--convert)", +} + +ALL_FLAGS = VALUE_FLAGS | PAIRED_BOOL_FLAGS | BOOL_FLAGS + + +def flag_to_env_name(flag: str) -> str: + return flag.lstrip("-").replace("-", "_").upper() + + +def _boolean_value(flag: str, raw_value: str) -> Optional[bool]: + value = str(raw_value).strip().lower() + if value in TRUE_VALUES: + return True + if value in FALSE_VALUES: + return False + raise ValueError(f"invalid boolean value {raw_value!r} for {flag}") + + +def _apply(flag: str, raw_value, flags: dict) -> None: + value = str(raw_value).strip() + if value == "": + return + if flag in PAIRED_BOOL_FLAGS or flag in BOOL_FLAGS: + is_true = _boolean_value(flag, raw_value) + if flag in PAIRED_BOOL_FLAGS: + # `--flag` / `--no-flag`, never with a value after it. + flags[flag if is_true else f"--no-{flag.lstrip('-')}"] = None + elif is_true: + flags[flag] = None # store-true: bare flag + else: + flags.pop(flag, None) # omitted entirely + elif flag in ZERO_MEANS_UNSET and value == "0": + logging.warning("%s=0 is invalid for vLLM; treating it as unset (vLLM default)", flag) + else: + flags[flag] = value + + +def build_vllm_args(env: Optional[Mapping[str, str]] = None) -> list[str]: + """Build the `vllm serve` argument list from environment variables.""" + env = os.environ if env is None else env + flags: dict[str, Optional[str]] = {} + + # 0. Warn about env vars for flags current vLLM no longer accepts. + for flag, note in sorted(REMOVED_FLAGS.items()): + env_name = flag_to_env_name(flag) + if env.get(env_name, "").strip(): + logging.warning( + "%s is set but vLLM no longer supports %s; ignoring it (%s).", + env_name, flag, note, + ) + + # 1. Generic scan (env name == UPPER_SNAKE of the flag). + for flag in sorted(ALL_FLAGS): + env_name = flag_to_env_name(flag) + if env_name in RESERVED_ENV_VARS: + continue + value = env.get(env_name) + if value is not None: + _apply(flag, value, flags) + + # 2. Historical aliases (win over the generic scan). + for env_name, flag in ENV_ALIASES.items(): + value = env.get(env_name) + if value is not None: + _apply(flag, value, flags) + + argv: list[str] = [] + for flag, value in flags.items(): + argv.append(flag) + if value is not None: + argv.append(value) + + # 3. Raw passthrough, appended last so argparse "last occurrence wins". + extra = env.get("VLLM_EXTRA_ARGS", "") + argv.extend(shlex.split(extra)) + return argv diff --git a/src/download_model.py b/src/download_model.py new file mode 100644 index 0000000..aa48521 --- /dev/null +++ b/src/download_model.py @@ -0,0 +1,47 @@ +"""Download a model (and tokenizer) into the image at build time. + +Used by the Dockerfile "Option 2" flow: when the MODEL_NAME build ARG is set, +the weights are fetched into the HF cache and the resolved snapshot paths are +written to LOCAL_MODEL_ARGS_PATH, which main.py reads at container start. +""" + +import json +import logging +import os + +from huggingface_hub import snapshot_download + +LOCAL_MODEL_ARGS_PATH = "/local_model_args.json" + + +def snapshot(repo_id: str, revision: str | None = None) -> str: + path = snapshot_download( + repo_id, + revision=revision, + cache_dir=os.getenv("HF_HOME"), + token=os.getenv("HF_TOKEN"), + ) + logging.info("Downloaded %s to %s", repo_id, path) + return path + + +def main() -> None: + model_name = os.environ["MODEL_NAME"] + model_revision = os.getenv("MODEL_REVISION") or None + tokenizer_name = os.getenv("TOKENIZER_NAME") or model_name + tokenizer_revision = os.getenv("TOKENIZER_REVISION") or model_revision + + metadata = { + "MODEL_NAME": snapshot(model_name, model_revision), + "TOKENIZER_NAME": snapshot(tokenizer_name, tokenizer_revision), + "QUANTIZATION": os.getenv("QUANTIZATION"), + } + + with open(LOCAL_MODEL_ARGS_PATH, "w") as f: + json.dump({k: v for k, v in metadata.items() if v not in (None, "")}, f) + logging.info("Wrote %s", LOCAL_MODEL_ARGS_PATH) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + main() diff --git a/src/handler.py b/src/handler.py new file mode 100644 index 0000000..c359654 --- /dev/null +++ b/src/handler.py @@ -0,0 +1,164 @@ +"""RunPod Serverless handler that proxies jobs to the local vLLM OpenAI server. + +`main.py` starts `vllm serve` on 127.0.0.1:VLLM_PORT and only then starts the +RunPod serverless loop, so by the time a job arrives the vLLM HTTP server is up +and fully backwards/forwards compatible — we never import vLLM. The one +exception: when vLLM could not start for a reason a restart cannot fix, +`main.py` sets `startup_error` and every job is answered with that message. + +Accepted job input shapes (all under job["input"]): + +1. RunPod OpenAI passthrough (what the platform sends on /openai/v1/...): + {"openai_route": "/v1/chat/completions", "openai_input": {...}} +2. Generic proxy to any vLLM route: + {"route": "/v1/completions", "body": {...}, "method": "POST"} + ("method" is optional — bodies POST, body-less requests GET) +3. Legacy shorthand: + {"prompt": "...", "sampling_params": {...}, "stream": false} + {"messages": [...], "sampling_params": {...}, "stream": true} + +When the request body has "stream": true, raw SSE chunks are yielded as they +arrive; otherwise the parsed JSON response is yielded once. +""" + +import logging +import os +from typing import Any, AsyncGenerator, Optional, Tuple + +import aiohttp + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + +VLLM_PORT = os.getenv("VLLM_PORT", "8000") +VLLM_BASE_URL = os.getenv("VLLM_BASE_URL", f"http://127.0.0.1:{VLLM_PORT}") +REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", "3600")) + +DEFAULT_CHAT_ROUTE = "/v1/chat/completions" +DEFAULT_COMPLETION_ROUTE = "/v1/completions" + +# Set by main.py once the vLLM subprocess is running, so we can fail fast +# instead of hanging on a dead server. +vllm_process = None + +# Set by main.py when vLLM failed to start for a reason a restart cannot fix +# (see startup_errors.py). Every job then answers with the cause instead of +# dialling a server that is not there. +startup_error: Optional[str] = None + +_default_model_cache: Optional[str] = None + + +def _is_vllm_alive() -> bool: + return vllm_process is None or vllm_process.poll() is None + + +async def _default_model(session: aiohttp.ClientSession) -> Optional[str]: + """Best-effort model id for legacy shorthand requests that omit "model".""" + global _default_model_cache + if _default_model_cache is not None: + return _default_model_cache + + served = os.getenv("SERVED_MODEL_NAME") or os.getenv("OPENAI_SERVED_MODEL_NAME_OVERRIDE") + if served: + _default_model_cache = served.split(",")[0].strip() + return _default_model_cache + + try: + async with session.get(f"{VLLM_BASE_URL}/v1/models") as resp: + data = await resp.json(content_type=None) + model_id = (data.get("data") or [{}])[0].get("id") + if model_id: + _default_model_cache = model_id + return model_id + except Exception as e: + logging.warning("Could not resolve default model from /v1/models: %s", e) + + return os.getenv("MODEL_NAME") + + +def _normalize_job_input(job_input: dict) -> Tuple[str, str, Optional[dict]]: + """Return (route, method, body) for any accepted job input shape.""" + if job_input.get("openai_input"): + return job_input.get("openai_route") or DEFAULT_CHAT_ROUTE, "POST", job_input["openai_input"] + + if job_input.get("openai_route"): + # Bare route without a body (e.g. /v1/models) is proxied as a GET. + return job_input["openai_route"], "GET", None + + if job_input.get("route"): + body = job_input.get("body") + # The verb defaults to the payload: bodies POST, body-less requests GET + # so read-only routes (/v1/models, /health) don't need "method" spelled out. + method = (job_input.get("method") or ("POST" if body else "GET")).upper() + return job_input["route"], method, body + + # Legacy shorthand: prompt/messages + sampling_params at the top level. + messages = job_input.get("messages") + prompt = job_input.get("prompt") + if messages is None and prompt is None: + raise ValueError( + "Job input must contain one of: openai_input (+openai_route), " + "route (+body), or prompt/messages." + ) + + sampling_params = dict(job_input.get("sampling_params") or {}) + body = {**sampling_params, "stream": bool(job_input.get("stream", False))} + if messages is not None: + body["messages"] = messages + return DEFAULT_CHAT_ROUTE, "POST", body + body["prompt"] = prompt + return DEFAULT_COMPLETION_ROUTE, "POST", body + + +def _error(message: str, error_type: str = "worker_error") -> dict: + return {"error": {"message": message, "type": error_type, "code": None}} + + +async def handler(job: dict) -> AsyncGenerator[Any, None]: + if startup_error: + yield _error(startup_error, error_type="startup_error") + return + + job_input = job.get("input") or {} + + try: + route, method, body = _normalize_job_input(job_input) + except ValueError as e: + yield _error(str(e)) + return + + if not _is_vllm_alive(): + yield _error("vLLM server process is not running; worker is unhealthy") + return + + headers = {"Content-Type": "application/json"} + + if method != "GET" and body is not None and "model" not in body: + timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) + try: + async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session: + model = await _default_model(session) + if model: + body = {**body, "model": model} + except Exception: + pass + + timeout = aiohttp.ClientTimeout(total=REQUEST_TIMEOUT) + try: + async with aiohttp.ClientSession(timeout=timeout, headers=headers) as session: + async with session.request(method, f"{VLLM_BASE_URL}{route}", json=body) as resp: + if resp.status >= 400: + detail = await resp.text() + logging.error("vLLM %s %s returned HTTP %s: %s", method, route, resp.status, detail) + yield _error(f"vLLM returned HTTP {resp.status}: {detail}") + return + + wants_stream = isinstance(body, dict) and body.get("stream") is True + if wants_stream: + async for chunk in resp.content.iter_any(): + yield chunk.decode("utf-8", errors="replace") + else: + yield await resp.json(content_type=None) + except aiohttp.ClientError as e: + logging.exception("Request to vLLM failed") + yield _error(f"Request to vLLM failed: {e}") diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..b04a100 --- /dev/null +++ b/src/main.py @@ -0,0 +1,193 @@ +"""Container entrypoint: spawn `vllm serve`, wait for it, start RunPod serverless. + +The worker never imports vLLM. We build the CLI from environment variables +(see args_builder.py), launch `vllm serve` on the loopback interface, poll +/health until the server (and model) is ready, and only then start the RunPod +serverless job loop so no job is pulled before the backend can serve it. + +If vLLM dies during startup for a reason a restart cannot fix (CUDA OOM, a +MAX_MODEL_LEN the GPU cannot hold, a bad flag, a gated model), the worker stays +up and answers every job with the cause instead of crash-looping; see +startup_errors.py. Unrecognised failures still exit non-zero so the platform +retries them. +""" + +import collections +import json +import logging +import os +import signal +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request + +import startup_errors +from args_builder import build_vllm_args +from download_model import LOCAL_MODEL_ARGS_PATH + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + +VLLM_HOST = "127.0.0.1" +VLLM_PORT = os.getenv("VLLM_PORT", "8000") +STARTUP_TIMEOUT = int(os.getenv("VLLM_STARTUP_TIMEOUT", "1200")) # seconds +HEALTH_POLL_INTERVAL = 2 # seconds +# How long to wait for vLLM to exit after SIGTERM before SIGKILL. +SHUTDOWN_GRACE = 30 # seconds + +vllm_process: subprocess.Popen | None = None + +# Enough of vLLM's output to recognise why it died. The traceback that matters +# is always the last thing it prints. +recent_output: collections.deque[str] = collections.deque(maxlen=400) +output_pump: threading.Thread | None = None + + +def apply_local_model_args() -> None: + """Load args baked into the image by download_model.py (Option 2 builds). + + The baked model path wins over MODEL_NAME env vars, and HF hub access is + forced offline since the weights are already on disk. + """ + if not os.path.exists(LOCAL_MODEL_ARGS_PATH): + return + try: + with open(LOCAL_MODEL_ARGS_PATH) as f: + local_args = json.load(f) + except (OSError, json.JSONDecodeError) as e: + logging.error("Failed to read %s: %s", LOCAL_MODEL_ARGS_PATH, e) + return + + logging.info("Using baked-in model args: %s", local_args) + for key, value in local_args.items(): + if value not in (None, ""): + os.environ[key] = str(value) + os.environ["TRANSFORMERS_OFFLINE"] = "1" + os.environ["HF_HUB_OFFLINE"] = "1" + + +def pump_output(proc: subprocess.Popen) -> None: + """Echo vLLM's output to our stdout while keeping the tail for diagnosis.""" + for line in proc.stdout: # type: ignore[union-attr] + sys.stdout.write(line) + sys.stdout.flush() + recent_output.append(line) + + +def start_vllm() -> subprocess.Popen: + global output_pump + + argv = ["vllm", "serve", "--host", VLLM_HOST, "--port", VLLM_PORT] + argv += build_vllm_args() + + logging.info("Starting vLLM: %s", " ".join(argv)) + # vLLM's stdout+stderr flow through a pipe so we can both forward them to the + # worker logs and keep the tail to classify a startup failure. Unbuffered so + # the child's log lines arrive as they are written, not when its buffer fills. + env = {**os.environ, "PYTHONUNBUFFERED": "1"} + proc = subprocess.Popen( + argv, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + errors="replace", + env=env, + ) + output_pump = threading.Thread(target=pump_output, args=(proc,), daemon=True) + output_pump.start() + return proc + + +def stop_vllm(proc: subprocess.Popen) -> None: + """Make sure a failed vLLM is gone (and its GPU memory released).""" + if proc.poll() is None: + proc.send_signal(signal.SIGTERM) + try: + proc.wait(timeout=SHUTDOWN_GRACE) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + # Let the pump drain whatever the pipe still holds so the traceback tail + # is in recent_output before we look at it. + if output_pump is not None: + output_pump.join(timeout=10) + + +def wait_for_vllm(proc: subprocess.Popen) -> None: + """Poll GET /health until vLLM is ready; fail fast if it crashes or times out.""" + url = f"http://{VLLM_HOST}:{VLLM_PORT}/health" + + deadline = time.monotonic() + STARTUP_TIMEOUT + while time.monotonic() < deadline: + if proc.poll() is not None: + raise RuntimeError(f"vLLM serve exited during startup with code {proc.returncode}") + try: + request = urllib.request.Request(url) + with urllib.request.urlopen(request, timeout=10) as resp: + if resp.status == 200: + logging.info("vLLM is healthy") + return + except (urllib.error.URLError, ConnectionError, OSError): + pass + time.sleep(HEALTH_POLL_INTERVAL) + + raise RuntimeError(f"vLLM did not become healthy within {STARTUP_TIMEOUT}s") + + +def _forward_signal(signum, _frame): + logging.info("Received signal %s, shutting down vLLM", signum) + if vllm_process and vllm_process.poll() is None: + vllm_process.send_signal(signum) + sys.exit(128 + signum) + + +def main() -> None: + global vllm_process + + apply_local_model_args() + + if not (os.getenv("MODEL_NAME") or os.getenv("VLLM_CONFIG_FILE") or os.getenv("MODEL")): + logging.warning("MODEL_NAME is not set; `vllm serve` will fail without a --model argument") + + for sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(sig, _forward_signal) + + vllm_process = start_vllm() + startup_error = None + try: + wait_for_vllm(vllm_process) + except RuntimeError as e: + logging.error("%s", e) + stop_vllm(vllm_process) + startup_error = startup_errors.classify("".join(recent_output), model=os.getenv("MODEL_NAME")) + if startup_error is None: + # Nothing recognisable, so let the platform restart us: a failed + # download or a bad host is worth another attempt. + sys.exit(1) + # A restart cannot fix this one. Exiting would crash-loop the worker, + # paying for the download and the load on every attempt and showing the + # user a traceback instead of a cause, so stay up and answer jobs with it. + logging.error("vLLM cannot start on this configuration; answering jobs with the cause: %s", startup_error) + + # Import here (not at module import time) so the RunPod SDK and handler + # start only after the backend is confirmed healthy (or confirmed dead). + import handler as proxy_handler + import runpod + + proxy_handler.vllm_process = vllm_process + proxy_handler.startup_error = startup_error + + max_concurrency = int(os.getenv("MAX_CONCURRENCY", "30")) + runpod.serverless.start( + { + "handler": proxy_handler.handler, + "concurrency_modifier": lambda _current: max_concurrency, + "return_aggregate_stream": True, + } + ) + + +if __name__ == "__main__": + main() diff --git a/src/startup_errors.py b/src/startup_errors.py new file mode 100644 index 0000000..6b9fd1e --- /dev/null +++ b/src/startup_errors.py @@ -0,0 +1,118 @@ +"""Turn a failed `vllm serve` start into a sentence the user can act on. + +vLLM reports its problems as Python tracebacks a hundred lines long, and the +worker only sees them as stdout. When a known one goes by we would rather answer +the next job with the cause and the fix than let the container exit and be +restarted into the same wall: a crash-loop pays for the download and the model +load on every attempt and the console shows nothing but silent restarts. + +Only failures that a restart cannot fix belong here. Anything unrecognised is +left alone (``classify`` returns None) so the platform retries it, which is the +right move for a flaky download or a host that lost its GPU. +""" + +import re +from typing import Optional + +# --- GPU memory ------------------------------------------------------------- +# torch prints both the request and the card's capacity; the capacity is the +# useful half, since the request is whatever happened to be next. +_OOM = re.compile(r"torch\.OutOfMemoryError|CUDA out of memory", re.I) +_OOM_CAPACITY = re.compile(r"total capacity of ([\d.]+) GiB", re.I) +# vllm/v1/core/kv_cache_utils.py: the weights fit but nothing is left for the +# KV cache, or not enough is left to hold a single request of max_model_len. +_NO_KV_MEMORY = re.compile(r"No available memory for the cache blocks", re.I) +_KV_TOO_SMALL = re.compile( + r"larger than the available KV cache memory" + r"|larger than the maximum number of tokens that can be stored in KV cache", + re.I, +) +_KV_ESTIMATED_LEN = re.compile(r"estimated maximum model length is (\d+)", re.I) + +# --- Configuration ---------------------------------------------------------- +# vllm/config/model.py: MAX_MODEL_LEN above what the model's config allows. +_MAX_LEN_EXCEEDS_MODEL = re.compile( + r"User-specified max_model_len \((\d+)\) is greater than the derived max_model_len", re.I +) +# argparse: exit code 2 before the engine even starts. +_BAD_ARGS = re.compile(r"(?:vllm serve|vllm): error: (.+)") + +# --- Model access ----------------------------------------------------------- +_GATED = re.compile(r"GatedRepoError|401 Client Error", re.I) +_NOT_FOUND = re.compile(r"RepositoryNotFoundError|404 Client Error", re.I) +_UNSUPPORTED_ARCH = re.compile(r"Model architectures \[.*?\] (?:are not supported|failed to be inspected)", re.I) + +# --- Disk ------------------------------------------------------------------- +_NO_SPACE = re.compile(r"No space left on device|ENOSPC|errno 28", re.I) + + +def classify(output: str, model: Optional[str] = None) -> Optional[str]: + """One actionable message for a known fatal failure, or None to let it retry.""" + named = model or "The model" + + # Memory problems first: an OOM traceback is often surrounded by secondary + # errors (engine core died, connection reset) that would match nothing. + if _OOM.search(output) or _NO_KV_MEMORY.search(output) or _KV_TOO_SMALL.search(output): + capacity = _OOM_CAPACITY.search(output) + card = f" This GPU has {capacity.group(1)} GiB." if capacity else "" + estimated = _KV_ESTIMATED_LEN.search(output) + hint = ( + f" vLLM estimates this GPU can serve a context of about {estimated.group(1)} tokens." + if estimated + else "" + ) + return ( + f"{named} ran out of GPU memory during startup.{card}{hint} " + f"Lower MAX_MODEL_LEN or MAX_NUM_SEQS, set ENFORCE_EAGER=true to skip " + f"CUDA graph capture, use a quantized checkpoint, or redeploy on a " + f"larger GPU (or more GPUs with TENSOR_PARALLEL_SIZE). If the model " + f"loaded but the KV cache did not fit, raising GPU_MEMORY_UTILIZATION " + f"a little (default 0.9) can also help." + ) + + match = _MAX_LEN_EXCEEDS_MODEL.search(output) + if match: + return ( + f"MAX_MODEL_LEN={match.group(1)} is larger than the context length " + f"{named} declares in its config.json. Lower MAX_MODEL_LEN (or unset it " + f"to use the model's own limit). To override the limit anyway, set " + f"VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 and expect degraded output past it." + ) + + match = _BAD_ARGS.search(output) + if match: + return ( + f"vLLM rejected its command line: {match.group(1).strip()} " + f"Check the environment variables that map to vLLM flags and " + f"VLLM_EXTRA_ARGS; run `vllm serve --help` in the image for the " + f"accepted flags and values." + ) + + if _GATED.search(output): + return ( + f"{named} is gated or private on Hugging Face and the worker was not " + f"allowed to download it. Set HF_TOKEN to a token whose account has " + f"accepted the model's license, or pick a public model." + ) + + if _NOT_FOUND.search(output): + return ( + f"{named} was not found on Hugging Face. Check MODEL_NAME for typos " + f"(it must be the full `org/repo` id) and MODEL_REVISION if set. Private " + f"repositories also return not-found until HF_TOKEN grants access." + ) + + if _UNSUPPORTED_ARCH.search(output): + return ( + f"{named} uses an architecture this vLLM version cannot serve. Check " + f"the vLLM supported-models list, or try a newer worker release." + ) + + if _NO_SPACE.search(output): + return ( + f"{named} ran out of disk while downloading. Increase the endpoint's " + f"container disk to comfortably exceed the size of the repository, or " + f"attach a network volume so the weights are cached there instead." + ) + + return None diff --git a/tests/requirements.txt b/tests/requirements.txt new file mode 100644 index 0000000..2742e2a --- /dev/null +++ b/tests/requirements.txt @@ -0,0 +1,4 @@ +# Tests are CPU-only and never import vllm/torch. aiohttp is needed because +# src/handler.py imports it at module load (the proxy code is not exercised). +pytest>=8,<10 +aiohttp diff --git a/tests/test_args_builder.py b/tests/test_args_builder.py new file mode 100644 index 0000000..39f3224 --- /dev/null +++ b/tests/test_args_builder.py @@ -0,0 +1,161 @@ +"""Unit tests for env-var -> `vllm serve` CLI translation (src/args_builder.py).""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from args_builder import ( # noqa: E402 + BOOL_FLAGS, + ENV_ALIASES, + PAIRED_BOOL_FLAGS, + REMOVED_FLAGS, + VALUE_FLAGS, + build_vllm_args, + flag_to_env_name, +) + + +def argv_to_dict(argv): + """Fold an argv list into {flag: value|None} (last occurrence wins, like argparse).""" + out = {} + for i, token in enumerate(argv): + if token.startswith("--"): + next_token = argv[i + 1] if i + 1 < len(argv) else None + if next_token is None or next_token.startswith("--"): + out[token] = None + else: + out[token] = next_token + return out + + +class TestGenericMapping: + def test_value_flag_from_env(self): + args = build_vllm_args({"MAX_MODEL_LEN": "4096", "GPU_MEMORY_UTILIZATION": "0.9"}) + assert ["--max-model-len", "4096"] == args[args.index("--max-model-len"):][:2] + assert ["--gpu-memory-utilization", "0.9"] == args[args.index("--gpu-memory-utilization"):][:2] + + def test_paired_bool_flag(self): + # vLLM booleans are `--flag/--no-flag` pairs that take no value; a + # falsy env value must emit the `--no-` form, never `--flag false`. + args = build_vllm_args({"ENFORCE_EAGER": "true", "TRUST_REMOTE_CODE": "false"}) + assert "--enforce-eager" in args + assert "true" not in args + assert "--trust-remote-code" not in args + assert "--no-trust-remote-code" in args + + def test_paired_bool_flag_invalid_value(self): + with pytest.raises(ValueError, match="invalid boolean value"): + build_vllm_args({"ENFORCE_EAGER": "maybe"}) + + def test_store_true_flag(self): + flag = next(iter(BOOL_FLAGS)) + env_name = flag_to_env_name(flag) + assert flag in build_vllm_args({env_name: "true"}) + assert flag not in build_vllm_args({env_name: "false"}) + + def test_json_values_pass_verbatim(self): + blob = '{"model": "org/draft", "num_speculative_tokens": 3}' + args = argv_to_dict(build_vllm_args({"SPECULATIVE_CONFIG": blob})) + assert args["--speculative-config"] == blob + + def test_unknown_env_vars_ignored(self): + args = build_vllm_args({"NOT_A_VLLM_FLAG": "x", "HOME": "/root", "RUNPOD_POD_ID": "abc"}) + assert args == [] + + def test_empty_values_skipped(self): + assert "--max-model-len" not in build_vllm_args({"MAX_MODEL_LEN": " "}) + + def test_reserved_flags_not_scanned(self): + # HOST/PORT/API_KEY/MODEL are owned by the wrapper (or pinned in main.py). + args = build_vllm_args({"HOST": "0.0.0.0", "PORT": "9000", "API_KEY": "secret", "MODEL": "org/m"}) + assert "--host" not in args + assert "--port" not in args + assert "--api-key" not in args + assert "--model" not in args + + +class TestAliases: + def test_model_name_alias(self): + assert ["--model", "org/model"] == build_vllm_args({"MODEL_NAME": "org/model"}) + + def test_aliases_map_to_real_flags(self): + for env_name, flag in ENV_ALIASES.items(): + assert env_name not in {flag_to_env_name(flag)}, f"{env_name} should not be needed as alias" + args = argv_to_dict(build_vllm_args({env_name: "v"})) + assert flag in args, f"{env_name} should produce {flag}" + + def test_alias_wins_over_generic_scan(self): + # SERVED_MODEL_NAME scans generically; the override alias must win. + args = argv_to_dict( + build_vllm_args({"SERVED_MODEL_NAME": "base", "OPENAI_SERVED_MODEL_NAME_OVERRIDE": "override"}) + ) + assert args["--served-model-name"] == "override" + + +class TestExtraArgs: + def test_extra_args_appended_last(self): + args = build_vllm_args({"MAX_MODEL_LEN": "4096", "VLLM_EXTRA_ARGS": "--max-model-len 2048 --new-flag x"}) + assert args[-4:] == ["--max-model-len", "2048", "--new-flag", "x"] + # argparse last-occurrence-wins: the extra args entry must be the effective one + assert argv_to_dict(args)["--max-model-len"] == "2048" + + def test_extra_args_quoting(self): + args = build_vllm_args({"VLLM_EXTRA_ARGS": "--override-generation-config '{\"max_new_tokens\": 512}'"}) + assert args == ["--override-generation-config", '{"max_new_tokens": 512}'] + + +class TestRemovedFlags: + def test_removed_flags_warn_and_emit_nothing(self, caplog): + # Flags deleted upstream must not crash startup; they warn and emit nothing. + args = build_vllm_args({"SWAP_SPACE": "4", "ROPE_SCALING": '{"factor": 2.0}', "TASK": "generate"}) + assert args == [] + assert "SWAP_SPACE" in caplog.text + assert "ROPE_SCALING" in caplog.text + assert "TASK" in caplog.text + + +class TestZeroMeansUnset: + def test_zero_max_num_batched_tokens_is_omitted(self): + # 0 is rejected by vLLM's SchedulerConfig (ge=1) but templates carry it; + # omitting the flag lets vLLM size the batch for the GPU. + assert "--max-num-batched-tokens" not in build_vllm_args({"MAX_NUM_BATCHED_TOKENS": "0"}) + + def test_nonzero_value_passes_through(self): + args = argv_to_dict(build_vllm_args({"MAX_NUM_BATCHED_TOKENS": "8192"})) + assert args["--max-num-batched-tokens"] == "8192" + + def test_zero_not_dropped_for_other_flags(self): + # SEED=0 et al. are legitimate values; only ZERO_MEANS_UNSET flags drop 0. + args = argv_to_dict(build_vllm_args({"SEED": "0"})) + assert args["--seed"] == "0" + + +class TestAllowlistSanity: + def test_expected_flags_present(self): + # Regression: these env vars back the e2e configs and docs. + expected_value = [ + "--max-model-len", "--tensor-parallel-size", "--gpu-memory-utilization", + "--dtype", "--quantization", "--kv-cache-dtype", "--speculative-config", + "--tool-call-parser", "--reasoning-parser", + ] + for flag in expected_value: + assert flag in VALUE_FLAGS, flag + expected_bool = ["--enforce-eager", "--enable-prefix-caching", "--enable-lora", "--trust-remote-code"] + for flag in expected_bool: + assert flag in PAIRED_BOOL_FLAGS, flag + + def test_flag_categories_do_not_overlap(self): + assert not (VALUE_FLAGS & PAIRED_BOOL_FLAGS) + assert not (VALUE_FLAGS & BOOL_FLAGS) + assert not (PAIRED_BOOL_FLAGS & BOOL_FLAGS) + assert not (set(REMOVED_FLAGS) & (VALUE_FLAGS | PAIRED_BOOL_FLAGS | BOOL_FLAGS)) + + def test_every_flag_has_valid_env_name(self): + all_flags = VALUE_FLAGS | BOOL_FLAGS | PAIRED_BOOL_FLAGS | set(REMOVED_FLAGS) + for flag in all_flags: + assert flag.startswith("--") + name = flag_to_env_name(flag) + assert name == name.upper() and not name.startswith("_") diff --git a/tests/test_handler.py b/tests/test_handler.py new file mode 100644 index 0000000..b79183d --- /dev/null +++ b/tests/test_handler.py @@ -0,0 +1,103 @@ +"""Unit tests for job-input normalization (src/handler.py). No HTTP is made.""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from handler import _normalize_job_input # noqa: E402 + + +class TestOpenAIPassthrough: + def test_openai_input_posts_chat_by_default(self): + route, method, body = _normalize_job_input({"openai_input": {"messages": []}}) + assert (route, method) == ("/v1/chat/completions", "POST") + assert body == {"messages": []} + + def test_bare_openai_route_is_get(self): + assert _normalize_job_input({"openai_route": "/v1/models"}) == ("/v1/models", "GET", None) + + def test_empty_openai_input_falls_back_to_get(self): + # A RunPod GET on /openai/* must not turn into a POST (405 on /v1/models). + assert _normalize_job_input({"openai_route": "/v1/models", "openai_input": {}}) == ( + "/v1/models", "GET", None) + + +class TestGenericRoute: + def test_route_without_body_defaults_to_get(self): + assert _normalize_job_input({"route": "/v1/models"}) == ("/v1/models", "GET", None) + + def test_route_with_body_defaults_to_post(self): + route, method, body = _normalize_job_input({"route": "/v1/completions", "body": {"prompt": "hi"}}) + assert (route, method, body) == ("/v1/completions", "POST", {"prompt": "hi"}) + + def test_explicit_method_wins(self): + route, method, _ = _normalize_job_input({"route": "/v1/models", "method": "post"}) + assert method == "POST" + + +class TestLegacyShorthand: + def test_messages_maps_to_chat_completions(self): + route, method, body = _normalize_job_input({"messages": [{"role": "user", "content": "hi"}]}) + assert (route, method) == ("/v1/chat/completions", "POST") + assert body["messages"][0]["content"] == "hi" + assert body["stream"] is False + + def test_prompt_maps_to_completions(self): + route, method, body = _normalize_job_input({"prompt": "hi", "sampling_params": {"temperature": 0.5}}) + assert (route, method) == ("/v1/completions", "POST") + assert body["prompt"] == "hi" + assert body["temperature"] == 0.5 + + def test_empty_input_raises(self): + with pytest.raises(ValueError): + _normalize_job_input({}) + + +def _collect(job: dict) -> list: + """Drain the async-generator handler synchronously.""" + import asyncio + + import handler as handler_module + + async def run(): + return [item async for item in handler_module.handler(job)] + + return asyncio.run(run()) + + +class TestStartupError: + """main.py sets handler.startup_error when vLLM died for a reason a restart cannot fix.""" + + def test_answers_every_job_with_the_cause(self, monkeypatch): + import handler as handler_module + + monkeypatch.setattr(handler_module, "startup_error", "Ran out of GPU memory.") + + outputs = _collect({"input": {"prompt": "x"}}) + + assert outputs == [{"error": {"message": "Ran out of GPU memory.", "type": "startup_error", "code": None}}] + + def test_wins_over_a_dead_process_and_bad_input(self, monkeypatch): + import handler as handler_module + + monkeypatch.setattr(handler_module, "startup_error", "Ran out of GPU memory.") + # Either of these would otherwise be reported instead, and both say less. + monkeypatch.setattr(handler_module, "_is_vllm_alive", lambda: False) + + outputs = _collect({"input": {}}) + + assert "GPU memory" in outputs[0]["error"]["message"] + + def test_a_healthy_start_leaves_jobs_alone(self, monkeypatch): + import handler as handler_module + + monkeypatch.setattr(handler_module, "startup_error", None) + monkeypatch.setattr(handler_module, "_is_vllm_alive", lambda: False) + + outputs = _collect({"input": {"prompt": "x"}}) + + # Falls through to the normal liveness check rather than short-circuiting. + assert "not running" in outputs[0]["error"]["message"] diff --git a/tests/test_startup_errors.py b/tests/test_startup_errors.py new file mode 100644 index 0000000..3f31deb --- /dev/null +++ b/tests/test_startup_errors.py @@ -0,0 +1,128 @@ +"""Which vLLM startup failures are worth answering, and which are worth a restart.""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from startup_errors import classify # noqa: E402 + +TORCH_OOM = """ +(EngineCore_DP0 pid=245) ERROR 08-19 04:02:11 [core.py:1346] EngineCore failed to start. +(EngineCore_DP0 pid=245) ERROR 08-19 04:02:11 [core.py:1346] torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 108.00 MiB. +GPU 0 has a total capacity of 19.57 GiB of which 105.81 MiB is free. +""" +NO_KV_MEMORY = ( + "ValueError: No available memory for the cache blocks. Try increasing " + "`gpu_memory_utilization` when initializing the engine." +) +KV_TOO_SMALL = ( + "ValueError: To serve at least one request with the model's max seq len (131072), " + "(16.00 GiB KV cache is needed, which is larger than the available KV cache memory " + "(4.21 GiB). Based on the available memory, the estimated maximum model length is 34480. " + "Try increasing `gpu_memory_utilization` or decreasing `max_model_len` when initializing the engine." +) +KV_TOO_SMALL_OLD = ( + "ValueError: The model's max seq len (131072) is larger than the maximum number of " + "tokens that can be stored in KV cache (34480)." +) +MAX_LEN_EXCEEDS_MODEL = ( + "ValueError: User-specified max_model_len (65536) is greater than the derived " + "max_model_len (max_position_embeddings=32768 or model_max_length=None in model's config.json)." +) +BAD_ARGS = "vllm serve: error: unrecognized arguments: --foo-bar 3" +BAD_VALUE = "vllm serve: error: argument --max-model-len: invalid int value: 'lots'" +GATED = ( + "huggingface_hub.errors.GatedRepoError: 401 Client Error. Cannot access gated repo " + "for url https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/resolve/main/config.json." +) +NOT_FOUND = ( + "huggingface_hub.errors.RepositoryNotFoundError: 404 Client Error. Repository Not Found " + "for url https://huggingface.co/org/nope/resolve/main/config.json." +) +UNSUPPORTED = "ValueError: Model architectures ['FooForCausalLM'] are not supported for now." +NO_SPACE = "OSError: [Errno 28] No space left on device" + + +class TestOutOfMemory: + def test_torch_oom_names_the_model_and_the_card(self): + message = classify(TORCH_OOM, model="Qwen/Qwen3-30B-A3B") + + assert message.startswith("Qwen/Qwen3-30B-A3B ran out of GPU memory") + assert "19.57 GiB" in message + + def test_advice_covers_the_knobs_the_worker_exposes(self): + message = classify(TORCH_OOM) + + for knob in ("MAX_MODEL_LEN", "ENFORCE_EAGER", "GPU_MEMORY_UTILIZATION", "TENSOR_PARALLEL_SIZE"): + assert knob in message + + def test_no_kv_cache_memory_is_out_of_memory(self): + assert "ran out of GPU memory" in classify(NO_KV_MEMORY) + + def test_kv_cache_too_small_passes_on_the_estimated_context(self): + message = classify(KV_TOO_SMALL) + + assert "ran out of GPU memory" in message + assert "34480 tokens" in message + + def test_older_kv_cache_wording_is_recognised(self): + assert classify(KV_TOO_SMALL_OLD) is not None + + def test_oom_wins_over_surrounding_noise(self): + assert "GPU memory" in classify(f"Connection reset by peer\n{TORCH_OOM}\nEngineCore died") + + +class TestConfiguration: + def test_max_model_len_above_the_model_limit_names_the_value(self): + message = classify(MAX_LEN_EXCEEDS_MODEL, model="org/model") + + assert "MAX_MODEL_LEN=65536" in message + assert "VLLM_ALLOW_LONG_MAX_MODEL_LEN" in message + + @pytest.mark.parametrize("output", [BAD_ARGS, BAD_VALUE]) + def test_argparse_rejections_quote_the_reason(self, output): + message = classify(output) + + assert message.startswith("vLLM rejected its command line") + assert output.split("error: ", 1)[1] in message + assert "VLLM_EXTRA_ARGS" in message + + +class TestModelAccess: + def test_gated_repo_points_at_hf_token(self): + assert "HF_TOKEN" in classify(GATED, model="meta-llama/Llama-3.1-8B-Instruct") + + def test_missing_repo_points_at_model_name(self): + assert "MODEL_NAME" in classify(NOT_FOUND) + + def test_unsupported_architecture(self): + assert "architecture" in classify(UNSUPPORTED) + + +def test_out_of_disk_points_at_container_disk(): + assert "container disk" in classify(NO_SPACE) + + +def test_falls_back_to_a_generic_subject(): + assert classify(NO_SPACE).startswith("The model") + + +class TestUnknownFailures: + @pytest.mark.parametrize( + "output", + [ + "", + "Connection reset by peer while downloading", + "RuntimeError: something nobody has seen before", + "vLLM serve exited during startup with code 1", + # Only vLLM's own argparse errors count, not any line with "error:". + "INFO 08-19 [launcher.py] error: none", + ], + ) + def test_are_left_for_the_platform_to_retry(self, output): + # Answering these forever would turn a flaky download into a dead + # endpoint; a restart is the right response. + assert classify(output) is None From 150e006a8491803bae0c3ba80336eefd8279ae3b Mon Sep 17 00:00:00 2001 From: danstotts-ops Date: Tue, 29 Sep 2026 22:36:19 -0700 Subject: [PATCH 2/2] ci: verify Holo4 worker bundle and source --- .github/workflows/verify.yml | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 .github/workflows/verify.yml diff --git a/.github/workflows/verify.yml b/.github/workflows/verify.yml new file mode 100644 index 0000000..c296788 --- /dev/null +++ b/.github/workflows/verify.yml @@ -0,0 +1,21 @@ +name: Verify Holo4 worker + +on: + pull_request: + push: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - run: python -m pip install pytest aiohttp + - run: python -m pytest -q tests + - run: python -m compileall -q src handler.py + - run: cmp handler.py src/handler.py + - run: python -m json.tool .runpod/hub.json > /dev/null + - run: python -m json.tool .runpod/tests.json > /dev/null