← back to Exo
runner process checks (#1592)
2261014715346a9c4125ab63e16be04de5f56001 · 2026-02-23 16:11:18 +0000 · Evan Quiney
partial solve to some of the more mysterious failures
notably, runner now switches to failed after EXO RUNNER MUST OOM
Files touched
M src/exo/utils/channels.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/runner/bootstrap.pyM src/exo/worker/runner/runner_supervisor.py
Diff
commit 2261014715346a9c4125ab63e16be04de5f56001
Author: Evan Quiney <evanev7@gmail.com>
Date: Mon Feb 23 16:11:18 2026 +0000
runner process checks (#1592)
partial solve to some of the more mysterious failures
notably, runner now switches to failed after EXO RUNNER MUST OOM
---
src/exo/utils/channels.py | 4 +-
src/exo/worker/engines/mlx/utils_mlx.py | 2 +-
src/exo/worker/runner/bootstrap.py | 6 +--
src/exo/worker/runner/runner_supervisor.py | 61 ++++++++++++++++++++----------
4 files changed, 48 insertions(+), 25 deletions(-)
diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py
index 6ce4849d..aa50410c 100644
--- a/src/exo/utils/channels.py
+++ b/src/exo/utils/channels.py
@@ -5,7 +5,7 @@ from math import inf
from multiprocessing.synchronize import Event
from queue import Empty, Full
from types import TracebackType
-from typing import Self
+from typing import Any, Self
from anyio import (
CapacityLimiter,
@@ -157,7 +157,7 @@ class MpSender[T]:
) -> None:
self.close()
- def __getstate__(self):
+ def __getstate__(self) -> dict[str, Any]:
d = self.__dict__.copy()
d.pop("__orig_class__", None)
return d
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 2668f06f..30f9afb4 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -642,7 +642,7 @@ class NullKVCache(KVCache):
raise NotImplementedError("We should not be setting a NullKVCache.")
-def mlx_force_oom(size: int = 40000) -> None:
+def mlx_force_oom(size: int = 200000) -> None:
"""
Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
"""
diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py
index 61521f8e..3a77fdba 100644
--- a/src/exo/worker/runner/bootstrap.py
+++ b/src/exo/worker/runner/bootstrap.py
@@ -18,15 +18,15 @@ def entrypoint(
cancel_receiver: MpReceiver[TaskId],
_logger: "loguru.Logger",
) -> None:
+ global logger
+ logger = _logger
+
fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
if fast_synch_override != "off":
os.environ["MLX_METAL_FAST_SYNCH"] = "1"
else:
os.environ["MLX_METAL_FAST_SYNCH"] = "0"
- global logger
- logger = _logger
-
logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
# Import main after setting global logger - this lets us just import logger from this module
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index e8a06a77..4aa4e6f7 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -1,7 +1,7 @@
import contextlib
+import multiprocessing as mp
import signal
from dataclasses import dataclass, field
-from multiprocessing import Process
from typing import Self
import anyio
@@ -32,6 +32,7 @@ from exo.shared.types.worker.runners import (
)
from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
+from exo.utils.task_group import TaskGroup
from exo.worker.runner.bootstrap import entrypoint
PREFILL_TIMEOUT_SECONDS = 60
@@ -42,16 +43,20 @@ DECODE_TIMEOUT_SECONDS = 5
class RunnerSupervisor:
shard_metadata: ShardMetadata
bound_instance: BoundInstance
- runner_process: Process
+ runner_process: mp.Process
initialize_timeout: float
_ev_recv: MpReceiver[Event]
_task_sender: MpSender[Task]
_event_sender: Sender[Event]
_cancel_sender: MpSender[TaskId]
+ _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)
completed: set[TaskId] = field(default_factory=set, init=False)
cancelled: set[TaskId] = field(default_factory=set, init=False)
+ _cancel_watch_runner: anyio.CancelScope = field(
+ default_factory=anyio.CancelScope, init=False
+ )
@classmethod
def create(
@@ -65,7 +70,7 @@ class RunnerSupervisor:
task_sender, task_recv = mp_channel[Task]()
cancel_sender, cancel_recv = mp_channel[TaskId]()
- runner_process = Process(
+ runner_process = mp.Process(
target=entrypoint,
args=(
bound_instance,
@@ -94,12 +99,17 @@ class RunnerSupervisor:
async def run(self):
self.runner_process.start()
- await self._forward_events()
+ async with self._tg as tg:
+ tg.start_soon(self._watch_runner)
+ tg.start_soon(self._forward_events)
def shutdown(self):
logger.info("Runner supervisor shutting down")
+ self._tg.cancel_tasks()
self._ev_recv.close()
self._task_sender.close()
+ if not self._cancel_watch_runner.cancel_called:
+ self._cancel_watch_runner.cancel()
with contextlib.suppress(ClosedResourceError):
self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
self._cancel_sender.close()
@@ -151,8 +161,8 @@ class RunnerSupervisor:
await self._check_runner(TimeoutError("cancel pipe blocked"))
async def _forward_events(self):
- with self._ev_recv as events:
- try:
+ try:
+ with self._ev_recv as events:
async for event in events:
if isinstance(event, RunnerStatusUpdated):
self.status = event.runner_status
@@ -176,25 +186,34 @@ class RunnerSupervisor:
)
self.completed.add(event.task_id)
await self._event_sender.send(event)
- except (ClosedResourceError, BrokenResourceError) as e:
- await self._check_runner(e)
- for tid in self.pending:
- self.pending[tid].set()
- self._event_sender.close()
+ except (ClosedResourceError, BrokenResourceError) as e:
+ await self._check_runner(e)
+ finally:
+ for tid in self.pending:
+ self.pending[tid].set()
def __del__(self) -> None:
if self.runner_process.is_alive():
- logger.warning("RunnerSupervisor was not stopped cleanly.")
+ logger.critical("RunnerSupervisor was not stopped cleanly.")
with contextlib.suppress(ValueError):
self.runner_process.kill()
+ async def _watch_runner(self) -> None:
+ with self._cancel_watch_runner:
+ while True:
+ await anyio.sleep(5)
+ if not self.runner_process.is_alive():
+ await self._check_runner(RuntimeError("Runner found to be dead"))
+
async def _check_runner(self, e: Exception) -> None:
+ if not self._cancel_watch_runner.cancel_called:
+ self._cancel_watch_runner.cancel()
logger.info("Checking runner's status")
if self.runner_process.is_alive():
logger.info("Runner was found to be alive, attempting to join process")
await to_thread.run_sync(self.runner_process.join, 5)
rc = self.runner_process.exitcode
- logger.info(f"RunnerSupervisor exited with exit code {rc}")
+ logger.info(f"Runner exited with exit code {rc}")
if rc == 0:
return
@@ -207,15 +226,19 @@ class RunnerSupervisor:
else:
cause = f"exitcode={rc}"
- logger.opt(exception=e).error(f"Runner terminated ({cause})")
+ logger.opt(exception=e).error(f"Runner terminated with {cause}")
try:
- await self._event_sender.send(
- RunnerStatusUpdated(
- runner_id=self.bound_instance.bound_runner_id,
- runner_status=RunnerFailed(error_message=f"Terminated ({cause})"),
+ self.status = RunnerFailed(error_message=f"Terminated ({cause})")
+ with anyio.CancelScope(shield=True):
+ await self._event_sender.send(
+ RunnerStatusUpdated(
+ runner_id=self.bound_instance.bound_runner_id,
+ runner_status=RunnerFailed(
+ error_message=f"Terminated ({cause})"
+ ),
+ )
)
- )
except (ClosedResourceError, BrokenResourceError):
logger.warning(
"Event sender already closed, unable to report runner failure"
← 61d2a2b6 add lazy task group (#1569)
·
back to Exo
·
fix: handle gossipsub MessageTooLarge error to prevent silen dab7ed48 →