From d93be914c7187861bdb0bd6031abee53acb993d2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 20 Nov 2024 17:09:54 -0800 Subject: [PATCH] Fix stream order --- libs/langgraph/langgraph/func/__init__.py | 7 +- libs/langgraph/langgraph/pregel/call.py | 25 ++--- libs/langgraph/langgraph/pregel/runner.py | 94 +++++++++-------- libs/langgraph/langgraph/types.py | 18 ---- libs/langgraph/langgraph/utils/future.py | 121 ++++++++++++++++++++++ libs/langgraph/tests/test_pregel.py | 44 +++++++- libs/langgraph/tests/test_pregel_async.py | 84 +++++++++++++-- 7 files changed, 292 insertions(+), 101 deletions(-) create mode 100644 libs/langgraph/langgraph/utils/future.py diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 6826a8740..b48cf32e2 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -23,7 +23,7 @@ from langgraph.pregel.call import get_runnable_for_func from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import RetryPolicy, acall, call +from langgraph.types import RetryPolicy, call P = ParamSpec("P") T = TypeVar("T") @@ -48,10 +48,7 @@ def task( Callable[[Callable[P, T]], Callable[P, concurrent.futures.Future[T]]], ]: def _task(func: Callable[P, T]) -> Callable[P, concurrent.futures.Future[T]]: - if asyncio.iscoroutinefunction(func): - return update_wrapper(partial(acall, func), func) - else: - return update_wrapper(partial(call, func), func) + return update_wrapper(partial(call, func), func) return _task diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index a837ff4a8..a9986102d 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -1,11 +1,10 @@ -import asyncio import sys import types from typing import Any, Callable, Optional from langgraph.constants import RETURN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry -from langgraph.utils.runnable import RunnableCallable, RunnableSeq +from langgraph.utils.runnable import RunnableSeq, coerce_to_runnable """ Utilities borrowed from cloudpickle. @@ -110,25 +109,15 @@ def _lookup_module_and_qualname( def get_runnable_for_func(func: Callable[..., Any]) -> RunnableSeq: if func in CACHE: return CACHE[func] - elif not _lookup_module_and_qualname(func): - return RunnableSeq( - RunnableCallable(None, func, trace=False) - if asyncio.iscoroutinefunction(func) - else RunnableCallable(func, trace=False), + else: + seq = RunnableSeq( + coerce_to_runnable(func, name=None, trace=False), ChannelWrite([ChannelWriteEntry(RETURN)]), name=func.__name__, ) - else: - return CACHE.setdefault( - func, - RunnableSeq( - RunnableCallable(None, func, trace=False) - if asyncio.iscoroutinefunction(func) - else RunnableCallable(func, trace=False), - ChannelWrite([ChannelWriteEntry(RETURN)]), - name=func.__name__, - ), - ) + if not _lookup_module_and_qualname(func): + return seq + return CACHE.setdefault(func, seq) CACHE: dict[Callable[..., Any], RunnableSeq] = {} diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 7c10a72ac..8decbedec 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -34,6 +34,7 @@ from langgraph.pregel.algo import Call from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry from langgraph.types import PregelExecutableTask, RetryPolicy +from langgraph.utils.future import chain_future class PregelRunner: @@ -133,6 +134,7 @@ class PregelRunner: }, __reraise_on_exit__=reraise, ) + fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] @@ -165,7 +167,7 @@ class PregelRunner: ) self.commit(t, None) except Exception as exc: - self.commit(t, exc) + self.commit(t, None, exc) if reraise and futures: # will be re-raised after futures are done fut: concurrent.futures.Future = concurrent.futures.Future() @@ -183,18 +185,18 @@ class PregelRunner: # yield updates/debug output as each task finishes for t in tasks: if not t.writes: - futures[ - self.submit( - run_with_retry, - t, - retry_policy, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - __reraise_on_exit__=reraise, - ) - ] = t + fut = self.submit( + run_with_retry, + t, + retry_policy, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, + __reraise_on_exit__=reraise, + ) + fut.add_done_callback(partial(self.commit, t)) + futures[fut] = t end_time = timeout + time.monotonic() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = concurrent.futures.wait( @@ -213,8 +215,6 @@ class PregelRunner: else: # store for panic check done_futures.add(fut) - # task finished, commit writes - self.commit(task, _exception(fut)) else: # remove references to loop vars del fut, task @@ -264,11 +264,8 @@ class PregelRunner: if w[0] != PUSH: continue # schedule the next task, if the callback returns one - if next_task := self.schedule_task( - task, - idx, - calls[idx - prev_length] if calls is not None else None, - ): + wcall = calls[idx - prev_length] if calls is not None else None + if next_task := self.schedule_task(task, idx, wcall): # if the parent task was retried, # the next task might already be running if fut := next( @@ -314,6 +311,7 @@ class PregelRunner: __reraise_on_exit__=reraise, ), ) + fut.add_done_callback(partial(self.commit, next_task)) futures[fut] = next_task rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] @@ -322,10 +320,15 @@ class PregelRunner: task: PregelExecutableTask, func: Callable[[Any], Union[Awaitable[Any], Any]], input: Any, - ) -> asyncio.Future[Any]: + ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: (fut,) = writer(task, [(PUSH, None)], calls=[Call(func, input)]) assert fut is not None, "writer did not return a future for call" - return fut + if asyncio.iscoroutinefunction(func): + return fut + # adapted from asyncio.run_coroutine_threadsafe + sfut = concurrent.futures.Future() + loop.call_soon_threadsafe(chain_future, fut, sfut) + return sfut loop = asyncio.get_event_loop() tasks = tuple(tasks) @@ -348,7 +351,7 @@ class PregelRunner: ) self.commit(t, None) except Exception as exc: - self.commit(t, exc) + self.commit(t, None, exc) if reraise and futures: # will be re-raised after futures are done fut: asyncio.Future = loop.create_future() @@ -366,24 +369,24 @@ class PregelRunner: # yield updates/debug output as each task finishes for t in tasks: if not t.writes: - futures[ - cast( - asyncio.Future, - self.submit( - arun_with_retry, - t, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, t), - CONFIG_KEY_CALL: partial(call, t), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, - ), - ) - ] = t + fut = cast( + asyncio.Future, + self.submit( + arun_with_retry, + t, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, t), + CONFIG_KEY_CALL: partial(call, t), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), + ) + fut.add_done_callback(partial(self.commit, t)) + futures[fut] = t end_time = timeout + loop.time() if timeout else None while len(futures) > (1 if get_waiter is not None else 0): done, inflight = await asyncio.wait( @@ -402,8 +405,6 @@ class PregelRunner: else: # store for panic check done_futures.add(fut) - # task finished, commit writes - self.commit(task, _exception(fut)) else: # remove references to loop vars del fut, task @@ -423,8 +424,13 @@ class PregelRunner: ) def commit( - self, task: PregelExecutableTask, exception: Optional[BaseException] + self, + task: PregelExecutableTask, + fut: Union[None, concurrent.futures.Future[Any], asyncio.Future[Any]], + exception: Optional[BaseException] = None, ) -> None: + if fut is not None: + exception = _exception(fut) if exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3bdbf2049..5ce64b52b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -1,4 +1,3 @@ -import asyncio import concurrent import concurrent.futures import dataclasses @@ -7,7 +6,6 @@ from collections import deque from typing import ( TYPE_CHECKING, Any, - Awaitable, Callable, ClassVar, Generic, @@ -382,20 +380,4 @@ def call( conf = get_configurable() impl = conf[CONFIG_KEY_CALL] fut = impl(func, *args, **kwargs) - if not isinstance(fut, concurrent.futures.Future): - raise RuntimeError("In an async context, use acall() instead of call()") - return fut - - -def acall( - func: str | Callable[P, Union[T, Awaitable[T]]], *args: P.args, **kwargs: P.kwargs -) -> asyncio.Future[T]: - from langgraph.constants import CONFIG_KEY_CALL - from langgraph.utils.config import get_configurable - - conf = get_configurable() - impl = conf[CONFIG_KEY_CALL] - fut = impl(func, *args, **kwargs) - if isinstance(fut, concurrent.futures.Future): - fut = asyncio.wrap_future(fut) return fut diff --git a/libs/langgraph/langgraph/utils/future.py b/libs/langgraph/langgraph/utils/future.py new file mode 100644 index 000000000..d7b8718f6 --- /dev/null +++ b/libs/langgraph/langgraph/utils/future.py @@ -0,0 +1,121 @@ +import asyncio +import concurrent.futures +from typing import Union + +AnyFuture = Union[asyncio.Future, concurrent.futures.Future] + + +def _get_loop(fut: asyncio.Future) -> asyncio.AbstractEventLoop: + # Tries to call Future.get_loop() if it's available. + # Otherwise fallbacks to using the old '_loop' property. + try: + get_loop = fut.get_loop + except AttributeError: + pass + else: + return get_loop() + return fut._loop + + +def _convert_future_exc(exc): + exc_class = type(exc) + if exc_class is concurrent.futures.CancelledError: + return asyncio.CancelledError(*exc.args) + elif exc_class is concurrent.futures.TimeoutError: + return asyncio.TimeoutError(*exc.args) + elif exc_class is concurrent.futures.InvalidStateError: + return asyncio.InvalidStateError(*exc.args) + else: + return exc + + +def _set_concurrent_future_state(concurrent, source): + """Copy state from a future to a concurrent.futures.Future.""" + assert source.done() + if source.cancelled(): + concurrent.cancel() + if not concurrent.set_running_or_notify_cancel(): + return + exception = source.exception() + if exception is not None: + concurrent.set_exception(_convert_future_exc(exception)) + else: + result = source.result() + concurrent.set_result(result) + + +def _copy_future_state(source, dest): + """Internal helper to copy state from another Future. + + The other Future may be a concurrent.futures.Future. + """ + assert source.done() + if dest.cancelled(): + return + assert not dest.done() + if source.cancelled(): + dest.cancel() + else: + exception = source.exception() + if exception is not None: + dest.set_exception(_convert_future_exc(exception)) + else: + result = source.result() + dest.set_result(result) + + +def _chain_future(source: AnyFuture, destination: AnyFuture) -> None: + """Chain two futures so that when one completes, so does the other. + + The result (or exception) of source will be copied to destination. + If destination is cancelled, source gets cancelled too. + Compatible with both asyncio.Future and concurrent.futures.Future. + """ + if not asyncio.isfuture(source) and not isinstance( + source, concurrent.futures.Future + ): + raise TypeError("A future is required for source argument") + if not asyncio.isfuture(destination) and not isinstance( + destination, concurrent.futures.Future + ): + raise TypeError("A future is required for destination argument") + source_loop = _get_loop(source) if asyncio.isfuture(source) else None + dest_loop = _get_loop(destination) if asyncio.isfuture(destination) else None + + def _set_state(future, other): + if asyncio.isfuture(future): + _copy_future_state(other, future) + else: + _set_concurrent_future_state(future, other) + + def _call_check_cancel(destination): + if destination.cancelled(): + if source_loop is None or source_loop is dest_loop: + source.cancel() + else: + source_loop.call_soon_threadsafe(source.cancel) + + def _call_set_state(source): + if destination.cancelled() and dest_loop is not None and dest_loop.is_closed(): + return + if dest_loop is None or dest_loop is source_loop: + _set_state(destination, source) + else: + if dest_loop.is_closed(): + return + dest_loop.call_soon_threadsafe(_set_state, destination, source) + + destination.add_done_callback(_call_check_cancel) + source.add_done_callback(_call_set_state) + + +def chain_future(source: AnyFuture, destination: AnyFuture) -> None: + # adapted from asyncio.run_coroutine_threadsafe + try: + _chain_future(source, destination) + except (SystemExit, KeyboardInterrupt): + raise + except BaseException as exc: + if destination.set_running_or_notify_cancel(): + destination.set_exception(exc) + raise diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 49ee6bec5..ab675ea79 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1976,21 +1976,21 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non mapper_calls = 0 @task() - def mapper(input: str) -> str: + def mapper(input: int) -> str: nonlocal mapper_calls mapper_calls += 1 - return input * 2 + time.sleep(input / 100) + return str(input) * 2 @imp(checkpointer=checkpointer) - def graph(input: list[str]) -> list[str]: + def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] answer = interrupt("question") return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert [*graph.stream(["0", "1"], thread1)] == [ - # TODO make test not depend on order of execution (which is not guaranteed) + assert [*graph.stream([0, 1], thread1)] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -2013,6 +2013,40 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non assert mapper_calls == 2 +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_imp_stream_order( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + @task() + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task() + def bar(state: dict) -> dict: + return {"a": state["a"] + state["b"], "c": "bark"} + + @task() + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @imp(checkpointer=checkpointer) + def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(fut_foo.result()) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c for c in graph.stream({"a": "0"}, thread1)] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_send_dedupe_on_resume( request: pytest.FixtureRequest, checkpointer_name: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 450606d80..07e16a538 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2650,26 +2650,24 @@ async def test_send_sequences(checkpointer_name: str) -> None: @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str) -> None: - mapper_calls = 0 - - @task() - async def mapper(input: str) -> str: - nonlocal mapper_calls - mapper_calls += 1 - return input * 2 - async with awith_checkpointer(checkpointer_name) as checkpointer: + mapper_calls = 0 + + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls + mapper_calls += 1 + return str(input) * 2 @imp(checkpointer=checkpointer) - async def graph(input: list[str]) -> list[str]: + async def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = await asyncio.gather(*futures) answer = interrupt("question") return [m + answer for m in mapped] thread1 = {"configurable": {"thread_id": "1"}} - assert [c async for c in graph.astream(["0", "1"], thread1)] == [ - # TODO make test not depend on order of execution (which is not guaranteed) + assert [c async for c in graph.astream([0, 1], thread1)] == [ {"mapper": "00"}, {"mapper": "11"}, { @@ -2692,6 +2690,70 @@ async def test_imp_task(checkpointer_name: str) -> None: assert mapper_calls == 2 +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_sync_from_async(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @task() + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task() + def bar(state: dict) -> dict: + return {"a": state["a"] + state["b"], "c": "bark"} + + @task() + def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @imp(checkpointer=checkpointer) + def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(fut_foo.result()) + fut_baz = baz(fut_bar.result()) + return fut_baz.result() + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_stream_order(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @task() + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo", "b": "bar"} + + @task() + async def bar(state: dict) -> dict: + return {"a": state["a"] + state["b"], "c": "bark"} + + @task() + async def baz(state: dict) -> dict: + return {"a": state["a"] + "baz", "c": "something else"} + + @imp(checkpointer=checkpointer) + async def graph(state: dict) -> dict: + fut_foo = foo(state) + fut_bar = bar(await fut_foo) + fut_baz = baz(await fut_bar) + return await fut_baz + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream({"a": "0"}, thread1)] == [ + {"foo": {"a": "0foo", "b": "bar"}}, + {"bar": {"a": "0foobar", "c": "bark"}}, + {"baz": {"a": "0foobarbaz", "c": "something else"}}, + {"graph": {"a": "0foobarbaz", "c": "something else"}}, + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_send_dedupe_on_resume(checkpointer_name: str) -> None: if not FF_SEND_V2: