← back to Exo

fix(inference): prevent TP collective deadlock via agree_on_tasks order (#2048)

89d20c1888359c8d1db4956211f72b6a2b36199e · 2026-05-06 07:24:58 -0500 · Drifter4242

If you have two machines and make two requests at the same time, it can
crash. This is because the tasks can sometimes end up in different
orders on different machines. We need to sort the tasks and
mx_all_gather_tasks already sorts the tasks but the code ignores that
ordering. The fix is to make sure the sort order is preserved.

The rest is written by Sonnet (reviewed by me):

Tensor-parallel inference requires that every rank enqueues tasks in the
same order before running agree_on_tasks collectives. The old
implementation filtered from _maybe_queue:

self._queue.extend(task for task in self._maybe_queue if task in agreed)
self._maybe_queue = [task for task in self._maybe_queue if task in
different]

Because _maybe_queue is independently ordered per-rank (tasks arrive via
gRPC in whatever order the API server sends them), two concurrent
requests could produce different _maybe_queue orderings on rank 0 vs
rank 1. The filter then preserved those different orders into _queue, so
each rank started processing tasks in a different sequence. The next mlx
collective (all_reduce, all_gather, etc.) on rank 0 corresponded to a
different task than on rank 1 → permanent deadlock.

Fix: extend from agreed directly. mx_all_gather_tasks returns agreed as
a list sorted by task_id on all ranks, so every rank appends the same
sequence regardless of local arrival order.

Applies to both SequentialGenerator and BatchGenerator.

## Motivation

`agree_on_tasks` is called on every rank after accumulating new requests
in
`_maybe_queue`. Its job is to run an `all_gather` collective so all
ranks agree
on which tasks to promote to `_queue` before the next inference step.

The old implementation re-imposed **local arrival order** when extending
`_queue`:

```python
self._queue.extend(task for task in self._maybe_queue if task in agreed)
```

`mx_all_gather_tasks` already returns `agreed` sorted by `task_id` — the
same
deterministic order on every rank. But iterating `self._maybe_queue`
instead of
`agreed` discarded that sort and substituted the local gRPC arrival
order, which
differs per rank under concurrent load. Two concurrent requests arriving
in
`[A, B]` order on rank 0 and `[B, A]` on rank 1 caused the first MLX
collective
in the next step to hang permanently: each rank was executing a
different task's
collective and would never match.

## Changes

`SequentialGenerator.agree_on_tasks` and
`BatchGenerator.agree_on_tasks`:

```python
# Before
self._queue.extend(task for task in self._maybe_queue if task in agreed)
self._maybe_queue = [task for task in self._maybe_queue if task in different]

# After
self._queue.extend(agreed)          # preserves mx_all_gather_tasks sort order
self._maybe_queue = list(different) # already in local order; filter was redundant
```

## Why It Works

`mx_all_gather_tasks` (in `utils_mlx.py`) computes the agreed set then
sorts by
`task_id`:

```python
agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
```

Because `task_id` is a UUID and the sort is lexicographic, every rank
produces
the same `agreed` list regardless of local arrival order. Using `agreed`
directly
preserves this guarantee. The `different` list (tasks not yet seen on
all ranks)
is built by iterating `tasks` in local order, which is already correct.

## Test Plan

### Manual Testing

**Hardware:** 2× Mac Studio M3 Ultra 512 GB, Thunderbolt 5 direct
bridge,
`MlxJaccl` RDMA tensor-parallel (`moonshotai/Kimi-K2.6`, 595 GB INT4, 61
layers).

- Sent concurrent streaming requests; confirmed all complete without
deadlock.
- This hardware configuration (sub-millisecond inter-node latency) is
the most
likely to trigger the race, as requests from separate HTTP connections
can
reach rank 0 and rank 1 in opposite order before `agree_on_tasks` runs.

### Automated Testing

All existing tests pass: `pytest src -m "not slow"
--import-mode=importlib`
— 422/422 passed. The existing `test_event_ordering.py` covers the
`agree_on_tasks` call path with a mock that returns tasks in consistent
order;
the race requires real distributed hardware to reproduce
deterministically.

Files touched

• M src/exo/worker/runner/llm_inference/batch_generator.py

Diff

commit 89d20c1888359c8d1db4956211f72b6a2b36199e
Author: Drifter4242 <davehind@yahoo.co.uk>
Date:   Wed May 6 07:24:58 2026 -0500

    fix(inference): prevent TP collective deadlock via agree_on_tasks order (#2048)
    
    If you have two machines and make two requests at the same time, it can
    crash. This is because the tasks can sometimes end up in different
    orders on different machines. We need to sort the tasks and
    mx_all_gather_tasks already sorts the tasks but the code ignores that
    ordering. The fix is to make sure the sort order is preserved.
    
    The rest is written by Sonnet (reviewed by me):
    
    Tensor-parallel inference requires that every rank enqueues tasks in the
    same order before running agree_on_tasks collectives. The old
    implementation filtered from _maybe_queue:
    
    self._queue.extend(task for task in self._maybe_queue if task in agreed)
    self._maybe_queue = [task for task in self._maybe_queue if task in
    different]
    
    Because _maybe_queue is independently ordered per-rank (tasks arrive via
    gRPC in whatever order the API server sends them), two concurrent
    requests could produce different _maybe_queue orderings on rank 0 vs
    rank 1. The filter then preserved those different orders into _queue, so
    each rank started processing tasks in a different sequence. The next mlx
    collective (all_reduce, all_gather, etc.) on rank 0 corresponded to a
    different task than on rank 1 → permanent deadlock.
    
    Fix: extend from agreed directly. mx_all_gather_tasks returns agreed as
    a list sorted by task_id on all ranks, so every rank appends the same
    sequence regardless of local arrival order.
    
    Applies to both SequentialGenerator and BatchGenerator.
    
    ## Motivation
    
    `agree_on_tasks` is called on every rank after accumulating new requests
    in
    `_maybe_queue`. Its job is to run an `all_gather` collective so all
    ranks agree
    on which tasks to promote to `_queue` before the next inference step.
    
    The old implementation re-imposed **local arrival order** when extending
    `_queue`:
    
    ```python
    self._queue.extend(task for task in self._maybe_queue if task in agreed)
    ```
    
    `mx_all_gather_tasks` already returns `agreed` sorted by `task_id` — the
    same
    deterministic order on every rank. But iterating `self._maybe_queue`
    instead of
    `agreed` discarded that sort and substituted the local gRPC arrival
    order, which
    differs per rank under concurrent load. Two concurrent requests arriving
    in
    `[A, B]` order on rank 0 and `[B, A]` on rank 1 caused the first MLX
    collective
    in the next step to hang permanently: each rank was executing a
    different task's
    collective and would never match.
    
    ## Changes
    
    `SequentialGenerator.agree_on_tasks` and
    `BatchGenerator.agree_on_tasks`:
    
    ```python
    # Before
    self._queue.extend(task for task in self._maybe_queue if task in agreed)
    self._maybe_queue = [task for task in self._maybe_queue if task in different]
    
    # After
    self._queue.extend(agreed)          # preserves mx_all_gather_tasks sort order
    self._maybe_queue = list(different) # already in local order; filter was redundant
    ```
    
    ## Why It Works
    
    `mx_all_gather_tasks` (in `utils_mlx.py`) computes the agreed set then
    sorts by
    `task_id`:
    
    ```python
    agreed = [local_tasks[tid] for tid in sorted(agreed_ids)]
    ```
    
    Because `task_id` is a UUID and the sort is lexicographic, every rank
    produces
    the same `agreed` list regardless of local arrival order. Using `agreed`
    directly
    preserves this guarantee. The `different` list (tasks not yet seen on
    all ranks)
    is built by iterating `tasks` in local order, which is already correct.
    
    ## Test Plan
    
    ### Manual Testing
    
    **Hardware:** 2× Mac Studio M3 Ultra 512 GB, Thunderbolt 5 direct
    bridge,
    `MlxJaccl` RDMA tensor-parallel (`moonshotai/Kimi-K2.6`, 595 GB INT4, 61
    layers).
    
    - Sent concurrent streaming requests; confirmed all complete without
    deadlock.
    - This hardware configuration (sub-millisecond inter-node latency) is
    the most
    likely to trigger the race, as requests from separate HTTP connections
    can
    reach rank 0 and rank 1 in opposite order before `agree_on_tasks` runs.
    
    ### Automated Testing
    
    All existing tests pass: `pytest src -m "not slow"
    --import-mode=importlib`
    — 422/422 passed. The existing `test_event_ordering.py` covers the
    `agree_on_tasks` call path with a mock that returns tasks in consistent
    order;
    the race requires real distributed hardware to reproduce
    deterministically.
---
 src/exo/worker/runner/llm_inference/batch_generator.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py
index 3d898d83..098c829e 100644
--- a/src/exo/worker/runner/llm_inference/batch_generator.py
+++ b/src/exo/worker/runner/llm_inference/batch_generator.py
@@ -138,8 +138,10 @@ class SequentialGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -373,8 +375,10 @@ class BatchGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""

← dbcceaa5 Initialise _cancelled_tasks in ImageEngine (#2051)  ·  back to Exo  ·  fix(placement): gate RDMA on nodeRdmaCtl.enabled at both end a0c00f9d →