Fix runner-reconnect cascade: reset retry backoff on success, wait for teardown before recreate

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)
                         )

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
                 ):

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

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(

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

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()