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

@@ -84,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),
@@ -96,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,
@@ -251,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
@@ -261,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/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()