super secret fix

This commit is contained in:
Evan
2025-10-08 18:41:07 +01:00
parent 151c89f5f7
commit 36ff53e1f0
2 changed files with 25 additions and 10 deletions

View File

@@ -62,17 +62,29 @@ def get_instance_placements_after_create(
elif smallest_tb_cycles != []:
smallest_cycles = smallest_tb_cycles
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
selected_cycle = None
for cycle in smallest_cycles:
for node in cycle:
if topology.node_is_leaf(node.node_id):
selected_cycle = cycle
break
else:
continue
break
if not selected_cycle:
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
),
start=Memory(),
),
start=Memory(),
),
)
)
shard_assignments = get_shard_assignments(command.model_meta, selected_cycle)

View File

@@ -49,6 +49,9 @@ class Topology:
self._node_id_to_rx_id_map[node.node_id] = rx_id
self._rx_id_to_node_id_map[rx_id] = node.node_id
def node_is_leaf(self, node_id: NodeId) -> bool:
return node_id in self._node_id_to_rx_id_map and len(self._graph.neighbors(self._node_id_to_rx_id_map[node_id])) == 1
def contains_node(self, node_id: NodeId) -> bool:
return node_id in self._node_id_to_rx_id_map