feat: bootstrap api state from snapshot

src/exo/api/main.py

@@ -121,6 +121,8 @@ from exo.api.types.openai_responses import (
 )
 from exo.master.image_store import ImageStore
 from exo.master.placement import place_instance as get_instance_placements
+from exo.routing.event_router import EventRouter
+from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.apply import apply
 from exo.shared.constants import (
     DASHBOARD_DIR,
@@ -164,6 +166,7 @@ from exo.shared.types.commands import (
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
+    RequestSnapshot,
     SendInputChunk,
     SetInstanceLink,
     StartDownload,
@@ -171,7 +174,7 @@ from exo.shared.types.commands import (
     TaskFinished,
     TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SystemId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -181,6 +184,7 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.memory import Memory
+from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -207,6 +211,8 @@ from exo.utils.task_group import TaskGroup
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
 ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
 
+_SNAPSHOT_FETCH_TIMEOUT_SECONDS = 30
+
 
 def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
     return f"image/{image_format or 'png'}"
@@ -236,9 +242,12 @@ class API:
     def __init__(
         self,
         node_id: NodeId,
+        session_id: SessionId,
         *,
         port: int,
+        event_router: EventRouter,
         event_receiver: Receiver[IndexedEvent],
+        snapshot_chunk_receiver: Receiver[SnapshotChunk],
         command_sender: Sender[ForwarderCommand],
         download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
@@ -247,9 +256,12 @@ class API:
         self.state = State()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self._system_id = SystemId()
+        self.session_id = session_id
+        self.event_router = event_router
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
         self.event_receiver = event_receiver
+        self.snapshot_chunk_receiver = snapshot_chunk_receiver
         self.election_receiver = election_receiver
         self.node_id: NodeId = node_id
         self.last_completed_election: int = 0
@@ -291,18 +303,29 @@ class API:
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup = TaskGroup()
 
-    def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
+    def reset(
+        self,
+        result_clock: int,
+        session_id: SessionId,
+        event_router: EventRouter,
+        event_receiver: Receiver[IndexedEvent],
+        snapshot_chunk_receiver: Receiver[SnapshotChunk],
+    ):
         logger.info("Resetting API State")
         self._event_log.close()
         self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
         self.state = State()
         self._system_id = SystemId()
+        self.session_id = session_id
+        self.event_router = event_router
         self._text_generation_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
         self.event_receiver.close()
         self.event_receiver = event_receiver
-        self._tg.start_soon(self._apply_state)
+        self.snapshot_chunk_receiver.close()
+        self.snapshot_chunk_receiver = snapshot_chunk_receiver
+        self._tg.start_soon(self._bootstrap_then_apply_state)
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -1836,7 +1859,7 @@ class API:
         try:
             async with self._tg as tg:
                 logger.info("Starting API")
-                tg.start_soon(self._apply_state)
+                tg.start_soon(self._bootstrap_then_apply_state)
                 tg.start_soon(self._pause_on_new_election)
                 tg.start_soon(self._cleanup_expired_images)
                 print_startup_banner(self.port)
@@ -1850,6 +1873,7 @@ class API:
             self._event_log.close()
             self.command_sender.close()
             self.event_receiver.close()
+            self.snapshot_chunk_receiver.close()
 
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
@@ -1865,9 +1889,43 @@ class API:
                 shutdown_trigger=ev.wait,
             )
 
+    async def _bootstrap_then_apply_state(self):
+        await self._fetch_snapshot()
+        await self._apply_state()
+
+    async def _fetch_snapshot(self) -> None:
+        receiver = SnapshotReceiver(self.node_id, self.session_id)
+        await self.command_sender.send(
+            ForwarderCommand(
+                origin=self._system_id,
+                command=RequestSnapshot(requester_node_id=self.node_id),
+            )
+        )
+
+        with anyio.move_on_after(_SNAPSHOT_FETCH_TIMEOUT_SECONDS):
+            with self.snapshot_chunk_receiver as chunks:
+                async for chunk in chunks:
+                    received = receiver.ingest(chunk)
+                    if received is None:
+                        continue
+                    self.state = received.state
+                    self.event_router.set_buffer_start(
+                        received.last_event_applied_idx + 1
+                    )
+                    logger.info(
+                        f"API bootstrapped from snapshot at idx "
+                        f"{received.last_event_applied_idx}"
+                    )
+                    return
+        logger.info(
+            "API: no snapshot received before timeout; falling back to full event-log replay"
+        )
+
     async def _apply_state(self):
         with self.event_receiver as events:
             async for i_event in events:
+                if i_event.idx <= self.state.last_event_applied_idx:
+                    continue
                 self._event_log.append(i_event.event)
                 self.state = apply(self.state, i_event)
                 event = i_event.event

src/exo/api/tests/test_api_snapshot_bootstrap.py

@@ -0,0 +1,136 @@
+# pyright: reportPrivateUsage=false
+
+import hashlib
+
+import anyio
+import pytest
+import zstandard
+
+from exo.api.main import API
+from exo.routing.event_router import EventRouter
+from exo.shared.types.commands import ForwarderCommand, RequestSnapshot
+from exo.shared.types.common import NodeId, SessionId, SystemId
+from exo.shared.types.events import (
+    Event,
+    GlobalForwarderEvent,
+    IndexedEvent,
+    LocalForwarderEvent,
+    TestEvent,
+)
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+from exo.shared.types.state import State
+from exo.utils.channels import Receiver, Sender, channel
+
+
+class _FakeEventLog:
+    def __init__(self) -> None:
+        self.appended: list[Event] = []
+
+    def append(self, event: Event) -> None:
+        self.appended.append(event)
+
+
+def _snapshot_chunk(
+    state: State, *, requester_node_id: NodeId, session_id: SessionId
+) -> SnapshotChunk:
+    body = zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
+    return SnapshotChunk.from_data(
+        data=body,
+        transfer_id=SnapshotTransferId("transfer-1"),
+        requester_node_id=requester_node_id,
+        session_id=session_id,
+        schema_version=state.schema_version,
+        last_event_applied_idx=state.last_event_applied_idx,
+        chunk_index=0,
+        total_chunks=1,
+        sha256_hex=hashlib.sha256(body).hexdigest(),
+    )
+
+
+def _api(
+    node_id: NodeId, session_id: SessionId
+) -> tuple[
+    API,
+    EventRouter,
+    Receiver[ForwarderCommand],
+    Sender[SnapshotChunk],
+    Sender[IndexedEvent],
+    _FakeEventLog,
+]:
+    router_command_sender, _router_command_receiver = channel[ForwarderCommand]()
+    _global_event_sender, global_event_receiver = channel[GlobalForwarderEvent]()
+    local_event_sender, _local_event_receiver = channel[LocalForwarderEvent]()
+    event_router = EventRouter(
+        session_id=session_id,
+        command_sender=router_command_sender,
+        external_inbound=global_event_receiver,
+        external_outbound=local_event_sender,
+    )
+
+    event_sender, event_receiver = channel[IndexedEvent]()
+    command_sender, command_receiver = channel[ForwarderCommand]()
+    snapshot_sender, snapshot_receiver = channel[SnapshotChunk]()
+
+    api = object.__new__(API)
+    api.node_id = node_id
+    api.session_id = session_id
+    api.event_router = event_router
+    api.event_receiver = event_receiver
+    api.snapshot_chunk_receiver = snapshot_receiver
+    api.command_sender = command_sender
+    api._system_id = SystemId("api-system")
+    api.state = State()
+    event_log = _FakeEventLog()
+    api._event_log = event_log  # pyright: ignore[reportAttributeAccessIssue]
+    api._image_generation_queues = {}
+    api._text_generation_queues = {}
+    return api, event_router, command_receiver, snapshot_sender, event_sender, event_log
+
+
+@pytest.mark.asyncio
+async def test_api_fetch_snapshot_applies_state_and_fast_forwards_router() -> None:
+    node_id = NodeId("api")
+    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
+    api, event_router, command_receiver, snapshot_sender, _event_sender, _event_log = (
+        _api(node_id, session_id)
+    )
+    state = State(last_event_applied_idx=7)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(api._fetch_snapshot)
+        command = await command_receiver.receive()
+        assert isinstance(command.command, RequestSnapshot)
+        assert command.command.requester_node_id == node_id
+
+        await snapshot_sender.send(
+            _snapshot_chunk(state, requester_node_id=node_id, session_id=session_id)
+        )
+
+    assert api.state.last_event_applied_idx == 7
+    assert event_router.event_buffer.next_idx_to_release == 8
+
+
+@pytest.mark.asyncio
+async def test_api_apply_state_ignores_events_covered_by_snapshot() -> None:
+    node_id = NodeId("api")
+    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
+    (
+        api,
+        _event_router,
+        _command_receiver,
+        _snapshot_sender,
+        event_sender,
+        event_log,
+    ) = _api(node_id, session_id)
+    api.state = State(last_event_applied_idx=7)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(api._apply_state)
+        await event_sender.send(IndexedEvent(idx=7, event=TestEvent()))
+        await event_sender.send(IndexedEvent(idx=8, event=TestEvent()))
+
+        while api.state.last_event_applied_idx != 8:
+            await anyio.sleep(0.001)
+        tg.cancel_scope.cancel()
+
+    assert len(event_log.appended) == 1

src/exo/main.py

@@ -59,6 +59,7 @@ class Node:
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
         await router.register_topic(topics.DOWNLOAD_COMMANDS)
+        await router.register_topic(topics.SNAPSHOT_RESPONSES)
         event_router = EventRouter(
             session_id,
             command_sender=router.sender(topics.COMMANDS),
@@ -83,8 +84,11 @@ class Node:
         if args.spawn_api:
             api = API(
                 node_id,
+                session_id,
                 port=args.api_port,
+                event_router=event_router,
                 event_receiver=event_router.receiver(),
+                snapshot_chunk_receiver=router.receiver(topics.SNAPSHOT_RESPONSES),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -95,8 +99,11 @@ class Node:
         if not args.no_worker:
             worker = Worker(
                 node_id,
+                session_id,
+                event_router=event_router,
                 event_receiver=event_router.receiver(),
                 event_sender=event_router.sender(),
+                snapshot_chunk_receiver=router.receiver(topics.SNAPSHOT_RESPONSES),
                 command_sender=router.sender(topics.COMMANDS),
                 download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 api_port=args.api_port,
@@ -112,6 +119,7 @@ class Node:
             global_event_sender=router.sender(topics.GLOBAL_EVENTS),
             local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
             command_receiver=router.receiver(topics.COMMANDS),
+            snapshot_chunk_sender=router.sender(topics.SNAPSHOT_RESPONSES),
             download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
         )
 
@@ -210,6 +218,9 @@ class Node:
                         global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
                         local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
                         command_receiver=self.router.receiver(topics.COMMANDS),
+                        snapshot_chunk_sender=self.router.sender(
+                            topics.SNAPSHOT_RESPONSES
+                        ),
                         download_command_sender=self.router.sender(
                             topics.DOWNLOAD_COMMANDS
                         ),
@@ -246,8 +257,13 @@ class Node:
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
+                            result.session_id,
+                            event_router=self.event_router,
                             event_receiver=self.event_router.receiver(),
                             event_sender=self.event_router.sender(),
+                            snapshot_chunk_receiver=self.router.receiver(
+                                topics.SNAPSHOT_RESPONSES
+                            ),
                             command_sender=self.router.sender(topics.COMMANDS),
                             download_command_sender=self.router.sender(
                                 topics.DOWNLOAD_COMMANDS
@@ -256,7 +272,13 @@ class Node:
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
-                        self.api.reset(result.won_clock, self.event_router.receiver())
+                        self.api.reset(
+                            result.won_clock,
+                            result.session_id,
+                            self.event_router,
+                            self.event_router.receiver(),
+                            self.router.receiver(topics.SNAPSHOT_RESPONSES),
+                        )
                     self._tg.start_soon(self.event_router.run)
                 else:
                     if self.api:

src/exo/master/main.py

@@ -1,6 +1,8 @@
+import hashlib
 from datetime import datetime, timedelta, timezone
 
 import anyio
+from anyio import to_thread
 from loguru import logger
 
 from exo.master.placement import (
@@ -25,6 +27,7 @@ from exo.shared.types.commands import (
     ImageGeneration,
     PlaceInstance,
     RequestEventLog,
+    RequestSnapshot,
     SendInputChunk,
     SetInstanceLink,
     TaskCancelled,
@@ -54,6 +57,7 @@ from exo.shared.types.events import (
     TracesMerged,
 )
 from exo.shared.types.instance_link import InstanceLink
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
@@ -74,6 +78,15 @@ from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
 from exo.utils.task_group import TaskGroup
 
+_SNAPSHOT_CHUNK_BYTES = 512 * 1024
+_MAX_EVENT_LOG_REPLAY_BATCH = 1000
+
+
+def _encode_state_for_transfer(state: State) -> bytes:
+    import zstandard
+
+    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
+
 
 def _prefill_endpoint_for(state: State, decode_instance_id: InstanceId) -> str | None:
     decode = state.instances.get(decode_instance_id)
@@ -125,6 +138,7 @@ class Master:
         event_sender: Sender[Event],
         local_event_receiver: Receiver[LocalForwarderEvent],
         global_event_sender: Sender[GlobalForwarderEvent],
+        snapshot_chunk_sender: Sender[SnapshotChunk],
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
         self.node_id = node_id
@@ -135,6 +149,7 @@ class Master:
         self.command_receiver = command_receiver
         self.local_event_receiver = local_event_receiver
         self.global_event_sender = global_event_sender
+        self.snapshot_chunk_sender = snapshot_chunk_sender
         self.download_command_sender = download_command_sender
         self.event_sender = event_sender
         self._system_id = SystemId()
@@ -155,6 +170,7 @@ class Master:
             self._event_log.close()
             self.global_event_sender.close()
             self.local_event_receiver.close()
+            self.snapshot_chunk_sender.close()
             self.command_receiver.close()
 
     async def shutdown(self):
@@ -441,12 +457,19 @@ class Master:
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
-                            end = min(command.since_idx + 1000, len(self._event_log))
+                            end = min(
+                                command.since_idx + _MAX_EVENT_LOG_REPLAY_BATCH,
+                                len(self._event_log),
+                            )
                             for i, event in enumerate(
                                 self._event_log.read_range(command.since_idx, end),
                                 start=command.since_idx,
                             ):
                                 await self._send_event(IndexedEvent(idx=i, event=event))
+                        case RequestSnapshot():
+                            self._tg.start_soon(
+                                self._serve_snapshot, command.requester_node_id
+                            )
                     for event in generated_events:
                         await self.event_sender.send(event)
                 except ValueError as e:
@@ -506,6 +529,42 @@ class Master:
                     self._event_log.append(event)
                     await self._send_event(indexed)
 
+    async def _serve_snapshot(self, requester_node_id: NodeId) -> None:
+        state = self.state
+        if state.last_event_applied_idx < 0:
+            logger.info(
+                f"RequestSnapshot from {requester_node_id} but master has no events yet"
+            )
+            return
+
+        body = await to_thread.run_sync(_encode_state_for_transfer, state)
+        sha256 = hashlib.sha256(body).hexdigest()
+        chunks = [
+            body[i : i + _SNAPSHOT_CHUNK_BYTES]
+            for i in range(0, len(body), _SNAPSHOT_CHUNK_BYTES)
+        ] or [b""]
+        transfer_id = SnapshotTransferId()
+
+        logger.info(
+            f"Serving snapshot to {requester_node_id}: "
+            f"idx={state.last_event_applied_idx}, "
+            f"{len(chunks)} chunk(s), {len(body)} bytes total"
+        )
+        for index, chunk in enumerate(chunks):
+            await self.snapshot_chunk_sender.send(
+                SnapshotChunk.from_data(
+                    data=chunk,
+                    transfer_id=transfer_id,
+                    requester_node_id=requester_node_id,
+                    session_id=self.session_id,
+                    schema_version=state.schema_version,
+                    last_event_applied_idx=state.last_event_applied_idx,
+                    chunk_index=index,
+                    total_chunks=len(chunks),
+                    sha256_hex=sha256,
+                )
+            )
+
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly

src/exo/master/tests/test_master.py

@@ -7,12 +7,14 @@ from loguru import logger
 
 from exo.master.main import Master
 from exo.routing.router import get_node_id_keypair
+from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.commands import (
     CommandId,
     ForwarderCommand,
     ForwarderDownloadCommand,
     PlaceInstance,
+    RequestSnapshot,
     TextGeneration,
 )
 from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
@@ -29,6 +31,7 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryUsage,
 )
+from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.tasks import TaskStatus
 from exo.shared.types.tasks import TextGeneration as TextGenerationTask
 from exo.shared.types.text_generation import (
@@ -56,6 +59,7 @@ async def test_master():
     local_event_sender, le_receiver = channel[LocalForwarderEvent]()
     fcds, _fcdr = channel[ForwarderDownloadCommand]()
     ev_send, ev_recv = channel[Event]()
+    snapshot_chunk_send, _snapshot_chunk_recv = channel[SnapshotChunk]()
 
     async def mock_event_router():
         idx = 0
@@ -92,6 +96,7 @@ async def test_master():
         global_event_sender=ge_sender,
         local_event_receiver=le_receiver,
         command_receiver=co_receiver,
+        snapshot_chunk_sender=snapshot_chunk_send,
         download_command_sender=fcds,
     )
     logger.info("run the master")
@@ -229,3 +234,52 @@ async def test_master():
 
         ev_send.close()
         await master.shutdown()
+
+
+@pytest.mark.asyncio
+async def test_master_serves_snapshot_for_current_state():
+    node_id = NodeId("master")
+    requester_node_id = NodeId("worker")
+    session_id = SessionId(master_node_id=node_id, election_clock=0)
+
+    ge_sender, _global_event_receiver = channel[GlobalForwarderEvent]()
+    command_sender, command_receiver = channel[ForwarderCommand]()
+    _local_event_sender, local_event_receiver = channel[LocalForwarderEvent]()
+    download_command_sender, _download_command_receiver = channel[
+        ForwarderDownloadCommand
+    ]()
+    event_sender, _event_receiver = channel[Event]()
+    snapshot_chunk_sender, snapshot_chunk_receiver = channel[SnapshotChunk]()
+
+    master = Master(
+        node_id,
+        session_id,
+        event_sender=event_sender,
+        global_event_sender=ge_sender,
+        local_event_receiver=local_event_receiver,
+        command_receiver=command_receiver,
+        snapshot_chunk_sender=snapshot_chunk_sender,
+        download_command_sender=download_command_sender,
+    )
+    master.state = master.state.model_copy(update={"last_event_applied_idx": 12})
+
+    receiver = SnapshotReceiver(requester_node_id, session_id)
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(master.run)
+        await command_sender.send(
+            ForwarderCommand(
+                origin=SystemId("api"),
+                command=RequestSnapshot(requester_node_id=requester_node_id),
+            )
+        )
+
+        received = None
+        while received is None:
+            chunk = await snapshot_chunk_receiver.receive()
+            received = receiver.ingest(chunk)
+
+        assert received.last_event_applied_idx == 12
+        assert received.state.last_event_applied_idx == 12
+
+        await master.shutdown()
+        tg.cancel_scope.cancel()

src/exo/routing/tests/test_snapshot_routing_types.py

@@ -0,0 +1,37 @@
+from exo.routing import topics
+from exo.shared.types.commands import ForwarderCommand, RequestSnapshot
+from exo.shared.types.common import NodeId, SessionId, SystemId
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+
+
+def test_request_snapshot_round_trips_through_forwarder_command() -> None:
+    command = ForwarderCommand(
+        origin=SystemId("system-1"),
+        command=RequestSnapshot(requester_node_id=NodeId("worker-1")),
+    )
+
+    restored = ForwarderCommand.model_validate_json(command.model_dump_json())
+
+    assert isinstance(restored.command, RequestSnapshot)
+    assert restored.command.requester_node_id == NodeId("worker-1")
+
+
+def test_snapshot_response_topic_round_trips_chunk() -> None:
+    chunk = SnapshotChunk.from_data(
+        data=b"snapshot-bytes",
+        transfer_id=SnapshotTransferId("transfer-1"),
+        requester_node_id=NodeId("worker-1"),
+        session_id=SessionId(master_node_id=NodeId("master"), election_clock=1),
+        schema_version=1,
+        last_event_applied_idx=42,
+        chunk_index=0,
+        total_chunks=1,
+        sha256_hex="unused",
+    )
+
+    restored = topics.SNAPSHOT_RESPONSES.deserialize(
+        topics.SNAPSHOT_RESPONSES.serialize(chunk)
+    )
+
+    assert restored == chunk
+    assert restored.data == b"snapshot-bytes"

src/exo/routing/topics.py

@@ -8,6 +8,7 @@ from exo.shared.types.events import (
     GlobalForwarderEvent,
     LocalForwarderEvent,
 )
+from exo.shared.types.snapshots import SnapshotChunk
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -49,3 +50,6 @@ CONNECTION_MESSAGES = TypedTopic(
 DOWNLOAD_COMMANDS = TypedTopic(
     "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
 )
+SNAPSHOT_RESPONSES = TypedTopic(
+    "snapshot_responses", PublishPolicy.Always, SnapshotChunk
+)

src/exo/shared/types/commands.py

@@ -67,6 +67,12 @@ class RequestEventLog(BaseCommand):
     since_idx: int
 
 
+class RequestSnapshot(BaseCommand):
+    """Ask the current master to send a State snapshot to this node."""
+
+    requester_node_id: NodeId
+
+
 class StartDownload(BaseCommand):
     target_node_id: NodeId
     shard_metadata: ShardMetadata
@@ -106,6 +112,7 @@ DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 Command = (
     TestCommand
     | RequestEventLog
+    | RequestSnapshot
     | TextGeneration
     | ImageGeneration
     | ImageEdits

src/exo/worker/main.py

@@ -8,6 +8,8 @@ from loguru import logger
 
 from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
+from exo.routing.event_router import EventRouter
+from exo.routing.snapshot_receiver import SnapshotReceiver
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
 from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
@@ -16,9 +18,10 @@ from exo.shared.types.commands import (
     DeleteInstance,
     ForwarderCommand,
     ForwarderDownloadCommand,
+    RequestSnapshot,
     StartDownload,
 )
-from exo.shared.types.common import CommandId, NodeId, SystemId
+from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
     CustomModelCardAdded,
     CustomModelCardDeleted,
@@ -33,6 +36,7 @@ from exo.shared.types.events import (
     TopologyEdgeDeleted,
 )
 from exo.shared.types.multiaddr import Multiaddr
+from exo.shared.types.snapshots import SnapshotChunk
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     CancelTask,
@@ -58,14 +62,19 @@ from exo.utils.task_group import TaskGroup
 from exo.worker.plan import plan
 from exo.worker.runner.supervisor import RunnerSupervisor
 
+_SNAPSHOT_FETCH_TIMEOUT_SECONDS = 30
+
 
 class Worker:
     def __init__(
         self,
         node_id: NodeId,
+        session_id: SessionId,
         *,
+        event_router: EventRouter,
         event_receiver: Receiver[IndexedEvent],
         event_sender: Sender[Event],
+        snapshot_chunk_receiver: Receiver[SnapshotChunk],
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
@@ -73,8 +82,11 @@ class Worker:
         api_port: int,
     ):
         self.node_id: NodeId = node_id
+        self.session_id: SessionId = session_id
+        self.event_router = event_router
         self.event_receiver = event_receiver
         self.event_sender = event_sender
+        self.snapshot_chunk_receiver = snapshot_chunk_receiver
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
         self.api_port = api_port
@@ -104,21 +116,57 @@ class Worker:
 
         try:
             async with self._tg as tg:
-                tg.start_soon(info_gatherer.run)
-                tg.start_soon(self._forward_info, info_recv)
-                tg.start_soon(self.plan_step)
-                tg.start_soon(self._event_applier)
-                tg.start_soon(self._poll_connection_updates)
+                tg.start_soon(self._bootstrap_then_run, info_gatherer, info_recv)
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
             self.event_sender.close()
+            self.snapshot_chunk_receiver.close()
             self.command_sender.close()
             self.download_command_sender.close()
             for runner in self.runners.values():
                 runner.shutdown()
             self._stopped.set()
 
+    async def _bootstrap_then_run(
+        self, info_gatherer: InfoGatherer, info_recv: Receiver[GatheredInfo]
+    ) -> None:
+        await self._fetch_snapshot()
+        self._sync_input_views_from_state()
+        self._tg.start_soon(info_gatherer.run)
+        self._tg.start_soon(self._forward_info, info_recv)
+        self._tg.start_soon(self.plan_step)
+        self._tg.start_soon(self._event_applier)
+        self._tg.start_soon(self._poll_connection_updates)
+
+    async def _fetch_snapshot(self) -> None:
+        receiver = SnapshotReceiver(self.node_id, self.session_id)
+        await self.command_sender.send(
+            ForwarderCommand(
+                origin=self._system_id,
+                command=RequestSnapshot(requester_node_id=self.node_id),
+            )
+        )
+
+        with anyio.move_on_after(_SNAPSHOT_FETCH_TIMEOUT_SECONDS):
+            with self.snapshot_chunk_receiver as chunks:
+                async for chunk in chunks:
+                    received = receiver.ingest(chunk)
+                    if received is None:
+                        continue
+                    self.state = received.state
+                    self.event_router.set_buffer_start(
+                        received.last_event_applied_idx + 1
+                    )
+                    logger.info(
+                        f"Worker bootstrapped from snapshot at idx "
+                        f"{received.last_event_applied_idx}"
+                    )
+                    return
+        logger.info(
+            "No snapshot received before timeout; falling back to full event-log replay"
+        )
+
     async def _forward_info(self, recv: Receiver[GatheredInfo]):
         with recv as info_stream:
             async for info in info_stream:
@@ -133,6 +181,8 @@ class Worker:
     async def _event_applier(self):
         with self.event_receiver as events:
             async for event in events:
+                if event.idx <= self.state.last_event_applied_idx:
+                    continue
                 # 2. for each event, apply it to the state
                 self.state = apply(self.state, event=event)
                 event = event.event

src/exo/worker/tests/unittests/test_worker_snapshot_bootstrap.py

@@ -0,0 +1,126 @@
+# pyright: reportPrivateUsage=false
+
+import hashlib
+
+import anyio
+import pytest
+import zstandard
+
+from exo.routing.event_router import EventRouter
+from exo.shared.types.commands import (
+    ForwarderCommand,
+    ForwarderDownloadCommand,
+    RequestSnapshot,
+)
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.events import (
+    Event,
+    GlobalForwarderEvent,
+    IndexedEvent,
+    LocalForwarderEvent,
+    TestEvent,
+)
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+from exo.shared.types.state import State
+from exo.utils.channels import Receiver, Sender, channel
+from exo.worker.main import Worker
+
+
+def _snapshot_chunk(
+    state: State, *, requester_node_id: NodeId, session_id: SessionId
+) -> SnapshotChunk:
+    body = zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
+    return SnapshotChunk.from_data(
+        data=body,
+        transfer_id=SnapshotTransferId("transfer-1"),
+        requester_node_id=requester_node_id,
+        session_id=session_id,
+        schema_version=state.schema_version,
+        last_event_applied_idx=state.last_event_applied_idx,
+        chunk_index=0,
+        total_chunks=1,
+        sha256_hex=hashlib.sha256(body).hexdigest(),
+    )
+
+
+def _worker(
+    node_id: NodeId, session_id: SessionId
+) -> tuple[
+    Worker,
+    EventRouter,
+    Receiver[ForwarderCommand],
+    Sender[SnapshotChunk],
+    Sender[IndexedEvent],
+]:
+    router_command_sender, _router_command_receiver = channel[ForwarderCommand]()
+    _global_event_sender, global_event_receiver = channel[GlobalForwarderEvent]()
+    local_event_sender, _local_event_receiver = channel[LocalForwarderEvent]()
+    event_router = EventRouter(
+        session_id=session_id,
+        command_sender=router_command_sender,
+        external_inbound=global_event_receiver,
+        external_outbound=local_event_sender,
+    )
+
+    event_sender, event_receiver = channel[IndexedEvent]()
+    local_event_output_sender, _local_event_output_receiver = channel[Event]()
+    command_sender, command_receiver = channel[ForwarderCommand]()
+    download_command_sender, _download_command_receiver = channel[
+        ForwarderDownloadCommand
+    ]()
+    snapshot_sender, snapshot_receiver = channel[SnapshotChunk]()
+
+    worker = Worker(
+        node_id,
+        session_id,
+        event_router=event_router,
+        event_receiver=event_receiver,
+        event_sender=local_event_output_sender,
+        snapshot_chunk_receiver=snapshot_receiver,
+        command_sender=command_sender,
+        download_command_sender=download_command_sender,
+        api_port=52415,
+    )
+    return worker, event_router, command_receiver, snapshot_sender, event_sender
+
+
+@pytest.mark.asyncio
+async def test_worker_fetch_snapshot_applies_state_and_fast_forwards_router() -> None:
+    node_id = NodeId("worker")
+    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
+    worker, event_router, command_receiver, snapshot_sender, _event_sender = _worker(
+        node_id, session_id
+    )
+    state = State(last_event_applied_idx=7)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(worker._fetch_snapshot)
+        command = await command_receiver.receive()
+        assert isinstance(command.command, RequestSnapshot)
+        assert command.command.requester_node_id == node_id
+
+        await snapshot_sender.send(
+            _snapshot_chunk(state, requester_node_id=node_id, session_id=session_id)
+        )
+
+    assert worker.state.last_event_applied_idx == 7
+    assert event_router.event_buffer.next_idx_to_release == 8
+
+
+@pytest.mark.asyncio
+async def test_worker_event_applier_ignores_events_covered_by_snapshot() -> None:
+    node_id = NodeId("worker")
+    session_id = SessionId(master_node_id=NodeId("master"), election_clock=1)
+    worker, _event_router, _command_receiver, _snapshot_sender, event_sender = _worker(
+        node_id, session_id
+    )
+    worker.state = State(last_event_applied_idx=7)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(worker._event_applier)
+        await event_sender.send(IndexedEvent(idx=7, event=TestEvent()))
+        await event_sender.send(IndexedEvent(idx=8, event=TestEvent()))
+
+        while worker.state.last_event_applied_idx != 8:
+            await anyio.sleep(0.001)
+        tg.cancel_scope.cancel()