← back to Exo

Prioritise tb for ring instances (#1556)

f662c129dd228a5de66bf2db759db8386b4a114e · 2026-02-19 21:32:48 +0000 · rltakashige

## Motivation

TB has better bandwidth and latency than ethernet. We should prioritise
TB5 where possible. This drastically improves distributed image
generation performance.

## Test Plan

### Manual Testing
Saw on the dashboard that TB (169.254) addresses were prioritised.

Tested that image models scale much better.

### Automated Testing
No regression on Kimi K2.5

Files touched

• M src/exo/master/placement_utils.py
• M src/exo/worker/runner/runner_supervisor.py

Diff

commit f662c129dd228a5de66bf2db759db8386b4a114e
Author: rltakashige <rl.takashige@gmail.com>
Date:   Thu Feb 19 21:32:48 2026 +0000

    Prioritise tb for ring instances (#1556)
    
    ## Motivation
    
    TB has better bandwidth and latency than ethernet. We should prioritise
    TB5 where possible. This drastically improves distributed image
    generation performance.
    
    ## Test Plan
    
    ### Manual Testing
    Saw on the dashboard that TB (169.254) addresses were prioritised.
    
    Tested that image models scale much better.
    
    ### Automated Testing
    No regression on Kimi K2.5
---
 src/exo/master/placement_utils.py          | 35 ++++++++++++++++++++++--------
 src/exo/worker/runner/runner_supervisor.py |  3 ++-
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py
index b80d70c0..eb78c8fa 100644
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -341,6 +341,7 @@ def _find_ip_prioritised(
     other_node_id: NodeId,
     cycle_digraph: Topology,
     node_network: Mapping[NodeId, NodeNetworkInfo],
+    ring: bool,
 ) -> str | None:
     """Find an IP address between nodes with prioritization.
 
@@ -353,13 +354,27 @@ def _find_ip_prioritised(
     ip_to_type = {
         iface.ip_address: iface.interface_type for iface in other_network.interfaces
     }
-    priority = {
-        "ethernet": 0,
-        "wifi": 1,
-        "unknown": 2,
-        "maybe_ethernet": 3,
-        "thunderbolt": 4,
-    }
+
+    # Ring should prioritise fastest connection. As a best-effort, we prioritise TB.
+    # TODO: Profile and get actual connection speeds.
+    if ring:
+        priority = {
+            "thunderbolt": 0,
+            "maybe_ethernet": 1,
+            "ethernet": 2,
+            "wifi": 3,
+            "unknown": 4,
+        }
+
+    # RDMA prefers ethernet coordinator
+    else:
+        priority = {
+            "ethernet": 0,
+            "wifi": 1,
+            "unknown": 2,
+            "maybe_ethernet": 3,
+            "thunderbolt": 4,
+        }
     return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
 
 
@@ -399,7 +414,7 @@ def get_mlx_ring_hosts_by_node(
                 continue
 
             connection_ip = _find_ip_prioritised(
-                node_id, other_node_id, cycle_digraph, node_network
+                node_id, other_node_id, cycle_digraph, node_network, ring=True
             )
             if connection_ip is None:
                 raise ValueError(
@@ -430,7 +445,9 @@ def get_mlx_jaccl_coordinators(
         if n == coordinator:
             return "0.0.0.0"
 
-        ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network)
+        ip = _find_ip_prioritised(
+            n, coordinator, cycle_digraph, node_network, ring=False
+        )
         if ip is not None:
             return ip
 
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index 2cee74fd..e8a06a77 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/src/exo/worker/runner/runner_supervisor.py
@@ -100,7 +100,8 @@ class RunnerSupervisor:
         logger.info("Runner supervisor shutting down")
         self._ev_recv.close()
         self._task_sender.close()
-        self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
+        with contextlib.suppress(ClosedResourceError):
+            self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
         self._cancel_sender.close()
         self.runner_process.join(5)
         if not self.runner_process.is_alive():

← c45ff9ad memory tidy (#1558)  ·  back to Exo  ·  Ensure coordinator is rank 0 (#1559) 3006c8ea →