From 70153ceba22013d34b51317cd76c976a62af80cb Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 23 May 2025 12:26:27 -0700 Subject: [PATCH 1/5] Fix exception handling for imperative tasks - These exceptions should not be re-raised at end of tick, given they're handled explicitly by the developer in their entrypoint --- libs/langgraph/langgraph/pregel/runner.py | 29 +++++++++++------------ libs/langgraph/tests/test_pregel_async.py | 28 ++++++++++++++++++++++ 2 files changed, 42 insertions(+), 15 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 96a256ad9..22d7a290a 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -56,6 +56,8 @@ EXCLUDED_FRAME_FNAMES = ( "concurrent/futures/_base.py", ) +SKIP_RERAISE_SET = weakref.WeakSet() + class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): event: E @@ -165,7 +167,6 @@ class PregelRunner: futures=weakref.ref(futures), schedule_task=schedule_task, submit=self.submit, - reraise=reraise, ), }, ) @@ -207,7 +208,6 @@ class PregelRunner: futures=weakref.ref(futures), schedule_task=schedule_task, submit=self.submit, - reraise=reraise, ), }, __reraise_on_exit__=reraise, @@ -302,7 +302,6 @@ class PregelRunner: futures=weakref.ref(futures), schedule_task=schedule_task, submit=self.submit, - reraise=reraise, loop=loop, ), }, @@ -349,7 +348,6 @@ class PregelRunner: futures=weakref.ref(futures), schedule_task=schedule_task, submit=self.submit, - reraise=reraise, loop=loop, ), }, @@ -494,7 +492,8 @@ def _panic_or_proceed( interrupts: list[GraphInterrupt] = [] while done: # if any task failed - if exc := _exception(done.pop()): + fut = done.pop() + if exc := _exception(fut): # cancel all pending tasks while inflight: inflight.pop().cancel() @@ -503,7 +502,7 @@ def _panic_or_proceed( if isinstance(exc, GraphInterrupt): # collect interrupts interrupts.append(exc) - else: + elif fut not in SKIP_RERAISE_SET: raise exc # raise combined interrupts if interrupts: @@ -530,7 +529,6 @@ def _call( [PregelExecutableTask, int, Optional[Call]], Optional[PregelExecutableTask] ], submit: weakref.ref[Submit], - reraise: bool, ) -> concurrent.futures.Future[Any]: if asyncio.iscoroutinefunction(func): raise RuntimeError("In an sync context async tasks cannot be called") @@ -582,14 +580,16 @@ def _call( callbacks=callbacks, schedule_task=schedule_task, submit=submit, - reraise=reraise, ), }, - __reraise_on_exit__=reraise, + __reraise_on_exit__=False, # starting a new task in the next tick ensures # updates from this tick are committed/streamed first __next_tick__=True, ) + # exceptions for call() tasks are raised into the parent task + # so we should not re-raise at the end of the tick + SKIP_RERAISE_SET.add(fut) futures()[fut] = next_task # type: ignore[index] fut = cast(Union[asyncio.Future, concurrent.futures.Future], fut) # return a chained future to ensure commit() callback is called @@ -613,7 +613,6 @@ def _acall( ], submit: weakref.ref[Submit], loop: asyncio.AbstractEventLoop, - reraise: bool = False, stream: bool = False, ) -> Union[asyncio.Future[Any], concurrent.futures.Future[Any]]: # return a chained future to ensure commit() callback is called @@ -643,7 +642,6 @@ def _acall( schedule_task=schedule_task, submit=submit, loop=loop, - reraise=reraise, stream=stream, ), loop, @@ -669,7 +667,6 @@ async def _acall_impl( ], submit: weakref.ref[Submit], loop: asyncio.AbstractEventLoop, - reraise: bool = False, stream: bool = False, ) -> None: try: @@ -726,17 +723,19 @@ async def _acall_impl( schedule_task=schedule_task, submit=submit, loop=loop, - reraise=reraise, ), }, - __name__=task().name, # type: ignore[union-attr] + __name__=next_task.name, __cancel_on_exit__=True, - __reraise_on_exit__=reraise, + __reraise_on_exit__=False, # starting a new task in the next tick ensures # updates from this tick are committed/streamed first __next_tick__=True, ), ) + # exceptions for call() tasks are raised into the parent task + # so we should not re-raise at the end of the tick + SKIP_RERAISE_SET.add(fut) futures()[fut] = next_task # type: ignore[index] if fut is not None: chain_future(fut, destination) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3c95a2cec..58745ec25 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9148,3 +9148,31 @@ async def test_draw_invalid(): {"source": "nothing", "target": "__end__"}, ], } + + +@NEEDS_CONTEXTVARS +async def test_imp_exception( + async_checkpointer: BaseCheckpointSaver, +) -> None: + @task() + async def my_task(number: int): + await asyncio.sleep(1) + return number * 2 + + @task() + async def task_with_exception(number: int): + await asyncio.sleep(1) + raise Exception("This is a test exception") + + @entrypoint(checkpointer=async_checkpointer) + async def my_workflow(number: int): + await my_task(number) + try: + await task_with_exception(number) + except Exception as e: + print(f"Exception caught: {e}") + await my_task(number) + return "done" + + thread1 = {"configurable": {"thread_id": "1"}} + assert await my_workflow.ainvoke(1, thread1) == "done" From 0ebb78d9b202782d65da35e03619c9f65e19d532 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 23 May 2025 13:50:11 -0700 Subject: [PATCH 2/5] Fix --- libs/langgraph/langgraph/pregel/loop.py | 1 - libs/langgraph/langgraph/pregel/runner.py | 7 +++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 412266ad3..d4c075a99 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -135,7 +135,6 @@ P = ParamSpec("P") INPUT_DONE = object() INPUT_RESUMING = object() INPUT_SHOULD_VALIDATE = object() -SPECIAL_CHANNELS = (ERROR, INTERRUPT, SCHEDULED) WritesT = Sequence[tuple[str, Any]] diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 22d7a290a..bbdb90bae 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -56,7 +56,9 @@ EXCLUDED_FRAME_FNAMES = ( "concurrent/futures/_base.py", ) -SKIP_RERAISE_SET = weakref.WeakSet() +SKIP_RERAISE_SET: weakref.WeakSet[Union[concurrent.futures.Future, asyncio.Future]] = ( + weakref.WeakSet() +) class FuturesDict(Generic[F, E], dict[F, Optional[PregelExecutableTask]]): @@ -432,7 +434,8 @@ class PregelRunner: raise exception else: # save error to checkpointer - self.put_writes()(task.id, [(ERROR, exception)]) # type: ignore[misc] + task.writes.append((ERROR, exception)) + self.put_writes()(task.id, task.writes) # type: ignore[misc] else: if self.node_finished and ( task.config is None or TAG_HIDDEN not in task.config.get("tags", []) From 3c0d9346c278f628ad1b340f4579eef63e08dda6 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 23 May 2025 13:51:13 -0700 Subject: [PATCH 3/5] Add sync test --- libs/langgraph/tests/test_pregel.py | 27 +++++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 4 ++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 5558daa68..1a6edb2bb 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8769,3 +8769,30 @@ def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None: assert json.dumps(graph.get_graph().to_json(), indent=2) == snapshot assert graph.get_graph().draw_mermaid(with_styles=False) == snapshot + + +def test_imp_exception( + checkpointer: BaseCheckpointSaver, +) -> None: + @task() + def my_task(number: int): + time.sleep(0.1) + return number * 2 + + @task() + def task_with_exception(number: int): + time.sleep(0.1) + raise Exception("This is a test exception") + + @entrypoint(checkpointer=checkpointer) + def my_workflow(number: int): + my_task(number) + try: + task_with_exception(number) + except Exception as e: + print(f"Exception caught: {e}") + my_task(number) + return "done" + + thread1 = {"configurable": {"thread_id": "1"}} + assert my_workflow.invoke(1, thread1) == "done" diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 58745ec25..f77612241 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -9156,12 +9156,12 @@ async def test_imp_exception( ) -> None: @task() async def my_task(number: int): - await asyncio.sleep(1) + await asyncio.sleep(0.1) return number * 2 @task() async def task_with_exception(number: int): - await asyncio.sleep(1) + await asyncio.sleep(0.1) raise Exception("This is a test exception") @entrypoint(checkpointer=async_checkpointer) From 913b8d5e9547d78523e75304d473555d7347546f Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 23 May 2025 14:32:45 -0700 Subject: [PATCH 4/5] Lint --- libs/langgraph/tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 1a6edb2bb..45b6ec1d4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8772,7 +8772,7 @@ def test_get_graph_root_channel(snapshot: SnapshotAssertion) -> None: def test_imp_exception( - checkpointer: BaseCheckpointSaver, + sync_checkpointer: BaseCheckpointSaver, ) -> None: @task() def my_task(number: int): @@ -8784,7 +8784,7 @@ def test_imp_exception( time.sleep(0.1) raise Exception("This is a test exception") - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=sync_checkpointer) def my_workflow(number: int): my_task(number) try: From bece43dc67024bcd0d4d7649323df3104f56949c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 23 May 2025 15:06:18 -0700 Subject: [PATCH 5/5] Fix --- libs/langgraph/tests/test_pregel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 45b6ec1d4..81b5eadf8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8786,12 +8786,12 @@ def test_imp_exception( @entrypoint(checkpointer=sync_checkpointer) def my_workflow(number: int): - my_task(number) + my_task(number).result() try: - task_with_exception(number) + task_with_exception(number).result() except Exception as e: print(f"Exception caught: {e}") - my_task(number) + my_task(number).result() return "done" thread1 = {"configurable": {"thread_id": "1"}}