From 9dbcb031858125271dc4b14f08947ae08ccb4519 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Feb 2025 10:36:02 -0800 Subject: [PATCH 1/5] Fix busy loop in AsyncBatchedBaseStore - while loop w asyncio.sleep(0) takes up cpu --- libs/checkpoint/langgraph/store/base/batch.py | 32 ++++++++++++------- 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 6cfc11419..2e0435f15 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -1,7 +1,8 @@ import asyncio import functools import weakref -from typing import Any, Callable, Iterable, Literal, Optional, TypeVar, Union +from collections.abc import Iterable +from typing import Any, Callable, Literal, Optional, TypeVar, Union from langgraph.store.base import ( BaseStore, @@ -55,9 +56,11 @@ class AsyncBatchedBaseStore(BaseStore): super().__init__() self._loop = asyncio.get_running_loop() self._aqueue: dict[asyncio.Future, Op] = {} + self._aqueue_ = asyncio.Queue() self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) def __del__(self) -> None: + print("cancelling task") self._task.cancel() async def aget( @@ -66,7 +69,7 @@ class AsyncBatchedBaseStore(BaseStore): key: str, ) -> Optional[Item]: fut = self._loop.create_future() - self._aqueue[fut] = GetOp(namespace, key) + self._aqueue_.put((fut, GetOp(namespace, key))) return await fut async def asearch( @@ -80,7 +83,9 @@ class AsyncBatchedBaseStore(BaseStore): offset: int = 0, ) -> list[SearchItem]: fut = self._loop.create_future() - self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query) + self._aqueue_.put( + (fut, SearchOp(namespace_prefix, filter, limit, offset, query)) + ) return await fut async def aput( @@ -92,7 +97,7 @@ class AsyncBatchedBaseStore(BaseStore): ) -> None: _validate_namespace(namespace) fut = self._loop.create_future() - self._aqueue[fut] = PutOp(namespace, key, value, index) + self._aqueue_.put((fut, PutOp(namespace, key, value, index))) return await fut async def adelete( @@ -101,7 +106,7 @@ class AsyncBatchedBaseStore(BaseStore): key: str, ) -> None: fut = self._loop.create_future() - self._aqueue[fut] = PutOp(namespace, key, None) + self._aqueue_.put((fut, PutOp(namespace, key, None))) return await fut async def alist_namespaces( @@ -126,7 +131,7 @@ class AsyncBatchedBaseStore(BaseStore): limit=limit, offset=offset, ) - self._aqueue[fut] = op + self._aqueue_.put((fut, op)) return await fut @_check_loop @@ -250,16 +255,19 @@ def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]: async def _run( - aqueue: dict[asyncio.Future, Op], + aqueue: asyncio.Queue[tuple[asyncio.Future, Op]], store: weakref.ReferenceType[BaseStore], ) -> None: - while True: - await asyncio.sleep(0) - if not aqueue: - continue + while item := await aqueue.get(): + # check if store is still alive if s := store(): + # accumulate operations scheduled "concurrently" + items = [item] + await asyncio.sleep(0) + while item := aqueue.get_nowait(): + items.append(item) # get the operations to run - taken = aqueue.copy() + taken = dict(items) # action each operation try: values = list(taken.values()) From b5479b48bf8bb1258be1731751c9680d69dc51ce Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Feb 2025 10:39:34 -0800 Subject: [PATCH 2/5] Lint --- libs/checkpoint/langgraph/store/base/batch.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 2e0435f15..9af2b519b 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -60,7 +60,6 @@ class AsyncBatchedBaseStore(BaseStore): self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) def __del__(self) -> None: - print("cancelling task") self._task.cancel() async def aget( From a064ccdca13f75d87486dd74584d94751c303117 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Feb 2025 10:40:09 -0800 Subject: [PATCH 3/5] Lint --- libs/checkpoint/langgraph/store/base/batch.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 9af2b519b..d3e56eeb3 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -55,8 +55,7 @@ class AsyncBatchedBaseStore(BaseStore): def __init__(self) -> None: super().__init__() self._loop = asyncio.get_running_loop() - self._aqueue: dict[asyncio.Future, Op] = {} - self._aqueue_ = asyncio.Queue() + self._aqueue: asyncio.Queue[tuple[asyncio.Future, Op]] = asyncio.Queue() self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) def __del__(self) -> None: @@ -68,7 +67,7 @@ class AsyncBatchedBaseStore(BaseStore): key: str, ) -> Optional[Item]: fut = self._loop.create_future() - self._aqueue_.put((fut, GetOp(namespace, key))) + self._aqueue.put((fut, GetOp(namespace, key))) return await fut async def asearch( @@ -82,7 +81,7 @@ class AsyncBatchedBaseStore(BaseStore): offset: int = 0, ) -> list[SearchItem]: fut = self._loop.create_future() - self._aqueue_.put( + self._aqueue.put( (fut, SearchOp(namespace_prefix, filter, limit, offset, query)) ) return await fut @@ -96,7 +95,7 @@ class AsyncBatchedBaseStore(BaseStore): ) -> None: _validate_namespace(namespace) fut = self._loop.create_future() - self._aqueue_.put((fut, PutOp(namespace, key, value, index))) + self._aqueue.put((fut, PutOp(namespace, key, value, index))) return await fut async def adelete( @@ -105,7 +104,7 @@ class AsyncBatchedBaseStore(BaseStore): key: str, ) -> None: fut = self._loop.create_future() - self._aqueue_.put((fut, PutOp(namespace, key, None))) + self._aqueue.put((fut, PutOp(namespace, key, None))) return await fut async def alist_namespaces( @@ -130,7 +129,7 @@ class AsyncBatchedBaseStore(BaseStore): limit=limit, offset=offset, ) - self._aqueue_.put((fut, op)) + self._aqueue.put((fut, op)) return await fut @_check_loop From 6baf320d8e3ace9e9de1ac27d11323aee450fcfa Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Feb 2025 10:43:03 -0800 Subject: [PATCH 4/5] Fix --- libs/checkpoint/langgraph/store/base/batch.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index d3e56eeb3..d15364f24 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -67,7 +67,7 @@ class AsyncBatchedBaseStore(BaseStore): key: str, ) -> Optional[Item]: fut = self._loop.create_future() - self._aqueue.put((fut, GetOp(namespace, key))) + self._aqueue.put_nowait((fut, GetOp(namespace, key))) return await fut async def asearch( @@ -81,7 +81,7 @@ class AsyncBatchedBaseStore(BaseStore): offset: int = 0, ) -> list[SearchItem]: fut = self._loop.create_future() - self._aqueue.put( + self._aqueue.put_nowait( (fut, SearchOp(namespace_prefix, filter, limit, offset, query)) ) return await fut @@ -95,7 +95,7 @@ class AsyncBatchedBaseStore(BaseStore): ) -> None: _validate_namespace(namespace) fut = self._loop.create_future() - self._aqueue.put((fut, PutOp(namespace, key, value, index))) + self._aqueue.put_nowait((fut, PutOp(namespace, key, value, index))) return await fut async def adelete( @@ -104,7 +104,7 @@ class AsyncBatchedBaseStore(BaseStore): key: str, ) -> None: fut = self._loop.create_future() - self._aqueue.put((fut, PutOp(namespace, key, None))) + self._aqueue.put_nowait((fut, PutOp(namespace, key, None))) return await fut async def alist_namespaces( @@ -129,7 +129,7 @@ class AsyncBatchedBaseStore(BaseStore): limit=limit, offset=offset, ) - self._aqueue.put((fut, op)) + self._aqueue.put_nowait((fut, op)) return await fut @_check_loop From da97d2e1baa9c47581ad0e8647cfafd14abec35b Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Feb 2025 12:03:03 -0800 Subject: [PATCH 5/5] Fix --- libs/checkpoint/langgraph/store/base/batch.py | 54 ++++++++++--------- 1 file changed, 30 insertions(+), 24 deletions(-) diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index d15364f24..1bf2ac79e 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -66,6 +66,7 @@ class AsyncBatchedBaseStore(BaseStore): namespace: tuple[str, ...], key: str, ) -> Optional[Item]: + assert not self._task.done() fut = self._loop.create_future() self._aqueue.put_nowait((fut, GetOp(namespace, key))) return await fut @@ -80,6 +81,7 @@ class AsyncBatchedBaseStore(BaseStore): limit: int = 10, offset: int = 0, ) -> list[SearchItem]: + assert not self._task.done() fut = self._loop.create_future() self._aqueue.put_nowait( (fut, SearchOp(namespace_prefix, filter, limit, offset, query)) @@ -93,6 +95,7 @@ class AsyncBatchedBaseStore(BaseStore): value: dict[str, Any], index: Optional[Union[Literal[False], list[str]]] = None, ) -> None: + assert not self._task.done() _validate_namespace(namespace) fut = self._loop.create_future() self._aqueue.put_nowait((fut, PutOp(namespace, key, value, index))) @@ -103,6 +106,7 @@ class AsyncBatchedBaseStore(BaseStore): namespace: tuple[str, ...], key: str, ) -> None: + assert not self._task.done() fut = self._loop.create_future() self._aqueue.put_nowait((fut, PutOp(namespace, key, None))) return await fut @@ -116,6 +120,7 @@ class AsyncBatchedBaseStore(BaseStore): limit: int = 100, offset: int = 0, ) -> list[tuple[str, ...]]: + assert not self._task.done() fut = self._loop.create_future() match_conditions = [] if prefix: @@ -259,31 +264,32 @@ async def _run( while item := await aqueue.get(): # check if store is still alive if s := store(): - # accumulate operations scheduled "concurrently" - items = [item] - await asyncio.sleep(0) - while item := aqueue.get_nowait(): - items.append(item) - # get the operations to run - taken = dict(items) - # action each operation try: - values = list(taken.values()) - listen, dedupped = _dedupe_ops(values) - results = await s.abatch(dedupped) - if listen is not None: - results = [results[ix] for ix in listen] + # accumulate operations scheduled in same tick + items = [item] + try: + while item := aqueue.get_nowait(): + 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(taken, results): - fut.set_result(result) - except Exception as e: - for fut in taken: - fut.set_exception(e) - # remove the operations from the queue - for fut in taken: - del aqueue[fut] + # set the results of each operation + for fut, result in zip(futs, results): + fut.set_result(result) + except Exception as e: + for fut in futs: + fut.set_exception(e) + finally: + # remove strong ref to store + del s else: break - # remove strong ref to store - del s