From 3a2a4e44b2dfd9c6e93644026fbc921f89110771 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 17 Sep 2024 12:50:29 -0700 Subject: [PATCH] Stream output from subgraphs while in-progress for sync stream - previously implemented only for async stream --- libs/langgraph/Makefile | 6 +- libs/langgraph/langgraph/pregel/__init__.py | 38 ++++-- libs/langgraph/langgraph/pregel/runner.py | 46 ++++--- libs/langgraph/langgraph/utils/aio.py | 35 ------ libs/langgraph/langgraph/utils/queue.py | 127 ++++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 80 ++++++++++++ 6 files changed, 270 insertions(+), 62 deletions(-) delete mode 100644 libs/langgraph/langgraph/utils/aio.py create mode 100644 libs/langgraph/langgraph/utils/queue.py diff --git a/libs/langgraph/Makefile b/libs/langgraph/Makefile index 0173351d6..3d8a175e7 100644 --- a/libs/langgraph/Makefile +++ b/libs/langgraph/Makefile @@ -40,16 +40,16 @@ start-postgres: stop-postgres: docker compose -f tests/compose-postgres.yml down -v -TEST_PATH ?= . +TEST ?= . test: - make start-postgres && poetry run pytest $(TEST_PATH); \ + make start-postgres && poetry run pytest $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE test_watch: - make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST_PATH); \ + make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST); \ EXIT_CODE=$$?; \ make stop-postgres; \ exit $$EXIT_CODE diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 315934429..1cf5ec982 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1,6 +1,9 @@ from __future__ import annotations import asyncio +import concurrent +import concurrent.futures +import queue from collections import deque from functools import partial from typing import ( @@ -83,7 +86,6 @@ from langgraph.pregel.utils import get_new_channel_versions from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.utils.aio import Queue from langgraph.utils.config import ( ensure_config, merge_configs, @@ -92,6 +94,7 @@ from langgraph.utils.config import ( patch_configurable, ) from langgraph.utils.pydantic import create_model +from langgraph.utils.queue import AsyncQueue, SyncQueue from langgraph.utils.runnable import RunnableCallable WriteValue = Union[Callable[[Input], Output], Any] @@ -1162,11 +1165,14 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): ``` """ - stream = deque() + stream = SyncQueue() def output() -> Iterator: - while stream: - ns, mode, payload = stream.popleft() + while True: + try: + ns, mode, payload = stream.get(block=False) + except queue.Empty: + break if subgraphs and isinstance(stream_mode, list): yield (ns, mode, payload) elif isinstance(stream_mode, list): @@ -1210,7 +1216,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): with SyncPregelLoop( input, - stream=StreamProtocol(stream.append, stream_modes), + stream=StreamProtocol(stream.put, stream_modes), config=config, store=self.store, checkpointer=checkpointer, @@ -1228,6 +1234,22 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): # enable subgraph streaming if subgraphs: loop.config["configurable"][CONFIG_KEY_STREAM] = loop.stream + # we are careful to have a single waiter live at any one time + # because on exit we increment semaphore count by exactly 1 + waiter: Optional[concurrent.futures.Future] = None + # because sync futures cannot be cancelled, we instead + # release the stream semaphore on exit, which will cause + # a pending waiter to return immediately + loop.stack.callback(stream._count.release) + + def get_waiter() -> asyncio.Task[None]: + nonlocal waiter + if waiter is None or waiter.done(): + return (waiter := loop.submit(stream.wait)) + else: + return waiter + else: + get_waiter = None # Similarly to Bulk Synchronous Parallel / Pregel model # computation proceeds in steps, while there are channel updates # channel updates from step N are only visible in step N+1 @@ -1243,10 +1265,10 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): loop.tasks.values(), timeout=self.step_timeout, retry_policy=self.retry_policy, + get_waiter=get_waiter, ): # emit output - for o in output(): - yield o + yield from output() # emit output yield from output() # handle exit @@ -1342,7 +1364,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): ``` """ - stream = Queue() + stream = AsyncQueue() aioloop = asyncio.get_running_loop() def output() -> Iterator: diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index fc2d7e464..1f282c584 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -38,6 +38,7 @@ class PregelRunner: reraise: bool = True, timeout: Optional[float] = None, retry_policy: Optional[RetryPolicy] = None, + get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: tasks = tuple(tasks) # give control back to the caller @@ -53,23 +54,30 @@ class PregelRunner: if reraise: raise return + # add waiter task if requested + if get_waiter is not None: + futures: dict[concurrent.futures.Future, Optional[PregelExecutableTask]] = { + get_waiter(): None + } + else: + futures = {} # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes - futures = { - self.submit( - run_with_retry, - task, - retry_policy, - __reraise_on_exit__=reraise, - ): task - for task in tasks - if not task.writes - } + for task in tasks: + if not task.writes: + futures[ + self.submit( + run_with_retry, + task, + retry_policy, + __reraise_on_exit__=reraise, + ) + ] = task all_futures = futures.copy() end_time = timeout + time.monotonic() if timeout else None - while futures: - done, _ = concurrent.futures.wait( + while len(futures) > (1 if get_waiter is not None else 0): + done, inflight = concurrent.futures.wait( futures, return_when=concurrent.futures.FIRST_COMPLETED, timeout=(max(0, end_time - time.monotonic()) if end_time else None), @@ -78,8 +86,13 @@ class PregelRunner: break # timed out for fut in done: task = futures.pop(fut) - # task finished, commit writes - self.commit(task, _exception(fut)) + if task is None: + # waiter task finished, schedule another + if inflight: + futures[get_waiter()] = None + else: + # task finished, commit writes + self.commit(task, _exception(fut)) else: # remove references to loop vars del fut, task @@ -141,7 +154,7 @@ class PregelRunner: all_futures = futures.copy() end_time = timeout + loop.time() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): - done, _ = await asyncio.wait( + done, inflight = await asyncio.wait( futures, return_when=asyncio.FIRST_COMPLETED, timeout=(max(0, end_time - loop.time()) if end_time else None), @@ -152,7 +165,8 @@ class PregelRunner: task = futures.pop(fut) if task is None: # waiter task finished, schedule another - futures[get_waiter()] = None + if inflight: + futures[get_waiter()] = None else: # task finished, commit writes self.commit(task, _exception(fut)) diff --git a/libs/langgraph/langgraph/utils/aio.py b/libs/langgraph/langgraph/utils/aio.py deleted file mode 100644 index afe02c050..000000000 --- a/libs/langgraph/langgraph/utils/aio.py +++ /dev/null @@ -1,35 +0,0 @@ -import asyncio -import sys - -PY_310 = sys.version_info >= (3, 10) - - -class Queue(asyncio.Queue): - async def wait(self): - """If queue is empty, wait until an item is available. - - Copied from Queue.get(), removing the call to .get_nowait(), - ie. this doesn't consume the item, just waits for it. - """ - while self.empty(): - if PY_310: - getter = self._get_loop().create_future() - else: - getter = self._loop.create_future() - self._getters.append(getter) - try: - await getter - except: - getter.cancel() # Just in case getter is not done yet. - try: - # Clean self._getters from canceled getters. - self._getters.remove(getter) - except ValueError: - # The getter could be removed from self._getters by a - # previous put_nowait call. - pass - if not self.empty() and not getter.cancelled(): - # We were woken up by put_nowait(), but can't take - # the call. Wake up the next in line. - self._wakeup_next(self._getters) - raise diff --git a/libs/langgraph/langgraph/utils/queue.py b/libs/langgraph/langgraph/utils/queue.py new file mode 100644 index 000000000..99d94f5c4 --- /dev/null +++ b/libs/langgraph/langgraph/utils/queue.py @@ -0,0 +1,127 @@ +import asyncio +import queue +import sys +import threading +import types +from collections import deque +from time import monotonic + +PY_310 = sys.version_info >= (3, 10) + + +class AsyncQueue(asyncio.Queue): + """Async unbounded FIFO queue with a wait() method. + + Subclassed from asyncio.Queue, adding a wait() method.""" + + async def wait(self): + """If queue is empty, wait until an item is available. + + Copied from Queue.get(), removing the call to .get_nowait(), + ie. this doesn't consume the item, just waits for it. + """ + while self.empty(): + if PY_310: + getter = self._get_loop().create_future() + else: + getter = self._loop.create_future() + self._getters.append(getter) + try: + await getter + except: + getter.cancel() # Just in case getter is not done yet. + try: + # Clean self._getters from canceled getters. + self._getters.remove(getter) + except ValueError: + # The getter could be removed from self._getters by a + # previous put_nowait call. + pass + if not self.empty() and not getter.cancelled(): + # We were woken up by put_nowait(), but can't take + # the call. Wake up the next in line. + self._wakeup_next(self._getters) + raise + + +class Semaphore(threading.Semaphore): + """Semaphore subclass with a wait() method.""" + + def wait(self, blocking: bool = True, timeout: float = None): + """Block until the semaphore can be acquired, but don't acquire it.""" + if not blocking and timeout is not None: + raise ValueError("can't specify timeout for non-blocking acquire") + rc = False + endtime = None + with self._cond: + while self._value == 0: + if not blocking: + break + if timeout is not None: + if endtime is None: + endtime = monotonic() + timeout + else: + timeout = endtime - monotonic() + if timeout <= 0: + break + self._cond.wait(timeout) + else: + rc = True + return rc + + +class SyncQueue: + """Unbounded FIFO queue with a wait() method. + Adapted from pure Python implementation of queue.SimpleQueue. + """ + + def __init__(self): + self._queue = deque() + self._count = Semaphore(0) + + def put(self, item, block=True, timeout=None): + """Put the item on the queue. + + The optional 'block' and 'timeout' arguments are ignored, as this method + never blocks. They are provided for compatibility with the Queue class. + """ + self._queue.append(item) + self._count.release() + + def get(self, block=True, timeout=None): + """Remove and return an item from the queue. + + If optional args 'block' is true and 'timeout' is None (the default), + block if necessary until an item is available. If 'timeout' is + a non-negative number, it blocks at most 'timeout' seconds and raises + the Empty exception if no item was available within that time. + Otherwise ('block' is false), return an item if one is immediately + available, else raise the Empty exception ('timeout' is ignored + in that case). + """ + if timeout is not None and timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + if not self._count.acquire(block, timeout): + raise queue.Empty + try: + return self._queue.popleft() + except IndexError: + raise queue.Empty + + def wait(self, block=True, timeout=None): + """If queue is empty, wait until an item maybe is available, + but don't consume it. + """ + if timeout is not None and timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + self._count.wait(block, timeout) + + def empty(self): + """Return True if the queue is empty, False otherwise (not reliable!).""" + return len(self._queue) == 0 + + def qsize(self): + """Return the approximate size of the queue (not reliable!).""" + return len(self._queue) + + __class_getitem__ = classmethod(types.GenericAlias) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 3e3ebc72f..9096803a5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8362,6 +8362,86 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_stream_subgraphs_during_execution( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": "got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + time.sleep(0.5) + return { + "my_key": " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: Annotated[str, operator.add] + + def outer_1(state: State): + time.sleep(0.2) + return {"my_key": " and parallel"} + + def outer_2(state: State): + return {"my_key": " and back again"} + + graph = StateGraph(State) + graph.add_node("inner", inner.compile()) + graph.add_node("outer_1", outer_1) + graph.add_node("outer_2", outer_2) + + graph.add_edge(START, "inner") + graph.add_edge(START, "outer_1") + graph.add_edge(["inner", "outer_1"], "outer_2") + graph.add_edge("outer_2", END) + + app = graph.compile(checkpointer=checkpointer) + + start = time.perf_counter() + chunks: list[tuple[float, Any]] = [] + config = {"configurable": {"thread_id": "2"}} + for c in app.stream({"my_key": ""}, config, subgraphs=True): + chunks.append((round(time.perf_counter() - start, 1), c)) + for idx in range(len(chunks)): + elapsed, c = chunks[idx] + chunks[idx] = (round(elapsed - chunks[0][0], 1), c) + + assert chunks == [ + # arrives before "inner" finishes + ( + 0.0, + ( + (AnyStr("inner:"),), + {"inner_1": {"my_key": "got here", "my_other_key": ""}}, + ), + ), + (0.2, ((), {"outer_1": {"my_key": " and parallel"}})), + ( + 0.5, + ( + (AnyStr("inner:"),), + {"inner_2": {"my_key": " and there", "my_other_key": "got here"}}, + ), + ), + (0.5, ((), {"inner": {"my_key": "got here and there"}})), + (0.5, ((), {"outer_2": {"my_key": " and back again"}})), + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str