fix warmup order. should be rank!=0 then rank=0

ade3ee7ec5c4e8233fd4f04d033c87fc924c7369 · 2025-12-30 02:20:51 +0000 · Alex Cheema

Files touched

M src/exo/worker/plan.py

Diff

commit ade3ee7ec5c4e8233fd4f04d033c87fc924c7369
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Tue Dec 30 02:20:51 2025 +0000

    fix warmup order. should be rank!=0 then rank=0
---
 src/exo/worker/plan.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index e313ca73..6c621e72 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -235,9 +235,8 @@ def _ready_to_warmup(
         assert device_rank < world_size
         assert device_rank >= 0
 
-        # TODO: Ensure these align with MLX distributeds expectations.
-        # Rank < n-1
-        accepting_ranks_ready = device_rank < world_size - 1 and all(
+        # Rank != 0
+        accepting_ranks_ready = device_rank > 0 and all(
             isinstance(
                 all_runners.get(global_runner_id, None),
                 (RunnerLoaded, RunnerWarmingUp),
@@ -245,8 +244,8 @@ def _ready_to_warmup(
             for global_runner_id in shard_assignments.runner_to_shard
         )
 
-        # Rank = n-1
-        connecting_rank_ready = device_rank == world_size - 1 and all(
+        # Rank = 0
+        connecting_rank_ready = device_rank == 0 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

← fea42473 Place local node at the top of the dashboard. (#1033) · back to Exo · log EXO_LIBP2P_NAMESPACE on start 16e2bfd3 →