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),給所有通道發「這是最後一個超步 (superstep)」信號。永久性通道(如 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