mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-25 11:05:12 +02:00
Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3be96b7d73 | ||
|
|
302ce795a5 | ||
|
|
0552703e21 | ||
|
|
c4df893d54 | ||
|
|
02c91dccf7 | ||
|
|
dffb68192d |
@@ -80,12 +80,14 @@ def get_updated_channels_from_tasks(
|
|||||||
|
|
||||||
def get_delta_channels_from_all_channels(
|
def get_delta_channels_from_all_channels(
|
||||||
channels: Mapping[str, BaseChannel],
|
channels: Mapping[str, BaseChannel],
|
||||||
|
*,
|
||||||
|
include_unavailable: bool = False,
|
||||||
) -> set[str]:
|
) -> set[str]:
|
||||||
"""DeltaChannels to snapshot on the first update_state of a fresh thread."""
|
"""DeltaChannels to snapshot on the first update_state of a fresh thread or fork."""
|
||||||
return {
|
return {
|
||||||
k
|
k
|
||||||
for k, ch in channels.items()
|
for k, ch in channels.items()
|
||||||
if isinstance(ch, DeltaChannel) and ch.is_available()
|
if isinstance(ch, DeltaChannel) and (include_unavailable or ch.is_available())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
@@ -122,15 +124,22 @@ def create_checkpoint_plan_for_update_state_api(
|
|||||||
parents: dict[str, Any],
|
parents: dict[str, Any],
|
||||||
saved_metadata: Mapping[str, Any] | None,
|
saved_metadata: Mapping[str, Any] | None,
|
||||||
is_fresh_thread: bool,
|
is_fresh_thread: bool,
|
||||||
|
is_fork: bool,
|
||||||
) -> tuple[set[str], dict[str, Any]]:
|
) -> tuple[set[str], dict[str, Any]]:
|
||||||
"""Return ``(channels_to_snapshot, metadata)`` for an update_state head."""
|
"""Return ``(channels_to_snapshot, metadata)`` for an update_state head.
|
||||||
|
|
||||||
|
A fork snapshots everything, like a fresh thread: its base also holds the
|
||||||
|
writes of the branch it abandons, so the ancestor walk must stop here.
|
||||||
|
"""
|
||||||
metadata: dict[str, Any] = {
|
metadata: dict[str, Any] = {
|
||||||
"source": "update",
|
"source": "update",
|
||||||
"step": step,
|
"step": step,
|
||||||
"parents": parents,
|
"parents": parents,
|
||||||
}
|
}
|
||||||
if is_fresh_thread:
|
if is_fresh_thread or is_fork:
|
||||||
return get_delta_channels_from_all_channels(channels), metadata
|
return get_delta_channels_from_all_channels(
|
||||||
|
channels, include_unavailable=is_fork
|
||||||
|
), metadata
|
||||||
|
|
||||||
new_counters = create_metadata_for_update_state_api(
|
new_counters = create_metadata_for_update_state_api(
|
||||||
channels,
|
channels,
|
||||||
@@ -146,6 +155,34 @@ def create_checkpoint_plan_for_update_state_api(
|
|||||||
return channels_to_snapshot, metadata
|
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 the update_state paths that skip the plan.
|
||||||
|
|
||||||
|
The fork has to be sealed by its first checkpoint: any later superstep
|
||||||
|
has already rebuilt its delta channels through the shared base. These
|
||||||
|
paths never write the delta channel, so its version must be bumped here
|
||||||
|
or ``put`` drops the blob; derive ``new_versions`` from the result.
|
||||||
|
"""
|
||||||
|
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(
|
def create_checkpoint(
|
||||||
checkpoint: Checkpoint,
|
checkpoint: Checkpoint,
|
||||||
channels: Mapping[str, BaseChannel] | None,
|
channels: Mapping[str, BaseChannel] | None,
|
||||||
@@ -174,14 +211,23 @@ def create_checkpoint(
|
|||||||
values = {}
|
values = {}
|
||||||
channel_versions = dict(checkpoint["channel_versions"])
|
channel_versions = dict(checkpoint["channel_versions"])
|
||||||
for k in channels:
|
for k in channels:
|
||||||
if k not in channel_versions:
|
|
||||||
continue
|
|
||||||
ch = channels[k]
|
ch = channels[k]
|
||||||
|
if k not in channel_versions:
|
||||||
|
# A forced snapshot of a never-written channel still has to
|
||||||
|
# land to stop the ancestor walk, and `put` only stores blobs
|
||||||
|
# for versioned channels.
|
||||||
|
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
|
||||||
if k in channels_to_snapshot:
|
if k in channels_to_snapshot:
|
||||||
# Callers force a full snapshot blob here: exit mode when a
|
# Callers force a full snapshot blob here: exit mode when a
|
||||||
# delta channel reaches its snapshot cadence, and update_state
|
# delta channel reaches its snapshot cadence, update_state on
|
||||||
# on a fresh thread (no ancestor to replay writes from). The
|
# a fresh thread (no ancestor to replay writes from), and a
|
||||||
# manual version-bump below only applies to the exit-mode case.
|
# fork. The manual version-bump below only applies to the
|
||||||
|
# exit-mode case.
|
||||||
#
|
#
|
||||||
# In exit mode, the snapshot decision is deferred to exit
|
# In exit mode, the snapshot decision is deferred to exit
|
||||||
# time (intermediate steps have do_checkpoint=False). The
|
# time (intermediate steps have do_checkpoint=False). The
|
||||||
|
|||||||
@@ -222,10 +222,16 @@ class PregelLoop:
|
|||||||
# under the saver's `ORDER BY task_id, idx` sorting.
|
# under the saver's `ORDER BY task_id, idx` sorting.
|
||||||
_exit_delta_writes: list[tuple[int, str, str, Any]] | None = None
|
_exit_delta_writes: list[tuple[int, str, str, Any]] | None = None
|
||||||
|
|
||||||
# Delta channels that saw an Overwrite since the last checkpoint. These
|
# Delta channels that must snapshot at the next checkpoint, whatever their
|
||||||
# channels must snapshot after live update applies overwrite semantics so
|
# cadence counters say:
|
||||||
# sparse replay starts from the same post-overwrite value.
|
# * an Overwrite arrived since the last checkpoint, so sparse replay has to
|
||||||
_delta_channels_with_overwrite: set[str]
|
# start from the post-overwrite value;
|
||||||
|
# * this run forked off an explicitly addressed checkpoint. That base also
|
||||||
|
# holds the writes of the branch the fork abandons, and nothing records
|
||||||
|
# which child consumed which, so the ancestor walk must stop inside the
|
||||||
|
# fork. Any addressed checkpoint counts, because telling a real fork
|
||||||
|
# apart would mean trusting the base's `pending_writes` to be complete.
|
||||||
|
_delta_channels_forced_snapshot: set[str]
|
||||||
|
|
||||||
# The checkpoint_config that points at the parent loaded at `__enter__`
|
# The checkpoint_config that points at the parent loaded at `__enter__`
|
||||||
# (or the synthetic-empty checkpoint, on first run). We capture it
|
# (or the synthetic-empty checkpoint, on first run). We capture it
|
||||||
@@ -369,6 +375,13 @@ class PregelLoop:
|
|||||||
if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
|
if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)
|
||||||
else ()
|
else ()
|
||||||
)
|
)
|
||||||
|
# Value, not key presence like `is_replaying`: subgraph task configs
|
||||||
|
# always carry an explicit `None` checkpoint_id.
|
||||||
|
self._delta_channels_forced_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
|
self.prev_checkpoint_config = None
|
||||||
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
|
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
|
||||||
self.control = runtime.control if isinstance(runtime, Runtime) else None
|
self.control = runtime.control if isinstance(runtime, Runtime) else None
|
||||||
@@ -683,7 +696,7 @@ class PregelLoop:
|
|||||||
def after_tick(self) -> None:
|
def after_tick(self) -> None:
|
||||||
# finish superstep
|
# finish superstep
|
||||||
writes = [w for t in self.tasks.values() for w in t.writes]
|
writes = [w for t in self.tasks.values() for w in t.writes]
|
||||||
self._delta_channels_with_overwrite.update(
|
self._delta_channels_forced_snapshot.update(
|
||||||
ch
|
ch
|
||||||
for ch, v in writes
|
for ch, v in writes
|
||||||
if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0]
|
if isinstance(self.specs.get(ch), DeltaChannel) and _get_overwrite(v)[0]
|
||||||
@@ -991,7 +1004,7 @@ class PregelLoop:
|
|||||||
manager=None,
|
manager=None,
|
||||||
updated_channels=updated_channels,
|
updated_channels=updated_channels,
|
||||||
)
|
)
|
||||||
self._delta_channels_with_overwrite.update(
|
self._delta_channels_forced_snapshot.update(
|
||||||
c
|
c
|
||||||
for c, v in input_writes
|
for c, v in input_writes
|
||||||
if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0]
|
if isinstance(self.specs.get(c), DeltaChannel) and _get_overwrite(v)[0]
|
||||||
@@ -1136,7 +1149,7 @@ class PregelLoop:
|
|||||||
# create new checkpoint
|
# create new checkpoint
|
||||||
channels_to_snapshot = (
|
channels_to_snapshot = (
|
||||||
delta_channels_to_snapshot(self.channels, new_counters)
|
delta_channels_to_snapshot(self.channels, new_counters)
|
||||||
| self._delta_channels_with_overwrite
|
| self._delta_channels_forced_snapshot
|
||||||
if do_checkpoint
|
if do_checkpoint
|
||||||
else set()
|
else set()
|
||||||
)
|
)
|
||||||
@@ -1154,7 +1167,7 @@ class PregelLoop:
|
|||||||
for k in channels_to_snapshot:
|
for k in channels_to_snapshot:
|
||||||
new_counters[k] = (0, 0)
|
new_counters[k] = (0, 0)
|
||||||
if do_checkpoint:
|
if do_checkpoint:
|
||||||
self._delta_channels_with_overwrite.difference_update(channels_to_snapshot)
|
self._delta_channels_forced_snapshot.difference_update(channels_to_snapshot)
|
||||||
non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)}
|
non_zero = {k: v for k, v in new_counters.items() if v != (0, 0)}
|
||||||
if non_zero:
|
if non_zero:
|
||||||
self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero
|
self.checkpoint_metadata["counters_since_delta_snapshot"] = non_zero
|
||||||
@@ -1239,7 +1252,7 @@ class PregelLoop:
|
|||||||
)
|
)
|
||||||
channels_to_snapshot = (
|
channels_to_snapshot = (
|
||||||
delta_channels_to_snapshot(self.channels, counters)
|
delta_channels_to_snapshot(self.channels, counters)
|
||||||
| self._delta_channels_with_overwrite
|
| self._delta_channels_forced_snapshot
|
||||||
)
|
)
|
||||||
|
|
||||||
pending = [
|
pending = [
|
||||||
@@ -1684,7 +1697,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
|||||||
)
|
)
|
||||||
self._delta_write_futs = []
|
self._delta_write_futs = []
|
||||||
self._error_handler_write_futs = []
|
self._error_handler_write_futs = []
|
||||||
self._delta_channels_with_overwrite = set()
|
|
||||||
self._exit_delta_writes = (
|
self._exit_delta_writes = (
|
||||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||||
)
|
)
|
||||||
@@ -1942,7 +1954,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
|||||||
)
|
)
|
||||||
self._delta_write_futs = []
|
self._delta_write_futs = []
|
||||||
self._error_handler_write_futs = []
|
self._error_handler_write_futs = []
|
||||||
self._delta_channels_with_overwrite = set()
|
|
||||||
self._exit_delta_writes = (
|
self._exit_delta_writes = (
|
||||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -108,6 +108,7 @@ from langgraph.callbacks import (
|
|||||||
get_sync_graph_callback_manager_for_config,
|
get_sync_graph_callback_manager_for_config,
|
||||||
)
|
)
|
||||||
from langgraph.channels.base import BaseChannel
|
from langgraph.channels.base import BaseChannel
|
||||||
|
from langgraph.channels.delta import DeltaChannel
|
||||||
from langgraph.channels.topic import Topic
|
from langgraph.channels.topic import Topic
|
||||||
from langgraph.config import get_config
|
from langgraph.config import get_config
|
||||||
from langgraph.constants import END
|
from langgraph.constants import END
|
||||||
@@ -133,6 +134,7 @@ from langgraph.pregel._checkpoint import (
|
|||||||
copy_checkpoint,
|
copy_checkpoint,
|
||||||
create_checkpoint,
|
create_checkpoint,
|
||||||
create_checkpoint_plan_for_update_state_api,
|
create_checkpoint_plan_for_update_state_api,
|
||||||
|
create_fork_checkpoint,
|
||||||
empty_checkpoint,
|
empty_checkpoint,
|
||||||
get_updated_channels_from_tasks,
|
get_updated_channels_from_tasks,
|
||||||
)
|
)
|
||||||
@@ -1637,8 +1639,21 @@ class Pregel(
|
|||||||
else:
|
else:
|
||||||
raise ValueError(f"Subgraph {recast} not found")
|
raise ValueError(f"Subgraph {recast} not found")
|
||||||
|
|
||||||
|
# Read once from the caller's config: every later superstep receives
|
||||||
|
# the config of the checkpoint just written, which always names one.
|
||||||
|
# Cleared by the first checkpoint that carries the snapshots, which
|
||||||
|
# `__copy__` does not write.
|
||||||
|
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(
|
def perform_superstep(
|
||||||
input_config: RunnableConfig, updates: Sequence[StateUpdate]
|
input_config: RunnableConfig,
|
||||||
|
updates: Sequence[StateUpdate],
|
||||||
|
*,
|
||||||
|
is_fork: bool,
|
||||||
) -> RunnableConfig:
|
) -> RunnableConfig:
|
||||||
# get last checkpoint
|
# get last checkpoint
|
||||||
config = ensure_config(self.config, input_config)
|
config = ensure_config(self.config, input_config)
|
||||||
@@ -1726,9 +1741,17 @@ class Pregel(
|
|||||||
self.trigger_to_nodes,
|
self.trigger_to_nodes,
|
||||||
)
|
)
|
||||||
# save checkpoint
|
# 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(
|
next_config = checkpointer.put(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
create_checkpoint(checkpoint, channels, step),
|
next_checkpoint,
|
||||||
{
|
{
|
||||||
"source": "update",
|
"source": "update",
|
||||||
"step": step + 1,
|
"step": step + 1,
|
||||||
@@ -1736,7 +1759,7 @@ class Pregel(
|
|||||||
},
|
},
|
||||||
get_new_channel_versions(
|
get_new_channel_versions(
|
||||||
checkpoint_previous_versions,
|
checkpoint_previous_versions,
|
||||||
checkpoint["channel_versions"],
|
next_checkpoint["channel_versions"],
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
return patch_checkpoint_map(
|
return patch_checkpoint_map(
|
||||||
@@ -1765,9 +1788,17 @@ class Pregel(
|
|||||||
if saved and saved.metadata.get("step") is not None
|
if saved and saved.metadata.get("step") is not None
|
||||||
else -1
|
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(
|
next_config = checkpointer.put(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
create_checkpoint(checkpoint, channels, next_step),
|
next_checkpoint,
|
||||||
{
|
{
|
||||||
"source": "input",
|
"source": "input",
|
||||||
"step": next_step,
|
"step": next_step,
|
||||||
@@ -1777,7 +1808,7 @@ class Pregel(
|
|||||||
},
|
},
|
||||||
get_new_channel_versions(
|
get_new_channel_versions(
|
||||||
checkpoint_previous_versions,
|
checkpoint_previous_versions,
|
||||||
checkpoint["channel_versions"],
|
next_checkpoint["channel_versions"],
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1873,6 +1904,7 @@ class Pregel(
|
|||||||
return perform_superstep(
|
return perform_superstep(
|
||||||
patch_checkpoint_map(next_config, saved.metadata),
|
patch_checkpoint_map(next_config, saved.metadata),
|
||||||
[item for lst in user_group_by.values() for item in lst],
|
[item for lst in user_group_by.values() for item in lst],
|
||||||
|
is_fork=is_fork,
|
||||||
)
|
)
|
||||||
|
|
||||||
return patch_checkpoint_map(next_config, saved.metadata)
|
return patch_checkpoint_map(next_config, saved.metadata)
|
||||||
@@ -2020,6 +2052,7 @@ class Pregel(
|
|||||||
parents=saved.metadata.get("parents", {}) if saved else {},
|
parents=saved.metadata.get("parents", {}) if saved else {},
|
||||||
saved_metadata=saved.metadata if saved else None,
|
saved_metadata=saved.metadata if saved else None,
|
||||||
is_fresh_thread=saved is None,
|
is_fresh_thread=saved is None,
|
||||||
|
is_fork=is_fork,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
checkpoint = create_checkpoint(
|
checkpoint = create_checkpoint(
|
||||||
@@ -2032,6 +2065,8 @@ class Pregel(
|
|||||||
else None,
|
else None,
|
||||||
channels_to_snapshot=channels_to_snapshot,
|
channels_to_snapshot=channels_to_snapshot,
|
||||||
)
|
)
|
||||||
|
if is_fork:
|
||||||
|
fork_pending.difference_update(checkpoint["channel_values"])
|
||||||
next_config = checkpointer.put(
|
next_config = checkpointer.put(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
checkpoint,
|
checkpoint,
|
||||||
@@ -2050,7 +2085,9 @@ class Pregel(
|
|||||||
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
|
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
|
||||||
)
|
)
|
||||||
for superstep in supersteps:
|
for superstep in supersteps:
|
||||||
current_config = perform_superstep(current_config, superstep)
|
current_config = perform_superstep(
|
||||||
|
current_config, superstep, is_fork=bool(fork_pending)
|
||||||
|
)
|
||||||
return current_config
|
return current_config
|
||||||
|
|
||||||
async def abulk_update_state(
|
async def abulk_update_state(
|
||||||
@@ -2103,8 +2140,21 @@ class Pregel(
|
|||||||
else:
|
else:
|
||||||
raise ValueError(f"Subgraph {recast} not found")
|
raise ValueError(f"Subgraph {recast} not found")
|
||||||
|
|
||||||
|
# Read once from the caller's config: every later superstep receives
|
||||||
|
# the config of the checkpoint just written, which always names one.
|
||||||
|
# Cleared by the first checkpoint that carries the snapshots, which
|
||||||
|
# `__copy__` does not write.
|
||||||
|
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(
|
async def aperform_superstep(
|
||||||
input_config: RunnableConfig, updates: Sequence[StateUpdate]
|
input_config: RunnableConfig,
|
||||||
|
updates: Sequence[StateUpdate],
|
||||||
|
*,
|
||||||
|
is_fork: bool,
|
||||||
) -> RunnableConfig:
|
) -> RunnableConfig:
|
||||||
# get last checkpoint
|
# get last checkpoint
|
||||||
config = ensure_config(self.config, input_config)
|
config = ensure_config(self.config, input_config)
|
||||||
@@ -2190,16 +2240,25 @@ class Pregel(
|
|||||||
self.trigger_to_nodes,
|
self.trigger_to_nodes,
|
||||||
)
|
)
|
||||||
# save checkpoint
|
# 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(
|
next_config = await checkpointer.aput(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
create_checkpoint(checkpoint, channels, step),
|
next_checkpoint,
|
||||||
{
|
{
|
||||||
"source": "update",
|
"source": "update",
|
||||||
"step": step + 1,
|
"step": step + 1,
|
||||||
"parents": saved.metadata.get("parents", {}) if saved else {},
|
"parents": saved.metadata.get("parents", {}) if saved else {},
|
||||||
},
|
},
|
||||||
get_new_channel_versions(
|
get_new_channel_versions(
|
||||||
checkpoint_previous_versions, checkpoint["channel_versions"]
|
checkpoint_previous_versions,
|
||||||
|
next_checkpoint["channel_versions"],
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
return patch_checkpoint_map(
|
return patch_checkpoint_map(
|
||||||
@@ -2228,9 +2287,17 @@ class Pregel(
|
|||||||
if saved and saved.metadata.get("step") is not None
|
if saved and saved.metadata.get("step") is not None
|
||||||
else -1
|
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(
|
next_config = await checkpointer.aput(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
create_checkpoint(checkpoint, channels, next_step),
|
next_checkpoint,
|
||||||
{
|
{
|
||||||
"source": "input",
|
"source": "input",
|
||||||
"step": next_step,
|
"step": next_step,
|
||||||
@@ -2240,7 +2307,7 @@ class Pregel(
|
|||||||
},
|
},
|
||||||
get_new_channel_versions(
|
get_new_channel_versions(
|
||||||
checkpoint_previous_versions,
|
checkpoint_previous_versions,
|
||||||
checkpoint["channel_versions"],
|
next_checkpoint["channel_versions"],
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -2335,6 +2402,7 @@ class Pregel(
|
|||||||
return await aperform_superstep(
|
return await aperform_superstep(
|
||||||
patch_checkpoint_map(next_config, saved.metadata),
|
patch_checkpoint_map(next_config, saved.metadata),
|
||||||
[item for lst in user_group_by.values() for item in lst],
|
[item for lst in user_group_by.values() for item in lst],
|
||||||
|
is_fork=is_fork,
|
||||||
)
|
)
|
||||||
|
|
||||||
return patch_checkpoint_map(
|
return patch_checkpoint_map(
|
||||||
@@ -2480,6 +2548,7 @@ class Pregel(
|
|||||||
parents=saved.metadata.get("parents", {}) if saved else {},
|
parents=saved.metadata.get("parents", {}) if saved else {},
|
||||||
saved_metadata=saved.metadata if saved else None,
|
saved_metadata=saved.metadata if saved else None,
|
||||||
is_fresh_thread=saved is None,
|
is_fresh_thread=saved is None,
|
||||||
|
is_fork=is_fork,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
checkpoint = create_checkpoint(
|
checkpoint = create_checkpoint(
|
||||||
@@ -2492,6 +2561,8 @@ class Pregel(
|
|||||||
else None,
|
else None,
|
||||||
channels_to_snapshot=channels_to_snapshot,
|
channels_to_snapshot=channels_to_snapshot,
|
||||||
)
|
)
|
||||||
|
if is_fork:
|
||||||
|
fork_pending.difference_update(checkpoint["channel_values"])
|
||||||
next_config = await checkpointer.aput(
|
next_config = await checkpointer.aput(
|
||||||
checkpoint_config,
|
checkpoint_config,
|
||||||
checkpoint,
|
checkpoint,
|
||||||
@@ -2509,7 +2580,9 @@ class Pregel(
|
|||||||
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
|
config, {CONFIG_KEY_THREAD_ID: str(config[CONF][CONFIG_KEY_THREAD_ID])}
|
||||||
)
|
)
|
||||||
for superstep in supersteps:
|
for superstep in supersteps:
|
||||||
current_config = await aperform_superstep(current_config, superstep)
|
current_config = await aperform_superstep(
|
||||||
|
current_config, superstep, is_fork=bool(fork_pending)
|
||||||
|
)
|
||||||
return current_config
|
return current_config
|
||||||
|
|
||||||
def update_state(
|
def update_state(
|
||||||
|
|||||||
@@ -85,11 +85,13 @@ class MemorySaverAssertImmutable(InMemorySaver):
|
|||||||
)
|
)
|
||||||
== saved
|
== saved
|
||||||
), config["configurable"]["checkpoint_ns"]
|
), config["configurable"]["checkpoint_ns"]
|
||||||
|
next_config = super().put(config, checkpoint, metadata, new_versions)
|
||||||
|
# Read back, not the object handed in: a DeltaChannel a step did not
|
||||||
|
# write is refilled on read from the blob its inherited version points at.
|
||||||
self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = (
|
self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = (
|
||||||
self.serde.dumps_typed(checkpoint)
|
self.serde.dumps_typed(super().get(next_config))
|
||||||
)
|
)
|
||||||
# call super to write checkpoint
|
return next_config
|
||||||
return super().put(config, checkpoint, metadata, new_versions)
|
|
||||||
|
|
||||||
|
|
||||||
class MemorySaverNoPending(InMemorySaver):
|
class MemorySaverNoPending(InMemorySaver):
|
||||||
|
|||||||
@@ -0,0 +1,346 @@
|
|||||||
|
"""Forking a thread must not replay the abandoned branch into the fork.
|
||||||
|
|
||||||
|
Every graph carries a ``DeltaChannel`` and a plain reducer channel fed the same
|
||||||
|
values; the plain channel needs no replay, so it is the oracle.
|
||||||
|
"""
|
||||||
|
|
||||||
|
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:
|
||||||
|
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]
|
||||||
|
other: Annotated[list, add]
|
||||||
|
|
||||||
|
|
||||||
|
def _build(checkpointer: BaseCheckpointSaver, tag: str) -> Any:
|
||||||
|
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:
|
||||||
|
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:
|
||||||
|
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]:
|
||||||
|
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:
|
||||||
|
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)
|
||||||
|
abandoned_head = graph.get_state(config)
|
||||||
|
|
||||||
|
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"]
|
||||||
|
|
||||||
|
abandoned = graph.get_state(abandoned_head.config).values
|
||||||
|
assert abandoned["log"] == abandoned["plain"] == abandoned_head.values["log"]
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
abandoned_head = await graph.aget_state(config)
|
||||||
|
|
||||||
|
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"]
|
||||||
|
|
||||||
|
abandoned = (await graph.aget_state(abandoned_head.config)).values
|
||||||
|
assert abandoned["log"] == abandoned["plain"] == abandoned_head.values["log"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_fork_off_checkpoint_before_first_input(
|
||||||
|
sync_checkpointer: BaseCheckpointSaver, durability: Durability
|
||||||
|
) -> None:
|
||||||
|
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:
|
||||||
|
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:
|
||||||
|
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:
|
||||||
|
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:
|
||||||
|
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:
|
||||||
|
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:
|
||||||
|
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)
|
||||||
|
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:
|
||||||
|
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