Avoid creating checkpoint unless we're saving it (#4106)

- When checkpointing is disabled don't call create_checkpoint in
PregelLoop
- In local_read apply writes directly to copies of updated channels
- Add BaseChannel.copy() method to create channel copies with less
overhead

libs/langgraph/langgraph/channels/any_value.py

@@ -30,9 +30,14 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.key)
+        empty.value = self.value
+        return empty
+
     def from_checkpoint(self, checkpoint: Value) -> Self:
-        empty = self.__class__(self.typ)
-        empty.key = self.key
+        empty = self.__class__(self.typ, self.key)
         if checkpoint is not MISSING:
             empty.value = checkpoint
         return empty

libs/langgraph/langgraph/channels/base.py

@@ -30,6 +30,12 @@ class BaseChannel(Generic[Value, Update, C], ABC):
 
     # serialize/deserialize methods
 
+    def copy(self) -> Self:
+        """Return a copy of the channel.
+        By default, delegates to checkpoint() and from_checkpoint().
+        Subclasses can override this method with a more efficient implementation."""
+        return self.from_checkpoint(self.checkpoint())
+
     def checkpoint(self) -> C:
         """Return a serializable representation of the channel's current state.
         Raises EmptyChannelError if the channel is empty (never updated yet),

libs/langgraph/langgraph/channels/binop.py

@@ -66,6 +66,13 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.operator)
+        empty.key = self.key
+        empty.value = self.value
+        return empty
+
     def from_checkpoint(self, checkpoint: Value) -> Self:
         empty = self.__class__(self.typ, self.operator)
         empty.key = self.key

libs/langgraph/langgraph/channels/dynamic_barrier_value.py

@@ -46,6 +46,14 @@ class DynamicBarrierValue(
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ)
+        empty.key = self.key
+        empty.names = self.names
+        empty.seen = self.seen.copy()
+        return empty
+
     def checkpoint(self) -> tuple[Optional[set[Value]], set[Value]]:
         return (self.names, self.seen)
 

libs/langgraph/langgraph/channels/ephemeral_value.py

@@ -30,6 +30,13 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.guard)
+        empty.key = self.key
+        empty.value = self.value
+        return empty
+
     def from_checkpoint(self, checkpoint: Value) -> Self:
         empty = self.__class__(self.typ, self.guard)
         empty.key = self.key

libs/langgraph/langgraph/channels/last_value.py

@@ -34,9 +34,14 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.key)
+        empty.value = self.value
+        return empty
+
     def from_checkpoint(self, checkpoint: Value) -> Self:
-        empty = self.__class__(self.typ)
-        empty.key = self.key
+        empty = self.__class__(self.typ, self.key)
         if checkpoint is not MISSING:
             empty.value = checkpoint
         return empty

libs/langgraph/langgraph/channels/named_barrier_value.py

@@ -33,6 +33,13 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]):
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.names)
+        empty.key = self.key
+        empty.seen = self.seen.copy()
+        return empty
+
     def checkpoint(self) -> set[Value]:
         return self.seen
 

libs/langgraph/langgraph/channels/topic.py

@@ -48,6 +48,13 @@ class Topic(
         """The type of the update received by the channel."""
         return Union[self.typ, list[self.typ]]  # type: ignore[name-defined]
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.accumulate)
+        empty.key = self.key
+        empty.values = self.values.copy()
+        return empty
+
     def checkpoint(self) -> list[Value]:
         return self.values
 

libs/langgraph/langgraph/channels/untracked_value.py

@@ -30,6 +30,13 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]):
         """The type of the update received by the channel."""
         return self.typ
 
+    def copy(self) -> Self:
+        """Return a copy of the channel."""
+        empty = self.__class__(self.typ, self.guard)
+        empty.key = self.key
+        empty.value = self.value
+        return empty
+
     def checkpoint(self) -> Value:
         return MISSING
 

libs/langgraph/langgraph/pregel/__init__.py

@@ -1535,12 +1535,9 @@ class Pregel(PregelProtocol):
                             ),
                             CONFIG_KEY_READ: partial(
                                 local_read,
-                                step + 1,
-                                checkpoint,
                                 channels,
                                 managed,
                                 task,
-                                config,
                             ),
                         },
                     ),
@@ -1944,12 +1941,9 @@ class Pregel(PregelProtocol):
                             ),
                             CONFIG_KEY_READ: partial(
                                 local_read,
-                                step + 1,
-                                checkpoint,
                                 channels,
                                 managed,
                                 task,
-                                config,
                             ),
                         },
                     ),

libs/langgraph/langgraph/pregel/algo.py

@@ -33,7 +33,6 @@ from langgraph.checkpoint.base import (
     Checkpoint,
     PendingWrite,
     V,
-    copy_checkpoint,
 )
 from langgraph.constants import (
     CONF,
@@ -69,12 +68,10 @@ from langgraph.managed.base import ManagedValueMapping
 from langgraph.pregel.call import get_runnable_for_task
 from langgraph.pregel.io import read_channel, read_channels
 from langgraph.pregel.log import logger
-from langgraph.pregel.manager import ChannelsManager
 from langgraph.pregel.read import PregelNode
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
-    LoopProtocol,
     PregelExecutableTask,
     PregelScratchpad,
     PregelTask,
@@ -169,39 +166,39 @@ def should_interrupt(
 
 
 def local_read(
-    step: int,
-    checkpoint: Checkpoint,
     channels: Mapping[str, BaseChannel],
     managed: ManagedValueMapping,
     task: WritesProtocol,
-    config: RunnableConfig,
     select: Union[list[str], str],
     fresh: bool = False,
 ) -> Union[dict[str, Any], Any]:
     """Function injected under CONFIG_KEY_READ in task config, to read current state.
     Used by conditional edges to read a copy of the state with reflecting the writes
     from that node only."""
+    updated: dict[str, list[Any]] = defaultdict(list)
     if isinstance(select, str):
         managed_keys = []
-        for c, _ in task.writes:
+        for c, v in task.writes:
             if c == select:
-                updated = {c}
-                break
-        else:
-            updated = set()
+                updated[c].append(v)
     else:
         managed_keys = [k for k in select if k in managed]
         select = [k for k in select if k not in managed]
-        updated = set(select).intersection(c for c, _ in task.writes)
+        for c, v in task.writes:
+            if c in select:
+                updated[c].append(v)
     if fresh and updated:
-        with ChannelsManager(
-            {k: v for k, v in channels.items() if k in updated},
-            checkpoint,
-            LoopProtocol(config=config, step=step, stop=step + 1),
-            skip_context=True,
-        ) as (local_channels, _):
-            apply_writes(copy_checkpoint(checkpoint), local_channels, [task], None)
-            values = read_channels({**channels, **local_channels}, select)
+        # apply writes
+        local_channels: dict[str, BaseChannel] = {}
+        for k in channels:
+            if k in updated:
+                cc = channels[k].copy()
+                cc.update(updated[k])
+            else:
+                cc = channels[k]
+            local_channels[k] = cc
+        # read fresh values
+        values = read_channels(local_channels, select)
     else:
         values = read_channels(channels, select)
     if managed_keys:
@@ -335,6 +332,17 @@ def apply_writes(
     return pending_writes_by_managed, updated_channels
 
 
+def has_next_tasks(
+    trigger_to_nodes: Mapping[str, Sequence[str]],
+    updated_channels: set[str],
+    checkpoint: Checkpoint,
+) -> bool:
+    """Check if there are any tasks that should be run in the next step."""
+    return bool(checkpoint["pending_sends"]) or not updated_channels.isdisjoint(
+        trigger_to_nodes
+    )
+
+
 @overload
 def prepare_next_tasks(
     checkpoint: Checkpoint,
@@ -562,12 +570,9 @@ def prepare_single_task(
                         ),
                         CONFIG_KEY_READ: partial(
                             local_read,
-                            step,
-                            checkpoint,
                             channels,
                             managed,
                             PregelTaskWrites(task_path[:3], name, writes, triggers),
-                            config,
                         ),
                         CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)),
                         CONFIG_KEY_CHECKPOINTER: (
@@ -667,14 +672,11 @@ def prepare_single_task(
                             ),
                             CONFIG_KEY_READ: partial(
                                 local_read,
-                                step,
-                                checkpoint,
                                 channels,
                                 managed,
                                 PregelTaskWrites(
                                     task_path[:3], packet.node, writes, triggers
                                 ),
-                                config,
                             ),
                             CONFIG_KEY_STORE: (
                                 store or configurable.get(CONFIG_KEY_STORE)
@@ -789,8 +791,6 @@ def prepare_single_task(
                                 ),
                                 CONFIG_KEY_READ: partial(
                                     local_read,
-                                    step,
-                                    checkpoint,
                                     channels,
                                     managed,
                                     PregelTaskWrites(
@@ -799,7 +799,6 @@ def prepare_single_task(
                                         writes,
                                         triggers,
                                     ),
-                                    config,
                                 ),
                                 CONFIG_KEY_STORE: (
                                     store or configurable.get(CONFIG_KEY_STORE)

libs/langgraph/langgraph/pregel/loop.py

@@ -155,6 +155,8 @@ class PregelLoop(LoopProtocol):
     manager: Union[None, AsyncParentRunManager, ParentRunManager]
     interrupt_after: Union[All, Sequence[str]]
     interrupt_before: Union[All, Sequence[str]]
+    checkpoint_every_step: bool
+    debug: bool
 
     checkpointer_get_next_version: GetNextVersion
     checkpointer_put_writes: Optional[
@@ -211,6 +213,7 @@ class PregelLoop(LoopProtocol):
         input_model: Optional[Type[BaseModel]] = None,
         debug: bool = False,
         trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None,
+        checkpoint_every_step: bool = True,
     ) -> None:
         super().__init__(
             step=0,
@@ -235,6 +238,7 @@ class PregelLoop(LoopProtocol):
             or CONFIG_KEY_DEDUPE_TASKS in config[CONF]
         )
         self.trigger_to_nodes = trigger_to_nodes
+        self.checkpoint_every_step = checkpoint_every_step
         self.debug = debug
         if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]:
             self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM])
@@ -703,8 +707,6 @@ class PregelLoop(LoopProtocol):
         return updated_channels
 
     def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
-        for k, v in self.config["metadata"].items():
-            metadata.setdefault(k, v)  # type: ignore
         # assign step and parents
         metadata["step"] = self.step
         metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {})
@@ -719,10 +721,15 @@ class PregelLoop(LoopProtocol):
                     else self.stream_keys
                 ),
             )
-        # create new checkpoint
-        self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step)
         # bail if no checkpointer
         if self._checkpointer_put_after_previous is not None:
+            for k, v in self.config["metadata"].items():
+                metadata.setdefault(k, v)  # type: ignore
+
+            # create new checkpoint
+            self.checkpoint = create_checkpoint(
+                self.checkpoint, self.channels, self.step
+            )
             self.checkpoint_metadata = metadata
 
             self.prev_checkpoint_config = (