fix: persist vLLM's torch.compile cache on the network volume

.runpod/hub.json

@@ -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
         }

Dockerfile

@@ -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
 

README.md

@@ -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`).

docs/configuration.md

@@ -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.           |
@@ -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
 
@@ -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.

docs/conventions.md

@@ -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,
@@ -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.
@@ -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
 
@@ -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)
 ```
@@ -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**
 

src/compile_cache.py

@@ -0,0 +1,119 @@
+"""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):
+                logging.warning(
+                    "VLLM_CACHE_ROOT=%s is %s, but the fallback %s is on the same "
+                    "disk; not moving the compile cache.",
+                    root,
+                    reason,
+                    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

src/main.py

@@ -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.
 """
@@ -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
@@ -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")
@@ -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:
@@ -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)

src/startup_errors.py

@@ -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
@@ -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.
 
@@ -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