Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .runpod/hub.json
Original file line number Diff line number Diff line change
Expand Up @@ -710,7 +710,7 @@
"input": {
"name": "Base Path",
"type": "string",
"description": "Storage directory for Huggingface cache and model",
"description": "Set when the image is built; changing it here moves neither the Hugging Face cache nor the vLLM compile cache (set HF_HUB_CACHE / VLLM_CACHE_ROOT instead)",
"default": "/runpod-volume",
"advanced": true
}
Expand Down
8 changes: 8 additions & 0 deletions Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,14 @@ ENV MODEL_NAME=$MODEL_NAME \
HF_HOME="${BASE_PATH}/huggingface-cache/hub" \
HUGGINGFACE_HUB_CACHE="${BASE_PATH}/huggingface-cache/hub" \
HF_DATASETS_CACHE="${BASE_PATH}/huggingface-cache/datasets" \
# Same for vLLM's torch.compile cache: otherwise every cold start recompiles
# (~30 s for Qwen3-8B on an RTX 4090). With this image's torch (AOT + mega
# artifact on by default) the cache key includes the GPU name and the
# torch/Triton/CUDA builds, so GPU types sharing a volume get separate
# entries. Without a volume BASE_PATH is a plain container directory, as
# before. src/compile_cache.py falls back to vLLM's default root if the
# volume is full or read-only.
VLLM_CACHE_ROOT="${BASE_PATH}/vllm-cache" \
HF_HUB_ENABLE_HF_TRANSFER=0 \
TOKENIZERS_PARALLELISM=false

Expand Down
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ To build an image with the model baked in, you must specify the following docker
- **Optional**
- `MODEL_REVISION`: Model revision to load (default: `main`).
- `VLLM_VERSION`: Tag of the official [`vllm/vllm-openai`](https://hub.docker.com/r/vllm/vllm-openai) base image to use (default: `v0.23.0`).
- `BASE_PATH`: Storage directory where huggingface cache and model will be located. (default: `/runpod-volume`, which will utilize network storage if you attach it or create a local directory within the image if you don't. If your intention is to bake the model into the image, you should set this to something like `/models` to make sure there are no issues if you were to accidentally attach network storage.)
- `BASE_PATH`: Storage directory where huggingface cache and model will be located, along with vLLM's torch.compile cache (`$BASE_PATH/vllm-cache`). (default: `/runpod-volume`, which will utilize network storage if you attach it or create a local directory within the image if you don't. If your intention is to bake the model into the image, you should set this to something like `/models` to make sure there are no issues if you were to accidentally attach network storage.)
- `QUANTIZATION`
- `TOKENIZER_NAME`: Tokenizer repository if you would like to use a different tokenizer than the one that comes with the model. (default: `None`, which uses the model's tokenizer)
- `TOKENIZER_REVISION`: Tokenizer revision to load (default: `main`).
Expand Down
15 changes: 12 additions & 3 deletions docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ These are consumed by the wrapper itself, not passed to vLLM:
| --------------------- | ------- | ---------------------------------------------------------------------- |
| `MODEL_NAME` | — | Required. HF repo id or local path of the model. |
| `HF_TOKEN` | — | Hugging Face token for gated/private models. Passed to vLLM via the environment, never as `--hf-token`. |
| `BASE_PATH` | `/runpod-volume` | Root for the HF cache (persists on a network volume). |
| `BASE_PATH` | `/runpod-volume` | Build arg, not read at runtime: the image derives `HF_HOME` and `VLLM_CACHE_ROOT` (`$BASE_PATH/vllm-cache`) from it at build time, so both caches persist on a network volume. Setting it on an endpoint moves neither; set `HF_HUB_CACHE` (model downloads) and `VLLM_CACHE_ROOT` (compile cache) directly. |
| `MAX_CONCURRENCY` | `30` | Max concurrent jobs per worker (RunPod concurrency modifier). vLLM queues internally beyond this. |
| `VLLM_PORT` | `8000` | Loopback port the internal `vllm serve` binds to. |
| `VLLM_STARTUP_TIMEOUT`| `1200` | Seconds to wait for vLLM `/health` before failing the worker. |
Expand All @@ -79,7 +79,16 @@ These are consumed by the wrapper itself, not passed to vLLM:

vLLM-native env vars (e.g. `VLLM_USE_DEEP_GEMM`, `PYTORCH_ALLOC_CONF`) are **not**
converted to flags — they reach the `vllm serve` subprocess directly through the
environment, exactly as upstream supports them.
environment, exactly as upstream supports them. This worker's Dockerfile sets one
of them: `VLLM_CACHE_ROOT=$BASE_PATH/vllm-cache`, so the torch.compile cache survives
cold starts on a network volume. If that folder isn't writable or the volume has less
than 1 GiB free, the worker compiles into vLLM's default root (in the container) for
that start instead, and if the volume still runs out during the start, it relaunches
once that way (see `src/compile_cache.py` and `src/main.py`). vLLM never prunes the
cache: each vLLM version, model config and GPU type adds an entry.
`$BASE_PATH/vllm-cache` is safe to delete whenever no worker is starting (e.g. with the
endpoint scaled to zero); vLLM rebuilds it. Set `VLLM_CACHE_ROOT=/root/.cache/vllm` to
keep the cache in the container.

## Speculative decoding

Expand Down Expand Up @@ -113,7 +122,7 @@ LORA_MODULES='[{"name":"my-adapter","path":"org/adapter-repo"}]'
| `TOKENIZER_NAME` | same as model | Optional separate tokenizer repo. |
| `TOKENIZER_REVISION`| model revision | Tokenizer revision. |
| `QUANTIZATION` | — | Recorded for baked models. |
| `BASE_PATH` | `/runpod-volume` | HF cache location baked into the image env. |
| `BASE_PATH` | `/runpod-volume` | HF cache and vLLM compile cache (`VLLM_CACHE_ROOT`) location baked into the image env. |

The build secret `HF_TOKEN` (`--secret id=HF_TOKEN`) is used during the bake step for
gated models and never lands in the image.
Expand Down
14 changes: 14 additions & 0 deletions docs/conventions.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ loopback HTTP.
Container start:
ENTRYPOINT python3 /src/main.py
├─ (optional) reads /local_model_args.json for baked-in models
├─ compile_cache.ensure_usable() → keep VLLM_CACHE_ROOT on the volume only
│ if it is writable with room (else vLLM's default root)
├─ args_builder.build_vllm_args() → env vars → CLI flags
├─ spawns `vllm serve --host 127.0.0.1 --port $VLLM_PORT ...`
├─ polls GET /health until ready (exit non-zero if vLLM dies or times out,
Expand All @@ -47,6 +49,10 @@ boundary or the CLI; that is what makes vLLM upgrades a one-line `VLLM_VERSION`
- `src/args_builder.py`: pure-Python env var → CLI flag translation (allowlist of
`vllm serve` flags + legacy aliases + `VLLM_EXTRA_ARGS` passthrough). No third-party
imports, so it is fully unit-testable on CPU runners.
- `src/compile_cache.py`: checks before launch that `VLLM_CACHE_ROOT` (on the network
volume) is writable with room to spare, else falls back to vLLM's default root (in
the container) for that start; `main.py` also relaunches once that way if the
volume runs out of disk mid-start. Pure Python.
- `src/startup_errors.py`: regexes over the tail of vLLM's output that turn a known
fatal startup failure (CUDA OOM, KV cache too small for `MAX_MODEL_LEN`, rejected
flag, gated/missing model, out of disk) into one actionable sentence. Pure Python.
Expand Down Expand Up @@ -80,6 +86,10 @@ VLLM_EXTRA_ARGS > env aliases (MODEL_NAME, ...) > env flag scan
- **Configuration**: Entirely via environment variables
- **Model Loading**: Downloaded at container start by `vllm serve` into the HF cache on
the network volume (`BASE_PATH`, default `/runpod-volume`)
- **Compile cache**: vLLM's torch.compile cache (`VLLM_CACHE_ROOT`) sits on the same
volume at `$BASE_PATH/vllm-cache`, so only the first start for each model config, GPU
type and image version compiles; `src/compile_cache.py` falls back to the container
when the volume is full or read-only

### Option 2: Baked Model Images

Expand All @@ -99,6 +109,7 @@ src/
├── main.py # Entrypoint: vLLM subprocess + RunPod loop lifecycle
├── args_builder.py # env vars → `vllm serve` CLI flags (pure Python)
├── startup_errors.py # fatal startup failure → actionable message (pure Python)
├── compile_cache.py # full/read-only volume → compile cache in the container
├── handler.py # RunPod handler: aiohttp proxy to the vLLM server
└── download_model.py # Build-time model download (Option 2)
```
Expand Down Expand Up @@ -202,6 +213,9 @@ workers at container start.
- `tests/test_startup_errors.py` pins which vLLM failure messages are answered and
which are left for a restart; `tests/test_handler.py` checks the `startup_error`
short-circuit in the handler.
- `tests/test_dockerfile_env.py` keeps the HF cache and vLLM's compile cache under
`BASE_PATH` in the image env, so both persist on a network volume;
`tests/test_compile_cache.py` checks the full/read-only volume fallback.

### 2. **Local Smoke Testing**

Expand Down
112 changes: 112 additions & 0 deletions src/compile_cache.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
"""Keep vLLM's compile cache on the network volume only when the volume can take it.

The image sets VLLM_CACHE_ROOT=$BASE_PATH/vllm-cache so torch.compile output
survives cold starts on a network volume. vLLM does not handle write errors on
most of its compile-cache writes, so a full or read-only volume would fail
startup where compiling into vLLM's default root used to work. Two guards:

- ensure_usable(), before launch: if the directory isn't writable or has less
than MIN_FREE_BYTES free, use vLLM's default root for this start.
- fall_back(), from main.py when vLLM still runs out of disk mid-start (the
check runs before the weights download onto the same volume): relaunch once
with vLLM's default root.

Pure Python, no vLLM import.
"""

import logging
import os
import shutil
import tempfile

# Free space required on the volume before compiling onto it. A heuristic, not
# a measurement: running out later in the start is handled by the relaunch.
MIN_FREE_BYTES = 1 << 30


def vllm_default_root(env=os.environ) -> str:
"""Where vLLM puts its cache when VLLM_CACHE_ROOT is unset (vllm/envs.py)."""
base = env.get("XDG_CACHE_HOME") or os.path.expanduser("~/.cache")
return os.path.join(base, "vllm")


def _configured_root(env) -> str | None:
"""VLLM_CACHE_ROOT as vLLM resolves it, or None if unset or already the default."""
root = env.get("VLLM_CACHE_ROOT")
if not root:
return None
root = os.path.expanduser(root)
if os.path.realpath(root) == os.path.realpath(vllm_default_root(env)):
return None
return root


def _device(path: str) -> int:
"""st_dev of the path, or of its nearest existing ancestor."""
while not os.path.exists(path):
parent = os.path.dirname(path)
if parent == path:
break
path = parent
return os.stat(path).st_dev


def fall_back(reason: str, env=os.environ, out_of_space: bool = False) -> bool:
"""Point VLLM_CACHE_ROOT at vLLM's default root.

Returns False (and changes nothing) when that can't help: VLLM_CACHE_ROOT is
unset or already the default, or, for ``out_of_space``, the default is on
the same filesystem (no volume attached, so the same disk would run out).
"""
root = _configured_root(env)
fallback = vllm_default_root(env)
if root is None:
return False
if out_of_space:
try:
if _device(root) == _device(fallback):
return False
except OSError:
pass
logging.warning(
"VLLM_CACHE_ROOT=%s is %s; compiling into %s for this start instead. "
"The vllm-cache folder is safe to delete while no worker is starting "
"(e.g. endpoint scaled to zero); vLLM rebuilds it.",
root,
reason,
fallback,
)
env["VLLM_CACHE_ROOT"] = fallback
return True


def ensure_usable(env=os.environ, min_free_bytes: int | None = None) -> str | None:
"""Fall back to vLLM's default root if VLLM_CACHE_ROOT can't take writes.

Returns why it fell back, or None when VLLM_CACHE_ROOT is unset, already the
default, or writable with at least ``min_free_bytes`` (default
MIN_FREE_BYTES) free.
"""
if min_free_bytes is None:
min_free_bytes = MIN_FREE_BYTES
root = _configured_root(env)
if root is None:
return None
try:
os.makedirs(root, exist_ok=True)
# A zero-byte file still fits on a full volume or under an exhausted
# quota (statvfs may report the storage pool, not the quota); one
# written and synced block does not.
with tempfile.NamedTemporaryFile(dir=root, prefix=".write-check-") as probe:
probe.write(b"\0" * 65536)
probe.flush()
os.fsync(probe.fileno())
free = shutil.disk_usage(root).free
except OSError as e:
reason = f"not writable ({e.strerror or e})"
else:
if free >= min_free_bytes:
return None
reason = f"nearly full ({free >> 20} MiB free)"
return reason if fall_back(reason, env, out_of_space=True) else None
return reason if fall_back(reason, env) else None
36 changes: 29 additions & 7 deletions src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,14 @@
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. Two startup failures get a single relaunch first, because
startup_errors.py. Three startup failures get a single relaunch first, because
retrying them can genuinely succeed: a Hugging Face revision that no longer
exists (vLLM pins refs to commit hashes since v0.28, so a force-pushed repo
invalidates the pin) is retried against the repo's current state, and a
startup OOM is retried with a smaller memory footprint (CUDA graphs disabled
and a reduced MAX_NUM_BATCHED_TOKENS — the two settings whose defaults grew at
invalidates the pin) is retried against the repo's current state, running out
of disk with vLLM's compile cache on the network volume is retried with the
cache in the container (see compile_cache.py), and a startup OOM is retried
with a smaller memory footprint (CUDA graphs disabled and a reduced
MAX_NUM_BATCHED_TOKENS — the two settings whose defaults grew at
init time in recent vLLM releases). Unrecognised failures still exit non-zero
so the platform retries them.
"""
Expand All @@ -35,6 +37,7 @@
import urllib.error
import urllib.request

import compile_cache
import model_preflight
import startup_errors
from args_builder import TRUE_VALUES, build_vllm_args, redact_argv
Expand Down Expand Up @@ -222,6 +225,9 @@ def main() -> None:
global vllm_process

apply_local_model_args()
# The compile cache lives on the network volume (Dockerfile); a full or
# read-only volume would fail the start, so compile in the container instead.
compile_cache.ensure_usable()

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")
Expand All @@ -239,12 +245,16 @@ def main() -> None:
"Model pre-flight failed; answering jobs with the cause instead of starting vLLM: %s",
startup_error,
)
# At most three launches: each of the two recoverable startup failures gets
# At most four launches: each of the three recoverable startup failures gets
# one relaunch — a vanished Hugging Face revision (vLLM re-resolves the pin
# on every launch; the pre-flight catches a *configured* revision that never
# existed, this catches one gone by load time or resolved by vLLM itself)
# and a startup OOM (a smaller footprint can fit where the defaults did not).
# existed, this catches one gone by load time or resolved by vLLM itself),
# running out of disk with the compile cache on the network volume (the
# volume can fill after the pre-launch check, during the weight download; a
# relaunch compiles into the container instead), and a startup OOM (a
# smaller footprint can fit where the defaults did not).
can_retry_revision = True
can_retry_disk = True
can_retry_oom = True
oom_retried = False
while startup_error is None:
Expand All @@ -267,6 +277,18 @@ def main() -> None:
drop_pinned_revisions()
recent_output.clear()
continue
if (
can_retry_disk
and startup_errors.out_of_disk(output)
and compile_cache.fall_back("out of space during startup", out_of_space=True)
):
can_retry_disk = False
logging.warning(
"vLLM ran out of disk during startup with its compile cache on "
"the network volume. Relaunching once with it in the container."
)
recent_output.clear()
continue
if (
can_retry_oom
and startup_errors.memory_shortfall(output)
Expand Down
23 changes: 19 additions & 4 deletions src/startup_errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@
_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)
_NO_SPACE = re.compile(r"No space left on device|ENOSPC|errno 28|Disk quota exceeded|EDQUOT|errno 122", re.I)


# The model-access messages are shared with model_preflight.py, which detects
Expand Down Expand Up @@ -92,6 +92,17 @@ def memory_shortfall(output: str) -> bool:
return bool(_OOM.search(output) or _NO_KV_MEMORY.search(output) or _KV_TOO_SMALL.search(output))


def out_of_disk(output: str) -> bool:
"""True when startup died for lack of disk space (or a full disk quota).

Split out from classify() because main.py relaunches once with vLLM's
compile cache moved off the network volume: the pre-launch check in
compile_cache.py runs before the weights download onto the same volume, so
the volume can still fill during the start.
"""
return bool(_NO_SPACE.search(output))


def revision_not_found(output: str) -> bool:
"""True when the failure is a Hugging Face revision that no longer exists.

Expand Down Expand Up @@ -196,9 +207,13 @@ def classify(output: str, model: Optional[str] = None, oom_retried: bool = False

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."
f"{named} ran out of disk during startup. With a network volume "
f"attached, the model download (huggingface-cache) and vLLM's compile "
f"cache (vllm-cache) are stored on it: free space there (vllm-cache is "
f"safe to delete while no worker is starting) or use a larger volume. "
f"Without one, increase the endpoint's container disk to comfortably "
f"exceed the size of the repository, or attach a network volume so the "
f"weights are cached there instead."
)

return None
Loading
Loading