mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-25 00:52:25 +02:00
lib: Performance improvements
- don't create contextvars.Context/asyncio.Task in RunnableSeq (not needed as each step creates it if necessary) - don't run in-memory-saver methods in background threads (no point as they hold the gil) - avoid calling should_interrupt when no interrupts set
This commit is contained in:
@@ -1,4 +1,3 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
@@ -6,7 +5,6 @@ import random
|
||||
import shutil
|
||||
from collections import defaultdict
|
||||
from contextlib import AbstractAsyncContextManager, AbstractContextManager, ExitStack
|
||||
from functools import partial
|
||||
from types import TracebackType
|
||||
from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple, Type
|
||||
|
||||
@@ -395,9 +393,7 @@ class MemorySaver(
|
||||
Returns:
|
||||
Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
||||
"""
|
||||
return await asyncio.get_running_loop().run_in_executor(
|
||||
None, self.get_tuple, config
|
||||
)
|
||||
return self.get_tuple(config)
|
||||
|
||||
async def alist(
|
||||
self,
|
||||
@@ -418,24 +414,8 @@ class MemorySaver(
|
||||
Yields:
|
||||
AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
|
||||
"""
|
||||
loop = asyncio.get_running_loop()
|
||||
iter = await loop.run_in_executor(
|
||||
None,
|
||||
partial(
|
||||
self.list,
|
||||
before=before,
|
||||
limit=limit,
|
||||
filter=filter,
|
||||
),
|
||||
config,
|
||||
)
|
||||
while True:
|
||||
# handling StopIteration exception inside coroutine won't work
|
||||
# as expected, so using next() with default value to break the loop
|
||||
if item := await loop.run_in_executor(None, next, iter, None):
|
||||
yield item
|
||||
else:
|
||||
break
|
||||
for item in self.list(config, filter=filter, before=before, limit=limit):
|
||||
yield item
|
||||
|
||||
async def aput(
|
||||
self,
|
||||
@@ -455,9 +435,7 @@ class MemorySaver(
|
||||
Returns:
|
||||
RunnableConfig: The updated config containing the saved checkpoint's timestamp.
|
||||
"""
|
||||
return await asyncio.get_running_loop().run_in_executor(
|
||||
None, self.put, config, checkpoint, metadata, new_versions
|
||||
)
|
||||
return self.put(config, checkpoint, metadata, new_versions)
|
||||
|
||||
async def aput_writes(
|
||||
self,
|
||||
@@ -474,10 +452,8 @@ class MemorySaver(
|
||||
config (RunnableConfig): The config to associate with the writes.
|
||||
writes (List[Tuple[str, Any]]): The writes to save, each as a (channel, value) pair.
|
||||
task_id (str): Identifier for the task creating the writes.
|
||||
return self.put_writes(config, writes, task_id)
|
||||
"""
|
||||
return await asyncio.get_running_loop().run_in_executor(
|
||||
None, self.put_writes, config, writes, task_id
|
||||
)
|
||||
|
||||
def get_next_version(self, current: Optional[str], channel: ChannelProtocol) -> str:
|
||||
if current is None:
|
||||
|
||||
@@ -311,7 +311,9 @@ class PregelLoop(LoopProtocol):
|
||||
) -> Optional[PregelExecutableTask]:
|
||||
"""Accept a PUSH from a task, potentially returning a new task to start."""
|
||||
# don't start if we should interrupt *after* the original task
|
||||
if should_interrupt(self.checkpoint, self.interrupt_after, [task]):
|
||||
if self.interrupt_after and should_interrupt(
|
||||
self.checkpoint, self.interrupt_after, [task]
|
||||
):
|
||||
self.to_interrupt.append(task)
|
||||
return
|
||||
if pushed := cast(
|
||||
@@ -333,7 +335,9 @@ class PregelLoop(LoopProtocol):
|
||||
),
|
||||
):
|
||||
# don't start if we should interrupt *before* the new task
|
||||
if should_interrupt(self.checkpoint, self.interrupt_before, [pushed]):
|
||||
if self.interrupt_before and should_interrupt(
|
||||
self.checkpoint, self.interrupt_before, [pushed]
|
||||
):
|
||||
self.to_interrupt.append(pushed)
|
||||
return
|
||||
# produce debug output
|
||||
@@ -409,7 +413,7 @@ class PregelLoop(LoopProtocol):
|
||||
}
|
||||
)
|
||||
# after execution, check if we should interrupt
|
||||
if should_interrupt(
|
||||
if self.interrupt_after and should_interrupt(
|
||||
self.checkpoint, self.interrupt_after, self.tasks.values()
|
||||
):
|
||||
self.status = "interrupt_after"
|
||||
@@ -481,7 +485,7 @@ class PregelLoop(LoopProtocol):
|
||||
return self.tick(input_keys=input_keys)
|
||||
|
||||
# before execution, check if we should interrupt
|
||||
if should_interrupt(
|
||||
if self.interrupt_before and should_interrupt(
|
||||
self.checkpoint, self.interrupt_before, self.tasks.values()
|
||||
):
|
||||
self.status = "interrupt_before"
|
||||
|
||||
@@ -404,12 +404,10 @@ class RunnableSeq(Runnable):
|
||||
config = patch_config(
|
||||
config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
|
||||
)
|
||||
context = copy_context()
|
||||
context.run(_set_config_context, config)
|
||||
if i == 0:
|
||||
input = context.run(step.invoke, input, config, **kwargs)
|
||||
input = step.invoke(input, config, **kwargs)
|
||||
else:
|
||||
input = context.run(step.invoke, input, config)
|
||||
input = step.invoke(input, config)
|
||||
# finish the root run
|
||||
except BaseException as e:
|
||||
run_manager.on_chain_error(e)
|
||||
@@ -443,16 +441,10 @@ class RunnableSeq(Runnable):
|
||||
config = patch_config(
|
||||
config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
|
||||
)
|
||||
context = copy_context()
|
||||
context.run(_set_config_context, config)
|
||||
if i == 0:
|
||||
coro = step.ainvoke(input, config, **kwargs)
|
||||
input = await step.ainvoke(input, config, **kwargs)
|
||||
else:
|
||||
coro = step.ainvoke(input, config)
|
||||
if ASYNCIO_ACCEPTS_CONTEXT:
|
||||
input = await asyncio.create_task(coro, context=context)
|
||||
else:
|
||||
input = await asyncio.create_task(coro)
|
||||
input = await step.ainvoke(input, config)
|
||||
# finish the root run
|
||||
except BaseException as e:
|
||||
await run_manager.on_chain_error(e)
|
||||
|
||||
Reference in New Issue
Block a user