From 90343001635bd9a0406f5115198ceea397c2aeff Mon Sep 17 00:00:00 2001 From: Michael Harrigan <9045677+michaelharrigan@users.noreply.github.com> Date: Thu, 26 Mar 2026 13:28:47 -0400 Subject: [PATCH] [Fix] Node hang on reelection (#1801) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation During master reelection, `_elect_loop` called `worker.shutdown()` (fire & forget) then immediately created and started a new Worker. This caused the old runner subprocess's Metal/GPU teardown to race with the new worker's startup, resulting in `IOConnectUnmapMemory failed: kr=0xe00002bc` errors and a full node hang requiring `^C`. Same issue existed for `DownloadCoordinator`. ## Changes - Added `anyio.Event`-based `_stopped` signal to `Worker` and `DownloadCoordinator`, set at the end of their `run()` finally blocks - Added `wait_stopped()` async method to both classes - Updated `_elect_loop` to `await wait_stopped()` after calling `shutdown()` on the old Worker and DownloadCoordinator before creating replacements ## Why It Works The old Worker's task group contains the RunnerSupervisor tasks, whose finally blocks join the runner subprocess (with 5s timeout + SIGTERM + SIGKILL escalation). By awaiting `wait_stopped()`, we guarantee the old runner process has fully exited — including GPU memory cleanup — before a new Worker can start and potentially access the GPU. This eliminates the race without changing the shutdown mechanics themselves. ## Test Plan ### Manual Testing Hardware: M4 Pro Mac Mini 24GB + M3 Ultra Mac Studio 96GB, connected via Thunderbolt **Repro steps:** 1. Start exo on two nodes with a model sharded across both (e.g. `Josiefied-Qwen3-14B-abliterated-v3-4bit`) 2. Wait for "runner ready" on both 3. `kill -9` the master node 4. Observe the surviving node's re-election behavior **Before fix (original crash):** ``` [ 11:02:39.0896AM ] Runner supervisor shutting down [ 11:02:39.0905AM ] bye from the runner [ 11:02:39.1052AM ] Stopping Worker IOConnectUnmapMemory failed: kr=0xe00002bc IOConnectUnmapMemory failed: kr=0xe00002bc IOConnectUnmapMemory failed: kr=0xe00002bc IOConnectUnmapMemory failed: kr=0xe00002bc ^C[ 11:03:45 ] ← hung for over a minute, required manual kill ``` **After fix (clean re-election):** ``` [ 12:15:22.4703PM ] runner loaded [ 12:15:24.1672PM ] runner ready [ 12:15:33.5393PM ] Waiting for other campaign to finish [ 12:15:36.5409PM ] Node elected Master [ 12:15:36.5413PM ] Unpausing API ``` No `IOConnectUnmapMemory` errors, no hang, no `^C` needed. ### Automated Testing - No existing tests cover the `_elect_loop` re-election path; this is an integration-level flow requiring a live router/election/worker stack - All existing tests pass (307/308, 1 pre-existing Rust binding failure) - basedpyright: 0 errors, ruff: all checks passed --------- Co-authored-by: Evan --- src/exo/download/coordinator.py | 13 +++++++++---- src/exo/download/tests/test_re_download.py | 2 +- src/exo/main.py | 4 ++-- src/exo/worker/main.py | 5 ++++- 4 files changed, 16 insertions(+), 8 deletions(-) diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 5c55970f..06ff54ee 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -49,6 +49,7 @@ class DownloadCoordinator: active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict) _tg: TaskGroup = field(init=False, default_factory=TaskGroup) + _stopped: anyio.Event = field(init=False, default_factory=anyio.Event) # Per-model throttle for download progress events _last_progress_time: dict[ModelId, float] = field(default_factory=dict) @@ -100,12 +101,16 @@ class DownloadCoordinator: logger.info( f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}" ) - async with self._tg as tg: - tg.start_soon(self._command_processor) - tg.start_soon(self._emit_existing_download_progress) + try: + async with self._tg as tg: + tg.start_soon(self._command_processor) + tg.start_soon(self._emit_existing_download_progress) + finally: + self._stopped.set() - def shutdown(self) -> None: + async def shutdown(self) -> None: self._tg.cancel_tasks() + await self._stopped.wait() async def _command_processor(self) -> None: with self.download_command_receiver as commands: diff --git a/src/exo/download/tests/test_re_download.py b/src/exo/download/tests/test_re_download.py index 3f159814..75617997 100644 --- a/src/exo/download/tests/test_re_download.py +++ b/src/exo/download/tests/test_re_download.py @@ -186,7 +186,7 @@ async def test_re_download_after_delete_completes() -> None: "Re-download after deletion should complete" ) finally: - coordinator.shutdown() + await coordinator.shutdown() coordinator_task.cancel() with contextlib.suppress(asyncio.CancelledError): await coordinator_task diff --git a/src/exo/main.py b/src/exo/main.py index 93922cc2..68cd70d4 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -228,7 +228,7 @@ class Node: ) if result.is_new_master: if self.download_coordinator: - self.download_coordinator.shutdown() + await self.download_coordinator.shutdown() self.download_coordinator = DownloadCoordinator( self.node_id, exo_shard_downloader(offline=self.offline), @@ -240,7 +240,7 @@ class Node: ) self._tg.start_soon(self.download_coordinator.run) if self.worker: - self.worker.shutdown() + await self.worker.shutdown() # TODO: add profiling etc to resource monitor self.worker = Worker( self.node_id, diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 4bfed825..833f9134 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -80,6 +80,7 @@ class Worker: self.input_chunk_counts: dict[CommandId, int] = {} self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0) + self._stopped: anyio.Event = anyio.Event() async def run(self): logger.info("Starting Worker") @@ -102,6 +103,7 @@ class Worker: self.download_command_sender.close() for runner in self.runners.values(): runner.shutdown() + self._stopped.set() async def _forward_info(self, recv: Receiver[GatheredInfo]): with recv as info_stream: @@ -280,8 +282,9 @@ class Worker: case task: await self._start_runner_task(task) - def shutdown(self): + async def shutdown(self): self._tg.cancel_tasks() + await self._stopped.wait() async def _start_runner_task(self, task: Task): if (instance := self.state.instances.get(task.instance_id)) is not None: