[object Object]

← back to Exo

fix(worker): emit error chunks when a runner dies mid-command (#1645)

848580504275b05290f7b35d60c25040339c4234 · 2026-03-04 18:15:58 +0100 · Miguel Miranda Dias

Closes #1586

## Summary
- track in-flight tasks in `RunnerSupervisor` (not only unacknowledged
pending tasks)
- when `_check_runner()` detects a crashed runner, emit
`ChunkGenerated(ErrorChunk)` for each in-flight command task
(`TextGeneration`, `ImageGeneration`, `ImageEdits`)
- keep existing `RunnerStatusUpdated(RunnerFailed)` emission so
planner/state still transition correctly
- add a unit test for supervisor crash path to ensure an error chunk is
emitted before failed runner status

## Why
`#1586` reports streams that can hang forever when runners crash during
warmup/loading. This keeps failure signaling at the runner-supervisor
layer, matching maintainer guidance in the issue thread.

## Validation
- attempted: `uv run pytest
src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py`
- blocked locally by environment disk exhaustion while uv tried to
materialize heavy CUDA wheels (`No space left on device` during
`nvidia-cudnn-cu13` extraction)

I kept the change scoped and added a targeted unit test for the failure
path.

---------

Co-authored-by: Evan <evanev7@gmail.com>

Files touched

Diff

commit 848580504275b05290f7b35d60c25040339c4234
Author: Miguel Miranda Dias <7780875+pandego@users.noreply.github.com>
Date:   Wed Mar 4 18:15:58 2026 +0100

    fix(worker): emit error chunks when a runner dies mid-command (#1645)
    
    Closes #1586
    
    ## Summary
    - track in-flight tasks in `RunnerSupervisor` (not only unacknowledged
    pending tasks)
    - when `_check_runner()` detects a crashed runner, emit
    `ChunkGenerated(ErrorChunk)` for each in-flight command task
    (`TextGeneration`, `ImageGeneration`, `ImageEdits`)
    - keep existing `RunnerStatusUpdated(RunnerFailed)` emission so
    planner/state still transition correctly
    - add a unit test for supervisor crash path to ensure an error chunk is
    emitted before failed runner status
    
    ## Why
    `#1586` reports streams that can hang forever when runners crash during
    warmup/loading. This keeps failure signaling at the runner-supervisor
    layer, matching maintainer guidance in the issue thread.
    
    ## Validation
    - attempted: `uv run pytest
    src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py`
    - blocked locally by environment disk exhaustion while uv tried to
    materialize heavy CUDA wheels (`No space left on device` during
    `nvidia-cudnn-cu13` extraction)
    
    I kept the change scoped and added a targeted unit test for the failure
    path.
    
    ---------
    
    Co-authored-by: Evan <evanev7@gmail.com>
---
 src/exo/worker/runner/runner_supervisor.py         | 34 +++++++-
 .../test_runner/test_runner_supervisor.py          | 94 +++++++++++++++++++++-
 2 files changed, 123 insertions(+), 5 deletions(-)

diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index 6bea66ce..8ada028d 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -12,13 +12,22 @@ from anyio import (
 )
 from loguru import logger
 
+from exo.shared.types.chunks import ErrorChunk
 from exo.shared.types.events import (
+    ChunkGenerated,
     Event,
     RunnerStatusUpdated,
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
+from exo.shared.types.tasks import (
+    ImageEdits,
+    ImageGeneration,
+    Task,
+    TaskId,
+    TaskStatus,
+    TextGeneration,
+)
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import (
     RunnerConnecting,
@@ -52,7 +61,7 @@ class RunnerSupervisor:
     _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
-    in_progress: set[TaskId] = field(default_factory=set, init=False)
+    in_progress: dict[TaskId, Task] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
     cancelled: set[TaskId] = field(default_factory=set, init=False)
     _cancel_watch_runner: anyio.CancelScope = field(
@@ -148,9 +157,11 @@ class RunnerSupervisor:
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
+        self.in_progress[task.task_id] = task
         try:
             await self._task_sender.send_async(task)
         except ClosedResourceError:
+            self.in_progress.pop(task.task_id, None)
             logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()
@@ -181,7 +192,6 @@ class RunnerSupervisor:
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
                         self.pending.pop(event.task_id).set()
-                        self.in_progress.add(event.task_id)
                         continue
                     if (
                         isinstance(event, TaskStatusUpdated)
@@ -198,7 +208,7 @@ class RunnerSupervisor:
                                 RunnerShuttingDown,
                             ),
                         )
-                        self.in_progress.discard(event.task_id)
+                        self.in_progress.pop(event.task_id, None)
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
         except (ClosedResourceError, BrokenResourceError) as e:
@@ -243,6 +253,22 @@ class RunnerSupervisor:
 
         logger.opt(exception=e).error(f"Runner terminated with {cause}")
 
+        for task in self.in_progress.values():
+            if isinstance(task, (TextGeneration, ImageGeneration, ImageEdits)):
+                with anyio.CancelScope(shield=True):
+                    await self._event_sender.send(
+                        ChunkGenerated(
+                            command_id=task.command_id,
+                            chunk=ErrorChunk(
+                                model=self.shard_metadata.model_card.model_id,
+                                error_message=(
+                                    "Runner shutdown before completing command "
+                                    f"({cause})"
+                                ),
+                            ),
+                        )
+                    )
+
         try:
             self.status = RunnerFailed(error_message=f"Terminated ({cause})")
             with anyio.CancelScope(shield=True):
diff --git a/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py b/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
index e151d4aa..ecb07f94 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
@@ -1 +1,93 @@
-# TODO:
+import multiprocessing as mp
+from typing import cast
+
+import anyio
+import pytest
+
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.chunks import ErrorChunk
+from exo.shared.types.common import CommandId, NodeId
+from exo.shared.types.events import ChunkGenerated, Event, RunnerStatusUpdated
+from exo.shared.types.tasks import Task, TaskId, TextGeneration
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.worker.instances import BoundInstance, InstanceId
+from exo.shared.types.worker.runners import RunnerFailed, RunnerId
+from exo.utils.channels import channel, mp_channel
+from exo.worker.runner.runner_supervisor import RunnerSupervisor
+from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
+
+
+class _DeadProcess:
+    exitcode = -6
+
+    def start(self) -> None:
+        return None
+
+    def is_alive(self) -> bool:
+        return False
+
+    def join(self, _timeout: float | None = None) -> None:
+        return None
+
+    def terminate(self) -> None:
+        return None
+
+    def kill(self) -> None:
+        return None
+
+
+@pytest.mark.asyncio
+async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None:
+    event_sender, event_receiver = channel[Event]()
+    task_sender, _ = mp_channel[Task]()
+    cancel_sender, _ = mp_channel[TaskId]()
+    _, ev_recv = mp_channel[Event]()
+
+    bound_instance: BoundInstance = get_bound_mlx_ring_instance(
+        instance_id=InstanceId("instance-a"),
+        model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
+        runner_id=RunnerId("runner-a"),
+        node_id=NodeId("node-a"),
+    )
+
+    supervisor = RunnerSupervisor(
+        shard_metadata=bound_instance.bound_shard,
+        bound_instance=bound_instance,
+        runner_process=cast("mp.Process", cast(object, _DeadProcess())),
+        initialize_timeout=400,
+        _ev_recv=ev_recv,
+        _task_sender=task_sender,
+        _event_sender=event_sender,
+        _cancel_sender=cancel_sender,
+    )
+
+    command_id = CommandId("cmd-a")
+    task = TextGeneration(
+        task_id=TaskId("task-a"),
+        instance_id=bound_instance.instance.instance_id,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=bound_instance.bound_shard.model_card.model_id,
+            input=[InputMessage(role="user", content="hi")],
+            stream=True,
+        ),
+    )
+    supervisor.in_progress[task.task_id] = task
+    supervisor.shutdown = lambda: None
+
+    await supervisor._check_runner(RuntimeError("boom"))  # pyright: ignore[reportPrivateUsage]
+
+    got_chunk = await event_receiver.receive()
+    got_status = await event_receiver.receive()
+
+    assert isinstance(got_chunk, ChunkGenerated)
+    assert got_chunk.command_id == command_id
+    assert isinstance(got_chunk.chunk, ErrorChunk)
+    assert "Runner shutdown before completing command" in got_chunk.chunk.error_message
+
+    assert isinstance(got_status, RunnerStatusUpdated)
+    assert isinstance(got_status.runner_status, RunnerFailed)
+
+    event_sender.close()
+    with anyio.move_on_after(0.1):
+        await event_receiver.aclose()

← 4de8f801 #Add reasoning parms to chat completion and responses APIs (  ·  back to Exo  ·  Fix copy code button not working in dashboard (#1659) 3a4d635d →