Compare commits

...
13 changed files with 879 additions and 12 deletions

No files matched your search

+62 -4
View File
@@ -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
@@ -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
+23 -1
View File
@@ -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:
+60 -1
View File
@@ -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
+54
View File
@@ -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()
+109
View File
@@ -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
)
@@ -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
@@ -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"
+4
View File
@@ -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
)
+7
View File
@@ -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
+54
View File
@@ -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"]
+56 -6
View File
@@ -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
@@ -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()