共计 3513 个字符,预计需要花费 9 分钟才能阅读完成。
背景痛点
在构建 AI Agent 工作流时,开发者常遇到三大核心问题:

- 流程割裂 :不同任务模块间通过手工脚本或临时接口调用,缺乏统一编排机制
- 状态追踪困难 :分布式环境下任务执行进度和中间结果难以全局监控
- 错误恢复复杂 :失败任务的重试常导致重复计算或状态不一致
以对话型 Agent 为例,一个完整流程可能包含意图识别、数据库查询、大模型生成、敏感词过滤等步骤。当某个环节失败时,传统方案往往需要人工介入检查日志并重新触发流程。
技术选型
主流工作流引擎对比:
| 方案 | 优势 | 局限性 |
|---|---|---|
| Airflow | 丰富的 Operator 生态,可视化界面 | 调度粒度较粗,Python 动态 DAG 有学习成本 |
| Luigi | 依赖解析简洁,适合数据管道场景 | 监控能力较弱 |
| 自研 DAG 引擎 | 轻量级定制,深度适配 AI 场景需求 | 需要实现基础组件 |
为什么选择 DAG 模型 :
- 直观表达任务间的拓扑关系,如图像预处理必须先于模型推理
- 天然支持并行执行独立分支,如同时进行语音转文本和情感分析
- 可通过拓扑排序检测循环依赖,避免死锁
核心实现
最小化工作流引擎架构
class TaskNode:
"""基础任务节点"""
def __init__(self, task_id: str):
self.task_id = task_id
self.dependencies = [] # 前置任务 ID 列表
self.status = TaskStatus.PENDING # 使用 Enum 定义状态
async def execute(self, context: Dict):
"""异步执行抽象方法"""
raise NotImplementedError
依赖解析与调度
- 使用邻接表存储 DAG 结构:
def build_adjacency_list(tasks: List[TaskNode]) -> Dict[str, List[str]]:
return {t.task_id: t.dependencies for t in tasks}
- 基于 Kahn 算法实现拓扑排序(时间复杂度 O(V+E)):
def topological_sort(adj_list):
in_degree = {u: 0 for u in adj_list}
for u in adj_list:
for v in adj_list[u]:
in_degree[v] += 1
queue = deque([u for u in in_degree if in_degree[u] == 0])
sorted_order = []
while queue:
u = queue.popleft()
sorted_order.append(u)
for v in adj_list.get(u, []):
in_degree[v] -= 1
if in_degree[v] == 0:
queue.append(v)
if len(sorted_order) != len(adj_list):
raise ValueError("存在循环依赖")
return sorted_order
状态持久化实现
使用 SQLAlchemy 定义任务运行记录:
class TaskRun(Base):
__tablename__ = 'task_runs'
id = Column(String(36), primary_key=True)
task_id = Column(String(64), index=True)
status = Column(Enum(TaskStatus))
started_at = Column(DateTime)
finished_at = Column(DateTime)
context_json = Column(JSON) # 存储输入输出上下文
def to_dict(self):
return {c.name: getattr(self, c.name) for c in self.__table__.columns}
生产考量
幂等性设计
关键策略:
- 每个任务分配唯一 run_id
- 执行前检查是否已有成功记录
- 使用 SELECT FOR UPDATE 避免并发冲突
async def execute_with_retry(task: TaskNode, max_retries=3):
for attempt in range(max_retries):
try:
async with db.begin() as tx:
# 悲观锁确保唯一执行
run = await tx.run_sql(
"SELECT * FROM task_runs"
"WHERE task_id=:task_id AND status='SUCCESS'""FOR UPDATE NOWAIT",
task_id=task.task_id)
if run:
return run.context_json
return await task.execute()
except DatabaseError as e:
if attempt == max_retries - 1:
raise
await asyncio.sleep(2 ** attempt)
资源隔离方案
使用 cgroups v2 实现 CPU 限制:
# 创建 AI 任务控制组
sudo cgcreate -g cpu:/ai_agent
# 限制 CPU 用量为 2 核
echo "200000" > /sys/fs/cgroup/ai_agent/cpu.max
监控指标埋点
Prometheus 客户端示例:
from prometheus_client import Gauge
TASK_DURATION = Gauge(
'workflow_task_duration_seconds',
'Task execution time',
['task_type'])
class MetricsTaskWrapper(TaskNode):
async def execute(self, context):
start = time.time()
try:
result = await super().execute(context)
TASK_DURATION.labels(self.task_type).set(time.time() - start)
return result
except Exception as e:
TASK_FAILURES.labels(self.task_type).inc()
raise
避坑指南
循环依赖检测
除拓扑排序外,推荐在 DAG 构建时进行强校验:
def validate_dag(tasks: List[TaskNode]):
visited = set()
def check_cycle(task, path):
if task.task_id in path:
raise ValueError(f"循环依赖: {' -> '.join(path)}")
if task.task_id in visited:
return
visited.add(task.task_id)
for dep in task.dependencies:
check_cycle(dep, path + [task.task_id])
for task in tasks:
check_cycle(task, [])
结构化日志
推荐使用 loguru 的 JSON 格式输出:
from loguru import logger
import json
logger.add(
"workflow.log",
format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {message}",
serialize=True # 自动转 JSON
)
logger.info("Task completed",
extra={
"task_id": "text_processing",
"duration": 2.18,
"metrics": {"input_len": 512, "output_len": 128}
})
冷启动优化
- 预热关键模型:在系统启动时加载高频使用的 AI 模型
- 连接池预建:数据库 /API 客户端提前初始化
- 渐进式调度:优先执行关键路径任务
延伸思考
在 Serverless 架构下,工作流引擎可进一步优化:
- 动态伸缩 :根据队列深度自动调整 Lambda 并发度
- 按需计费 :短时任务使用 Spot 实例降低成本
- 事件驱动 :通过 SQS/SNS 触发下游任务
示例架构:
flowchart LR
APIGateway --> LambdaTrigger
LambdaTrigger --> StepFunctions
StepFunctions --> SQS
SQS --> WorkerLambda
总结
构建生产级 AI Agent 工作流需要兼顾灵活性和可靠性。通过本文介绍的方法,开发者可以:
- 使用 DAG 清晰定义任务拓扑
- 基于异步 IO 实现高效调度
- 通过持久化和监控保障稳定性
- 针对 AI 场景优化资源利用率
实际部署时,建议先从简单流程开始验证,逐步扩展到复杂场景。完整的示例代码已开源在 GitHub 仓库中。
正文完
