plan fix

7312a7e00029973c05a4e64be3ce33fbd0689387 · 2025-12-05 16:43:11 +0000 · Evan Quiney

Files touched

M src/exo/worker/plan.py

Diff

commit 7312a7e00029973c05a4e64be3ce33fbd0689387
Author: Evan Quiney <evanev7@gmail.com>
Date:   Fri Dec 5 16:43:11 2025 +0000

    plan fix
---
 src/exo/worker/plan.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 63ab9cfd..9d1806ad 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -174,11 +174,15 @@ def _ready_to_warmup(
         shard = runner.bound_instance.bound_shard
         device_rank = shard.device_rank
         runner_id = runner.bound_instance.bound_runner_id
+        world_size = shard.world_size
 
         is_runner_loaded = isinstance(runner.status, RunnerLoaded)
 
-        # Rank != 0
-        all_runners_loaded_or_warming_up = all(
+        assert device_rank < world_size
+        assert device_rank >= 0
+
+        # Rank != n-1
+        accepting_ranks_ready = device_rank != world_size - 1 and all(
             isinstance(
                 all_runners.get(global_runner_id, None),
                 (RunnerLoaded, RunnerWarmingUp),
@@ -186,17 +190,14 @@ def _ready_to_warmup(
             for global_runner_id in shard_assignments.runner_to_shard
         )
 
-        # Rank= 0
-        all_other_runners_warming_up = all(
+        # Rank = n-1
+        connecting_rank_ready = device_rank == world_size - 1 and all(
             isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp)
             for global_runner_id in shard_assignments.runner_to_shard
             if global_runner_id != runner_id
         )
 
-        nonzero_rank_ready = device_rank != 0 and all_runners_loaded_or_warming_up
-        zero_rank_ready = device_rank == 0 and all_other_runners_warming_up
-
-        if is_runner_loaded and (nonzero_rank_ready or zero_rank_ready):
+        if is_runner_loaded and (accepting_ranks_ready or connecting_rank_ready):
             return StartWarmup(instance_id=instance.instance_id)
 
     return None

← 9e0a1c23 rename ibv to jaccl inline with mlx · back to Exo · fmt: format all python/rust/nix files 56299838 →