Ensure coordinator is rank 0 (#1559)

3006c8ea4ee43addf90e59e71a44c34135afba05 · 2026-02-20 11:46:24 +0000 · rltakashige

## Motivation

Coordinator can be a random rank. Let's just fix this to rank 0 as
that's what we typically assume.

## Test Plan

### Manual Testing
Works as normal on 2 nodes.


Let's wait for a little more testing to merge this.

---------

Co-authored-by: Evan <evanev7@gmail.com>

Files touched

M src/exo/master/placement.py

Diff

commit 3006c8ea4ee43addf90e59e71a44c34135afba05
Author: rltakashige <rl.takashige@gmail.com>
Date:   Fri Feb 20 11:46:24 2026 +0000

    Ensure coordinator is rank 0 (#1559)
    
    ## Motivation
    
    Coordinator can be a random rank. Let's just fix this to rank 0 as
    that's what we typically assume.
    
    ## Test Plan
    
    ### Manual Testing
    Works as normal on 2 nodes.
    
    
    Let's wait for a little more testing to merge this.
    
    ---------
    
    Co-authored-by: Evan <evanev7@gmail.com>
---
 src/exo/master/placement.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index cf31ca78..a0604055 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -141,15 +141,29 @@ def place_instance(
     if len(selected_cycle) == 1:
         command.instance_meta = InstanceMeta.MlxRing
 
-    # TODO: Single node instances
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
+            # TODO(evan): shard assignments should contain information about ranks, this is ugly
+            def get_device_rank(node_id: NodeId) -> int:
+                runner_id = shard_assignments.node_to_runner[node_id]
+                shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
+                assert shard_metadata is not None
+                return shard_metadata.device_rank
+
+            zero_node_ids = [
+                node_id
+                for node_id in selected_cycle.node_ids
+                if get_device_rank(node_id) == 0
+            ]
+            assert len(zero_node_ids) == 1
+            coordinator_node_id = zero_node_ids[0]
+
             mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
                 [node_id for node_id in selected_cycle],
                 cycle_digraph,
             )
             mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
-                coordinator=selected_cycle.node_ids[0],
+                coordinator=coordinator_node_id,
                 coordinator_port=random_ephemeral_port(),
                 cycle_digraph=cycle_digraph,
                 node_network=node_network,

← f662c129 Prioritise tb for ring instances (#1556) · back to Exo · fix: correct misleading docstring in seed_models (#1561) a16ff2c0 →