mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 19:29:43 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e857fede04 | ||
|
|
ecc420e7fe |
@@ -538,7 +538,10 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
||||
seeded: set[str] = set()
|
||||
|
||||
with self.cursor(transaction=False) as cur:
|
||||
cur.execute(DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id))
|
||||
cur.execute(
|
||||
DELTA_STAGE1_SQL,
|
||||
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
|
||||
)
|
||||
for row in cur:
|
||||
cid, parent_cid, type_tag, blob = row
|
||||
if step_walk_with_row(
|
||||
|
||||
@@ -26,16 +26,38 @@ from typing import Any
|
||||
|
||||
from langgraph.checkpoint.base import DeltaChannelHistory, PendingWrite
|
||||
|
||||
# Stage 1 streams ancestors of `target_cid` newest-first. The `<=`
|
||||
# predicate keeps target itself in the stream so we can read its
|
||||
# `parent_checkpoint_id` from the first row without a separate lookup;
|
||||
# the caller skips target's own writes/seed (matches the
|
||||
# `BaseCheckpointSaver` contract).
|
||||
# Stage 1 streams target followed by its ancestors, oldest step last, by
|
||||
# following `parent_checkpoint_id` in a recursive CTE. Target is the anchor
|
||||
# row, so the caller reads its `parent_checkpoint_id` without a separate
|
||||
# lookup and skips its own writes/seed (matches the `BaseCheckpointSaver`
|
||||
# contract).
|
||||
#
|
||||
# Ancestry is defined by `parent_checkpoint_id` alone. An earlier form
|
||||
# filtered `checkpoint_id <= target` and ordered by `checkpoint_id DESC`,
|
||||
# which additionally required every child's id to sort above its parent's.
|
||||
# Nothing in the contract promises that, and a parent sorting above its
|
||||
# child was dropped from the stream entirely, silently costing that
|
||||
# parent's seed and writes. See #8550.
|
||||
#
|
||||
# A parent chain can cycle: `put` writes with `INSERT OR REPLACE`, so
|
||||
# re-putting an existing checkpoint id under a descendant's config rewrites
|
||||
# its parent. The old id-range scan read a finite row set and could not
|
||||
# loop; this one can, so `step_walk_with_row` stops on a repeated
|
||||
# checkpoint_id. sqlite yields recursive rows lazily, so abandoning the
|
||||
# cursor ends the recursion rather than waiting on it.
|
||||
DELTA_STAGE1_SQL = (
|
||||
"WITH RECURSIVE ancestors(checkpoint_id, parent_checkpoint_id, type, "
|
||||
"checkpoint) AS ("
|
||||
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint "
|
||||
"FROM checkpoints "
|
||||
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id <= ? "
|
||||
"ORDER BY checkpoint_id DESC"
|
||||
"WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ? "
|
||||
"UNION ALL "
|
||||
"SELECT c.checkpoint_id, c.parent_checkpoint_id, c.type, c.checkpoint "
|
||||
"FROM checkpoints c JOIN ancestors a "
|
||||
"ON c.checkpoint_id = a.parent_checkpoint_id "
|
||||
"WHERE c.thread_id = ? AND c.checkpoint_ns = ?"
|
||||
") "
|
||||
"SELECT checkpoint_id, parent_checkpoint_id, type, checkpoint FROM ancestors"
|
||||
)
|
||||
|
||||
|
||||
@@ -95,14 +117,20 @@ def step_walk_with_row(
|
||||
Off-path rows (different branch on the same thread) advance the
|
||||
cursor without doing any work.
|
||||
|
||||
Returns True when every requested channel is seeded — the caller
|
||||
can stop iterating and close the cursor.
|
||||
Returns True when the caller can stop iterating and close the cursor,
|
||||
either because every requested channel is seeded or because the parent
|
||||
chain revisited a checkpoint it had already walked. `put` writes with
|
||||
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
|
||||
descendant's config points it at its own descendant and makes the chain
|
||||
a loop; without this check the recursive stage-1 query would feed rows
|
||||
forever.
|
||||
"""
|
||||
if "started" not in walk_state:
|
||||
if cid == target_id:
|
||||
walk_state["started"] = True
|
||||
walk_state["cur_cid"] = parent_cid
|
||||
walk_state["active"] = {ch for ch in channels if ch not in seeded}
|
||||
walk_state["walked"] = {cid}
|
||||
# Not target yet (or target not present): keep streaming.
|
||||
return False
|
||||
active: set[str] = walk_state["active"]
|
||||
@@ -111,6 +139,10 @@ def step_walk_with_row(
|
||||
if cid != walk_state["cur_cid"]:
|
||||
# Off-path row from a sibling branch — skip without deserializing.
|
||||
return False
|
||||
walked: set[str] = walk_state["walked"]
|
||||
if cid in walked:
|
||||
return True
|
||||
walked.add(cid)
|
||||
for ch in active:
|
||||
chain_by_ch[ch].append(cid)
|
||||
ckpt = serde.loads_typed((type_tag, blob))
|
||||
|
||||
@@ -650,7 +650,8 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
||||
|
||||
async with self.lock, self.conn.cursor() as cur:
|
||||
await cur.execute(
|
||||
DELTA_STAGE1_SQL, (thread_id, checkpoint_ns, checkpoint_id)
|
||||
DELTA_STAGE1_SQL,
|
||||
(thread_id, checkpoint_ns, checkpoint_id, thread_id, checkpoint_ns),
|
||||
)
|
||||
async for row in cur:
|
||||
cid, parent_cid, type_tag, blob = row
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""The stage-1 ancestor walk follows `parent_checkpoint_id`, not id order.
|
||||
|
||||
Ancestry is defined by the `parent_checkpoint_id` column. An earlier stage-1
|
||||
query filtered `checkpoint_id <= target` and streamed `ORDER BY checkpoint_id
|
||||
DESC`, so it also required every child's id to sort above its parent's. Real
|
||||
ids are `uuid6` and happen to satisfy that, but the contract never promised it,
|
||||
and a parent sorting above its child was dropped from the stream: its seed and
|
||||
its writes vanished with no error. See #8550.
|
||||
|
||||
Following parent pointers in a recursive CTE removes both assumptions, at the
|
||||
cost of needing a cycle guard, which the bounded id scan got for free.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import (
|
||||
BaseCheckpointSaver,
|
||||
Checkpoint,
|
||||
DeltaChannelHistory,
|
||||
empty_checkpoint,
|
||||
)
|
||||
|
||||
from langgraph.checkpoint.sqlite import SqliteSaver
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
CHANNEL = "ch"
|
||||
CONFIG: dict[str, Any] = {"configurable": {"thread_id": "t", "checkpoint_ns": ""}}
|
||||
EXPECTED: DeltaChannelHistory = {
|
||||
"writes": [("task", CHANNEL, "write-root")],
|
||||
"seed": "seed",
|
||||
}
|
||||
|
||||
|
||||
def _checkpoint(checkpoint_id: str, values: dict[str, Any]) -> Checkpoint:
|
||||
value = empty_checkpoint()
|
||||
value["id"] = checkpoint_id
|
||||
value["channel_values"] = values
|
||||
return value
|
||||
|
||||
|
||||
# `z` sorts above `a`, so the parent's id sorts above its child's. Real uuid6
|
||||
# ids never do this; nothing in the contract stops a caller supplying ids that
|
||||
# do, and clock skew between two processes writing one thread produces it.
|
||||
PARENT_ID_ORDERS = [
|
||||
pytest.param("z-older", "a-newer", id="parent_id_sorts_above_child"),
|
||||
pytest.param("a-older", "z-newer", id="parent_id_sorts_below_child"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
|
||||
def test_sync_walk_reaches_parent_whatever_the_id_order(
|
||||
root_id: str, child_id: str
|
||||
) -> None:
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = saver.put(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
|
||||
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
child = saver.put(root, _checkpoint(child_id, {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(config=child, channels=[CHANNEL])
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
# The unoptimised implementation is the contract; the fast path must
|
||||
# agree with it on the same rows.
|
||||
assert (
|
||||
got[CHANNEL]
|
||||
== BaseCheckpointSaver.get_delta_channel_history(
|
||||
saver, config=child, channels=[CHANNEL]
|
||||
)[CHANNEL]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("root_id", "child_id"), PARENT_ID_ORDERS)
|
||||
async def test_async_walk_reaches_parent_whatever_the_id_order(
|
||||
root_id: str, child_id: str
|
||||
) -> None:
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = await saver.aput(CONFIG, _checkpoint(root_id, {CHANNEL: "seed"}), {}, {})
|
||||
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
child = await saver.aput(root, _checkpoint(child_id, {}), {}, {})
|
||||
|
||||
got = await saver.aget_delta_channel_history(config=child, channels=[CHANNEL])
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
def test_long_chain_walks_the_whole_thread() -> None:
|
||||
"""A migrated thread can hold its only stored value at the root.
|
||||
|
||||
The walk then legitimately runs the length of the thread, so nothing in
|
||||
the cycle check may cut it short. Ids descend as the chain grows here, so
|
||||
id order fights the walk at every step.
|
||||
"""
|
||||
steps = 40
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
parent = saver.put(
|
||||
CONFIG, _checkpoint(f"id-{steps:03d}", {CHANNEL: "seed"}), {}, {}
|
||||
)
|
||||
saver.put_writes(parent, [(CHANNEL, "write-root")], "task")
|
||||
for step in range(steps - 1, 0, -1):
|
||||
parent = saver.put(parent, _checkpoint(f"id-{step:03d}", {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(config=parent, channels=[CHANNEL])
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
def test_cyclic_parent_chain_terminates() -> None:
|
||||
"""A parent chain that loops must not feed the walk forever.
|
||||
|
||||
Reachable through `put` alone, no corruption needed: it writes with
|
||||
`INSERT OR REPLACE`, so re-putting an existing checkpoint id under a
|
||||
descendant's config repoints that checkpoint at its own descendant. The
|
||||
old id-range scan read a finite row set and could not loop; a recursive
|
||||
parent-pointer query can, so the walk stops on a repeated id.
|
||||
|
||||
Note this one fails by hanging, not by asserting, since a regression
|
||||
means the row stream never ends. The package has no timeout plugin, so
|
||||
the CI job timeout is the backstop.
|
||||
"""
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
a = saver.put(CONFIG, _checkpoint("cid-a", {}), {}, {})
|
||||
b = saver.put(a, _checkpoint("cid-b", {}), {}, {})
|
||||
# Re-put "cid-a" with "cid-b" as its parent: a -> b -> a.
|
||||
saver.put(b, _checkpoint("cid-a", {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(config=b, channels=[CHANNEL])
|
||||
# Nothing on the cycle stores a value, so no seed. The point of the
|
||||
# test is that the call returns at all.
|
||||
assert "seed" not in got[CHANNEL]
|
||||
@@ -80,20 +80,12 @@ def get_updated_channels_from_tasks(
|
||||
|
||||
def get_delta_channels_from_all_channels(
|
||||
channels: Mapping[str, BaseChannel],
|
||||
*,
|
||||
include_unavailable: bool = False,
|
||||
) -> set[str]:
|
||||
"""Every available DeltaChannel.
|
||||
|
||||
The set to snapshot whenever no ancestor walk can reconstruct these
|
||||
channels: the first update_state of a fresh thread (no ancestors at all),
|
||||
and the first checkpoint of a fork (whose base also holds the writes of the
|
||||
branch the fork abandons).
|
||||
"""
|
||||
"""DeltaChannels to snapshot on the first update_state of a fresh thread."""
|
||||
return {
|
||||
k
|
||||
for k, ch in channels.items()
|
||||
if isinstance(ch, DeltaChannel) and (include_unavailable or ch.is_available())
|
||||
if isinstance(ch, DeltaChannel) and ch.is_available()
|
||||
}
|
||||
|
||||
|
||||
@@ -130,27 +122,15 @@ def create_checkpoint_plan_for_update_state_api(
|
||||
parents: dict[str, Any],
|
||||
saved_metadata: Mapping[str, Any] | None,
|
||||
is_fresh_thread: bool,
|
||||
is_fork: bool,
|
||||
) -> tuple[set[str], dict[str, Any]]:
|
||||
"""Return ``(channels_to_snapshot, metadata)`` for an update_state head.
|
||||
|
||||
``is_fork`` (the update was addressed at an explicit checkpoint) forces a
|
||||
full snapshot for the same reason ``is_fresh_thread`` does: the ancestor
|
||||
walk cannot reconstruct this head. The base a fork branches off keeps the
|
||||
pending writes of the branch being abandoned, and nothing records which
|
||||
child consumed which write, so the walk would replay them here too.
|
||||
Snapshotting terminates the walk at this checkpoint. Every delta channel
|
||||
snapshots, so no counters carry over.
|
||||
"""
|
||||
"""Return ``(channels_to_snapshot, metadata)`` for an update_state head."""
|
||||
metadata: dict[str, Any] = {
|
||||
"source": "update",
|
||||
"step": step,
|
||||
"parents": parents,
|
||||
}
|
||||
if is_fresh_thread or is_fork:
|
||||
return get_delta_channels_from_all_channels(
|
||||
channels, include_unavailable=is_fork
|
||||
), metadata
|
||||
if is_fresh_thread:
|
||||
return get_delta_channels_from_all_channels(channels), metadata
|
||||
|
||||
new_counters = create_metadata_for_update_state_api(
|
||||
channels,
|
||||
@@ -166,43 +146,6 @@ def create_checkpoint_plan_for_update_state_api(
|
||||
return channels_to_snapshot, metadata
|
||||
|
||||
|
||||
def create_fork_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, BaseChannel],
|
||||
step: int,
|
||||
*,
|
||||
is_fork: bool,
|
||||
get_next_version: GetNextVersion,
|
||||
) -> Checkpoint:
|
||||
"""``create_checkpoint`` for an update_state path that bypasses the plan.
|
||||
|
||||
The ``as_node`` INPUT and END paths write the fork's first checkpoint and
|
||||
return before ``create_checkpoint_plan_for_update_state_api`` runs. Left
|
||||
without a snapshot that checkpoint does not seal the fork, and the next
|
||||
superstep reconstructs its delta channels by walking through the shared
|
||||
base, picking up the abandoned branch's writes and then baking them into
|
||||
whatever it snapshots. Sealing has to happen on the fork's *first*
|
||||
checkpoint, which is this one.
|
||||
|
||||
``get_next_version`` is required for the same reason exit mode needs it:
|
||||
these paths apply writes to the input channel, not to the delta channel,
|
||||
so nothing bumps the delta channel's version and ``put`` would drop the
|
||||
blob as not-a-new-version. Callers must derive ``new_versions`` from the
|
||||
returned checkpoint rather than the one they passed in.
|
||||
"""
|
||||
if not is_fork:
|
||||
return create_checkpoint(checkpoint, channels, step)
|
||||
return create_checkpoint(
|
||||
checkpoint,
|
||||
channels,
|
||||
step,
|
||||
get_next_version=get_next_version,
|
||||
channels_to_snapshot=get_delta_channels_from_all_channels(
|
||||
channels, include_unavailable=True
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def create_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
channels: Mapping[str, BaseChannel] | None,
|
||||
@@ -231,27 +174,14 @@ def create_checkpoint(
|
||||
values = {}
|
||||
channel_versions = dict(checkpoint["channel_versions"])
|
||||
for k in channels:
|
||||
ch = channels[k]
|
||||
if k not in channel_versions:
|
||||
# Nothing was ever written to this channel on this branch, so
|
||||
# it has no version and `put` would drop any blob stored for
|
||||
# it. A *forced* snapshot still has to land: it is the only
|
||||
# thing that stops the ancestor walk running past this
|
||||
# checkpoint into a fork base that holds another branch's
|
||||
# writes. Mint a first version so the blob survives.
|
||||
if k in channels_to_snapshot and get_next_version is not None:
|
||||
channel_versions[k] = get_next_version(None, None)
|
||||
values[k] = _DeltaSnapshot(
|
||||
ch.get() if ch.is_available() else ch.typ()
|
||||
)
|
||||
continue
|
||||
ch = channels[k]
|
||||
if k in channels_to_snapshot:
|
||||
# Callers force a full snapshot blob here: exit mode when a
|
||||
# delta channel reaches its snapshot cadence, update_state on
|
||||
# a fresh thread (no ancestor to replay writes from), and a
|
||||
# fork (whose base also holds the abandoned branch's writes).
|
||||
# The manual version-bump below only applies to the exit-mode
|
||||
# case.
|
||||
# delta channel reaches its snapshot cadence, and update_state
|
||||
# on a fresh thread (no ancestor to replay writes from). The
|
||||
# manual version-bump below only applies to the exit-mode case.
|
||||
#
|
||||
# In exit mode, the snapshot decision is deferred to exit
|
||||
# time (intermediate steps have do_checkpoint=False). The
|
||||
|
||||
@@ -222,32 +222,10 @@ class PregelLoop:
|
||||
# under the saver's `ORDER BY task_id, idx` sorting.
|
||||
_exit_delta_writes: list[tuple[int, str, str, Any]] | None = None
|
||||
|
||||
# Delta channels that must snapshot at the next checkpoint, whatever their
|
||||
# cadence counters say. Two sources:
|
||||
# * an Overwrite arrived since the last checkpoint, so the snapshot has to
|
||||
# happen after live update applied overwrite semantics and sparse replay
|
||||
# starts from the same post-overwrite value;
|
||||
# * this run forked off an explicitly addressed checkpoint, see
|
||||
# `_delta_channels_awaiting_fork_snapshot`.
|
||||
_delta_channels_forced_snapshot: set[str]
|
||||
|
||||
# Delta channels still owed a fork snapshot, when this run was launched
|
||||
# against an explicitly addressed checkpoint (time travel / fork). That
|
||||
# base keeps the pending writes of the branch the fork abandons, and
|
||||
# nothing records which child consumed which write, so the ancestor walk
|
||||
# would replay them into this branch too. Snapshotting terminates the walk
|
||||
# inside the fork instead of at the shared base. Names drop out once the
|
||||
# blob has landed; a channel with no value yet has nothing to snapshot, so
|
||||
# it waits for the superstep that gives it one.
|
||||
#
|
||||
# The trigger is deliberately coarse: any addressed checkpoint, not only
|
||||
# one that turns out to have abandoned writes on it. Which writes belong to
|
||||
# which child is exactly what is not recorded, so a narrower test would
|
||||
# have to trust the base's `pending_writes` to be complete, and a saver
|
||||
# that leaves them out would silently go back to leaking. Snapshotting when
|
||||
# it was not needed costs one blob per addressed run; not snapshotting when
|
||||
# it was needed is silent corruption.
|
||||
_delta_channels_awaiting_fork_snapshot: set[str]
|
||||
# Delta channels that saw an Overwrite since the last checkpoint. These
|
||||
# channels must snapshot after live update applies overwrite semantics so
|
||||
# sparse replay starts from the same post-overwrite value.
|
||||
_delta_channels_with_overwrite: set[str]
|
||||
|
||||
# The checkpoint_config that points at the parent loaded at `__enter__`
|
||||
# (or the synthetic-empty checkpoint, on first run). We capture it
|
||||
@@ -391,16 +369,6 @@ class PregelLoop:
|
||||
if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
|
||||
else ()
|
||||
)
|
||||
# Checks the value, not just key presence like `is_replaying` above:
|
||||
# subgraph task configs always carry an explicit `None` here, and only
|
||||
# a real id means the caller addressed one specific checkpoint. Read
|
||||
# off `checkpoint_config` so subgraphs resolved through a checkpoint
|
||||
# map during time travel are covered too, matching `__enter__`.
|
||||
self._delta_channels_awaiting_fork_snapshot = (
|
||||
{k for k, spec in specs.items() if isinstance(spec, DeltaChannel)}
|
||||
if self.checkpoint_config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
|
||||
else set()
|
||||
)
|
||||
self.prev_checkpoint_config = None
|
||||
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
|
||||
self.control = runtime.control if isinstance(runtime, Runtime) else None
|
||||
@@ -715,7 +683,7 @@ class PregelLoop:
|
||||
def after_tick(self) -> None:
|
||||
# finish superstep
|
||||
writes = [w for t in self.tasks.values() for w in t.writes]
|
||||
self._delta_channels_forced_snapshot.update(
|
||||
self._delta_channels_with_overwrite.update(
|
||||
ch
|
||||
for ch, v in writes
|
||||
if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0]
|
||||
@@ -1023,7 +991,7 @@ class PregelLoop:
|
||||
manager=None,
|
||||
updated_channels=updated_channels,
|
||||
)
|
||||
self._delta_channels_forced_snapshot.update(
|
||||
self._delta_channels_with_overwrite.update(
|
||||
c
|
||||
for c, v in input_writes
|
||||
if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0]
|
||||
@@ -1165,21 +1133,10 @@ class PregelLoop:
|
||||
do_checkpoint = self._checkpointer_put_after_previous is not None and (
|
||||
exiting or self.durability != "exit"
|
||||
)
|
||||
# Fork: make this checkpoint self-contained, so the ancestor walk stops
|
||||
# inside the fork instead of reaching the base this run forked off and
|
||||
# collecting the abandoned branch's writes from it. Resolved here
|
||||
# rather than in `_first` so channels that only got a value this
|
||||
# superstep are covered too.
|
||||
if self._delta_channels_awaiting_fork_snapshot:
|
||||
self._delta_channels_forced_snapshot.update(
|
||||
k
|
||||
for k in self._delta_channels_awaiting_fork_snapshot
|
||||
if k in self.channels
|
||||
)
|
||||
# create new checkpoint
|
||||
channels_to_snapshot = (
|
||||
delta_channels_to_snapshot(self.channels, new_counters)
|
||||
| self._delta_channels_forced_snapshot
|
||||
| self._delta_channels_with_overwrite
|
||||
if do_checkpoint
|
||||
else set()
|
||||
)
|
||||
@@ -1197,13 +1154,7 @@ class PregelLoop:
|
||||
for k in channels_to_snapshot:
|
||||
new_counters[k] = (0, 0)
|
||||
if do_checkpoint:
|
||||
self._delta_channels_forced_snapshot.difference_update(channels_to_snapshot)
|
||||
# `create_checkpoint` drops a requested snapshot for a channel with
|
||||
# no version in this checkpoint yet (nothing was ever written to it
|
||||
# on this branch), so keep asking until the blob really landed.
|
||||
self._delta_channels_awaiting_fork_snapshot.difference_update(
|
||||
self.checkpoint["channel_values"]
|
||||
)
|
||||
self._delta_channels_with_overwrite.difference_update(channels_to_snapshot)
|
||||
non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)}
|
||||
if non_zero:
|
||||
self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero
|
||||
@@ -1288,7 +1239,7 @@ class PregelLoop:
|
||||
)
|
||||
channels_to_snapshot = (
|
||||
delta_channels_to_snapshot(self.channels, counters)
|
||||
| self._delta_channels_forced_snapshot
|
||||
| self._delta_channels_with_overwrite
|
||||
)
|
||||
|
||||
pending = [
|
||||
@@ -1733,7 +1684,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
)
|
||||
self._delta_write_futs = []
|
||||
self._error_handler_write_futs = []
|
||||
self._delta_channels_forced_snapshot = set()
|
||||
self._delta_channels_with_overwrite = set()
|
||||
self._exit_delta_writes = (
|
||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||
)
|
||||
@@ -1991,7 +1942,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
)
|
||||
self._delta_write_futs = []
|
||||
self._error_handler_write_futs = []
|
||||
self._delta_channels_forced_snapshot = set()
|
||||
self._delta_channels_with_overwrite = set()
|
||||
self._exit_delta_writes = (
|
||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||
)
|
||||
|
||||
@@ -108,7 +108,6 @@ from langgraph.callbacks import (
|
||||
get_sync_graph_callback_manager_for_config,
|
||||
)
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.topic import Topic
|
||||
from langgraph.config import get_config
|
||||
from langgraph.constants import END
|
||||
@@ -134,7 +133,6 @@ from langgraph.pregel._checkpoint import (
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
create_checkpoint_plan_for_update_state_api,
|
||||
create_fork_checkpoint,
|
||||
empty_checkpoint,
|
||||
get_updated_channels_from_tasks,
|
||||
)
|
||||
@@ -1639,24 +1637,8 @@ class Pregel(
|
||||
else:
|
||||
raise ValueError(f"Subgraph {recast} not found")
|
||||
|
||||
# Delta channels still owed a fork snapshot. Mirrors
|
||||
# `_delta_channels_awaiting_fork_snapshot` in `_loop.py`: a fork is only
|
||||
# sealed once a checkpoint actually carries the blob. Several superstep
|
||||
# paths (`as_node` of INPUT, END or `__copy__`) write a checkpoint and
|
||||
# return before reaching the plan, so a flag cleared after the first
|
||||
# superstep would leave the branch unsealed and the next superstep would
|
||||
# reconstruct through the shared base.
|
||||
fork_pending: set[str] = (
|
||||
{k for k, v in self.channels.items() if isinstance(v, DeltaChannel)}
|
||||
if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
|
||||
else set()
|
||||
)
|
||||
|
||||
def perform_superstep(
|
||||
input_config: RunnableConfig,
|
||||
updates: Sequence[StateUpdate],
|
||||
*,
|
||||
is_fork: bool,
|
||||
input_config: RunnableConfig, updates: Sequence[StateUpdate]
|
||||
) -> RunnableConfig:
|
||||
# get last checkpoint
|
||||
config = ensure_config(self.config, input_config)
|
||||
@@ -1744,17 +1726,9 @@ class Pregel(
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
# save checkpoint
|
||||
next_checkpoint = create_fork_checkpoint(
|
||||
checkpoint,
|
||||
channels,
|
||||
step,
|
||||
is_fork=is_fork,
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
)
|
||||
fork_pending.difference_update(next_checkpoint["channel_values"])
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
next_checkpoint,
|
||||
create_checkpoint(checkpoint, channels, step),
|
||||
{
|
||||
"source": "update",
|
||||
"step": step + 1,
|
||||
@@ -1762,7 +1736,7 @@ class Pregel(
|
||||
},
|
||||
get_new_channel_versions(
|
||||
checkpoint_previous_versions,
|
||||
next_checkpoint["channel_versions"],
|
||||
checkpoint["channel_versions"],
|
||||
),
|
||||
)
|
||||
return patch_checkpoint_map(
|
||||
@@ -1791,17 +1765,9 @@ class Pregel(
|
||||
if saved and saved.metadata.get("step") is not None
|
||||
else -1
|
||||
)
|
||||
next_checkpoint = create_fork_checkpoint(
|
||||
checkpoint,
|
||||
channels,
|
||||
next_step,
|
||||
is_fork=is_fork,
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
)
|
||||
fork_pending.difference_update(next_checkpoint["channel_values"])
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
next_checkpoint,
|
||||
create_checkpoint(checkpoint, channels, next_step),
|
||||
{
|
||||
"source": "input",
|
||||
"step": next_step,
|
||||
@@ -1811,7 +1777,7 @@ class Pregel(
|
||||
},
|
||||
get_new_channel_versions(
|
||||
checkpoint_previous_versions,
|
||||
next_checkpoint["channel_versions"],
|
||||
checkpoint["channel_versions"],
|
||||
),
|
||||
)
|
||||
|
||||
@@ -1907,9 +1873,6 @@ class Pregel(
|
||||
return perform_superstep(
|
||||
patch_checkpoint_map(next_config, saved.metadata),
|
||||
[item for lst in user_group_by.values() for item in lst],
|
||||
# The checkpoint just written clears tasks and carries
|
||||
# no delta snapshot, so a fork is still unsealed here.
|
||||
is_fork=is_fork,
|
||||
)
|
||||
|
||||
return patch_checkpoint_map(next_config, saved.metadata)
|
||||
@@ -2057,7 +2020,6 @@ class Pregel(
|
||||
parents=saved.metadata.get("parents", {}) if saved else {},
|
||||
saved_metadata=saved.metadata if saved else None,
|
||||
is_fresh_thread=saved is None,
|
||||
is_fork=is_fork,
|
||||
)
|
||||
)
|
||||
checkpoint = create_checkpoint(
|
||||
@@ -2070,8 +2032,6 @@ class Pregel(
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
)
|
||||
if is_fork:
|
||||
fork_pending.difference_update(checkpoint["channel_values"])
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
checkpoint,
|
||||
@@ -2089,14 +2049,8 @@ class Pregel(
|
||||
current_config = patch_configurable(
|
||||
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
|
||||
)
|
||||
# The flag cannot be derived from `current_config`: `perform_superstep`
|
||||
# returns the config of the checkpoint it just wrote and the loop feeds
|
||||
# that back in, so from the second superstep on it always names a
|
||||
# checkpoint whether or not the caller addressed one.
|
||||
for superstep in supersteps:
|
||||
current_config = perform_superstep(
|
||||
current_config, superstep, is_fork=bool(fork_pending)
|
||||
)
|
||||
current_config = perform_superstep(current_config, superstep)
|
||||
return current_config
|
||||
|
||||
async def abulk_update_state(
|
||||
@@ -2149,24 +2103,8 @@ class Pregel(
|
||||
else:
|
||||
raise ValueError(f"Subgraph {recast} not found")
|
||||
|
||||
# Delta channels still owed a fork snapshot. Mirrors
|
||||
# `_delta_channels_awaiting_fork_snapshot` in `_loop.py`: a fork is only
|
||||
# sealed once a checkpoint actually carries the blob. Several superstep
|
||||
# paths (`as_node` of INPUT, END or `__copy__`) write a checkpoint and
|
||||
# return before reaching the plan, so a flag cleared after the first
|
||||
# superstep would leave the branch unsealed and the next superstep would
|
||||
# reconstruct through the shared base.
|
||||
fork_pending: set[str] = (
|
||||
{k for k, v in self.channels.items() if isinstance(v, DeltaChannel)}
|
||||
if config[CONF].get(CONFIG_KEY_CHECKPOINT_ID)
|
||||
else set()
|
||||
)
|
||||
|
||||
async def aperform_superstep(
|
||||
input_config: RunnableConfig,
|
||||
updates: Sequence[StateUpdate],
|
||||
*,
|
||||
is_fork: bool,
|
||||
input_config: RunnableConfig, updates: Sequence[StateUpdate]
|
||||
) -> RunnableConfig:
|
||||
# get last checkpoint
|
||||
config = ensure_config(self.config, input_config)
|
||||
@@ -2252,25 +2190,16 @@ class Pregel(
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
# save checkpoint
|
||||
next_checkpoint = create_fork_checkpoint(
|
||||
checkpoint,
|
||||
channels,
|
||||
step,
|
||||
is_fork=is_fork,
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
)
|
||||
fork_pending.difference_update(next_checkpoint["channel_values"])
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
next_checkpoint,
|
||||
create_checkpoint(checkpoint, channels, step),
|
||||
{
|
||||
"source": "update",
|
||||
"step": step + 1,
|
||||
"parents": saved.metadata.get("parents", {}) if saved else {},
|
||||
},
|
||||
get_new_channel_versions(
|
||||
checkpoint_previous_versions,
|
||||
next_checkpoint["channel_versions"],
|
||||
checkpoint_previous_versions, checkpoint["channel_versions"]
|
||||
),
|
||||
)
|
||||
return patch_checkpoint_map(
|
||||
@@ -2299,17 +2228,9 @@ class Pregel(
|
||||
if saved and saved.metadata.get("step") is not None
|
||||
else -1
|
||||
)
|
||||
next_checkpoint = create_fork_checkpoint(
|
||||
checkpoint,
|
||||
channels,
|
||||
next_step,
|
||||
is_fork=is_fork,
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
)
|
||||
fork_pending.difference_update(next_checkpoint["channel_values"])
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
next_checkpoint,
|
||||
create_checkpoint(checkpoint, channels, next_step),
|
||||
{
|
||||
"source": "input",
|
||||
"step": next_step,
|
||||
@@ -2319,7 +2240,7 @@ class Pregel(
|
||||
},
|
||||
get_new_channel_versions(
|
||||
checkpoint_previous_versions,
|
||||
next_checkpoint["channel_versions"],
|
||||
checkpoint["channel_versions"],
|
||||
),
|
||||
)
|
||||
|
||||
@@ -2414,9 +2335,6 @@ class Pregel(
|
||||
return await aperform_superstep(
|
||||
patch_checkpoint_map(next_config, saved.metadata),
|
||||
[item for lst in user_group_by.values() for item in lst],
|
||||
# The checkpoint just written clears tasks and carries
|
||||
# no delta snapshot, so a fork is still unsealed here.
|
||||
is_fork=is_fork,
|
||||
)
|
||||
|
||||
return patch_checkpoint_map(
|
||||
@@ -2562,7 +2480,6 @@ class Pregel(
|
||||
parents=saved.metadata.get("parents", {}) if saved else {},
|
||||
saved_metadata=saved.metadata if saved else None,
|
||||
is_fresh_thread=saved is None,
|
||||
is_fork=is_fork,
|
||||
)
|
||||
)
|
||||
checkpoint = create_checkpoint(
|
||||
@@ -2575,8 +2492,6 @@ class Pregel(
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
)
|
||||
if is_fork:
|
||||
fork_pending.difference_update(checkpoint["channel_values"])
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
checkpoint,
|
||||
@@ -2593,14 +2508,8 @@ class Pregel(
|
||||
current_config = patch_configurable(
|
||||
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
|
||||
)
|
||||
# The flag cannot be derived from `current_config`: `aperform_superstep`
|
||||
# returns the config of the checkpoint it just wrote and the loop feeds
|
||||
# that back in, so from the second superstep on it always names a
|
||||
# checkpoint whether or not the caller addressed one.
|
||||
for superstep in supersteps:
|
||||
current_config = await aperform_superstep(
|
||||
current_config, superstep, is_fork=bool(fork_pending)
|
||||
)
|
||||
current_config = await aperform_superstep(current_config, superstep)
|
||||
return current_config
|
||||
|
||||
def update_state(
|
||||
|
||||
@@ -85,19 +85,11 @@ class MemorySaverAssertImmutable(InMemorySaver):
|
||||
)
|
||||
== saved
|
||||
), config["configurable"]["checkpoint_ns"]
|
||||
# call super to write checkpoint
|
||||
next_config = super().put(config, checkpoint, metadata, new_versions)
|
||||
# Record the checkpoint as the saver reads it back, not the object it
|
||||
# was handed. `channel_values` are stored per (channel, version), so a
|
||||
# channel a step did not write is refilled from the blob its inherited
|
||||
# version still points at. A `DeltaChannel` omits its value except at a
|
||||
# snapshot, which makes the two representations differ for reasons that
|
||||
# are not mutation. Comparing read-back against read-back still catches
|
||||
# a checkpoint whose stored data actually changed.
|
||||
self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = (
|
||||
self.serde.dumps_typed(super().get(next_config))
|
||||
self.serde.dumps_typed(checkpoint)
|
||||
)
|
||||
return next_config
|
||||
# call super to write checkpoint
|
||||
return super().put(config, checkpoint, metadata, new_versions)
|
||||
|
||||
|
||||
class MemorySaverNoPending(InMemorySaver):
|
||||
|
||||
@@ -1,414 +0,0 @@
|
||||
"""Forking a thread must not replay the abandoned branch into the fork.
|
||||
|
||||
Regression suite for #8443. Addressing an older checkpoint creates a fork: the
|
||||
shared base ends up with two children, and it keeps the ``checkpoint_writes``
|
||||
of the branch the fork abandons. Nothing in the stored data records which child
|
||||
consumed which write, so the ``DeltaChannel`` ancestor walk used to collect the
|
||||
abandoned branch's writes as well.
|
||||
|
||||
Every graph here carries a ``DeltaChannel`` and a plain reducer channel fed the
|
||||
same values. ``full`` channels store complete ``channel_values`` and need no
|
||||
replay, so the plain channel is the oracle: after a fork the two must agree.
|
||||
|
||||
Coverage: fork by ``invoke`` with new input (sync/async, all durabilities),
|
||||
fork off the checkpoint that predates the thread's first input (sync/async),
|
||||
fork by ``update_state`` / ``aupdate_state``, fork before the delta channel
|
||||
ever had a value, and guards that neither an
|
||||
unaddressed run nor an unaddressed multi-superstep ``bulk_update_state``
|
||||
departs from the normal ``snapshot_frequency`` cadence.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from operator import add
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph._internal._constants import INPUT
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.types import Durability, StateSnapshot, StateUpdate
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _append(current: list | None, writes: Sequence[Any]) -> list:
|
||||
"""DeltaChannel reducer: extend the list with every batched write."""
|
||||
out = list(current or [])
|
||||
for write in writes:
|
||||
out.extend(write if isinstance(write, list) else [write])
|
||||
return out
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
log: Annotated[list, DeltaChannel(_append, snapshot_frequency=1000)]
|
||||
plain: Annotated[list, add]
|
||||
# Written only by the tests that fork before `log` ever has a value, so the
|
||||
# fork advances without touching the delta channel.
|
||||
other: Annotated[list, add]
|
||||
|
||||
|
||||
def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any:
|
||||
"""Compile a one-node graph whose node appends ``{tag}-out`` to both channels.
|
||||
|
||||
``snapshot_frequency=1000`` keeps the cadence from masking the bug: without
|
||||
a forced snapshot the fork's ancestor walk always runs past the fork base.
|
||||
"""
|
||||
|
||||
def node(state: _State) -> dict:
|
||||
return {"log": [f"{tag}-out"], "plain": [f"{tag}-out"]}
|
||||
|
||||
builder = StateGraph(_State)
|
||||
builder.add_node("n", node)
|
||||
builder.set_entry_point("n")
|
||||
builder.set_finish_point("n")
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _build_without_delta_writes(checkpointer: BaseCheckpointSaver, tag: str) -> Any:
|
||||
"""Compile a graph whose node writes only ``other``, never the delta channel."""
|
||||
|
||||
def node(state: _State) -> dict:
|
||||
return {"other": [f"{tag}-other"]}
|
||||
|
||||
builder = StateGraph(_State)
|
||||
builder.add_node("n", node)
|
||||
builder.set_entry_point("n")
|
||||
builder.set_finish_point("n")
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
def _thread(thread_id: str) -> RunnableConfig:
|
||||
return {"configurable": {"thread_id": thread_id}}
|
||||
|
||||
|
||||
def _at(config: RunnableConfig, snapshot: StateSnapshot) -> RunnableConfig:
|
||||
"""Config addressing one specific checkpoint of ``config``'s thread."""
|
||||
return {
|
||||
"configurable": {
|
||||
**config["configurable"],
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": snapshot.config["configurable"]["checkpoint_id"],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _input(marker: str) -> dict:
|
||||
return {"log": [marker], "plain": [marker]}
|
||||
|
||||
|
||||
def _snapshotted_checkpoints(
|
||||
checkpointer: BaseCheckpointSaver, config: RunnableConfig
|
||||
) -> list[str]:
|
||||
"""Ids of this thread's checkpoints carrying a ``log`` snapshot blob."""
|
||||
return [
|
||||
tuple_.config["configurable"]["checkpoint_id"]
|
||||
for tuple_ in checkpointer.list(config)
|
||||
if isinstance(tuple_.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
|
||||
]
|
||||
|
||||
|
||||
def _assert_fork_is_clean(state: StateSnapshot, abandoned: str) -> None:
|
||||
"""The delta channel must match the plain channel and drop ``abandoned``."""
|
||||
assert state.values["log"] == state.values["plain"], (
|
||||
f"delta channel diverged from the plain channel: "
|
||||
f"{state.values['log']} != {state.values['plain']}"
|
||||
)
|
||||
assert abandoned not in state.values["log"], (
|
||||
f"{abandoned!r} belongs to the branch the fork replaced, "
|
||||
f"but was replayed into {state.values['log']}"
|
||||
)
|
||||
|
||||
|
||||
def test_fork_by_invoke(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
config = _thread("t")
|
||||
_build(sync_checkpointer, "first").invoke(
|
||||
_input("in-1"), config, durability=durability
|
||||
)
|
||||
graph = _build(sync_checkpointer, "second")
|
||||
graph.invoke(_input("in-2"), config, durability=durability)
|
||||
|
||||
# The last checkpoint that predates "in-2" entering state: forking here
|
||||
# abandons the "in-2" branch, whose writes still hang off this checkpoint.
|
||||
base = next(
|
||||
snapshot
|
||||
for snapshot in graph.get_state_history(config)
|
||||
if "in-2" not in snapshot.values["log"]
|
||||
)
|
||||
_build(sync_checkpointer, "third").invoke(
|
||||
_input("in-3"), _at(config, base), durability=durability
|
||||
)
|
||||
|
||||
state = graph.get_state(config)
|
||||
_assert_fork_is_clean(state, "in-2")
|
||||
assert state.values["log"] == [*base.values["log"], "in-3", "third-out"]
|
||||
|
||||
|
||||
async def test_afork_by_invoke(
|
||||
async_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
config = _thread("t")
|
||||
await _build(async_checkpointer, "first").ainvoke(
|
||||
_input("in-1"), config, durability=durability
|
||||
)
|
||||
graph = _build(async_checkpointer, "second")
|
||||
await graph.ainvoke(_input("in-2"), config, durability=durability)
|
||||
|
||||
base = await anext(
|
||||
snapshot
|
||||
async for snapshot in graph.aget_state_history(config)
|
||||
if "in-2" not in snapshot.values["log"]
|
||||
)
|
||||
await _build(async_checkpointer, "third").ainvoke(
|
||||
_input("in-3"), _at(config, base), durability=durability
|
||||
)
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
_assert_fork_is_clean(state, "in-2")
|
||||
assert state.values["log"] == [*base.values["log"], "in-3", "third-out"]
|
||||
|
||||
|
||||
def test_fork_off_checkpoint_before_first_input(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
"""Fork off the root checkpoint, which predates any value for ``log``.
|
||||
|
||||
``create_checkpoint`` drops a requested snapshot for a channel absent from
|
||||
``channel_versions``, so the fork's own first checkpoint cannot carry the
|
||||
blob. The request has to stay queued until a superstep gives the channel a
|
||||
value, otherwise the root's ``in-1`` write still leaks into the fork.
|
||||
"""
|
||||
config = _thread("t")
|
||||
graph = _build(sync_checkpointer, "first")
|
||||
graph.invoke(_input("in-1"), config, durability=durability)
|
||||
|
||||
root = list(graph.get_state_history(config))[-1]
|
||||
assert root.values["log"] == []
|
||||
|
||||
_build(sync_checkpointer, "third").invoke(
|
||||
_input("in-9"), _at(config, root), durability=durability
|
||||
)
|
||||
|
||||
state = graph.get_state(config)
|
||||
_assert_fork_is_clean(state, "in-1")
|
||||
assert state.values["log"] == ["in-9", "third-out"]
|
||||
|
||||
|
||||
async def test_afork_off_checkpoint_before_first_input(
|
||||
async_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
"""Async twin of ``test_fork_off_checkpoint_before_first_input``."""
|
||||
config = _thread("t")
|
||||
graph = _build(async_checkpointer, "first")
|
||||
await graph.ainvoke(_input("in-1"), config, durability=durability)
|
||||
|
||||
root = [snapshot async for snapshot in graph.aget_state_history(config)][-1]
|
||||
assert root.values["log"] == []
|
||||
|
||||
await _build(async_checkpointer, "third").ainvoke(
|
||||
_input("in-9"), _at(config, root), durability=durability
|
||||
)
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
_assert_fork_is_clean(state, "in-1")
|
||||
assert state.values["log"] == ["in-9", "third-out"]
|
||||
|
||||
|
||||
def test_fork_by_update_state(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
config = _thread("t")
|
||||
_build(sync_checkpointer, "first").invoke(_input("in-1"), config)
|
||||
graph = _build(sync_checkpointer, "second")
|
||||
graph.invoke(_input("in-2"), config)
|
||||
|
||||
base = next(
|
||||
snapshot
|
||||
for snapshot in graph.get_state_history(config)
|
||||
if "in-2" not in snapshot.values["log"]
|
||||
)
|
||||
forked = graph.update_state(_at(config, base), _input("patched"))
|
||||
|
||||
state = graph.get_state(forked)
|
||||
_assert_fork_is_clean(state, "in-2")
|
||||
assert state.values["log"] == [*base.values["log"], "patched"]
|
||||
|
||||
|
||||
async def test_afork_by_update_state(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
config = _thread("t")
|
||||
await _build(async_checkpointer, "first").ainvoke(_input("in-1"), config)
|
||||
graph = _build(async_checkpointer, "second")
|
||||
await graph.ainvoke(_input("in-2"), config)
|
||||
|
||||
base = await anext(
|
||||
snapshot
|
||||
async for snapshot in graph.aget_state_history(config)
|
||||
if "in-2" not in snapshot.values["log"]
|
||||
)
|
||||
forked = await graph.aupdate_state(_at(config, base), _input("patched"))
|
||||
|
||||
state = await graph.aget_state(forked)
|
||||
_assert_fork_is_clean(state, "in-2")
|
||||
assert state.values["log"] == [*base.values["log"], "patched"]
|
||||
|
||||
|
||||
def test_unaddressed_run_keeps_snapshot_cadence(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
"""A run with no explicitly addressed checkpoint writes no snapshot blob.
|
||||
|
||||
Guards the cost of the fix: the forced snapshot is one per addressed run,
|
||||
not a change to the normal ``snapshot_frequency`` cadence.
|
||||
"""
|
||||
config = _thread("t")
|
||||
graph = _build(sync_checkpointer, "first")
|
||||
graph.invoke(_input("in-1"), config, durability=durability)
|
||||
graph.invoke(_input("in-2"), config, durability=durability)
|
||||
|
||||
assert not _snapshotted_checkpoints(sync_checkpointer, config)
|
||||
|
||||
|
||||
def test_fork_before_first_value_when_fork_never_writes_the_channel(
|
||||
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
"""Seal the fork even when the channel has no value to snapshot.
|
||||
|
||||
Forking before ``log`` was ever written leaves nothing to copy into the
|
||||
fork's first checkpoint, so without a minted version and an empty blob the
|
||||
boundary goes unrecorded and the walk runs into the base. The fork here
|
||||
never writes ``log`` at all, so no later superstep can seal it either.
|
||||
"""
|
||||
config = _thread("t")
|
||||
graph = _build(sync_checkpointer, "first")
|
||||
graph.invoke(_input("in-1"), config, durability=durability)
|
||||
|
||||
root = list(graph.get_state_history(config))[-1]
|
||||
assert root.values["log"] == []
|
||||
|
||||
_build_without_delta_writes(sync_checkpointer, "third").invoke(
|
||||
{"other": ["in-9"]}, _at(config, root), durability=durability
|
||||
)
|
||||
|
||||
state = graph.get_state(config)
|
||||
_assert_fork_is_clean(state, "in-1")
|
||||
assert state.values["log"] == []
|
||||
|
||||
|
||||
async def test_afork_before_first_value_when_fork_never_writes_the_channel(
|
||||
async_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||
) -> None:
|
||||
"""Async twin of ``test_fork_before_first_value_when_fork_never_writes_the_channel``."""
|
||||
config = _thread("t")
|
||||
graph = _build(async_checkpointer, "first")
|
||||
await graph.ainvoke(_input("in-1"), config, durability=durability)
|
||||
|
||||
root = [snapshot async for snapshot in graph.aget_state_history(config)][-1]
|
||||
assert root.values["log"] == []
|
||||
|
||||
await _build_without_delta_writes(async_checkpointer, "third").ainvoke(
|
||||
{"other": ["in-9"]}, _at(config, root), durability=durability
|
||||
)
|
||||
|
||||
state = await graph.aget_state(config)
|
||||
_assert_fork_is_clean(state, "in-1")
|
||||
assert state.values["log"] == []
|
||||
|
||||
|
||||
def test_fork_before_first_value_by_bulk_update(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""The fork's first superstep touches another key, the delta key comes later.
|
||||
|
||||
The first checkpoint has to seal the boundary on its own. Deferring until
|
||||
the superstep that finally writes ``log`` is too late, because that
|
||||
superstep reconstructs through the unsealed checkpoint first.
|
||||
"""
|
||||
config = _thread("t")
|
||||
graph = _build(sync_checkpointer, "first")
|
||||
graph.invoke(_input("in-1"), config)
|
||||
|
||||
root = list(graph.get_state_history(config))[-1]
|
||||
assert root.values["log"] == []
|
||||
|
||||
forked = graph.bulk_update_state(
|
||||
_at(config, root),
|
||||
[
|
||||
[StateUpdate({"other": ["s1"]}, "n")],
|
||||
[StateUpdate(_input("s2"), "n")],
|
||||
],
|
||||
)
|
||||
|
||||
state = graph.get_state(forked)
|
||||
_assert_fork_is_clean(state, "in-1")
|
||||
assert state.values["log"] == ["s2"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("first_as_node", [INPUT, END, "__copy__"])
|
||||
def test_fork_by_bulk_update_whose_first_superstep_skips_the_plan(
|
||||
sync_checkpointer: BaseCheckpointSaver, first_as_node: str
|
||||
) -> None:
|
||||
"""The fork must be sealed by whichever checkpoint the fork writes first.
|
||||
|
||||
``as_node`` of INPUT, END or ``__copy__`` writes a checkpoint and returns
|
||||
before ``create_checkpoint_plan_for_update_state_api`` runs. If that
|
||||
checkpoint carries no snapshot the branch is still unsealed, so the next
|
||||
superstep reconstructs through the shared base, picks up the abandoned
|
||||
writes, and bakes them into whatever it snapshots. Sealing later is too
|
||||
late: by then the in-memory value is already wrong.
|
||||
"""
|
||||
config = _thread("t")
|
||||
_build(sync_checkpointer, "first").invoke(_input("in-1"), config)
|
||||
graph = _build(sync_checkpointer, "second")
|
||||
graph.invoke(_input("in-2"), config)
|
||||
|
||||
base = next(
|
||||
snapshot
|
||||
for snapshot in graph.get_state_history(config)
|
||||
if "in-2" not in snapshot.values["log"]
|
||||
)
|
||||
first = (
|
||||
StateUpdate(_input("first-step"), first_as_node)
|
||||
if first_as_node == INPUT
|
||||
else StateUpdate(None, first_as_node)
|
||||
)
|
||||
forked = graph.bulk_update_state(
|
||||
_at(config, base),
|
||||
[[first], [StateUpdate(_input("second-step"), "n")]],
|
||||
)
|
||||
|
||||
state = graph.get_state(forked)
|
||||
# END legitimately absorbs the base's already-run task writes, so "in-2"
|
||||
# belongs there; the plain channel is the oracle for which is which.
|
||||
assert state.values["log"] == state.values["plain"], (
|
||||
f"delta channel diverged from the plain channel: "
|
||||
f"{state.values['log']} != {state.values['plain']}"
|
||||
)
|
||||
|
||||
|
||||
def test_unaddressed_bulk_update_keeps_snapshot_cadence(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
) -> None:
|
||||
"""A multi-superstep ``bulk_update_state`` forks at most once, at the head.
|
||||
|
||||
``perform_superstep`` returns the config of the checkpoint it just wrote
|
||||
and the driver feeds that back in, so every superstep after the first
|
||||
receives a config naming a checkpoint even when the caller addressed none.
|
||||
Deriving the fork flag from that config snapshots the whole growing value
|
||||
once per superstep.
|
||||
"""
|
||||
config = _thread("t")
|
||||
graph = _build(sync_checkpointer, "first")
|
||||
graph.invoke(_input("in-1"), config)
|
||||
|
||||
graph.bulk_update_state(
|
||||
config,
|
||||
[[StateUpdate(_input(f"u{i}"), "n")] for i in range(4)],
|
||||
)
|
||||
|
||||
assert not _snapshotted_checkpoints(sync_checkpointer, config)
|
||||
Reference in New Issue
Block a user