并发执行:@task 的 future 编排
职责
@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信号量:异步路径里AsyncBackgroundExecutor接config["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— 同步线程池上下文管理器,submit用copy_context()在子线程里继承 contextvar。AsyncBackgroundExecutor:122— 异步上下文管理器,基于asyncio.get_running_loop(),可选max_concurrency信号量。gated:214—async def gated(semaphore, coro):先acquire()再release(),把并发上限套在 task coro 外面。
数据流
_call 是 task 调用的核心,把「在 entrypoint 函数里调 task_fn(arg)」翻译成「在 runner 里调度一个 PUSH 任务」:
# 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 丢给 BackgroundExecutor 或 AsyncBackgroundExecutor,并把递归的 _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,父 taskawait时才会感知(skip reraise:781)。这样错误传播路径跟调用层级一致,而不是被 runner 直接抛出打断整个 tick。 __next_tick__在同步模式是延迟执行:BackgroundExecutor.submit用next_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.submit用copy_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。