Use incremental storage in memory checkpointer

- This makes our checkpoint benchmarks more closely resemble the behavior of our prod checkpointers
- Also found and fixed a bug w multiple subgraphs in same node accidentally sharing checkpoints

libs/checkpoint/langgraph/checkpoint/memory/__init__.py

@@ -70,6 +70,12 @@ class InMemorySaver(
         tuple[str, str, str],
         dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]],
     ]
+    blobs: dict[
+        tuple[
+            str, str, str, str | int | float
+        ],  # thread id, checkpoint ns, channel, version
+        tuple[str, bytes],
+    ]
 
     def __init__(
         self,
@@ -80,6 +86,7 @@ class InMemorySaver(
         super().__init__(serde=serde)
         self.storage = factory(lambda: defaultdict(dict))
         self.writes = factory(dict)
+        self.blobs = factory()
         self.stack = ExitStack()
         if factory is not defaultdict:
             self.stack.enter_context(self.storage)  # type: ignore[arg-type]
@@ -107,6 +114,18 @@ class InMemorySaver(
     ) -> Optional[bool]:
         return self.stack.__exit__(__exc_type, __exc_value, __traceback)
 
+    def _load_blobs(
+        self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions
+    ) -> dict[str, Any]:
+        channel_values: dict[str, Any] = {}
+        for k, v in versions.items():
+            kk = (thread_id, checkpoint_ns, k, v)
+            if kk in self.blobs:
+                vv = self.blobs[kk]
+                if vv[0] != "empty":
+                    channel_values[k] = self.serde.loads_typed(vv)
+        return channel_values
+
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         """Get a checkpoint tuple from the in-memory storage.
 
@@ -121,8 +140,8 @@ class InMemorySaver(
         Returns:
             Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
         """
-        thread_id = config["configurable"]["thread_id"]
-        checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
+        thread_id: str = config["configurable"]["thread_id"]
+        checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "")
         if checkpoint_id := get_checkpoint_id(config):
             if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id):
                 checkpoint, metadata, parent_checkpoint_id = saved
@@ -140,10 +159,14 @@ class InMemorySaver(
                     )
                 else:
                     sends = []
+                checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
                 return CheckpointTuple(
                     config=config,
                     checkpoint={
-                        **self.serde.loads_typed(checkpoint),
+                        **checkpoint_,
+                        "channel_values": self._load_blobs(
+                            thread_id, checkpoint_ns, checkpoint_["channel_versions"]
+                        ),
                         "pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
                     },
                     metadata=self.serde.loads_typed(metadata),
@@ -180,6 +203,9 @@ class InMemorySaver(
                     )
                 else:
                     sends = []
+
+                checkpoint_ = self.serde.loads_typed(checkpoint)
+
                 return CheckpointTuple(
                     config={
                         "configurable": {
@@ -189,7 +215,10 @@ class InMemorySaver(
                         }
                     },
                     checkpoint={
-                        **self.serde.loads_typed(checkpoint),
+                        **checkpoint_,
+                        "channel_values": self._load_blobs(
+                            thread_id, checkpoint_ns, checkpoint_["channel_versions"]
+                        ),
                         "pending_sends": [self.serde.loads_typed(s[2]) for s in sends],
                     },
                     metadata=self.serde.loads_typed(metadata),
@@ -297,6 +326,8 @@ class InMemorySaver(
                     else:
                         sends = []
 
+                    checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint)
+
                     yield CheckpointTuple(
                         config={
                             "configurable": {
@@ -306,7 +337,12 @@ class InMemorySaver(
                             }
                         },
                         checkpoint={
-                            **self.serde.loads_typed(checkpoint),
+                            **checkpoint_,
+                            "channel_values": self._load_blobs(
+                                thread_id,
+                                checkpoint_ns,
+                                checkpoint_["channel_versions"],
+                            ),
                             "pending_sends": [
                                 self.serde.loads_typed(s[2]) for s in sends
                             ],
@@ -353,6 +389,11 @@ class InMemorySaver(
         c.pop("pending_sends")  # type: ignore[misc]
         thread_id = config["configurable"]["thread_id"]
         checkpoint_ns = config["configurable"]["checkpoint_ns"]
+        values: dict[str, Any] = c.pop("channel_values")  # type: ignore[misc]
+        for k, v in new_versions.items():
+            self.blobs[(thread_id, checkpoint_ns, k, v)] = (
+                self.serde.dumps_typed(values[k]) if k in values else ("empty", b"")
+            )
         self.storage[thread_id][checkpoint_ns].update(
             {
                 checkpoint["id"]: (

libs/checkpoint/tests/test_memory.py

@@ -68,7 +68,9 @@ class TestMemorySaver:
             },
             "metadata": {"run_id": "my_run_id"},
         }
-        self.memory_saver.put(config, self.chkpnt_2, self.metadata_2, {})
+        self.memory_saver.put(
+            config, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"]
+        )
         checkpoint = self.memory_saver.get_tuple(config)
         assert checkpoint is not None
         assert checkpoint.metadata == {
@@ -80,9 +82,24 @@ class TestMemorySaver:
     async def test_search(self) -> None:
         # set up test
         # save checkpoints
-        self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
-        self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
-        self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
+        self.memory_saver.put(
+            self.config_1,
+            self.chkpnt_1,
+            self.metadata_1,
+            self.chkpnt_1["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_2,
+            self.chkpnt_2,
+            self.metadata_2,
+            self.chkpnt_2["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_3,
+            self.chkpnt_3,
+            self.metadata_3,
+            self.chkpnt_3["channel_versions"],
+        )
 
         # call method / assertions
         query_1 = {"source": "input"}  # search by 1 key
@@ -129,9 +146,24 @@ class TestMemorySaver:
     async def test_asearch(self) -> None:
         # set up test
         # save checkpoints
-        self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
-        self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
-        self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
+        self.memory_saver.put(
+            self.config_1,
+            self.chkpnt_1,
+            self.metadata_1,
+            self.chkpnt_1["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_2,
+            self.chkpnt_2,
+            self.metadata_2,
+            self.chkpnt_2["channel_versions"],
+        )
+        self.memory_saver.put(
+            self.config_3,
+            self.chkpnt_3,
+            self.metadata_3,
+            self.chkpnt_3["channel_versions"],
+        )
 
         # call method / assertions
         query_1 = {"source": "input"}  # search by 1 key

libs/langgraph/langgraph/pregel/loop.py

@@ -269,13 +269,13 @@ class PregelLoop(LoopProtocol):
             self.checkpoint_config = patch_configurable(
                 self.config,
                 {
-                    CONFIG_KEY_CHECKPOINT_ID: config[CONF][CONFIG_KEY_CHECKPOINT_MAP][
-                        self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]
-                    ]
+                    CONFIG_KEY_CHECKPOINT_ID: self.config[CONF][
+                        CONFIG_KEY_CHECKPOINT_MAP
+                    ][self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]]
                 },
             )
         else:
-            self.checkpoint_config = config
+            self.checkpoint_config = self.config
         self.checkpoint_ns = (
             tuple(cast(str, self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP))
             if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)

libs/langgraph/tests/memory_assert.py

@@ -1,4 +1,3 @@
-import asyncio
 import os
 import tempfile
 from collections import defaultdict
@@ -13,7 +12,6 @@ from langgraph.checkpoint.base import (
     CheckpointMetadata,
     CheckpointTuple,
     SerializerProtocol,
-    copy_checkpoint,
 )
 from langgraph.checkpoint.memory import InMemorySaver, PersistentDict
 
@@ -63,69 +61,14 @@ class MemorySaverAssertImmutable(InMemorySaver):
                     self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]]
                 )
                 == saved
-            )
+            ), config["configurable"]["checkpoint_ns"]
         self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = (
-            self.serde.dumps_typed(copy_checkpoint(checkpoint))
+            self.serde.dumps_typed(checkpoint)
         )
         # call super to write checkpoint
         return super().put(config, checkpoint, metadata, new_versions)
 
 
-class MemorySaverAssertCheckpointMetadata(InMemorySaver):
-    """This custom checkpointer is for verifying that a run's configurable
-    fields are merged with the previous checkpoint config for each step in
-    the run. This is the desired behavior. Because the checkpointer's (a)put()
-    method is called for each step, the implementation of this checkpointer
-    should produce a side effect that can be asserted.
-    """
-
-    def put(
-        self,
-        config: RunnableConfig,
-        checkpoint: Checkpoint,
-        metadata: CheckpointMetadata,
-        new_versions: ChannelVersions,
-    ) -> None:
-        """The implementation of put() merges config["configurable"] (a run's
-        configurable fields) with the metadata field. The state of the
-        checkpoint metadata can be asserted to confirm that the run's
-        configurable fields were merged with the previous checkpoint config.
-        """
-        configurable = config["configurable"].copy()
-
-        # remove checkpoint_id to make testing simpler
-        checkpoint_id = configurable.pop("checkpoint_id", None)
-        thread_id = config["configurable"]["thread_id"]
-        checkpoint_ns = config["configurable"]["checkpoint_ns"]
-        self.storage[thread_id][checkpoint_ns].update(
-            {
-                checkpoint["id"]: (
-                    self.serde.dumps_typed(checkpoint),
-                    # merge configurable fields and metadata
-                    self.serde.dumps_typed({**configurable, **metadata}),
-                    checkpoint_id,
-                )
-            }
-        )
-        return {
-            "configurable": {
-                "thread_id": config["configurable"]["thread_id"],
-                "checkpoint_id": checkpoint["id"],
-            }
-        }
-
-    async def aput(
-        self,
-        config: RunnableConfig,
-        checkpoint: Checkpoint,
-        metadata: CheckpointMetadata,
-        new_versions: ChannelVersions,
-    ) -> RunnableConfig:
-        return await asyncio.get_running_loop().run_in_executor(
-            None, self.put, config, checkpoint, metadata, new_versions
-        )
-
-
 class MemorySaverNoPending(InMemorySaver):
     def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
         result = super().get_tuple(config)

libs/langgraph/tests/test_pregel.py

@@ -84,7 +84,6 @@ from tests.conftest import (
     REGULAR_CHECKPOINTERS_SYNC,
     SHOULD_CHECK_SNAPSHOTS,
 )
-from tests.memory_assert import MemorySaverAssertCheckpointMetadata
 from tests.messages import (
     _AnyIdAIMessage,
     _AnyIdAIMessageChunk,
@@ -4213,11 +4212,11 @@ def test_checkpoint_metadata() -> None:
     workflow.add_edge("tools", "agent")
 
     # graph w/o interrupt
-    checkpointer_1 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_1 = InMemorySaver()
     app = workflow.compile(checkpointer=checkpointer_1)
 
     # graph w/ interrupt
-    checkpointer_2 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_2 = InMemorySaver()
     app_w_interrupt = workflow.compile(
         checkpointer=checkpointer_2, interrupt_before=["tools"]
     )
@@ -4635,59 +4634,6 @@ def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None:
     assert app.get_graph(xray=True).draw_mermaid() == snapshot
 
 
-def test_subgraph_retries():
-    class State(TypedDict):
-        count: int
-
-    class ChildState(State):
-        some_list: Annotated[list, operator.add]
-
-    called_times = 0
-
-    class RandomError(ValueError):
-        """This will be retried on."""
-
-    def parent_node(state: State):
-        return {"count": state["count"] + 1}
-
-    def child_node_a(state: ChildState):
-        nonlocal called_times
-        # We want it to retry only on node_b
-        # NOT re-compute the whole graph.
-        assert not called_times
-        called_times += 1
-        return {"some_list": ["val"]}
-
-    def child_node_b(state: ChildState):
-        raise RandomError("First attempt fails")
-
-    child = StateGraph(ChildState)
-    child.add_node(child_node_a)
-    child.add_node(child_node_b)
-    child.add_edge("__start__", "child_node_a")
-    child.add_edge("child_node_a", "child_node_b")
-
-    parent = StateGraph(State)
-    parent.add_node("parent_node", parent_node)
-    parent.add_node(
-        "child_graph",
-        child.compile(),
-        retry=RetryPolicy(
-            max_attempts=3,
-            retry_on=(RandomError,),
-            backoff_factor=0.0001,
-            initial_interval=0.0001,
-        ),
-    )
-    parent.add_edge("parent_node", "child_graph")
-    parent.set_entry_point("parent_node")
-
-    checkpointer = InMemorySaver()
-    app = parent.compile(checkpointer=checkpointer)
-    with pytest.raises(RandomError):
-        app.invoke({"count": 0}, {"configurable": {"thread_id": "foo"}})
-
-
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 @pytest.mark.parametrize("store_name", ALL_STORES_SYNC)
 def test_store_injected(
@@ -6294,6 +6240,7 @@ def test_double_interrupt_subgraph(
     def invoke_sub_agent(state: AgentState):
         return subgraph.invoke(state)
 
+    thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
     parent_agent = (
         StateGraph(AgentState)
         .add_node("invoke_sub_agent", invoke_sub_agent)

libs/langgraph/tests/test_pregel_async.py

@@ -77,10 +77,7 @@ from tests.conftest import (
     awith_store,
 )
 from tests.fake_tracer import FakeTracer
-from tests.memory_assert import (
-    MemorySaverAssertCheckpointMetadata,
-    MemorySaverNoPending,
-)
+from tests.memory_assert import MemorySaverNoPending
 from tests.messages import (
     _AnyIdAIMessage,
     _AnyIdAIMessageChunk,
@@ -5762,11 +5759,11 @@ async def test_checkpoint_metadata() -> None:
     workflow.add_edge("tools", "agent")
 
     # graph w/o interrupt
-    checkpointer_1 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_1 = InMemorySaver()
     app = workflow.compile(checkpointer=checkpointer_1)
 
     # graph w/ interrupt
-    checkpointer_2 = MemorySaverAssertCheckpointMetadata()
+    checkpointer_2 = InMemorySaver()
     app_w_interrupt = workflow.compile(
         checkpointer=checkpointer_2, interrupt_before=["tools"]
     )
@@ -7012,6 +7009,8 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None:
         def invoke_sub_agent(state: AgentState):
             return subgraph.invoke(state)
 
+        thread = {"configurable": {"thread_id": str(uuid.uuid4())}}
+
         parent_agent = (
             StateGraph(AgentState)
             .add_node("invoke_sub_agent", invoke_sub_agent)