From 960e3e63a90bce5fa7854ced1e7d2368ea71fbc4 Mon Sep 17 00:00:00 2001 From: Alyta Phoenix Date: Thu, 30 Jul 2026 10:13:21 +0100 Subject: [PATCH] Fix runner-reconnect cascade: reset retry backoff on success, wait for teardown before recreate Investigated a report of a model failing to load on a two-node TB5 cluster. The instance actually loaded and served chat requests fine, then got dragged through a Shutdown/CreateRunner/ConnectToGroup cycle roughly every couple of minutes (not user-initiated), which sometimes crashed with "[jaccl] Changing queue pair to RTR failed with errno 16" (EBUSY) and was eventually left permanently unloaded. Root cause, as far as could be verified on this node's log alone: - plan.py's _kill_runner correctly shuts down every rank of an instance when any one rank fails (a half-pipeline is useless) -- this is by-design and left alone. - But main.py's per-instance retry-backoff counter (_instance_backoff) only ever reset on InstanceDeleted, never on a successful reconnect. So a node dragged into a restart cycle by a *sibling's* crash silently burned its own retry budget for a fault that was never its own, and could independently hit the 5-attempt cap and request deletion of an instance it was otherwise serving fine. - Separately, main.py's Shutdown handling popped the runner out of self.runners (unblocking the next CreateRunner) before the runner's underlying OS process had actually finished tearing down -- runner.shutdown() only requests cancellation; the real process kill and RDMA/queue-pair release happens asynchronously in the RunnerSupervisor.run() task, which nothing awaited. A fast Shutdown->CreateRunner cycle for the same instance could plausibly race that teardown, producing exactly the EBUSY symptom. Fixes: 1. worker/plan.py: new instance_to_reset_backoff() -- resets an instance's retry backoff when its local runner reaches RunnerReady/RunnerRunning (actually serving), not earlier states like RunnerConnected. Resetting at Connected would defeat the circuit breaker: a rank that connects fine but crashes every LoadModel (bad weights, OOM) would loop forever instead of eventually giving up. Wired into worker/main.py's _event_applier alongside the existing InstanceDeleted reset. 2. worker/runner/supervisor.py: RunnerSupervisor now exposes wait_stopped(), backed by an anyio.Event set once run()'s teardown (including the actual runner_process.stop()) has finished. worker/main.py's Shutdown handling now awaits this (bounded to 15s) before considering the slot free, closing the race between tearing down the old runner and creating its replacement. 3. master/main.py: the "kill broken instances" topology-eviction path had zero logging before silently sending InstanceDeleted. Added a warning naming the instance and the missing node, so a recurrence is attributable to this path instead of indistinguishable from the worker-side backoff path. Known limitations / behavior changes, called out for review: - I could not confirm which mechanism actually killed the reported instance. Master runs on a different physical node in this cluster, and its log (which would show whether master's topology-eviction path fired) isn't available from here. Fix 3 (the new logging) is what makes a recurrence diagnosable; this PR should not be read as a confirmed fix for "the model won't load," only for the two concrete bugs found by code inspection. - Behavior change: resetting the backoff on every successful Ready/Running means a *flapping* instance (repeatedly reaches serving, then dies, over and over) will now retry forever instead of giving up after 5 total attempts -- since each successful reconnect resets the counter back to 0. This seems like better UX (a transient hardware hiccup shouldn't cause permanent data loss of a model that keeps mostly working) but is a real change from "give up after 5 lifetime attempts" to "give up after 5 *consecutive* failures." Flagging in case the lifetime cap was intentional. - Fix 2 (wait_stopped) is unit-tested for its actual contract (blocks until the OS process exits, safe to await twice), but the EBUSY race it targets is RDMA-timing on hardware unavailable in this environment -- the fix is unverified against the real failure it's meant to close. Also, in the crash-cascade case, the healthy rank's generator.close() tears down a distributed group whose peer just died, which could itself stall -- the 15s timeout is what bounds that, not a guarantee it resolves quickly. Worth a look from someone with two-node TB5 hardware. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_011rjSfwDBTkmySmfU6NgHKF --- src/exo/master/main.py | 4 + src/exo/worker/main.py | 13 ++- src/exo/worker/plan.py | 29 +++++++ src/exo/worker/runner/supervisor.py | 24 ++++-- .../unittests/test_plan/test_backoff_reset.py | 82 +++++++++++++++++++ .../test_runner/test_runner_supervisor.py | 66 +++++++++++++++ 6 files changed, 211 insertions(+), 7 deletions(-) create mode 100644 src/exo/worker/tests/unittests/test_plan/test_backoff_reset.py diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 485ede30f7..2dc7de9139 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -475,6 +475,10 @@ async def _plan(self) -> None: for instance_id, instance in self.state.instances.items(): for node_id in instance.shard_assignments.node_to_runner: if node_id not in connected_node_ids: + logger.warning( + f"Deleting instance {instance_id}: node {node_id} " + "is no longer in topology" + ) await self.event_sender.send( InstanceDeleted(instance_id=instance_id) ) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index a641bacfb6..d1c515483c 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -58,7 +58,7 @@ from exo.utils.info_gatherer.net_profile import check_reachable from exo.utils.keyed_backoff import KeyedBackoff from exo.utils.task_group import TaskGroup -from exo.worker.plan import plan +from exo.worker.plan import instance_to_reset_backoff, plan from exo.worker.runner.supervisor import RunnerSupervisor @@ -147,6 +147,9 @@ async def _event_applier(self): if isinstance(event, InstanceDeleted): self._instance_backoff.reset(event.instance_id) + if (iid := instance_to_reset_backoff(event, self.runners)) is not None: + self._instance_backoff.reset(iid) + # Buffer input image chunks for image editing if isinstance(event, InputChunkReceived): cmd_id = event.command_id @@ -291,6 +294,14 @@ async def plan_step(self): ) finally: runner.shutdown() + # Wait for the runner's process (and whatever OS-level + # resources it held, e.g. an RDMA queue pair) to + # actually go away before the next plan() tick is free + # to create a replacement for this same instance -- + # otherwise a fast Shutdown->CreateRunner cycle can + # race the old process's teardown. + with anyio.move_on_after(15): + await runner.wait_stopped() case CancelTask( cancelled_task_id=cancelled_task_id, runner_id=runner_id ): diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 3824e4bb7a..d472d22440 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -4,6 +4,7 @@ from exo.shared.types.chunks import InputImageChunk from exo.shared.types.common import CommandId, ModelId, NodeId +from exo.shared.types.events import Event, RunnerStatusUpdated from exo.shared.types.tasks import ( CancelTask, ConnectToGroup, @@ -44,6 +45,34 @@ from exo.worker.runner.supervisor import RunnerSupervisor +def instance_to_reset_backoff( + event: Event, + runners: Mapping[RunnerId, RunnerSupervisor], +) -> InstanceId | None: + """Return the instance whose retry backoff should be cleared, if this + event shows our local runner reached a fully-serving state. + + The backoff must reset on success, not just on InstanceDeleted -- + otherwise a node dragged into repeated restarts by a *sibling* runner's + failure (see _kill_runner) silently burns its own retry budget for a + fault that was never its own, and can eventually request deletion of an + instance that this node was serving just fine. + + Reset happens at RunnerReady/RunnerRunning rather than earlier states + (e.g. RunnerConnected) so a rank that connects but keeps crashing later + (bad weights, OOM during load) still trips the retry-exhaustion circuit + breaker instead of looping forever. + """ + if not isinstance(event, RunnerStatusUpdated): + return None + if not isinstance(event.runner_status, (RunnerReady, RunnerRunning)): + return None + runner = runners.get(event.runner_id) + if runner is None: + return None + return runner.bound_instance.instance.instance_id + + def plan( node_id: NodeId, # Runners is expected to be FRESH and so should not come from state diff --git a/src/exo/worker/runner/supervisor.py b/src/exo/worker/runner/supervisor.py index 9611262473..83bd7e0c28 100644 --- a/src/exo/worker/runner/supervisor.py +++ b/src/exo/worker/runner/supervisor.py @@ -198,6 +198,7 @@ class RunnerSupervisor: _cancel_watch_runner: anyio.CancelScope = field( default_factory=anyio.CancelScope, init=False ) + _stopped: anyio.Event = field(default_factory=anyio.Event, init=False) @classmethod async def create( @@ -206,13 +207,14 @@ async def create( bound_instance: BoundInstance, event_sender: Sender[Event], initialize_timeout: float = 400, + target: Callable[..., object] = entrypoint, ) -> Self: ev_send, ev_recv = mp_channel[Event | RunnerTerminationError]() task_sender, task_recv = mp_channel[Task]() cancel_sender, cancel_recv = mp_channel[TaskId]() runner_process = AsyncProcess( - target=entrypoint, + target=target, args=( bound_instance, ev_send, @@ -266,15 +268,25 @@ async def run(self): with contextlib.suppress(ClosedResourceError): self._cancel_sender.close() - with anyio.CancelScope(shield=True): - await self.runner_process.stop() - logger.info( - f"Runner process successfully terminated: {self.runner_process.exitcode}" - ) + try: + with anyio.CancelScope(shield=True): + await self.runner_process.stop() + logger.info( + f"Runner process successfully terminated: {self.runner_process.exitcode}" + ) + finally: + self._stopped.set() def shutdown(self): self._tg.cancel_tasks() + async def wait_stopped(self) -> None: + """Wait until run() has fully finished, including the OS process + actually exiting. Used to make sure a runner's resources (e.g. an + RDMA queue pair) are released before a replacement is created for + the same slot.""" + await self._stopped.wait() + async def start_task(self, task: Task): if task.task_id in self.pending: logger.warning( diff --git a/src/exo/worker/tests/unittests/test_plan/test_backoff_reset.py b/src/exo/worker/tests/unittests/test_plan/test_backoff_reset.py new file mode 100644 index 0000000000..9d25be3e6e --- /dev/null +++ b/src/exo/worker/tests/unittests/test_plan/test_backoff_reset.py @@ -0,0 +1,82 @@ +from exo.shared.types.events import InstanceDeleted, RunnerStatusUpdated +from exo.shared.types.worker.instances import BoundInstance +from exo.shared.types.worker.runners import ( + RunnerConnected, + RunnerId, + RunnerReady, + RunnerRunning, +) +from exo.worker.plan import instance_to_reset_backoff +from exo.worker.tests.constants import ( + INSTANCE_1_ID, + MODEL_A_ID, + NODE_A, + RUNNER_1_ID, + RUNNER_2_ID, +) +from exo.worker.tests.unittests.conftest import ( + FakeRunnerSupervisor, + get_mlx_ring_instance, + get_pipeline_shard_metadata, +) + + +def _make_runners() -> dict[RunnerId, FakeRunnerSupervisor]: + shard = get_pipeline_shard_metadata(model_id=MODEL_A_ID, device_rank=0) + instance = get_mlx_ring_instance( + instance_id=INSTANCE_1_ID, + model_id=MODEL_A_ID, + node_to_runner={NODE_A: RUNNER_1_ID}, + runner_to_shard={RUNNER_1_ID: shard}, + ) + bound_instance = BoundInstance( + instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A + ) + return { + RUNNER_1_ID: FakeRunnerSupervisor( + bound_instance=bound_instance, status=RunnerConnected() + ) + } + + +def test_resets_backoff_when_local_runner_becomes_ready(): + event = RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerReady()) + + result = instance_to_reset_backoff(event, _make_runners()) # type: ignore[arg-type] + + assert result == INSTANCE_1_ID + + +def test_resets_backoff_when_local_runner_becomes_running(): + event = RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerRunning()) + + result = instance_to_reset_backoff(event, _make_runners()) # type: ignore[arg-type] + + assert result == INSTANCE_1_ID + + +def test_does_not_reset_on_runner_connected(): + """RunnerConnected is reached before LoadModel -- a rank that connects but + keeps crashing during load (bad weights, OOM) must still trip the + retry-exhaustion circuit breaker, so resetting here would be wrong.""" + event = RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerConnected()) + + result = instance_to_reset_backoff(event, _make_runners()) # type: ignore[arg-type] + + assert result is None + + +def test_does_not_reset_for_unknown_runner(): + event = RunnerStatusUpdated(runner_id=RUNNER_2_ID, runner_status=RunnerReady()) + + result = instance_to_reset_backoff(event, _make_runners()) # type: ignore[arg-type] + + assert result is None + + +def test_ignores_non_runner_status_events(): + event = InstanceDeleted(instance_id=INSTANCE_1_ID) + + result = instance_to_reset_backoff(event, _make_runners()) # type: ignore[arg-type] + + assert result is None 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 87cb9c7441..6836f5b1e3 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,3 +1,4 @@ +import time from typing import cast import anyio @@ -22,6 +23,10 @@ from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance +def _sleep_forever(*_args: object) -> None: + time.sleep(1000) + + class _DeadProcess: def __init__(self): rx1, _ = channel[bytes]() @@ -95,3 +100,64 @@ async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> event_sender.close() with anyio.move_on_after(0.1): await event_receiver.aclose() + + +@pytest.mark.anyio +async def test_wait_stopped_resolves_only_after_process_actually_exits() -> None: + """Regression test: main.py's Shutdown handling awaits wait_stopped() + before the next plan() tick is allowed to create a replacement runner for + the same instance. If wait_stopped() resolved before the OS process (and + whatever resources it held, e.g. an RDMA queue pair) actually went away, + a fast Shutdown->CreateRunner cycle could race the old process's + teardown.""" + event_sender, event_receiver = channel[Event]() + task_sender, _ = mp_channel[Task]() + cancel_sender, _ = mp_channel[TaskId]() + _, ev_recv = mp_channel[Event | RunnerTerminationError]() + + bound_instance: BoundInstance = get_bound_mlx_ring_instance( + instance_id=InstanceId("instance-a"), + model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"), + runner_id=RunnerId("runner-a"), + node_id=NodeId("node-a"), + ) + + runner_process = AsyncProcess(target=_sleep_forever, args=(), daemon=True) + handler = await RunnerStdioHandler.create( + stdout_rx=runner_process.stdout, stderr_rx=runner_process.stderr + ) + supervisor = RunnerSupervisor( + shard_metadata=bound_instance.bound_shard, + bound_instance=bound_instance, + runner_process=runner_process, + _runner_stdio_handler=handler, + initialize_timeout=400, + _ev_recv=ev_recv, + _task_sender=task_sender, + _event_sender=event_sender, + _cancel_sender=cancel_sender, + ) + + async with anyio.create_task_group() as tg: + tg.start_soon(supervisor.run) + + with anyio.fail_after(5): + while not runner_process.is_alive(): + await anyio.sleep(0.01) + + assert not supervisor._stopped.is_set() # pyright: ignore[reportPrivateUsage] + + supervisor.shutdown() + + with anyio.fail_after(10): + await supervisor.wait_stopped() + + assert not runner_process.is_alive() + + # Safe to await again once already stopped (level-triggered event). + with anyio.fail_after(1): + await supervisor.wait_stopped() + + event_sender.close() + with anyio.move_on_after(0.1): + await event_receiver.aclose()