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:
Nuno Campos
2024-12-10 11:40:24 -08:00
parent 60d742ea48
commit 11e80210a2
3 changed files with 17 additions and 45 deletions
@@ -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:
+8 -4
View File
@@ -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"
+4 -12
View File
@@ -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)