mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-29 11:19:54 +02:00
60385e452f
Cleaning up CI for #6701 ```md I'm trying to help a customer with some issues related to their checkpointing in postgres. They have some timeouts and retries around the langgraph checkpoint queries and I suspect the queue may be filling up with cancelled tasks. This PR adds some eager checks to not execute an operation if the future was already cancelled. This is to prevent the queue executing tasks that might have already timed out, which would otherwise cause more tasks to timeout due to the longer execution delay. Thank you for contributing to LangGraph! Follow these steps to mark your pull request as ready for review. **If any of these steps are not completed, your PR will not be considered for review.** ``` Co-authored-by: Conrad Ludgate <conradludgate@gmail.com>
372 lines
11 KiB
Python
372 lines
11 KiB
Python
"""Utilities for batching operations in a background task."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import functools
|
|
import weakref
|
|
from collections.abc import Callable, Iterable
|
|
from typing import Any, Literal, TypeVar
|
|
|
|
from langgraph.store.base import (
|
|
NOT_PROVIDED,
|
|
BaseStore,
|
|
GetOp,
|
|
Item,
|
|
ListNamespacesOp,
|
|
MatchCondition,
|
|
NamespacePath,
|
|
NotProvided,
|
|
Op,
|
|
PutOp,
|
|
Result,
|
|
SearchItem,
|
|
SearchOp,
|
|
_ensure_refresh,
|
|
_ensure_ttl,
|
|
_validate_namespace,
|
|
)
|
|
|
|
F = TypeVar("F", bound=Callable)
|
|
|
|
|
|
def _check_loop(func: F) -> F:
|
|
@functools.wraps(func)
|
|
def wrapper(store: AsyncBatchedBaseStore, *args: Any, **kwargs: Any) -> Any:
|
|
method_name: str = func.__name__
|
|
try:
|
|
current_loop = asyncio.get_running_loop()
|
|
if current_loop is store._loop:
|
|
replacement_str = (
|
|
f"Specifically, replace `store.{method_name}(...)` with `await store.a{method_name}(...)"
|
|
if method_name
|
|
else "For example, replace `store.get(...)` with `await store.aget(...)`"
|
|
)
|
|
raise asyncio.InvalidStateError(
|
|
f"Synchronous calls to {store.__class__.__name__} detected in the main event loop. "
|
|
"This can lead to deadlocks or performance issues. "
|
|
"Please use the asynchronous interface for main thread operations. "
|
|
f"{replacement_str} "
|
|
)
|
|
except RuntimeError:
|
|
pass
|
|
return func(store, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
class AsyncBatchedBaseStore(BaseStore):
|
|
"""Efficiently batch operations in a background task."""
|
|
|
|
__slots__ = ("_loop", "_aqueue", "_task")
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._loop = asyncio.get_running_loop()
|
|
self._aqueue: asyncio.Queue[tuple[asyncio.Future, Op]] = asyncio.Queue()
|
|
self._task: asyncio.Task | None = None
|
|
self._ensure_task()
|
|
|
|
def __del__(self) -> None:
|
|
try:
|
|
if self._task:
|
|
self._task.cancel()
|
|
except RuntimeError:
|
|
pass
|
|
|
|
def _ensure_task(self) -> None:
|
|
"""Ensure the background processing loop is running."""
|
|
if self._task is None or self._task.done():
|
|
self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self)))
|
|
|
|
async def aget(
|
|
self,
|
|
namespace: tuple[str, ...],
|
|
key: str,
|
|
*,
|
|
refresh_ttl: bool | None = None,
|
|
) -> Item | None:
|
|
self._ensure_task()
|
|
fut = self._loop.create_future()
|
|
self._aqueue.put_nowait(
|
|
(
|
|
fut,
|
|
GetOp(
|
|
namespace,
|
|
key,
|
|
refresh_ttl=_ensure_refresh(self.ttl_config, refresh_ttl),
|
|
),
|
|
)
|
|
)
|
|
return await fut
|
|
|
|
async def asearch(
|
|
self,
|
|
namespace_prefix: tuple[str, ...],
|
|
/,
|
|
*,
|
|
query: str | None = None,
|
|
filter: dict[str, Any] | None = None,
|
|
limit: int = 10,
|
|
offset: int = 0,
|
|
refresh_ttl: bool | None = None,
|
|
) -> list[SearchItem]:
|
|
self._ensure_task()
|
|
fut = self._loop.create_future()
|
|
self._aqueue.put_nowait(
|
|
(
|
|
fut,
|
|
SearchOp(
|
|
namespace_prefix,
|
|
filter,
|
|
limit,
|
|
offset,
|
|
query,
|
|
refresh_ttl=_ensure_refresh(self.ttl_config, refresh_ttl),
|
|
),
|
|
)
|
|
)
|
|
return await fut
|
|
|
|
async def aput(
|
|
self,
|
|
namespace: tuple[str, ...],
|
|
key: str,
|
|
value: dict[str, Any],
|
|
index: Literal[False] | list[str] | None = None,
|
|
*,
|
|
ttl: float | None | NotProvided = NOT_PROVIDED,
|
|
) -> None:
|
|
self._ensure_task()
|
|
_validate_namespace(namespace)
|
|
fut = self._loop.create_future()
|
|
self._aqueue.put_nowait(
|
|
(
|
|
fut,
|
|
PutOp(
|
|
namespace, key, value, index, ttl=_ensure_ttl(self.ttl_config, ttl)
|
|
),
|
|
)
|
|
)
|
|
return await fut
|
|
|
|
async def adelete(
|
|
self,
|
|
namespace: tuple[str, ...],
|
|
key: str,
|
|
) -> None:
|
|
self._ensure_task()
|
|
fut = self._loop.create_future()
|
|
self._aqueue.put_nowait((fut, PutOp(namespace, key, None)))
|
|
return await fut
|
|
|
|
async def alist_namespaces(
|
|
self,
|
|
*,
|
|
prefix: NamespacePath | None = None,
|
|
suffix: NamespacePath | None = None,
|
|
max_depth: int | None = None,
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
) -> list[tuple[str, ...]]:
|
|
self._ensure_task()
|
|
fut = self._loop.create_future()
|
|
match_conditions = []
|
|
if prefix:
|
|
match_conditions.append(MatchCondition(match_type="prefix", path=prefix))
|
|
if suffix:
|
|
match_conditions.append(MatchCondition(match_type="suffix", path=suffix))
|
|
|
|
op = ListNamespacesOp(
|
|
match_conditions=tuple(match_conditions),
|
|
max_depth=max_depth,
|
|
limit=limit,
|
|
offset=offset,
|
|
)
|
|
self._aqueue.put_nowait((fut, op))
|
|
return await fut
|
|
|
|
@_check_loop
|
|
def batch(self, ops: Iterable[Op]) -> list[Result]:
|
|
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self._loop).result()
|
|
|
|
@_check_loop
|
|
def get(
|
|
self,
|
|
namespace: tuple[str, ...],
|
|
key: str,
|
|
*,
|
|
refresh_ttl: bool | None = None,
|
|
) -> Item | None:
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.aget(namespace, key=key, refresh_ttl=refresh_ttl), self._loop
|
|
).result()
|
|
|
|
@_check_loop
|
|
def search(
|
|
self,
|
|
namespace_prefix: tuple[str, ...],
|
|
/,
|
|
*,
|
|
query: str | None = None,
|
|
filter: dict[str, Any] | None = None,
|
|
limit: int = 10,
|
|
offset: int = 0,
|
|
refresh_ttl: bool | None = None,
|
|
) -> list[SearchItem]:
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.asearch(
|
|
namespace_prefix,
|
|
query=query,
|
|
filter=filter,
|
|
limit=limit,
|
|
offset=offset,
|
|
refresh_ttl=refresh_ttl,
|
|
),
|
|
self._loop,
|
|
).result()
|
|
|
|
@_check_loop
|
|
def put(
|
|
self,
|
|
namespace: tuple[str, ...],
|
|
key: str,
|
|
value: dict[str, Any],
|
|
index: Literal[False] | list[str] | None = None,
|
|
*,
|
|
ttl: float | None | NotProvided = NOT_PROVIDED,
|
|
) -> None:
|
|
_validate_namespace(namespace)
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.aput(
|
|
namespace,
|
|
key=key,
|
|
value=value,
|
|
index=index,
|
|
ttl=_ensure_ttl(self.ttl_config, ttl),
|
|
),
|
|
self._loop,
|
|
).result()
|
|
|
|
@_check_loop
|
|
def delete(
|
|
self,
|
|
namespace: tuple[str, ...],
|
|
key: str,
|
|
) -> None:
|
|
asyncio.run_coroutine_threadsafe(
|
|
self.adelete(namespace, key=key), self._loop
|
|
).result()
|
|
|
|
@_check_loop
|
|
def list_namespaces(
|
|
self,
|
|
*,
|
|
prefix: NamespacePath | None = None,
|
|
suffix: NamespacePath | None = None,
|
|
max_depth: int | None = None,
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
) -> list[tuple[str, ...]]:
|
|
return asyncio.run_coroutine_threadsafe(
|
|
self.alist_namespaces(
|
|
prefix=prefix,
|
|
suffix=suffix,
|
|
max_depth=max_depth,
|
|
limit=limit,
|
|
offset=offset,
|
|
),
|
|
self._loop,
|
|
).result()
|
|
|
|
|
|
def _dedupe_ops(values: list[Op]) -> tuple[list[int] | None, list[Op]]:
|
|
"""Dedupe operations while preserving order for results.
|
|
|
|
Args:
|
|
values: List of operations to dedupe
|
|
|
|
Returns:
|
|
Tuple of (listen indices, deduped operations)
|
|
where listen indices map deduped operation results back to original positions
|
|
"""
|
|
if len(values) <= 1:
|
|
return None, list(values)
|
|
|
|
dedupped: list[Op] = []
|
|
listen: list[int] = []
|
|
puts: dict[tuple[tuple[str, ...], str], int] = {}
|
|
|
|
for op in values:
|
|
if isinstance(op, (GetOp, SearchOp, ListNamespacesOp)):
|
|
try:
|
|
listen.append(dedupped.index(op))
|
|
except ValueError:
|
|
listen.append(len(dedupped))
|
|
dedupped.append(op)
|
|
elif isinstance(op, PutOp):
|
|
putkey = (op.namespace, op.key)
|
|
if putkey in puts:
|
|
# Overwrite previous put
|
|
ix = puts[putkey]
|
|
dedupped[ix] = op
|
|
listen.append(ix)
|
|
else:
|
|
puts[putkey] = len(dedupped)
|
|
listen.append(len(dedupped))
|
|
dedupped.append(op)
|
|
|
|
else: # Any new ops will be treated regularly
|
|
listen.append(len(dedupped))
|
|
dedupped.append(op)
|
|
|
|
return listen, dedupped
|
|
|
|
|
|
async def _run(
|
|
aqueue: asyncio.Queue[tuple[asyncio.Future, Op]],
|
|
store: weakref.ReferenceType[BaseStore],
|
|
) -> None:
|
|
while item := await aqueue.get():
|
|
# don't run batch if the future is done (e.g. cancelled)
|
|
if item[0].done():
|
|
continue
|
|
# check if store is still alive
|
|
if s := store():
|
|
try:
|
|
# accumulate operations scheduled in same tick
|
|
items = [item]
|
|
try:
|
|
while item := aqueue.get_nowait():
|
|
# don't insert if the future is done (e.g. cancelled)
|
|
if item[0].done():
|
|
continue
|
|
items.append(item)
|
|
except asyncio.QueueEmpty:
|
|
pass
|
|
# get the operations to run
|
|
futs = [item[0] for item in items]
|
|
values = [item[1] for item in items]
|
|
# action each operation
|
|
try:
|
|
listen, dedupped = _dedupe_ops(values)
|
|
results = await s.abatch(dedupped)
|
|
if listen is not None:
|
|
results = [results[ix] for ix in listen]
|
|
|
|
# set the results of each operation
|
|
for fut, result in zip(futs, results, strict=False):
|
|
# guard against future being done (e.g. cancelled)
|
|
if not fut.done():
|
|
fut.set_result(result)
|
|
except Exception as e:
|
|
for fut in futs:
|
|
# guard against future being done (e.g. cancelled)
|
|
if not fut.done():
|
|
fut.set_exception(e)
|
|
finally:
|
|
# remove strong ref to store
|
|
del s
|
|
else:
|
|
break
|