← back to Exo
Fix tests broken by 283c (#1063)
c1be5184b2787bb11bc16c72789d0afd7b6b5510 · 2025-12-31 01:53:55 +0000 · Evan Quiney
Some tests were broken by #1058 and #1046 - this fixes them.
Files touched
M src/exo/master/placement.pyM src/exo/master/tests/test_placement.pyM src/exo/worker/tests/constants.pyM src/exo/worker/tests/unittests/test_plan/test_warmup.py
Diff
commit c1be5184b2787bb11bc16c72789d0afd7b6b5510
Author: Evan Quiney <evanev7@gmail.com>
Date: Wed Dec 31 01:53:55 2025 +0000
Fix tests broken by 283c (#1063)
Some tests were broken by #1058 and #1046 - this fixes them.
---
src/exo/master/placement.py | 19 +++++---
src/exo/master/tests/test_placement.py | 2 +-
src/exo/worker/tests/constants.py | 2 +
.../tests/unittests/test_plan/test_warmup.py | 54 +++++++++++++++-------
4 files changed, 52 insertions(+), 25 deletions(-)
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index 04a78bd0..666bf282 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -61,28 +61,33 @@ def place_instance(
candidate_cycles = list(
filter(lambda it: len(it) >= command.min_nodes, cycles + singleton_cycles)
)
+ cycles_with_sufficient_memory = filter_cycles_by_memory(
+ candidate_cycles, command.model_meta.storage_size
+ )
+ if not cycles_with_sufficient_memory:
+ raise ValueError("No cycles found with sufficient memory")
+
if command.sharding == Sharding.Tensor:
if not command.model_meta.supports_tensor:
raise ValueError(
f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_meta.model_id}"
)
# TODO: the condition here for tensor parallel is not correct, but it works good enough for now.
- candidate_cycles = [
+ cycles_with_sufficient_memory = [
cycle
- for cycle in candidate_cycles
+ for cycle in cycles_with_sufficient_memory
if command.model_meta.hidden_size % len(cycle) == 0
]
+ if not cycles_with_sufficient_memory:
+ raise ValueError(
+ f"No tensor sharding found for model with hidden_size {command.model_meta.hidden_size} candidate cycles"
+ )
if command.sharding == Sharding.Pipeline and command.model_meta.model_id == ModelId(
"mlx-community/DeepSeek-V3.1-8bit"
):
raise ValueError(
"Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)"
)
- cycles_with_sufficient_memory = filter_cycles_by_memory(
- candidate_cycles, command.model_meta.storage_size
- )
- if not cycles_with_sufficient_memory:
- raise ValueError("No cycles found with sufficient memory")
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 423dcfde..f32ac7b1 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -50,7 +50,7 @@ def model_meta() -> ModelMetadata:
storage_size=Memory.from_kb(1000),
pretty_name="Test Model",
n_layers=10,
- hidden_size=10,
+ hidden_size=30,
supports_tensor=True,
)
diff --git a/src/exo/worker/tests/constants.py b/src/exo/worker/tests/constants.py
index 6d7fabe7..55b26177 100644
--- a/src/exo/worker/tests/constants.py
+++ b/src/exo/worker/tests/constants.py
@@ -9,9 +9,11 @@ MASTER_NODE_ID = NodeId("ffffffff-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
NODE_A: Final[NodeId] = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
NODE_B: Final[NodeId] = NodeId("bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb")
+NODE_C: Final[NodeId] = NodeId("cccccccc-cccc-4ccc-8ccc-cccccccccccc")
RUNNER_1_ID: Final[RunnerId] = RunnerId("11111111-1111-4111-8111-111111111111")
RUNNER_2_ID: Final[RunnerId] = RunnerId("33333333-3333-4333-8333-333333333333")
+RUNNER_3_ID: Final[RunnerId] = RunnerId("Runner3")
INSTANCE_1_ID: Final[InstanceId] = InstanceId("22222222-2222-4222-8222-222222222222")
INSTANCE_2_ID: Final[InstanceId] = InstanceId("44444444-4444-4444-8444-444444444444")
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 9978436e..b42a5afd 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_warmup.py
@@ -12,8 +12,10 @@ from exo.worker.tests.constants import (
MODEL_A_ID,
NODE_A,
NODE_B,
+ NODE_C,
RUNNER_1_ID,
RUNNER_2_ID,
+ RUNNER_3_ID,
)
from exo.worker.tests.unittests.conftest import (
FakeRunnerSupervisor,
@@ -24,37 +26,39 @@ from exo.worker.tests.unittests.conftest import (
def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
"""
- For non-final device_rank shards, StartWarmup should be emitted when all
+ For non-zero 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)
- shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2)
+ shard0 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=3)
+ shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=3)
+ shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=2, world_size=3)
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},
+ node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID, NODE_C: RUNNER_3_ID},
+ runner_to_shard={RUNNER_1_ID: shard0, RUNNER_2_ID: shard1, RUNNER_3_ID: shard2},
)
bound_instance = BoundInstance(
- instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
+ instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B
)
local_runner = FakeRunnerSupervisor(
bound_instance=bound_instance, status=RunnerLoaded()
)
- runners = {RUNNER_1_ID: local_runner}
+ runners = {RUNNER_2_ID: local_runner}
instances = {INSTANCE_1_ID: instance}
all_runners = {
RUNNER_1_ID: RunnerLoaded(),
RUNNER_2_ID: RunnerLoaded(),
+ RUNNER_3_ID: RunnerWarmingUp(),
}
result = plan_mod.plan(
- node_id=NODE_A,
+ node_id=NODE_B,
runners=runners, # type: ignore
download_status={},
- global_download_status={NODE_B: []},
+ global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
tasks={},
@@ -150,9 +154,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
+ For accepting ranks (device_rank != 0), StartWarmup should be
emitted when all shards in the instance are Loaded/WarmingUp.
- In a 2-node setup, rank 0 is the accepting rank.
+ In a 2-node setup, rank 1 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)
@@ -163,7 +167,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
+ # Rank 1 is the accepting rank
bound_instance = BoundInstance(
instance=instance, bound_runner_id=RUNNER_1_ID, bound_node_id=NODE_A
)
@@ -188,6 +192,23 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
tasks={},
)
+ assert result is None
+
+ all_runners = {
+ RUNNER_1_ID: RunnerLoaded(),
+ RUNNER_2_ID: RunnerWarmingUp(),
+ }
+
+ result = plan_mod.plan(
+ 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
@@ -280,9 +301,8 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
"""
- Connecting rank (device_rank == world_size - 1) should not start warmup
+ Connecting rank (device_rank == 0) 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)
@@ -295,13 +315,13 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
# Rank 1 is the connecting rank
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(),
@@ -309,7 +329,7 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
}
result = plan_mod.plan(
- node_id=NODE_B,
+ node_id=NODE_A,
runners=runners, # type: ignore
download_status={},
global_download_status={NODE_A: [], NODE_B: []},
← 1ec550df Emit download progress on start, and change downloads to be
·
back to Exo
·
fix: prevent form submission during IME composition (#1069) 844bcc7c →