Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,055 Bytes
df6c67d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import abc
from collections import defaultdict
from queue import Queue
from typing import List, Optional, Set
import networkx as nx
from inference.enterprise.workflows.complier.utils import get_nodes_of_specific_kind
from inference.enterprise.workflows.constants import STEP_NODE_KIND
class StepExecutionCoordinator(metaclass=abc.ABCMeta):
@classmethod
@abc.abstractmethod
def init(cls, execution_graph: nx.DiGraph) -> "StepExecutionCoordinator":
pass
@abc.abstractmethod
def get_steps_to_execute_next(
self, steps_to_discard: Set[str]
) -> Optional[List[str]]:
pass
class SerialExecutionCoordinator(StepExecutionCoordinator):
@classmethod
def init(cls, execution_graph: nx.DiGraph) -> "StepExecutionCoordinator":
return cls(execution_graph=execution_graph)
def __init__(self, execution_graph: nx.DiGraph):
self._execution_graph = execution_graph.copy()
self._discarded_steps: Set[str] = set()
self.__order: Optional[List[str]] = None
self.__step_pointer = 0
def get_steps_to_execute_next(
self, steps_to_discard: Set[str]
) -> Optional[List[str]]:
if self.__order is None:
self.__establish_execution_order()
self._discarded_steps.update(steps_to_discard)
next_step = None
while self.__step_pointer < len(self.__order):
candidate_step = self.__order[self.__step_pointer]
self.__step_pointer += 1
if candidate_step in self._discarded_steps:
continue
return [candidate_step]
return next_step
def __establish_execution_order(self) -> None:
step_nodes = get_nodes_of_specific_kind(
execution_graph=self._execution_graph, kind=STEP_NODE_KIND
)
self.__order = [
n for n in nx.topological_sort(self._execution_graph) if n in step_nodes
]
self.__step_pointer = 0
class ParallelStepExecutionCoordinator(StepExecutionCoordinator):
@classmethod
def init(cls, execution_graph: nx.DiGraph) -> "StepExecutionCoordinator":
return cls(execution_graph=execution_graph)
def __init__(self, execution_graph: nx.DiGraph):
self._execution_graph = execution_graph.copy()
self._discarded_steps: Set[str] = set()
self.__execution_order: Optional[List[List[str]]] = None
self.__execution_pointer = 0
def get_steps_to_execute_next(
self, steps_to_discard: Set[str]
) -> Optional[List[str]]:
if self.__execution_order is None:
self.__execution_order = establish_execution_order(
execution_graph=self._execution_graph
)
self.__execution_pointer = 0
self._discarded_steps.update(steps_to_discard)
next_step = None
while self.__execution_pointer < len(self.__execution_order):
candidate_steps = [
e
for e in self.__execution_order[self.__execution_pointer]
if e not in self._discarded_steps
]
self.__execution_pointer += 1
if len(candidate_steps) == 0:
continue
return candidate_steps
return next_step
def establish_execution_order(
execution_graph: nx.DiGraph,
) -> List[List[str]]:
steps_flow_graph = construct_steps_flow_graph(execution_graph=execution_graph)
steps_flow_graph = assign_max_distances_from_start(
steps_flow_graph=steps_flow_graph
)
return get_groups_execution_order(steps_flow_graph=steps_flow_graph)
def construct_steps_flow_graph(execution_graph: nx.DiGraph) -> nx.DiGraph:
steps_flow_graph = nx.DiGraph()
steps_flow_graph.add_node("start")
steps_flow_graph.add_node("end")
step_nodes = get_nodes_of_specific_kind(
execution_graph=execution_graph, kind=STEP_NODE_KIND
)
for step_node in step_nodes:
for predecessor in execution_graph.predecessors(step_node):
start_node = predecessor if predecessor in step_nodes else "start"
steps_flow_graph.add_edge(start_node, step_node)
for successor in execution_graph.successors(step_node):
end_node = successor if successor in step_nodes else "end"
steps_flow_graph.add_edge(step_node, end_node)
return steps_flow_graph
def assign_max_distances_from_start(steps_flow_graph: nx.DiGraph) -> nx.DiGraph:
nodes_to_consider = Queue()
nodes_to_consider.put("start")
while nodes_to_consider.qsize() > 0:
node_to_consider = nodes_to_consider.get()
predecessors = list(steps_flow_graph.predecessors(node_to_consider))
if not all(
steps_flow_graph.nodes[p].get("distance") is not None for p in predecessors
):
# we can proceed to establish distance, only if all parents have distances established
continue
if len(predecessors) == 0:
distance_from_start = 0
else:
distance_from_start = (
max(steps_flow_graph.nodes[p]["distance"] for p in predecessors) + 1
)
steps_flow_graph.nodes[node_to_consider]["distance"] = distance_from_start
for neighbour in steps_flow_graph.successors(node_to_consider):
nodes_to_consider.put(neighbour)
return steps_flow_graph
def get_groups_execution_order(steps_flow_graph: nx.DiGraph) -> List[List[str]]:
distance2steps = defaultdict(list)
for node_name, node_data in steps_flow_graph.nodes(data=True):
if node_name in {"start", "end"}:
continue
distance2steps[node_data["distance"]].append(node_name)
sorted_distances = sorted(list(distance2steps.keys()))
return [distance2steps[d] for d in sorted_distances]
def get_next_steps_to_execute(
execution_order: List[List[str]],
execution_pointer: int,
discarded_steps: Set[str],
) -> List[str]:
return [e for e in execution_order[execution_pointer] if e not in discarded_steps]
|