From aed313ebb576f85312fc96c35309665649ab4300 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 22 Jul 2024 17:09:47 -0700 Subject: [PATCH] Fix stack not being unwound when suppressing interrupt --- libs/langgraph/langgraph/pregel/executor.py | 55 +++++++----- libs/langgraph/langgraph/pregel/loop.py | 98 ++++++++++----------- 2 files changed, 79 insertions(+), 74 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 948711508..69d3b5d31 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/libs/langgraph/langgraph/pregel/executor.py @@ -80,18 +80,18 @@ class BackgroundExecutor(ContextManager): if cancel: task.cancel() # wait for all tasks to finish - concurrent.futures.wait({t for t in self.tasks if not t.done()}) + if tasks := {t for t in self.tasks if not t.done()}: + concurrent.futures.wait(tasks) # shutdown the executor self.stack.__exit__(exc_type, exc_value, traceback) - # raise caught exception - if exc_type is not None: - raise exc_value # re-raise the first exception that occurred in a task - for task in self.tasks: - try: - task.result() - except concurrent.futures.CancelledError: - pass + if exc_type is None: + # if there's already an exception being raised, don't raise another one + for task in self.tasks: + try: + task.result() + except concurrent.futures.CancelledError: + pass class AsyncBackgroundExecutor(AsyncContextManager): @@ -132,16 +132,27 @@ class AsyncBackgroundExecutor(AsyncContextManager): 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 exit( + self, + exc_type: Optional[type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + # cancel all tasks that should be cancelled + for task, cancel in self.tasks.items(): + if cancel: + task.cancel(self.sentinel) + # wait for all tasks to finish + if self.tasks: + await asyncio.wait(self.tasks) + # re-raise the first exception that occurred in a task + if exc_type is None: + # if there's already an exception being raised, don't raise another one + for task in self.tasks: + try: + task.result() + except asyncio.CancelledError: + pass async def __aexit__( self, @@ -151,8 +162,6 @@ class AsyncBackgroundExecutor(AsyncContextManager): ) -> Optional[bool]: # we cannot use `await` outside of asyncio.shield, as this code can run # after owning task is cancelled, so pulling async logic to separate method - 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()) + await asyncio.shield(self.exit(exc_type, exc_value, traceback)) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index da7537cba..bfb8020b0 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -102,6 +102,23 @@ class PregelLoop: # public + def __init__( + self, + input: Optional[Any], + *, + config: RunnableConfig, + checkpointer: Optional[BaseCheckpointSaver], + graph: "Pregel", + ) -> None: + self.stream = deque() + self.input = input + self.config = config + self.checkpointer = checkpointer + self.graph = graph + # TODO if managed values no longer needs graph we can replace with + # managed_specs, channel_specs + self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) + def mark_tasks_scheduled(self, tasks: Sequence[PregelExecutableTask]) -> None: """Mark tasks as scheduled, to be used by queue-based executors.""" raise NotImplementedError @@ -219,7 +236,7 @@ class PregelLoop: if should_interrupt(self.checkpoint, interrupt_before, self.tasks): self.status = "interrupt_before" if self.is_nested: - raise GraphInterrupt(self) + raise GraphInterrupt() else: return False @@ -318,6 +335,15 @@ class PregelLoop: # increment step self.step += 1 + def _suppress_interrupt( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> Optional[bool]: + if exc_type is GraphInterrupt and not self.is_nested: + return True + class SyncPregelLoop(PregelLoop, ContextManager): def __init__( @@ -328,24 +354,21 @@ class SyncPregelLoop(PregelLoop, ContextManager): checkpointer: Optional[BaseCheckpointSaver], graph: "Pregel", ) -> None: - self.stream = deque() + super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) self.stack = ExitStack() - self.input = input - self.config = config - self.checkpointer = checkpointer - self.checkpointer_get_next_version = ( - checkpointer.get_next_version if checkpointer else increment - ) - self.checkpointer_put_writes = checkpointer.put_writes if checkpointer else None - self.checkpointer_put = checkpointer.put if checkpointer else None - self.graph = graph - # TODO if managed values no longer needs graph we can replace with - # managed_specs, channel_specs + self.stack.push(self._suppress_interrupt) + if checkpointer: + self.checkpointer_get_next_version = checkpointer.get_next_version + self.checkpointer_put_writes = checkpointer.put_writes + self.checkpointer_put = checkpointer.put + else: + self.checkpointer_get_next_version = increment + self.checkpointer_put_writes = None + self.checkpointer_put = None # context manager def __enter__(self) -> Self: - self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) saved = ( self.checkpointer.get_tuple(self.config) if self.checkpointer else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) @@ -381,16 +404,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - # handle interrupt - if exc_type is GraphInterrupt: - if exc_value.args[0] is self: - # interrupt raised by this loop - exc_value.args = (object(),) - if not self.is_nested: - # in outer graph, catch interrupt - del self.graph - return True or self.stack.__exit__(None, None, None) - # unwind stack del self.graph return self.stack.__exit__(exc_type, exc_value, traceback) @@ -405,26 +418,21 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): checkpointer: Optional[BaseCheckpointSaver], graph: "Pregel", ) -> None: - self.stream = deque() + super().__init__(input, config=config, checkpointer=checkpointer, graph=graph) self.stack = AsyncExitStack() - self.input = input - self.config = config - self.checkpointer = checkpointer - self.checkpointer_get_next_version = ( - checkpointer.get_next_version if checkpointer else increment - ) - self.checkpointer_put_writes = ( - checkpointer.aput_writes if checkpointer else None - ) - self.checkpointer_put = checkpointer.aput if checkpointer else None - self.graph = graph - # TODO if managed values no longer needs graph we can replace with - # managed_specs, channel_specs + self.stack.push(self._suppress_interrupt) + if checkpointer: + self.checkpointer_get_next_version = checkpointer.get_next_version + self.checkpointer_put_writes = checkpointer.aput_writes + self.checkpointer_put = checkpointer.aput + else: + self.checkpointer_get_next_version = increment + self.checkpointer_put_writes = None + self.checkpointer_put = None # context manager async def __aenter__(self) -> Self: - self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) saved = ( await self.checkpointer.aget_tuple(self.config) if self.checkpointer @@ -462,18 +470,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: - # handle interrupt - if exc_type is GraphInterrupt: - if exc_value.args[0] is self: - # interrupt raised by this loop - exc_value.args = (object(),) - if not self.is_nested: - # in outer graph, catch interrupt - del self.graph - return True or await asyncio.shield( - self.stack.__aexit__(None, None, None) - ) - # unwind stack del self.graph return await asyncio.shield(