Skip to content

並行執行:@task 的 future 編排

源码版本1.2.9

職責

@task 的執行期語意由 PregelRunner 接管——task 呼叫本身不會立即執行,而是把函式和參數交給當前 runner,runner 決定何時、在哪個執行緒/協程上跑。這一層負責三件事:把 task 呼叫轉成 PUSH 任務塞進佇列、並行執行同一超步 (superstep) 內被觸發的多個任務、用 future 把結果回傳給呼叫方。PregelRunner(PregelRunner:135)是這套編排的中心。

它的位置在 PregelLoop 之下、BackgroundExecutor / AsyncBackgroundExecutor 之上。loop.tick 算出本步該跑哪些任務,把它們交給 runner.tick / atick(PregelRunner.tick:176);runner 透過 submit 把每個任務丟給 executor,自己用 FuturesDict 追蹤所有 in-flight future,等它們 done 後寫回 writes、把例外路由到錯誤處理節點。

設計動機

  • task 呼叫即 PUSH 任務:在 entrypoint 函式體裡寫 task_a(); task_b(),這兩個 task 不需要顯式宣告邊——_call_with_options 從 config 取出 CONFIG_KEY_CALL 回調(就是 runner 注入的 _call / _acall),它把 task 包成 Call 資料結構,透過 schedule_task 塞進 TASKS 通道(schedule PUSH task:713),下一輪 prepare_next_tasks 自然會發現它。
  • 同步與非同步統一抽象:SyncAsyncFuture(SyncAsyncFuture:253)同時是 concurrent.futures.Future 和有 __await__ 的物件——它的 __await__ 就是 yield 一次,把控制權讓給外層 runner。同步 entrypoint 裡 task().result() 阻塞取值,非同步 entrypoint 裡 await task() 讓出協程,兩者共用同一個未來物件。
  • next_tick 排程:子 task 被排程時不立即跑,而是用 __next_tick__=True(next_tick:768)告訴 executor「下一個 tick 才跑」——這樣當前 tick 的 writes 先 commit、先流出 stream 事件,再開始子任務,保證 stream 順序符合使用者在程式碼裡寫的呼叫順序。
  • 重入去重:同一 task 被多次呼叫時,_call 先在 futures 裡找有沒有同 id 的 future,有就復用;如果 task 已經跑完(next_task.writes 非空),直接從 writes 裡取 RETURN 值,構造一個已完成的 future 回傳(dedup already ran:724)——避免重試父任務時子任務被重複執行。
  • max_concurrency 訊號量:非同步路徑裡 AsyncBackgroundExecutorconfig["max_concurrency"],起一個 asyncio.Semaphore,所有 task coro 套一層 gated(semaphore, coro)(gated:214),控制同一圖內同時 in-flight 的 task 數量上限,防止 LLM 呼叫把事件loop 拖垮。

關鍵檔案

  • PregelRunner:135 — 並行執行器,持有 submit / put_writes / node_error_handler_map / schedule_error_handler
  • PregelRunner.tick:176 — 同步執行一批任務,concurrent.futures.wait(FIRST_COMPLETED) 迴圈拿 done future、commit writes。
  • PregelRunner.atick:360 — 非同步版本,asyncio.wait 替換同步 wait,迭代器語意相同。
  • _call:700 — 同步路徑的 CONFIG_KEY_CALL 回調,負責把 task 呼叫轉 PUSH 任務、去重、鏈 future。
  • _acall:789 — 非同步路徑回調,語意同上,處理 coroutine。
  • schedule PUSH task:713 — 透過 schedule_task(實為 loop.accept_push / aaccept_push)把 Call 資料結構塞進 TASKS 通道。
  • SyncAsyncFuture:253 — 同時是 sync Future 和 async awaitable 的橋接物件,__await__ 單次 yield 讓出控制權。
  • _call_with_options:276_TaskFunction.__call__ 的實作,從 config 取 CONFIG_KEY_CALL 呼叫之。
  • BackgroundExecutor:40 — 同步執行緒池上下文管理器,submitcopy_context() 在子執行緒裡繼承 contextvar。
  • AsyncBackgroundExecutor:122 — 非同步上下文管理器,基於 asyncio.get_running_loop(),可選 max_concurrency 訊號量。
  • gated:214async def gated(semaphore, coro):先 acquire()release(),把並行上限套在 task coro 外面。

資料流

_call 是 task 呼叫的核心,把「在 entrypoint 函式裡呼叫 task_fn(arg)」翻譯成「在 runner 裡排程一個 PUSH 任務」:

python
# schedule the next task, if the callback returns one
if next_task := schedule_task(
    task(),
    scratchpad.call_counter(),
    Call(
        func,
        input,
        retry_policy=retry_policy,
        cache_policy=cache_policy,
        callbacks=callbacks,
        timeout=timeout,
    ),
):
    if fut := next(
        (
            f
            for f, t in list(futures().items())
            if t is not None and t == next_task.id
        ),
        None,
    ):
        # if the parent task was retried,
        # the next task might already be running
        pass
    elif next_task.writes:
        # if it already ran, return the result
        fut = concurrent.futures.Future()
        ret = next((v for c, v in next_task.writes if c == RETURN), MISSING)
        # ...
    else:
        # schedule the next task
        fut = submit()(
            run_with_retry,
            next_task,
            retry_policy,
            configurable={
                CONFIG_KEY_CALL: partial(
                    _call,
                    weakref.ref(next_task),
                    futures=futures,
                    # ...
                ),
            },
            __reraise_on_exit__=False,
            __next_tick__=True,
        )
        SKIP_RERAISE_SET.add(fut)
        futures()[fut] = next_task

三條分支依次判斷:(1) 同 id task 已在跑——直接復用 future;(2) task 已經跑過(writes 非空)——從 writes 裡取 RETURN 值構造已完成 future;(3) 全新 task——submit()run_with_retry 丟給 BackgroundExecutorAsyncBackgroundExecutor,並把遞迴的 _call 作為 CONFIG_KEY_CALL 注入到子任務的 config 裡,這樣子任務內部再呼叫 task() 會走同一套排程。

PregelRunner.tick 自己的核心是 FuturesDict + concurrent.futures.wait(FIRST_COMPLETED) 迴圈(wait loop:281)。每當有 future done,彈出對應的 task,呼叫 commit 把它的 writes 寫回 loop,如果該 task 配了錯誤處理節點,再 submit 一個 error handler task 進來。整個迴圈跑到所有 in-flight future 都完成(或只剩 get_waiter 占位 future)才回傳,把控制權交還給 loop.after_tick 落通道。

邊界與失敗

  • 同步上下文呼叫 async task 報錯:_call 開頭 if inspect.iscoroutinefunction(func): raise RuntimeError(sync async mismatch:716)——同步 entrypoint 裡不能 await,非同步 task 沒法跑。
  • 子 task 例外冒泡到父:SKIP_RERAISE_SET.add(fut) 把子 task 的 future 標記成「不要在 tick 結束時 re-raise」,例外透過 chainFuture 鏈回父 future,父 task await 時才會感知(skip reraise:781)。這樣錯誤傳播路徑跟呼叫層級一致,而不是被 runner 直接拋出打斷整個 tick。
  • __next_tick__ 在同步模式是延遲執行:BackgroundExecutor.submitnext_tick(ctx.run, fn, ...)(next_tick sync:66)把任務排到下一輪 tick;非同步 executor 裡 __next_tick__ 是 noop,因為 asyncio 任務天然非同步排程。
  • max_concurrency 不影響節點級並行:AsyncBackgroundExecutor 的訊號量只約束 task 級 coro(gated),不影響 PregelRunner.tick 同時排程多少節點——節點並行由 prepare_next_tasks 決定,訊號量只是節流(semaphore:153)。
  • 父任務重試時子任務不重跑:如果父任務 retry,_call 走到「dedup already ran」分支,直接回傳已有 future 或從 writes 取結果(dedup already ran:724),保證冪等。
  • contextvar 繼承:BackgroundExecutor.submitcopy_context()(copy_context:62)讓子執行緒繼承父執行緒的 contextvar(RunnableConfig 等),避免子任務拿不到 config。

小結

@task 的並行性來自 PregelRunner 的 futures 編排,核心抽象是 SyncAsyncFuture 橋接同步/非同步、_call / _acall 把 task 呼叫轉 PUSH 任務、BackgroundExecutor / AsyncBackgroundExecutor 提供執行執行緒/協程。要追 PUSH 任務的排程細節看 /subgraph/send-command,要追 task 在 entrypoint 裡的語意看 /func/entrypoint-task。對照官方資料:LangGraph 文件 · README