Pregel 算法层:中断、写回、调度

源码版本1.2.9

职责

_algo.py 是 PregelLoop 和 PregelRunner 的「数学内核」。loop 负责「跑几圈」,runner 负责「怎么并发执行」,而 algo 负责「每一步的语义」:什么时候该中断、节点写回的 writes 怎么合进通道 (channel)、下一步要触发哪些节点。它没有状态、没有 IO——全是纯函数,输入 checkpoint + 通道 + tasks,输出新的 channel 版本表 / updated_channels 集合 / 一份新的 task 字典。

三个核心函数正好对应 PregelLoop.tick + after_tick 的三段调用:should_interrupt(should_interrupt:155)在 tick 中段被调,判断是否要在执行前抛 GraphInterrupt;apply_writes(apply_writes:232)在 after_tick 开头被调,把本步所有 task 的 writes 合进通道并返回 updated_channels;prepare_next_tasks(prepare_next_tasks:392)在 tick 开头被调,根据当前 checkpoint 算出本步要执行哪些 task。

换句话说,这三个函数决定了「Pregel 这套 BSP 循环里,每一步看到什么、做什么、收尾后通道变成什么样」。理解了它们,就理解了 Pregel 的「状态变化规则」。

设计动机

为什么把算法逻辑拆成纯函数?

• 可测试性:should_interrupt / apply_writes / prepare_next_tasks 都是纯函数,输入是 checkpoint dict + channel mapping,输出是数据结构。可以脱离 Pregel 的 IO、submit、checkpointer 单独做单元测试——这正是 libs/langgraph/tests/unit/ 里大量测试的形态。
• 同步/异步共用:PregelLoop 有 SyncPregelLoop 和 AsyncPregelLoop 两个子类,但 algo 层完全无 side effect,两个子类共用同一份逻辑(tick calls prepare_next_tasks:612)。否则同步异步得各写一份算法,bug 修两边。
• updated_channels 优化:apply_writes 返回 updated_channels 集合,prepare_next_tasks 接收它作为 hint,跳过「扫描所有节点看 trigger 是否触发」(updated_channels hint:475-486)。大图上这一步能把 O(N×M) 的 trigger 扫描降到 O(updated×triggered),是最重要的性能优化。
• should_interrupt 用 versions_seen 做去重:INTERRUPT 通道的 versions_seen(versions_seen INTERRUPT:163)记录上次 interrupt 时看到的通道版本号。只有「自上次 interrupt 以来有新通道更新」才考虑再 interrupt——避免循环触发 interrupt 后又 resume 又 interrupt 的死循环。
• apply_writes 末尾的 finish 信号:bump_step and updated_channels.isdisjoint(trigger_to_nodes) 时调 channels[chan].finish()(finish:336-342),给所有通道发「这是最后一个超步」信号。永久性通道(如 LastValue)用这个机会把自己标记为不可用,让 prepare_next_tasks 不再触发任何节点——这是图的「自然死亡」。

关键文件

• should_interrupt:155-185 — 检查 graph 是否该中断,先看 versions_seen[INTERRUPT] 判断有没有新更新,再看 triggered task 是否落在 interrupt_nodes 列表里。
• any_updates_since_prev_interrupt:161-168 — 用通道版本号比较「自上次 interrupt 以来是否有更新」,这是 interrupt 去重的核心。
• interrupt_nodes filter:170-184 — interrupt_nodes == "*" 表示所有非 hidden task 都算,否则按 task.name in interrupt_nodes 过滤。
• apply_writes:232-345 — 把本步 task 的 writes 合进通道,返回 updated_channels 集合。
• bump_step:256-259 — 任一 task 有 triggers 时 bump_step=True,意味着这一步要消耗通道并推进 step 计数。
• update seen versions:262-269 — 每个 task 把它读过的 trigger 通道的版本号记进 versions_seen,下次 _triggers 判断时就是「自上次执行以来有没有新更新」。
• consume channels:284-292 — 调 channels[chan].consume() 把本步读过的通道标记为已消费,这是 PULL 通道「读一次就清」语义的实现。
• apply writes to channels:315-323 — channels[chan].update(vals) 真正把 writes 合进通道,只有 is_available() 的通道才进 updated_channels。
• bump_step notify:326-333 — bump_step=True 时没被本步更新的可用通道也 update(EMPTY_SEQ),让它们的版本号推进——这样订阅这些通道但没在本步触发的节点,下次 tick 也不会被错误触发。
• finish signal:336-342 — 没有节点被本步更新触发时,给所有通道发 finish() 信号,永久性通道借机变 unavailable,图自然结束。
• prepare_next_tasks:392-513 — 算出下一步要跑哪些 task,PUSH(来自 Send 扇出)和 PULL(来自边触发)都从这里出。
• updated_channels optimization:475-486 — 用上一步的 updated_channels + trigger_to_nodes 反查触发的节点集合,跳过全表扫描。
• prepare_single_task:524 — 单个 task 构造,PUSH 走 prepare_push_task_functional / prepare_push_task_send,PULL 走 _triggers 判断 + _proc_input 取输入。
• _triggers:1260-1277 — 判断节点的 trigger 通道是否「有未读的新版本」,这是 PULL 调度的核心条件。
• local_read:188-224 — 节点函数读取通道值的入口,支持 fresh=True 读「自己刚写但还没合进通道」的本地视图,用于条件边 (conditional edge) 的即写即读语义。

数据流

should_interrupt 是这三个里最简单的,逻辑是「有新更新 + 当前 task 命中列表」:

def should_interrupt(
    checkpoint: Checkpoint,
    interrupt_nodes: All | Sequence[str],
    tasks: Iterable[PregelExecutableTask],
) -> list[PregelExecutableTask]:
    """Check if the graph should be interrupted based on current state."""
    version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
    null_version = version_type()  # type: ignore[misc]
    seen = checkpoint["versions_seen"].get(INTERRUPT, {})
    # interrupt if any channel has been updated since last interrupt
    any_updates_since_prev_interrupt = any(
        version > seen.get(chan, null_version)  # type: ignore[operator]
        for chan, version in checkpoint["channel_versions"].items()
    )
    # and any triggered node is in interrupt_nodes list
    return (
        [task for task in tasks if ...]
        if any_updates_since_prev_interrupt
        else []
    )

这段来自 should_interrupt:155-185。关键是 versions_seen[INTERRUPT] 这条特殊记录——它不是某个节点读过的版本号,而是「上次发生 interrupt 时,所有通道的版本号快照」。any_updates_since_prev_interrupt 比的就是这个快照和当前 channel_versions,只要有任意通道版本号更高就认为「有新东西」,才允许进入 interrupt 判断。这样 resume 后即使再 tick 触发了同样名字的节点,也不会立刻又 interrupt——必须等通道真的有新写入。

apply_writes 是收尾阶段的核心,先更新 versions_seen,再消费读过的通道,最后写新值:

# update seen versions
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"]
        }
    )

# Consume all channels that were read
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

# Apply writes to channels
updated_channels: set[str] = set()
for chan, vals in pending_writes_by_channel.items():
    if chan in channels:
        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)

这段来自 apply_writes 核心:262-323。三段语义很清晰:

1. 更新 versions_seen——把每个 task 当前的 trigger 通道版本号记下来,作为「这个节点已经看到这版通道」的证据。下次 _triggers 判断时会和未来版本号比,只有更新了才会再触发。
2. consume() 读过的通道——Topic 这类「读一次就清」的通道在 consume() 里把内部 buffer 清空,只保留版本号推进。这就是 Send 扇出的消息被消费后就不会再触发其它节点的原理。
3. update(vals) 写新值——把本步所有 task 对每个通道的 writes 收集起来一次性合入,触发 reducer / LastValue 等通道的合并逻辑。只有 is_available()(还有值且没被 finish)的通道才进 updated_channels——finish() 后的通道版本号会推进,但不会触发任何节点。

prepare_next_tasks 在 tick 开头被调,决定本步要跑哪些 task:

# Consume pending tasks
tasks_channel = cast(Topic[Send] | None, 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), None, ..., for_execution=for_execution, ...
        ):
            tasks.append(task)

# This section is an optimization that allows which nodes will be active
# during the next step.
if updated_channels and trigger_to_nodes:
    triggered_nodes: set[str] = set()
    for channel in updated_channels:
        if node_ids := trigger_to_nodes.get(channel):
            triggered_nodes.update(node_ids)
    candidate_nodes: Iterable[str] = sorted(triggered_nodes)
elif not checkpoint["channel_versions"]:
    candidate_nodes = ()
else:
    candidate_nodes = processes.keys()

for name in candidate_nodes:
    if task := prepare_single_task((PULL, name), None, ..., for_execution=for_execution, ...):
        tasks.append(task)
return {t.id: t for t in tasks}

这段来自 prepare_next_tasks 核心:441-513。两类 task 在这里合流:

• PUSH task:从 TASKS 通道取 Send 对象,这是上一步某个节点 return Send(...) 扇出来的,每个 Send 对应一个 PUSH task。
• PULL task:从 updated_channels 反查 trigger_to_nodes——这个 mapping 是编译期算好的「通道 → 订阅它的节点列表」,直接取并集就是本步该触发的候选节点。sorted 是为了保证 task 顺序确定,不影响执行结果但影响 checkpoint 重放。

每个候选节点再走 _triggers 二次确认(_triggers call:606-612)——「通道有新版本号且节点还没看过这版」才会真正生成 task。这一步是 PULL 节点去重的关键:就算 updated_channels 里包含了某节点的 trigger 通道,如果 versions_seen[name] 已经记过这版,task 不会被重新调度。

整个 algo 层在 Pregel 循环里的位置:

边界与失败

• versions_seen[INTERRUPT] 不更新会死循环:should_interrupt 只在 any_updates_since_prev_interrupt 为真时才返回非空列表,而 versions_seen[INTERRUPT] 的更新发生在 apply_writes 之后的 _put_checkpoint 隐式逻辑里。如果 interrupt 后 resume 时没正确推进 INTERRUPT 通道的 seen 版本,同一个更新会被反复判定为「新」,导致 interrupt-resume-interrupt 死循环(versions_seen INTERRUPT:163-168)。
• bump_step 是 step 推进的唯一信号:any(t.triggers for t in tasks) 为假时不 bump——这是「null task 只写通道不推进 step」的语义(bump_step:259)。null task 用来在 tick 开始前注入 input writes,不应该让 step 计数器跳。
• finish() 只在「没有节点被触发」时调:bump_step and updated_channels.isdisjoint(trigger_to_nodes)(finish condition:336)——只要 updated_channels 里还有任意通道被某节点订阅,就别 finish。这意味着「最后一步图死亡时」才发 finish,不是每步都发。
• 未知通道写入只 warn 不报错:pending_writes_by_channel 里发现 task 写了不在 channels 里的通道名,只是 logger.warning(unknown channel warn:310-313)。这是为了容错——比如节点写了已经 finish 掉的通道,不该让整个图崩。
• local_read 的 fresh=True 路径会拷贝通道:fresh 读取时为每个通道 channels[k].copy()(fresh read copy:213-219),在副本上应用本 task 的 writes——这样条件边 (conditional edge) 节点能读到自己刚写但还没 apply_writes 的本地视图,不影响其它 task 看到的全局状态。
• prepare_single_task 的 task id 哈希算法:checkpoint["v"] > 1 时用 _xxhash_str,否则用 _uuid5_str(task_id_func:550),这是 checkpoint 版本之间的兼容性——旧 checkpoint 用 uuid5 生成的 task id 必须能重放出来。

小结

_algo.py 三个纯函数 should_interrupt / apply_writes / prepare_next_tasks 把 Pregel 的语义层抽干净:何时中断、写回怎么合、下一步跑谁。它们没有状态、没有 IO,被 PregelLoop.tick + after_tick 调用驱动整个 BSP 循环。循环驱动层见 /pregel/loop;执行层见 /pregel/runner;整体组装见 /pregel/pregel;通道的 update / consume / finish / is_available 这几个方法的语义见 /channel/base-channel。

对照官方资料:LangGraph 文档 · README。