← back to Exo
placement: generate per-node host lists for MLX ring backend
b5d424b6587923fa0f5a61e37d8e24499ece5209 · 2025-12-23 22:18:38 +0000 · Jake Hillion
Pipeline + MLX Ring worked with 2 nodes but failed to initialize with
3 or more nodes. The MLX ring backend requires each node to know its
specific left and right neighbors in the ring, but the previous
implementation provided a single flat host list shared by all nodes.
With 2 nodes, a flat list [host0, host1] accidentally worked because
each node could find its only neighbor. With 3+ nodes, each node needs
a customized view:
- Rank 0: [self, right_neighbor, placeholder]
- Rank 1: [left_neighbor, self, right_neighbor]
- Rank 2: [placeholder, left_neighbor, self]
Changed MlxRingInstance from `hosts: list[Host]` to
`hosts_by_node: dict[NodeId, list[Host]]` with `ephemeral_port: int`.
Added `get_mlx_ring_hosts_by_node()` which generates per-node host
lists where:
- Self position uses 0.0.0.0 for local binding
- Left/right neighbors use actual connection IPs
- Non-neighbors use 198.51.100.1 (RFC 5737 TEST-NET-2 placeholder)
Also added IP prioritization (en0 > en1 > non-Thunderbolt > any) to
prefer stable network interfaces.
Fixed topology discovery recording loopback addresses (127.0.0.1) as
valid connections to remote nodes. The reachability check now verifies
node identity via HTTP GET /node_id rather than just checking if the
port is open.
Test plan:
- Built a DMG [0]
- Installed on all Macs and started cluster.
- Requested a 3 node Pipeline + MLX Ring Llama 3.3 70B (FP16).
- It started and I was able to send a few chat messages.
Eventually my instance seemed to get into a broken state and chat
stopped working, but this commit is a clear step forward.
[0] https://github.com/exo-explore/exo/actions/runs/20473983471/job/58834969418
Files touched
M src/exo/master/placement.pyM src/exo/master/placement_utils.pyM src/exo/master/tests/test_master.pyM src/exo/master/tests/test_placement.pyM src/exo/shared/types/worker/instances.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/tests/unittests/conftest.pyM src/exo/worker/tests/unittests/test_plan/test_warmup.pyM src/exo/worker/utils/net_profile.pyM uv.lock
Diff
commit b5d424b6587923fa0f5a61e37d8e24499ece5209
Author: Jake Hillion <jake@hillion.co.uk>
Date: Tue Dec 23 22:18:38 2025 +0000
placement: generate per-node host lists for MLX ring backend
Pipeline + MLX Ring worked with 2 nodes but failed to initialize with
3 or more nodes. The MLX ring backend requires each node to know its
specific left and right neighbors in the ring, but the previous
implementation provided a single flat host list shared by all nodes.
With 2 nodes, a flat list [host0, host1] accidentally worked because
each node could find its only neighbor. With 3+ nodes, each node needs
a customized view:
- Rank 0: [self, right_neighbor, placeholder]
- Rank 1: [left_neighbor, self, right_neighbor]
- Rank 2: [placeholder, left_neighbor, self]
Changed MlxRingInstance from `hosts: list[Host]` to
`hosts_by_node: dict[NodeId, list[Host]]` with `ephemeral_port: int`.
Added `get_mlx_ring_hosts_by_node()` which generates per-node host
lists where:
- Self position uses 0.0.0.0 for local binding
- Left/right neighbors use actual connection IPs
- Non-neighbors use 198.51.100.1 (RFC 5737 TEST-NET-2 placeholder)
Also added IP prioritization (en0 > en1 > non-Thunderbolt > any) to
prefer stable network interfaces.
Fixed topology discovery recording loopback addresses (127.0.0.1) as
valid connections to remote nodes. The reachability check now verifies
node identity via HTTP GET /node_id rather than just checking if the
port is open.
Test plan:
- Built a DMG [0]
- Installed on all Macs and started cluster.
- Requested a 3 node Pipeline + MLX Ring Llama 3.3 70B (FP16).
- It started and I was able to send a few chat messages.
Eventually my instance seemed to get into a broken state and chat
stopped working, but this commit is a clear step forward.
[0] https://github.com/exo-explore/exo/actions/runs/20473983471/job/58834969418
---
src/exo/master/placement.py | 19 ++-
src/exo/master/placement_utils.py | 119 ++++++++++++++++-
src/exo/master/tests/test_master.py | 52 ++++----
src/exo/master/tests/test_placement.py | 51 +++----
src/exo/shared/types/worker/instances.py | 3 +-
src/exo/worker/engines/mlx/utils_mlx.py | 5 +-
src/exo/worker/main.py | 7 +-
src/exo/worker/plan.py | 8 +-
src/exo/worker/tests/unittests/conftest.py | 3 +-
.../tests/unittests/test_plan/test_warmup.py | 147 ++++++++++++++++++++-
src/exo/worker/utils/net_profile.py | 71 +++++++---
uv.lock | 2 +-
12 files changed, 383 insertions(+), 104 deletions(-)
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index 68567b87..e7efa51e 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -7,9 +7,9 @@ from loguru import logger
from exo.master.placement_utils import (
filter_cycles_by_memory,
- get_hosts_from_subgraph,
get_mlx_ibv_devices_matrix,
get_mlx_jaccl_coordinators,
+ get_mlx_ring_hosts_by_node,
get_shard_assignments,
get_smallest_cycles,
)
@@ -19,7 +19,6 @@ from exo.shared.types.commands import (
DeleteInstance,
PlaceInstance,
)
-from exo.shared.types.common import Host
from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
from exo.shared.types.memory import Memory
from exo.shared.types.topology import NodeInfo
@@ -130,17 +129,17 @@ def place_instance(
jaccl_coordinators=mlx_jaccl_coordinators,
)
case InstanceMeta.MlxRing:
- hosts: list[Host] = get_hosts_from_subgraph(cycle_digraph)
+ ephemeral_port = random_ephemeral_port()
+ hosts_by_node = get_mlx_ring_hosts_by_node(
+ selected_cycle=selected_cycle,
+ cycle_digraph=cycle_digraph,
+ ephemeral_port=ephemeral_port,
+ )
target_instances[instance_id] = MlxRingInstance(
instance_id=instance_id,
shard_assignments=shard_assignments,
- hosts=[
- Host(
- ip=host.ip,
- port=random_ephemeral_port(),
- )
- for host in hosts
- ],
+ hosts_by_node=hosts_by_node,
+ ephemeral_port=ephemeral_port,
)
return target_instances
diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py
index f4acfb52..f7bff6f8 100644
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -215,9 +215,11 @@ def get_mlx_ibv_devices_matrix(
continue
# Find the IP J uses to talk to I
- for connection_ip in _find_connection_ip(node_j, node_i, cycle_digraph):
+ for connection_ip, _ in _find_connection_ip(node_j, node_i, cycle_digraph):
# This is a local IP on I, which is attached to an interface: find that interface
- if interface_name := _find_interface_name_for_ip(connection_ip, node_i):
+ if interface_name := _find_rdma_interface_name_for_ip(
+ connection_ip, node_i
+ ):
matrix[i][j] = interface_name
logger.info(
f"Interface name for {connection_ip} on {node_i.node_id}: {interface_name}"
@@ -238,17 +240,17 @@ def _find_connection_ip(
node_i: NodeInfo,
node_j: NodeInfo,
cycle_digraph: Topology,
-) -> Generator[str]:
- """Find all IP addresses that connect node i to node j."""
+) -> Generator[tuple[str, bool]]:
+ """Find all IP addresses that connect node i to node j, with thunderbolt flag."""
for connection in cycle_digraph.list_connections():
if (
connection.local_node_id == node_i.node_id
and connection.send_back_node_id == node_j.node_id
):
- yield connection.send_back_multiaddr.ip_address
+ yield connection.send_back_multiaddr.ip_address, connection.is_thunderbolt()
-def _find_interface_name_for_ip(
+def _find_rdma_interface_name_for_ip(
ip_address: str,
node_info: NodeInfo,
) -> str | None:
@@ -269,6 +271,109 @@ def _find_interface_name_for_ip(
return None
+def _find_interface_name_for_ip(
+ ip_address: str,
+ node_info: NodeInfo,
+) -> str | None:
+ """Find the interface name for an IP address on a node (any interface)."""
+ if node_info.node_profile is None:
+ return None
+
+ for interface in node_info.node_profile.network_interfaces:
+ if interface.ip_address == ip_address:
+ return interface.name
+
+ return None
+
+
+def _find_ip_prioritised(
+ node: NodeInfo, other_node: NodeInfo, cycle_digraph: Topology
+) -> str | None:
+ # TODO: Actually prioritize in the correct Ethernet > Wifi > Non-TB > TB order.
+ """Find an IP address between nodes with prioritization.
+
+ Priority order:
+ 1. en0 (Ethernet on Mac Studio, WiFi on MacBook)
+ 2. en1 (WiFi on Mac Studio, Ethernet on MacBook)
+ 3. Non-Thunderbolt connections
+ 4. Any other IP address
+ """
+ ips = list(_find_connection_ip(node, other_node, cycle_digraph))
+ # We expect a unique iface -> ip mapping
+ iface_map = {_find_interface_name_for_ip(ip, other_node): ip for ip, _ in ips}
+
+ en0_ip = iface_map.get("en0")
+ if en0_ip:
+ return en0_ip
+
+ en1_ip = iface_map.get("en1")
+ if en1_ip:
+ return en1_ip
+
+ non_thunderbolt_ip = next(
+ (ip for (ip, is_thunderbolt) in ips if not is_thunderbolt), None
+ )
+
+ if non_thunderbolt_ip:
+ return non_thunderbolt_ip
+
+ if ips:
+ return ips[0][0]
+
+ return None
+
+
+def get_mlx_ring_hosts_by_node(
+ selected_cycle: list[NodeInfo],
+ cycle_digraph: Topology,
+ ephemeral_port: int,
+) -> dict[NodeId, list[Host]]:
+ """Generate per-node host lists for MLX ring backend.
+
+ Each node gets a list where:
+ - Self position: Host(ip="0.0.0.0", port=ephemeral_port)
+ - Left/right neighbors: actual connection IPs
+ - Non-neighbors: Host(ip="198.51.100.1", port=0) placeholder (RFC 5737 TEST-NET-2)
+ """
+ world_size = len(selected_cycle)
+ if world_size == 0:
+ return {}
+
+ hosts_by_node: dict[NodeId, list[Host]] = {}
+
+ for rank, node in enumerate(selected_cycle):
+ node_id = node.node_id
+ left_rank = (rank - 1) % world_size
+ right_rank = (rank + 1) % world_size
+
+ hosts_for_node: list[Host] = []
+
+ for idx, other_node in enumerate(selected_cycle):
+ if idx == rank:
+ hosts_for_node.append(Host(ip="0.0.0.0", port=ephemeral_port))
+ continue
+
+ if idx not in {left_rank, right_rank}:
+ # Placeholder IP from RFC 5737 TEST-NET-2
+ hosts_for_node.append(Host(ip="198.51.100.1", port=0))
+ continue
+
+ connection_ip = _find_ip_prioritised(node, other_node, cycle_digraph)
+ if connection_ip is None:
+ logger.warning(
+ f"Failed to find prioritised connection IP between {node_id} and {other_node.node_id}"
+ )
+ raise ValueError(
+ "MLX ring backend requires connectivity between neighbouring nodes"
+ )
+
+ hosts_for_node.append(Host(ip=connection_ip, port=ephemeral_port))
+
+ hosts_by_node[node_id] = hosts_for_node
+
+ return hosts_by_node
+
+
def get_mlx_jaccl_coordinators(
selected_cycle: list[NodeInfo],
coordinator_port: int,
@@ -286,7 +391,7 @@ def get_mlx_jaccl_coordinators(
if n.node_id == rank_0_node.node_id:
return "0.0.0.0"
- for ip in _find_connection_ip(n, rank_0_node, cycle_digraph):
+ for ip, _ in _find_connection_ip(n, rank_0_node, cycle_digraph):
return ip
logger.warning(
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index 0380e715..35128165 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -163,32 +163,36 @@ async def test_master():
assert events[2].idx == 2
assert isinstance(events[0].event, NodePerformanceMeasured)
assert isinstance(events[1].event, InstanceCreated)
- runner_id = list(
- events[1].event.instance.shard_assignments.runner_to_shard.keys()
- )[0]
- assert events[1].event.instance == MlxRingInstance(
- instance_id=events[1].event.instance.instance_id,
- shard_assignments=ShardAssignments(
- model_id=ModelId("llama-3.2-1b"),
- runner_to_shard={
- (runner_id): PipelineShardMetadata(
- start_layer=0,
- end_layer=16,
+ created_instance = events[1].event.instance
+ assert isinstance(created_instance, MlxRingInstance)
+ runner_id = list(created_instance.shard_assignments.runner_to_shard.keys())[0]
+ # Validate the shard assignments
+ expected_shard_assignments = ShardAssignments(
+ model_id=ModelId("llama-3.2-1b"),
+ runner_to_shard={
+ (runner_id): PipelineShardMetadata(
+ start_layer=0,
+ end_layer=16,
+ n_layers=16,
+ model_meta=ModelMetadata(
+ model_id=ModelId("llama-3.2-1b"),
+ pretty_name="Llama 3.2 1B",
n_layers=16,
- model_meta=ModelMetadata(
- model_id=ModelId("llama-3.2-1b"),
- pretty_name="Llama 3.2 1B",
- n_layers=16,
- storage_size=Memory.from_bytes(678948),
- ),
- device_rank=0,
- world_size=1,
- )
- },
- node_to_runner={node_id: runner_id},
- ),
- hosts=[],
+ storage_size=Memory.from_bytes(678948),
+ ),
+ device_rank=0,
+ world_size=1,
+ )
+ },
+ node_to_runner={node_id: runner_id},
)
+ assert created_instance.shard_assignments == expected_shard_assignments
+ # For single-node, hosts_by_node should have one entry with self-binding
+ assert len(created_instance.hosts_by_node) == 1
+ assert node_id in created_instance.hosts_by_node
+ assert len(created_instance.hosts_by_node[node_id]) == 1
+ assert created_instance.hosts_by_node[node_id][0].ip == "0.0.0.0"
+ assert created_instance.ephemeral_port > 0
assert isinstance(events[2].event, TaskCreated)
assert events[2].event.task.task_status == TaskStatus.Pending
assert isinstance(events[2].event.task, ChatCompletionTask)
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 9d410275..282ca1e0 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -38,7 +38,8 @@ def instance() -> Instance:
shard_assignments=ShardAssignments(
model_id=ModelId("test-model"), runner_to_shard={}, node_to_runner={}
),
- hosts=[],
+ hosts_by_node={},
+ ephemeral_port=50000,
)
@@ -92,9 +93,13 @@ def test_get_instance_placements_create_instance(
topology.add_node(create_node(available_memory[0], node_id_a))
topology.add_node(create_node(available_memory[1], node_id_b))
topology.add_node(create_node(available_memory[2], node_id_c))
+ # Add bidirectional connections for ring topology
topology.add_connection(create_connection(node_id_a, node_id_b))
+ topology.add_connection(create_connection(node_id_b, node_id_a))
topology.add_connection(create_connection(node_id_b, node_id_c))
+ topology.add_connection(create_connection(node_id_c, node_id_b))
topology.add_connection(create_connection(node_id_c, node_id_a))
+ topology.add_connection(create_connection(node_id_a, node_id_c))
# act
placements = place_instance(cic, topology, {})
@@ -234,17 +239,15 @@ def test_get_transition_events_delete_instance(instance: Instance):
assert events[0].instance_id == instance_id
-def test_placement_prioritizes_leaf_cycle_with_less_memory(
+def test_placement_selects_cycle_with_most_memory(
topology: Topology,
model_meta: ModelMetadata,
create_node: Callable[[int, NodeId | None], NodeInfo],
create_connection: Callable[[NodeId, NodeId], Connection],
):
- # Arrange two 3-node cycles. The A-B-C cycle has a leaf node (only one outgoing
- # neighbor per node). The D-E-F cycle has extra outgoing edges making its nodes
- # non-leaves. Ensure both cycles have sufficient total memory, with the A-B-C
- # cycle having LESS total memory than D-E-F. The algorithm should still choose
- # the cycle that contains a leaf node.
+ # Arrange two 3-node cycles with different total memory.
+ # With bidirectional connections for ring topology, both cycles have non-leaf nodes.
+ # The algorithm should select the cycle with the most available memory.
# Model requires more than any single node but fits within a 3-node cycle
model_meta.storage_size.in_bytes = 1500
@@ -258,11 +261,6 @@ def test_placement_prioritizes_leaf_cycle_with_less_memory(
node_id_e = NodeId()
node_id_f = NodeId()
- # Extra sink nodes to make D/E/F non-leaf via additional outgoing edges
- node_id_x = NodeId()
- node_id_y = NodeId()
- node_id_z = NodeId()
-
# A-B-C cycle total memory = 1600 (< D-E-F total)
topology.add_node(create_node(400, node_id_a))
topology.add_node(create_node(400, node_id_b))
@@ -273,24 +271,20 @@ def test_placement_prioritizes_leaf_cycle_with_less_memory(
topology.add_node(create_node(600, node_id_e))
topology.add_node(create_node(600, node_id_f))
- # Extra nodes with tiny memory so they can't form singleton placements
- topology.add_node(create_node(10, node_id_x))
- topology.add_node(create_node(10, node_id_y))
- topology.add_node(create_node(10, node_id_z))
-
- # Build directed cycles
+ # Build bidirectional cycles for ring topology
topology.add_connection(create_connection(node_id_a, node_id_b))
+ topology.add_connection(create_connection(node_id_b, node_id_a))
topology.add_connection(create_connection(node_id_b, node_id_c))
+ topology.add_connection(create_connection(node_id_c, node_id_b))
topology.add_connection(create_connection(node_id_c, node_id_a))
+ topology.add_connection(create_connection(node_id_a, node_id_c))
topology.add_connection(create_connection(node_id_d, node_id_e))
+ topology.add_connection(create_connection(node_id_e, node_id_d))
topology.add_connection(create_connection(node_id_e, node_id_f))
+ topology.add_connection(create_connection(node_id_f, node_id_e))
topology.add_connection(create_connection(node_id_f, node_id_d))
-
- # Add extra outgoing edges from D/E/F so none of them are leaves
- topology.add_connection(create_connection(node_id_d, node_id_x))
- topology.add_connection(create_connection(node_id_e, node_id_y))
- topology.add_connection(create_connection(node_id_f, node_id_z))
+ topology.add_connection(create_connection(node_id_d, node_id_f))
cic = place_instance_command(
model_meta=model_meta,
@@ -299,18 +293,17 @@ def test_placement_prioritizes_leaf_cycle_with_less_memory(
# Act
placements = place_instance(cic, topology, {})
- # Assert the chosen cycle is A-B-C (contains at least one leaf node), even though
- # D-E-F has more total memory.
+ # Assert: D-E-F cycle should be selected as it has more total memory
assert len(placements) == 1
instance_id = list(placements.keys())[0]
instance = placements[instance_id]
assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
- expected_leaf_cycle_nodes = {node_id_a, node_id_b, node_id_c}
- non_leaf_cycle_nodes = {node_id_d, node_id_e, node_id_f}
+ less_memory_cycle_nodes = {node_id_a, node_id_b, node_id_c}
+ more_memory_cycle_nodes = {node_id_d, node_id_e, node_id_f}
- assert expected_leaf_cycle_nodes.issubset(assigned_nodes)
- assert assigned_nodes.isdisjoint(non_leaf_cycle_nodes)
+ assert more_memory_cycle_nodes.issubset(assigned_nodes)
+ assert assigned_nodes.isdisjoint(less_memory_cycle_nodes)
def test_tensor_rdma_backend_connectivity_matrix(
diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py
index 6bb1a75e..0f6273af 100644
--- a/src/exo/shared/types/worker/instances.py
+++ b/src/exo/shared/types/worker/instances.py
@@ -25,7 +25,8 @@ class BaseInstance(TaggedModel):
class MlxRingInstance(BaseInstance):
- hosts: list[Host]
+ hosts_by_node: dict[NodeId, list[Host]]
+ ephemeral_port: int
class MlxJacclInstance(BaseInstance):
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index b53b3297..1aa9b827 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -112,11 +112,12 @@ def mlx_distributed_init(
try:
# TODO: singleton instances
match bound_instance.instance:
- case MlxRingInstance(hosts=hosts):
+ case MlxRingInstance(hosts_by_node=hosts_by_node, ephemeral_port=_):
coordination_file = (
f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
)
- hosts_json = HostList.from_hosts(hosts).model_dump_json()
+ hosts_for_node = hosts_by_node[bound_instance.bound_node_id]
+ hosts_json = HostList.from_hosts(hosts_for_node).model_dump_json()
with open(coordination_file, "w") as f:
_ = f.write(hosts_json)
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 6fc6da87..a9560853 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -414,9 +414,14 @@ class Worker:
while True:
# TODO: EdgeDeleted
edges = set(self.state.topology.list_connections())
- conns = await check_reachable(self.state.topology)
+ conns = await check_reachable(self.state.topology, self.node_id)
for nid in conns:
for ip in conns[nid]:
+ if "127.0.0.1" in ip or "localhost" in ip:
+ logger.warning(
+ f"Loopback connection should not happen: {ip=} for {nid=}"
+ )
+
edge = Connection(
local_node_id=self.node_id,
send_back_node_id=nid,
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 20275623..e313ca73 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -236,8 +236,8 @@ def _ready_to_warmup(
assert device_rank >= 0
# TODO: Ensure these align with MLX distributeds expectations.
- # Rank != 0
- accepting_ranks_ready = device_rank > 0 and all(
+ # Rank < n-1
+ accepting_ranks_ready = device_rank < world_size - 1 and all(
isinstance(
all_runners.get(global_runner_id, None),
(RunnerLoaded, RunnerWarmingUp),
@@ -245,8 +245,8 @@ def _ready_to_warmup(
for global_runner_id in shard_assignments.runner_to_shard
)
- # Rank = 0
- connecting_rank_ready = device_rank == 0 and 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
diff --git a/src/exo/worker/tests/unittests/conftest.py b/src/exo/worker/tests/unittests/conftest.py
index 83631d38..e736fa2a 100644
--- a/src/exo/worker/tests/unittests/conftest.py
+++ b/src/exo/worker/tests/unittests/conftest.py
@@ -72,7 +72,8 @@ def get_mlx_ring_instance(
shard_assignments=get_shard_assignments(
model_id, node_to_runner, runner_to_shard
),
- hosts=[],
+ hosts_by_node={},
+ ephemeral_port=50000,
)
diff --git a/src/exo/worker/tests/unittests/test_plan/test_warmup.py b/src/exo/worker/tests/unittests/test_plan/test_warmup.py
index f5073c41..9978436e 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_warmup.py
@@ -4,6 +4,7 @@ from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runners import (
RunnerIdle,
RunnerLoaded,
+ RunnerLoading,
RunnerWarmingUp,
)
from exo.worker.tests.constants import (
@@ -21,9 +22,9 @@ from exo.worker.tests.unittests.conftest import (
)
-def test_plan_starts_warmup_for_non_zero_rank_when_all_loaded_or_warming():
+def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
"""
- For non-zero device_rank shards, StartWarmup should be emitted when all
+ For non-final device_rank shards, StartWarmup should be emitted when all
shards in the instance are Loaded/WarmingUp.
"""
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
@@ -36,13 +37,13 @@ def test_plan_starts_warmup_for_non_zero_rank_when_all_loaded_or_warming():
)
bound_instance = BoundInstance(
- instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B
+ instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
)
local_runner = FakeRunnerSupervisor(
bound_instance=bound_instance, status=RunnerLoaded()
)
- runners = {RUNNER_2_ID: local_runner}
+ runners = {RUNNER_1_ID: local_runner}
instances = {INSTANCE_1_ID: instance}
all_runners = {
RUNNER_1_ID: RunnerLoaded(),
@@ -50,10 +51,10 @@ def test_plan_starts_warmup_for_non_zero_rank_when_all_loaded_or_warming():
}
result = plan_mod.plan(
- node_id=NODE_B,
+ node_id=NODE_A,
runners=runners, # type: ignore
download_status={},
- global_download_status={NODE_A: []},
+ global_download_status={NODE_B: []},
instances=instances,
all_runners=all_runners,
tasks={},
@@ -149,6 +150,9 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
"""
Rank-zero shard should not start warmup until all non-zero ranks are
already WarmingUp.
+ For accepting ranks (device_rank != world_size - 1), StartWarmup should be
+ emitted when all shards in the instance are Loaded/WarmingUp.
+ In a 2-node setup, rank 0 is the accepting rank.
"""
shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
@@ -159,6 +163,7 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
)
+ # Rank 0 is the accepting rank
bound_instance = BoundInstance(
instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
)
@@ -177,6 +182,136 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
node_id=NODE_A,
runners=runners, # type: ignore
download_status={},
+ global_download_status={NODE_A: []},
+ instances=instances,
+ all_runners=all_runners,
+ tasks={},
+ )
+
+ assert isinstance(result, StartWarmup)
+ assert result.instance_id == INSTANCE_1_ID
+
+
+def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
+ """
+ For connecting rank (device_rank == world_size - 1), StartWarmup should
+ only be emitted once all the other runners are already warming up.
+ In a 2-node setup, rank 1 is the connecting rank.
+ """
+ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+ instance = get_mlx_ring_instance(
+ instance_id=INSTANCE_1_ID,
+ model_id=MODEL_A_ID,
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+ runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
+ )
+
+ # Rank 1 is the connecting rank
+ bound_instance = BoundInstance(
+ instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B
+ )
+ local_runner = FakeRunnerSupervisor(
+ bound_instance=bound_instance, status=RunnerLoaded()
+ )
+
+ runners = {RUNNER_2_ID: local_runner}
+ instances = {INSTANCE_1_ID: instance}
+ all_runners = {
+ RUNNER_1_ID: RunnerWarmingUp(),
+ RUNNER_2_ID: RunnerLoaded(),
+ }
+
+ result = plan_mod.plan(
+ node_id=NODE_B,
+ runners=runners, # type: ignore
+ download_status={},
+ global_download_status={NODE_B: []},
+ instances=instances,
+ all_runners=all_runners,
+ tasks={},
+ )
+
+ assert isinstance(result, StartWarmup)
+ assert result.instance_id == INSTANCE_1_ID
+
+
+def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warming():
+ """
+ Accepting rank should not start warmup while any shard is not Loaded/WarmingUp.
+ In a 2-node setup, rank 0 is the accepting rank.
+ """
+ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+ instance = get_mlx_ring_instance(
+ instance_id=INSTANCE_1_ID,
+ model_id=MODEL_A_ID,
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+ runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
+ )
+
+ # Rank 0 is the accepting rank
+ bound_instance = BoundInstance(
+ instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+ )
+ local_runner = FakeRunnerSupervisor(
+ bound_instance=bound_instance, status=RunnerLoaded()
+ )
+
+ runners = {RUNNER_1_ID: local_runner}
+ instances = {INSTANCE_1_ID: instance}
+ all_runners = {
+ RUNNER_1_ID: RunnerLoaded(),
+ RUNNER_2_ID: RunnerLoading(),
+ }
+
+ result = plan_mod.plan(
+ node_id=NODE_A,
+ runners=runners, # type: ignore
+ download_status={},
+ global_download_status={NODE_A: [], NODE_B: []},
+ instances=instances,
+ all_runners=all_runners,
+ tasks={},
+ )
+
+ assert result is None
+
+
+def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
+ """
+ Connecting rank (device_rank == world_size - 1) should not start warmup
+ until all other ranks are already WarmingUp.
+ In a 2-node setup, rank 1 is the connecting rank.
+ """
+ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
+ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+ instance = get_mlx_ring_instance(
+ instance_id=INSTANCE_1_ID,
+ model_id=MODEL_A_ID,
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID},
+ runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1},
+ )
+
+ # Rank 1 is the connecting rank
+ bound_instance = BoundInstance(
+ instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B
+ )
+ local_runner = FakeRunnerSupervisor(
+ bound_instance=bound_instance, status=RunnerLoaded()
+ )
+
+ runners = {RUNNER_2_ID: local_runner}
+ instances = {INSTANCE_1_ID: instance}
+ all_runners = {
+ RUNNER_1_ID: RunnerLoaded(),
+ RUNNER_2_ID: RunnerLoaded(),
+ }
+
+ result = plan_mod.plan(
+ node_id=NODE_B,
+ runners=runners, # type: ignore
+ download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
diff --git a/src/exo/worker/utils/net_profile.py b/src/exo/worker/utils/net_profile.py
index b7ef0a1e..b45891f4 100644
--- a/src/exo/worker/utils/net_profile.py
+++ b/src/exo/worker/utils/net_profile.py
@@ -1,33 +1,64 @@
-import socket
+import http.client
from anyio import create_task_group, to_thread
+from loguru import logger
from exo.shared.topology import Topology
from exo.shared.types.common import NodeId
-# TODO: ref. api port
async def check_reachability(
- target_ip: str, target_node_id: NodeId, out: dict[NodeId, set[str]]
+ target_ip: str,
+ expected_node_id: NodeId,
+ self_node_id: NodeId,
+ out: dict[NodeId, set[str]],
) -> None:
- sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
- sock.settimeout(1) # 1 second timeout
- try:
- result = await to_thread.run_sync(sock.connect_ex, (target_ip, 52415))
- except socket.gaierror:
- # seems to throw on ipv6 loopback. oh well
- # logger.warning(f"invalid {target_ip=}")
+ """Check if a node is reachable at the given IP and verify its identity."""
+
+ def _fetch_remote_node_id() -> NodeId | None:
+ connection = http.client.HTTPConnection(target_ip, 52415, timeout=1)
+ try:
+ connection.request("GET", "/node_id")
+ response = connection.getresponse()
+ if response.status != 200:
+ return None
+
+ body = response.read().decode("utf-8").strip()
+
+ # Strip quotes if present (JSON string response)
+ if body.startswith('"') and body.endswith('"') and len(body) >= 2:
+ body = body[1:-1]
+
+ return NodeId(body) or None
+ except OSError:
+ return None
+ finally:
+ connection.close()
+
+ remote_node_id = await to_thread.run_sync(_fetch_remote_node_id)
+ if remote_node_id is None:
+ return
+
+ if remote_node_id == self_node_id:
+ return
+
+ if remote_node_id != expected_node_id:
+ logger.warning(
+ f"Discovered node with unexpected node_id; "
+ f"ip={target_ip}, expected_node_id={expected_node_id}, "
+ f"remote_node_id={remote_node_id}"
+ )
return
- finally:
- sock.close()
- if result == 0:
- if target_node_id not in out:
- out[target_node_id] = set()
- out[target_node_id].add(target_ip)
+ if remote_node_id not in out:
+ out[remote_node_id] = set()
+ out[remote_node_id].add(target_ip)
-async def check_reachable(topology: Topology) -> dict[NodeId, set[str]]:
+async def check_reachable(
+ topology: Topology, self_node_id: NodeId
+) -> dict[NodeId, set[str]]:
+ """Check which nodes are reachable and return their IPs."""
reachable: dict[NodeId, set[str]] = {}
async with create_task_group() as tg:
for node in topology.list_nodes():
@@ -35,7 +66,11 @@ async def check_reachable(topology: Topology) -> dict[NodeId, set[str]]:
continue
for iface in node.node_profile.network_interfaces:
tg.start_soon(
- check_reachability, iface.ip_address, node.node_id, reachable
+ check_reachability,
+ iface.ip_address,
+ node.node_id,
+ self_node_id,
+ reachable,
)
return reachable
diff --git a/uv.lock b/uv.lock
index c7a9b689..c09a0bd2 100644
--- a/uv.lock
+++ b/uv.lock
@@ -375,7 +375,7 @@ requires-dist = [
{ name = "huggingface-hub", specifier = ">=0.33.4" },
{ name = "hypercorn", specifier = ">=0.18.0" },
{ name = "loguru", specifier = ">=0.7.3" },
- { name = "mlx", marker = "sys_platform != 'linux'", specifier = ">=0.30.1" },
+ { name = "mlx", marker = "sys_platform == 'darwin'", specifier = ">=0.30.1" },
{ name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = ">=0.30.1" },
{ name = "mlx-lm", specifier = ">=0.28.3" },
{ name = "networkx", specifier = ">=3.5" },
← b4651340 Fix Kimi K2 Thinking download by adding tiktoken.model to do
·
back to Exo
·
some dashboard updates (#1017) 9d9e24f9 →