WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-28 10:49:56 +02:00
Code Issues Packages Projects Releases Wiki Activity

Move all checkpoint edits to apply_writes

Browse Source
This commit is contained in:
Nuno Campos
2024-07-19 16:37:27 -07:00
parent 27cd4e2221
commit 1dbf7a3392
6 changed files with 109 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/checkpoint/base.py
+3 -7
View File
@@ -1,5 +1,4 @@
from abc import ABC
from collections import defaultdict
from datetime import datetime, timezone
from typing import (
Any,
@@ -79,7 +78,7 @@ class Checkpoint(TypedDict):
The keys are channel names and the values are the logical time step
at which the channel was last updated.
"""
versions_seen: defaultdict[str, dict[str, Union[str, int, float]]]
versions_seen: dict[str, dict[str, Union[str, int, float]]]
"""Map from node ID to map from channel name to version seen.
This keeps track of the versions of the channels that each node has seen.
@@ -100,7 +99,7 @@ def empty_checkpoint() -> Checkpoint:
ts=datetime.now(timezone.utc).isoformat(),
channel_values={},
channel_versions={},
versions_seen=defaultdict(dict),
versions_seen={},
pending_sends=[],
current_tasks={},
)
@@ -113,10 +112,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint:
id=checkpoint["id"],
channel_values=checkpoint["channel_values"].copy(),
channel_versions=checkpoint["channel_versions"].copy(),
versions_seen=defaultdict(
dict,
{k: v.copy() for k, v in checkpoint["versions_seen"].items()},
),
versions_seen={k: v.copy() for k, v in checkpoint["versions_seen"].items()},
pending_sends=checkpoint.get("pending_sends", []).copy(),
current_tasks=checkpoint.get("current_tasks", {}).copy(),
)
libs/langgraph/langgraph/constants.py
+2 -1
View File
@@ -1,10 +1,11 @@
from typing import Any
INPUT = "__input__"
CONFIG_KEY_SEND = "__pregel_send"
CONFIG_KEY_READ = "__pregel_read"
INTERRUPT = "__interrupt__"
TASKS = "__pregel_tasks"
RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ}
RESERVED = {INTERRUPT, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, INPUT}
TAG_HIDDEN = "langsmith:hidden"
START = "__start__"
libs/langgraph/langgraph/pregel/__init__.py
+8 -9
View File
@@ -355,7 +355,7 @@ class Pregel(
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
_, next_tasks = prepare_next_tasks(
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
@@ -387,7 +387,7 @@ class Pregel(
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
_, next_tasks = prepare_next_tasks(
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
@@ -429,7 +429,7 @@ class Pregel(
) as channels, ManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
_, next_tasks = prepare_next_tasks(
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
@@ -475,7 +475,7 @@ class Pregel(
) as channels, AsyncManagedValuesManager(
self.managed_values_dict, ensure_config(config), self
) as managed:
_, next_tasks = prepare_next_tasks(
next_tasks = prepare_next_tasks(
checkpoint,
self.nodes,
channels,
@@ -560,14 +560,14 @@ class Pregel(
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
CONFIG_KEY_READ: partial(
local_read, checkpoint, channels, task.writes, config
local_read, checkpoint, channels, task, config
),
},
),
)
# apply to checkpoint and save
apply_writes(
checkpoint, channels, task.writes, self.checkpointer.get_next_version
checkpoint, channels, [task], self.checkpointer.get_next_version
)
step = saved.metadata.get("step", -2) + 1 if saved else -1
@@ -652,14 +652,14 @@ class Pregel(
# deque.extend is thread-safe
CONFIG_KEY_SEND: task.writes.extend,
CONFIG_KEY_READ: partial(
local_read, checkpoint, channels, task.writes, config
local_read, checkpoint, channels, task, config
),
},
),
)
# apply to checkpoint and save
apply_writes(
checkpoint, channels, task.writes, self.checkpointer.get_next_version
checkpoint, channels, [task], self.checkpointer.get_next_version
)
step = saved.metadata.get("step", -2) + 1 if saved else -1
@@ -1145,7 +1145,6 @@ class Pregel(
# exception will be handled in panic_or_proceed
futures.clear()
else:
print(loop.step, task.name, stream_modes)
# save task writes to checkpointer
loop.put_writes(task.id, task.writes)
# yield updates output for the finished task
libs/langgraph/langgraph/pregel/algo.py
+68 -53
View File
@@ -7,7 +7,9 @@ from typing import (
Iterator,
Literal,
Mapping,
NamedTuple,
Optional,
Protocol,
Sequence,
Union,
overload,
@@ -29,6 +31,7 @@ from langgraph.constants import (
CONFIG_KEY_READ,
CONFIG_KEY_SEND,
INTERRUPT,
RESERVED,
TAG_HIDDEN,
TASKS,
Send,
@@ -41,6 +44,18 @@ from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription
class WritesProtocol(Protocol):
name: str
writes: Sequence[tuple[str, Any]]
triggers: Sequence[str]
class PregelTaskWrites(NamedTuple):
name: str
writes: Sequence[tuple[str, Any]]
triggers: Sequence[str]
def should_interrupt(
checkpoint: Checkpoint,
interrupt_nodes: Union[All, Sequence[str]],
@@ -48,8 +63,7 @@ def should_interrupt(
) -> bool:
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
null_version = version_type()
# defaultdicts are mutated on access :( so we need to copy
seen = checkpoint["versions_seen"].copy()[INTERRUPT]
seen = checkpoint["versions_seen"].get(INTERRUPT, {})
return (
# interrupt if any channel has been updated since last interrupt
any(
@@ -72,21 +86,21 @@ def should_interrupt(
def local_read(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
writes: Sequence[tuple[str, Any]],
task: WritesProtocol,
config: RunnableConfig,
select: Union[list[str], str],
fresh: bool = False,
) -> Union[dict[str, Any], Any]:
if fresh:
checkpoint = create_checkpoint(checkpoint, channels, -1)
new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1)
context_channels = {k: v for k, v in channels.items() if isinstance(v, Context)}
with ChannelsManager(
{k: v for k, v in channels.items() if k not in context_channels},
checkpoint,
new_checkpoint,
config,
) as channels:
all_channels = {**channels, **context_channels}
apply_writes(copy_checkpoint(checkpoint), all_channels, writes, None)
apply_writes(new_checkpoint, all_channels, [task], None)
return read_channels(all_channels, select)
else:
return read_channels(channels, select)
@@ -118,19 +132,46 @@ def increment(current: Optional[int], channel: BaseChannel) -> int:
def apply_writes(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
pending_writes: Sequence[tuple[str, Any]],
tasks: Sequence[WritesProtocol],
get_next_version: Optional[Callable[[int, BaseChannel], int]],
) -> None:
# update seen versions
for task in tasks:
checkpoint["versions_seen"].setdefault(task.name, {}).update(
{
chan: checkpoint["channel_versions"][chan]
for chan in task.triggers
if chan in checkpoint["channel_versions"]
}
)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = None
# Consume all channels that were read
for chan in {
chan for task in tasks for chan in task.triggers if chan not in RESERVED
}:
if channels[chan].consume():
if get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, channels[chan]
)
# clear pending sends
if checkpoint["pending_sends"]:
checkpoint["pending_sends"].clear()
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
# Group writes by channel
for chan, val in pending_writes:
if chan == TASKS:
checkpoint["pending_sends"].append(val)
else:
pending_writes_by_channel[chan].append(val)
pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list)
for task in tasks:
for chan, val in task.writes:
if chan == TASKS:
checkpoint["pending_sends"].append(val)
else:
pending_writes_by_channel[chan].append(val)
# Find the highest version of all channels
if checkpoint["channel_versions"]:
@@ -138,8 +179,8 @@ def apply_writes(
else:
max_version = None
updated_channels: set[str] = set()
# Apply writes to channels
updated_channels: set[str] = set()
for chan, vals in pending_writes_by_channel.items():
if chan in channels:
try:
@@ -153,6 +194,7 @@ def apply_writes(
max_version, channels[chan]
)
updated_channels.add(chan)
# Channels that weren't updated in this step are notified of a new step
for chan in channels:
if chan not in updated_channels:
@@ -171,9 +213,8 @@ def prepare_next_tasks(
config: RunnableConfig,
step: int,
for_execution: Literal[False],
get_next_version: Literal[None] = None,
manager: Literal[None] = None,
) -> tuple[Checkpoint, list[PregelTaskDescription]]:
) -> list[PregelTaskDescription]:
...
@@ -186,9 +227,8 @@ def prepare_next_tasks(
config: RunnableConfig,
step: int,
for_execution: Literal[True],
get_next_version: Callable[[int, BaseChannel], int],
manager: Union[None, ParentRunManager, AsyncParentRunManager],
) -> tuple[Checkpoint, list[PregelExecutableTask]]:
) -> list[PregelExecutableTask]:
...
@@ -201,10 +241,8 @@ def prepare_next_tasks(
step: int,
*,
for_execution: bool,
get_next_version: Union[None, Callable[[int, BaseChannel], int]] = None,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]:
checkpoint = copy_checkpoint(checkpoint)
) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]:
tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = []
# Consume pending packets
for packet in checkpoint["pending_sends"]:
@@ -247,7 +285,11 @@ def prepare_next_tasks(
local_write, writes.extend, processes, channels
),
CONFIG_KEY_READ: partial(
local_read, checkpoint, channels, writes, config
local_read,
checkpoint,
channels,
PregelTaskWrites(packet.node, writes, triggers),
config,
),
},
),
@@ -258,18 +300,14 @@ def prepare_next_tasks(
)
else:
tasks.append(PregelTaskDescription(packet.node, packet.arg))
if for_execution:
checkpoint["pending_sends"].clear()
# Collect channels to consume
channels_to_consume = set()
# Check if any processes should be run in next step
# If so, prepare the values to be passed to them
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
null_version = version_type()
if null_version is None:
return checkpoint, tasks
return tasks
for name, proc in processes.items():
seen = checkpoint["versions_seen"][name]
seen = checkpoint["versions_seen"].get(name, {})
# If any of the channels read by this process were updated
if triggers := sorted(
chan
@@ -280,22 +318,11 @@ def prepare_next_tasks(
and checkpoint["channel_versions"].get(chan, null_version)
> seen.get(chan, null_version)
):
channels_to_consume.update(triggers)
try:
val = next(_proc_input(step, name, proc, managed, channels))
except StopIteration:
continue
# update seen versions
if for_execution:
seen.update(
{
chan: checkpoint["channel_versions"][chan]
for chan in proc.triggers
if chan in checkpoint["channel_versions"]
}
)
if for_execution:
if node := proc.get_node():
metadata = {
@@ -333,7 +360,7 @@ def prepare_next_tasks(
local_read,
checkpoint,
channels,
writes,
PregelTaskWrites(name, writes, triggers),
config,
),
},
@@ -345,19 +372,7 @@ def prepare_next_tasks(
)
else:
tasks.append(PregelTaskDescription(name, val))
# Find the highest version of all channels
if checkpoint["channel_versions"]:
max_version = max(checkpoint["channel_versions"].values())
else:
max_version = None
# Consume all channels that were read
if for_execution:
for chan in channels_to_consume:
if channels[chan].consume():
checkpoint["channel_versions"][chan] = get_next_version(
max_version, channels[chan]
)
return checkpoint, tasks
return tasks
def _proc_input(
libs/langgraph/langgraph/pregel/loop.py
+18 -20
View File
@@ -38,13 +38,14 @@ from langgraph.checkpoint.base import (
copy_checkpoint,
empty_checkpoint,
)
from langgraph.constants import INTERRUPT
from langgraph.constants import INPUT, INTERRUPT
from langgraph.managed.base import (
AsyncManagedValuesManager,
ManagedValueMapping,
ManagedValuesManager,
)
from langgraph.pregel.algo import (
PregelTaskWrites,
apply_writes,
increment,
prepare_next_tasks,
@@ -88,6 +89,7 @@ class PregelLoop:
checkpoint_metadata: CheckpointMetadata
checkpoint_pending_writes: Optional[List[PendingWrite]]
step: int
status: Literal[
"pending", "done", "interrupt_before", "interrupt_after", "out_of_steps"
]
@@ -133,19 +135,13 @@ class PregelLoop:
if self.input is not INPUT_DONE:
self._first()
elif len({tid for tid, _, _ in self.checkpoint_pending_writes}) == len(
self.tasks
):
# assign writes to tasks, apply them in order
grouped: dict[str, list[tuple[str, Any]]] = {}
for tid, k, v in self.checkpoint_pending_writes:
grouped.setdefault(tid, []).append((k, v))
writes = [(k, v) for t in self.tasks for k, v in grouped.get(t.id, [])]
elif all(task.writes for task in self.tasks):
writes = [w for t in self.tasks for w in t.writes]
# all tasks have finished
apply_writes(
self.checkpoint,
self.channels,
writes,
self.tasks,
self.checkpointer_get_next_version,
)
# produce values output
@@ -179,8 +175,7 @@ class PregelLoop:
return False
# prepare next tasks
prev_checkpoint = self.checkpoint
self.checkpoint, self.tasks = prepare_next_tasks(
self.tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.channels,
@@ -188,7 +183,6 @@ class PregelLoop:
self.config,
self.step,
for_execution=True,
get_next_version=self.checkpointer_get_next_version,
manager=manager,
)
@@ -202,9 +196,14 @@ class PregelLoop:
for tid, k, v in self.checkpoint_pending_writes:
if task := next((t for t in self.tasks if t.id == tid), None):
task.writes.append((k, v))
# TODO clear checkpoint_pending_writes
# if all tasks have finished, re-tick
if all(task.writes for task in self.tasks):
return self.tick()
# before execution, check if we should interrupt
if should_interrupt(prev_checkpoint, interrupt_before, self.tasks):
if should_interrupt(self.checkpoint, interrupt_before, self.tasks):
self.status = "interrupt_before"
return False
@@ -219,7 +218,7 @@ class PregelLoop:
# map inputs to channel updates
if input_writes := deque(map_input(self.graph.input_channels, self.input)):
# discard any unfinished tasks from previous checkpoint
self.checkpoint, _ = prepare_next_tasks(
discard_tasks = prepare_next_tasks(
self.checkpoint,
self.graph.nodes,
self.channels,
@@ -227,20 +226,19 @@ class PregelLoop:
self.config,
self.step,
for_execution=True,
get_next_version=self.checkpointer_get_next_version,
)
# apply input writes
apply_writes(
self.checkpoint,
self.channels,
input_writes,
discard_tasks + [PregelTaskWrites(INPUT, input_writes, [])],
self.checkpointer_get_next_version,
)
# save input checkpoint
self._put_checkpoint({"source": "input", "writes": self.input})
else:
# no input is taken as signal to proceed past previous interrupt
self.checkpoint = copy_checkpoint(self.checkpoint)
self.checkpoint["versions_seen"].setdefault(INTERRUPT, {})
for k in self.channels:
if k in self.checkpoint["channel_versions"]:
version = self.checkpoint["channel_versions"][k]
@@ -327,7 +325,7 @@ class SyncPregelLoop(PregelLoop, ContextManager):
**saved.config.get("configurable", {}),
},
}
self.checkpoint = saved.checkpoint
self.checkpoint = copy_checkpoint(saved.checkpoint)
self.checkpoint_metadata = saved.metadata
self.checkpoint_pending_writes = saved.pending_writes
@@ -396,7 +394,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
**saved.config.get("configurable", {}),
},
}
self.checkpoint = saved.checkpoint
self.checkpoint = copy_checkpoint(saved.checkpoint)
self.checkpoint_metadata = saved.metadata
self.checkpoint_pending_writes = saved.pending_writes
libs/langgraph/tests/test_prebuilt.py
+10 -17
View File
@@ -1,4 +1,3 @@
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional, Sequence, Type, Union
import pytest
@@ -100,14 +99,11 @@ def test_no_modifier(checkpointer: Optional[BaseCheckpointSaver]):
"start:agent": 3,
"agent": 3,
},
"versions_seen": defaultdict(
dict,
{
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
"tools": {},
},
),
"versions_seen": {
"__input__": {},
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
},
"pending_sends": [],
"current_tasks": {},
}
@@ -159,14 +155,11 @@ async def test_no_modifier_async(checkpointer: Optional[BaseCheckpointSaver]):
"start:agent": 3,
"agent": 3,
},
"versions_seen": defaultdict(
dict,
{
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
"tools": {},
},
),
"versions_seen": {
"__input__": {},
"__start__": {"__start__": 1},
"agent": {"start:agent": 2},
},
"pending_sends": [],
"current_tasks": {},
}