From 1dbf7a33925e9c26c4d955cd82fa2f07d88eea47 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 19 Jul 2024 16:37:27 -0700 Subject: [PATCH] Move all checkpoint edits to apply_writes --- libs/langgraph/langgraph/checkpoint/base.py | 10 +- libs/langgraph/langgraph/constants.py | 3 +- libs/langgraph/langgraph/pregel/__init__.py | 17 ++- libs/langgraph/langgraph/pregel/algo.py | 121 +++++++++++--------- libs/langgraph/langgraph/pregel/loop.py | 38 +++--- libs/langgraph/tests/test_prebuilt.py | 27 ++--- 6 files changed, 109 insertions(+), 107 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index 8cb278f32..7f1ffab1f 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -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(), ) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 41970ae16..4a50e8833 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -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__" diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 41b517e0a..14edeacbd 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -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 diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 46950b9f8..e93d6e14a 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -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( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index c49d20697..62bf6c806 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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 diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 2e3c2098b..a0a8436e5 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -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": {}, }