| """ |
| 拓扑 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) |
