← back to Exo
Fix placement single node
98f204d14a0812acf0e62a97c380ca5ecb5fe7c6 · 2025-07-26 20:08:37 +0100 · Alex Cheema
Files touched
M master/placement.pyM master/tests/test_placement.pyM shared/apply/apply.pyM shared/topology.py
Diff
commit 98f204d14a0812acf0e62a97c380ca5ecb5fe7c6
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Sat Jul 26 20:08:37 2025 +0100
Fix placement single node
---
master/placement.py | 12 ++++---
master/tests/test_placement.py | 71 ++++++++++++++++++++++++++++++++++++++++++
shared/apply/apply.py | 5 ++-
shared/topology.py | 2 +-
4 files changed, 84 insertions(+), 6 deletions(-)
diff --git a/master/placement.py b/master/placement.py
index 82730472..7137938f 100644
--- a/master/placement.py
+++ b/master/placement.py
@@ -26,12 +26,16 @@ def get_instance_placements(
if command.model_meta.model_id in available_models:
raise ValueError(f"Instance for {command.model_meta.model_id} already exists")
- candidate_cycles = topology.get_cycles()
- cycles = filter_cycles_by_memory(candidate_cycles, command.model_meta.storage_size_kilobytes)
- if not cycles:
+ all_nodes = list(topology.list_nodes())
+ cycles = topology.get_cycles()
+ nodes_in_cycles = {node.node_id for cycle in cycles for node in cycle}
+ singleton_cycles = [[node] for node in all_nodes if node.node_id not in nodes_in_cycles]
+ candidate_cycles = cycles + singleton_cycles
+ cycles_with_sufficient_memory = filter_cycles_by_memory(candidate_cycles, command.model_meta.storage_size_kilobytes)
+ if not cycles_with_sufficient_memory:
raise ValueError("No cycles found with sufficient memory")
- smallest_cycles = get_smallest_cycles(cycles)
+ smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
selected_cycle = max(smallest_cycles, key=lambda cycle: sum(node.node_profile.memory.ram_available for node in cycle if node.node_profile is not None))
shard_assignments = get_shard_assignments(command.model_meta, selected_cycle)
diff --git a/master/tests/test_placement.py b/master/tests/test_placement.py
index 9bef8116..d51d16b1 100644
--- a/master/tests/test_placement.py
+++ b/master/tests/test_placement.py
@@ -108,6 +108,77 @@ def test_get_instance_placements_create_instance(
assert shards_sorted[0].start_layer == 0
assert shards_sorted[-1].end_layer == total_layers
+def test_get_instance_placements_one_node_exact_fit(
+ create_node: Callable[[int, NodeId | None], Node],
+) -> None:
+ topology = Topology()
+ node_id = NodeId()
+ topology.add_node(create_node(1000, node_id))
+ create_instance_command = CreateInstanceCommand(
+ command_id=CommandId(),
+ model_meta=ModelMetadata(
+ model_id="test-model",
+ storage_size_kilobytes=1000,
+ pretty_name="Test Model",
+ n_layers=10
+ ),
+ instance_id=InstanceId(),
+ )
+ placements = get_instance_placements(create_instance_command, topology, {})
+
+ assert len(placements) == 1
+ instance_id = list(placements.keys())[0]
+ instance = placements[instance_id]
+ assert instance.shard_assignments.model_id == "test-model"
+ assert len(instance.shard_assignments.node_to_runner) == 1
+ assert len(instance.shard_assignments.runner_to_shard) == 1
+ assert len(instance.shard_assignments.runner_to_shard) == 1
+
+def test_get_instance_placements_one_node_fits_with_extra_memory(
+ create_node: Callable[[int, NodeId | None], Node],
+) -> None:
+ topology = Topology()
+ node_id = NodeId()
+ topology.add_node(create_node(1001, node_id))
+ create_instance_command = CreateInstanceCommand(
+ command_id=CommandId(),
+ model_meta=ModelMetadata(
+ model_id="test-model",
+ storage_size_kilobytes=1000,
+ pretty_name="Test Model",
+ n_layers=10
+ ),
+ instance_id=InstanceId(),
+ )
+ placements = get_instance_placements(create_instance_command, topology, {})
+
+ assert len(placements) == 1
+ instance_id = list(placements.keys())[0]
+ instance = placements[instance_id]
+ assert instance.shard_assignments.model_id == "test-model"
+ assert len(instance.shard_assignments.node_to_runner) == 1
+ assert len(instance.shard_assignments.runner_to_shard) == 1
+ assert len(instance.shard_assignments.runner_to_shard) == 1
+
+def test_get_instance_placements_one_node_not_fit(
+ create_node: Callable[[int, NodeId | None], Node],
+) -> None:
+ topology = Topology()
+ node_id = NodeId()
+ topology.add_node(create_node(1000, node_id))
+ create_instance_command = CreateInstanceCommand(
+ command_id=CommandId(),
+ model_meta=ModelMetadata(
+ model_id="test-model",
+ storage_size_kilobytes=1001,
+ pretty_name="Test Model",
+ n_layers=10
+ ),
+ instance_id=InstanceId(),
+ )
+
+ with pytest.raises(ValueError, match="No cycles found with sufficient memory"):
+ get_instance_placements(create_instance_command, topology, {})
def test_get_transition_events_no_change(topology: Topology, instance: Instance):
# arrange
diff --git a/shared/apply/apply.py b/shared/apply/apply.py
index 0cf79e40..b5f49538 100644
--- a/shared/apply/apply.py
+++ b/shared/apply/apply.py
@@ -113,7 +113,10 @@ def apply_runner_deleted(event: RunnerDeleted, state: State) -> State:
@event_apply.register(NodePerformanceMeasured)
def apply_node_performance_measured(event: NodePerformanceMeasured, state: State) -> State:
new_profiles: Mapping[NodeId, NodePerformanceProfile] = {**state.node_profiles, event.node_id: event.node_profile}
- return state.model_copy(update={"node_profiles": new_profiles})
+ state = state.model_copy(update={"node_profiles": new_profiles})
+ topology = copy.copy(state.topology)
+ topology.update_node_profile(event.node_id, event.node_profile)
+ return state.model_copy(update={"topology": topology})
@event_apply.register(WorkerStatusUpdated)
def apply_worker_status_updated(event: WorkerStatusUpdated, state: State) -> State:
diff --git a/shared/topology.py b/shared/topology.py
index 52e2f9cd..7b5cde1d 100644
--- a/shared/topology.py
+++ b/shared/topology.py
@@ -58,7 +58,7 @@ class Topology(TopologyProto):
def add_node(self, node: Node) -> None:
if node.node_id in self._node_id_to_rx_id_map:
- raise ValueError("Node already exists")
+ return
rx_id = self._graph.add_node(node)
self._node_id_to_rx_id_map[node.node_id] = rx_id
self._rx_id_to_node_id_map[rx_id] = node.node_id
← 93330f02 Inference Integration Test
·
back to Exo
·
Discovery integration master b687dec6 →