From fd64ada9de39b54e49ce454910ba89599eefb85f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 31 Mar 2025 17:48:19 -0700 Subject: [PATCH] Avoid raise-catch strategy in BaseChannel.checkpoint() - This mirrors the work done earlier on BaseChannel.get() - Comparing to a sentinel value is significantly faster than raising and catching an exception --- .../langgraph/checkpoint/base/__init__.py | 1 + libs/langgraph/bench/sequential.py | 2 +- .../langgraph/langgraph/channels/any_value.py | 3 ++ libs/langgraph/langgraph/channels/base.py | 6 ++- libs/langgraph/langgraph/channels/binop.py | 3 ++ .../langgraph/channels/ephemeral_value.py | 3 ++ .../langgraph/channels/last_value.py | 3 ++ .../langgraph/channels/untracked_value.py | 2 +- libs/langgraph/langgraph/pregel/__init__.py | 2 +- libs/langgraph/langgraph/pregel/checkpoint.py | 37 +++++++++++++++++++ libs/langgraph/langgraph/pregel/loop.py | 2 +- 11 files changed, 59 insertions(+), 5 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/checkpoint.py diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 2d995e291..aa53306bb 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -124,6 +124,7 @@ def copy_checkpoint(checkpoint: Checkpoint) -> Checkpoint: ) +# Kept for backwards compat, newer versions of LangGraph no longer use this. def create_checkpoint( checkpoint: Checkpoint, channels: Optional[Mapping[str, ChannelProtocol]], diff --git a/libs/langgraph/bench/sequential.py b/libs/langgraph/bench/sequential.py index ef3ce6329..bfdad823e 100644 --- a/libs/langgraph/bench/sequential.py +++ b/libs/langgraph/bench/sequential.py @@ -34,7 +34,7 @@ if __name__ == "__main__": import uvloop - graph = create_sequential(2000).compile() + graph = create_sequential(3000).compile() input = {"messages": []} # Empty list of messages config = {"recursion_limit": 20000000000} diff --git a/libs/langgraph/langgraph/channels/any_value.py b/libs/langgraph/langgraph/channels/any_value.py index 0276030aa..58b923a03 100644 --- a/libs/langgraph/langgraph/channels/any_value.py +++ b/libs/langgraph/langgraph/channels/any_value.py @@ -55,3 +55,6 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): def is_available(self) -> bool: return self.value is not MISSING + + def checkpoint(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 82fd059d1..33143fffb 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -3,6 +3,7 @@ from typing import Any, Generic, Sequence, TypeVar from typing_extensions import Self +from langgraph.constants import MISSING from langgraph.errors import EmptyChannelError, InvalidUpdateError Value = TypeVar("Value") @@ -33,7 +34,10 @@ class BaseChannel(Generic[Value, Update, C], ABC): """Return a serializable representation of the channel's current state. Raises EmptyChannelError if the channel is empty (never updated yet), or doesn't support checkpoints.""" - return self.get() + try: + return self.get() + except EmptyChannelError: + return MISSING @abstractmethod def from_checkpoint(self, checkpoint: C) -> Self: diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index eb90cae8c..59cc11550 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/libs/langgraph/langgraph/channels/binop.py @@ -90,3 +90,6 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): def is_available(self) -> bool: return self.value is not MISSING + + def checkpoint(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/channels/ephemeral_value.py b/libs/langgraph/langgraph/channels/ephemeral_value.py index 242149fe6..29471afb4 100644 --- a/libs/langgraph/langgraph/channels/ephemeral_value.py +++ b/libs/langgraph/langgraph/channels/ephemeral_value.py @@ -59,3 +59,6 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): def is_available(self) -> bool: return self.value is not MISSING + + def checkpoint(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py index 32a951a4b..90707575c 100644 --- a/libs/langgraph/langgraph/channels/last_value.py +++ b/libs/langgraph/langgraph/channels/last_value.py @@ -61,3 +61,6 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): def is_available(self) -> bool: return self.value is not MISSING + + def checkpoint(self) -> Value: + return self.value diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py index cc9c99bee..24ffbd425 100644 --- a/libs/langgraph/langgraph/channels/untracked_value.py +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -31,7 +31,7 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): return self.typ def checkpoint(self) -> Value: - raise EmptyChannelError() + return MISSING def from_checkpoint(self, checkpoint: Value) -> Self: empty = self.__class__(self.typ, self.guard) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 45efe5c17..a550de23d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -50,7 +50,6 @@ from langgraph.checkpoint.base import ( BaseCheckpointSaver, CheckpointTuple, copy_checkpoint, - create_checkpoint, empty_checkpoint, ) from langgraph.constants import ( @@ -91,6 +90,7 @@ from langgraph.pregel.algo import ( local_write, prepare_next_tasks, ) +from langgraph.pregel.checkpoint import create_checkpoint from langgraph.pregel.debug import tasks_w_writes from langgraph.pregel.io import map_input, read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop diff --git a/libs/langgraph/langgraph/pregel/checkpoint.py b/libs/langgraph/langgraph/pregel/checkpoint.py new file mode 100644 index 000000000..3f0af208b --- /dev/null +++ b/libs/langgraph/langgraph/pregel/checkpoint.py @@ -0,0 +1,37 @@ +from datetime import datetime, timezone +from typing import Mapping, Optional + +from langgraph.channels.base import BaseChannel +from langgraph.checkpoint.base import LATEST_VERSION, Checkpoint +from langgraph.checkpoint.base.id import uuid6 +from langgraph.constants import MISSING + + +def create_checkpoint( + checkpoint: Checkpoint, + channels: Optional[Mapping[str, BaseChannel]], + step: int, + *, + id: Optional[str] = None, +) -> Checkpoint: + """Create a checkpoint for the given channels.""" + ts = datetime.now(timezone.utc).isoformat() + if channels is None: + values = checkpoint["channel_values"] + else: + values = {} + for k in channels: + if k not in checkpoint["channel_versions"]: + continue + v = channels[k].checkpoint() + if v is not MISSING: + values[k] = v + return Checkpoint( + v=LATEST_VERSION, + ts=ts, + id=id or str(uuid6(clock_seq=step)), + channel_values=values, + channel_versions=checkpoint["channel_versions"], + versions_seen=checkpoint["versions_seen"], + pending_sends=checkpoint.get("pending_sends", []), + ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 3449ccde3..031d4d2a2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -38,7 +38,6 @@ from langgraph.checkpoint.base import ( CheckpointTuple, PendingWrite, copy_checkpoint, - create_checkpoint, empty_checkpoint, ) from langgraph.constants import ( @@ -88,6 +87,7 @@ from langgraph.pregel.algo import ( should_interrupt, task_path_str, ) +from langgraph.pregel.checkpoint import create_checkpoint from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results,