AI 任务调度引擎：从串行等待到 DAG 并行编排
AI 任务调度引擎：从串行等待到 DAG 并行编排
一、串行执行的瓶颈：GPU 空转与 CPU 排队
AI 应用里有个挺常见的问题：任务执行顺序不合理，导致资源浪费。拿典型的 RAG 请求来说，流程通常包括意图识别、查询改写、向量检索、重排序和答案生成。如果按顺序串行执行，总延迟就是这五步耗时之和，大概 3 到 5 秒。但实际上，意图识别和查询改写可以并行，向量检索和关键词检索也能并行，只有重排序必须等检索完成，答案生成必须等重排序完成。
用 DAG（有向无环图）调度把无依赖的步骤并行起来，端到端延迟能从 5 秒压到 2.5 秒，基本接近关键路径的耗时。在高并发场景下，这意味着同样的硬件资源，吞吐量能翻倍。
多 Agent 系统里情况更复杂。比如规划 Agent 拆出 5 个子任务，其中 3 个能并行，2 个得等前 3 个的输出。如果用简单的 FIFO 队列，那 2 个依赖任务会提前占着队列位置，导致调度器空转等待。DAG 调度器能自动识别依赖，只有依赖满足时才把任务投入执行。
二、DAG 调度的核心：拓扑排序、关键路径与动态依赖
DAG 调度的本质，就是把一组有依赖关系的任务组织成图，按拓扑顺序执行——每个节点得等所有前驱节点完成才能开始。
graph TB subgraph RAG 请求处理 DAG A[意图识别<br/>200ms] --> C[查询改写<br/>300ms] B[用户画像加载<br/>150ms] --> C C --> D[向量检索<br/>500ms] C --> E[关键词检索<br/>200ms] D --> F[结果融合与重排序<br/>300ms] E --> F F --> G[答案生成<br/>1000ms] end subgraph 关键路径 A -.-> C -.-> D -.-> F -.-> G end style A fill:#4CAF50,color:#fff style B fill:#4CAF50,color:#fff style C fill:#2196F3,color:#fff style D fill:#2196F3,color:#fff style E fill:#2196F3,color:#fff style F fill:#FF9800,color:#fff style G fill:#F44336,color:#fff2.1 拓扑排序：动态计算可执行集合
拓扑排序是 DAG 调度的第一步。理论上它把节点排成线性序列，让每条边 (u, v) 里的 u 都排在 v 前面。实际调度里，不需要一次性排好所有节点，而是动态计算“当前可执行的节点集合”——也就是所有前驱都已完成、且自身还没执行的节点。
这种动态计算有两个好处：一是支持运行时动态加节点（有些任务执行过程中才发现需要新子任务）；二是支持条件分支（根据前驱输出决定走哪个分支）。
2.2 关键路径：找到真正的瓶颈
关键路径是 DAG 里耗时最长的路径，它决定了整个任务图的最短完成时间。上图里，关键路径是 A→C→D→F→G，总耗时 2300ms。哪怕 E（关键词检索）只要 200ms，也缩短不了总延迟——因为 D（向量检索）要 500ms，F 得等 D 和 E 都完成。
关键路径分析的价值在于：优化非关键路径上的节点对总延迟没帮助。只有缩短关键路径上的节点耗时，才能减少端到端延迟。这直接指导了性能优化的投入方向——优先优化关键路径上的慢节点。
2.3 动态依赖：运行时决定后续任务
静态 DAG 在编译期就确定了所有节点和边。但 AI 应用里很多决策是运行时才做的——意图识别的结果决定走“知识检索”还是“工具调用”，检索结果数量决定要不要额外精排。
动态 DAG 调度器支持节点执行完后，根据输出动态加新节点和边。这比静态 DAG 灵活，但也更难调试——执行图在运行时才完全确定，没法提前验证有没有环或死锁。
三、生产级 DAG 调度引擎实现
下面是一个支持动态依赖、超时控制和资源约束的 DAG 调度引擎实现。
import asyncio import time from dataclasses import dataclass, field from enum import Enum from typing import Any, Callable, Coroutine, Optional class TaskStatus(Enum): PENDING = "pending" RUNNING = "running" SUCCESS = "success" FAILED = "failed" TIMEOUT = "timeout" SKIPPED = "skipped" # 前驱失败时跳过 @dataclass class TaskNode: """DAG 中的任务节点——携带执行逻辑与元数据""" task_id: str executor: Callable[..., Coroutine] dependencies: list[str] = field(default_factory=list) timeout: float = 30.0 # 单任务超时（秒） retry_count: int = 0 max_retries: int = 2 status: TaskStatus = TaskStatus.PENDING result: Any = None error: Optional[str] = None started_at: float = 0.0 finished_at: float = 0.0 # started_at / finished_at 用于计算实际耗时，识别性能瓶颈 @property def duration(self) -> float: if self.finished_at and self.started_at: return self.finished_at - self.started_at return 0.0 class DAGScheduler: """DAG 调度引擎——拓扑执行 + 并行调度 + 动态扩展""" def __init__(self, max_parallel: int = 10): self.nodes: dict[str, TaskNode] = {} self.max_parallel = max_parallel self._semaphore = asyncio.Semaphore(max_parallel) self._completed_events: dict[str, asyncio.Event] = {} # _completed_events 用于跨任务等待：依赖任务完成时设置事件 def add_node(self, node: TaskNode): """添加任务节点""" self.nodes[node.task_id] = node self._completed_events[node.task_id] = asyncio.Event() def add_dependency(self, task_id: str, depends_on: str): """添加依赖关系——task_id 依赖 depends_on""" if task_id in self.nodes and depends_on in self.nodes: self.nodes[task_id].dependencies.append(depends_on) async def _wait_dependencies(self, node: TaskNode) -> bool: """等待所有依赖任务完成——返回 True 表示依赖全部成功""" for dep_id in node.dependencies: event = self._completed_events.get(dep_id) if event: await event.wait() dep_node = self.nodes.get(dep_id) if dep_node and dep_node.status != TaskStatus.SUCCESS: # 依赖任务失败，当前任务跳过 node.status = TaskStatus.SKIPPED node.error = f"依赖任务 {dep_id} 状态为 {dep_node.status.value}" return False return True async def _execute_node(self, node: TaskNode) -> Any: """执行单个任务节点——带超时与重试""" # 等待依赖完成 deps_ok = await self._wait_dependencies(node) if not deps_ok: self._completed_events[node.task_id].set() return None async with self._semaphore: node.status = TaskStatus.RUNNING node.started_at = time.monotonic() for attempt in range(node.max_retries + 1): try: result = await asyncio.wait_for( node.executor(node.result if node.result else {}), timeout=node.timeout, ) node.status = TaskStatus.SUCCESS node.result = result node.finished_at = time.monotonic() self._completed_events[node.task_id].set() return result except asyncio.TimeoutError: node.retry_count = attempt + 1 if attempt < node.max_retries: await asyncio.sleep(0.5 * (attempt + 1)) else: node.status = TaskStatus.TIMEOUT node.error = f"任务超时（{node.timeout}s）" node.finished_at = time.monotonic() except Exception as exc: node.retry_count = attempt + 1 if attempt < node.max_retries: await asyncio.sleep(0.5 * (attempt + 1)) else: node.status = TaskStatus.FAILED node.error = str(exc) node.finished_at = time.monotonic() # 所有重试耗尽，标记完成事件（让依赖此节点的任务能继续判断） self._completed_events[node.task_id].set() return None async def run(self) -> dict[str, Any]: """执行整个 DAG——并行调度无依赖节点""" # 收集所有无入度的节点作为起始执行集合 tasks = [ self._execute_node(node) for node in self.nodes.values() ] await asyncio.gather(*tasks) return { node_id: { "status": node.status.value, "duration_ms": round(node.duration * 1000, 1), "error": node.error, } for node_id, node in self.nodes.items() } def get_critical_path(self) -> list[str]: """计算关键路径——耗时最长的依赖链""" # 动态规划：每个节点的最早完成时间 = max(前驱完成时间) + 自身耗时 earliest_finish: dict[str, float] = {} def _calc(node_id: str) -> float: if node_id in earliest_finish: return earliest_finish[node_id] node = self.nodes[node_id] if not node.dependencies: earliest_finish[node_id] = node.duration else: max_dep = max(_calc(dep) for dep in node.dependencies) earliest_finish[node_id] = max_dep + node.duration return earliest_finish[node_id] for nid in self.nodes: _calc(nid) # 从终点回溯，每步选择耗时最长的前驱 end_node = max(earliest_finish, key=earliest_finish.get) path = [end_node] while True: current = path[-1] deps = self.nodes[current].dependencies if not deps: break # 选择最早完成时间最大的前驱（关键路径上的前驱） next_node = max(deps, key=lambda d: earliest_finish.get(d, 0)) path.append(next_node) path.reverse() return path3.1 RAG 场景的 DAG 编排示例
async def intent_recognition(context: dict) -> dict: """意图识别——判断用户查询类型""" await asyncio.sleep(0.2) # 模拟 LLM 调用 return {"intent": "knowledge_query", "confidence": 0.95} async def query_rewrite(context: dict) -> dict: """查询改写——优化检索效果""" await asyncio.sleep(0.3) return {"rewritten_query": "优化后的查询语句"} async def vector_search(context: dict) -> dict: """向量检索——从知识库召回相关文档""" await asyncio.sleep(0.5) return {"candidates": [{"doc_id": f"doc_{i}", "score": 0.9 - i * 0.05} for i in range(5)]} async def keyword_search(context: dict) -> dict: """关键词检索——补充精确匹配结果""" await asyncio.sleep(0.2) return {"keyword_results": [{"doc_id": "doc_kw_1", "score": 0.85}]} async def rerank(context: dict) -> dict: """重排序——融合多路检索结果""" await asyncio.sleep(0.3) return {"ranked_results": [{"doc_id": "doc_0", "final_score": 0.92}]} async def generate_answer(context: dict) -> dict: """答案生成——基于检索结果生成最终回答""" await asyncio.sleep(1.0) return {"answer": "基于检索结果的生成回答", "sources": ["doc_0"]} # 构建 DAG scheduler = DAGScheduler(max_parallel=5) scheduler.add_node(TaskNode(task_id="intent", executor=intent_recognition, timeout=5.0)) scheduler.add_node(TaskNode(task_id="rewrite", executor=query_rewrite, timeout=5.0)) scheduler.add_node(TaskNode(task_id="vector_search", executor=vector_search, timeout=10.0)) scheduler.add_node(TaskNode(task_id="keyword_search", executor=keyword_search, timeout=5.0)) scheduler.add_node(TaskNode(task_id="rerank", executor=rerank, timeout=8.0)) scheduler.add_node(TaskNode(task_id="generate", executor=generate_answer, timeout=30.0)) # 定义依赖关系 scheduler.add_dependency("rewrite", "intent") # 改写依赖意图识别 scheduler.add_dependency("vector_search", "rewrite") # 检索依赖改写结果 scheduler.add_dependency("keyword_search", "rewrite") # 关键词检索也依赖改写 scheduler.add_dependency("rerank", "vector_search") # 重排序依赖向量检索 scheduler.add_dependency("rerank", "keyword_search") # 重排序依赖关键词检索 scheduler.add_dependency("generate", "rerank") # 生成依赖重排序 # 执行 # results = await scheduler.run() # critical_path = scheduler.get_critical_path()3.2 几个关键设计决策
为什么用asyncio.Event而不是轮询？轮询会浪费 CPU 周期，而且轮询间隔很难权衡——太短浪费资源，太长增加延迟。Event 是零开销的等待机制：依赖任务完成时 set 事件，等待任务立即被唤醒。
为什么失败任务也要设置 Event？如果失败任务不设置 Event，依赖它的任务会永远等下去，导致调度器死锁。设置 Event 后，依赖任务在_wait_dependencies里检测到前驱失败，把自己标记为 SKIPPED，避免无意义的执行。
为什么信号量放在节点执行层而不是调度层？信号量控制的是实际执行并发度，不是调度并发度。如果放在调度层，一个等待依赖的任务也会占用信号量，导致其他无依赖的任务没法执行。放在执行层，只有真正在运行的任务才占信号量。
四、DAG 调度的代价与边界
DAG 调度不是万能的，它引入的复杂度在某些场景下得不偿失。
动态依赖的不可预测性：运行时动态加节点意味着执行图在调度前没法完整验证。可能出现循环依赖（A 依赖 B，B 又动态依赖 A），导致死锁。解决办法是设置全局超时和深度限制，但这只是兜底，不是根治。
资源竞争与优先级：DAG 调度器默认所有就绪节点平等竞争执行资源。但实际场景里，关键路径上的任务应该优先拿 GPU/CPU 资源。实现优先级调度需要额外的资源分配器，系统复杂度就上去了。
状态管理的爆炸：每个节点的状态（PENDING/RUNNING/SUCCESS/FAILED/TIMEOUT/SKIPPED）组合起来，整个 DAG 的状态空间呈指数增长。节点数超过 20 个时，穷举所有可能的执行路径来测试就不现实了。
适用边界：DAG 调度适合任务步骤多（>3 步）、存在可并行的独立子路径、且总延迟敏感的场景。RAG 请求处理、多 Agent 任务编排、数据 ETL 流水线都是典型适用场景。
禁用场景：任务步骤少（2-3 步）且全部串行的简单流程——DAG 调度引入的抽象层反而增加理解和调试成本；任务执行时间差异极大（1ms 与 10min 混合）的场景——短任务被长任务阻塞在信号量上，需要更细粒度的资源隔离；对执行顺序有严格确定性要求的场景——异步并行执行的时序在不同运行间可能不同。
五、总结
DAG 调度引擎通过拓扑排序识别可并行执行的任务集合，通过关键路径分析定位延迟瓶颈，通过动态依赖支持运行时的条件分支。核心设计要素包括：依赖等待机制（Event 驱动而非轮询）、并发控制（信号量限制实际执行并发度）、失败传播（前驱失败时跳过后继而非死锁等待）。
落地路线建议：先把现有串行流程建模为静态 DAG，识别可并行的子路径并验证延迟收益；然后实现基础 DAG 调度器，支持依赖等待、超时控制和重试；接着引入关键路径分析，指导性能优化投入方向；最后对需要运行时决策的场景，逐步开放动态依赖能力，同时设置深度限制和全局超时防止失控。始终记住：DAG 调度的价值在于缩短关键路径耗时，而不是让所有任务并行——并行只是手段，降低端到端延迟才是目的。