From 45149d2dc5eef8b13266612c0104989883058e77 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Fri, 3 May 2024 10:54:32 -0700
Subject: [PATCH] Make end-of-step chckpointing costfree

- instead of waiting for checkpoint put before proceeding, submit to executor / event loop and await all at the end
---
 langgraph/pregel/__init__.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py
index 143052d0f..d83eb906b 100644
--- a/langgraph/pregel/__init__.py
+++ b/langgraph/pregel/__init__.py
@@ -799,9 +799,15 @@ class Pregel(
                     and self.checkpointer.at == CheckpointAt.END_OF_RUN
                 ):
                     checkpoint = create_checkpoint(checkpoint, channels)
-                    checkpoint_config = self.checkpointer.put(
-                        checkpoint_config, checkpoint
+                    executor.submit(
+                        self.checkpointer.put(checkpoint_config, checkpoint)
                     )
+                    checkpoint_config = {
+                        "configurable": {
+                            "thread_id": checkpoint_config["configurable"]["thread_id"],
+                            "thread_ts": checkpoint["ts"],
+                        }
+                    }
                     if stream_mode == "debug":
                         yield map_debug_checkpoint(
                             step,
@@ -854,6 +860,7 @@ class Pregel(
             None,
         )
         try:
+            tasks: list[asyncio.Task] = []
             if config["recursion_limit"] < 1:
                 raise ValueError("recursion_limit must be at least 1")
             if self.checkpointer and not config.get("configurable"):
@@ -1058,9 +1065,17 @@ class Pregel(
                     and self.checkpointer.at == CheckpointAt.END_OF_RUN
                 ):
                     checkpoint = create_checkpoint(checkpoint, channels)
-                    checkpoint_config = await self.checkpointer.aput(
-                        checkpoint_config, checkpoint
+                    tasks.append(
+                        asyncio.create_task(
+                            self.checkpointer.aput(checkpoint_config, checkpoint)
+                        )
                     )
+                    checkpoint_config = {
+                        "configurable": {
+                            "thread_id": checkpoint_config["configurable"]["thread_id"],
+                            "thread_ts": checkpoint["ts"],
+                        }
+                    }
                     if stream_mode == "debug":
                         yield map_debug_checkpoint(
                             step, checkpoint_config, channels, self.stream_channels_asis
@@ -1077,8 +1092,11 @@ class Pregel(
             try:
                 for task in futures:
                     task.cancel()
+                    tasks.append(task)
             except NameError:
                 pass
+            # wait for all tasks to finish
+            await asyncio.gather(*tasks, return_exceptions=True)
 
     def invoke(
         self,