electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: add health check and heartbeat to RunnerSupervisor

Browse Source 
Add proactive monitoring to detect runner process death and unresponsiveness:

- Health check loop polls is_alive() every 1s, detects unexpected exits
- Counter-based heartbeat detects frozen/unresponsive processes
- Emits RunnerFailed event and releases pending task waiters on failure
- Add EXO_RUNNER_MUST_DIE debug trigger for testing abrupt process death
- Add chaos E2E test that kills runner mid-inference

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Alex Cheema
2026-02-16 10:01:08 -08:00
co-authored by Claude Opus 4.6
parent 199fb1c9fa
commit 54a036223b
5 changed files with 412 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
e2e/test_runner_chaos.py
+65
View File
@@ -0,0 +1,65 @@
"""Test: Runner chaos — abrupt runner death detection.
slow
Sends a chat completion with the EXO_RUNNER_MUST_DIE trigger, which causes
the runner process to call os._exit(1) (simulating an OOM kill). Verifies that
the RunnerSupervisor health check detects the death and the system doesn't hang.
Requires a machine that can run MLX inference at reasonable speed (Apple Silicon).
Run with: python3 e2e/run_all.py --slow or E2E_SLOW=1 python3 e2e/run_all.py
"""
import asyncio
import contextlib
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent))
from conftest import Cluster
MODEL = "mlx-community/Qwen3-0.6B-4bit"
async def main():
async with Cluster("runner_chaos") as cluster:
await cluster.build()
await cluster.start()
await cluster.assert_healthy()
# Place the model so a runner is loaded and ready
print(f" Launching model {MODEL}...")
await cluster.place_model(MODEL)
# Send a chat request with the die trigger.
# The runner will call os._exit(1) mid-inference, simulating OOM kill.
# The chat request itself will fail — that's expected.
print(" Sending EXO_RUNNER_MUST_DIE trigger...")
with contextlib.suppress(Exception):
await cluster.chat(
model=MODEL,
messages=[{"role": "user", "content": "EXO RUNNER MUST DIE"}],
timeout=60,
)
# Wait for the health check to detect the death and emit RunnerFailed
async def health_check_detected():
log = await cluster.logs()
return "runner process died unexpectedly" in log
await cluster.wait_for(
"Health check detected runner death",
health_check_detected,
timeout=30,
)
# Verify RunnerFailed was emitted (visible in logs)
log = await cluster.logs()
assert "runner process died unexpectedly" in log, (
"Expected health check to detect runner death but it didn't.\n"
f"Logs:\n{log}"
)
print("PASSED: runner_chaos")
if __name__ == "__main__":
asyncio.run(main())
src/exo/worker/runner/bootstrap.py
+28
View File
@@ -1,4 +1,8 @@
from __future__ import annotations
import os
import threading
from multiprocessing.sharedctypes import Synchronized
import loguru
@@ -10,6 +14,15 @@ from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
logger: "loguru.Logger" = loguru.logger
HEARTBEAT_INTERVAL_SECONDS = 0.5
def _heartbeat_loop(heartbeat: Synchronized[int], stop: threading.Event) -> None:
"""Daemon thread that periodically increments the heartbeat counter."""
while not stop.is_set():
heartbeat.value += 1
stop.wait(HEARTBEAT_INTERVAL_SECONDS)
def entrypoint(
bound_instance: BoundInstance,
@@ -17,6 +30,7 @@ def entrypoint(
task_receiver: MpReceiver[Task],
cancel_receiver: MpReceiver[TaskId],
_logger: "loguru.Logger",
heartbeat: Synchronized[int] | None = None,
) -> None:
fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
if fast_synch_override == "on" or (
@@ -35,6 +49,17 @@ def entrypoint(
logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
# Start heartbeat thread so the supervisor can detect if we freeze.
stop_heartbeat = threading.Event()
heartbeat_thread: threading.Thread | None = None
if heartbeat is not None:
heartbeat_thread = threading.Thread(
target=_heartbeat_loop,
args=(heartbeat, stop_heartbeat),
daemon=True,
)
heartbeat_thread.start()
# Import main after setting global logger - this lets us just import logger from this module
try:
from exo.worker.runner.runner import main
@@ -53,6 +78,9 @@ def entrypoint(
)
)
finally:
stop_heartbeat.set()
if heartbeat_thread is not None:
heartbeat_thread.join(timeout=1)
try:
event_sender.close()
task_receiver.close()
src/exo/worker/runner/runner.py
+5
View File
@@ -1,6 +1,7 @@
import base64
import json
import math
import os
import resource
import time
from collections.abc import Generator
@@ -999,6 +1000,7 @@ def _validate_single_tool(obj: dict[str, Any]) -> ToolCallItem:
EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"
EXO_RUNNER_MUST_DIE = "EXO RUNNER MUST DIE"
def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
@@ -1014,6 +1016,9 @@ def _check_for_debug_prompts(task_params: TextGenerationTaskParams) -> None:
if not prompt:
return
if EXO_RUNNER_MUST_DIE in prompt:
logger.info("Abrupt process death triggered (simulates OOM kill)")
os._exit(1)
if EXO_RUNNER_MUST_FAIL in prompt:
logger.info("raising exception")
raise Exception("Artificial runner exception - for testing purposes only.")
src/exo/worker/runner/runner_supervisor.py
+112 -4
View File
@@ -1,12 +1,17 @@
from __future__ import annotations
import contextlib
import multiprocessing
import signal
from dataclasses import dataclass, field
from multiprocessing import Process
from multiprocessing.sharedctypes import Synchronized
from typing import Self
import anyio
from anyio import (
BrokenResourceError,
CancelScope,
ClosedResourceError,
to_thread,
)
@@ -26,6 +31,7 @@ from exo.shared.types.worker.runners import (
RunnerIdle,
RunnerLoading,
RunnerRunning,
RunnerShutdown,
RunnerShuttingDown,
RunnerStatus,
RunnerWarmingUp,
@@ -36,6 +42,8 @@ from exo.worker.runner.bootstrap import entrypoint
PREFILL_TIMEOUT_SECONDS = 60
DECODE_TIMEOUT_SECONDS = 5
HEALTH_CHECK_INTERVAL_SECONDS = 1
HEARTBEAT_STALE_CHECKS = 10
@dataclass(eq=False)
@@ -48,10 +56,14 @@ class RunnerSupervisor:
_task_sender: MpSender[Task]
_event_sender: Sender[Event]
_cancel_sender: MpSender[TaskId]
_heartbeat: Synchronized[int]
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)
_death_handled: bool = field(default=False, init=False)
_last_heartbeat_value: int = field(default=0, init=False)
_heartbeat_stale_count: int = field(default=0, init=False)
@classmethod
def create(
@@ -65,6 +77,8 @@ class RunnerSupervisor:
task_sender, task_recv = mp_channel[Task]()
cancel_sender, cancel_recv = mp_channel[TaskId]()
heartbeat: Synchronized[int] = multiprocessing.Value("Q", 0)
runner_process = Process(
target=entrypoint,
args=(
@@ -73,6 +87,7 @@ class RunnerSupervisor:
task_recv,
cancel_recv,
logger,
heartbeat,
),
daemon=True,
)
@@ -88,13 +103,16 @@ class RunnerSupervisor:
_task_sender=task_sender,
_cancel_sender=cancel_sender,
_event_sender=event_sender,
_heartbeat=heartbeat,
)
return self
async def run(self):
self.runner_process.start()
await self._forward_events()
async with anyio.create_task_group() as tg:
tg.start_soon(self._forward_events)
tg.start_soon(self._health_check, tg.cancel_scope)
def shutdown(self):
logger.info("Runner supervisor shutting down")
@@ -177,9 +195,99 @@ 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()
if not self._death_handled:
self._death_handled = True
await self._check_runner(e)
for tid in self.pending:
self.pending[tid].set()
async def _health_check(self, cancel_scope: CancelScope) -> None:
"""Periodically check if the runner process is alive and responsive.
Detects two failure modes:
1. Process death (e.g. OOM kill) without cleanly closing the event
channel, which would leave _forward_events blocked on queue.get().
2. Unresponsive process (e.g. frozen by OS memory pressure, deadlock)
detected via a stale heartbeat counter.
"""
while True:
await anyio.sleep(HEALTH_CHECK_INTERVAL_SECONDS)
if not self.runner_process.is_alive():
self._handle_process_exit(cancel_scope)
return
# Check heartbeat counter — if it hasn't changed between
# consecutive checks, the subprocess may be frozen.
current = self._heartbeat.value
if current > 0:
if current == self._last_heartbeat_value:
self._heartbeat_stale_count += 1
if self._heartbeat_stale_count >= HEARTBEAT_STALE_CHECKS:
logger.error(
f"Health check: runner process unresponsive "
f"(heartbeat stale for {self._heartbeat_stale_count} checks), killing"
)
self._handle_unresponsive(cancel_scope)
return
else:
self._heartbeat_stale_count = 0
self._last_heartbeat_value = current
def _handle_process_exit(self, cancel_scope: CancelScope) -> None:
"""Handle runner process that has exited."""
if not self._death_handled:
self._death_handled = True
if isinstance(
self.status, (RunnerShutdown, RunnerShuttingDown, RunnerFailed)
):
logger.info("Health check: runner process exited (expected)")
else:
rc = self.runner_process.exitcode
if isinstance(rc, int) and rc < 0:
sig = -rc
try:
cause = f"signal={sig} ({signal.strsignal(sig)})"
except Exception:
cause = f"signal={sig}"
else:
cause = f"exitcode={rc}"
logger.error(
f"Health check: runner process died unexpectedly ({cause})"
)
self._event_sender.send_nowait(
RunnerStatusUpdated(
runner_id=self.bound_instance.bound_runner_id,
runner_status=RunnerFailed(
error_message=f"Terminated ({cause})"
),
)
)
self.shutdown()
for tid in self.pending:
self.pending[tid].set()
cancel_scope.cancel()
def _handle_unresponsive(self, cancel_scope: CancelScope) -> None:
"""Handle runner process that is alive but unresponsive."""
if not self._death_handled:
self._death_handled = True
self._event_sender.send_nowait(
RunnerStatusUpdated(
runner_id=self.bound_instance.bound_runner_id,
runner_status=RunnerFailed(
error_message="Runner process unresponsive (heartbeat timeout)"
),
)
)
for tid in self.pending:
self.pending[tid].set()
self.shutdown()
cancel_scope.cancel()
def __del__(self) -> None:
if self.runner_process.is_alive():
src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py
+202 -1
View File
@@ -1 +1,202 @@
# TODO:
from __future__ import annotations
import multiprocessing
import os
import signal as signal_module
from collections.abc import Callable
from multiprocessing.sharedctypes import Synchronized
from typing import Any
import anyio
from exo.shared.types.events import Event, RunnerStatusUpdated
from exo.shared.types.tasks import Task, TaskId
from exo.shared.types.worker.runners import RunnerFailed, RunnerIdle, RunnerShutdown
from exo.utils.channels import Receiver, Sender, channel, mp_channel
from exo.worker.runner.runner_supervisor import (
HEALTH_CHECK_INTERVAL_SECONDS,
HEARTBEAT_STALE_CHECKS,
RunnerSupervisor,
)
from ...constants import (
INSTANCE_1_ID,
MODEL_A_ID,
NODE_A,
RUNNER_1_ID,
)
from ..conftest import get_bound_mlx_ring_instance
def _die_immediately() -> None:
"""Subprocess target that exits with a non-zero code."""
os._exit(1)
def _die_with_signal() -> None:
"""Subprocess target that kills itself with SIGKILL (simulates OOM)."""
os.kill(os.getpid(), signal_module.SIGKILL)
def _exit_cleanly() -> None:
"""Subprocess target that exits with code 0."""
os._exit(0)
def _hang_forever() -> None:
"""Subprocess target that hangs without updating heartbeat (simulates freeze)."""
import time
# Write one heartbeat so the supervisor starts tracking, then stop.
time.sleep(100000)
def _build_supervisor(
event_sender: Sender[Event],
target: Callable[..., Any],
) -> RunnerSupervisor:
"""Build a RunnerSupervisor with a custom subprocess target.
Uses a clone of event_sender (matching real Worker behavior) so that
closing the supervisor's copy doesn't close the test's receiver.
"""
bound_instance = get_bound_mlx_ring_instance(
instance_id=INSTANCE_1_ID,
model_id=MODEL_A_ID,
runner_id=RUNNER_1_ID,
node_id=NODE_A,
)
_ev_send, ev_recv = mp_channel[Event]()
task_sender, _task_recv = mp_channel[Task]()
runner_process = multiprocessing.Process(target=target, daemon=True)
heartbeat: Synchronized[int] = multiprocessing.Value("Q", 0)
return RunnerSupervisor(
bound_instance=bound_instance,
shard_metadata=bound_instance.bound_shard,
runner_process=runner_process,
initialize_timeout=10,
_ev_recv=ev_recv,
_task_sender=task_sender,
_event_sender=event_sender.clone(),
_heartbeat=heartbeat,
)
def _collect_failed_events(
event_receiver: Receiver[Event],
) -> list[RunnerFailed]:
"""Drain the receiver and return all RunnerFailed statuses."""
out: list[RunnerFailed] = []
while True:
try:
event = event_receiver.receive_nowait()
except Exception:
break
if isinstance(event, RunnerStatusUpdated) and isinstance(
event.runner_status, RunnerFailed
):
out.append(event.runner_status)
return out
async def test_health_check_detects_dead_process():
"""When the runner process dies with a non-zero exit code, the health check
should emit a RunnerFailed event and run() should return."""
event_sender, event_receiver = channel[Event]()
supervisor = _build_supervisor(event_sender, _die_immediately)
with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
await supervisor.run()
failures = _collect_failed_events(event_receiver)
assert len(failures) == 1
assert failures[0].error_message is not None
assert "exitcode=1" in failures[0].error_message
async def test_health_check_detects_signal_death():
"""When the runner process is killed by a signal (e.g. OOM -> SIGKILL),
the health check should report the signal in the failure message."""
event_sender, event_receiver = channel[Event]()
supervisor = _build_supervisor(event_sender, _die_with_signal)
with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
await supervisor.run()
failures = _collect_failed_events(event_receiver)
assert len(failures) == 1
assert failures[0].error_message is not None
assert "signal=9" in failures[0].error_message
async def test_health_check_releases_pending_tasks():
"""When the runner dies, any pending start_task() waiters should be unblocked."""
event_sender, _event_receiver = channel[Event]()
supervisor = _build_supervisor(event_sender, _die_immediately)
# Register a pending waiter as if start_task() was waiting for acknowledgement
task_event = anyio.Event()
tid = TaskId("pending-task")
supervisor.pending[tid] = task_event
with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
await supervisor.run()
assert task_event.is_set()
async def test_clean_exit_no_failure_when_shutdown_status():
"""When the runner was in RunnerShutdown status and exits with code 0,
no RunnerFailed event should be emitted."""
event_sender, event_receiver = channel[Event]()
supervisor = _build_supervisor(event_sender, _exit_cleanly)
# Simulate that the runner had already reported shutdown via events
supervisor.status = RunnerShutdown()
with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
await supervisor.run()
failures = _collect_failed_events(event_receiver)
assert len(failures) == 0
async def test_unexpected_exit_code_zero_emits_failure():
"""When the runner exits with code 0 but was NOT in a shutdown state,
this is unexpected and should still emit RunnerFailed."""
event_sender, event_receiver = channel[Event]()
supervisor = _build_supervisor(event_sender, _exit_cleanly)
assert isinstance(supervisor.status, RunnerIdle)
with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
await supervisor.run()
failures = _collect_failed_events(event_receiver)
assert len(failures) == 1
assert failures[0].error_message is not None
assert "exitcode=0" in failures[0].error_message
async def test_heartbeat_timeout_detects_unresponsive_process():
"""When the runner process is alive but its heartbeat goes stale,
the health check should kill it and emit RunnerFailed."""
event_sender, event_receiver = channel[Event]()
supervisor = _build_supervisor(event_sender, _hang_forever)
# Pre-seed the heartbeat counter with a non-zero value and set the
# supervisor's last-seen value to match so it appears stale immediately.
# Set stale count to HEARTBEAT_STALE_CHECKS - 1 so a single check triggers.
supervisor._heartbeat.value = 42 # pyright: ignore[reportPrivateUsage]
supervisor._last_heartbeat_value = 42 # pyright: ignore[reportPrivateUsage]
supervisor._heartbeat_stale_count = HEARTBEAT_STALE_CHECKS - 1 # pyright: ignore[reportPrivateUsage]
with anyio.fail_after(HEALTH_CHECK_INTERVAL_SECONDS + 5):
await supervisor.run()
failures = _collect_failed_events(event_receiver)
assert len(failures) == 1
assert failures[0].error_message is not None
assert "unresponsive" in failures[0].error_message.lower()