WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-13 13:17:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Make end-of-step chckpointing costfree

Browse Source 
- instead of waiting for checkpoint put before proceeding, submit to executor / event loop and await all at the end
This commit is contained in:
Nuno Campos
2024-05-03 10:54:32 -07:00
parent 42d3215e34
commit 45149d2dc5
1 changed files with 22 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+22 -4
View File
@@ -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,