From 287c29fbdc641607ff71e362a5d359a8599e94a2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 20 Nov 2024 15:38:43 -0800 Subject: [PATCH] Fix async --- libs/langgraph/langgraph/pregel/runner.py | 119 +++++++++++++++------- libs/langgraph/langgraph/types.py | 1 - 2 files changed, 80 insertions(+), 40 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index a7a17fc09..7c10a72ac 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -1,5 +1,6 @@ import asyncio import concurrent.futures +import threading import time from functools import partial from typing import ( @@ -66,18 +67,26 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], concurrent.futures.Future[None]]] = None, ) -> Iterator[None]: + locks: dict[str, threading.Lock] = {} + def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], *, calls: Optional[Sequence[Call]] = None, ) -> Sequence[Optional[concurrent.futures.Future]]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - # TODO could use a lock here instead, if writes can come from many threads - assert len(task.writes) == prev_length + len(writes) + if all(w[0] != PUSH for w in writes): + return task.config[CONF][CONFIG_KEY_SEND](writes) + + if task.id not in locks: + locks[task.id] = threading.Lock() + with locks[task.id]: + prev_length = len(task.writes) + # delegate to the underlying writer + task.config[CONF][CONFIG_KEY_SEND](writes) + # confirm no other concurrent writes were added + assert len(task.writes) == prev_length + len(writes) + # schedule PUSH tasks, collect futures rtn: dict[int, Optional[concurrent.futures.Future]] = {} for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write @@ -100,7 +109,6 @@ class PregelRunner: rtn[idx - prev_length] = fut elif next_task.writes: # if it already ran, return the result - # TODO we could also set the result for non-RETURN writes fut = concurrent.futures.Future() if val := next(v for c, v in next_task.writes if c == RETURN): fut.set_result(val) @@ -212,7 +220,6 @@ class PregelRunner: del fut, task # maybe stop other tasks if _should_stop_others(done): - print("stopping others") break # give control back to the caller yield @@ -231,24 +238,31 @@ class PregelRunner: retry_policy: Optional[RetryPolicy] = None, get_waiter: Optional[Callable[[], asyncio.Future[None]]] = None, ) -> AsyncIterator[None]: + locks: dict[str, threading.Lock] = {} + def writer( task: PregelExecutableTask, writes: Sequence[tuple[str, Any]], *, calls: Optional[Sequence[Call]] = None, ) -> Sequence[Optional[asyncio.Future]]: - prev_length = len(task.writes) - # delegate to the underlying writer - task.config[CONF][CONFIG_KEY_SEND](writes) - # confirm no other concurrent writes were added - # TODO could use a lock here instead, if writes can come from many threads - assert len(task.writes) == prev_length + len(writes) + if all(w[0] != PUSH for w in writes): + return task.config[CONF][CONFIG_KEY_SEND](writes) + + if task.id not in locks: + locks[task.id] = threading.Lock() + with locks[task.id]: + prev_length = len(task.writes) + # delegate to the underlying writer + task.config[CONF][CONFIG_KEY_SEND](writes) + # confirm no other concurrent writes were added + assert len(task.writes) == prev_length + len(writes) + # schedule PUSH tasks, collect futures rtn: dict[int, Optional[asyncio.Future]] = {} for idx, w in enumerate(writes, start=prev_length): # bail if not a PUSH write if w[0] != PUSH: continue - # TODO apply changes from sync version # schedule the next task, if the callback returns one if next_task := self.schedule_task( task, @@ -257,29 +271,51 @@ class PregelRunner: ): # if the parent task was retried, # the next task might already be running - if any( - t == next_task.id for t in futures.values() if t is not None - ): - continue - # schedule the next task - fut = cast( - asyncio.Future, - self.submit( - arun_with_retry, - next_task, - retry_policy, - stream=self.use_astream, - configurable={ - CONFIG_KEY_SEND: partial(writer, next_task), - CONFIG_KEY_CALL: partial(call, next_task), - }, - __name__=t.name, - __cancel_on_exit__=True, - __reraise_on_exit__=reraise, + if fut := next( + ( + f + for f, t in futures.items() + if t is not None and t == next_task.id ), - ) - futures[fut] = next_task - rtn[idx - prev_length] = fut + None, + ): + # if the parent task was retried, + # the next task might already be running + rtn[idx - prev_length] = fut + elif next_task.writes: + # if it already ran, return the result + fut = asyncio.Future() + if val := next(v for c, v in next_task.writes if c == RETURN): + fut.set_result(val) + elif exc := next(v for c, v in next_task.writes if c == ERROR): + fut.set_exception( + exc + if isinstance(exc, BaseException) + else Exception(exc) + ) + else: + fut.set_result(None) + rtn[idx - prev_length] = fut + else: + # schedule the next task + fut = cast( + asyncio.Future, + self.submit( + arun_with_retry, + next_task, + retry_policy, + stream=self.use_astream, + configurable={ + CONFIG_KEY_SEND: partial(writer, next_task), + CONFIG_KEY_CALL: partial(call, next_task), + }, + __name__=t.name, + __cancel_on_exit__=True, + __reraise_on_exit__=reraise, + ), + ) + futures[fut] = next_task + rtn[idx - prev_length] = fut return [rtn.get(i) for i in range(len(writes))] def call( @@ -294,6 +330,7 @@ class PregelRunner: loop = asyncio.get_event_loop() tasks = tuple(tasks) futures: dict[asyncio.Future, Optional[PregelExecutableTask]] = {} + done_futures: set[asyncio.Future] = set() # give control back to the caller yield # fast path if single task with no waiter and no timeout @@ -312,7 +349,12 @@ class PregelRunner: self.commit(t, None) except Exception as exc: self.commit(t, exc) - if reraise: + if reraise and futures: + # will be re-raised after futures are done + fut: asyncio.Future = loop.create_future() + fut.set_exception(exc) + done_futures.add(fut) + elif reraise: raise if not futures: # maybe `t` schuduled another task return @@ -342,7 +384,6 @@ class PregelRunner: ), ) ] = t - done_futures: set[asyncio.Future] = set() 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( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index a0949613c..3bdbf2049 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -398,5 +398,4 @@ def acall( fut = impl(func, *args, **kwargs) if isinstance(fut, concurrent.futures.Future): fut = asyncio.wrap_future(fut) - fut.cancel return fut