mirror of
https://github.com/exo-explore/exo.git
synced 2026-09-08 11:35:40 -04:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ea04549692 | ||
|
|
a2c2a0fc92 | ||
|
|
24c56138e3 | ||
|
|
c89caaf87a | ||
|
|
343d5bc6d4 |
No files matched your search
+62
-4
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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()
|
||||
Reference in new issue
Block a user