From 0ae8870432c482a5d9b88c17f6962414980ad9d2 Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Mon, 4 May 2026 17:20:15 -0700 Subject: [PATCH] chore: minor clean up around checkpoint and delta channel (#7705) --- .../langgraph/checkpoint/base/__init__.py | 21 ++- libs/langgraph/langgraph/pregel/_algo.py | 101 +++++++++++-- .../langgraph/langgraph/pregel/_checkpoint.py | 136 ++++++++++++------ libs/langgraph/langgraph/pregel/_loop.py | 6 +- libs/langgraph/langgraph/pregel/main.py | 16 +-- 5 files changed, 213 insertions(+), 67 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index b34509f00..ac2e32f01 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -726,22 +726,29 @@ def empty_checkpoint() -> Checkpoint: def create_checkpoint( - checkpoint: Checkpoint, + prev_checkpoint: Checkpoint, channels: Mapping[str, ChannelProtocol] | None, step: int, *, id: str | None = None, ) -> Checkpoint: - """Create a checkpoint for the given channels.""" + """Simplified checkpoint constructor used only by saver test suites. + + The real runtime version lives in + ``langgraph.pregel._checkpoint.create_checkpoint`` and handles + DeltaChannel snapshots, version tracking, etc. This copy exists so + that ``libs/checkpoint-*`` tests can build test checkpoints without + depending on the ``langgraph`` main package. + """ from datetime import datetime, timezone ts = datetime.now(timezone.utc).isoformat() if channels is None: - values = checkpoint["channel_values"] + values = prev_checkpoint["channel_values"] else: values = {} for k, v in channels.items(): - if k not in checkpoint["channel_versions"]: + if k not in prev_checkpoint["channel_versions"]: continue try: values[k] = v.checkpoint() @@ -752,8 +759,8 @@ def create_checkpoint( ts=ts, id=id or str(uuid6(clock_seq=step)), channel_values=values, - channel_versions=checkpoint["channel_versions"], - versions_seen=checkpoint["versions_seen"], - pending_sends=checkpoint.get("pending_sends", []), + channel_versions=prev_checkpoint["channel_versions"], + versions_seen=prev_checkpoint["versions_seen"], + pending_sends=prev_checkpoint.get("pending_sends", []), updated_channels=None, ) diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index 103f6cce0..85d6fec77 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/libs/langgraph/langgraph/pregel/_algo.py @@ -236,19 +236,98 @@ def apply_writes( get_next_version: GetNextVersion | None, trigger_to_nodes: Mapping[str, Sequence[str]], ) -> set[str]: - """Apply writes from a set of tasks (usually the tasks from a Pregel step) - to the checkpoint and channels, and return managed values writes to be applied - externally. + """Apply task writes to channels and update checkpoint version metadata. + + Mutates two objects in-place: + + - ``checkpoint["channel_versions"]`` and ``checkpoint["versions_seen"]`` + (version tracking for scheduling — see below). + - ``channels[k].value`` (in-memory channel state via ``update`` / + ``consume`` / ``finish``). + + Does NOT touch ``checkpoint["channel_values"]`` (serialized blobs). + Serialization is deferred to ``create_checkpoint``, which reads the + final in-memory state from ``channels``. This avoids serializing + intermediate states that may be overwritten by later phases (e.g. + ``finish()`` can change a channel after ``update()``). + + Version system recap + --------------------- + ``channel_versions`` and ``versions_seen`` form the scheduling + mechanism that prevents nodes from re-executing on stale triggers. + + - ``channel_versions[chan]``: monotonically increasing stamp, bumped + whenever ``chan`` is modified. All channels modified in the same + superstep share one ``next_version`` (computed from the current + max). This works because writes within a superstep are atomic — + no ordering between them — so a single stamp suffices. + - ``versions_seen[node_name][chan]``: the version of ``chan`` that + ``node_name`` last consumed. Sparse — only the node's trigger + channels (defined by graph edges) are recorded. + ``prepare_next_tasks`` triggers a node only when + ``channel_versions[chan] > versions_seen[node][chan]``. + + :: + + superstep 1: next_version = 2 + channel_versions = {"messages": 2, "count": 2} + versions_seen = {"A": {"messages": 2}, "B": {"count": 2}} + + superstep 2: only "count" written, next_version = 3 + channel_versions = {"messages": 2, "count": 3} + Node A trigger "messages": 2 == seen 2 -> NOT triggered + Node B trigger "count": 3 > seen 2 -> triggered + + Pipeline (4 phases) + ------------------- + 1. **versions_seen**: record each task's trigger-channel versions so + ``prepare_next_tasks`` won't re-fire the same node on the same + data. + 2. **consume**: call ``ch.consume()`` on channels that triggered + tasks. Most channels no-op (base returns ``False``). + ``NamedBarrierValue`` resets its seen set; + ``LastValueAfterFinish`` clears its value after a finish cycle. + 3. **update** (two sub-phases): + - 3a: apply actual writes — ``ch.update(vals)`` for each channel. + - 3b: notify un-written channels of a new step via + ``ch.update([])``. Most return ``False``; ephemeral channels + use this to clear themselves. + 4. **finish**: if no updated channel can trigger further nodes (check + via ``updated_channels.isdisjoint(trigger_to_nodes)``), call + ``ch.finish()`` on all channels. ``LastValueAfterFinish`` uses + this to become available. + + Regular vs DeltaChannel + ----------------------- + Both channel types go through the same ``update()`` path and get the + same ``channel_versions`` bump. The difference is at checkpoint time: + + - ``BinaryOperatorAggregate.checkpoint()`` returns the full value, + stored in ``channel_values``. + - ``DeltaChannel.checkpoint()`` returns ``MISSING`` — omitted from + ``channel_values``. Its state is reconstructed on load via + ancestor replay of ``checkpoint_writes``. Args: - checkpoint: The checkpoint to update. - channels: The channels to update. - tasks: The tasks to apply writes from. - get_next_version: Optional function to determine the next version of a channel. - trigger_to_nodes: Mapping of channel names to the set of nodes that can be triggered by updates to that channel. + checkpoint: Mutated in-place — ``channel_versions`` (stamps + bumped for modified channels) and ``versions_seen`` (records + which versions each node consumed, preventing re-trigger). + channels: Mutated in-place (channel values via ``update`` / + ``consume`` / ``finish``). + tasks: The writes to apply, sorted deterministically by path. + get_next_version: Generates the next monotonic version stamp. + ``None`` = lightweight mode (skip version tracking, used for + graph drawing / dry runs). + trigger_to_nodes: Maps channel names to triggerable nodes + (derived from graph edges). Used in phase 4 to check if any + updated channel can trigger further nodes; if not, + ``finish()`` is called. Returns: - Set of channels that were updated in this step. + Set of channel names whose ``update()`` returned ``True`` and are + ``is_available()``. Used by ``_loop.py`` to populate + ``updated_channels`` for ``create_checkpoint`` and + ``prepare_next_tasks``. """ # sort tasks on path, to ensure deterministic order for update application # any path parts after the 3rd are ignored for sorting @@ -268,7 +347,9 @@ def apply_writes( } ) - # Find the highest version of all channels + # Lightweight mode: when get_next_version is None (e.g. graph drawing, + # dry runs without a checkpointer) we still apply channel state updates + # but skip all channel_versions tracking and updated_channels collection. if get_next_version is None: next_version = None else: diff --git a/libs/langgraph/langgraph/pregel/_checkpoint.py b/libs/langgraph/langgraph/pregel/_checkpoint.py index 4e3bbffd6..e1d8438e9 100644 --- a/libs/langgraph/langgraph/pregel/_checkpoint.py +++ b/libs/langgraph/langgraph/pregel/_checkpoint.py @@ -2,7 +2,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping from datetime import datetime, timezone -from typing import Any, cast +from typing import Any, NamedTuple, cast from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -23,6 +23,16 @@ LATEST_VERSION = 4 GetNextVersion = Callable[[Any, None], Any] +class CreateCheckpointResult(NamedTuple): + """Return value of :func:`create_checkpoint`.""" + + checkpoint: Checkpoint + """The checkpoint to persist via ``saver.put()``.""" + snapshotted: set[str] + """DeltaChannel names that were snapshotted this step. The caller + should reset their ``updates_since_snapshot`` counters to ``0``.""" + + def empty_checkpoint() -> Checkpoint: return Checkpoint( v=LATEST_VERSION, @@ -59,7 +69,7 @@ def _should_snapshot_delta( def create_checkpoint( - checkpoint: Checkpoint, + mutated_checkpoint: Checkpoint, channels: Mapping[str, BaseChannel] | None, step: int, *, @@ -68,39 +78,64 @@ def create_checkpoint( get_next_version: GetNextVersion | None = None, force_delta_snapshot: bool = False, updates_since_snapshot: Mapping[str, int] | None = None, - new_updates_since_snapshot: dict[str, int] | None = None, -) -> Checkpoint: - """Create a checkpoint for the given channels. +) -> CreateCheckpointResult: + """Build a new ``Checkpoint`` from the previous one and live channel state. - For each `DeltaChannel`, a `_DeltaSnapshot(value)` blob is written into - `channel_values[k]` when this channel has accumulated at least - `snapshot_frequency` updates since its last snapshot (counter supplied - via `updates_since_snapshot`). Otherwise the channel is omitted from - `channel_values`; its `channel_versions` entry still bumps so that the - saver tracks the channel and the ancestor walk can replay writes. + Args: + mutated_checkpoint: The checkpoint that has been mutated in-place by + ``apply_writes`` (its ``channel_versions`` and ``versions_seen`` + already reflect this superstep's writes). The new checkpoint + inherits these mutated values and builds new ``channel_values`` + from the live channels. + channels: In-memory channel objects whose state was updated by + ``apply_writes``. Needed because ``mutated_checkpoint`` only + carries updated version metadata — its ``channel_values`` still + holds stale blobs from the previous checkpoint load. This + function calls ``ch.checkpoint()`` (or wraps in ``_DeltaSnapshot`` + for DeltaChannels) to produce fresh serialised ``channel_values``. + Pass ``None`` to skip serialisation — the new checkpoint reuses + ``mutated_checkpoint``'s ``channel_values`` as-is (used by + ``durability="exit"`` intermediate steps or when no checkpointer + is present). + step: Superstep number, used to generate the checkpoint ``id`` via + ``uuid6(clock_seq=step)`` when *id* is not provided explicitly. + id: Explicit checkpoint id. When supplied (e.g. during ``exiting``), + overrides the ``uuid6``-based generation. + updated_channels: Set of channel names written during this superstep + (produced by ``apply_writes``). Persisted as a sorted list in + the new checkpoint for efficient cold-start scheduling. + get_next_version: Version generator (e.g. ``saver.get_next_version``). + ``None`` means version tracking is skipped (lightweight / no- + checkpointer mode). + force_delta_snapshot: When ``True``, every DeltaChannel is snapshotted + regardless of ``snapshot_frequency``. Used by ``durability="exit"`` + where intermediate ``checkpoint_writes`` are not stored, so ancestor + replay would have nothing to replay from. + updates_since_snapshot: *Read-only* counters — maps each DeltaChannel + name to the number of updates since its last snapshot. Used by + ``_should_snapshot_delta`` to decide whether to snapshot now. - Snapshots are eager: even if the channel had no write this step, a - version bump is forced (via `get_next_version`) so `put()` includes - the channel in `new_versions` and stores the blob. - - `force_delta_snapshot` ignores the cadence and always snapshots — - used by `durability="exit"` where intermediate writes are not stored - as ancestor `checkpoint_writes`. - - If `new_updates_since_snapshot` is provided, the function resets the - counter to `0` for any channel that snapshotted this step. Counters - for channels that did not snapshot are left untouched (the caller is - responsible for incrementing them based on `updated_channels`). + Returns: + A tuple of: + - The checkpoint to persist via ``saver.put()``, with a fresh + ``id``, ``ts``, and serialised ``channel_values`` / + ``channel_versions``. + - Names of DeltaChannels that were snapshotted this step. The + caller should reset their ``updates_since_snapshot`` counters + to ``0``. """ ts = datetime.now(timezone.utc).isoformat() counts = updates_since_snapshot or {} + snapshotted: set[str] = set() if channels is None: - values = checkpoint["channel_values"] - channel_versions = checkpoint["channel_versions"] + values = mutated_checkpoint["channel_values"] + channel_versions = mutated_checkpoint["channel_versions"] else: values = {} - channel_versions = dict(checkpoint["channel_versions"]) + channel_versions = dict(mutated_checkpoint["channel_versions"]) for k in channels: + # Channel has never been written to (no version entry from + # apply_writes), so there is no meaningful state to checkpoint. if k not in channel_versions: continue ch = channels[k] @@ -114,27 +149,48 @@ def create_checkpoint( force=force_delta_snapshot, ) ): - # Eager snapshot: bump version if not already written this step - # so put() includes this channel in new_versions and stores blob. - if get_next_version is not None and ( - updated_channels is None or k not in updated_channels + # Force-snapshot (durability="exit"): some channels may not + # have been written this step, so apply_writes didn't bump + # their version. Manually bump so saver.put() persists the + # blob. Other channels in the same force-snapshot batch + # *were* written this step and already bumped by + # apply_writes — those are skipped below to avoid + # double-bumping. In the normal count-based path the + # channel was necessarily written (otherwise count can't + # reach snapshot_frequency), so this branch never fires. + # TODO: force-snapshot on every exit is wasteful for short + # runs — a 1-message run still serialises the full state. + # The right fix is to persist checkpoint_writes for delta + # channels in durability="exit" mode so ancestor replay + # works, eliminating the need for force-snapshot entirely. + if ( + force_delta_snapshot + and get_next_version is not None + # Even in force-snapshot mode, some channels were + # written this step and already bumped by apply_writes. + # Skip those to avoid double-bumping. + and k not in (updated_channels or ()) ): channel_versions[k] = get_next_version(channel_versions[k], None) values[k] = _DeltaSnapshot(ch.get()) - if new_updates_since_snapshot is not None: - new_updates_since_snapshot[k] = 0 + snapshotted.add(k) else: v = ch.checkpoint() if v is not MISSING: values[k] = v - return Checkpoint( - v=LATEST_VERSION, - ts=ts, - id=id or str(uuid6(clock_seq=step)), - channel_values=values, - channel_versions=channel_versions, - versions_seen=checkpoint["versions_seen"], - updated_channels=None if updated_channels is None else sorted(updated_channels), + return CreateCheckpointResult( + checkpoint=Checkpoint( + v=LATEST_VERSION, + ts=ts, + id=id or str(uuid6(clock_seq=step)), + channel_values=values, + channel_versions=channel_versions, + versions_seen=mutated_checkpoint["versions_seen"], + updated_channels=None + if updated_channels is None + else sorted(updated_channels), + ), + snapshotted=snapshotted, ) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index ad772016a..c848a85d3 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -935,7 +935,7 @@ class PregelLoop: exiting or self.durability != "exit" ) # create new checkpoint - self.checkpoint = create_checkpoint( + result = create_checkpoint( self.checkpoint, self.channels if do_checkpoint else None, self.step, @@ -946,8 +946,10 @@ class PregelLoop: else None, force_delta_snapshot=exiting and self.durability == "exit", updates_since_snapshot=new_counts, - new_updates_since_snapshot=new_counts, ) + self.checkpoint = result.checkpoint + for k in result.snapshotted: + new_counts[k] = 0 if new_counts: self.checkpoint_metadata["delta_updates_since_snapshot"] = new_counts elif "delta_updates_since_snapshot" in self.checkpoint_metadata: diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 1550e5a92..d96a167dc 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -1726,7 +1726,7 @@ class Pregel( # save checkpoint next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, step), + create_checkpoint(checkpoint, channels, step).checkpoint, { "source": "update", "step": step + 1, @@ -1765,7 +1765,7 @@ class Pregel( ) next_config = checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), + create_checkpoint(checkpoint, channels, next_step).checkpoint, { "source": "input", "step": next_step, @@ -1804,7 +1804,7 @@ class Pregel( if saved is None: raise InvalidUpdateError("Cannot copy a non-existent checkpoint") - next_checkpoint = create_checkpoint(checkpoint, None, step) + next_checkpoint = create_checkpoint(checkpoint, None, step).checkpoint # copy checkpoint next_config = checkpointer.put( @@ -2009,7 +2009,7 @@ class Pregel( checkpointer.get_next_version, self.trigger_to_nodes, ) - checkpoint = create_checkpoint(checkpoint, channels, step + 1) + checkpoint = create_checkpoint(checkpoint, channels, step + 1).checkpoint next_config = checkpointer.put( checkpoint_config, checkpoint, @@ -2175,7 +2175,7 @@ class Pregel( # save checkpoint next_config = await checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, channels, step), + create_checkpoint(checkpoint, channels, step).checkpoint, { "source": "update", "step": step + 1, @@ -2213,7 +2213,7 @@ class Pregel( ) next_config = await checkpointer.aput( checkpoint_config, - create_checkpoint(checkpoint, channels, next_step), + create_checkpoint(checkpoint, channels, next_step).checkpoint, { "source": "input", "step": next_step, @@ -2252,7 +2252,7 @@ class Pregel( if saved is None: raise InvalidUpdateError("Cannot copy a non-existent checkpoint") - next_checkpoint = create_checkpoint(checkpoint, None, step) + next_checkpoint = create_checkpoint(checkpoint, None, step).checkpoint # copy checkpoint next_config = await checkpointer.aput( @@ -2456,7 +2456,7 @@ class Pregel( checkpointer.get_next_version, self.trigger_to_nodes, ) - checkpoint = create_checkpoint(checkpoint, channels, step + 1) + checkpoint = create_checkpoint(checkpoint, channels, step + 1).checkpoint # save checkpoint, after applying writes next_config = await checkpointer.aput( checkpoint_config,