Catch errors thrown by nodes run by Executor

libs/langgraph/langgraph/pregel/executor.py

@@ -31,6 +31,7 @@ class Submit(Protocol[P, T]):
         *args: P.args,
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]: ...
 
@@ -39,7 +40,7 @@ class BackgroundExecutor(ContextManager):
     def __init__(self, config: RunnableConfig) -> None:
         self.stack = ExitStack()
         self.executor = self.stack.enter_context(get_executor_for_config(config))
-        self.tasks: dict[concurrent.futures.Future, bool] = {}
+        self.tasks: dict[concurrent.futures.Future, tuple[bool, bool]] = {}
 
     def submit(
         self,
@@ -47,10 +48,11 @@ class BackgroundExecutor(ContextManager):
         *args: P.args,
         __name__: Optional[str] = None,  # currently not used in sync version
         __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> concurrent.futures.Future[T]:
         task = self.executor.submit(fn, *args, **kwargs)
-        self.tasks[task] = __cancel_on_exit__
+        self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task
 
@@ -76,7 +78,7 @@ class BackgroundExecutor(ContextManager):
         traceback: Optional[TracebackType],
     ) -> Optional[bool]:
         # cancel all tasks that should be cancelled
-        for task, cancel in self.tasks.items():
+        for task, (cancel, _) in self.tasks.items():
             if cancel:
                 task.cancel()
         # wait for all tasks to finish
@@ -87,7 +89,9 @@ class BackgroundExecutor(ContextManager):
         # 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:
+            for task, (_, reraise) in self.tasks.items():
+                if not reraise:
+                    continue
                 try:
                     task.result()
                 except concurrent.futures.CancelledError:
@@ -97,7 +101,7 @@ class BackgroundExecutor(ContextManager):
 class AsyncBackgroundExecutor(AsyncContextManager):
     def __init__(self) -> None:
         self.context_not_supported = sys.version_info < (3, 11)
-        self.tasks: dict[asyncio.Task, bool] = {}
+        self.tasks: dict[asyncio.Task, tuple[bool, bool]] = {}
         self.sentinel = object()
         self.loop = asyncio.get_running_loop()
 
@@ -107,6 +111,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         *args: P.args,
         __name__: Optional[str] = None,
         __cancel_on_exit__: bool = False,
+        __reraise_on_exit__: bool = True,
         **kwargs: P.kwargs,
     ) -> asyncio.Task[T]:
         coro = fn(*args, **kwargs)
@@ -114,7 +119,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
             task = self.loop.create_task(coro, name=__name__)
         else:
             task = self.loop.create_task(coro, name=__name__, context=copy_context())
-        self.tasks[task] = __cancel_on_exit__
+        self.tasks[task] = (__cancel_on_exit__, __reraise_on_exit__)
         task.add_done_callback(self.done)
         return task
 
@@ -140,7 +145,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         traceback: Optional[TracebackType],
     ) -> None:
         # cancel all tasks that should be cancelled
-        for task, cancel in self.tasks.items():
+        for task, (cancel, _) in self.tasks.items():
             if cancel:
                 task.cancel(self.sentinel)
         # wait for all tasks to finish
@@ -149,7 +154,9 @@ class AsyncBackgroundExecutor(AsyncContextManager):
         # if there's already an exception being raised, don't raise another one
         if exc_type is None:
             # re-raise the first exception that occurred in a task
-            for task in self.tasks:
+            for task, (_, reraise) in self.tasks.items():
+                if not reraise:
+                    continue
                 try:
                     if exc := task.exception():
                         raise exc

libs/langgraph/langgraph/pregel/runner.py

@@ -35,6 +35,7 @@ class PregelRunner:
         self,
         tasks: list[PregelExecutableTask],
         *,
+        reraise: bool = True,
         timeout: Optional[float] = None,
         retry_policy: Optional[RetryPolicy] = None,
     ) -> Iterator[None]:
@@ -48,6 +49,7 @@ class PregelRunner:
                 run_with_retry,
                 task,
                 retry_policy,
+                __reraise_on_exit__=reraise,
             ): task
             for task in tasks
             if not task.writes
@@ -84,12 +86,13 @@ class PregelRunner:
             # give control back to the caller
             yield
         # panic on failure or timeout
-        _panic_or_proceed(all_futures)
+        _panic_or_proceed(all_futures, panic=reraise)
 
     async def atick(
         self,
         tasks: list[PregelExecutableTask],
         *,
+        reraise: bool = True,
         timeout: Optional[float] = None,
         retry_policy: Optional[RetryPolicy] = None,
     ) -> AsyncIterator[None]:
@@ -107,6 +110,7 @@ class PregelRunner:
                 stream=self.use_astream,
                 __name__=task.name,
                 __cancel_on_exit__=True,
+                __reraise_on_exit__=reraise,
             ): task
             for task in tasks
             if not task.writes
@@ -142,7 +146,9 @@ class PregelRunner:
             # give control back to the caller
             yield
         # panic on failure or timeout
-        _panic_or_proceed(all_futures, asyncio.TimeoutError)
+        _panic_or_proceed(
+            all_futures, timeout_exc_cls=asyncio.TimeoutError, panic=reraise
+        )
 
 
 def _should_stop_others(
@@ -171,7 +177,9 @@ def _exception(
 
 def _panic_or_proceed(
     futs: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]],
+    *,
     timeout_exc_cls: Type[Exception] = TimeoutError,
+    panic: bool = True,
 ) -> None:
     done: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set()
     inflight: set[Union[concurrent.futures.Future[Any], asyncio.Task[Any]]] = set()
@@ -187,7 +195,10 @@ def _panic_or_proceed(
             while inflight:
                 inflight.pop().cancel()
             # raise the exception
-            raise exc
+            if panic:
+                raise exc
+            else:
+                return
     if inflight:
         # if we got here means we timed out
         while inflight:

libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py

@@ -60,7 +60,7 @@ class KafkaExecutor(AbstractAsyncContextManager):
         return self
 
     async def __aexit__(self, *args: Any) -> None:
-        await self.stack.__aexit__(*args)
+        return await self.stack.__aexit__(*args)
 
     def __aiter__(self) -> Self:
         return self
@@ -109,7 +109,7 @@ class KafkaExecutor(AbstractAsyncContextManager):
                     submit=submit,
                     put_writes=partial(self._put_writes, submit, msg["config"]),
                 )
-                async for _ in runner.atick([task]):
+                async for _ in runner.atick([task], reraise=False):
                     pass
             else:
                 # task was not found

libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py

@@ -53,7 +53,7 @@ class KafkaOrchestrator(AbstractAsyncContextManager):
         return self
 
     async def __aexit__(self, *args: Any) -> None:
-        await self.stack.__aexit__(*args)
+        return await self.stack.__aexit__(*args)
 
     def __aiter__(self) -> Self:
         return self

libs/scheduler-kafka/pyproject.toml

@@ -53,5 +53,5 @@ lint.ignore = ["E501", "B008", "UP007", "UP006"]
 [tool.pytest-watcher]
 now = true
 delay = 0.1
-runner_args = ["--ff", "-v", "--tb", "short"]
+runner_args = ["--ff", "-v", "--tb", "short", "-s"]
 patterns = ["*.py"]

libs/scheduler-kafka/tests/test_scheduler.py

@@ -3,6 +3,7 @@ import functools
 import operator
 from typing import Annotated, Callable, ParamSpec, TypedDict, TypeVar, Union
 
+import anyio
 import pytest
 from aiokafka import AIOKafkaProducer
 
@@ -48,7 +49,6 @@ def mk_fanout_graph(checkpointer: BaseCheckpointSaver) -> Pregel:
         docs: Annotated[list[str], sorted_add]
 
     async def rewrite_query(data: State) -> State:
-        print("rewrite_query", data)
         return {"query": f'query: {data["query"]}'}
 
     async def retriever_picker(data: State) -> list[str]:
@@ -105,8 +105,8 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -
         nonlocal n_orch_msgs
         async with KafkaOrchestrator(graph, topics) as orch:
             async for msgs in orch:
-                n_orch_msgs += len(msgs)
                 print("orch", msgs)
+                n_orch_msgs += len(msgs)
                 if n_orch_msgs == expected:
                     break
 
@@ -114,8 +114,8 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -
         nonlocal n_exec_msgs
         async with KafkaExecutor(graph, topics) as exec:
             async for msgs in exec:
-                n_exec_msgs += len(msgs)
                 print("exec", msgs)
+                n_exec_msgs += len(msgs)
                 if n_exec_msgs == expected:
                     break
 
@@ -130,9 +130,9 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -
         )
 
     # run the orchestrator and executor
-    async with asyncio.TaskGroup() as tg:
-        tg.create_task(orchestrator(13), name="orchestrator")
-        tg.create_task(executor(12), name="executor")
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(orchestrator, 13, name="orchestrator")
+        tg.start_soon(executor, 12, name="executor")
 
     assert n_orch_msgs == 13
     assert n_exec_msgs == 12