[object Object]

← back to Exo

[Fix] Node hang on reelection (#1801)

90343001635bd9a0406f5115198ceea397c2aeff · 2026-03-26 13:28:47 -0400 · Michael Harrigan

## 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 <evanev7@gmail.com>

Files touched

Diff

commit 90343001635bd9a0406f5115198ceea397c2aeff
Author: Michael Harrigan <9045677+michaelharrigan@users.noreply.github.com>
Date:   Thu Mar 26 13:28:47 2026 -0400

    [Fix]  Node hang on reelection (#1801)
    
    ## 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 <evanev7@gmail.com>
---
 src/exo/download/coordinator.py            | 15 ++++++++++-----
 src/exo/download/tests/test_re_download.py |  2 +-
 src/exo/main.py                            |  4 ++--
 src/exo/worker/main.py                     |  5 ++++-
 4 files changed, 17 insertions(+), 9 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)
-
-    def shutdown(self) -> None:
+        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()
+
+    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:

← 1d1dfaa1 Don't download original/ and metal/ folders from HF (#1800)  ·  back to Exo  ·  Rework model storage directory management (for external stor 15f1b61f →