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