feat: serve state snapshots from master

src/exo/api/main.py

@@ -254,7 +254,6 @@ class API:
         self.node_id: NodeId = node_id
         self.last_completed_election: int = 0
         self.port = port
-        self._sent_image_hashes: set[str] = set()
 
         self.paused: bool = False
         self.paused_ev: anyio.Event = anyio.Event()
@@ -304,7 +303,6 @@ class API:
         self.event_receiver.close()
         self.event_receiver = event_receiver
         self._tg.start_soon(self._apply_state)
-        self._sent_image_hashes = set()
 
     def unpause(self, result_clock: int):
         logger.info("Unpausing API")
@@ -826,18 +824,8 @@ class API:
         )
         command = TextGeneration(task_params=task_params)
 
-        new_images: list[tuple[int, str]] = []
-        for idx, (img, h) in enumerate(zip(images, hashes, strict=True)):
-            if h not in self._sent_image_hashes:
-                self._sent_image_hashes.add(h)
-                new_images.append((idx, img))
-
-        if not new_images:
-            await self._send(command)
-            return command
-
         all_chunks: list[tuple[int, str]] = []
-        for img_idx, img_data in new_images:
+        for img_idx, img_data in enumerate(images):
             for i in range(0, len(img_data), EXO_MAX_CHUNK_SIZE):
                 all_chunks.append((img_idx, img_data[i : i + EXO_MAX_CHUNK_SIZE]))
 

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),
@@ -112,6 +113,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 +212,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
                         ),

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/event_router.py

@@ -80,6 +80,9 @@ class EventRouter:
     def shutdown(self) -> None:
         self._tg.cancel_tasks()
 
+    def set_buffer_start(self, idx: int) -> None:
+        self.event_buffer.fast_forward_to(idx)
+
     async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
         idx = 0
         with recv as events:
@@ -95,7 +98,6 @@ class EventRouter:
                 self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
 
     async def _run_ext_in(self):
-        buf = OrderedBuffer[Event]()
         with self.external_inbound as events:
             async for event in events:
                 if event.session != self.session_id:
@@ -103,12 +105,12 @@ class EventRouter:
                 if event.origin != self.session_id.master_node_id:
                     continue
 
-                buf.ingest(event.origin_idx, event.event)
+                self.event_buffer.ingest(event.origin_idx, event.event)
                 event_id = event.event.event_id
                 if event_id in self.out_for_delivery:
                     self.out_for_delivery.pop(event_id)
 
-                drained = buf.drain_indexed()
+                drained = self.event_buffer.drain_indexed()
                 if drained:
                     self._nack_attempts = 0
                     if self._nack_cancel_scope:
@@ -119,7 +121,9 @@ class EventRouter:
                     or self._nack_cancel_scope.cancel_called
                 ):
                     # Request the next index.
-                    self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
+                    self._tg.start_soon(
+                        self._nack_request, self.event_buffer.next_idx_to_release
+                    )
                     continue
 
                 for idx, event in drained:

src/exo/routing/snapshot_receiver.py

@@ -0,0 +1,109 @@
+"""Reassembles a snapshot from a stream of `SnapshotChunk`s.
+
+A receiver belongs to one node; it ignores chunks addressed to other
+requesters and chunks from prior sessions. Once a transfer's chunks have
+all been collected and the SHA-256 checks out, the snapshot is decoded into
+a `State`. Concurrent transfers (for the same requester) are tolerated:
+each is keyed by `transfer_id`.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, field
+from typing import final
+
+import zstandard
+from loguru import logger
+
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+from exo.shared.types.state import State
+
+
+@final
+@dataclass
+class _Assembly:
+    """Partial state for one in-flight snapshot transfer."""
+
+    total_chunks: int
+    sha256_hex: str
+    schema_version: int
+    last_event_applied_idx: int
+    chunks: dict[int, bytes] = field(default_factory=dict)
+
+    def is_complete(self) -> bool:
+        return len(self.chunks) == self.total_chunks
+
+    def assemble(self) -> bytes:
+        return b"".join(self.chunks[i] for i in range(self.total_chunks))
+
+
+@dataclass
+class ReceivedSnapshot:
+    last_event_applied_idx: int
+    state: State
+
+
+class SnapshotReceiver:
+    """Filters and reassembles inbound chunks into a `ReceivedSnapshot`.
+
+    Stateless w.r.t. delivery: callers feed `SnapshotChunk`s in via `ingest`
+    and check the return value for completion.
+    """
+
+    def __init__(self, my_node_id: NodeId, session_id: SessionId) -> None:
+        self._my_node_id = my_node_id
+        self._session_id = session_id
+        self._assemblies: dict[SnapshotTransferId, _Assembly] = {}
+
+    def ingest(self, chunk: SnapshotChunk) -> ReceivedSnapshot | None:
+        """Absorb a chunk; return the snapshot once a transfer completes.
+
+        Returns None for partial transfers, mismatched recipients, stale
+        sessions, version mismatches, or corrupt payloads.
+        """
+        if chunk.requester_node_id != self._my_node_id:
+            return None
+        if chunk.session_id != self._session_id:
+            return None
+
+        existing = self._assemblies.get(chunk.transfer_id)
+        if existing is None:
+            existing = _Assembly(
+                total_chunks=chunk.total_chunks,
+                sha256_hex=chunk.sha256_hex,
+                schema_version=chunk.schema_version,
+                last_event_applied_idx=chunk.last_event_applied_idx,
+            )
+            self._assemblies[chunk.transfer_id] = existing
+        existing.chunks[chunk.chunk_index] = chunk.data
+
+        if not existing.is_complete():
+            return None
+
+        # Transfer complete — finalise and remove from the in-flight map.
+        del self._assemblies[chunk.transfer_id]
+        body = existing.assemble()
+        if hashlib.sha256(body).hexdigest() != existing.sha256_hex:
+            logger.warning(f"Snapshot {chunk.transfer_id} failed checksum; discarding")
+            return None
+        try:
+            decompressed = zstandard.ZstdDecompressor().decompress(body)
+            state = State.model_validate_json(decompressed.decode("utf-8"))
+        except (zstandard.ZstdError, ValueError) as e:
+            logger.opt(exception=e).warning(
+                f"Snapshot {chunk.transfer_id} could not be decoded; discarding"
+            )
+            return None
+        if state.schema_version != existing.schema_version:
+            # Should not happen — the master writes schema_version into both
+            # the chunk meta and the State payload — but treat it as corrupt.
+            logger.warning(
+                f"Snapshot {chunk.transfer_id} schema version mismatch "
+                f"(chunk={existing.schema_version}, state={state.schema_version})"
+            )
+            return None
+        return ReceivedSnapshot(
+            last_event_applied_idx=existing.last_event_applied_idx, state=state
+        )

src/exo/routing/tests/test_event_buffer.py

@@ -141,3 +141,28 @@ async def test_drain_and_ingest_with_new_sequence(buffer: OrderedBuffer[Event]):
     assert [e[0] for e in drained] == [2]
     assert buffer.next_idx_to_release == 3
     assert 4 in buffer.store
+
+
+@pytest.mark.asyncio
+async def test_fast_forward_discards_buffered_stale_events(
+    buffer: OrderedBuffer[Event],
+):
+    buffer.ingest(*make_indexed_event(0))
+    buffer.ingest(*make_indexed_event(2))
+    buffer.ingest(*make_indexed_event(4))
+
+    buffer.fast_forward_to(3)
+
+    assert buffer.next_idx_to_release == 3
+    assert set(buffer.store) == {4}
+
+
+@pytest.mark.asyncio
+async def test_fast_forward_only_moves_forward(buffer: OrderedBuffer[Event]):
+    buffer.ingest(*make_indexed_event(0))
+    buffer.ingest(*make_indexed_event(1))
+    buffer.drain()
+
+    buffer.fast_forward_to(1)
+
+    assert buffer.next_idx_to_release == 2

src/exo/routing/tests/test_snapshot_receiver.py

@@ -0,0 +1,151 @@
+import hashlib
+
+import pytest
+import zstandard
+
+from exo.routing.snapshot_receiver import SnapshotReceiver
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.snapshots import SnapshotChunk, SnapshotTransferId
+from exo.shared.types.state import State
+
+
+@pytest.fixture
+def session_id() -> SessionId:
+    return SessionId(master_node_id=NodeId("master"), election_clock=0)
+
+
+@pytest.fixture
+def my_node() -> NodeId:
+    return NodeId("worker-1")
+
+
+def _encode(state: State) -> bytes:
+    return zstandard.ZstdCompressor().compress(state.model_dump_json().encode("utf-8"))
+
+
+def _make_chunks(
+    body: bytes,
+    *,
+    chunk_size: int,
+    requester_node_id: NodeId,
+    session_id: SessionId,
+    state: State,
+    transfer_id: SnapshotTransferId | None = None,
+) -> list[SnapshotChunk]:
+    sha256 = hashlib.sha256(body).hexdigest()
+    transfer_id = transfer_id or SnapshotTransferId()
+    pieces = [body[i : i + chunk_size] for i in range(0, len(body), chunk_size)] or [
+        b""
+    ]
+    return [
+        SnapshotChunk.from_data(
+            data=piece,
+            transfer_id=transfer_id,
+            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=i,
+            total_chunks=len(pieces),
+            sha256_hex=sha256,
+        )
+        for i, piece in enumerate(pieces)
+    ]
+
+
+def test_completes_on_full_transfer(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=42)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=my_node,
+        session_id=session_id,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+
+    received = None
+    for chunk in chunks:
+        received = receiver.ingest(chunk)
+    assert received is not None
+    assert received.last_event_applied_idx == 42
+    assert received.state.last_event_applied_idx == 42
+
+
+def test_handles_out_of_order_chunks(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=99)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=32,
+        requester_node_id=my_node,
+        session_id=session_id,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+
+    # Reverse them.
+    received = None
+    for chunk in reversed(chunks):
+        received = receiver.ingest(chunk)
+    assert received is not None
+    assert received.last_event_applied_idx == 99
+
+
+def test_ignores_chunks_for_other_recipients(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=1)
+    other = NodeId("worker-2")
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=other,
+        session_id=session_id,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+    for chunk in chunks:
+        assert receiver.ingest(chunk) is None
+
+
+def test_ignores_chunks_from_stale_session(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=1)
+    other_session = SessionId(master_node_id=NodeId("other-master"), election_clock=99)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=my_node,
+        session_id=other_session,
+        state=state,
+    )
+    receiver = SnapshotReceiver(my_node, session_id)
+    for chunk in chunks:
+        assert receiver.ingest(chunk) is None
+
+
+def test_discards_on_checksum_mismatch(my_node: NodeId, session_id: SessionId):
+    state = State(last_event_applied_idx=1)
+    chunks = _make_chunks(
+        _encode(state),
+        chunk_size=64,
+        requester_node_id=my_node,
+        session_id=session_id,
+        state=state,
+    )
+    # Corrupt the last byte of the last chunk.
+    original = chunks[-1]
+    chunks[-1] = SnapshotChunk.from_data(
+        data=original.data + b"\x00garbage",
+        transfer_id=original.transfer_id,
+        requester_node_id=original.requester_node_id,
+        session_id=original.session_id,
+        schema_version=original.schema_version,
+        last_event_applied_idx=original.last_event_applied_idx,
+        chunk_index=original.chunk_index,
+        total_chunks=original.total_chunks,
+        sha256_hex=original.sha256_hex,
+    )
+
+    receiver = SnapshotReceiver(my_node, session_id)
+    received = None
+    for chunk in chunks:
+        received = receiver.ingest(chunk)
+    assert received is None

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/apply.py

@@ -40,7 +40,14 @@ from exo.shared.types.profiling import (
     ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
+from exo.shared.types.tasks import (
+    ImageEdits,
+    ImageGeneration,
+    Task,
+    TaskId,
+    TaskStatus,
+    TextGeneration,
+)
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -72,7 +79,6 @@ def event_apply(event: Event, state: State) -> State:
             TestEvent()
             | ChunkGenerated()
             | TaskAcknowledged()
-            | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
             | CustomModelCardAdded()
@@ -93,6 +99,8 @@ def event_apply(event: Event, state: State) -> State:
             return apply_runner_status_updated(event, state)
         case TaskCreated():
             return apply_task_created(event, state)
+        case InputChunkReceived():
+            return apply_input_chunk_received(event, state)
         case TaskDeleted():
             return apply_task_deleted(event, state)
         case TaskFailed():
@@ -157,10 +165,32 @@ def apply_task_created(event: TaskCreated, state: State) -> State:
     return state.model_copy(update={"tasks": new_tasks})
 
 
+def apply_input_chunk_received(event: InputChunkReceived, state: State) -> State:
+    command_chunks = {
+        **state.input_chunks.get(event.command_id, {}),
+        event.chunk.chunk_index: event.chunk,
+    }
+    return state.model_copy(
+        update={
+            "input_chunks": {**state.input_chunks, event.command_id: command_chunks}
+        }
+    )
+
+
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
+    task = state.tasks.get(event.task_id)
     new_tasks: Mapping[TaskId, Task] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
+    if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
+        new_input_chunks = {
+            command_id: chunks
+            for command_id, chunks in state.input_chunks.items()
+            if command_id != task.command_id
+        }
+        return state.model_copy(
+            update={"tasks": new_tasks, "input_chunks": new_input_chunks}
+        )
     return state.model_copy(update={"tasks": new_tasks})
 
 

src/exo/shared/tests/test_apply/test_apply_input_chunks.py

@@ -0,0 +1,85 @@
+from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.chunks import InputImageChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.events import (
+    IndexedEvent,
+    InputChunkReceived,
+    TaskCreated,
+    TaskDeleted,
+)
+from exo.shared.types.state import State
+from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
+from exo.shared.types.text_generation import (
+    InputMessage,
+    InputMessageContent,
+    TextGenerationTaskParams,
+)
+from exo.shared.types.worker.instances import InstanceId
+
+
+def test_apply_input_chunk_received_stores_chunk_in_state() -> None:
+    command_id = CommandId("command-1")
+    chunk = InputImageChunk(
+        model=ModelId("mlx-community/test-model"),
+        command_id=command_id,
+        data="abc",
+        chunk_index=0,
+        total_chunks=1,
+        image_index=0,
+    )
+
+    state = apply(
+        State(),
+        IndexedEvent(
+            idx=0,
+            event=InputChunkReceived(command_id=command_id, chunk=chunk),
+        ),
+    )
+
+    assert state.input_chunks == {command_id: {0: chunk}}
+
+
+def test_apply_task_deleted_removes_chunks_for_generation_command() -> None:
+    command_id = CommandId("command-1")
+    task_id = TaskId("task-1")
+    chunk = InputImageChunk(
+        model=ModelId("mlx-community/test-model"),
+        command_id=command_id,
+        data="abc",
+        chunk_index=0,
+        total_chunks=1,
+        image_index=0,
+    )
+    task = TextGeneration(
+        task_id=task_id,
+        instance_id=InstanceId("instance-1"),
+        task_status=TaskStatus.Pending,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=ModelId("mlx-community/test-model"),
+            input=[
+                InputMessage(role="user", content=InputMessageContent("hello")),
+            ],
+        ),
+    )
+
+    state = State()
+    state = apply(
+        state,
+        IndexedEvent(
+            idx=0,
+            event=InputChunkReceived(command_id=command_id, chunk=chunk),
+        ),
+    )
+    state = apply(
+        state,
+        IndexedEvent(idx=1, event=TaskCreated(task_id=task_id, task=task)),
+    )
+    state = apply(
+        state,
+        IndexedEvent(idx=2, event=TaskDeleted(task_id=task_id)),
+    )
+
+    assert state.tasks == {}
+    assert state.input_chunks == {}

src/exo/shared/tests/test_state_serialization.py

@@ -25,6 +25,7 @@ def test_state_serialization_roundtrip() -> None:
     json_repr = state.model_dump_json()
     restored_state = State.model_validate_json(json_repr)
 
+    assert restored_state.schema_version == state.schema_version
     assert (
         state.topology.to_snapshot().nodes
         == restored_state.topology.to_snapshot().nodes

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/shared/types/snapshots.py

@@ -0,0 +1,54 @@
+"""Wire types for snapshot transfer between master and a joining node.
+
+Snapshots can be tens of MB; the gossipsub message ceiling is around 1 MB.
+We slice the compressed snapshot body into chunks and publish each chunk on
+the SNAPSHOT_RESPONSES topic. The receiver collects chunks for its own
+`requester_node_id`, validates the SHA-256 of the reassembled body, and
+materialises the State.
+"""
+
+import base64
+
+from exo.shared.types.common import Id, NodeId, SessionId
+from exo.utils.pydantic_ext import FrozenModel
+
+
+class SnapshotTransferId(Id):
+    """Identifies a single snapshot transfer (one master response to one
+    `RequestSnapshot`). Distinct transfers may interleave; the id lets
+    receivers keep them apart."""
+
+
+class SnapshotChunk(FrozenModel):
+    """One slice of a snapshot in flight.
+
+    `data_b64` carries a base64-encoded slice of the zstd-compressed JSON
+    dump of State. Concatenating the *decoded* bytes of all chunks for a
+    `transfer_id` in order of `chunk_index` yields the full compressed
+    body; `sha256_hex` is the SHA-256 of that decoded blob.
+
+    We use base64 explicitly because the topic layer JSON-encodes messages,
+    and JSON can't carry raw binary. Helpers `from_data` / `data` keep the
+    base64 detail at the boundaries.
+    """
+
+    transfer_id: SnapshotTransferId
+    requester_node_id: NodeId
+    session_id: SessionId
+    schema_version: int
+    last_event_applied_idx: int
+    chunk_index: int
+    total_chunks: int
+    sha256_hex: str
+    data_b64: str
+
+    @classmethod
+    def from_data(cls, *, data: bytes, **kwargs: object) -> "SnapshotChunk":
+        return cls(data_b64=base64.b64encode(data).decode("ascii"), **kwargs)  # pyright: ignore[reportArgumentType]
+
+    @property
+    def data(self) -> bytes:
+        return base64.b64decode(self.data_b64)
+
+
+__all__ = ["SnapshotChunk", "SnapshotTransferId"]

src/exo/shared/types/state.py

@@ -6,7 +6,8 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import NodeId
+from exo.shared.types.chunks import InputImageChunk
+from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -41,10 +42,16 @@ class State(FrozenModel):
         strict=True,
         arbitrary_types_allowed=True,
     )
+    # Bump when a State change makes older snapshots unsafe to restore.
+    schema_version: int = Field(default=1, ge=1)
+
     instances: Mapping[InstanceId, Instance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
+    # Durable request input chunks for active image requests. Workers rebuild
+    # local image caches from this state instead of reading events directly.
+    input_chunks: Mapping[CommandId, Mapping[int, InputImageChunk]] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)

src/exo/utils/event_buffer.py

@@ -47,6 +47,18 @@ class OrderedBuffer[T]:
         logger.trace(f"Releasing event {ret}")
         return ret
 
+    def fast_forward_to(self, idx: int) -> None:
+        """Skip every event before idx.
+
+        Snapshot restore uses this after applying state that already includes
+        events before idx. Any buffered or future event below idx is stale.
+        """
+        if idx <= self.next_idx_to_release:
+            return
+        self.next_idx_to_release = idx
+        for stale_idx in [i for i in self.store if i < idx]:
+            del self.store[stale_idx]
+
 
 class MultiSourceBuffer[SourceId, T]:
     """

src/exo/worker/main.py

@@ -24,7 +24,6 @@ from exo.shared.types.events import (
     CustomModelCardDeleted,
     Event,
     IndexedEvent,
-    InputChunkReceived,
     InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
@@ -141,37 +140,6 @@ class Worker:
                 if isinstance(event, InstanceDeleted):
                     self._instance_backoff.reset(event.instance_id)
 
-                # Buffer input image chunks for image editing
-                if isinstance(event, InputChunkReceived):
-                    cmd_id = event.command_id
-                    if cmd_id not in self.input_chunk_buffer:
-                        self.input_chunk_buffer[cmd_id] = {}
-                        self.input_chunk_counts[cmd_id] = event.chunk.total_chunks
-
-                    self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
-                        event.chunk
-                    )
-
-                    if (
-                        len(self.input_chunk_buffer[cmd_id])
-                        == self.input_chunk_counts[cmd_id]
-                    ):
-                        per_image: defaultdict[int, list[InputImageChunk]] = (
-                            defaultdict(list)
-                        )
-                        for chunk in self.input_chunk_buffer[cmd_id].values():
-                            per_image[chunk.image_index].append(chunk)
-                        for chunks_for_image in per_image.values():
-                            sorted_chunks = sorted(
-                                chunks_for_image, key=lambda c: c.chunk_index
-                            )
-                            img = Base64Image("".join(c.data for c in sorted_chunks))
-                            self.image_cache[
-                                Base64ImageHash(
-                                    hashlib.sha256(img.encode("ascii")).hexdigest()
-                                )
-                            ] = img
-
                 if isinstance(event, CustomModelCardAdded):
                     await event.model_card.save_to_custom_dir()
                     add_to_card_cache(event.model_card)
@@ -179,6 +147,35 @@ class Worker:
                 if isinstance(event, CustomModelCardDeleted):
                     await delete_custom_card(event.model_id)
 
+                self._sync_input_views_from_state()
+
+    def _sync_input_views_from_state(self) -> None:
+        self.input_chunk_buffer = {
+            command_id: dict(chunks)
+            for command_id, chunks in self.state.input_chunks.items()
+        }
+        self.input_chunk_counts = {
+            command_id: next(iter(chunks.values())).total_chunks
+            for command_id, chunks in self.input_chunk_buffer.items()
+            if chunks
+        }
+
+        self.image_cache = {}
+        for command_id, chunks in self.input_chunk_buffer.items():
+            expected_chunks = self.input_chunk_counts.get(command_id)
+            if expected_chunks is None or len(chunks) != expected_chunks:
+                continue
+
+            per_image: defaultdict[int, list[InputImageChunk]] = defaultdict(list)
+            for chunk in chunks.values():
+                per_image[chunk.image_index].append(chunk)
+            for chunks_for_image in per_image.values():
+                sorted_chunks = sorted(chunks_for_image, key=lambda c: c.chunk_index)
+                image = Base64Image("".join(chunk.data for chunk in sorted_chunks))
+                self.image_cache[
+                    Base64ImageHash(hashlib.sha256(image.encode("ascii")).hexdigest())
+                ] = image
+
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
@@ -189,7 +186,7 @@ class Worker:
                 self.state.instances,
                 self.state.runners,
                 self.state.tasks,
-                self.input_chunk_buffer,
+                self.state.input_chunks,
                 self.image_cache,
                 self._instance_backoff,
                 self._download_backoff,
@@ -321,15 +318,9 @@ class Worker:
                             advanced_params=task.task_params.advanced_params,
                         ),
                     )
-                    # Cleanup buffers
-                    if cmd_id in self.input_chunk_buffer:
-                        del self.input_chunk_buffer[cmd_id]
-                    if cmd_id in self.input_chunk_counts:
-                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
 
                 case TextGeneration() if task.task_params.image_hashes:
-                    cmd_id = task.command_id
                     resolved_images = [
                         self.image_cache[h]
                         for _, h in sorted(task.task_params.image_hashes.items())
@@ -341,10 +332,6 @@ class Worker:
                             )
                         }
                     )
-                    if cmd_id in self.input_chunk_buffer:
-                        del self.input_chunk_buffer[cmd_id]
-                    if cmd_id in self.input_chunk_counts:
-                        del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
                 case LoadModel(instance_id=instance_id):
                     if (instance := self.state.instances.get(instance_id)) is not None: