Skip to content

vllm.snapshot.engine

Classes:

  • SnapshotCanaryError –

    The initialized engine did not produce a valid snapshot oracle.

  • SnapshotMemoryPolicy –

    Paired memory operations; neither operation controls public admission.

  • SnapshotSession –

    One capture/recovery cycle on an exclusively owned startup engine.

SnapshotCanaryError

Bases: RuntimeError

The initialized engine did not produce a valid snapshot oracle.

Source code in vllm/snapshot/engine.py 
15
16
class SnapshotCanaryError(RuntimeError):
    """The initialized engine did not produce a valid snapshot oracle."""

SnapshotMemoryPolicy dataclass

Paired memory operations; neither operation controls public admission.

Source code in vllm/snapshot/engine.py 
31
32
33
34
35
36
@dataclass(frozen=True)
class SnapshotMemoryPolicy:
    """Paired memory operations; neither operation controls public admission."""

    release: Callable[[Any], Awaitable[None]]
    restore: Callable[[Any], Awaitable[None]]

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.

Methods:

  • prepare –

    Return only after rehearsal and every capture prerequisite complete.

  • recover –

    Complete communicator, memory and output validation before serving.

Source code in vllm/snapshot/engine.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
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")

prepare() async

Return only after rehearsal and every capture prerequisite complete.

Source code in vllm/snapshot/engine.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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

recover() async

Complete communicator, memory and output validation before serving.

Source code in vllm/snapshot/engine.py 
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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"

_release_reloadable_state(engine) async

Discard model and KV state before the process image is captured.

Source code in vllm/snapshot/engine.py 
19
20
21
async def _release_reloadable_state(engine: Any) -> None:
    """Discard model and KV state before the process image is captured."""
    await engine.sleep(level=2)

_restore_reloadable_state(engine) async

Rebuild state discarded by _release_reloadable_state.

Source code in vllm/snapshot/engine.py 
24
25
26
27
28
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"])