← back to Exo
Prefer higher model download % for placement (#1767)
e06e70a835ad4dd2b7f77a4cefc045f2fb456cfc · 2026-03-24 12:11:56 +0000 · ciaranbor
## Motivation
When placing a model instance across the cluster, the master previously
only considered available RAM. This meant it could pick a node that
hasn't downloaded the model yet, even when another node already has it
(or is further along in downloading it).
## Changes
- Added download_status parameter to place_instance() in placement.py
- Added _get_node_download_fraction() to compute 0.0–1.0 download
progress per node/model
- Added _cycle_download_score() to sum download fractions across a
cycle's nodes
- Cycle selection now uses a (download_score, available_ram) tuple key —
download progress is the primary sort, RAM is the tiebreaker
- Passed self.state.downloads into place_instance() from master/main.py
## Why It Works
Python's tuple comparison gives download progress strict priority over
RAM, so a node with the model already downloaded will always be
preferred over one with more free RAM but no download.
## Test Plan
### Automated Testing
3 new tests cover: completed download preferred, higher partial progress
preferred, failed download not preferred over no-download node
Files touched
M src/exo/master/main.pyM src/exo/master/placement.pyM src/exo/master/tests/test_placement.py
Diff
commit e06e70a835ad4dd2b7f77a4cefc045f2fb456cfc
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Tue Mar 24 12:11:56 2026 +0000
Prefer higher model download % for placement (#1767)
## Motivation
When placing a model instance across the cluster, the master previously
only considered available RAM. This meant it could pick a node that
hasn't downloaded the model yet, even when another node already has it
(or is further along in downloading it).
## Changes
- Added download_status parameter to place_instance() in placement.py
- Added _get_node_download_fraction() to compute 0.0–1.0 download
progress per node/model
- Added _cycle_download_score() to sum download fractions across a
cycle's nodes
- Cycle selection now uses a (download_score, available_ram) tuple key —
download progress is the primary sort, RAM is the tiebreaker
- Passed self.state.downloads into place_instance() from master/main.py
## Why It Works
Python's tuple comparison gives download progress strict priority over
RAM, so a node with the model already downloaded will always be
preferred over one with more free RAM but no download.
## Test Plan
### Automated Testing
3 new tests cover: completed download preferred, higher partial progress
preferred, failed download not preferred over no-download node
---
src/exo/master/main.py | 1 +
src/exo/master/placement.py | 61 ++++++++++-
src/exo/master/tests/test_placement.py | 188 ++++++++++++++++++++++++++++++++-
3 files changed, 245 insertions(+), 5 deletions(-)
diff --git a/src/exo/master/main.py b/src/exo/master/main.py
index dc347741..94701c6c 100644
--- a/src/exo/master/main.py
+++ b/src/exo/master/main.py
@@ -294,6 +294,7 @@ class Master:
self.state.instances,
self.state.node_memory,
self.state.node_network,
+ download_status=self.state.downloads,
)
transition_events = get_transition_events(
self.state.instances, placement, self.state.tasks
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index f9e9de73..f9380693 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -32,7 +32,10 @@ from exo.shared.types.memory import Memory
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import (
+ DownloadCompleted,
+ DownloadFailed,
DownloadOngoing,
+ DownloadPending,
DownloadProgress,
)
from exo.shared.types.worker.instances import (
@@ -60,6 +63,45 @@ def add_instance_to_placements(
return {**current_instances, command.instance.instance_id: command.instance}
+def _get_node_download_fraction(
+ node_id: NodeId,
+ model_id: ModelId,
+ download_status: Mapping[NodeId, Sequence[DownloadProgress]],
+) -> float:
+ """Return the download fraction (0.0–1.0) for a model on a given node."""
+ for progress in download_status.get(node_id, []):
+ if progress.shard_metadata.model_card.model_id != model_id:
+ continue
+ match progress:
+ case DownloadCompleted():
+ return 1.0
+ case DownloadOngoing():
+ total = progress.download_progress.total.in_bytes
+ return (
+ progress.download_progress.downloaded.in_bytes / total
+ if total > 0
+ else 0.0
+ )
+ case DownloadPending():
+ total = progress.total.in_bytes
+ return progress.downloaded.in_bytes / total if total > 0 else 0.0
+ case DownloadFailed():
+ return 0.0
+ return 0.0
+
+
+def _cycle_download_score(
+ cycle: Cycle,
+ model_id: ModelId,
+ download_status: Mapping[NodeId, Sequence[DownloadProgress]],
+) -> float:
+ """Sum of download fractions across all nodes in a cycle."""
+ return sum(
+ _get_node_download_fraction(node_id, model_id, download_status)
+ for node_id in cycle
+ )
+
+
def place_instance(
command: PlaceInstance,
topology: Topology,
@@ -67,6 +109,7 @@ def place_instance(
node_memory: Mapping[NodeId, MemoryUsage],
node_network: Mapping[NodeId, NodeNetworkInfo],
required_nodes: set[NodeId] | None = None,
+ download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
) -> dict[InstanceId, Instance]:
cycles = topology.get_cycles()
candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -130,11 +173,21 @@ def place_instance(
if any(topology.node_is_leaf(node_id) for node_id in cycle)
]
+ resolved_download_status = download_status or {}
+ candidate_cycles = (
+ cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles
+ )
+
selected_cycle = max(
- cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
- key=lambda cycle: sum(
- (node_memory[node_id].ram_available for node_id in cycle),
- start=Memory(),
+ candidate_cycles,
+ key=lambda cycle: (
+ _cycle_download_score(
+ cycle, command.model_card.model_id, resolved_download_status
+ ),
+ sum(
+ (node_memory[node_id].ram_available for node_id in cycle),
+ start=Memory(),
+ ),
),
)
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 390a56d2..3a8a36c0 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -25,6 +25,12 @@ from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.topology import Connection, SocketConnection
+from exo.shared.types.worker.downloads import (
+ DownloadCompleted,
+ DownloadFailed,
+ DownloadOngoing,
+ DownloadProgressData,
+)
from exo.shared.types.worker.instances import (
Instance,
InstanceId,
@@ -33,7 +39,7 @@ from exo.shared.types.worker.instances import (
MlxRingInstance,
)
from exo.shared.types.worker.runners import ShardAssignments
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
@pytest.fixture
@@ -576,3 +582,183 @@ def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
assert cancel_events[0].task_status == TaskStatus.Cancelled
assert len(delete_events) == 1
assert delete_events[0].instance_id == instance_id_a
+
+
+def _make_shard_metadata(model_card: ModelCard) -> PipelineShardMetadata:
+ return PipelineShardMetadata(
+ model_card=model_card,
+ device_rank=0,
+ world_size=1,
+ start_layer=0,
+ end_layer=model_card.n_layers,
+ n_layers=model_card.n_layers,
+ )
+
+
+def test_placement_prefers_cycle_with_downloaded_model(
+ model_card: ModelCard,
+) -> None:
+ """When two cycles are otherwise equal, prefer the one with the model already downloaded."""
+ topology = Topology()
+
+ model_card.storage_size = Memory.from_bytes(500)
+
+ node_a = NodeId()
+ node_b = NodeId()
+
+ node_memory = {
+ node_a: create_node_memory(1000),
+ node_b: create_node_memory(1000),
+ }
+ node_network = {
+ node_a: create_node_network(),
+ node_b: create_node_network(),
+ }
+
+ topology.add_node(node_a)
+ topology.add_node(node_b)
+ # No connections between them — two single-node cycles
+
+ shard_meta = _make_shard_metadata(model_card)
+
+ # node_b has the model fully downloaded, node_a does not
+ download_status = {
+ node_b: [
+ DownloadCompleted(
+ node_id=node_b,
+ shard_metadata=shard_meta,
+ total=model_card.storage_size,
+ ),
+ ],
+ }
+
+ cic = place_instance_command(model_card)
+ placements = place_instance(
+ cic, topology, {}, node_memory, node_network, download_status=download_status
+ )
+
+ assert len(placements) == 1
+ instance = list(placements.values())[0]
+ assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+ assert assigned_nodes == {node_b}
+
+
+def test_placement_prefers_cycle_with_higher_download_progress(
+ model_card: ModelCard,
+) -> None:
+ """When two cycles are otherwise equal, prefer the one with more download progress."""
+ topology = Topology()
+
+ model_card.storage_size = Memory.from_bytes(1000)
+
+ node_a = NodeId()
+ node_b = NodeId()
+
+ node_memory = {
+ node_a: create_node_memory(1000),
+ node_b: create_node_memory(1000),
+ }
+ node_network = {
+ node_a: create_node_network(),
+ node_b: create_node_network(),
+ }
+
+ topology.add_node(node_a)
+ topology.add_node(node_b)
+
+ shard_meta = _make_shard_metadata(model_card)
+
+ # node_a: 30% downloaded, node_b: 80% downloaded
+ download_status = {
+ node_a: [
+ DownloadOngoing(
+ node_id=node_a,
+ shard_metadata=shard_meta,
+ download_progress=DownloadProgressData(
+ total=Memory.from_bytes(1000),
+ downloaded=Memory.from_bytes(300),
+ downloaded_this_session=Memory.from_bytes(300),
+ completed_files=0,
+ total_files=1,
+ speed=0.0,
+ eta_ms=0,
+ files={},
+ ),
+ ),
+ ],
+ node_b: [
+ DownloadOngoing(
+ node_id=node_b,
+ shard_metadata=shard_meta,
+ download_progress=DownloadProgressData(
+ total=Memory.from_bytes(1000),
+ downloaded=Memory.from_bytes(800),
+ downloaded_this_session=Memory.from_bytes(800),
+ completed_files=0,
+ total_files=1,
+ speed=0.0,
+ eta_ms=0,
+ files={},
+ ),
+ ),
+ ],
+ }
+
+ cic = place_instance_command(model_card)
+ placements = place_instance(
+ cic, topology, {}, node_memory, node_network, download_status=download_status
+ )
+
+ assert len(placements) == 1
+ instance = list(placements.values())[0]
+ assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+ assert assigned_nodes == {node_b}
+
+
+def test_placement_does_not_prefer_cycle_with_failed_download(
+ model_card: ModelCard,
+) -> None:
+ """A failed download should count as 0% — not preferred over a node with no download history."""
+ topology = Topology()
+
+ model_card.storage_size = Memory.from_bytes(500)
+
+ node_a = NodeId()
+ node_b = NodeId()
+
+ # node_a has slightly more RAM so it would win on the RAM tiebreaker
+ node_memory = {
+ node_a: create_node_memory(1001),
+ node_b: create_node_memory(1000),
+ }
+ node_network = {
+ node_a: create_node_network(),
+ node_b: create_node_network(),
+ }
+
+ topology.add_node(node_a)
+ topology.add_node(node_b)
+
+ shard_meta = _make_shard_metadata(model_card)
+
+ # node_b has a failed download — should not be preferred
+ download_status = {
+ node_b: [
+ DownloadFailed(
+ node_id=node_b,
+ shard_metadata=shard_meta,
+ error_message="connection reset",
+ ),
+ ],
+ }
+
+ cic = place_instance_command(model_card)
+ placements = place_instance(
+ cic, topology, {}, node_memory, node_network, download_status=download_status
+ )
+
+ assert len(placements) == 1
+ instance = list(placements.values())[0]
+ assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+ # node_a should win on RAM tiebreaker since failed download scores 0.0
+ assert assigned_nodes == {node_a}
← e9fdd8d4 improve logging: add dates to verbose stderr, match file log
·
back to Exo
·
Sync custom model cards across nodes (#1768) 49951e1b →