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
2026-03-24 12:11:56 +00:00
parent e9fdd8d4af
commit e06e70a835
3 changed files with 245 additions and 5 deletions
@@ -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
@@ -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(),
+            ),
        ),
    )

@@ -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}