import asyncio, warnings, copy, time class BaseNode: def __init__(self): self.params, self.successors = {}, {} def set_params(self, params): self.params = params def next(self, node, action="default"): if action in self.successors: warnings.warn(f"Overwriting successor for action '{action}'") self.successors[action] = node return node def prep(self, shared): pass def exec(self, prep_res): pass def post(self, shared, prep_res, exec_res): pass def _exec(self, prep_res): return self.exec(prep_res) def _run(self, shared): p = self.prep(shared) e = self._exec(p) return self.post(shared, p, e) def run(self, shared): if self.successors: warnings.warn("Node won't run successors. Use Flow.") return self._run(shared) def __rshift__(self, other): return self.next(other) def __sub__(self, action): if isinstance(action, str): return _ConditionalTransition(self, action) raise TypeError("Action must be a string") class _ConditionalTransition: def __init__(self, src, action): self.src, self.action = src, action def __rshift__(self, tgt): return self.src.next(tgt, self.action) class Node(BaseNode): def __init__(self, max_retries=1, wait=0): super().__init__() self.max_retries, self.wait = max_retries, wait def exec_fallback(self, prep_res, exc): raise exc def _exec(self, prep_res): for self.cur_retry in range(self.max_retries): try: return self.exec(prep_res) except Exception as e: if self.cur_retry == self.max_retries - 1: return self.exec_fallback(prep_res, e) if self.wait > 0: time.sleep(self.wait) class BatchNode(Node): def _exec(self, items): return [super(BatchNode, self)._exec(i) for i in (items or [])] class Flow(BaseNode): def __init__(self, start=None): super().__init__() self.start_node = start def start(self, start): self.start_node = start return start def get_next_node(self, curr, action): nxt = curr.successors.get(action or "default") if not nxt and curr.successors: warnings.warn(f"Flow ends: '{action}' not found in {list(curr.successors)}") return nxt def _orch(self, shared, params=None): curr, p, last_action = ( copy.copy(self.start_node), (params or {**self.params}), None, ) while curr: curr.set_params(p) last_action = curr._run(shared) curr = copy.copy(self.get_next_node(curr, last_action)) return last_action def _run(self, shared): p = self.prep(shared) o = self._orch(shared) return self.post(shared, p, o) def post(self, shared, prep_res, exec_res): return exec_res class BatchFlow(Flow): def _run(self, shared): pr = self.prep(shared) or [] for bp in pr: self._orch(shared, {**self.params, **bp}) return self.post(shared, pr, None) class AsyncNode(Node): async def prep_async(self, shared): pass async def exec_async(self, prep_res): pass async def exec_fallback_async(self, prep_res, exc): raise exc async def post_async(self, shared, prep_res, exec_res): pass async def _exec(self, prep_res): for self.cur_retry in range(self.max_retries): try: return await self.exec_async(prep_res) except Exception as e: if self.cur_retry == self.max_retries - 1: return await self.exec_fallback_async(prep_res, e) if self.wait > 0: await asyncio.sleep(self.wait) async def run_async(self, shared): if self.successors: warnings.warn("Node won't run successors. Use AsyncFlow.") return await self._run_async(shared) async def _run_async(self, shared): p = await self.prep_async(shared) e = await self._exec(p) return await self.post_async(shared, p, e) def _run(self, shared): raise RuntimeError("Use run_async.") class AsyncBatchNode(AsyncNode, BatchNode): async def _exec(self, items): return [await super(AsyncBatchNode, self)._exec(i) for i in items] class AsyncParallelBatchNode(AsyncNode, BatchNode): async def _exec(self, items): return await asyncio.gather( *(super(AsyncParallelBatchNode, self)._exec(i) for i in items) ) class AsyncFlow(Flow, AsyncNode): async def _orch_async(self, shared, params=None): curr, p, last_action = ( copy.copy(self.start_node), (params or {**self.params}), None, ) while curr: curr.set_params(p) last_action = ( await curr._run_async(shared) if isinstance(curr, AsyncNode) else curr._run(shared) ) curr = copy.copy(self.get_next_node(curr, last_action)) return last_action async def _run_async(self, shared): p = await self.prep_async(shared) o = await self._orch_async(shared) return await self.post_async(shared, p, o) async def post_async(self, shared, prep_res, exec_res): return exec_res class AsyncBatchFlow(AsyncFlow, BatchFlow): async def _run_async(self, shared): pr = await self.prep_async(shared) or [] for bp in pr: await self._orch_async(shared, {**self.params, **bp}) return await self.post_async(shared, pr, None) class AsyncParallelBatchFlow(AsyncFlow, BatchFlow): async def _run_async(self, shared): pr = await self.prep_async(shared) or [] await asyncio.gather( *(self._orch_async(shared, {**self.params, **bp}) for bp in pr) ) return await self.post_async(shared, pr, None)