LangGraph设计与实现-第6章-Pregel 执行引擎

《LangGraph 设计与实现》完整目录

第6章 Pregel 执行引擎

6.1 引言

Pregel 是 LangGraph 的心脏。当你调用 compiled_graph.invoke()compiled_graph.stream() 时,真正驱动计算的是 Pregel 执行引擎------一个基于 Google Pregel 论文思想、专门为 AI 工作流定制的 BSP(Bulk Synchronous Parallel)运行时。

Google 在 2010 年发表的 Pregel 论文描述了一种用于大规模图计算的编程模型:计算以超步(superstep)为单位推进,每个超步中所有活跃节点并行执行,执行结果在超步之间通过消息传递可见。LangGraph 将这个模型适配到了 AI 工作流场景------"节点"是 AI Agent 或工具调用,"消息"是 Channel 中的状态更新,"超步"是一轮任务执行与状态同步。

本章将深入剖析以下核心组件:

  • Pregel 类(pregel/main.py)------ 执行引擎的公共接口
  • PregelLooppregel/_loop.py)------ 执行循环的核心状态机
  • SyncPregelLoop / AsyncPregelLoop ------ 同步和异步的具体实现
  • prepare_next_tasks 算法 ------ 决定每个超步执行哪些任务
  • apply_writes 算法 ------ 在超步之间更新 Channel 状态
  • 版本追踪机制 ------ 高效判定哪些节点需要被触发
  • max_steps 停止条件 ------ 防止无限循环的安全阀

:::tip 本章要点

  1. Pregel 类是 LangGraph 的统一运行时接口,invoke()stream() 都构建在同一个执行循环上
  2. PregelLoop 是一个状态机,核心循环为 tick() -> 执行任务 -> after_tick()
  3. prepare_next_tasks 通过 Channel 版本比较决定哪些节点在下一步执行
  4. apply_writes 在超步之间原子地更新所有 Channel,确保步内隔离
  5. 版本追踪使用 channel_versionsversions_seen 两张表实现高效的变更检测
  6. recursion_limit 通过 step > stop 条件提供安全停止保证 :::

6.2 Pregel 类:执行引擎的入口

Pregel 类定义在 pregel/main.py 中,是 CompiledStateGraph 的父类。它持有执行所需的全部配置:

python 复制代码
class Pregel(PregelProtocol, Generic[StateT, ContextT, InputT, OutputT]):
    nodes: dict[str, PregelNode]         # 编译后的节点
    channels: dict[str, BaseChannel | ManagedValueSpec]  # 通道定义
    input_channels: str | Sequence[str]  # 输入通道
    output_channels: str | Sequence[str] # 输出通道
    stream_channels: str | Sequence[str] | None
    trigger_to_nodes: Mapping[str, Sequence[str]]  # 优化映射
    checkpointer: Checkpointer           # 检查点存储
    store: BaseStore | None              # 持久化存储
    cache: BaseCache | None              # 节点结果缓存
    retry_policy: Sequence[RetryPolicy]  # 全局重试策略
    cache_policy: CachePolicy | None     # 全局缓存策略
    interrupt_before_nodes: All | Sequence[str]
    interrupt_after_nodes: All | Sequence[str]
    step_timeout: float | None           # 步超时

6.2.1 invoke() 和 stream() 的关系

在 Pregel 的设计中,invoke() 是基于 stream() 实现的------它调用 stream() 收集所有输出,然后返回最终值:

python 复制代码
def invoke(self, input, config=None, *, stream_mode="values", ...):
    latest = None
    for chunk in self.stream(
        input, config,
        stream_mode=["updates", "values"] if stream_mode == "values"
        else stream_mode,
        ...
    ):
        if stream_mode == "values":
            mode, payload = chunk
            if mode == "values":
                latest = payload
    return latest

这意味着 stream() 才是真正的执行入口。所有的执行逻辑都围绕流式输出构建。

6.2.2 stream() 的执行框架

stream() 方法的核心结构如下:

python 复制代码
def stream(self, input, config=None, *, stream_mode=None, ...):
    # 1. 准备配置
    stream_modes, output_keys, interrupt_before_, interrupt_after_, \
        checkpointer, store, cache, durability_ = self._defaults(...)

    # 2. 构建流式队列
    stream = SyncQueue()

    # 3. 进入 PregelLoop 上下文
    with SyncPregelLoop(
        input, stream=StreamProtocol(stream.put, stream_modes),
        config=config, checkpointer=checkpointer,
        nodes=self.nodes, specs=self.channels,
        ...
    ) as loop:
        # 4. 创建 Runner
        runner = PregelRunner(
            submit=weakref.WeakMethod(loop.submit),
            put_writes=weakref.WeakMethod(loop.put_writes),
        )
        # 5. BSP 主循环
        while loop.tick():
            for _ in runner.tick(
                [t for t in loop.tasks.values() if not t.writes],
                timeout=self.step_timeout,
                schedule_task=loop.accept_push,
            ):
                yield from _output(stream.get, ...)
            loop.after_tick()

这段代码精确地体现了 BSP 模型的三个阶段:

flowchart TB subgraph "BSP 超步循环" TICK["loop.tick()\n计划阶段"] --> EXEC["runner.tick()\n执行阶段"] EXEC --> AFTER["loop.after_tick()\n更新阶段"] AFTER -->|有新任务| TICK AFTER -->|无新任务| DONE[结束] end subgraph "计划阶段详情" T1[检查步数限制] --> T2[prepare_next_tasks] T2 --> T3[匹配已有写入] T3 --> T4[检查 interrupt_before] T4 --> T5[发出调试事件] end subgraph "更新阶段详情" A1[apply_writes 更新 Channel] --> A2[发出 values 事件] A2 --> A3[清空 pending_writes] A3 --> A4[保存 Checkpoint] A4 --> A5[检查 interrupt_after] A5 --> A6[步数 + 1] end TICK -.-> T1 AFTER -.-> A1

6.3 PregelLoop:执行循环的状态机

PregelLoop 是执行循环的核心,定义在 pregel/_loop.py 中。它不是一个简单的 while 循环,而是一个精心设计的状态机。

6.3.1 状态定义

python 复制代码
class PregelLoop:
    # 配置
    config: RunnableConfig
    nodes: Mapping[str, PregelNode]
    specs: Mapping[str, BaseChannel | ManagedValueSpec]
    input_keys: str | Sequence[str]
    output_keys: str | Sequence[str]
    stream_keys: str | Sequence[str]

    # 运行时状态
    step: int                    # 当前超步编号
    stop: int                    # 最大步数
    status: str                  # 状态机状态
    tasks: dict[str, PregelExecutableTask]  # 当前步的任务
    output: Any | None           # 最终输出
    updated_channels: set[str] | None  # 上一步更新的通道

    # Checkpoint 状态
    checkpoint: Checkpoint
    checkpoint_config: RunnableConfig
    checkpoint_metadata: CheckpointMetadata
    checkpoint_pending_writes: list[PendingWrite]
    checkpoint_previous_versions: dict[str, str | float | int]

    # Channel 和 Managed Values
    channels: Mapping[str, BaseChannel]
    managed: ManagedValueMapping

6.3.2 状态机转换

PregelLoop 的 status 字段可以是以下值之一:

stateDiagram-v2 [*] --> input : __enter__ input --> pending : _first() 成功 pending --> pending : tick() + after_tick() pending --> done : tick() 返回 False(无任务) pending --> out_of_steps : tick() 返回 False(超步数) pending --> interrupt_before : tick() 中触发中断 pending --> interrupt_after : after_tick() 中触发中断 done --> [*] : __exit__ out_of_steps --> [*] : 抛出 GraphRecursionError interrupt_before --> [*] : 抛出 GraphInterrupt interrupt_after --> [*] : 抛出 GraphInterrupt

6.3.3 SyncPregelLoop 的生命周期

SyncPregelLoop 实现为 Python 上下文管理器,其 __enter__ 方法执行完整的初始化:

python 复制代码
def __enter__(self) -> Self:
    # 1. 获取 Checkpoint
    if not self.checkpointer:
        saved = None
    elif self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID):
        saved = self.checkpointer.get_tuple(self.checkpoint_config)
    else:
        saved = self.checkpointer.get_tuple(self.checkpoint_config)

    if saved is None:
        saved = CheckpointTuple(
            self.checkpoint_config, empty_checkpoint(),
            {"step": -2}, None, []
        )
    elif self._migrate_checkpoint is not None:
        self._migrate_checkpoint(saved.checkpoint)

    # 2. 恢复 Checkpoint 状态
    self.checkpoint = saved.checkpoint
    self.checkpoint_metadata = saved.metadata
    self.checkpoint_pending_writes = [...]

    # 3. 初始化后台执行器
    self.submit = self.stack.enter_context(
        BackgroundExecutor(self.config)
    )

    # 4. 从 Checkpoint 恢复 Channel 状态
    self.channels, self.managed = channels_from_checkpoint(
        self.specs, self.checkpoint
    )

    # 5. 计算步数边界
    self.step = self.checkpoint_metadata["step"] + 1
    self.stop = self.step + self.config["recursion_limit"] + 1

    # 6. 处理首步输入
    self.updated_channels = self._first(
        input_keys=self.input_keys,
        updated_channels=...
    )
    return self

这里的关键点:

  • Checkpoint 恢复:如果存在之前的 Checkpoint(例如从中断点恢复),直接加载而非从空状态开始
  • step 的计算:从 Checkpoint 的元数据中恢复步数,确保恢复执行时步数连续
  • stop 的计算step + recursion_limit + 1+1 是因为比较条件是 step > stop(严格大于),所以需要多一步的余量
  • _first():处理首步输入------将用户输入写入 Channel,或者在恢复执行时跳过输入处理

6.4 _first():首步输入处理

_first() 方法是执行循环中最复杂的初始化逻辑,它需要区分三种场景:

flowchart TD START[_first 开始] --> CHECK{是否从 Checkpoint 恢复?} CHECK -->|是: is_resuming=True| RESUME[恢复模式] CHECK -->|否| INPUT{输入是 Command?} INPUT -->|是| CMD[处理 Command] INPUT -->|否| FRESH[新鲜输入模式] RESUME --> VSEEN["更新 versions_seen\n标记所有 Channel 为已见"] VSEEN --> EMIT_V[发出 values 事件] CMD --> MAP_CMD[map_command 解析] MAP_CMD --> WRITE_CMD[写入 Channel] FRESH --> MAP_INPUT[map_input 解析] MAP_INPUT --> DISCARD[丢弃未完成任务] DISCARD --> APPLY[apply_writes 应用输入] APPLY --> SAVE[保存输入 Checkpoint]

恢复模式

当从 Checkpoint 恢复时(is_resuming=True),_firstversions_seen[INTERRUPT] 设置为当前所有 Channel 的版本号。这告诉 should_interrupt() 函数:"我已经看到了所有当前的更新",防止恢复后立即再次触发中断。

python 复制代码
if is_resuming:
    self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
    for k in self.channels:
        if k in self.checkpoint["channel_versions"]:
            version = self.checkpoint["channel_versions"][k]
            self.checkpoint["versions_seen"][INTERRUPT][k] = version

新鲜输入模式

新输入通过 map_input 转化为 Channel 写入,然后通过 apply_writes 应用:

python 复制代码
elif input_writes := deque(map_input(input_keys, self.input)):
    # 丢弃任何未完成的任务
    discard_tasks = prepare_next_tasks(...)
    # 应用输入写入
    updated_channels = apply_writes(
        self.checkpoint, self.channels,
        [*discard_tasks.values(),
         PregelTaskWrites((), INPUT, input_writes, [])],
        self.checkpointer_get_next_version,
        self.trigger_to_nodes,
    )
    # 保存输入 Checkpoint
    self._put_checkpoint({"source": "input"})

注意 discard_tasks 的处理------如果 Checkpoint 中有之前未完成的任务,它们会在这里被"消费",确保不会在新执行中被重复触发。

6.5 tick():计划阶段

tick() 方法是 BSP 模型的计划阶段,它决定当前超步要执行哪些任务:

python 复制代码
def tick(self) -> bool:
    # 1. 检查步数限制
    if self.step > self.stop:
        self.status = "out_of_steps"
        return False

    # 2. 准备下一步任务
    self.tasks = prepare_next_tasks(
        self.checkpoint,
        self.checkpoint_pending_writes,
        self.nodes, self.channels, self.managed,
        self.config, self.step, self.stop,
        for_execution=True,
        manager=self.manager,
        store=self.store,
        checkpointer=self.checkpointer,
        trigger_to_nodes=self.trigger_to_nodes,
        updated_channels=self.updated_channels,
        retry_policy=self.retry_policy,
        cache_policy=self.cache_policy,
    )

    # 3. 如果没有任务,图执行完毕
    if not self.tasks:
        self.status = "done"
        return False

    # 4. 匹配已有写入(从 Checkpoint 恢复时)
    if not self.is_replaying and self.checkpoint_pending_writes:
        self._match_writes(self.tasks)

    # 5. 检查是否需要在执行前中断
    if self.interrupt_before and should_interrupt(
        self.checkpoint, self.interrupt_before, self.tasks.values()
    ):
        self.status = "interrupt_before"
        raise GraphInterrupt()

    return True

tick() 返回 True 表示"有任务需要执行",False 表示"执行结束"。调用者(stream() 方法)根据返回值决定是否继续循环。

6.6 prepare_next_tasks:任务准备算法

prepare_next_tasks 是整个执行引擎中最核心的调度算法,定义在 pregel/_algo.py 中。它的职责是确定下一个超步应该执行哪些任务。

6.6.1 算法概览

python 复制代码
def prepare_next_tasks(
    checkpoint, pending_writes, processes, channels, managed,
    config, step, stop, *, for_execution, store, checkpointer,
    manager, trigger_to_nodes, updated_channels,
    retry_policy, cache_policy,
):
    input_cache = {}
    tasks = []

    # 阶段一:处理 PUSH 任务(Send API 产生的动态任务)
    tasks_channel = channels.get(TASKS)
    if tasks_channel and tasks_channel.is_available():
        for idx, _ in enumerate(tasks_channel.get()):
            if task := prepare_single_task((PUSH, idx), ...):
                tasks.append(task)

    # 阶段二:确定候选节点(优化路径)
    if updated_channels and trigger_to_nodes:
        triggered_nodes = set()
        for channel in updated_channels:
            if node_ids := trigger_to_nodes.get(channel):
                triggered_nodes.update(node_ids)
        candidate_nodes = sorted(triggered_nodes)
    elif not checkpoint["channel_versions"]:
        candidate_nodes = ()
    else:
        candidate_nodes = processes.keys()

    # 阶段三:处理 PULL 任务(常规节点触发)
    for name in candidate_nodes:
        if task := prepare_single_task((PULL, name), ...):
            tasks.append(task)

    return {t.id: t for t in tasks}

6.6.2 PUSH 任务 vs PULL 任务

这两种任务类型对应了两种不同的触发机制:

flowchart LR subgraph "PULL 任务(常规触发)" CH_UPDATE["Channel 版本更新"] --> VER_CHECK["版本比较\n_triggers()"] VER_CHECK --> PULL_TASK["创建 PULL 任务\n节点名称 -> 任务"] end subgraph "PUSH 任务(Send API)" SEND["Send('node', data)"] --> TASKS_CH["__pregel_tasks\nTopic Channel"] TASKS_CH --> PUSH_TASK["创建 PUSH 任务\n带自定义输入"] end style PULL_TASK fill:#c8e6c9 style PUSH_TASK fill:#fff3e0
  • PULL 任务:由 Channel 版本变更触发。当节点订阅的 Channel 在上一步被更新时,该节点在下一步被"拉入"执行。这是 BSP 模型的标准触发方式。
  • PUSH 任务 :由 Send API 显式创建。节点可以通过返回 Send("target", data) 来动态创建任务,"推送"自定义数据到目标节点。PUSH 任务不经过 Channel 版本检查。

6.6.3 优化路径:trigger_to_nodes

阶段二的候选节点确定包含了一个重要的优化。当同时满足以下条件时,引擎使用 trigger_to_nodes 映射表快速定位需要检查的节点:

  1. updated_channels 不为空------知道上一步更新了哪些 Channel
  2. trigger_to_nodes 不为空------有预构建的映射表
python 复制代码
if updated_channels and trigger_to_nodes:
    triggered_nodes = set()
    for channel in updated_channels:
        if node_ids := trigger_to_nodes.get(channel):
            triggered_nodes.update(node_ids)
    candidate_nodes = sorted(triggered_nodes)

当这两个条件不满足时(例如首次执行时 updated_channels 可能为 None),退化为遍历所有节点。sorted() 确保了确定性的执行顺序。

6.6.4 版本比较:_triggers() 函数

对于每个候选 PULL 节点,prepare_single_task 调用 _triggers() 检查该节点是否真的需要被触发:

python 复制代码
def _triggers(channels, channel_versions, versions_seen,
              null_version, proc):
    """检查节点的任何触发 Channel 是否在上一步被更新过。"""
    if versions_seen is None:
        # 节点从未被执行过
        return any(
            channel_versions.get(chan, null_version) > null_version
            for chan in proc.triggers
            if chan in channels and channels[chan].is_available()
        )
    else:
        return any(
            channel_versions.get(chan, null_version)
            > versions_seen.get(chan, null_version)
            for chan in proc.triggers
            if chan in channels and channels[chan].is_available()
        )

这个函数是 Pregel 调度机制的核心。它使用两张版本表来做比较:

  • channel_versions[chan] ------ Channel 的当前版本号
  • versions_seen[node][chan] ------ 该节点上次执行时看到的 Channel 版本号

如果任何一个触发 Channel 的当前版本大于节点上次看到的版本,说明有新数据,节点需要被触发。

6.7 apply_writes:更新阶段算法

apply_writes 在每个超步结束后被调用(在 after_tick() 中),负责将所有任务的写入原子地应用到 Channel:

python 复制代码
def apply_writes(checkpoint, channels, tasks,
                 get_next_version, trigger_to_nodes):
    # 排序任务,确保确定性
    tasks = sorted(tasks, key=lambda t: task_path_str(t.path[:3]))
    bump_step = any(t.triggers for t in tasks)

    # 1. 更新 versions_seen
    for task in tasks:
        checkpoint["versions_seen"].setdefault(task.name, {}).update({
            chan: checkpoint["channel_versions"][chan]
            for chan in task.triggers
            if chan in checkpoint["channel_versions"]
        })

    # 2. 计算下一个版本号
    next_version = get_next_version(
        max(checkpoint["channel_versions"].values())
        if checkpoint["channel_versions"] else None,
        None,
    )

    # 3. 消费已读 Channel
    for chan in {chan for task in tasks for chan in task.triggers
                 if chan not in RESERVED and chan in channels}:
        if channels[chan].consume() and next_version is not None:
            checkpoint["channel_versions"][chan] = next_version

    # 4. 按 Channel 分组写入
    pending_writes_by_channel = defaultdict(list)
    for task in tasks:
        for chan, val in task.writes:
            if chan in channels:
                pending_writes_by_channel[chan].append(val)

    # 5. 应用写入到 Channel
    updated_channels = set()
    for chan, vals in pending_writes_by_channel.items():
        if channels[chan].update(vals) and next_version is not None:
            checkpoint["channel_versions"][chan] = next_version
            if channels[chan].is_available():
                updated_channels.add(chan)

    # 6. 通知未更新的 Channel 新步开始
    if bump_step:
        for chan in channels:
            if channels[chan].is_available() and chan not in updated_channels:
                if channels[chan].update(EMPTY_SEQ):
                    ...

    # 7. 尝试触发 finish()
    if bump_step and updated_channels.isdisjoint(trigger_to_nodes):
        for chan in channels:
            if channels[chan].finish():
                ...
                updated_channels.add(chan)

    return updated_channels

6.7.1 算法详解

让我们逐步解析这个算法:

步骤 1 - 更新 versions_seen :记录每个任务"看到"了它触发 Channel 的哪个版本。这是下一次 _triggers() 比较的基准。

步骤 2 - 计算 next_version :所有在本步中发生变更的 Channel 都会被赋予同一个版本号。默认的版本函数 increment 简单地将整数加 1。

步骤 3 - 消费已读 Channel :某些 Channel(如 EphemeralValue)在被消费后会清除自己的值。这确保了路由信号是一次性的。consume() 方法返回 True 表示 Channel 状态发生了变化。

步骤 4-5 - 分组并应用写入 :将同一 Channel 的所有写入收集在一起,然后一次性调用 channel.update(values)。这保证了 Channel 的更新是原子的------要么全部应用,要么因 InvalidUpdateError 全部拒绝。

步骤 6 - 空通知 :即使没有被写入的 Channel,也会收到一个空的 update([]) 调用。这允许 EphemeralValue 在超步之间清除自己------如果上一步有值但本步没有写入,update([]) 会将 value 设为 MISSING,Channel 变为不可用。

步骤 7 - finish() 触发 :当所有更新的 Channel 都不在 trigger_to_nodes 中时(即没有节点会被这些更新触发),算法认为这是最后一个超步,调用所有 Channel 的 finish() 方法。LastValueAfterFinish Channel 在此时变为可用,触发 defer 节点。

sequenceDiagram participant PL as PregelLoop participant Tasks as 任务集合 participant Channels as Channels participant CP as Checkpoint Note over PL: after_tick() 开始 PL->>Tasks: 收集所有 writes PL->>CP: 更新 versions_seen PL->>Channels: 消费已读 Channel (consume) Note over Channels: EphemeralValue 清除旧值 PL->>Channels: 分组应用写入 (update) Channels-->>Loop: 返回 updated_channels PL->>Channels: 空通知未更新 Channel (update([])) Note over Channels: EphemeralValue 清除未被写入的值 alt 无触发节点 PL->>Channels: finish() Note over Channels: LastValueAfterFinish 变为可用 end PL->>CP: 保存 Checkpoint PL->>Loop: step += 1

6.8 after_tick():超步结束处理

after_tick() 在每个超步的所有任务执行完成后被调用:

python 复制代码
def after_tick(self) -> None:
    # 1. 应用所有写入
    self.updated_channels = apply_writes(
        self.checkpoint, self.channels,
        self.tasks.values(),
        self.checkpointer_get_next_version,
        self.trigger_to_nodes,
    )

    # 2. 发出 values 流式事件(如果输出 Channel 被更新)
    if not self.updated_channels.isdisjoint(
        (self.output_keys,) if isinstance(self.output_keys, str)
        else self.output_keys
    ):
        self._emit("values", map_output_values,
                    self.output_keys, writes, self.channels)

    # 3. 清空 pending_writes
    self.checkpoint_pending_writes.clear()

    # 4. 标记不再是重放模式
    self.is_replaying = False

    # 5. 保存 Checkpoint
    self._put_checkpoint({"source": "loop"})

    # 6. 检查 interrupt_after
    if self.interrupt_after and should_interrupt(
        self.checkpoint, self.interrupt_after, self.tasks.values()
    ):
        self.status = "interrupt_after"
        raise GraphInterrupt()

after_tick 中有几个细节值得关注:

  • is_replaying = False:只在第一个 tick 中可能为 True(从 Checkpoint 恢复时重放),之后的 tick 都是新执行
  • Checkpoint 保存时机 :每个超步结束后都保存 Checkpoint(除非 durability="exit"),这确保了即使进程崩溃也能从最近的超步恢复
  • interrupt_after 检查:在写入应用之后、下一步开始之前检查,确保中断时状态已经更新

6.9 版本追踪机制

版本追踪是 Pregel 调度机制的基石。它通过两张表实现:

6.9.1 channel_versions

channel_versions 是 Checkpoint 的一部分,记录每个 Channel 的当前版本号:

python 复制代码
checkpoint["channel_versions"] = {
    "messages": 3,
    "count": 3,
    "branch:to:agent": 2,
    "branch:to:tool": 3,
    # ...
}

每次 apply_writes 更新 Channel 时,对应的版本号会被设置为 next_version(所有本步更新的 Channel 共享同一个版本号)。

6.9.2 versions_seen

versions_seen 记录每个节点(以及特殊的 INTERRUPT 标识)上次执行时看到的 Channel 版本:

python 复制代码
checkpoint["versions_seen"] = {
    "agent": {
        "branch:to:agent": 2,
    },
    "tool": {
        "branch:to:tool": 2,
    },
    INTERRUPT: {
        "messages": 3,
        "count": 3,
        # ...
    }
}

6.9.3 触发判定

节点是否需要被触发,取决于以下比较:

css 复制代码
channel_versions[trigger_chan] > versions_seen[node][trigger_chan]

如果节点的任何触发 Channel 的当前版本大于该节点上次看到的版本,说明有新数据需要处理,节点被触发。

flowchart TB subgraph "版本追踪示例" direction TB CV["channel_versions:\nbranch:to:agent = 4\nbranch:to:tool = 3"] VS["versions_seen:\nagent: branch:to:agent = 2\ntool: branch:to:tool = 3"] CV --> CMP{比较} VS --> CMP CMP --> R1["agent: 4 > 2 = True\n需要触发"] CMP --> R2["tool: 3 > 3 = False\n不触发"] end style R1 fill:#c8e6c9 style R2 fill:#fce4ec

这个机制的优雅之处在于:

  1. 无需全局锁:每个节点只关心自己订阅的 Channel 版本,不需要全局协调
  2. 幂等性:同一版本不会触发重复执行
  3. 自然停止:当没有任何 Channel 被更新时,所有版本比较都返回 False,循环自然终止

6.10 停止条件:max_steps 与 recursion_limit

Pregel 执行循环有两种正常停止条件和一种安全停止条件:

正常停止 1 - 无任务prepare_next_tasks 返回空字典,说明没有节点需要被触发,图执行完成。

正常停止 2 - 中断 :遇到 interrupt_beforeinterrupt_after 节点,抛出 GraphInterrupt

安全停止 - 步数限制

python 复制代码
def tick(self) -> bool:
    if self.step > self.stop:
        self.status = "out_of_steps"
        return False

self.stop 在初始化时计算:

python 复制代码
self.stop = self.step + self.config["recursion_limit"] + 1

默认的 recursion_limit 是 25,意味着最多执行 25 个超步。如果超过这个限制,stream() 方法会抛出 GraphRecursionError

python 复制代码
if loop.status == "out_of_steps":
    raise GraphRecursionError(
        f"Recursion limit of {config['recursion_limit']} reached "
        "without hitting a stop condition."
    )

这个安全阀防止了无限循环------在 AI 工作流中,循环依赖和非终止条件是常见的 bug,步数限制确保了图总会终止。

6.11 SyncPregelLoop vs AsyncPregelLoop

LangGraph 为同步和异步场景提供了两个 PregelLoop 实现。它们共享 PregelLoop 基类的所有逻辑(tickafter_tick_first 等),区别仅在于 I/O 操作的同步/异步形式:

组件 SyncPregelLoop AsyncPregelLoop
上下文管理器 AbstractContextManager AbstractAsyncContextManager
后台执行器 BackgroundExecutor(线程池) AsyncBackgroundExecutor(asyncio)
Checkpoint 读取 checkpointer.get_tuple() checkpointer.aget_tuple()
Checkpoint 写入 checkpointer.put() checkpointer.aput()
Writes 保存 checkpointer.put_writes() checkpointer.aput_writes()
缓存操作 cache.get() / cache.set() cache.aget() / cache.aset()
python 复制代码
class SyncPregelLoop(PregelLoop, AbstractContextManager):
    def __init__(self, ...):
        super().__init__(...)
        self.stack = ExitStack()
        if checkpointer:
            self.checkpointer_get_next_version = checkpointer.get_next_version
            self.checkpointer_put_writes = checkpointer.put_writes
        else:
            self.checkpointer_get_next_version = increment
            self.checkpointer_put_writes = None

class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
    def __init__(self, ...):
        super().__init__(...)
        self.stack = AsyncExitStack()
        if checkpointer:
            self.checkpointer_get_next_version = checkpointer.get_next_version
            self.checkpointer_put_writes = checkpointer.aput_writes  # 注意 a 前缀
        else:
            self.checkpointer_get_next_version = increment
            self.checkpointer_put_writes = None

注意:即使没有 Checkpointer,版本追踪仍然工作。increment 函数作为默认的版本生成器确保了 Channel 版本在没有持久化的情况下仍然正确递增。

6.12 Checkpoint 保存策略

PregelLoop 支持三种持久化模式(Durability),通过 durability 参数控制:

flowchart LR subgraph "sync 模式" S1[超步完成] --> S2[保存 Checkpoint] S2 --> S3[等待保存完成] S3 --> S4[开始下一步] end subgraph "async 模式(默认)" A1[超步完成] --> A2[异步保存 Checkpoint] A1 --> A3[开始下一步] A2 -.->|后台| A4[保存完成] end subgraph "exit 模式" E1[超步完成] --> E2[继续下一步] E2 --> E3[...] E3 --> E4[图退出时保存] end

_put_checkpoint 方法中,保存操作被提交给后台执行器:

python 复制代码
def _put_checkpoint(self, metadata):
    # ...
    self._put_checkpoint_fut = self.submit(
        self._checkpointer_put_after_previous,
        getattr(self, "_put_checkpoint_fut", None),  # 前一个保存的 Future
        self.checkpoint_config,
        copy_checkpoint(self.checkpoint),
        self.checkpoint_metadata,
        new_versions,
    )

_checkpointer_put_after_previous 确保 Checkpoint 按顺序保存------它会等待前一个保存操作完成后再执行当前保存。这个"链式等待"设计避免了并发写入导致的顺序问题,同时不阻塞主执行线程。

sync 持久化模式下,stream() 方法会在每步结束后显式等待保存完成:

python 复制代码
loop.after_tick()
if durability_ == "sync":
    loop._put_checkpoint_fut.result()  # 阻塞等待

6.13 put_writes:增量写入机制

put_writes 方法处理任务执行过程中产生的写入。它不仅更新内存中的 checkpoint_pending_writes,还会在非 exit 模式下即时保存到 Checkpointer:

python 复制代码
def put_writes(self, task_id, writes):
    if not writes:
        return
    # 去重特殊 Channel 的写入
    if all(w[0] in WRITES_IDX_MAP for w in writes):
        writes = list({w[0]: w for w in writes}.values())

    # 更新内存中的 pending writes
    self.checkpoint_pending_writes.extend(
        (task_id, c, v) for c, v in writes
    )

    # 即时保存到 Checkpointer(非 exit 模式)
    if self.durability != "exit" and self.checkpointer_put_writes:
        self.submit(
            self.checkpointer_put_writes,
            config, writes_to_save, task_id,
        )

    # 输出流式事件
    if hasattr(self, "tasks"):
        self.output_writes(task_id, writes)

这种"写入即保存"的策略确保了即使进程在超步中途崩溃,已完成任务的结果也不会丢失。恢复时,_match_writes 方法会将保存的 pending writes 重新匹配到对应的任务。

6.14 设计决策分析

为什么选择 BSP 模型而非 Actor 模型?

纯 Actor 模型中,消息发送是异步的,接收是即时的。但在 AI 工作流中,我们需要更强的一致性保证:

  1. 状态一致性:同一步中所有节点看到的是同一个状态快照,不会出现"读到半更新状态"
  2. 可重放性:BSP 的确定性执行顺序使得从 Checkpoint 恢复后能够精确重现执行过程
  3. 调试友好:步的概念使得"在第 3 步之后中断"这样的调试操作变得自然

为什么 apply_writes 中有 finish() 机制?

finish() 是为 defer 节点设计的。考虑一个场景:所有正常节点都执行完毕,但还有一个 defer 节点等待触发。此时 updated_channelstrigger_to_nodes 没有交集(正常节点的路由 Channel 不触发任何节点),finish() 被调用,LastValueAfterFinish Channel 变为可用,defer 节点在下一步被触发。

为什么版本号采用全局递增而非 Channel 独立递增?

统一的版本号简化了比较逻辑,同时支持一个重要特性:should_interrupt 函数检查"自上次中断以来是否有任何更新",这需要跨 Channel 比较,全局递增的版本号使这种比较成为可能。

6.15 小结

本章深入剖析了 Pregel 执行引擎的完整实现。核心要点回顾:

  1. BSP 超步模型 :计算以超步为单位推进,每步包含计划(tick)、执行(runner.tick)、更新(after_tick)三个阶段
  2. PregelLoop 状态机 :从 inputpendingdone/out_of_steps/interrupt,管理整个生命周期
  3. prepare_next_tasks :通过 Channel 版本比较和 trigger_to_nodes 优化,高效确定每步要执行的任务
  4. apply_writes :在超步之间原子地更新 Channel,同时维护版本追踪表,支持 consume/finish 等生命周期方法
  5. 版本追踪channel_versionsversions_seen 两张表协同工作,实现高效的变更检测和幂等调度
  6. 安全停止recursion_limit 通过步数比较提供硬性终止保证

Pregel 执行引擎本身不直接执行任务------它只负责调度和状态管理。真正的任务执行由 PregelRunner 负责,涉及线程池并行、重试策略、缓存匹配等复杂机制。下一章将深入这些内容。

相关推荐
Irissgwe5 小时前
LangChain之核心组件(输出解析器)
ai·langchain·llm·ai编程·输出解析器
KaneLogger7 小时前
如何提升模型编码能力
agent·ai编程
louiX7 小时前
初级 AI Agent 工程师
langchain·agent·客户端
阿珊和她的猫8 小时前
从实践中提炼的架构设计与工程规范
ai·agent·llama·cli·mcp
幸福巡礼8 小时前
【LangChain 1.2 实战(六)】 工具调用 (Function Calling)
langchain
大山同学9 小时前
Feynman—证据驱动的 AI 研究代理
人工智能·agent·智能体
欧雷殿9 小时前
跨设备自动化:家庭 AI 工作台的首个小目标
后端·agent·aiops
DigitalOcean9 小时前
AI变智能体,传统云不够用了:成本降67%,延迟降40%的新解法
aigc·agent
Irissgwe11 小时前
LangChain之核心组件(少样本提示词)
人工智能·langchain·llm·langgraph
python零基础入门小白11 小时前
从0到1:手把手教你用Coze打造AI Agent,小白也能转行AI!
人工智能·学习·程序员·大模型·agent·产品经理·ai大模型