diff --git a/libs/langgraph/langgraph/pregel/executor.py b/libs/langgraph/langgraph/pregel/executor.py index 981ebbf7d..e441c8f78 100644 --- a/libs/langgraph/langgraph/pregel/executor.py +++ b/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 diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index c895bf9e9..3f42fbd5f 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/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: diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 82549cad9..1cf55266f 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/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 diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index b8fe1f003..408499f70 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/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 diff --git a/libs/scheduler-kafka/pyproject.toml b/libs/scheduler-kafka/pyproject.toml index 1799acb44..8ca9618e0 100644 --- a/libs/scheduler-kafka/pyproject.toml +++ b/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"] diff --git a/libs/scheduler-kafka/tests/test_scheduler.py b/libs/scheduler-kafka/tests/test_scheduler.py index d4cda6210..b448c3c75 100644 --- a/libs/scheduler-kafka/tests/test_scheduler.py +++ b/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