diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 05d99d8d0..396fa6c30 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -41,7 +41,6 @@ from langchain_core.runnables.config import ( ensure_config, get_async_callback_manager_for_config, get_callback_manager_for_config, - get_executor_for_config, merge_configs, patch_config, ) @@ -94,6 +93,7 @@ from langgraph.pregel.debug import ( print_step_tasks, print_step_writes, ) +from langgraph.pregel.executor import AsyncBackgroundExecutor, BackgroundExecutor from langgraph.pregel.io import ( map_input, map_output_updates, @@ -836,7 +836,6 @@ class Pregel( run_id=config.get("run_id"), ) try: - bg: list[concurrent.futures.Future] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -880,29 +879,25 @@ class Pregel( start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint - with ChannelsManager( + with BackgroundExecutor(config) as submit, ChannelsManager( self.channels, checkpoint, config - ) as channels, get_executor_for_config( - config - ) as executor, ManagedValuesManager( + ) as channels, ManagedValuesManager( self.managed_values_dict, config, self ) as managed: def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None: if self.checkpointer is not None: - bg.append( - executor.submit( - self.checkpointer.put_writes, - { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, + submit( + self.checkpointer.put_writes, + { + **checkpoint_config, + "configurable": { + **checkpoint_config["configurable"], + "thread_ts": checkpoint["id"], }, - writes, - task_id, - ) + }, + writes, + task_id, ) def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: @@ -920,13 +915,11 @@ class Pregel( checkpoint, channels, metadata["step"] ) # save it, without blocking - bg.append( - executor.submit( - self.checkpointer.put, - checkpoint_config, - copy_checkpoint(checkpoint), - metadata, - ) + submit( + self.checkpointer.put, + checkpoint_config, + copy_checkpoint(checkpoint), + metadata, ) # update checkpoint config checkpoint_config = { @@ -1059,7 +1052,7 @@ class Pregel( # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - executor.submit(run_with_retry, task, self.retry_policy): task + submit(run_with_retry, task, self.retry_policy): task for task in next_tasks if not task.writes } @@ -1080,6 +1073,8 @@ class Pregel( else None ), ) + if not done: + break # timed out for fut in done: task = futures.pop(fut) if fut.exception() is not None: @@ -1182,19 +1177,6 @@ class Pregel( except BaseException as e: run_manager.on_chain_error(e) raise - finally: - # cancel any pending tasks when generator is interrupted - try: - for task in futures: - task.cancel() - except NameError: - pass - # wait for all background tasks to finish - done, _ = concurrent.futures.wait( - bg, return_when=concurrent.futures.ALL_COMPLETED - ) - for task in done: - task.result() async def astream( self, @@ -1294,7 +1276,6 @@ class Pregel( ) try: loop = asyncio.get_event_loop() - bg: list[asyncio.Task] = [] if config["recursion_limit"] < 1: raise ValueError("recursion_limit must be at least 1") if self.checkpointer and not config.get("configurable"): @@ -1342,7 +1323,7 @@ class Pregel( start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint - async with AsyncChannelsManager( + async with AsyncBackgroundExecutor() as submit, AsyncChannelsManager( self.channels, checkpoint, config ) as channels, AsyncManagedValuesManager( self.managed_values_dict, config, self @@ -1350,20 +1331,17 @@ class Pregel( def put_writes(task_id: str, writes: Sequence[tuple[str, Any]]) -> None: if self.checkpointer is not None: - bg.append( - asyncio.create_task( - self.checkpointer.aput_writes( - { - **checkpoint_config, - "configurable": { - **checkpoint_config["configurable"], - "thread_ts": checkpoint["id"], - }, - }, - writes, - task_id, - ) - ) + submit( + self.checkpointer.aput_writes, + { + **checkpoint_config, + "configurable": { + **checkpoint_config["configurable"], + "thread_ts": checkpoint["id"], + }, + }, + writes, + task_id, ) def put_checkpoint(metadata: CheckpointMetadata) -> Iterator[Any]: @@ -1381,13 +1359,13 @@ class Pregel( checkpoint, channels, metadata["step"] ) # save it, without blocking - bg.append( - asyncio.create_task( - self.checkpointer.aput( - checkpoint_config, copy_checkpoint(checkpoint), metadata - ) - ) + submit( + self.checkpointer.aput, + checkpoint_config, + copy_checkpoint(checkpoint), + metadata, ) + # update checkpoint config checkpoint_config = { **checkpoint_config, @@ -1517,8 +1495,13 @@ class Pregel( # each task is independent from all other concurrent tasks # yield updates/debug output as each task finishes futures = { - asyncio.create_task( - arun_with_retry(task, self.retry_policy, do_stream) + submit( + arun_with_retry, + task, + self.retry_policy, + do_stream, + __name__=task.name, + __cancel_on_exit__=True, ): task for task in next_tasks if not task.writes @@ -1536,6 +1519,8 @@ class Pregel( max(0, end_time - loop.time()) if end_time else None ), ) + if not done: + break # timed out for fut in done: task = futures.pop(fut) if fut.exception() is not None: @@ -1569,7 +1554,7 @@ class Pregel( del fut, task # panic on failure or timeout - _panic_or_proceed(done, inflight, step) + _panic_or_proceed(done, inflight, step, asyncio.TimeoutError) # don't keep futures around in memory longer than needed del done, inflight, futures @@ -1640,22 +1625,8 @@ class Pregel( # set final channel values as run output await run_manager.on_chain_end(read_channels(channels, output_keys)) except BaseException as e: - await run_manager.on_chain_error(e) + await asyncio.shield(run_manager.on_chain_error(e)) raise - finally: - # cancel any pending tasks when generator is interrupted - try: - for task in futures: - task.cancel() - bg.append(task) - except NameError: - pass - # wait for all background tasks to finish - fut = asyncio.gather(*bg) - # mark the exception as retrieved - fut.add_done_callback(_mark_cancelled_as_seen) - # wait for the future to finish, shielded from cancellation - await asyncio.shield(fut) def invoke( self, @@ -1773,6 +1744,7 @@ def _panic_or_proceed( done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], step: int, + timeout_exc_cls: Type[Exception] = TimeoutError, ) -> None: while done: # if any task failed @@ -1789,7 +1761,7 @@ def _panic_or_proceed( # cancel all pending tasks inflight.pop().cancel() # raise timeout error - raise TimeoutError(f"Timed out at step {step}") + raise timeout_exc_cls(f"Timed out at step {step}") def _should_interrupt( @@ -2167,10 +2139,3 @@ def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any] yield (mode, chunk) else: yield from iter - - -def _mark_cancelled_as_seen(fut: concurrent.futures.Future) -> None: - try: - fut.exception() - except (asyncio.CancelledError, concurrent.futures.CancelledError): - pass diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py new file mode 100644 index 000000000..4d43aab60 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -0,0 +1,128 @@ +import asyncio +import concurrent.futures +import sys +from contextlib import contextmanager +from contextvars import copy_context +from types import TracebackType +from typing import ( + AsyncContextManager, + Callable, + Iterator, + Optional, + Protocol, + TypeVar, +) + +from langchain_core.runnables import RunnableConfig +from langchain_core.runnables.config import get_executor_for_config +from typing_extensions import ParamSpec + +P = ParamSpec("P") +T = TypeVar("T") + + +class Submit(Protocol[P, T]): + def __call__( + self, + fn: Callable[P, T], + *args: P.args, + __name__: Optional[str] = None, + __cancel_on_exit__: bool = False, + **kwargs: P.kwargs, + ) -> concurrent.futures.Future[T]: + ... + + +@contextmanager +def BackgroundExecutor(config: RunnableConfig) -> Iterator[Submit]: + tasks: dict[concurrent.futures.Future, bool] = {} + with get_executor_for_config(config) as executor: + + def done(task: concurrent.futures.Future) -> None: + try: + task.result() + except BaseException: + pass + else: + tasks.pop(task) + + def submit( + fn: Callable[P, T], + *args: P.args, + __name__: Optional[str] = None, # currently not used in sync version + __cancel_on_exit__: bool = False, + **kwargs: P.kwargs, + ) -> concurrent.futures.Future: + task = executor.submit(fn, *args, **kwargs) + tasks[task] = __cancel_on_exit__ + task.add_done_callback(done) + return task + + try: + yield submit + finally: + for task, cancel in tasks.items(): + if cancel: + task.cancel() + # executor waits for all tasks to finish on exit + for task in tasks: + # the first task to have raised an exception will be re-raised here + task.result() + + +class AsyncBackgroundExecutor(AsyncContextManager): + def __init__(self) -> None: + self.context_not_supported = sys.version_info < (3, 11) + self.tasks: dict[asyncio.Task, bool] = {} + self.sentinel = object() + + def submit( + self, + fn: Callable[P, T], + *args: P.args, + __name__: Optional[str] = None, + __cancel_on_exit__: bool = False, + **kwargs: P.kwargs, + ) -> asyncio.Task[T]: + coro = fn(*args, **kwargs) + if self.context_not_supported: + task = asyncio.create_task(coro, name=__name__) + else: + task = asyncio.create_task(coro, name=__name__, context=copy_context()) + self.tasks[task] = __cancel_on_exit__ + task.add_done_callback(self.done) + return task + + def done(self, task: asyncio.Task) -> None: + try: + task.result() + except BaseException: + pass + else: + self.tasks.pop(task) + + async def __aenter__(self) -> Submit: + return self.submit + + async def exit(self) -> None: + fut = asyncio.gather(*self.tasks, return_exceptions=True) + try: + rtns = await asyncio.shield(fut) + finally: + del self.tasks + for rtn in rtns: + # if this is ever changed to BaseException, need to ignore CancelledError + if isinstance(rtn, Exception): + raise rtn + + async def __aexit__( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + for task, cancel in self.tasks.items(): + if cancel: + task.cancel(self.sentinel) + # wait for all background tasks to finish, shielded from cancellation + await asyncio.shield(self.exit()) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 6ec224012..0fe73a27a 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -41,6 +41,8 @@ def default_retry_on(exc: Exception) -> bool: class RetryPolicy(NamedTuple): + """Configuration for retrying nodes.""" + initial_interval: float = 0.5 """Amount of time that must elapse before the first retry occurs. In seconds.""" backoff_factor: float = 2.0 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 41f99d53f..5a0271586 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -220,7 +220,7 @@ async def test_step_timeout_on_stream_hang() -> None: graph = builder.compile() graph.step_timeout = 1 - with pytest.raises(asyncio.CancelledError): + with pytest.raises(asyncio.TimeoutError): async for chunk in graph.astream(1, stream_mode="updates"): assert chunk == {"alittlewhile": {"alittlewhile": "1"}} await asyncio.sleep(0.6) @@ -247,7 +247,7 @@ async def test_cancel_graph_astream( try: class State(TypedDict): - value: int + value: Annotated[int, operator.add] class AwhileMaker: def __init__(self) -> None: @@ -270,20 +270,31 @@ async def test_cancel_graph_astream( return {"value": 2} awhile = AwhileMaker() + aparallelwhile = AwhileMaker() builder = StateGraph(State) builder.add_node("awhile", awhile) + builder.add_node("aparallelwhile", aparallelwhile) builder.add_node(alittlewhile) builder.add_edge(START, "alittlewhile") + builder.add_edge(START, "aparallelwhile") builder.add_edge("alittlewhile", "awhile") graph = builder.compile(checkpointer=checkpointer) # test interrupting astream + got_event = False thread1: RunnableConfig = {"configurable": {"thread_id": 1}} async with aclosing(graph.astream({"value": 1}, thread1)) as stream: async for chunk in stream: assert chunk == {"alittlewhile": {"value": 2}} + got_event = True break + assert got_event + + # node aparallelwhile should start, but be cancelled + assert aparallelwhile.started is True + assert aparallelwhile.cancelled is True + # node "awhile" should never start assert awhile.started is False @@ -292,7 +303,10 @@ async def test_cancel_graph_astream( state = await graph.aget_state(thread1) assert state is not None assert state.values == {"value": 1} - assert state.next == ("alittlewhile",) + assert state.next == ( + "aparallelwhile", + "alittlewhile", + ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} finally: if getattr(checkpointer, "__aexit__", None): @@ -366,9 +380,10 @@ async def test_cancel_graph_astream_events_v2( # did break assert got_event - # node "awhile" starts but is cancelled - assert awhile.started is True - assert awhile.cancelled is True + # node "awhile" maybe starts (impl detail of astream_events) + # if it does start, it must be cancelled + if awhile.started: + assert awhile.cancelled is True # node "anotherwhile" should never start assert anotherwhile.started is False