PregelLoop: superstep loop driver

源码版本1.2.9

Responsibilities

PregelLoop is the "metronome" of the LangGraph runtime. Pregel itself only assembles resources and exposes invoke / astream entry points; what actually runs the graph round after round is a PregelLoop instance. It holds the current checkpoint, channel set, pending_writes, step counter, status state machine, and a tasks dict — the full context of a single superstep. The outer while loop.tick(): ... loop.after_tick() form is built on top of it.

Its two core methods are tick (tick:599) and after_tick (after_tick:683). The former handles "before running": checks whether the step limit is exceeded, calls prepare_next_tasks to load the nodes to execute this step into self.tasks, checks interrupts, handles drain_requested, and reapplies pending_writes left over from a previous resume to succeeded nodes. The latter handles "after running": collects all task.writes, calls apply_writes to merge them into channels, clears pending, persists the checkpoint, and then decides whether to interrupt_after. The two together sandwich the BSP loop of "read input -> schedule -> execute -> write back -> persist".

In other words, PregelLoop does not run nodes itself — that is PregelRunner's job; it only decides "who to look at this step, how to wrap up after running, and whether to continue next step".

Design motivation

Why split the loop logic into tick and after_tick instead of a single step() method?

• Execution control is handed back to the caller: after tick returns True, the caller (the implementation body of Pregel.astream, see sync main loop:2964-2984) drives runner.tick itself to run this batch of tasks, then calls back loop.after_tick. This way PregelLoop does not need to hold a reference to PregelRunner; the two objects collaborate in a loosely coupled way rather than being nested.
• Explicit state machine: the status field has seven values: "input" / "pending" / "done" / "draining" / "interrupt_before" / "interrupt_after" / "out_of_steps" (status Literal:256-264). Each value corresponds to a kind of exit reason, and external code uses it to decide whether to resume, raise, or break — no guessing about return value semantics.
• Clear resume semantics: is_replaying is true only on the first tick — it means this step is "replaying" tasks that had already completed before the last interruption, and is turned off once after_tick finishes (is_replaying = False:716). Combined with _reapply_writes_to_succeeded_nodes re-applying the successful writes remaining in the checkpoint to in-memory tasks, succeeded nodes do not rerun and failed nodes go into the error handler.
• Drain is the external intervention entry: RunControl.drain_requested is set by the stream consumer (e.g. the client wants to stop midway). As soon as tick sees it, it enters the draining state and returns False, letting the outer while exit naturally instead of hard-killing the node mid-execution.
• Exit-mode delta writes and checkpoint timing: _delta_write_futs (_delta_write_futs:207) collects all put_writes futures for delta channels; _checkpointer_put_after_previous must drain them before writing the next checkpoint — preserving the causal order that "writes produce the checkpoint, not the other way around".

Key files

• class PregelLoop:158 — class definition; fields are almost all bare attributes without property wrappers, making it easy for SyncPregelLoop / AsyncPregelLoop subclasses to override them directly.
• status Literal:256-264 — the seven-value enum; "out_of_steps" is triggered by the step limit, "draining" by an external stop request, "done" is normal completion.
• tick method:599-681 — the full logic of one tick: step check, prepare_next_tasks, empty-tasks done, drain, _reapply_writes_to_succeeded_nodes, should_interrupt.
• out_of_steps:607-609 — when self.step > self.stop, enters out_of_steps and returns False; the outer while exits.
• done branch:653-655 — when self.tasks is empty, enters done state; this is the graph's "natural death".
• draining branch:657-659 — when control.drain_requested is true, enters draining state.
• reapply + resume handlers:662-664 — on resume, first reapplies pending writes to succeeded nodes, then calls _resume_error_handlers_if_applicable to prepare handler tasks for failed nodes.
• interrupt_before:667-671 — calls should_interrupt to check whether nodes fall in the interrupt_before list; a hit raises GraphInterrupt.
• after_tick method:683-726 — collect writes -> apply_writes -> emit values -> clear pending -> _put_checkpoint -> interrupt_after check.
• _reapply_writes_to_succeeded_nodes:736-749 — reapplies pending writes to in-memory tasks but skips the four control signals ERROR / ERROR_SOURCE_NODE / INTERRUPT / RESUME.
• _resume_error_handlers_if_applicable:751-816 — scans ERROR_SOURCE_NODE markers and constructs error-handler tasks for failed nodes, so the runner skips the original task and runs the handler instead.
• put_writes:415-508 — entry point for task-produced writes; dedupes control channels, accumulates null tasks, starts the checkpointer's put_writes future.
• SyncPregelLoop:1469 / AsyncPregelLoop:1722 — the two subclasses, overriding __enter__/__exit__, accept_push, schedule_error_handler and other sync/async-specific parts.

Data flow

The entry of tick first checks the step limit, then calls prepare_next_tasks to compute "which tasks to run next". This step is the seed of the entire superstep — all subsequent actions process this batch of tasks.

def tick(self) -> bool:
    """Execute a single iteration of the Pregel loop.

    Returns:
        True if more iterations are needed.
    """

    # check if iteration limit is reached
    if self.step > self.stop:
        self.status = "out_of_steps"
        return False

    # prepare next tasks
    self.tasks = prepare_next_tasks(
        self.checkpoint,
        self.checkpoint_pending_writes,
        self.nodes,
        self.channels,
        self.managed,
        self.config,
        self.step,
        self.stop,
        for_execution=True,
        manager=self.manager,
        store=self.store,
        checkpointer=self.checkpointer,
        trigger_to_nodes=self.trigger_to_nodes,
        updated_channels=self.updated_channels,
        retry_policy=self.retry_policy,
        cache_policy=self.cache_policy,
    )

This comes from tick head:599-629. self.step > self.stop is the "step limit" exit condition — stop is set to recursion_limit in PregelLoop.__init__, and exceeding the recursion depth stops at out_of_steps. prepare_next_tasks takes the current checkpoint, pending_writes, node table, and channel table to compute which tasks to execute this step and stuffs them into self.tasks. for_execution=True means these tasks are actually run (not a dry-run used for streaming preview).

Next come three parallel exit conditions — empty tasks / drain / reapply writes on resume — followed by should_interrupt:

# if no more tasks, we're done
if not self.tasks:
    self.status = "done"
    return False

if self.control is not None and self.control.drain_requested:
    self.status = "draining"
    return False

# if there are pending writes from a previous loop, apply them
if not self.is_replaying and self.checkpoint_pending_writes:
    self._reapply_writes_to_succeeded_nodes(self.tasks)
    self._resume_error_handlers_if_applicable()

# before execution, check if we should interrupt
if self.interrupt_before and should_interrupt(
    self.checkpoint, self.interrupt_before, self.tasks.values()
):
    self.status = "interrupt_before"
    raise GraphInterrupt()

This comes from tick middle:652-671. The interrupt_before check is only done after the resume reapply of pending_writes — because on resume the channel versions already reflect the previous writes, so should_interrupt can correctly judge "whether there are new updates since the last interrupt".

after_tick is called after the runner has finished running all tasks, and handles wrap-up:

def after_tick(self) -> None:
    # finish superstep
    writes = [w for t in self.tasks.values() for w in t.writes]
    self._delta_channels_with_overwrite.update(
        ch
        for ch, v in writes
        if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0]
    )
    # all tasks have finished
    self.updated_channels = apply_writes(
        self.checkpoint,
        self.channels,
        self.tasks.values(),
        self.checkpointer_get_next_version,
        self.trigger_to_nodes,
    )

This comes from after_tick head:683-698. First all task writes are flattened, and the delta channels that had an overwrite are tracked (this affects the starting value of sparse replay); then apply_writes actually merges those writes into channels, and the returned updated_channels is used by the next step's prepare_next_tasks to speed up "finding the next node to trigger". The tail of after_tick also clears checkpoint_pending_writes, sets is_replaying to False, calls _put_checkpoint({"source": "loop"}) to persist the checkpoint, and then runs the interrupt_after check (after_tick tail:714-724).

The lifecycle of a whole superstep can be drawn like this:

Boundaries and failure modes

• out_of_steps is not an error: when step > stop, False is returned directly, status is set to out_of_steps, and the outer while exits naturally. The upper layer of Pregel decides based on status whether to throw RecursionError or silently wrap up — this is a "soft limit" semantics, not an exception path (out_of_steps:607-609).
• draining cannot emit lifecycle events: _emit_graph_lifecycle_event explicitly refuses to be called in the draining state (draining guard:384-385), because drain is a user-initiated stop and does not count as interrupt / resume semantics.
• Resume must skip control signals in pending writes: _reapply_writes_to_succeeded_nodes must skip ERROR / ERROR_SOURCE_NODE / INTERRUPT / RESUME (skip control signals:746-747), otherwise tasks that have failed would be treated as successful and handlers that have already run would be triggered again.
• Error handlers are rescheduled on resume: _resume_error_handlers_if_applicable only takes effect for nodes with an error_handler_node (handler_node check:791-793); failed nodes without a handler are left as-is and go through the panic path in the runner's _should_stop_others.
• The null-task accumulation semantics in put_writes: writes for NULL_TASK_ID accumulate rather than overwrite (null task accumulate:422-431), for input writes that "do not belong to any node but must enter the checkpoint".
• Delta channel futures must be flushed before the next checkpoint: _delta_write_futs collects all put_writes futures for delta channels, and the next checkpoint drains it first (_delta_write_futs:201-207) — reversing the order would make sparse replay miss the writes that produced the checkpoint.

Summary

PregelLoop is a state-machine-driven superstep metronome that sandwiches the three phases of "prepare tasks -> run -> wrap up" around PregelRunner without touching execution itself. Once you understand the two phases tick / after_tick, the seven status values, and the resume mechanism of is_replaying + _reapply_writes_to_succeeded_nodes, you have grasped the main axis of LangGraph's runtime model. See /pregel/runner for how the runner actually runs tasks, /pregel/algo for the internals of prepare_next_tasks / apply_writes / should_interrupt, and /pregel/pregel for the overall assembly. Channel write-back semantics are in /channel/base-channel.

See official docs: LangGraph docs · README.