def combine_two_partitions(partition_0: Partition, partition_1: Partition,
partitions: List[Partition]) -> None:
"""Given a list of partitions and its two partitions,
combine these two partitions into a new one appending to the partitions
and remove the previous two partitions from the list of partitions
"""
partition = Partition(len(partitions))
partition.nodes = partition_0.nodes.union(partition_1.nodes)
partition.recalculate_mem_size()
partitions.append(partition)
partitions.remove(partition_0)
partitions.remove(partition_1)
reorganize_partitions(partitions)
return
def aot_based_partition(self, node_to_partition_mapping, partition_to_logical_device_mapping):
"""This function helps to rebuild the partitions given the nodes and its
corresponding partition id
"""
partition_id_to_partition_mapping: Dict[int, Partition] = {}
self.node_to_partition = node_to_partition_mapping
for node in self.node_to_partition:
partition_id = self.node_to_partition[node]
# If the requested partition has not been created, create the partition
if partition_id not in partition_id_to_partition_mapping:
partition = Partition(partition_id)
self.partitions.append(partition)
partition_id_to_partition_mapping[partition_id] = partition
partition.logical_device_ids = partition_to_logical_device_mapping[partition_id]
else:
partition = partition_id_to_partition_mapping[self.node_to_partition[node]]
# Add the current node into the partition
partition.add_node(node)
def create_partition(self) -> Partition:
"""Create a partition and append it to self.partitions."""
partition_id = len(self.partitions)
partition = Partition(partition_id)
self.partitions.append(partition)
return partition