From 4406828c39b5b16f4f53217aeca91c2288429735 Mon Sep 17 00:00:00 2001 From: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com> Date: Thu, 5 Dec 2024 18:27:47 -0800 Subject: [PATCH] test --- .../langgraph/store/postgres/aio.py | 15 +- .../tests/test_async_store.py | 101 +++++++++ .../langgraph/store/base/__init__.py | 4 + libs/checkpoint/langgraph/store/base/batch.py | 207 +++++++++++++++++- 4 files changed, 313 insertions(+), 14 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index 4a516557a..2b4309d5e 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -19,7 +19,11 @@ from langgraph.store.base import ( Result, SearchOp, ) -from langgraph.store.base.batch import AsyncBatchedBaseStore +from langgraph.store.base.batch import ( + BatchedBaseStore, + AsyncBatchedBaseStore, + SyncBatchedBaseStore, +) from langgraph.store.postgres.base import ( _PLACEHOLDER, BasePostgresStore, @@ -156,8 +160,15 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con return results + # def batch(self, ops: Iterable[Op]) -> list[Result]: + # return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result() def batch(self, ops: Iterable[Op]) -> list[Result]: - return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result() + futures = [] + for op in ops: + fut = self._loop.create_future() + self._aqueue[fut] = op + futures.append(fut) + return [fut.result() for fut in asyncio.as_completed(futures)] @classmethod @asynccontextmanager diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index eda0e2820..7c9cd38ad 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -1,4 +1,5 @@ # type: ignore +import asyncio import itertools import sys import uuid @@ -63,6 +64,106 @@ async def store(request) -> AsyncIterator[AsyncPostgresStore]: await conn.execute(f"DROP DATABASE {database}") +async def test_large_batches(store: AsyncPostgresStore) -> None: + N = 100 + M = 10 + coros = [] + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=10) as executor: + for m in range(M): + ops = [] + for i in range(N): + for i in range(N): + coros.append( + executor.submit( + store.put, + ("test", "foo", "bar", "baz", str(m % 2)), + f"key{i}", + value={"foo": "bar" + str(i)}, + ) + ) + coros.append( + executor.submit( + store.get, + ("test", "foo", "bar", "baz", str(m % 2)), + f"key{i}", + ) + ) + coros.append( + executor.submit( + store.list_namespaces, + prefix=None, + max_depth=m + 1, + ) + ) + coros.append( + executor.submit( + store.search, + ("test",), + ) + ) + coros.append( + executor.submit( + store.put, + ("test", "foo", "bar", "baz", str(m % 2)), + f"key{i}", + value={"foo": "bar" + str(i)}, + ) + ) + coros.append( + executor.submit( + store.put, + ("test", "foo", "bar", "baz", str(m % 2)), + f"key{i}", + None, + ) + ) + # ops.extend( + # [ + # PutOp( + # namespace=("test", "foo", "bar", "baz", str(m % 2)), + # key=f"key{i}", # {m}", + # value=None, + # ), + # GetOp(namespace=("test",), key=f"key{i}{m}"), + # ListNamespacesOp( + # match_conditions=None, max_depth=i + 1, limit=m, offset=0 + # ), + # SearchOp( + # namespace_prefix=("test",), + # filter=None, + # limit=10, + # offset=0, + # ), + # ] + # + # ops.extend( + # [ + # # PutOp( + # # namespace=("test", "foo", "bar", "baz", str(m % 2)), + # # key=f"key{i}", # {m}", + # # value={"data": f"value{i}{m}"}, + # # ), + # # GetOp(namespace=("test",), key=f"key{i}{m}"), + # # ListNamespacesOp( + # # match_conditions=None, max_depth=i + 1, limit=m, offset=0 + # # ), + # # SearchOp( + # # namespace_prefix=("test",), + # # filter={"data": f"value{i}{m}"}, + # # limit=10, + # # offset=0, + # # ), + # ] + # ) + + # coros.extend(ops) + # executor.map(store.batch, [[op] for op in coros]) + + # await asyncio.gather(*coros) # *[store.abatch(ops) for ops in coros]) + + async def test_abatch_order(store: AsyncPostgresStore) -> None: # Setup test data await store.aput(("test", "foo"), "key1", {"data": "value1"}) diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index c635e43ed..06eeab9cc 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -808,6 +808,8 @@ class BaseStore(ABC): # [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] ``` """ + if max_depth is not None and max_depth <= 0: + raise ValueError("If provided, max_depth must be greater than 0") match_conditions = [] if prefix: match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) @@ -1004,6 +1006,8 @@ class BaseStore(ABC): # Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")] ``` """ + if max_depth is not None and max_depth <= 0: + raise ValueError("If provided, max_depth must be greater than 0") match_conditions = [] if prefix: match_conditions.append(MatchCondition(match_type="prefix", path=prefix)) diff --git a/libs/checkpoint/langgraph/store/base/batch.py b/libs/checkpoint/langgraph/store/base/batch.py index 33c502574..a3e2d343b 100644 --- a/libs/checkpoint/langgraph/store/base/batch.py +++ b/libs/checkpoint/langgraph/store/base/batch.py @@ -1,6 +1,9 @@ import asyncio +import threading +import time import weakref -from typing import Any, Literal, Optional, Union +from concurrent.futures import Future +from typing import Any, Iterable, Literal, Optional, Union from langgraph.store.base import ( BaseStore, @@ -17,19 +20,9 @@ from langgraph.store.base import ( ) -class AsyncBatchedBaseStore(BaseStore): +class AsyncBatchedBaseStoreMixin: """Efficiently batch operations in a background task.""" - __slots__ = ("_loop", "_aqueue", "_task") - - def __init__(self) -> None: - self._loop = asyncio.get_running_loop() - self._aqueue: dict[asyncio.Future, Op] = {} - self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) - - def __del__(self) -> None: - self._task.cancel() - async def aget( self, namespace: tuple[str, ...], @@ -99,6 +92,29 @@ class AsyncBatchedBaseStore(BaseStore): self._aqueue[fut] = op return await fut + # def batch(self, ops: Iterable[Op]) -> list[Result]: + # futures = [] + # for op in ops: + # fut = self._loop.create_future() + # self._aqueue[fut] = op + # futures.append(fut) + # return [fut.result() for fut in asyncio.as_completed(futures)] + + +class AsyncBatchedBaseStore(AsyncBatchedBaseStoreMixin, 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: dict[asyncio.Future, Op] = {} + self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) + + def __del__(self) -> None: + self._task.cancel() + def _dedupe_ops(values: list[Op]) -> tuple[Optional[list[int]], list[Op]]: """Dedupe operations while preserving order for results. @@ -174,3 +190,170 @@ async def _run( break # remove strong ref to store del s + + +class SyncBatchedBaseStoreMixin(BaseStore): + """Efficiently batch operations in a background thread.""" + + def get( + self, + namespace: tuple[str, ...], + key: str, + ) -> Optional[Item]: + fut = Future() + self._queue[fut] = GetOp(namespace, key) + return fut.result() + + def search( + self, + namespace_prefix: tuple[str, ...], + /, + *, + query: Optional[str] = None, + filter: Optional[dict[str, Any]] = None, + limit: int = 10, + offset: int = 0, + ) -> list[SearchItem]: + fut = Future() + self._queue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query) + return fut.result() + + def put( + self, + namespace: tuple[str, ...], + key: str, + value: dict[str, Any], + index: Optional[Union[Literal[False], list[str]]] = None, + ) -> None: + _validate_namespace(namespace) + fut = Future() + self._queue[fut] = PutOp(namespace, key, value, index) + return fut.result() + + def delete( + self, + namespace: tuple[str, ...], + key: str, + ) -> None: + fut = Future() + self._queue[fut] = PutOp(namespace, key, None) + return fut.result() + + def list_namespaces( + self, + *, + prefix: Optional[NamespacePath] = None, + suffix: Optional[NamespacePath] = None, + max_depth: Optional[int] = None, + limit: int = 100, + offset: int = 0, + ) -> list[tuple[str, ...]]: + fut = 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._queue[fut] = op + return fut.result() + + +class SyncBatchedBaseStore(SyncBatchedBaseStoreMixin, BaseStore): + """Efficiently batch operations in a background thread.""" + + __slots__ = ("_sync_queue", "_sync_thread") + + def __init__(self) -> None: + super().__init__() + self._sync_queue: dict[Future, Op] = {} + self._sync_thread = threading.Thread( + target=_sync_run, + args=(self._sync_queue, weakref.ref(self)), + daemon=True, + ) + self._sync_thread.start() + + def __del__(self) -> None: + # Signal the thread to stop + if self._sync_thread.is_alive(): + empty_future = Future() + self._sync_queue[empty_future] = None # type: ignore + self._sync_thread.join() + + +class BatchedBaseStore( + AsyncBatchedBaseStoreMixin, SyncBatchedBaseStoreMixin, BaseStore +): + __slots__ = ( + "_sync_queue", + "_sync_thread", + "_task", + "_loop", + "_aqueue", + ) + + def __init__(self) -> None: + super().__init__() + # Setup async processing + self._loop = asyncio.get_running_loop() + self._aqueue: dict[asyncio.Future, Op] = {} + self._task = self._loop.create_task(_run(self._aqueue, weakref.ref(self))) + + self._sync_queue: dict[Future, Op] = {} + self._sync_thread = threading.Thread( + target=_sync_run, + args=(self._sync_queue, weakref.ref(self)), + daemon=True, + ) + self._sync_thread.start() + + def __del__(self) -> None: + # Signal the thread to stop + if self._sync_thread.is_alive(): + empty_future = Future() + self._sync_queue[empty_future] = None # type: ignore + self._sync_thread.join() + + # Signal the thread to stop + if self._task is not None: + self._task.cancel() + + +def _sync_run(queue: dict[Future, Op], store: weakref.ReferenceType[BaseStore]) -> None: + while True: + time.sleep(0.001) # Yield to other threads + if not queue: + continue + if s := store(): + # get the operations to run + taken = queue.copy() + # action each operation + try: + values = list(taken.values()) + if None in values: # Exit signal + break + listen, dedupped = _dedupe_ops(values) + results = s.batch(dedupped) # Note: Using sync batch here + 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 queue[fut] + else: + break + # remove strong ref to store + del s