From dd6ffd2987fb0c28f08fb4a613b904c783095a34 Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Fri, 2 Oct 2026 14:24:22 -0700 Subject: [PATCH 1/2] Refactor snapshot engine preparation helpers Signed-off-by: Nils Matteson --- vllm/snapshot/engine.py | 66 ++++++++++++++++++++++++++++++++++ vllm/snapshot/server.py | 78 ++++++++++------------------------------- 2 files changed, 84 insertions(+), 60 deletions(-) create mode 100644 vllm/snapshot/engine.py diff --git a/vllm/snapshot/engine.py b/vllm/snapshot/engine.py new file mode 100644 index 000000000000..02f385a66a30 --- /dev/null +++ b/vllm/snapshot/engine.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +from vllm.snapshot.types import Oracle + +_CANARY_PROMPT = "The capital of France is" + + +class SnapshotCanaryError(RuntimeError): + """The initialized engine did not produce a valid snapshot oracle.""" + + +async def _release_reloadable_state(engine: Any) -> None: + """Discard model and KV state before the process image is captured.""" + await engine.sleep(level=2) + + +async def _restore_reloadable_state(engine: Any) -> None: + """Rebuild state discarded by ``_release_reloadable_state``.""" + await engine.wake_up(tags=["weights"]) + await engine.collective_rpc("reload_weights") + await engine.wake_up(tags=["kv_cache"]) + + +def oracle_from_request_output(request_output: Any) -> Oracle: + try: + candidate = request_output.outputs[0] + token_ids = tuple(candidate.token_ids) + (sampled_token_id,) = token_ids + return Oracle( + token_ids=token_ids, + text=candidate.text, + sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob, + ) + except ValueError as error: + raise SnapshotCanaryError( + "snapshot canary must produce exactly one finite token logprob" + ) from error + except (AttributeError, IndexError, KeyError, TypeError) as error: + raise SnapshotCanaryError( + "snapshot canary did not return sampled token logprob" + ) from error + + +async def run_engine_canary(engine: Any) -> Oracle: + from vllm import SamplingParams + + final_output = None + sampling_params = SamplingParams( + temperature=0, + min_tokens=1, + max_tokens=1, + seed=0, + logprobs=0, + ) + async for output in engine.generate( + _CANARY_PROMPT, + sampling_params, + request_id="vllm-snapshot-canary", + ): + final_output = output + if final_output is None: + raise SnapshotCanaryError("canary generation returned no output") + return oracle_from_request_output(final_output) diff --git a/vllm/snapshot/server.py b/vllm/snapshot/server.py index 2dc0f3384aeb..fac8f7a205da 100644 --- a/vllm/snapshot/server.py +++ b/vllm/snapshot/server.py @@ -13,15 +13,27 @@ from pydantic import ValidationError +from vllm.snapshot.engine import ( + _CANARY_PROMPT as _CANARY_PROMPT, +) +from vllm.snapshot.engine import ( + SnapshotCanaryError as SnapshotCanaryError, +) +from vllm.snapshot.engine import ( + _release_reloadable_state as _release_reloadable_state, +) +from vllm.snapshot.engine import ( + _restore_reloadable_state as _restore_reloadable_state, +) +from vllm.snapshot.engine import ( + oracle_from_request_output as oracle_from_request_output, +) +from vllm.snapshot.engine import ( + run_engine_canary as run_engine_canary, +) from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic from vllm.snapshot.types import Oracle, oracles_match -_CANARY_PROMPT = "The capital of France is" - - -class SnapshotCanaryError(RuntimeError): - """The initialized engine did not produce a valid snapshot oracle.""" - class SnapshotBarrierError(RuntimeError): """The controller supplied an invalid snapshot release marker.""" @@ -89,60 +101,6 @@ def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]: ) -async def _release_reloadable_state(engine: Any) -> None: - """Discard model and KV state before the process image is captured.""" - await engine.sleep(level=2) - - -async def _restore_reloadable_state(engine: Any) -> None: - """Rebuild state discarded by ``_release_reloadable_state``.""" - await engine.wake_up(tags=["weights"]) - await engine.collective_rpc("reload_weights") - await engine.wake_up(tags=["kv_cache"]) - - -def oracle_from_request_output(request_output: Any) -> Oracle: - try: - candidate = request_output.outputs[0] - token_ids = tuple(candidate.token_ids) - (sampled_token_id,) = token_ids - return Oracle( - token_ids=token_ids, - text=candidate.text, - sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob, - ) - except ValueError as error: - raise SnapshotCanaryError( - "snapshot canary must produce exactly one finite token logprob" - ) from error - except (AttributeError, IndexError, KeyError, TypeError) as error: - raise SnapshotCanaryError( - "snapshot canary did not return sampled token logprob" - ) from error - - -async def run_engine_canary(engine: Any) -> Oracle: - from vllm import SamplingParams - - final_output = None - sampling_params = SamplingParams( - temperature=0, - min_tokens=1, - max_tokens=1, - seed=0, - logprobs=0, - ) - async for output in engine.generate( - _CANARY_PROMPT, - sampling_params, - request_id="vllm-snapshot-canary", - ): - final_output = output - if final_output is None: - raise SnapshotCanaryError("canary generation returned no output") - return oracle_from_request_output(final_output) - - def detach_snapshot_streams() -> None: """Remove launch-log descriptors from the reusable process image.""" sys.stdout.flush() From 14903f61d1aa80f6b52b93ba9e866a4e21b0ff7d Mon Sep 17 00:00:00 2001 From: Nils Matteson Date: Fri, 2 Oct 2026 15:10:20 -0700 Subject: [PATCH 2/2] Complete snapshot recovery before serving The snapshot CLI bound HTTP before validating recovered output, and its engine sequence omitted communicator checkpoint hooks. Share a one-shot startup lifecycle that completes recovery and a private canary before serving. Failed or cancelled phases terminate the attempt. Preserve compact-state rehearsal and default server setup ordering. Keep memory policy separate, validate parser configuration before engine construction, and record the failed phase before worker shutdown. Signed-off-by: Nils Matteson --- docs/features/initialized_snapshots.md | 23 +- .../entrypoints/launchers/test_launch_cli.py | 301 ++++++++++++++++-- vllm/entrypoints/launchers/launcher.py | 15 +- vllm/snapshot/engine.py | 111 ++++++- vllm/snapshot/server.py | 106 +++--- 5 files changed, 478 insertions(+), 78 deletions(-) diff --git a/docs/features/initialized_snapshots.md b/docs/features/initialized_snapshots.md index 15cce7554bb0..49db08035db0 100644 --- a/docs/features/initialized_snapshots.md +++ b/docs/features/initialized_snapshots.md @@ -4,8 +4,9 @@ Initialized engine snapshots are an experimental way to trade local disk space and host privileges for a faster vLLM activation. Snapshot creation initializes the engine and records deterministic generation output before CRIU captures the process tree and CUDA state. Restore validates the saved environment, restores -the engine, binds the HTTP server, and reproduces the recorded token and sampled -token log probability before returning. +the engine, and checks the recorded token and sampled-token log probability +before binding HTTP. The controller repeats the check through the public API +before returning. This path is intended for repeatedly activating the same model and engine configuration on the same machine. It is not a portable model artifact. @@ -94,6 +95,7 @@ Stop and remove it only after the restored API server is no longer needed. Create initializes the engine, records a one-token canary, releases and reloads weights and KV cache to rehearse restore, then releases them again for capture. +Preparation and recovery also await the existing communicator checkpoint hooks. The manifest is published only after CRIU completes and the source tree stops. Creation is offline preparation and is not part of restore latency. @@ -103,19 +105,24 @@ Literal API keys and Hugging Face tokens are redacted from the manifest's engine arguments. The manifest records selected environment names plus deterministic name-bound fingerprints, not their values. The protected CRIU artifact can still contain process secrets, so treat it as sensitive data. Restore reloads -model files and KV cache before binding HTTP, then reproduces the canary or -tears down the restored tree. The inspect command prints that identity and -canary without executing the saved process. +model files and KV cache, then validates the canary before binding HTTP. +The inspect command prints that identity and canary without executing the saved +process. ## Restore behavior Restore fails before CRIU runs if the saved identity does not match the current host. It does not silently fall back to ordinary startup. After CRIU restores -the process tree, vLLM releases the saved engine to bind the requested HTTP -address and checks the first generated token and sampled-token log probability -against the snapshot canary. The command returns only after that check passes. +the process tree, vLLM completes communicator and memory recovery and checks +the snapshot canary on the private engine. Only then does it bind the requested +HTTP address. The command returns after a second canary check through HTTP. The restored API server continues to run as a detached process. +A failed engine preparation or recovery terminates the attempt without opening +HTTP. The artifact's `error.json` records the phase and error before worker +cleanup, since the restored process's launch-log streams are detached. Each +activation clears the previous recovery error before attempting recovery. + The pre-release port probe is best-effort only. It neither reserves the port nor authenticates the listener that appears afterward. diff --git a/tests/entrypoints/launchers/test_launch_cli.py b/tests/entrypoints/launchers/test_launch_cli.py index fc8301069858..daa1a5826482 100644 --- a/tests/entrypoints/launchers/test_launch_cli.py +++ b/tests/entrypoints/launchers/test_launch_cli.py @@ -3,6 +3,7 @@ """Unit tests for the `vllm launch` CLI subcommand.""" import argparse +import asyncio import hashlib import json import os @@ -11,11 +12,13 @@ import sys from contextlib import asynccontextmanager from pathlib import Path +from types import SimpleNamespace from typing import Any -from unittest.mock import create_autospec, patch +from unittest.mock import AsyncMock, Mock, create_autospec, patch import pytest +import vllm.snapshot.engine as snapshot_engine import vllm.snapshot.runtime as snapshot_runtime import vllm.snapshot.server as snapshot_server from vllm.entrypoints.cli import snapshot as snapshot_cli @@ -567,43 +570,289 @@ def test_snapshot_create_abort_rejects_surviving_process_group( tools.abort_create(100) -@pytest.mark.asyncio -async def test_snapshot_child_rejects_mismatch_before_publication( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): +@pytest.fixture +def snapshot_child(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + from vllm.entrypoints.launchers import launcher from vllm.entrypoints.launchers.api_server import entry as api_server - canaries = iter((_oracle(), _oracle((42,), " Lyon"))) + engine = SimpleNamespace( + get_num_unfinished_requests=Mock(return_value=0), + pause_generation=AsyncMock(), + resume_generation=AsyncMock(), + sleep=AsyncMock(), + wake_up=AsyncMock(), + collective_rpc=AsyncMock(), + shutdown=Mock(), + ) + canary = AsyncMock(return_value=_oracle()) + closed = Mock() @asynccontextmanager async def engine_context(): - yield object() - - async def next_canary(_engine: object) -> Oracle: - return next(canaries) - - async def return_none(_engine: object) -> None: - return None - - monkeypatch.setattr( - api_server, "build_async_engine_client", lambda _args: engine_context() - ) - monkeypatch.setattr(snapshot_server, "run_engine_canary", next_canary) - monkeypatch.setattr(snapshot_server, "_release_reloadable_state", return_none) - monkeypatch.setattr(snapshot_server, "_restore_reloadable_state", return_none) + try: + yield engine + finally: + closed() + + async def serve(*_args): + return asyncio.sleep(0) + + builder = Mock(side_effect=lambda _args: engine_context()) + bind = Mock(return_value=("http://127.0.0.1:8000", Mock())) + serving = AsyncMock(side_effect=serve) + monkeypatch.setattr(api_server, "build_async_engine_client", builder) + monkeypatch.setattr(api_server, "build_and_serve", serving) + monkeypatch.setattr(launcher, "bind_server_socket", bind) + monkeypatch.setattr(snapshot_engine, "run_engine_canary", canary) monkeypatch.setattr(snapshot_server, "detach_snapshot_streams", lambda: None) - monkeypatch.setattr( - snapshot_server, - "write_ready_atomic", - lambda *_args: pytest.fail("mismatched snapshot was published"), - ) control = snapshot_server.ControlArgs( ready_file=tmp_path / "ready.json", release_file=tmp_path / "release.json", release_timeout_s=1, ) + control.release_file.write_text(json.dumps({"release": True, "port": 8000})) + return SimpleNamespace( + engine=engine, + canary=canary, + control=control, + args=snapshot_server.parse_vllm_args(["test-model"]), + builder=builder, + bind=bind, + serving=serving, + closed=closed, + ) + + +@pytest.mark.asyncio +async def test_snapshot_child_rejects_mismatch_before_publication(snapshot_child): + child = snapshot_child + child.canary.side_effect = [_oracle(), _oracle((42,), " Lyon")] with pytest.raises(SnapshotCanaryError, match="rehearsal changed"): - await snapshot_server.run_vllm_snapshot_child(control, argparse.Namespace()) + await snapshot_server.run_vllm_snapshot_child(child.control, child.args) + assert not child.control.ready_file.exists() + child.bind.assert_not_called() + child.engine.shutdown.assert_called_once_with(timeout=0) + child.closed.assert_called_once() + + +@pytest.mark.asyncio +async def test_snapshot_child_waits_for_workers_and_canary_before_binding( + snapshot_child, +): + child = snapshot_child + prepare_entered, prepare_done = asyncio.Event(), asyncio.Event() + recover_entered, recover_done = asyncio.Event(), asyncio.Event() + canary_entered, canary_done = asyncio.Event(), asyncio.Event() + preparations = recoveries = canaries = 0 + + async def rpc(method, **_kwargs): + nonlocal preparations, recoveries + if method == "checkpoint_prepare": + preparations += 1 + if preparations == 2: + prepare_entered.set() + await prepare_done.wait() + elif method == "checkpoint_restore": + recoveries += 1 + if recoveries == 2: + recover_entered.set() + await recover_done.wait() + + async def canary(_engine): + nonlocal canaries + canaries += 1 + if canaries == 3: + canary_entered.set() + await canary_done.wait() + return _oracle() + + child.engine.collective_rpc.side_effect = rpc + child.canary.side_effect = canary + task = asyncio.create_task( + snapshot_server.run_vllm_snapshot_child(child.control, child.args) + ) + try: + await asyncio.wait_for(prepare_entered.wait(), 1) + assert not child.control.ready_file.exists() + child.bind.assert_not_called() + child.control.ready_file.with_name("error.json").write_text("old activation") + prepare_done.set() + await asyncio.wait_for(recover_entered.wait(), 1) + assert child.control.ready_file.exists() + assert not child.control.ready_file.with_name("error.json").exists() + child.bind.assert_not_called() + recover_done.set() + await asyncio.wait_for(canary_entered.wait(), 1) + # Memory is awake and may have resumed scheduling. It is still private. + assert child.engine.wake_up.await_count == 4 + child.bind.assert_not_called() + child.serving.assert_not_awaited() + canary_done.set() + await task + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + child.bind.assert_called_once() + child.serving.assert_awaited_once() + child.closed.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["communicator", "reload", "canary"]) +async def test_snapshot_recovery_failure_never_binds_and_keeps_diagnostic( + snapshot_child, failure: str +): + child = snapshot_child + restore_count = 0 + + async def rpc(method, **_kwargs): + nonlocal restore_count + if method == "checkpoint_restore": + restore_count += 1 + if ( + restore_count == 2 + and method + == { + "communicator": "checkpoint_restore", + "reload": "reload_weights", + "canary": "unused", + }[failure] + ): + raise RuntimeError("injected recovery failure") + + child.engine.collective_rpc.side_effect = rpc + if failure == "canary": + child.canary.side_effect = [_oracle(), _oracle(), _oracle((42,), " Lyon")] + # Cleanup failure must not hide the failure that made recovery unsafe. + child.engine.shutdown.side_effect = RuntimeError("cleanup failed") + with pytest.raises(RuntimeError, match="recovery failure|recovery changed"): + await snapshot_server.run_vllm_snapshot_child(child.control, child.args) + error = json.loads(child.control.ready_file.with_name("error.json").read_text()) + assert error["phase"].startswith("recovery:") + assert "cleanup failed" not in error["error"] + child.bind.assert_not_called() + child.serving.assert_not_awaited() + child.engine.shutdown.assert_called_once_with(timeout=0) + child.closed.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_snapshot_recovery_deadline_or_cancellation_is_terminal( + snapshot_child, cancel: bool +): + child = snapshot_child + session = snapshot_engine.SnapshotSession(child.engine, timeout_s=0.1) + await session.prepare() + entered = asyncio.Event() + + async def held_canary(_engine): + entered.set() + await asyncio.Event().wait() + + child.canary.side_effect = held_canary + recovery = asyncio.create_task(session.recover()) + await asyncio.wait_for(entered.wait(), 1) + if cancel: + recovery.cancel() + with pytest.raises(asyncio.CancelledError if cancel else asyncio.TimeoutError): + await recovery + assert session.state == "failed" + assert session.phase == "recovery: canary" + rpc_count = child.engine.collective_rpc.await_count + with pytest.raises(RuntimeError, match="state failed"): + await session.recover() + with pytest.raises(RuntimeError, match="state failed"): + await session.prepare() + assert child.engine.collective_rpc.await_count == rpc_count + + +@pytest.mark.asyncio +async def test_snapshot_partial_sleep_failure_cannot_be_recovered(snapshot_child): + child = snapshot_child + child.engine.sleep.side_effect = RuntimeError("failure after releasing memory") + session = snapshot_engine.SnapshotSession(child.engine, timeout_s=1) + with pytest.raises(RuntimeError, match="after releasing"): + await session.prepare() + with pytest.raises(RuntimeError, match="state failed"): + await session.recover() + child.engine.wake_up.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_snapshot_session_completed_operations_are_not_repeated(snapshot_child): + child = snapshot_child + session = snapshot_engine.SnapshotSession(child.engine, timeout_s=1) + assert await session.prepare() == await session.prepare() == _oracle() + assert child.engine.sleep.await_count == 2 + await session.recover() + await session.recover() + assert child.canary.await_count == 3 + with pytest.raises(RuntimeError, match="state ready"): + await session.prepare() + + +@pytest.mark.asyncio +async def test_snapshot_memory_policy_does_not_own_private_generation(snapshot_child): + child = snapshot_child + policy = snapshot_engine.SnapshotMemoryPolicy(AsyncMock(), AsyncMock()) + session = snapshot_engine.SnapshotSession( + child.engine, timeout_s=1, memory_policy=policy + ) + await session.prepare() + await session.recover() + child.engine.pause_generation.assert_awaited_with(mode="wait") + assert child.engine.resume_generation.await_count == 2 + assert policy.release.await_count == policy.restore.await_count == 2 + child.engine.sleep.assert_not_awaited() + child.engine.wake_up.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_snapshot_session_rejects_existing_requests(snapshot_child): + child = snapshot_child + child.engine.get_num_unfinished_requests.return_value = 1 + session = snapshot_engine.SnapshotSession(child.engine, timeout_s=1) + with pytest.raises(RuntimeError, match="idle startup engine"): + await session.prepare() + child.canary.assert_not_awaited() + child.engine.sleep.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_snapshot_parser_validation_precedes_engine_creation(snapshot_child): + child = snapshot_child + child.args.enable_auto_tool_choice = True + child.args.tool_call_parser = "not-registered-snapshot-parser" + with pytest.raises(KeyError, match="invalid tool call parser"): + await snapshot_server.run_vllm_snapshot_child(child.control, child.args) + child.builder.assert_not_called() + child.bind.assert_not_called() + assert not child.control.ready_file.exists() + + +def test_default_server_setup_still_validates_before_binding(monkeypatch): + from vllm.entrypoints.launchers import launcher + + events = [] + monkeypatch.setattr( + launcher, "prepare_server_args", lambda _args: events.append("validate") + ) + monkeypatch.setattr( + launcher, + "bind_server_socket", + lambda _args, **_kwargs: events.append("bind"), + ) + launcher.setup_server(argparse.Namespace(), reuse_port=False) + assert events == ["validate", "bind"] + + +@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf")]) +def test_snapshot_session_rejects_unbounded_deadline(timeout): + with pytest.raises(ValueError, match="positive and finite"): + snapshot_engine.SnapshotSession(object(), timeout_s=timeout) def test_snapshot_release_marker_ignores_unknown_keys(tmp_path: Path): diff --git a/vllm/entrypoints/launchers/launcher.py b/vllm/entrypoints/launchers/launcher.py index 9fabfa3277af..f3de23c156ca 100644 --- a/vllm/entrypoints/launchers/launcher.py +++ b/vllm/entrypoints/launchers/launcher.py @@ -248,9 +248,8 @@ def validate_api_server_args(args): ) -@instrument(span_name="API server setup") -def setup_server(args, *, reuse_port: bool): - """Validate API server args and create the server socket.""" +def prepare_server_args(args): + """Load parser plugins and validate serving arguments before engine startup.""" log_version_and_model(logger, VLLM_VERSION, args.model) log_non_default_args(args) @@ -262,9 +261,19 @@ def setup_server(args, *, reuse_port: bool): validate_api_server_args(args) + +@instrument(span_name="API server setup") +def setup_server(args, *, reuse_port: bool): + """Validate API server args and create the server socket.""" + prepare_server_args(args) # workaround to make sure that we bind the port before the engine is set up. # This avoids race conditions with ray. # see https://github.com/vllm-project/vllm/issues/8204 + return bind_server_socket(args, reuse_port=reuse_port) + + +def bind_server_socket(args, *, reuse_port: bool): + """Bind after argument validation; startup capture defers this until recovery.""" if args.uds: sock = create_server_unix_socket(args.uds) else: diff --git a/vllm/snapshot/engine.py b/vllm/snapshot/engine.py index 02f385a66a30..0844d02148e1 100644 --- a/vllm/snapshot/engine.py +++ b/vllm/snapshot/engine.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any +import asyncio +import math +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Any, Literal -from vllm.snapshot.types import Oracle +from vllm.snapshot.types import Oracle, oracles_match _CANARY_PROMPT = "The capital of France is" @@ -24,6 +28,109 @@ async def _restore_reloadable_state(engine: Any) -> None: await engine.wake_up(tags=["kv_cache"]) +@dataclass(frozen=True) +class SnapshotMemoryPolicy: + """Paired memory operations; neither operation controls public admission.""" + + release: Callable[[Any], Awaitable[None]] + restore: Callable[[Any], Awaitable[None]] + + +class SnapshotSession: + """One capture/recovery cycle on an exclusively owned startup engine. + + The caller must not expose the engine to request producers until recovery + succeeds. Memory wake can resume scheduling before validation finishes. + Failures are terminal: the owner must shut down the engine, not retry wake. + This is not an admission barrier for an already-serving engine. + """ + + def __init__( + self, + engine: Any, + *, + timeout_s: float, + memory_policy: SnapshotMemoryPolicy | None = None, + ) -> None: + if not math.isfinite(timeout_s) or timeout_s <= 0: + raise ValueError("snapshot phase timeout must be positive and finite") + self.engine = engine + self.timeout_s = timeout_s + self.memory_policy = memory_policy or SnapshotMemoryPolicy( + _release_reloadable_state, _restore_reloadable_state + ) + self.state: Literal[ + "new", "preparing", "prepared", "recovering", "ready", "failed" + ] = "new" + self.phase = "new" + self._oracle: Oracle | None = None + + async def prepare(self) -> Oracle: + """Return only after rehearsal and every capture prerequisite complete.""" + if self.state == "prepared": + assert self._oracle is not None + return self._oracle + if self.state != "new": + raise RuntimeError(f"cannot prepare snapshot session in state {self.state}") + self.state = "preparing" + try: + await asyncio.wait_for(self._prepare(), timeout=self.timeout_s) + except BaseException: + self.state = "failed" + raise + self.state = "prepared" + self.phase = "prepared" + assert self._oracle is not None + return self._oracle + + async def recover(self) -> None: + """Complete communicator, memory and output validation before serving.""" + if self.state == "ready": + return + if self.state != "prepared": + raise RuntimeError(f"cannot recover snapshot session in state {self.state}") + self.state = "recovering" + try: + await asyncio.wait_for(self._recover("recovery"), timeout=self.timeout_s) + except BaseException: + self.state = "failed" + raise + self.state = "ready" + self.phase = "ready" + + async def _prepare(self) -> None: + self.phase = "reference canary" + if self.engine.get_num_unfinished_requests(): + raise RuntimeError("snapshot preparation requires an idle startup engine") + self._oracle = await run_engine_canary(self.engine) + await self._release() + await self._recover("rehearsal") + await self._release() + + async def _release(self) -> None: + self.phase = "pause generation" + await self.engine.pause_generation(mode="wait") + self.phase = "release memory" + await self.memory_policy.release(self.engine) + self.phase = "prepare communicators" + await self.engine.collective_rpc("checkpoint_prepare", timeout=self.timeout_s) + + async def _recover(self, phase: str) -> None: + self.phase = f"{phase}: restore communicators" + await self.engine.collective_rpc("checkpoint_restore", timeout=self.timeout_s) + self.phase = f"{phase}: restore memory" + await self.memory_policy.restore(self.engine) + self.phase = f"{phase}: synchronize device" + await self.engine.collective_rpc("synchronize_device", timeout=self.timeout_s) + self.phase = f"{phase}: resume private generation" + await self.engine.resume_generation() + self.phase = f"{phase}: canary" + oracle = await run_engine_canary(self.engine) + assert self._oracle is not None + if not oracles_match(self._oracle, oracle): + raise SnapshotCanaryError(f"snapshot {phase} changed canary output") + + def oracle_from_request_output(request_output: Any) -> Oracle: try: candidate = request_output.outputs[0] diff --git a/vllm/snapshot/server.py b/vllm/snapshot/server.py index fac8f7a205da..a35c78bf464f 100644 --- a/vllm/snapshot/server.py +++ b/vllm/snapshot/server.py @@ -4,6 +4,7 @@ import argparse import asyncio import json +import math import os import stat import sys @@ -13,12 +14,14 @@ from pydantic import ValidationError +from vllm.logger import init_logger from vllm.snapshot.engine import ( _CANARY_PROMPT as _CANARY_PROMPT, ) from vllm.snapshot.engine import ( SnapshotCanaryError as SnapshotCanaryError, ) +from vllm.snapshot.engine import SnapshotSession from vllm.snapshot.engine import ( _release_reloadable_state as _release_reloadable_state, ) @@ -32,7 +35,9 @@ run_engine_canary as run_engine_canary, ) from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic -from vllm.snapshot.types import Oracle, oracles_match +from vllm.snapshot.types import Oracle + +logger = init_logger(__name__) class SnapshotBarrierError(RuntimeError): @@ -89,8 +94,10 @@ def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]: control_args, remaining = parser.parse_known_args(argv) if remaining and remaining[0] == "--": remaining = remaining[1:] - if control_args.release_timeout_s <= 0: - raise ValueError("release timeout must be positive") + if not math.isfinite(control_args.release_timeout_s) or ( + control_args.release_timeout_s <= 0 + ): + raise ValueError("release timeout must be positive and finite") return ( ControlArgs( ready_file=control_args.ready_file, @@ -119,10 +126,10 @@ async def wait_for_release_marker( timeout_s: float, poll_interval_s: float = 0.05, ) -> ListenerConfig: - if timeout_s <= 0: - raise ValueError("release timeout must be positive") - if poll_interval_s <= 0: - raise ValueError("release poll interval must be positive") + if not math.isfinite(timeout_s) or timeout_s <= 0: + raise ValueError("release timeout must be positive and finite") + if not math.isfinite(poll_interval_s) or poll_interval_s <= 0: + raise ValueError("release poll interval must be positive and finite") remaining_s = timeout_s while not path.exists(): if remaining_s <= 0: @@ -158,44 +165,65 @@ async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None: build_and_serve, build_async_engine_client, ) - from vllm.entrypoints.launchers.launcher import setup_server + from vllm.entrypoints.launchers.launcher import ( + bind_server_socket, + prepare_server_args, + ) args.enable_sleep_mode = True + prepare_server_args(args) async with build_async_engine_client(args) as engine: - oracle = await run_engine_canary(engine) + session = SnapshotSession(engine, timeout_s=control.release_timeout_s) + error_file = control.ready_file.with_name("error.json") + phase = "prepare" try: - await _release_reloadable_state(engine) - await _restore_reloadable_state(engine) - rehearsal_oracle = await run_engine_canary(engine) - if not oracles_match(oracle, rehearsal_oracle): - raise SnapshotCanaryError("snapshot rehearsal changed canary output") - await _release_reloadable_state(engine) - except SnapshotCanaryError: - raise - except Exception as error: - raise SnapshotCanaryError("snapshot rehearsal failed") from error - - detach_snapshot_streams() - write_ready_atomic(control.ready_file, oracle) - listener = await wait_for_release_marker( - control.release_file, - timeout_s=control.release_timeout_s, - ) - args.host = listener.host - args.port = listener.port - await _restore_reloadable_state(engine) - listen_address, sock = setup_server(args, reuse_port=False) - try: - shutdown_task = await build_and_serve( - engine, - listen_address, - sock, - args, + oracle = await session.prepare() + phase = "capture barrier" + detach_snapshot_streams() + write_ready_atomic(control.ready_file, oracle) + listener = await wait_for_release_marker( + control.release_file, + timeout_s=control.release_timeout_s, ) - await shutdown_task - finally: - sock.close() + args.host = listener.host + args.port = listener.port + phase = "recover" + error_file.unlink(missing_ok=True) + await session.recover() + phase = "serve" + listen_address, sock = bind_server_socket(args, reuse_port=False) + try: + shutdown_task = await build_and_serve( + engine, + listen_address, + sock, + args, + ) + await shutdown_task + finally: + sock.close() + except BaseException as error: + # Restored stdout/stderr point at /dev/null. Preserve the primary + # failure before shutdown, including failures in the release barrier. + try: + _write_json_atomic( + error_file, + { + "phase": session.phase if session.state == "failed" else phase, + "error_type": type(error).__name__, + "error": str(error), + }, + ) + except Exception: + logger.exception("Could not record snapshot child failure") + try: + # A cancelled RPC may still be running. Terminate owned workers + # without trying to drain or wake partially released memory. + engine.shutdown(timeout=0) + except Exception: + logger.exception("Could not shut down failed snapshot engine") + raise def main(argv: list[str] | None = None) -> None: