← back to Exo
Return error responses for Chat Completions (#1173)
745343c70562a137685c864f3d86fb117c1cbb9e · 2026-01-16 19:24:37 +0000 · rltakashige
- Error chunks
- Use error handling in exo_bench.py
## Motivation
Return when an error occurs so that generation stops. Adding timeouts is
a separate TODO for model loading and chat completions.
## Changes
- Return HTTP exceptions as JSON responses in an OpenAI compatible
format.
- Context manager for generation to catch and return error messages.
- Use error handling in exo_bench.py.
## Test Plan
### Manual Testing
Manually tested that exo_bench returns on failures within and outside
generation
### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
Files touched
M bench/exo_bench.pyM src/exo/master/api.pyA src/exo/master/tests/test_api_error_handling.pyM src/exo/shared/types/api.pyM src/exo/shared/types/chunks.pyM src/exo/worker/runner/runner.pyA src/exo/worker/tests/unittests/test_runner/test_error_handling.py
Diff
commit 745343c70562a137685c864f3d86fb117c1cbb9e
Author: rltakashige <rl.takashige@gmail.com>
Date: Fri Jan 16 19:24:37 2026 +0000
Return error responses for Chat Completions (#1173)
- Error chunks
- Use error handling in exo_bench.py
## Motivation
Return when an error occurs so that generation stops. Adding timeouts is
a separate TODO for model loading and chat completions.
## Changes
- Return HTTP exceptions as JSON responses in an OpenAI compatible
format.
- Context manager for generation to catch and return error messages.
- Use error handling in exo_bench.py.
## Test Plan
### Manual Testing
Manually tested that exo_bench returns on failures within and outside
generation
### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
---
bench/exo_bench.py | 33 ++++++-
src/exo/master/api.py | 69 ++++++++++---
src/exo/master/tests/test_api_error_handling.py | 107 ++++++++++++++++++++
src/exo/shared/types/api.py | 13 ++-
src/exo/shared/types/chunks.py | 1 +
src/exo/worker/runner/runner.py | 108 ++++++++++++++-------
.../unittests/test_runner/test_error_handling.py | 50 ++++++++++
7 files changed, 332 insertions(+), 49 deletions(-)
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 39d65553..6fbed737 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -3,6 +3,7 @@
from __future__ import annotations
import argparse
+import contextlib
import http.client
import json
import os
@@ -104,22 +105,46 @@ def runner_ready(runner: dict[str, Any]) -> bool:
return "RunnerReady" in runner
+def runner_failed(runner: dict[str, Any]) -> bool:
+ return "RunnerFailed" in runner
+
+
+def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
+ if "RunnerFailed" in runner:
+ return runner["RunnerFailed"].get("errorMessage")
+ return None
+
+
def wait_for_instance_ready(
client: ExoClient, instance_id: str, timeout: float = 24000.0
) -> None:
start_time = time.time()
+ instance_existed = False
while time.time() - start_time < timeout:
state = client.request_json("GET", "/state")
instances = state.get("instances", {})
if instance_id not in instances:
+ if instance_existed:
+ # Instance was deleted after being created - likely due to runner failure
+ raise RuntimeError(
+ f"Instance {instance_id} was deleted (runner may have failed)"
+ )
time.sleep(0.1)
continue
+ instance_existed = True
instance = instances[instance_id]
runner_ids = runner_ids_from_instance(instance)
runners = state.get("runners", {})
+ # Check for failed runners first
+ for rid in runner_ids:
+ runner = runners.get(rid, {})
+ if runner_failed(runner):
+ error_msg = get_runner_failed_message(runner) or "Unknown error"
+ raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+
if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
return
@@ -441,7 +466,13 @@ def main() -> int:
)
client.request_json("POST", "/instance", body={"instance": instance})
- wait_for_instance_ready(client, instance_id)
+ try:
+ wait_for_instance_ready(client, instance_id)
+ except (RuntimeError, TimeoutError) as e:
+ logger.error(f"Failed to initialize placement: {e}")
+ with contextlib.suppress(ExoHttpError):
+ client.request_json("DELETE", f"/instance/{instance_id}")
+ continue
time.sleep(1)
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index c24ab75b..3f36c3c8 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -1,13 +1,14 @@
import time
from collections.abc import AsyncGenerator
+from http import HTTPStatus
from typing import cast
import anyio
-from anyio import create_task_group
+from anyio import BrokenResourceError, create_task_group
from anyio.abc import TaskGroup
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType]
from hypercorn.config import Config
@@ -29,6 +30,8 @@ from exo.shared.types.api import (
CreateInstanceParams,
CreateInstanceResponse,
DeleteInstanceResponse,
+ ErrorInfo,
+ ErrorResponse,
FinishReason,
GenerationStats,
ModelList,
@@ -49,7 +52,12 @@ from exo.shared.types.commands import (
TaskFinished,
)
from exo.shared.types.common import CommandId, NodeId, SessionId
-from exo.shared.types.events import ChunkGenerated, Event, ForwarderEvent, IndexedEvent
+from exo.shared.types.events import (
+ ChunkGenerated,
+ Event,
+ ForwarderEvent,
+ IndexedEvent,
+)
from exo.shared.types.memory import Memory
from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.state import State
@@ -115,6 +123,7 @@ class API:
self.paused_ev: anyio.Event = anyio.Event()
self.app = FastAPI()
+ self._setup_exception_handlers()
self._setup_cors()
self._setup_routes()
@@ -145,6 +154,20 @@ class API:
self.paused_ev.set()
self.paused_ev = anyio.Event()
+ def _setup_exception_handlers(self) -> None:
+ @self.app.exception_handler(HTTPException)
+ async def http_exception_handler( # pyright: ignore[reportUnusedFunction]
+ _: Request, exc: HTTPException
+ ) -> JSONResponse:
+ err = ErrorResponse(
+ error=ErrorInfo(
+ message=exc.detail,
+ type=HTTPStatus(exc.status_code).phrase,
+ code=exc.status_code,
+ )
+ )
+ return JSONResponse(err.model_dump(), status_code=exc.status_code)
+
def _setup_cors(self) -> None:
self.app.add_middleware(
CORSMiddleware,
@@ -406,6 +429,18 @@ class API:
"""Generate chat completion stream as JSON strings."""
async for chunk in self._chat_chunk_stream(command_id):
+ if chunk.finish_reason == "error":
+ error_response = ErrorResponse(
+ error=ErrorInfo(
+ message=chunk.error_message or "Internal server error",
+ type="InternalServerError",
+ code=500,
+ )
+ )
+ yield f"data: {error_response.model_dump_json()}\n\n"
+ yield "data: [DONE]\n\n"
+ return
+
chunk_response: ChatCompletionResponse = chunk_to_response(
chunk, command_id
)
@@ -426,6 +461,12 @@ class API:
finish_reason: FinishReason | None = None
async for chunk in self._chat_chunk_stream(command_id):
+ if chunk.finish_reason == "error":
+ raise HTTPException(
+ status_code=500,
+ detail=chunk.error_message or "Internal server error",
+ )
+
if model is None:
model = chunk.model
@@ -463,6 +504,12 @@ class API:
stats: GenerationStats | None = None
async for chunk in self._chat_chunk_stream(command_id):
+ if chunk.finish_reason == "error":
+ raise HTTPException(
+ status_code=500,
+ detail=chunk.error_message or "Internal server error",
+ )
+
if model is None:
model = chunk.model
@@ -607,14 +654,14 @@ class API:
for idx, event in self.event_buffer.drain_indexed():
self._event_log.append(event)
self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
- if (
- isinstance(event, ChunkGenerated)
- and event.command_id in self._chat_completion_queues
- ):
+ if isinstance(event, ChunkGenerated):
assert isinstance(event.chunk, TokenChunk)
- await self._chat_completion_queues[event.command_id].send(
- event.chunk
- )
+ queue = self._chat_completion_queues.get(event.command_id)
+ if queue is not None:
+ try:
+ await queue.send(event.chunk)
+ except BrokenResourceError:
+ self._chat_completion_queues.pop(event.command_id, None)
async def _pause_on_new_election(self):
with self.election_receiver as ems:
diff --git a/src/exo/master/tests/test_api_error_handling.py b/src/exo/master/tests/test_api_error_handling.py
new file mode 100644
index 00000000..76001043
--- /dev/null
+++ b/src/exo/master/tests/test_api_error_handling.py
@@ -0,0 +1,107 @@
+# pyright: reportUnusedFunction=false, reportAny=false
+from typing import Any, get_args
+
+from fastapi import FastAPI, HTTPException
+from fastapi.testclient import TestClient
+
+from exo.shared.types.api import ErrorInfo, ErrorResponse, FinishReason
+from exo.shared.types.chunks import TokenChunk
+from exo.worker.tests.constants import MODEL_A_ID
+
+
+def test_http_exception_handler_formats_openai_style() -> None:
+ """Test that HTTPException is converted to OpenAI-style error format."""
+ from exo.master.api import API
+
+ app = FastAPI()
+
+ # Setup exception handler
+ api = object.__new__(API)
+ api.app = app
+ api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage]
+
+ # Add test routes that raise HTTPException
+ @app.get("/test-error")
+ async def _test_error() -> None:
+ raise HTTPException(status_code=500, detail="Test error message")
+
+ @app.get("/test-not-found")
+ async def _test_not_found() -> None:
+ raise HTTPException(status_code=404, detail="Resource not found")
+
+ client = TestClient(app)
+
+ # Test 500 error
+ response = client.get("/test-error")
+ assert response.status_code == 500
+ data: dict[str, Any] = response.json()
+ assert "error" in data
+ assert data["error"]["message"] == "Test error message"
+ assert data["error"]["type"] == "Internal Server Error"
+ assert data["error"]["code"] == 500
+
+ # Test 404 error
+ response = client.get("/test-not-found")
+ assert response.status_code == 404
+ data = response.json()
+ assert "error" in data
+ assert data["error"]["message"] == "Resource not found"
+ assert data["error"]["type"] == "Not Found"
+ assert data["error"]["code"] == 404
+
+
+def test_finish_reason_includes_error() -> None:
+ valid_reasons = get_args(FinishReason)
+ assert "error" in valid_reasons
+
+
+def test_token_chunk_with_error_fields() -> None:
+ chunk = TokenChunk(
+ idx=0,
+ model=MODEL_A_ID,
+ text="",
+ token_id=0,
+ finish_reason="error",
+ error_message="Something went wrong",
+ )
+
+ assert chunk.finish_reason == "error"
+ assert chunk.error_message == "Something went wrong"
+
+
+def test_token_chunk_without_error() -> None:
+ chunk = TokenChunk(
+ idx=1,
+ model=MODEL_A_ID,
+ text="Hello",
+ token_id=42,
+ finish_reason=None,
+ )
+
+ assert chunk.finish_reason is None
+ assert chunk.error_message is None
+
+
+def test_error_response_construction() -> None:
+ error_response = ErrorResponse(
+ error=ErrorInfo(
+ message="Generation failed",
+ type="InternalServerError",
+ code=500,
+ )
+ )
+
+ assert error_response.error.message == "Generation failed"
+ assert error_response.error.code == 500
+
+
+def test_normal_finish_reasons_still_work() -> None:
+ for reason in ["stop", "length", "tool_calls", "content_filter", "function_call"]:
+ chunk = TokenChunk(
+ idx=0,
+ model=MODEL_A_ID,
+ text="done",
+ token_id=100,
+ finish_reason=reason, # type: ignore[arg-type]
+ )
+ assert chunk.finish_reason == reason
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 5073f12b..3ae6917a 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -11,10 +11,21 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
FinishReason = Literal[
- "stop", "length", "tool_calls", "content_filter", "function_call"
+ "stop", "length", "tool_calls", "content_filter", "function_call", "error"
]
+class ErrorInfo(BaseModel):
+ message: str
+ type: str
+ param: str | None = None
+ code: int
+
+
+class ErrorResponse(BaseModel):
+ error: ErrorInfo
+
+
class ModelListModel(BaseModel):
id: str
object: str = "model"
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index 420e18eb..4785b172 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -22,6 +22,7 @@ class TokenChunk(BaseChunk):
token_id: int
finish_reason: FinishReason | None = None
stats: GenerationStats | None = None
+ error_message: str | None = None
class ImageChunk(BaseChunk):
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index fae4acf8..e195b9b3 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,6 +1,8 @@
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
@@ -13,6 +15,7 @@ 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,
@@ -20,6 +23,7 @@ from exo.shared.types.events import (
TaskAcknowledged,
TaskStatusUpdated,
)
+from exo.shared.types.models import ModelId
from exo.shared.types.tasks import (
ChatCompletion,
ConnectToGroup,
@@ -48,6 +52,7 @@ 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,
@@ -57,6 +62,33 @@ 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],
@@ -135,7 +167,7 @@ def main(
logger.info(f"warming up inference for instance: {instance}")
toks = warmup_inference(
- model=model,
+ model=cast(Model, model),
tokenizer=tokenizer,
# kv_prefix_cache=kv_prefix_cache, # supply for warmup-time prefix caching
)
@@ -148,8 +180,6 @@ def main(
case ChatCompletion(task_params=task_params, command_id=command_id) if (
isinstance(current_status, RunnerReady)
):
- assert model
- assert tokenizer
logger.info(f"received chat request: {str(task)[:500]}")
current_status = RunnerRunning()
logger.info("runner running")
@@ -158,41 +188,47 @@ def main(
runner_id=runner_id, runner_status=current_status
)
)
- assert task_params.messages[0].content is not None
- _check_for_debug_prompts(task_params.messages[0].content)
-
- # Generate responses using the actual MLX generation
- mlx_generator = mlx_generate(
- model=model,
- tokenizer=tokenizer,
- task=task_params,
- )
+ 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
+ _check_for_debug_prompts(task_params.messages[0].content)
+
+ # Generate responses using the actual MLX generation
+ mlx_generator = mlx_generate(
+ model=cast(Model, model),
+ tokenizer=tokenizer,
+ task=task_params,
+ )
- # GPT-OSS specific parsing to match other model formats.
- if isinstance(model, GptOssModel):
- mlx_generator = parse_gpt_oss(mlx_generator)
-
- # TODO: Add tool call parser here
-
- for response in mlx_generator:
- match response:
- case GenerationResponse():
- if shard_metadata.device_rank == 0:
- event_sender.send(
- ChunkGenerated(
- command_id=command_id,
- chunk=TokenChunk(
- idx=response.token,
- model=shard_metadata.model_meta.model_id,
- text=response.text,
- token_id=response.token,
- finish_reason=response.finish_reason,
- stats=response.stats,
- ),
+ # GPT-OSS specific parsing to match other model formats.
+ if isinstance(model, GptOssModel):
+ mlx_generator = parse_gpt_oss(mlx_generator)
+
+ # TODO: Add tool call parser here
+
+ for response in mlx_generator:
+ match response:
+ case GenerationResponse():
+ if shard_metadata.device_rank == 0:
+ event_sender.send(
+ ChunkGenerated(
+ command_id=command_id,
+ chunk=TokenChunk(
+ idx=response.token,
+ model=shard_metadata.model_meta.model_id,
+ text=response.text,
+ token_id=response.token,
+ finish_reason=response.finish_reason,
+ stats=response.stats,
+ ),
+ )
)
- )
- # case TokenizedResponse():
- # TODO: something here ig
current_status = RunnerReady()
logger.info("runner ready")
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
new file mode 100644
index 00000000..ed8039a1
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_runner/test_error_handling.py
@@ -0,0 +1,50 @@
+# 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()
← 5e28664c Fix draft release detection (attempt 3) (#1176)
·
back to Exo
·
Handle model timeouts (#1177) 5c8a2379 →