← back to Exo
fix(bug): EventRouter lifetime-handling fixed, no more process crashes (#2102)
74e9fe15e62fe189dc7e019db86e75c83eca2721 · 2026-05-22 14:20:04 +0100 · Andrei Cravtov
## Motivation
Trying to (partially) fix
[this](https://github.com/exo-explore/exo/issues/2101) issue.
## Changes
Changed channels (in channels.py) to support exception overriding.
Made EventRouter channels throw a subclass of the resource closed/broken
errors.
The current lifetime logic of EventRouter in event loop no longer blows
up because components that use channels from EventRouter now catch the
subclass exceptions in the run method: Worker, Master,
DownloadCoordinator, RunnerSupervisor.
Added logic to throw when API server exits without being asked to shut
down - this kill the sleep-forever in the task-group.
Files touched
M .gitignoreM src/exo/api/main.pyM src/exo/download/coordinator.pyM src/exo/main.pyM src/exo/master/main.pyM src/exo/routing/event_router.pyM src/exo/utils/channels.pyA src/exo/utils/tests/test_channels.pyD src/exo/utils/tests/test_mp_channel.pyM src/exo/worker/main.py
Diff
commit 74e9fe15e62fe189dc7e019db86e75c83eca2721
Author: Andrei Cravtov <the.andrei.cravtov@gmail.com>
Date: Fri May 22 14:20:04 2026 +0100
fix(bug): EventRouter lifetime-handling fixed, no more process crashes (#2102)
## Motivation
Trying to (partially) fix
[this](https://github.com/exo-explore/exo/issues/2101) issue.
## Changes
Changed channels (in channels.py) to support exception overriding.
Made EventRouter channels throw a subclass of the resource closed/broken
errors.
The current lifetime logic of EventRouter in event loop no longer blows
up because components that use channels from EventRouter now catch the
subclass exceptions in the run method: Worker, Master,
DownloadCoordinator, RunnerSupervisor.
Added logic to throw when API server exits without being asked to shut
down - this kill the sleep-forever in the task-group.
---
.gitignore | 2 +
src/exo/api/main.py | 6 +-
src/exo/download/coordinator.py | 11 +++
src/exo/main.py | 8 +-
src/exo/master/main.py | 15 ++-
src/exo/routing/event_router.py | 23 ++++-
src/exo/utils/channels.py | 163 +++++++++++++++++++++++++++++++--
src/exo/utils/tests/test_channels.py | 121 ++++++++++++++++++++++++
src/exo/utils/tests/test_mp_channel.py | 40 --------
src/exo/worker/main.py | 8 +-
10 files changed, 340 insertions(+), 57 deletions(-)
diff --git a/.gitignore b/.gitignore
index fa09fb01..3ca09b53 100644
--- a/.gitignore
+++ b/.gitignore
@@ -38,6 +38,8 @@ bench/**/*.json
# tmp
tmp/models
/build/exo
+/.agents
/.claude/skills
/.claude
/.codex
+skills-lock.json
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 90a93a4c..b836e149 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -20,7 +20,7 @@ from fastapi.staticfiles import StaticFiles
from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType]
from hypercorn.config import Config
from hypercorn.typing import ASGIFramework
-from hypercorn.utils import LifespanTimeoutError
+from hypercorn.utils import LifespanTimeoutError, ShutdownError
from loguru import logger
from exo.api.adapters.chat_completions import (
@@ -1914,6 +1914,10 @@ class API:
cfg,
shutdown_trigger=ev.wait,
)
+ if not ev.is_set():
+ raise ShutdownError(
+ "Server exited without shutdown trigger - exiting abnormally"
+ )
except LifespanTimeoutError as e:
logger.warning(
"Graceful server shutdown timed out, some connections forcebly closed"
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py
index de9c4722..5b3f3c35 100644
--- a/src/exo/download/coordinator.py
+++ b/src/exo/download/coordinator.py
@@ -15,6 +15,10 @@ from exo.download.download_utils import (
resolve_existing_model,
)
from exo.download.shard_downloader import ShardDownloader
+from exo.routing.event_router import (
+ EventRouterBrokenResourceError,
+ EventRouterClosedResourceError,
+)
from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
from exo.shared.models import model_cards
from exo.shared.models.model_cards import ModelId
@@ -139,7 +143,14 @@ class DownloadCoordinator:
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._emit_existing_download_progress)
+ except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
+ # Event router has been closed (try-star syntax handles error groups)
+ pass
finally:
+ # don't forget to clean up resources
+ self.download_command_receiver.close()
+ self.event_sender.close()
+
self._stopped.set()
async def shutdown(self) -> None:
diff --git a/src/exo/main.py b/src/exo/main.py
index 8e3220eb..2d788d81 100644
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
from typing import Self
import anyio
+from anyio.lowlevel import checkpoint as anyio_checkpoint
from loguru import logger
from pydantic import PositiveInt
@@ -190,7 +191,7 @@ class Node:
# - Shut down and re-create the API
if result.is_new_master:
- await anyio.sleep(0)
+ await anyio_checkpoint()
self.event_router.shutdown()
self.event_router = EventRouter(
result.session_id,
@@ -203,7 +204,10 @@ class Node:
result.session_id.master_node_id == self.node_id
and self.master is not None
):
- logger.info("Node elected Master")
+ assert not result.is_new_master, (
+ "cannot be new master if we remain master"
+ )
+ logger.info("Node elected Master - maintaining self")
elif (
result.session_id.master_node_id == self.node_id
and self.master is None
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index 655e6556..55fed7e6 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -11,6 +11,10 @@ from exo.master.placement import (
place_instance,
)
from exo.master.placement_utils import find_ip_prioritised
+from exo.routing.event_router import (
+ EventRouterBrokenResourceError,
+ EventRouterClosedResourceError,
+)
from exo.shared.apply import apply
from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
from exo.shared.types.commands import (
@@ -151,6 +155,9 @@ class Master:
tg.start_soon(self._event_processor)
tg.start_soon(self._command_processor)
tg.start_soon(self._plan)
+ except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
+ # Event router has been closed (try-star syntax handles error groups)
+ pass
finally:
self._event_log.close()
self.global_event_sender.close()
@@ -448,7 +455,9 @@ class Master:
self._event_log.read_range(command.since_idx, end),
start=command.since_idx,
):
- await self._send_event(IndexedEvent(idx=i, event=event))
+ await self._send_indexed_event(
+ IndexedEvent(idx=i, event=event)
+ )
for event in generated_events:
await self.event_sender.send(event)
except ValueError as e:
@@ -506,10 +515,10 @@ class Master:
self.state = apply(self.state, indexed)
self._event_log.append(event)
- await self._send_event(indexed)
+ await self._send_indexed_event(indexed)
# This function is re-entrant, take care!
- async def _send_event(self, event: IndexedEvent):
+ async def _send_indexed_event(self, event: IndexedEvent):
# Convenience method since this line is ugly
await self.global_event_sender.send(
GlobalForwarderEvent(
diff --git a/src/exo/routing/event_router.py b/src/exo/routing/event_router.py
index 4f99c152..1d6aecc0 100644
--- a/src/exo/routing/event_router.py
+++ b/src/exo/routing/event_router.py
@@ -15,11 +15,30 @@ from exo.shared.types.events import (
IndexedEvent,
LocalForwarderEvent,
)
+from exo.utils import channels
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
from exo.utils.task_group import TaskGroup
+class EventRouterClosedResourceError(ClosedResourceError):
+ pass
+
+
+class EventRouterBrokenResourceError(BrokenResourceError):
+ pass
+
+
+# Event Router is created and destroyed before consumers of its channels are,
+# hence its nice to have tagged errors for event-router channels being closed
+#
+# so consumers can catch specifically these errors, rather than the generic ones
+_ERROR_CFG = channels.ErrorOverride(
+ closed_resource_error=EventRouterClosedResourceError,
+ broken_resource_error=EventRouterBrokenResourceError,
+)
+
+
@dataclass
class EventRouter:
session_id: SessionId
@@ -64,7 +83,7 @@ class EventRouter:
await self.external_outbound.send(event)
def sender(self) -> Sender[Event]:
- send, recv = channel[Event]()
+ send, recv = channel[Event](error_override_config=_ERROR_CFG)
if self._tg.is_running():
self._tg.start_soon(self._ingest, SystemId(), recv)
else:
@@ -73,7 +92,7 @@ class EventRouter:
def receiver(self) -> Receiver[IndexedEvent]:
assert not self._tg.is_running()
- send, recv = channel[IndexedEvent]()
+ send, recv = channel[IndexedEvent](error_override_config=_ERROR_CFG)
self.internal_outbound.append(send)
return recv
diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py
index aa50410c..740b2eb4 100644
--- a/src/exo/utils/channels.py
+++ b/src/exo/utils/channels.py
@@ -1,13 +1,16 @@
import contextlib
import multiprocessing as mp
from dataclasses import dataclass, field
+from functools import wraps
+from inspect import iscoroutinefunction
from math import inf
from multiprocessing.synchronize import Event
from queue import Empty, Full
-from types import TracebackType
-from typing import Any, Self
+from types import CoroutineType, TracebackType
+from typing import Any, Callable, NoReturn, Self, cast, overload, override
from anyio import (
+ BrokenResourceError,
CapacityLimiter,
ClosedResourceError,
EndOfStream,
@@ -20,35 +23,172 @@ from anyio.streams.memory import (
from anyio.streams.memory import (
MemoryObjectSendStream as AnyioSender,
)
+from anyio.streams.memory import (
+ MemoryObjectStreamState,
+)
from anyio.streams.memory import (
MemoryObjectStreamState as AnyioState,
)
+@dataclass(eq=False)
+class ErrorOverride:
+ closed_resource_error: type[ClosedResourceError] = field(
+ default=ClosedResourceError,
+ )
+ broken_resource_error: type[BrokenResourceError] = field(
+ default=BrokenResourceError,
+ )
+ end_of_stream: type[EndOfStream] = field(
+ default=EndOfStream,
+ )
+ would_block: type[WouldBlock] = field(
+ default=WouldBlock,
+ )
+
+ @overload
+ def patch[**P, R](
+ self,
+ fn: Callable[P, CoroutineType[Any, Any, R]],
+ /,
+ ) -> Callable[P, CoroutineType[Any, Any, R]]: ...
+
+ @overload
+ def patch[**P, R](
+ self,
+ fn: Callable[P, R],
+ /,
+ ) -> Callable[P, R]: ...
+
+ def patch[**P, R](self, fn: Callable[P, Any], /) -> Callable[P, Any]:
+ """
+ Returns a function with all these exceptions replaced by their overrides
+ """
+
+ if iscoroutinefunction(fn):
+ async_fn = cast(Callable[P, CoroutineType[Any, Any, R]], fn)
+
+ @wraps(async_fn)
+ async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+ try:
+ return await async_fn(*args, **kwargs)
+ except ClosedResourceError as e:
+ self._raise_replace(self.closed_resource_error, e)
+ except BrokenResourceError as e:
+ self._raise_replace(self.broken_resource_error, e)
+ except EndOfStream as e:
+ self._raise_replace(self.end_of_stream, e)
+ except WouldBlock as e:
+ self._raise_replace(self.would_block, e)
+
+ return async_wrapper
+ else:
+ sync_fn = cast(Callable[P, R], fn)
+
+ @wraps(sync_fn)
+ def sync_wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+ try:
+ return sync_fn(*args, **kwargs)
+ except ClosedResourceError as e:
+ self._raise_replace(self.closed_resource_error, e)
+ except BrokenResourceError as e:
+ self._raise_replace(self.broken_resource_error, e)
+ except EndOfStream as e:
+ self._raise_replace(self.end_of_stream, e)
+ except WouldBlock as e:
+ self._raise_replace(self.would_block, e)
+
+ return sync_wrapper
+
+ @staticmethod
+ def _raise_replace(replacement: type[BaseException], e: BaseException) -> NoReturn:
+ if isinstance(e, replacement):
+ raise
+ raise replacement() from e
+
+
class Sender[T](AnyioSender[T]):
+ def __init__(
+ self,
+ state: MemoryObjectStreamState[T],
+ error_override_config: ErrorOverride | None,
+ ):
+ super().__init__(_state=state)
+
+ # patch the methods we want to override errors for
+ #
+ # NOTE: it is very important that new methods which are added,
+ # and which can throw, are patched in this block
+ if (e := error_override_config) is not None:
+ # new methods of this class
+ self.clone_receiver = e.patch(self.clone_receiver)
+
+ # overridden methods
+ self.clone = e.patch(self.clone)
+
+ # parent methods
+ self.send_nowait = e.patch(self.send_nowait)
+ self.send = e.patch(self.send)
+ self.close = e.patch(self.close)
+ self.aclose = e.patch(self.aclose)
+ self.statistics = e.patch(self.statistics)
+
+ self.err_config = error_override_config
+
+ @override
def clone(self) -> "Sender[T]":
if self._closed:
raise ClosedResourceError
- return Sender(_state=self._state)
+ return Sender(self._state, self.err_config)
def clone_receiver(self) -> "Receiver[T]":
"""Constructs a Receiver using a Senders shared state - similar to calling Receiver.clone() without needing the receiver"""
if self._closed:
raise ClosedResourceError
- return Receiver(_state=self._state)
+ return Receiver(self._state, self.err_config)
class Receiver[T](AnyioReceiver[T]):
+ def __init__(
+ self,
+ state: MemoryObjectStreamState[T],
+ error_override_config: ErrorOverride | None,
+ ):
+ super().__init__(_state=state)
+
+ # patch the methods we want to override errors for
+ #
+ # NOTE: it is very important that new methods which are added,
+ # and which can throw, are patched in this block
+ if (e := error_override_config) is not None:
+ # new methods of this class
+ self.clone_sender = e.patch(self.clone_sender)
+ self.collect = e.patch(self.collect)
+ self.receive_at_least = e.patch(self.receive_at_least)
+
+ # overridden methods
+ self.clone = e.patch(self.clone)
+
+ # parent methods
+ self.receive_nowait = e.patch(self.receive_nowait)
+ self.receive = e.patch(self.receive)
+ self.close = e.patch(self.close)
+ self.aclose = e.patch(self.aclose)
+ self.statistics = e.patch(self.statistics)
+
+ self.err_config = error_override_config
+
+ @override
def clone(self) -> "Receiver[T]":
if self._closed:
raise ClosedResourceError
- return Receiver(_state=self._state)
+ return Receiver(self._state, self.err_config)
def clone_sender(self) -> Sender[T]:
"""Constructs a Sender using a Receivers shared state - similar to calling Sender.clone() without needing the sender"""
if self._closed:
raise ClosedResourceError
- return Sender(_state=self._state)
+ return Sender(self._state, self.err_config)
def collect(self) -> list[T]:
"""Collect all currently available items from this receiver"""
@@ -70,6 +210,7 @@ class Receiver[T](AnyioReceiver[T]):
out.extend(self.collect())
return out
+ @override
def __enter__(self) -> Self:
return self
@@ -285,11 +426,17 @@ class MpReceiver[T]:
class channel[T]: # noqa: N801
"""Create a pair of asynchronous channels for communicating within the same process"""
- def __new__(cls, max_buffer_size: float = inf) -> tuple[Sender[T], Receiver[T]]:
+ def __new__(
+ cls,
+ max_buffer_size: float = inf,
+ error_override_config: ErrorOverride | None = None,
+ ) -> tuple[Sender[T], Receiver[T]]:
if max_buffer_size != inf and not isinstance(max_buffer_size, int):
raise ValueError("max_buffer_size must be either an integer or math.inf")
state = AnyioState[T](max_buffer_size)
- return Sender(_state=state), Receiver(_state=state)
+ return Sender(state, error_override_config), Receiver(
+ state, error_override_config
+ )
class mp_channel[T]: # noqa: N801
diff --git a/src/exo/utils/tests/test_channels.py b/src/exo/utils/tests/test_channels.py
new file mode 100644
index 00000000..a7ebbb48
--- /dev/null
+++ b/src/exo/utils/tests/test_channels.py
@@ -0,0 +1,121 @@
+import multiprocessing as mp
+import time
+
+import pytest
+from anyio import (
+ BrokenResourceError,
+ ClosedResourceError,
+ EndOfStream,
+ WouldBlock,
+ fail_after,
+)
+from loguru import logger
+
+from exo.utils.channels import ErrorOverride, MpReceiver, MpSender, channel, mp_channel
+
+
+class CustomClosedResourceError(ClosedResourceError):
+ pass
+
+
+class CustomBrokenResourceError(BrokenResourceError):
+ pass
+
+
+class CustomEndOfStream(EndOfStream):
+ pass
+
+
+class CustomWouldBlock(WouldBlock):
+ pass
+
+
+ERROR_OVERRIDE = ErrorOverride(
+ closed_resource_error=CustomClosedResourceError,
+ broken_resource_error=CustomBrokenResourceError,
+ end_of_stream=CustomEndOfStream,
+ would_block=CustomWouldBlock,
+)
+
+
+def foo(recv: MpReceiver[str]):
+ expected = ["hi", "hi 2", "bye"]
+ with recv as r:
+ for item in r:
+ assert item == expected.pop(0)
+
+
+def bar(send: MpSender[str]):
+ logger.warning("hi")
+ send.send("hi")
+ time.sleep(0.1)
+ logger.warning("hi 2")
+ send.send("hi 2")
+ time.sleep(0.1)
+ logger.warning("bye")
+ send.send("bye")
+ time.sleep(0.1)
+ send.close()
+
+
+@pytest.mark.anyio
+async def test_channel_ipc():
+ with fail_after(0.5):
+ s, r = mp_channel[str]()
+ p1 = mp.Process(target=foo, args=(r,))
+ p2 = mp.Process(target=bar, args=(s,))
+ p1.start()
+ p2.start()
+ p1.join()
+ p2.join()
+
+
+def test_channel_error_override_replaces_sync_errors_with_subclasses():
+ send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
+
+ with pytest.raises(CustomWouldBlock) as would_block_info:
+ send.send_nowait(1)
+ assert type(would_block_info.value.__cause__) is WouldBlock
+
+ recv.close()
+ with pytest.raises(CustomBrokenResourceError) as broken_resource_info:
+ send.send_nowait(1)
+ assert type(broken_resource_info.value.__cause__) is BrokenResourceError
+
+ send.close()
+ with pytest.raises(CustomClosedResourceError) as closed_resource_info:
+ send.send_nowait(1)
+ assert type(closed_resource_info.value.__cause__) is ClosedResourceError
+
+
+@pytest.mark.anyio
+async def test_channel_error_override_replaces_async_errors_with_subclasses():
+ send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
+ recv.close()
+
+ with pytest.raises(CustomBrokenResourceError) as broken_resource_info:
+ await send.send(1)
+ assert type(broken_resource_info.value.__cause__) is BrokenResourceError
+
+ send, recv = channel[int](error_override_config=ERROR_OVERRIDE)
+ send.close()
+ with pytest.raises(CustomEndOfStream) as end_of_stream_info:
+ await recv.receive()
+ assert type(end_of_stream_info.value.__cause__) is EndOfStream
+
+
+@pytest.mark.anyio
+async def test_channel_error_override_is_preserved_by_clones():
+ send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
+ send_clone = send.clone()
+ recv.close()
+
+ with pytest.raises(CustomBrokenResourceError):
+ await send_clone.send(1)
+
+ send, recv = channel[int](0, error_override_config=ERROR_OVERRIDE)
+ cloned_send = recv.clone_sender()
+ recv.close()
+
+ with pytest.raises(CustomBrokenResourceError):
+ await cloned_send.send(1)
diff --git a/src/exo/utils/tests/test_mp_channel.py b/src/exo/utils/tests/test_mp_channel.py
deleted file mode 100644
index 73a26adb..00000000
--- a/src/exo/utils/tests/test_mp_channel.py
+++ /dev/null
@@ -1,40 +0,0 @@
-import multiprocessing as mp
-import time
-
-import pytest
-from anyio import fail_after
-from loguru import logger
-
-from exo.utils.channels import MpReceiver, MpSender, mp_channel
-
-
-def foo(recv: MpReceiver[str]):
- expected = ["hi", "hi 2", "bye"]
- with recv as r:
- for item in r:
- assert item == expected.pop(0)
-
-
-def bar(send: MpSender[str]):
- logger.warning("hi")
- send.send("hi")
- time.sleep(0.1)
- logger.warning("hi 2")
- send.send("hi 2")
- time.sleep(0.1)
- logger.warning("bye")
- send.send("bye")
- time.sleep(0.1)
- send.close()
-
-
-@pytest.mark.anyio
-async def test_channel_ipc():
- with fail_after(0.5):
- s, r = mp_channel[str]()
- p1 = mp.Process(target=foo, args=(r,))
- p2 = mp.Process(target=bar, args=(s,))
- p1.start()
- p2.start()
- p1.join()
- p2.join()
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 5c34e7b4..a641bacf 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -8,6 +8,10 @@ from loguru import logger
from exo.api.types import ImageEditsTaskParams
from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
+from exo.routing.event_router import (
+ EventRouterBrokenResourceError,
+ EventRouterClosedResourceError,
+)
from exo.shared.apply import apply
from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
from exo.shared.models.model_cards import ModelId, card_cache
@@ -109,7 +113,9 @@ class Worker:
tg.start_soon(self._event_applier)
tg.start_soon(self._poll_connection_updates)
tg.start_soon(self._reconcile_custom_cards)
-
+ except* (EventRouterBrokenResourceError, EventRouterClosedResourceError):
+ # Event router has been closed (try-star syntax handles error groups)
+ pass
finally:
# Actual shutdown code - waits for all tasks to complete before executing.
logger.info("Stopping Worker")
← 90f24bef fix model cards not validating properly after #2071 (#2096)
·
back to Exo
·
feat: added background/daemon support (#2106) a1a22b5f →