File size: 2,916 Bytes
8ab53f1 | 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 | """
拓扑 DAG 工作流执行引擎
Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119)
"""
import asyncio
import logging
from collections import defaultdict, deque
from typing import Dict, Any, Set, Callable, Awaitable
logger = logging.getLogger("VerseFlow.DAGEngine")
class DAGNode:
def __init__(self, node_id: str, handler: Callable[..., Awaitable[Dict[str, Any]]]):
self.node_id = node_id
self.handler = handler
self.in_edges: Set[str] = set()
self.out_edges: Set[str] = set()
class DAGWorkflowEngine:
"""支持并发分裂与动态合并的拓扑 DAG 执行器"""
def __init__(self):
self.nodes: Dict[str, DAGNode] = {}
def add_node(self, node_id: str, handler: Callable[..., Awaitable[Dict[str, Any]]]):
if node_id not in self.nodes:
self.nodes[node_id] = DAGNode(node_id, handler)
def add_dependency(self, parent_id: str, child_id: str):
if parent_id in self.nodes and child_id in self.nodes:
self.nodes[parent_id].out_edges.add(child_id)
self.nodes[child_id].in_edges.add(parent_id)
async def execute(self, initial_inputs: Dict[str, Any]) -> Dict[str, Any]:
in_degree = {nid: len(node.in_edges) for nid, node in self.nodes.items()}
node_outputs: Dict[str, Dict[str, Any]] = defaultdict(dict)
ready_queue = deque([nid for nid, deg in in_degree.items() if deg == 0])
running_tasks: Dict[str, asyncio.Task] = {}
logger.info(f"开始执行 DAG 工作流, 节点总数: {len(self.nodes)}")
while ready_queue or running_tasks:
while ready_queue:
nid = ready_queue.popleft()
parent_data = {
pid: node_outputs[pid] for pid in self.nodes[nid].in_edges
}
merged_inputs = {**initial_inputs, **parent_data}
logger.info(f"启动 DAG 节点: {nid}")
task = asyncio.create_task(self.nodes[nid].handler(merged_inputs))
running_tasks[nid] = task
if not running_tasks:
break
done, _ = await asyncio.wait(
running_tasks.values(), return_when=asyncio.FIRST_COMPLETED
)
finished_nids = []
for nid, task in list(running_tasks.items()):
if task in done:
finished_nids.append(nid)
node_outputs[nid] = await task
logger.info(f"DAG 节点完成: {nid}")
for child_id in self.nodes[nid].out_edges:
in_degree[child_id] -= 1
if in_degree[child_id] == 0:
ready_queue.append(child_id)
for nid in finished_nids:
del running_tasks[nid]
return dict(node_outputs) |