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 fan-out、streaming chunk コールバックを処理し、最終的に writes を task オブジェクトに書き戻します。

設計動機

なぜ runner を普通のメソッドではなく generator にするのか?

• ストリーミング yield:runner は for _ in runner.tick(...) という generator 形態です(yield return type:188)。一つの task を走らせ終わるごとに yield して制御権を上層に返します——上層はすぐに streaming chunk をクライアントへ送れます。スーパーステップ全体を待つ必要はありません。これが LangGraph のストリーミング (streaming) の底層基盤です。
• 単 task fast path:len(tasks) == 1 and timeout is None and get_waiter is None のとき同期直接呼び出しでスレッドプールを立ち上げません(fast path:203-254)。大半のエージェント単分岐シナリオはこのパスに命中し、スレッドプールのオーバーヘッドを避けます。
• 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) ループ。一つ完成するごとに 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 fan-out の 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 から bubble up し、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