Route traces over transient events

src/exo/api/main.py

@@ -1939,16 +1939,14 @@ class API:
                     continue
                 self._event_log.append(i_event.event)
                 self.state = apply(self.state, i_event)
-                event = i_event.event
-
-                if isinstance(event, TracesMerged):
-                    self._save_merged_trace(event)
 
     async def _apply_transient(self) -> None:
         with self.transient_event_receiver as events:
             async for event in events:
                 if isinstance(event, ChunkGenerated):
                     await self._dispatch_chunk(event)
+                elif isinstance(event, TracesMerged):
+                    self._save_merged_trace(event)
 
     async def _dispatch_chunk(self, event: ChunkGenerated) -> None:
         if queue := self._image_generation_queues.get(event.command_id, None):

src/exo/master/main.py

@@ -169,6 +169,7 @@ class Master:
         try:
             async with self._tg as tg:
                 tg.start_soon(self._event_processor)
+                tg.start_soon(self._transient_event_processor)
                 tg.start_soon(self._command_processor)
                 tg.start_soon(self._plan)
         finally:
@@ -516,10 +517,6 @@ class Master:
                     local_event.origin,
                 )
                 for event in self._multi_buffer.drain():
-                    if isinstance(event, TracesCollected):
-                        await self._handle_traces_collected(event)
-                        continue
-
                     logger.debug(f"Master indexing event: {str(event)[:100]}")
 
                     event = event.model_copy(
@@ -572,6 +569,12 @@ class Master:
                 )
             )
 
+    async def _transient_event_processor(self) -> None:
+        with self.transient_event_receiver as transients:
+            async for event in transients:
+                if isinstance(event, TracesCollected):
+                    await self._handle_traces_collected(event)
+
     # This function is re-entrant, take care!
     async def _send_event(self, event: IndexedEvent):
         # Convenience method since this line is ugly
@@ -602,7 +605,7 @@ class Master:
         for trace_data in self._pending_traces[task_id].values():
             all_trace_data.extend(trace_data)
 
-        await self.event_sender.send(
+        await self.transient_event_sender.send(
             TracesMerged(task_id=task_id, traces=all_trace_data)
         )
 

src/exo/master/tests/test_master.py

@@ -26,6 +26,9 @@ from exo.shared.types.events import (
     LocalForwarderEvent,
     NodeGatheredInfo,
     TaskCreated,
+    TraceEventData,
+    TracesCollected,
+    TracesMerged,
     TransientEvent,
 )
 from exo.shared.types.memory import Memory
@@ -33,7 +36,7 @@ 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 TaskId, TaskStatus
 from exo.shared.types.tasks import TextGeneration as TextGenerationTask
 from exo.shared.types.text_generation import (
     InputMessage,
@@ -290,3 +293,57 @@ async def test_master_serves_snapshot_for_current_state():
 
         await master.shutdown()
         tg.cancel_scope.cancel()
+
+
+@pytest.mark.asyncio
+async def test_master_merges_traces_from_transient_events() -> None:
+    node_id = NodeId("master")
+    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]()
+    transient_input_sender, transient_input_receiver = channel[TransientEvent]()
+    transient_output_sender, transient_output_receiver = channel[TransientEvent]()
+    snapshot_chunk_sender, _snapshot_chunk_receiver = channel[SnapshotChunk]()
+
+    master = Master(
+        node_id,
+        session_id,
+        event_sender=event_sender,
+        transient_event_receiver=transient_input_receiver,
+        transient_event_sender=transient_output_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,
+    )
+
+    task_id = TaskId("task-a")
+    master._expected_ranks[task_id] = {0, 1}  # pyright: ignore[reportPrivateUsage]
+    trace_a = TraceEventData(
+        name="rank-0", start_us=1, duration_us=2, rank=0, category="test"
+    )
+    trace_b = TraceEventData(
+        name="rank-1", start_us=3, duration_us=4, rank=1, category="test"
+    )
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(master.run)
+        await transient_input_sender.send(
+            TracesCollected(task_id=task_id, rank=0, traces=[trace_a])
+        )
+        await transient_input_sender.send(
+            TracesCollected(task_id=task_id, rank=1, traces=[trace_b])
+        )
+
+        merged = await transient_output_receiver.receive()
+        assert isinstance(merged, TracesMerged)
+        assert merged.task_id == task_id
+        assert merged.traces == [trace_a, trace_b]
+        tg.cancel_scope.cancel()

src/exo/worker/runner/supervisor.py

@@ -19,6 +19,8 @@ from exo.shared.types.events import (
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
+    TracesCollected,
+    TracesMerged,
     TransientEvent,
 )
 from exo.shared.types.tasks import (
@@ -241,7 +243,9 @@ class RunnerSupervisor:
                         )
                         self.in_progress.pop(event.task_id, None)
                         self.completed.add(event.task_id)
-                    if isinstance(event, ChunkGenerated):
+                    if isinstance(
+                        event, (ChunkGenerated, TracesCollected, TracesMerged)
+                    ):
                         await self._transient_event_sender.send(event)
                     else:
                         await self._event_sender.send(event)