diff --git a/Dockerfile b/Dockerfile index 8e984633..1890b865 100644 --- a/Dockerfile +++ b/Dockerfile @@ -67,7 +67,5 @@ RUN --mount=type=secret,id=HF_TOKEN,required=false \ 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. +# Explicit ENTRYPOINT so the base image cannot swallow our command. ENTRYPOINT ["python3", "/src/main.py"] diff --git a/README.md b/README.md index 4c51c8f1..4314ec6d 100644 --- a/README.md +++ b/README.md @@ -119,7 +119,10 @@ The worker recognises these and stays up instead, answering every job with the c } ``` -The same message is logged by the worker, and the full vLLM traceback stays in the worker logs above it. Anything unrecognised still exits non-zero, because a failed download or a bad host is worth another attempt. +The same message is logged by the worker, and the full vLLM traceback stays in the worker logs above it. Anything unrecognised still exits non-zero, because a failed download or a bad host is worth another attempt; before exiting, the worker answers one queued job with a `prestart_failed` error that carries the end of vLLM's output, so the failure is visible without opening the logs. + +Startup runs in a RunPod prestart hook. Jobs stay queued while the model loads, +so their execution timeout does not start during the cold start. ## Option 2: Build Docker Image with Model Inside diff --git a/builder/requirements.txt b/builder/requirements.txt index a0cef110..11f338df 100644 --- a/builder/requirements.txt +++ b/builder/requirements.txt @@ -1,6 +1,6 @@ # vLLM itself comes from the vllm/vllm-openai base image; this is only what # the RunPod wrapper adds on top. -runpod~=1.12.0 +runpod~=2.0.0 aiohttp>=3.9 huggingface-hub>=0.24 # bitsandbytes moved out of vLLM core into an out-of-tree plugin in v0.28.0 diff --git a/docs/conventions.md b/docs/conventions.md index d0f84b40..45bbe875 100644 --- a/docs/conventions.md +++ b/docs/conventions.md @@ -27,12 +27,16 @@ Container start: ├─ (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, - │ unless startup_errors.classify() recognises the failure as one a restart - │ cannot fix — then stay up and answer every job with the cause) - └─ starts runpod.serverless loop with handler.handler + └─ starts runpod.serverless loop with handler.handler and the prestart hook + start_backend() (job intake stays closed until the hook settles): + ├─ 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. If vLLM dies or times out and + startup_errors.classify() recognises a failure a restart cannot fix, + stay up and answer every job with the cause; otherwise raise, so the + SDK answers one job with `prestart_failed` and exits non-zero. + (Realtime endpoints reject prestart hooks; there start_backend() runs + before the loop starts.) Per job: RunPod job → handler.py normalize input → aiohttp → 127.0.0.1:$VLLM_PORT/ @@ -44,8 +48,9 @@ boundary or the CLI; that is what makes vLLM upgrades a one-line `VLLM_VERSION` ### 2. **Components** -- `src/main.py`: container entrypoint — builds the `vllm serve` command, manages the - subprocess (health polling, signal forwarding), then starts the RunPod loop. +- `src/main.py`: container entrypoint — registers vLLM startup as a RunPod prestart + hook (builds the `vllm serve` command, manages the subprocess: health polling, + signal forwarding), then starts the RunPod loop. - `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. diff --git a/src/handler.py b/src/handler.py index c3596542..d9afa836 100644 --- a/src/handler.py +++ b/src/handler.py @@ -1,10 +1,8 @@ """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. +`main.py` starts vLLM in a prestart hook before this handler runs. Permanent +startup failures set `startup_error`, which is returned for every job. The +worker never imports vLLM. Accepted job input shapes (all under job["input"]): diff --git a/src/main.py b/src/main.py index d206fdc7..8129bbc8 100644 --- a/src/main.py +++ b/src/main.py @@ -1,9 +1,9 @@ -"""Container entrypoint: spawn `vllm serve`, wait for it, start RunPod serverless. +"""Container entrypoint: start RunPod serverless with vLLM startup as a prestart hook. -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. +The worker never imports vLLM. The prestart hook builds the CLI from environment +variables (see args_builder.py), launches `vllm serve` on the loopback interface, +and polls /health until the server (and model) is ready. Jobs stay queued until +startup settles. Before launching vLLM, a metadata-only pre-flight (model_preflight.py) checks that the configured model is actually fetchable from the Hugging Face Hub, so @@ -21,11 +21,12 @@ 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. +init time in recent vLLM releases). Unrecognised failures are reported by the SDK +before the worker exits non-zero, so the platform retries them. """ import collections +import contextvars import json import logging import os @@ -54,6 +55,8 @@ vllm_process: subprocess.Popen | None = None +backend_settled = threading.Event() + # 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) @@ -110,7 +113,10 @@ def start_vllm() -> subprocess.Popen: errors="replace", env=env, ) - output_pump = threading.Thread(target=pump_output, args=(proc,), daemon=True) + # Carry the hook's log capture into the output thread. + output_pump = threading.Thread( + target=contextvars.copy_context().run, args=(pump_output, proc), daemon=True + ) output_pump.start() return proc @@ -221,19 +227,10 @@ def apply_oom_recovery() -> bool: return relaxed -def main() -> None: +def start_backend() -> 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") - - for sig in (signal.SIGTERM, signal.SIGINT): - signal.signal(sig, _forward_signal) + import handler as proxy_handler # Ask the HF Hub whether the model is fetchable before paying for the # download: a typo'd MODEL_NAME, a gated model without HF_TOKEN, or a bad @@ -308,25 +305,45 @@ def main() -> None: 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) + raise # 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). + proxy_handler.vllm_process = vllm_process + proxy_handler.startup_error = startup_error + backend_settled.set() + + +def main() -> None: + 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") + + for sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(sig, _forward_signal) + + # SDK startup checks read the model environment at import time. import handler as proxy_handler import runpod - proxy_handler.vllm_process = vllm_process - proxy_handler.startup_error = startup_error + if int(os.getenv("RUNPOD_REALTIME_PORT", "0")): + # The SDK rejects prestart hooks in realtime mode. + start_backend() + else: + runpod.serverless.register_prestart_hook(start_backend) max_concurrency = int(os.getenv("MAX_CONCURRENCY", "30")) runpod.serverless.start( { "handler": proxy_handler.handler, - "concurrency_modifier": lambda _current: max_concurrency, + # Don't start job execution timeouts during model loading. + "concurrency_modifier": lambda _current: max_concurrency if backend_settled.is_set() else 0, "return_aggregate_stream": True, } ) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..17293d77 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,32 @@ +import sys +import threading +import types +from pathlib import Path +from unittest import mock + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +import main + + +@pytest.fixture +def fake_sdk(monkeypatch): + sdk = types.SimpleNamespace(hooks=[], handler=types.ModuleType("handler")) + sdk.handler.handler = lambda job: None + + def start(config): + for hook in sdk.hooks: + hook() + + sdk.start = mock.MagicMock(side_effect=start) + runpod_stub = types.ModuleType("runpod") + runpod_stub.serverless = types.SimpleNamespace( + register_prestart_hook=sdk.hooks.append, start=sdk.start + ) + monkeypatch.setitem(sys.modules, "handler", sdk.handler) + monkeypatch.setitem(sys.modules, "runpod", runpod_stub) + monkeypatch.setattr(main, "backend_settled", threading.Event()) + monkeypatch.delenv("RUNPOD_REALTIME_PORT", raising=False) + return sdk diff --git a/tests/test_compile_cache.py b/tests/test_compile_cache.py index 914d2479..38bd22d2 100644 --- a/tests/test_compile_cache.py +++ b/tests/test_compile_cache.py @@ -161,17 +161,11 @@ def test_fallback_root_matches_vllm_default(monkeypatch): @pytest.fixture -def run_main(monkeypatch): +def run_main(monkeypatch, fake_sdk): """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 = [] @@ -192,7 +186,7 @@ def fake_wait(proc): monkeypatch.setattr(main, "wait_for_vllm", fake_wait) main.recent_output.clear() main.main() - return roots, handler_stub + return roots, fake_sdk.handler return drive diff --git a/tests/test_main_oom_retry.py b/tests/test_main_oom_retry.py index a3ba8125..6d3e197a 100644 --- a/tests/test_main_oom_retry.py +++ b/tests/test_main_oom_retry.py @@ -2,7 +2,6 @@ import os import sys -import types from pathlib import Path from unittest import mock @@ -30,7 +29,7 @@ @pytest.fixture -def harness(monkeypatch): +def harness(monkeypatch, fake_sdk): """Run main() without vLLM, the RunPod SDK, or a GPU. Returns a driver: call it with the outputs each launch attempt should die @@ -47,13 +46,6 @@ def harness(monkeypatch): 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) - start_vllm = mock.MagicMock(return_value=mock.MagicMock()) monkeypatch.setattr(main, "start_vllm", start_vllm) @@ -70,8 +62,8 @@ def fake_wait(proc): monkeypatch.setattr(main, "wait_for_vllm", fake_wait) main.recent_output.clear() main.main() - handler_stub.launches = start_vllm.call_count - return handler_stub + fake_sdk.handler.launches = start_vllm.call_count + return fake_sdk.handler return drive diff --git a/tests/test_main_prestart.py b/tests/test_main_prestart.py new file mode 100644 index 00000000..e0e9717f --- /dev/null +++ b/tests/test_main_prestart.py @@ -0,0 +1,85 @@ +from unittest import mock + +import pytest + +import main + +GATED = "huggingface_hub.errors.GatedRepoError: 401 Client Error: Cannot access gated repo" +UNKNOWN = "RuntimeError: something nobody has seen before" + + +@pytest.fixture +def run_main(monkeypatch, fake_sdk): + monkeypatch.setenv("MODEL_NAME", "org/model") + monkeypatch.setenv("MAX_CONCURRENCY", "7") + monkeypatch.delenv("VLLM_CACHE_ROOT", raising=False) + monkeypatch.setattr(main.model_preflight, "check_model_access", lambda: None) + monkeypatch.setattr(main, "stop_vllm", lambda proc: None) + monkeypatch.setattr(main, "start_vllm", mock.MagicMock(return_value=mock.MagicMock())) + + seen = {} + real_start = fake_sdk.start.side_effect + + def start(config): + seen["before"] = config["concurrency_modifier"](1) + try: + real_start(config) + finally: + seen["after"] = config["concurrency_modifier"](1) + + fake_sdk.start.side_effect = start + + def drive(output): + def fake_wait(proc): + 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, "wait_for_vllm", fake_wait) + main.recent_output.clear() + main.main() + return seen + + return drive + + +def test_healthy_start_opens_job_intake(run_main, fake_sdk): + seen = run_main(None) + + assert fake_sdk.hooks == [main.start_backend] + assert seen == {"before": 0, "after": 7} + assert fake_sdk.handler.startup_error is None + + +def test_known_fatal_failure_opens_intake_to_answer_jobs_with_the_cause(run_main, fake_sdk): + seen = run_main(GATED) + + assert seen == {"before": 0, "after": 7} + assert "gated or private" in fake_sdk.handler.startup_error + + +def test_unknown_failure_raises_from_the_hook_and_never_opens_intake(run_main, fake_sdk): + with pytest.raises(RuntimeError, match="exited during startup"): + run_main(UNKNOWN) + + assert fake_sdk.hooks == [main.start_backend] + assert main.backend_settled.is_set() is False + + +def test_realtime_starts_vllm_before_the_sdk_without_a_hook(run_main, fake_sdk, monkeypatch): + monkeypatch.setenv("RUNPOD_REALTIME_PORT", "8888") + + seen = run_main(None) + + assert fake_sdk.hooks == [] + assert seen == {"before": 7, "after": 7} + + +def test_realtime_unknown_failure_exits_before_the_sdk_starts(run_main, fake_sdk, monkeypatch): + monkeypatch.setenv("RUNPOD_REALTIME_PORT", "8888") + + with pytest.raises(RuntimeError): + run_main(UNKNOWN) + + fake_sdk.start.assert_not_called() diff --git a/tests/test_main_revision_retry.py b/tests/test_main_revision_retry.py index 625db92e..0acf812f 100644 --- a/tests/test_main_revision_retry.py +++ b/tests/test_main_revision_retry.py @@ -5,7 +5,6 @@ """ import sys -import types from pathlib import Path from unittest import mock @@ -23,7 +22,7 @@ @pytest.fixture -def harness(monkeypatch): +def harness(monkeypatch, fake_sdk): """Run main() without vLLM, the RunPod SDK, or a GPU. Returns a driver: call it with the outputs each launch attempt should die @@ -38,13 +37,6 @@ def harness(monkeypatch): 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) - start_vllm = mock.MagicMock(return_value=mock.MagicMock()) monkeypatch.setattr(main, "start_vllm", start_vllm) @@ -61,8 +53,8 @@ def fake_wait(proc): monkeypatch.setattr(main, "wait_for_vllm", fake_wait) main.recent_output.clear() main.main() - handler_stub.launches = start_vllm.call_count - return handler_stub + fake_sdk.handler.launches = start_vllm.call_count + return fake_sdk.handler return drive