feat: run vLLM startup through prestart hooks

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"]

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
 

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

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

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"]):
 

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,
         }
     )

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

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
 

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
 

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()