diff --git a/.runpod/hub.json b/.runpod/hub.json index 03f3e4c7..8e3a8332 100644 --- a/.runpod/hub.json +++ b/.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 } diff --git a/Dockerfile b/Dockerfile index e07acc9a..9b31b568 100644 --- a/Dockerfile +++ b/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 diff --git a/README.md b/README.md index 32309e1c..cce51bbe 100644 --- a/README.md +++ b/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`). diff --git a/docs/configuration.md b/docs/configuration.md index b9e77554..bced49fd 100644 --- a/docs/configuration.md +++ b/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. diff --git a/docs/conventions.md b/docs/conventions.md index 8aa68d92..d0f84b40 100644 --- a/docs/conventions.md +++ b/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** diff --git a/src/compile_cache.py b/src/compile_cache.py new file mode 100644 index 00000000..bde20b39 --- /dev/null +++ b/src/compile_cache.py @@ -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 diff --git a/src/main.py b/src/main.py index ef623435..42698da0 100644 --- a/src/main.py +++ b/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) diff --git a/src/startup_errors.py b/src/startup_errors.py index 9b846557..eb2d0ca5 100644 --- a/src/startup_errors.py +++ b/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 diff --git a/tests/test_compile_cache.py b/tests/test_compile_cache.py new file mode 100644 index 00000000..bda93bec --- /dev/null +++ b/tests/test_compile_cache.py @@ -0,0 +1,254 @@ +"""The compile cache stays on the volume only when the volume can take the writes. + +A full or read-only volume must not fail the start: the worker falls back to +vLLM's default cache root (in the container), which is what every start used +before the cache moved to the volume — before launch when the check catches +it, or with one relaunch when the volume fills during the start. +""" + +import errno +import shutil +import sys +import types +from collections import namedtuple +from pathlib import Path +from unittest import mock + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +import compile_cache # noqa: E402 +import main # noqa: E402 + +DiskUsage = namedtuple("DiskUsage", "total used free") +NO_SPACE = "OSError: [Errno 28] No space left on device" + + +@pytest.fixture +def paths(tmp_path, monkeypatch): + """A volume and a container cache dir; the volume reports as its own disk.""" + volume = tmp_path / "volume" + container = tmp_path / "container-cache" + real_device = compile_cache._device + monkeypatch.setattr( + compile_cache, + "_device", + lambda p: -1 if str(p).startswith(str(volume)) else real_device(p), + ) + return types.SimpleNamespace( + root=str(volume / "vllm-cache"), + xdg=str(container), + fallback=str(container / "vllm"), + volume=volume, + ) + + +def plenty_free(monkeypatch): + monkeypatch.setattr(shutil, "disk_usage", lambda p: DiskUsage(100 << 30, 0, 100 << 30)) + + +def test_writable_volume_with_room_keeps_the_cache_there(paths, monkeypatch): + plenty_free(monkeypatch) + env = {"VLLM_CACHE_ROOT": paths.root, "XDG_CACHE_HOME": paths.xdg} + + assert compile_cache.ensure_usable(env) is None + assert env["VLLM_CACHE_ROOT"] == paths.root + assert list(Path(paths.root).iterdir()) == [] # probe file cleaned up + + +def test_unset_means_vllm_default_and_nothing_to_check(): + env = {} + + assert compile_cache.ensure_usable(env) is None + assert env == {} + + +def test_uncreatable_root_falls_back(paths): + # A path under a regular file can't be created, even as root. + paths.volume.write_text("") + env = {"VLLM_CACHE_ROOT": paths.root, "XDG_CACHE_HOME": paths.xdg} + + assert compile_cache.ensure_usable(env).startswith("not writable") + assert env["VLLM_CACHE_ROOT"] == paths.fallback + + +def test_existing_read_only_root_falls_back(paths, monkeypatch): + # makedirs(exist_ok=True) succeeds on an existing dir of a read-only + # volume; only the write probe catches it. + Path(paths.root).mkdir(parents=True) + + def read_only(*args, **kwargs): + raise OSError(errno.EROFS, "Read-only file system") + + monkeypatch.setattr(compile_cache.tempfile, "NamedTemporaryFile", read_only) + env = {"VLLM_CACHE_ROOT": paths.root, "XDG_CACHE_HOME": paths.xdg} + + assert compile_cache.ensure_usable(env) == "not writable (Read-only file system)" + assert env["VLLM_CACHE_ROOT"] == paths.fallback + + +def test_full_quota_is_caught_before_launch(paths, monkeypatch): + # statvfs may report the storage pool, not the volume's quota, and an empty + # file still fits under an exhausted quota; a synced block does not. + plenty_free(monkeypatch) + + def quota_full(fd): + raise OSError(errno.EDQUOT, "Disk quota exceeded") + + monkeypatch.setattr(compile_cache.os, "fsync", quota_full) + env = {"VLLM_CACHE_ROOT": paths.root, "XDG_CACHE_HOME": paths.xdg} + + assert compile_cache.ensure_usable(env) == "not writable (Disk quota exceeded)" + assert env["VLLM_CACHE_ROOT"] == paths.fallback + + +def test_device_lookup_failure_still_falls_back(paths, monkeypatch): + def broken(path): + raise OSError(errno.EIO, "Input/output error") + + monkeypatch.setattr(compile_cache, "_device", broken) + env = {"VLLM_CACHE_ROOT": paths.root, "XDG_CACHE_HOME": paths.xdg} + + assert compile_cache.fall_back("out of space during startup", env, out_of_space=True) + assert env["VLLM_CACHE_ROOT"] == paths.fallback + + +def test_nearly_full_volume_falls_back(paths, monkeypatch): + monkeypatch.setattr(shutil, "disk_usage", lambda p: DiskUsage(20 << 30, 0, 300 << 20)) + env = {"VLLM_CACHE_ROOT": paths.root, "XDG_CACHE_HOME": paths.xdg} + + assert compile_cache.ensure_usable(env) == "nearly full (300 MiB free)" + assert env["VLLM_CACHE_ROOT"] == paths.fallback + + +def test_nearly_full_without_a_volume_changes_nothing(tmp_path, monkeypatch): + # Root and fallback on the same disk: moving the cache can't free space. + monkeypatch.setattr(shutil, "disk_usage", lambda p: DiskUsage(20 << 30, 0, 300 << 20)) + env = {"VLLM_CACHE_ROOT": str(tmp_path / "vllm-cache"), "XDG_CACHE_HOME": str(tmp_path / "xdg")} + + assert compile_cache.ensure_usable(env) is None + assert env["VLLM_CACHE_ROOT"] == str(tmp_path / "vllm-cache") + + +def test_opt_out_is_not_probed(tmp_path): + env = {"VLLM_CACHE_ROOT": str(tmp_path / "xdg" / "vllm"), "XDG_CACHE_HOME": str(tmp_path / "xdg")} + + assert compile_cache.ensure_usable(env) is None + assert not (tmp_path / "xdg").exists() + + +def test_tilde_is_expanded_like_vllm_does(tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path / "home")) + monkeypatch.chdir(tmp_path) + plenty_free(monkeypatch) + env = {"VLLM_CACHE_ROOT": "~/vllm-cache"} + + assert compile_cache.ensure_usable(env) is None + assert (tmp_path / "home" / "vllm-cache").is_dir() + assert not (tmp_path / "~").exists() + + +def test_fallback_root_matches_vllm_default(monkeypatch): + monkeypatch.setenv("HOME", "/root") + + assert compile_cache.vllm_default_root({}) == "/root/.cache/vllm" + assert compile_cache.vllm_default_root({"XDG_CACHE_HOME": "/x"}) == "/x/vllm" + + +@pytest.fixture +def run_main(monkeypatch): + """Run main() without vLLM or the RunPod SDK; returns (cache roots per launch, handler).""" + monkeypatch.setenv("MODEL_NAME", "org/model") + monkeypatch.setattr(main.model_preflight, "check_model_access", lambda: None) + monkeypatch.setattr(main, "stop_vllm", lambda proc: None) + handler_stub = types.ModuleType("handler") + handler_stub.handler = lambda job: None + runpod_stub = types.ModuleType("runpod") + runpod_stub.serverless = types.SimpleNamespace(start=mock.MagicMock()) + monkeypatch.setitem(sys.modules, "handler", handler_stub) + monkeypatch.setitem(sys.modules, "runpod", runpod_stub) + + def drive(*attempt_outputs): + roots = [] + outputs = iter(attempt_outputs) + + def fake_start(): + roots.append(main.os.environ["VLLM_CACHE_ROOT"]) + return mock.MagicMock() + + def fake_wait(proc): + output = next(outputs) + if output is not None: + main.recent_output.clear() + main.recent_output.append(output) + raise RuntimeError("vLLM serve exited during startup with code 1") + + monkeypatch.setattr(main, "start_vllm", fake_start) + monkeypatch.setattr(main, "wait_for_vllm", fake_wait) + main.recent_output.clear() + main.main() + return roots, handler_stub + + return drive + + +def test_main_checks_the_cache_root_before_launching_vllm(paths, run_main, monkeypatch): + paths.volume.write_text("") + monkeypatch.setenv("VLLM_CACHE_ROOT", paths.root) + monkeypatch.setenv("XDG_CACHE_HOME", paths.xdg) + + roots, handler = run_main(None) + + assert roots == [paths.fallback] + assert handler.startup_error is None + + +def test_volume_filling_during_start_relaunches_once_in_the_container(paths, run_main, monkeypatch): + plenty_free(monkeypatch) # the pre-launch check passes; the download fills it later + monkeypatch.setenv("VLLM_CACHE_ROOT", paths.root) + monkeypatch.setenv("XDG_CACHE_HOME", paths.xdg) + + roots, handler = run_main(NO_SPACE, None) + + assert roots == [paths.root, paths.fallback] + assert handler.startup_error is None + + +def test_out_of_disk_without_a_volume_answers_jobs_instead_of_relaunching(tmp_path, run_main, monkeypatch): + plenty_free(monkeypatch) + monkeypatch.setenv("VLLM_CACHE_ROOT", str(tmp_path / "vllm-cache")) + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "xdg")) + + roots, handler = run_main(NO_SPACE) + + assert roots == [str(tmp_path / "vllm-cache")] + assert "vllm-cache" in handler.startup_error + + +def test_second_out_of_disk_answers_jobs(paths, run_main, monkeypatch): + plenty_free(monkeypatch) + monkeypatch.setenv("VLLM_CACHE_ROOT", paths.root) + monkeypatch.setenv("XDG_CACHE_HOME", paths.xdg) + + roots, handler = run_main(NO_SPACE, NO_SPACE) + + assert roots == [paths.root, paths.fallback] # never a third attempt + assert "vllm-cache" in handler.startup_error + + +def test_out_of_disk_is_relaunched_before_oom(paths, run_main, monkeypatch): + # The OOM relaunch sets ENFORCE_EAGER, which also turns torch.compile off; + # moving the cache keeps compile and CUDA graphs, so it must go first. + plenty_free(monkeypatch) + monkeypatch.setenv("VLLM_CACHE_ROOT", paths.root) + monkeypatch.setenv("XDG_CACHE_HOME", paths.xdg) + # setenv (not delenv) so teardown also undoes a write by the OOM path. + monkeypatch.setenv("ENFORCE_EAGER", "") + monkeypatch.setenv("MAX_NUM_BATCHED_TOKENS", "") + + roots, handler = run_main(NO_SPACE + "\ntorch.OutOfMemoryError: CUDA out of memory.", None) + + assert roots == [paths.root, paths.fallback] + assert main.os.environ["ENFORCE_EAGER"] == "" + assert handler.startup_error is None diff --git a/tests/test_dockerfile_env.py b/tests/test_dockerfile_env.py new file mode 100644 index 00000000..88162f75 --- /dev/null +++ b/tests/test_dockerfile_env.py @@ -0,0 +1,36 @@ +"""Caches that must live under BASE_PATH so they persist on a network volume. + +The RunPod network volume mounts at BASE_PATH (default /runpod-volume). Any +cache left at its default under ~/.cache is lost on every cold start, so the +image env points both the HF cache and vLLM's torch.compile cache there. +""" + +import re +import shlex +from pathlib import Path + +DOCKERFILE = Path(__file__).resolve().parent.parent / "Dockerfile" + + +def image_env() -> dict[str, str]: + """KEY=value pairs from the Dockerfile's ENV instructions (values unexpanded).""" + # Docker drops comment lines before joining continuations, so do the same. + lines = [ + line + for line in DOCKERFILE.read_text().splitlines() + if not line.lstrip().startswith("#") + ] + joined = re.sub(r"\\\n", " ", "\n".join(lines)) + env: dict[str, str] = {} + for line in joined.splitlines(): + if line.startswith("ENV "): + for pair in shlex.split(line[len("ENV "):]): + key, _, value = pair.partition("=") + env[key] = value + return env + + +def test_hf_and_vllm_caches_live_under_base_path(): + env = image_env() + assert env["HF_HOME"].startswith("${BASE_PATH}/") + assert env["VLLM_CACHE_ROOT"] == "${BASE_PATH}/vllm-cache" diff --git a/tests/test_main_oom_retry.py b/tests/test_main_oom_retry.py index 91d3eb8a..a3ba8125 100644 --- a/tests/test_main_oom_retry.py +++ b/tests/test_main_oom_retry.py @@ -38,6 +38,8 @@ def harness(monkeypatch): stub, whose startup_error records what jobs would be answered with. """ monkeypatch.setenv("MODEL_NAME", "org/model") + # main() probes VLLM_CACHE_ROOT (compile_cache.py); keep a developer's own out of it. + monkeypatch.delenv("VLLM_CACHE_ROOT", raising=False) monkeypatch.delenv("ENFORCE_EAGER", raising=False) monkeypatch.delenv("MAX_NUM_BATCHED_TOKENS", raising=False) # The merged pre-flight (model_preflight.py) would otherwise ask the real diff --git a/tests/test_main_revision_retry.py b/tests/test_main_revision_retry.py index 3079d712..625db92e 100644 --- a/tests/test_main_revision_retry.py +++ b/tests/test_main_revision_retry.py @@ -31,6 +31,8 @@ def harness(monkeypatch): stub, whose startup_error records what jobs would be answered with. """ monkeypatch.setenv("MODEL_NAME", "org/model") + # main() probes VLLM_CACHE_ROOT (compile_cache.py); keep a developer's own out of it. + monkeypatch.delenv("VLLM_CACHE_ROOT", raising=False) # The merged pre-flight (model_preflight.py) would otherwise ask the real # HF Hub about "org/model"; these tests are about what happens after it. monkeypatch.setattr(main.model_preflight, "check_model_access", lambda: None) diff --git a/tests/test_startup_errors.py b/tests/test_startup_errors.py index 5c9fb476..799a995e 100644 --- a/tests/test_startup_errors.py +++ b/tests/test_startup_errors.py @@ -7,6 +7,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) +import startup_errors # noqa: E402 from startup_errors import classify # noqa: E402 TORCH_OOM = """ @@ -106,6 +107,17 @@ def test_out_of_disk_points_at_container_disk(): assert "container disk" in classify(NO_SPACE) +def test_out_of_disk_names_the_volume_caches(): + # With a network volume, the HF cache and vLLM's compile cache live on it. + assert "vllm-cache" in classify(NO_SPACE) + assert "network volume" in classify(NO_SPACE) + + +@pytest.mark.parametrize("output", [NO_SPACE, "OSError: [Errno 122] Disk quota exceeded"]) +def test_out_of_disk_predicate_covers_quota(output): + assert startup_errors.out_of_disk(output) + + def test_falls_back_to_a_generic_subject(): assert classify(NO_SPACE).startswith("The model")