← back to Exo
re-raise exceptions in the runner (#1198)
d19bf0240441f4ba17a69e5528931c86f66f8f26 · 2026-01-19 10:35:23 +0000 · Evan Quiney
## Motivation
Runners that crash can swallow errors - we should re-raise. Also the
exception handler annoyed me.
## Changes
The try: except in the runner's chat now re-raises.
Files touched
M src/exo/worker/runner/runner.pyD src/exo/worker/tests/unittests/test_runner/test_error_handling.py
Diff
commit d19bf0240441f4ba17a69e5528931c86f66f8f26
Author: Evan Quiney <evanev7@gmail.com>
Date: Mon Jan 19 10:35:23 2026 +0000
re-raise exceptions in the runner (#1198)
## Motivation
Runners that crash can swallow errors - we should re-raise. Also the
exception handler annoyed me.
## Changes
The try: except in the runner's chat now re-raises.
---
src/exo/worker/runner/runner.py | 71 ++++++++--------------
.../unittests/test_runner/test_error_handling.py | 50 ---------------
2 files changed, 27 insertions(+), 94 deletions(-)
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 28f89d54..2b4e5b49 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,8 +1,6 @@
import time
from collections.abc import Generator
-from contextlib import contextmanager
from functools import cache
-from typing import cast
import mlx.core as mx
from mlx_lm.models.gpt_oss import Model as GptOssModel
@@ -15,7 +13,6 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
from exo.shared.types.api import ChatCompletionMessageText
from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.common import CommandId
from exo.shared.types.events import (
ChunkGenerated,
Event,
@@ -23,7 +20,6 @@ from exo.shared.types.events import (
TaskAcknowledged,
TaskStatusUpdated,
)
-from exo.shared.types.models import ModelId
from exo.shared.types.tasks import (
ChatCompletion,
ConnectToGroup,
@@ -52,7 +48,6 @@ from exo.shared.types.worker.runners import (
RunnerWarmingUp,
)
from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
from exo.worker.engines.mlx.utils_mlx import (
initialize_mlx,
@@ -62,33 +57,6 @@ from exo.worker.engines.mlx.utils_mlx import (
from exo.worker.runner.bootstrap import logger
-@contextmanager
-def send_error_chunk_on_exception(
- event_sender: MpSender[Event],
- command_id: CommandId,
- model_id: ModelId,
- device_rank: int,
-):
- try:
- yield
- except Exception as e:
- logger.error(e)
- if device_rank == 0:
- event_sender.send(
- ChunkGenerated(
- command_id=command_id,
- chunk=TokenChunk(
- idx=0,
- model=model_id,
- text="",
- token_id=0,
- finish_reason="error",
- error_message=str(e),
- ),
- )
- )
-
-
def main(
bound_instance: BoundInstance,
event_sender: MpSender[Event],
@@ -99,6 +67,7 @@ def main(
bound_instance.bound_runner_id,
bound_instance.bound_shard,
)
+ device_rank = shard_metadata.device_rank
logger.info("hello from the runner")
if getattr(shard_metadata, "immediate_exception", False):
raise Exception("Fake exception - runner failed to spin up.")
@@ -180,7 +149,7 @@ def main(
logger.info(f"warming up inference for instance: {instance}")
toks = warmup_inference(
- model=cast(Model, model),
+ model=model,
tokenizer=tokenizer,
# kv_prefix_cache=kv_prefix_cache, # supply for warmup-time prefix caching
)
@@ -201,20 +170,16 @@ def main(
runner_id=runner_id, runner_status=current_status
)
)
- with send_error_chunk_on_exception(
- event_sender,
- command_id,
- shard_metadata.model_meta.model_id,
- shard_metadata.device_rank,
- ):
- assert model
- assert tokenizer
- assert task_params.messages[0].content is not None
+ assert model
+ assert tokenizer
+ assert task_params.messages[0].content is not None
+
+ try:
_check_for_debug_prompts(task_params.messages[0].content)
# Generate responses using the actual MLX generation
mlx_generator = mlx_generate(
- model=cast(Model, model),
+ model=model,
tokenizer=tokenizer,
task=task_params,
)
@@ -228,7 +193,7 @@ def main(
for response in mlx_generator:
match response:
case GenerationResponse():
- if shard_metadata.device_rank == 0:
+ if device_rank == 0:
event_sender.send(
ChunkGenerated(
command_id=command_id,
@@ -243,6 +208,24 @@ def main(
)
)
+ # can we make this more explicit?
+ except Exception as e:
+ if device_rank == 0:
+ event_sender.send(
+ ChunkGenerated(
+ command_id=command_id,
+ chunk=TokenChunk(
+ idx=0,
+ model=shard_metadata.model_meta.model_id,
+ text="",
+ token_id=0,
+ finish_reason="error",
+ error_message=str(e),
+ ),
+ )
+ )
+ raise
+
current_status = RunnerReady()
logger.info("runner ready")
case Shutdown():
diff --git a/src/exo/worker/tests/unittests/test_runner/test_error_handling.py b/src/exo/worker/tests/unittests/test_runner/test_error_handling.py
deleted file mode 100644
index ed8039a1..00000000
--- a/src/exo/worker/tests/unittests/test_runner/test_error_handling.py
+++ /dev/null
@@ -1,50 +0,0 @@
-# pyright: reportAny=false
-from unittest.mock import MagicMock
-
-from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.common import CommandId
-from exo.shared.types.events import ChunkGenerated
-from exo.worker.runner.runner import send_error_chunk_on_exception
-from exo.worker.tests.constants import MODEL_A_ID
-
-
-def test_send_error_chunk_on_exception_no_error() -> None:
- event_sender = MagicMock()
- command_id = CommandId()
-
- with send_error_chunk_on_exception(
- event_sender, command_id, MODEL_A_ID, device_rank=0
- ):
- _ = 1 + 1
-
- event_sender.send.assert_not_called()
-
-
-def test_send_error_chunk_on_exception_catches_error() -> None:
- event_sender = MagicMock()
- command_id = CommandId()
-
- with send_error_chunk_on_exception(
- event_sender, command_id, MODEL_A_ID, device_rank=0
- ):
- raise ValueError("test error")
-
- event_sender.send.assert_called_once()
- call_args = event_sender.send.call_args[0][0]
- assert isinstance(call_args, ChunkGenerated)
- assert call_args.command_id == command_id
- assert isinstance(call_args.chunk, TokenChunk)
- assert call_args.chunk.finish_reason == "error"
- assert call_args.chunk.error_message == "test error"
-
-
-def test_send_error_chunk_on_exception_skips_non_rank_zero() -> None:
- event_sender = MagicMock()
- command_id = CommandId()
-
- with send_error_chunk_on_exception(
- event_sender, command_id, MODEL_A_ID, device_rank=1
- ):
- raise ValueError("test error")
-
- event_sender.send.assert_not_called()
← 618cee52 Resolve test event ordering flakiness (#1194)
·
back to Exo
·
Add dashboard screenshots to README (#1185) 7ff937d8 →