From e06e70a835ad4dd2b7f77a4cefc045f2fb456cfc Mon Sep 17 00:00:00 2001 From: ciaranbor <81697641+ciaranbor@users.noreply.github.com> Date: Tue, 24 Mar 2026 12:11:56 +0000 Subject: [PATCH] Prefer higher model download % for placement (#1767) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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}