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