PregelRunner:超步內的任務執行器

源码版本1.2.9

職責

PregelRunner 幹一件事:把 PregelLoop.tick 準備好的那批 PregelExecutableTask 真正跑起來,把它們的 writes 收齊交回 PregelLoop。它不關心圖長什麼樣,不關心 checkpoint,不關心通道怎麼 reduce——它只負責「並行執行 + 失敗處理 + writes 落盤」這三件事。

它的位置非常窄:Pregel.astream 的主循環裡(sync main loop:2964-2984),while loop.tick() 回傳 True 之後,緊跟著就是 for _ in runner.tick([t for t in loop.tasks.values() if not t.writes], ...):,跑完再 loop.after_tick()。也就是說,runner 的輸入是「tick 算出來的、還沒寫過 writes 的 task」,輸出是「這些 task 的 writes 已經進了 checkpoint_pending_writes」,中間沒有第二層緩衝。

PregelRunner 自身非常薄:兩個公開方法 tick(tick:176)和 commit(commit:574),外加一個 atick 非同步版。重活全在 _call / _acall(_call:700)/(_acall:789)這兩個底層呼叫器裡——它們負責把節點函式包成可排程的 future,處理 retry policy、子 task Send 扇出、串流 chunk 回呼,最終把 writes 寫回 task 物件。

設計動機

為什麼 runner 要做成 generator 而不是普通方法?

• 串流 yield:runner 是 for _ in runner.tick(...) 這種 generator 形態(yield return type:188),每跑完一個 task 就 yield 一次,把控制權交還上層——上層就能立刻把串流 chunk 發給客戶端,而不是等整個超步跑完。這是 LangGraph 串流 (streaming) 的底層基礎。
• 單 task 快路徑:len(tasks) == 1 and timeout is None and get_waiter is None 時走同步直接呼叫,不起執行緒池 (fast path:203-254)。絕大多數 agent 單分支場景命中這條路徑,避免執行緒池開銷。
• commit 作為 Future 回呼:FuturesDict.on_done(on_done:116)在每個 future 完成時被觸發,內部就是 weakref.WeakMethod(self.commit)——commit 是 per-task 的 writes 落盤掛鉤,而不是 batch 收尾。這樣 writes 進 pending_writes 的順序天然按完成順序,checkpoint 也能盡早寫。
• error handler 路由:_should_route_to_error_handler(_should_route_to_error_handler:171-174)判斷失敗節點是否配了 error_handler 節點。命中就把 exception id 加進 _handled_exception_ids,排程一個 handler task 替代原 task——這樣失敗不會立即 panic,而是走另一個節點收尾。
• _should_stop_others 取消機制:任何一個 task 拋非 GraphBubbleUp 異常時,runner 取消同批其它 in-flight task(_should_stop_others:616-634)——保證同超步內「要麼全成,要麼進 handler」的語意。GraphInterrupt 被顯式排除在外,因為中斷不算失敗。

關鍵檔案

• class PregelRunner:135-138 — 類別定義,docstring 一句話點明職責:執行 task、commit writes、yield 控制權、必要時中斷其它 task。
• __init__:140-169 — 持有 submit / put_writes 的 weakref、node_error_handler_map、schedule_error_handler/aschedule_error_handler 回呼,以及 _handled_exception_ids 這個跨 tick 的異常去重集合。
• FuturesDict:75-134 — 自訂 dict 子類別,on_done 回呼在每個 future 完成時觸發 commit;should_stop 欄位持有 _should_stop_others partial。
• tick signature:176-188 — 輸入是 Iterable[PregelExecutableTask],回傳 Iterator[None],signature 暴露 generator 語意。
• fast path:203-254 — 單 task + 無 timeout + 無 waiter 時直接同步跑,失敗時按需排程 error handler。
• schedule tasks:259-276 — 多 task 路徑:每個 task 起一個 future,透過 self.submit()(實際是 PregelLoop.submit)排程。
• concurrent wait loop:282-323 — concurrent.futures.wait(FIRST_COMPLETED) 循環,每完成一個 task 就 commit、emit 輸出、必要時起 handler task。
• commit method:574-613 — per-task 收尾:cancelled 寫 ERROR、GraphInterrupt 寫 INTERRUPT、普通 exception 寫 ERROR + ERROR_SOURCE_NODE、正常寫 task.writes + NO_WRITES marker。
• _should_stop_others:616-634 — 判斷是否要取消其它 task,顯式排除 GraphBubbleUp 和已 handled 的異常。
• _panic_or_proceed:650-697 — 收尾函式,取消所有 in-flight future,合併多個 GraphInterrupt 成一個,timeout 時拋 TimeoutError。
• _call:700-787 — 節點呼叫的底層包裝:retry、stream chunk emit、Send 扇出回呼 schedule_task,把 task.writes 累積起來。
• _should_route_to_error_handler:171-174 — task.name in self.node_error_handler_map 時回傳 True;error handler 節點本身不能被路由(防遞迴)。

資料流

tick 的開頭先構造 FuturesDict,這是個增強版 dict,每個 future 完成時自動調 commit:

def tick(
    self,
    tasks: Iterable[PregelExecutableTask],
    *,
    reraise: bool = True,
    timeout: float | None = None,
    retry_policy: Sequence[RetryPolicy] | None = None,
    get_waiter: Callable[[], concurrent.futures.Future[None]] | None = None,
    schedule_task: Callable[
        [PregelExecutableTask, int, Call | None],
        PregelExecutableTask | None,
    ],
) -> Iterator[None]:
    tasks = tuple(tasks)
    futures = FuturesDict(
        callback=weakref.WeakMethod(self.commit),
        event=threading.Event(),
        should_stop=partial(
            _should_stop_others, handled_exception_ids=self._handled_exception_ids
        ),
        future_type=concurrent.futures.Future,
    )
    # give control back to the caller
    yield

這段來自 tick 頭部:176-199。注意第一個 yield 立刻把控制權交還——上層 for _ in runner.tick(...) 第一次迭代只是「啟動」,真正的 task 執行發生在後續迭代中。這是 generator 協程的一種簡化形態,避免引入 async 關鍵字到同步路徑。

commit 是 per-task 的收尾函式,把 task 的 writes 落進 pending_writes:

def commit(
    self,
    task: PregelExecutableTask,
    exception: BaseException | None,
) -> None:
    if isinstance(exception, asyncio.CancelledError):
        task.writes.append((ERROR, exception))
        self.put_writes()(task.id, task.writes)
    elif exception:
        if isinstance(exception, GraphInterrupt):
            if exception.args[0]:
                writes = [(INTERRUPT, exception.args[0])]
                if resumes := [w for w in task.writes if w[0] == RESUME]:
                    writes.extend(resumes)
                self.put_writes()(task.id, writes)
        elif isinstance(exception, GraphBubbleUp):
            pass
        else:
            task.writes.append((ERROR, exception))
            if self._should_route_to_error_handler(task) and not isinstance(
                exception, GraphBubbleUp
            ):
                task.writes.append((ERROR_SOURCE_NODE, task.name))
                self._handled_exception_ids.add(id(exception))
            self.put_writes()(task.id, task.writes)
    else:
        if self.node_finished and (
            task.config is None or TAG_HIDDEN not in task.config.get("tags", [])
        ):
            self.node_finished(task.name)
        if not task.writes:
            task.writes.append((NO_WRITES, None))
        self.put_writes()(task.id, task.writes)

這段來自 commit:574-613。put_writes 是 PregelLoop.put_writes 的 weakref(put_writes:415),它把 writes 塞進 checkpoint_pending_writes 並起 checkpointer 的 future 持久化。注意幾種特殊 writes:

• CancelledError → 寫 (ERROR, exception),讓 after_tick 收尾時把 task 算成「失敗但已記錄」,不 panic。
• GraphInterrupt → 寫 (INTERRUPT, ...) + 任何 RESUME writes,中斷的恢復值要和 interrupt 一起落盤,下次 resume 時一起讀出來。
• 普通 exception + 配了 error handler → 多寫一條 (ERROR_SOURCE_NODE, task.name),這條標記會被 PregelLoop._resume_error_handlers_if_applicable 掃到(ERROR_SOURCE_NODE scan:775),用來在續跑時排程對應的 handler 節點。
• 正常完成但沒寫任何東西 → 補一條 (NO_WRITES, None),防止 prepare_next_tasks 把它當成「還沒跑過」(NO_WRITES marker:609-611)。

整個 runner 的執行流可以這樣畫:

邊界與失敗

• _handled_exception_ids 跨 tick 持久:__init__ 裡這個集合是 instance-level,不是 per-tick(_handled_exception_ids:169)。同一個異常物件 id 進了集合就不會再被 _should_stop_others 當成失敗,避免 error handler 跑完後又把原 exception 重拋。
• GraphBubbleUp 不算失敗:commit 對 GraphBubbleUp 顯式 pass,不寫任何 writes(GraphBubbleUp pass:592-594)。這類異常(如 ParentCommand)是給父圖發信號的,不是真錯誤。
• NO_WRITES 防止 task 重複執行:task 正常完成但沒寫任何東西,也要補個 (NO_WRITES, None)(NO_WRITES:609-611)。否則 after_tick 後 prepare_next_tasks 看到這個 task 沒 writes 會以為沒跑過,續跑時再跑一次。
• schedule_error_handler 可能回傳 None:handler 節點自己也可能沒設定或無法構造,這時 task 就走普通 panic 路徑(handler optional:230-248)。commit 裡的 ERROR_SOURCE_NODE write 會和這條路徑並存,_resume_error_handlers_if_applicable 續跑時再嘗試排程。
• timeout 撞了不拋 GraphInterrupt:_panic_or_proceed 的 if inflight: raise timeout_exc_cls("Timed out")(timeout:691-697)——timeout 是真異常,會冒泡出 Pregel,不像 GraphInterrupt 那樣被吞。
• traceback 修剪:reraise 之前會跳過 EXCLUDED_FRAME_FNAMES 列表裡的幀(tb trim:241-247),把 langgraph 內部堆疊幀過濾掉,讓使用者看到的 traceback 直接指向他們的節點程式碼。

小結

PregelRunner 是個薄殼:generator tick 排程並行執行 + per-task commit 把 writes 落進 pending_writes,中間塞了 error handler 路由和取消機制。理解了它,就理解了「一個超步裡 task 是怎麼真跑起來的」。上一層 PregelLoop 的循環驅動見 /pregel/loop;prepare_next_tasks / apply_writes / should_interrupt 這三個演算法函式見 /pregel/algo;整體組裝見 /pregel/pregel;串流輸出怎麼在 runner yield 之間被加工見 /stream/run-stream。

對照官方資料:LangGraph 文件 · README