From 007d7e72b124066dc1f96d2b203b0002b13d8f95 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 3 Dec 2024 15:04:55 -0800 Subject: [PATCH] Add test for cancellation --- libs/langgraph/langgraph/pregel/algo.py | 3 +- libs/langgraph/langgraph/pregel/runner.py | 24 +++++---- libs/langgraph/tests/test_pregel_async.py | 66 +++++++++++++++++++++++ 3 files changed, 83 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 4e17dd39d..b41424fe9 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -43,6 +43,7 @@ from langgraph.constants import ( CONFIG_KEY_TASK_ID, CONFIG_KEY_WRITES, EMPTY_SEQ, + ERROR, INTERRUPT, NO_WRITES, NS_END, @@ -270,7 +271,7 @@ def apply_writes( pending_writes_by_managed: dict[str, list[Any]] = defaultdict(list) for task in tasks: for chan, val in task.writes: - if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN): + if chan in (NO_WRITES, PUSH, RESUME, INTERRUPT, RETURN, ERROR): pass elif chan == TASKS: # TODO: remove branch in 1.0 checkpoint["pending_sends"].append(val) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6e68df66a..41c9e7e69 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -441,7 +441,12 @@ class PregelRunner: ) -> None: if fut is not None: exception = _exception(fut) - if exception: + if isinstance(exception, asyncio.CancelledError): + # for cancelled tasks, also save error in task, + # so loop can finish super-step + task.writes.append((ERROR, exception)) + self.put_writes(task.id, task.writes) + elif exception: if isinstance(exception, GraphInterrupt): # save interrupt to checkpointer if interrupts := [(INTERRUPT, i) for i in exception.args[0]]: @@ -472,11 +477,12 @@ def _should_stop_others( GraphInterrupts are not considered failures.""" for fut in done: if fut.cancelled(): - return True - if exc := fut.exception(): - return not isinstance(exc, GraphBubbleUp) - else: - return False + continue + elif exc := fut.exception(): + if not isinstance(exc, GraphBubbleUp): + return True + + return False def _exception( @@ -502,7 +508,9 @@ def _panic_or_proceed( done: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() inflight: set[Union[concurrent.futures.Future[Any], asyncio.Future[Any]]] = set() for fut in futs: - if fut.done(): + if fut.cancelled(): + continue + elif fut.done(): done.add(fut) else: inflight.add(fut) @@ -515,8 +523,6 @@ def _panic_or_proceed( # raise the exception if panic: raise exc - else: - return if inflight: # if we got here means we timed out while inflight: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c4e5174f6..9f568d5c3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -2648,6 +2648,10 @@ async def test_send_sequences(checkpointer_name: str) -> None: ] +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2690,6 +2694,64 @@ async def test_imp_task(checkpointer_name: str) -> None: assert mapper_calls == 2 +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_imp_task_cancel(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + mapper_calls = 0 + mapper_cancels = 0 + + @task() + async def mapper(input: int) -> str: + nonlocal mapper_calls, mapper_cancels + mapper_calls += 1 + try: + await asyncio.sleep(1) + except asyncio.CancelledError: + mapper_cancels += 1 + raise + return str(input) * 2 + + @entrypoint(checkpointer=checkpointer) + async def graph(input: list[int]) -> list[str]: + futures = [mapper(i) for i in input] + await asyncio.sleep(0.1) + futures.pop().cancel() # cancel one + mapped = await asyncio.gather(*futures) + answer = interrupt("question") + return [m + answer for m in mapped] + + thread1 = {"configurable": {"thread_id": "1"}} + assert [c async for c in graph.astream([0, 1], thread1)] == [ + {"mapper": "00"}, + { + "__interrupt__": ( + Interrupt( + value="question", + resumable=True, + ns=[AnyStr("graph:")], + when="during", + ), + ) + }, + ] + assert mapper_calls == 2 + assert mapper_cancels == 1 + + assert await graph.ainvoke(Command(resume="answer"), thread1) == [ + "00answer", + ] + assert mapper_calls == 3 + assert mapper_cancels == 2 + + +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_sync_from_async(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2722,6 +2784,10 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: ] +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_stream_order(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: