diff --git a/e2e/test_runner_chaos.py b/e2e/test_runner_chaos.py new file mode 100644 index 00000000..0052643d --- /dev/null +++ b/e2e/test_runner_chaos.py @@ -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()) diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index ed420aab..4e6db53c 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/src/exo/worker/runner/bootstrap.py @@ -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() diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index ad6c78f6..626f7208 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -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.") diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 5d39a881..31fa2d2c 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -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(): diff --git a/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py b/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py index e151d4aa..30b8f833 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py +++ b/src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py @@ -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()