WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-10-01 05:55:14 +02:00
Code Issues Packages Projects Releases Wiki Activity

chore: minor clean up around checkpoint and delta channel (#7705)

Browse Source
This commit is contained in:
Quanzheng Long
2026-05-04 17:20:15 -07:00
committed by GitHub
parent 60e305596d
commit 0ae8870432
5 changed files with 213 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/base/__init__.py
+14 -7
View File
@@ -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,
)
libs/langgraph/langgraph/pregel/_algo.py
+91 -10
View File
@@ -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:
libs/langgraph/langgraph/pregel/_checkpoint.py
+96 -40
View File
@@ -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,
)
libs/langgraph/langgraph/pregel/_loop.py
+4 -2
View File
@@ -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:
libs/langgraph/langgraph/pregel/main.py
+8 -8
View File
@@ -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,