From ed001f2409ef671a9d398d4aecdf259c9957cfb3 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Thu, 19 Feb 2026 18:23:28 +0000 Subject: [PATCH] remove prefillprogress event (#1550) this should never have been a separate event, but i didnt quite communicate that well when this was merged. convert PrefillProgress to a chunk like the rest of the runner responses. tested with Llama-3.3-70B, prefill progress events still show up in the dashboard as usual --- src/exo/master/api.py | 17 ----------------- src/exo/shared/apply.py | 2 -- src/exo/shared/types/events.py | 10 +--------- src/exo/worker/runner/runner.py | 19 +++++++++++++------ 4 files changed, 14 insertions(+), 34 deletions(-) diff --git a/src/exo/master/api.py b/src/exo/master/api.py index c3811072..319a8164 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -138,7 +138,6 @@ from exo.shared.types.events import ( Event, ForwarderEvent, IndexedEvent, - PrefillProgress, TracesMerged, ) from exo.shared.types.memory import Memory @@ -1455,22 +1454,6 @@ class API: await queue.send(event.chunk) except BrokenResourceError: self._text_generation_queues.pop(event.command_id, None) - - elif isinstance(event, PrefillProgress): - if queue := self._text_generation_queues.get( - event.command_id, None - ): - try: - await queue.send( - PrefillProgressChunk( - model=event.model, - processed_tokens=event.processed_tokens, - total_tokens=event.total_tokens, - ) - ) - except BrokenResourceError: - self._text_generation_queues.pop(event.command_id, None) - if isinstance(event, TracesMerged): self._save_merged_trace(event) diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index ee5b2229..94869dfe 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -15,7 +15,6 @@ from exo.shared.types.events import ( NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, - PrefillProgress, RunnerDeleted, RunnerStatusUpdated, TaskAcknowledged, @@ -65,7 +64,6 @@ def event_apply(event: Event, state: State) -> State: | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived() - | PrefillProgress() | TracesCollected() | TracesMerged() ): # Pass-through events that don't modify state diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 83e162ab..5cf93d0c 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -5,7 +5,7 @@ from pydantic import Field from exo.shared.topology import Connection from exo.shared.types.chunks import GenerationChunk, InputImageChunk -from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId +from exo.shared.types.common import CommandId, Id, NodeId, SessionId from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId @@ -102,13 +102,6 @@ class InputChunkReceived(BaseEvent): chunk: InputImageChunk -class PrefillProgress(BaseEvent): - command_id: CommandId - model: ModelId - processed_tokens: int - total_tokens: int - - class TopologyEdgeCreated(BaseEvent): conn: Connection @@ -155,7 +148,6 @@ Event = ( | NodeDownloadProgress | ChunkGenerated | InputChunkReceived - | PrefillProgress | TopologyEdgeCreated | TopologyEdgeDeleted | TracesCollected diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 2dc340c3..51950192 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -21,12 +21,17 @@ from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED from exo.shared.models.model_cards import ModelId, ModelTask from exo.shared.tracing import clear_trace_buffer, get_trace_buffer from exo.shared.types.api import ImageGenerationStats -from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk +from exo.shared.types.chunks import ( + ErrorChunk, + ImageChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) from exo.shared.types.common import CommandId from exo.shared.types.events import ( ChunkGenerated, Event, - PrefillProgress, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, @@ -315,11 +320,13 @@ def main( ) -> None: if device_rank == 0: event_sender.send( - PrefillProgress( + ChunkGenerated( command_id=command_id, - model=shard_metadata.model_card.model_id, - processed_tokens=processed, - total_tokens=total, + chunk=PrefillProgressChunk( + model=shard_metadata.model_card.model_id, + processed_tokens=processed, + total_tokens=total, + ), ) ) cancelled_tasks.update(cancel_receiver.collect())