def can_partition(dag: Dag, stored_with: set, top_available: set):
""" Returns whether the Dag passed to it can be partitioned. """
# copy so we don't overwrite global available nodes in this pass
available = deepcopy(top_available)
ordered = dag.top_sort()
unavailable = set()
for node in ordered:
if node in unavailable and get_stored_with(node) == stored_with:
for parent in node.parents:
if parent in available and not isinstance(parent, Persist):
return False
if is_correct_mode(node, available, stored_with):
available.add(node)
else:
# mark all descendants as unavailable
descendants = Dag(set([node])).get_all_nodes()
unavailable = unavailable.union(descendants)
return True
def find_new_roots(current_dag: Dag, available: set, stored_with: set):
# need topological ordering
ordered = current_dag.top_sort()
# roots of the next subdag, i.e., where the current subdag will end
new_roots = []
# traverse current condag until all boundary nodes are hit
for node in ordered:
if is_correct_mode(node, available, stored_with):
available.add(node)
elif (not node.parents) or (node.parents & available):
if node not in new_roots:
new_roots.append(node)
# roots of the next subdag
return new_roots
def part_dag(dag: Dag):
""" Exhaustive search partition function. """
sorted_nodes = dag.top_sort()
best = part.get_best_partition(sorted_nodes)
return best