Skip to content

Concurrency: @task future orchestration

源码版本1.2.9

Responsibilities

The runtime semantics of @task are taken over by PregelRunner — calling a task does not execute immediately. Instead the function and arguments are handed to the current runner, which decides when and on which thread/coroutine to run it. This layer does three things: turn task calls into PUSH tasks enqueued for execution, concurrently execute multiple tasks triggered within the same superstep, and return results to the caller via futures. PregelRunner (PregelRunner:135) is the center of this orchestration.

It sits below PregelLoop and above BackgroundExecutor / AsyncBackgroundExecutor. loop.tick computes which tasks should run this step and hands them to runner.tick / atick (PregelRunner.tick:176); the runner submits each task to the executor via submit, tracks all in-flight futures with a FuturesDict, and once they are done writes their writes back and routes exceptions to the error-handler node.

Design motivation

  • Task call as PUSH task: inside an entrypoint body task_a(); task_b() requires no explicit edges — _call_with_options pulls the CONFIG_KEY_CALL callback from config (the _call / _acall injected by the runner), wraps the task into a Call data structure, and pushes it into the TASKS channel via schedule_task (schedule PUSH task:713). The next prepare_next_tasks picks it up naturally.
  • Unified sync/async abstraction: SyncAsyncFuture (SyncAsyncFuture:253) is both a concurrent.futures.Future and an object with __await__ — its __await__ simply yields once, yielding control to the outer runner. In a sync entrypoint task().result() blocks for the value; in an async entrypoint await task() yields the coroutine. Both share the same future object.
  • next_tick scheduling: a sub-task is not run immediately when scheduled. Instead __next_tick__=True (next_tick:768) tells the executor "run this on the next tick" — so the current tick's writes commit and stream events flush first, then the sub-task starts. This keeps the stream order aligned with the call order the user wrote in code.
  • Reentry dedup: when the same task is called multiple times, _call first looks in futures for an existing future with the same id and reuses it; if the task has already run (next_task.writes non-empty), it takes the RETURN value directly from writes and constructs an already-completed future (dedup already ran:724) — this prevents sub-tasks from being re-executed when a parent task is retried.
  • max_concurrency semaphore: on the async path AsyncBackgroundExecutor reads config["max_concurrency"] and creates an asyncio.Semaphore; every task coro is wrapped in gated(semaphore, coro) (gated:214) to cap the number of in-flight tasks within a single graph and prevent LLM calls from starving the event loop.

Key files

  • PregelRunner:135 — concurrency executor; holds submit / put_writes / node_error_handler_map / schedule_error_handler.
  • PregelRunner.tick:176 — executes a batch of tasks synchronously; concurrent.futures.wait(FIRST_COMPLETED) loop pops done futures and commits writes.
  • PregelRunner.atick:360 — async version; asyncio.wait replaces the sync wait, iterator semantics identical.
  • _call:700 — the sync CONFIG_KEY_CALL callback; turns task calls into PUSH tasks, dedups, chains futures.
  • _acall:789 — async-path callback; same semantics, handles coroutines.
  • schedule PUSH task:713 — pushes a Call data structure into the TASKS channel via schedule_task (actually loop.accept_push / aaccept_push).
  • SyncAsyncFuture:253 — bridge object that is both a sync Future and an async awaitable; __await__ does a single yield to surrender control.
  • _call_with_options:276 — implementation of _TaskFunction.__call__; pulls CONFIG_KEY_CALL from config and invokes it.
  • BackgroundExecutor:40 — sync thread-pool context manager; submit uses copy_context() to inherit contextvars on sub-threads.
  • AsyncBackgroundExecutor:122 — async context manager based on asyncio.get_running_loop(); optional max_concurrency semaphore.
  • gated:214async def gated(semaphore, coro): acquire() then release(), wrapping the concurrency cap around a task coro.

Data flow

_call is the core of a task call — it translates "calling task_fn(arg) inside an entrypoint function" into "scheduling a PUSH task on the runner":

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

The three branches check in order: (1) same-id task is already running — reuse the future; (2) task already ran (writes non-empty) — construct a completed future from the RETURN value in writes; (3) brand-new task — submit() hands run_with_retry to BackgroundExecutor or AsyncBackgroundExecutor, and injects the recursive _call as CONFIG_KEY_CALL into the sub-task's config, so that a task() call inside the sub-task goes through the same scheduling path.

The core of PregelRunner.tick itself is a FuturesDict + concurrent.futures.wait(FIRST_COMPLETED) loop (wait loop:281). Each time a future completes, it pops the corresponding task, calls commit to write its writes back to the loop, and if that task had an error-handler node configured, submits an error-handler task. The loop runs until all in-flight futures are done (or only get_waiter placeholder futures remain), then returns control to loop.after_tick to drain the channel.

Boundaries and failure modes

  • Calling an async task in a sync context raises: at the start of _call, if inspect.iscoroutinefunction(func): raise RuntimeError (sync async mismatch:716) — a sync entrypoint cannot await, so an async task cannot run.
  • Sub-task exceptions bubble to the parent: SKIP_RERAISE_SET.add(fut) marks the sub-task future as "do not re-raise at the end of tick"; the exception is chained back to the parent future via chainFuture, and the parent task only observes it when it awaits. (skip reraise:781) This keeps the error-propagation path aligned with the call hierarchy, instead of being thrown by the runner and aborting the whole tick.
  • __next_tick__ is deferred on the sync path: BackgroundExecutor.submit uses next_tick(ctx.run, fn, ...) (next_tick sync:66) to queue the task onto the next tick; on the async executor __next_tick__ is a noop, because asyncio tasks are scheduled asynchronously anyway.
  • max_concurrency does not affect node-level concurrency: the AsyncBackgroundExecutor semaphore only constrains task-level coros (gated); it does not affect how many nodes PregelRunner.tick schedules in parallel — node concurrency is decided by prepare_next_tasks, the semaphore only throttles it (semaphore:153).
  • Sub-tasks are not re-run when the parent retries: if the parent task is retried, _call hits the "dedup already ran" branch and returns the existing future or the result from writes (dedup already ran:724), guaranteeing idempotency.
  • contextvar inheritance: BackgroundExecutor.submit uses copy_context() (copy_context:62) so sub-threads inherit the parent thread's contextvars (RunnableConfig and friends), preventing sub-tasks from losing access to config.

Summary

@task concurrency comes from PregelRunner's future orchestration. The core abstractions are SyncAsyncFuture bridging sync/async, _call / _acall turning task calls into PUSH tasks, and BackgroundExecutor / AsyncBackgroundExecutor providing the execution threads/coroutines. For PUSH-task scheduling details see /subgraph/send-command; for task semantics inside an entrypoint see /func/entrypoint-task. See official docs: LangGraph · README.