From 2414e0ab5e0bd1fa4e4f910acf65bdf0ea92f4a8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 16 May 2024 15:44:30 -0700 Subject: [PATCH 1/3] Add retry_policy for graphs - default is to retry, fully configurable - configuration options follow temporal https://docs.temporal.io/retry-policies#properties --- langgraph/pregel/__init__.py | 98 +++++++++---------------- langgraph/pregel/retry.py | 138 +++++++++++++++++++++++++++++++++++ tests/test_pregel.py | 9 +++ tests/test_pregel_async.py | 9 +++ 4 files changed, 191 insertions(+), 63 deletions(-) create mode 100644 langgraph/pregel/retry.py diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 210d75028..145b67fef 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -21,6 +21,7 @@ from typing import ( overload, ) +from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager from langchain_core.globals import get_debug from langchain_core.load.dump import dumpd from langchain_core.pydantic_v1 import BaseModel, Field, root_validator @@ -92,6 +93,7 @@ from langgraph.pregel.io import ( ) from langgraph.pregel.log import logger from langgraph.pregel.read import PregelNode +from langgraph.pregel.retry import RetryPolicy, arun_with_retry, run_with_retry from langgraph.pregel.types import ( All, PregelExecutableTask, @@ -212,6 +214,9 @@ class Pregel( checkpointer: Optional[BaseCheckpointSaver] = None """Checkpointer used to save and load graph state. Defaults to None.""" + retry_policy: Optional[RetryPolicy] = RetryPolicy() + """Retry policy to use when running tasks. Set to None to disable.""" + config_type: Optional[Type[Any]] = None name: str = "LangGraph" @@ -487,7 +492,9 @@ class Pregel( saved = self.checkpointer.get_tuple(config) checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() # find last node that updated the state, if not provided - if as_node is None and not saved: + if as_node is None and not any( + v for vv in checkpoint["versions_seen"].values() for v in vv.values() + ): if ( isinstance(self.input_channels, str) and self.input_channels in self.nodes @@ -783,6 +790,7 @@ class Pregel( config, step, for_execution=True, + manager=run_manager, ) # if no more tasks, we're done @@ -809,30 +817,9 @@ class Pregel( for chunk in map_debug_tasks(step, next_tasks): yield chunk - # prepare tasks with config - tasks_w_config = [ - ( - proc, - input, - patch_config( - proc_config, - run_name=name, - callbacks=run_manager.get_child(f"graph:step:{step}"), - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: writes.extend, - CONFIG_KEY_READ: partial( - _local_read, checkpoint, channels, writes - ), - }, - ), - ) - for name, input, proc, writes, proc_config, _ in next_tasks - ] - futures = [ - executor.submit(proc.invoke, input, config) - for proc, input, config in tasks_w_config + executor.submit(run_with_retry, task, self.retry_policy) + for task in next_tasks ] # execute tasks, and wait for one to fail or all to finish. @@ -1074,6 +1061,7 @@ class Pregel( config, step, for_execution=True, + manager=run_manager, ) # if no more tasks, we're done @@ -1100,39 +1088,13 @@ class Pregel( for chunk in map_debug_tasks(step, next_tasks): yield chunk - # prepare tasks with config - tasks_w_config = [ - ( - proc, - input, - patch_config( - proc_config, - run_name=name, - callbacks=run_manager.get_child(f"graph:step:{step}"), - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: writes.extend, - CONFIG_KEY_READ: partial( - _local_read, checkpoint, channels, writes - ), - }, - ), + futures = [ + asyncio.create_task( + arun_with_retry(task, self.retry_policy, do_stream) ) - for name, input, proc, writes, proc_config, _ in next_tasks + for task in next_tasks ] - futures = ( - [ - asyncio.create_task(_aconsume(proc.astream(input, config))) - for proc, input, config in tasks_w_config - ] - if do_stream - else [ - asyncio.create_task(proc.ainvoke(input, config)) - for proc, input, config in tasks_w_config - ] - ) - # execute tasks, and wait for one to fail or all to finish. # each task is independent from all other concurrent tasks done, inflight = await asyncio.wait( @@ -1370,7 +1332,6 @@ def _panic_or_proceed( inflight.pop().cancel() # raise the exception raise exc - # TODO this is where retry of an entire step would happen if inflight: # if we got here means we timed out @@ -1469,6 +1430,7 @@ def _prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[False], + manager: Literal[None] = None, ) -> tuple[Checkpoint, list[PregelTaskDescription]]: ... @@ -1482,6 +1444,7 @@ def _prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[True], + manager: Union[ParentRunManager, AsyncParentRunManager], ) -> tuple[Checkpoint, list[PregelExecutableTask]]: ... @@ -1495,6 +1458,7 @@ def _prepare_next_tasks( step: int, *, for_execution: bool, + manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, ) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: checkpoint = copy_checkpoint(checkpoint) tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] @@ -1560,22 +1524,30 @@ def _prepare_next_tasks( if for_execution: if node := proc.get_node(): + writes = deque() tasks.append( PregelExecutableTask( name, val, node, - deque(), - merge_configs(config, proc.config), + writes, + patch_config( + merge_configs(config, proc.config), + run_name=name, + callbacks=manager.get_child(f"graph:step:{step}") + if manager + else None, + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: writes.extend, + CONFIG_KEY_READ: partial( + _local_read, checkpoint, channels, writes + ), + }, + ), triggers, ) ) else: tasks.append(PregelTaskDescription(name, val)) return checkpoint, tasks - - -async def _aconsume(iterator: AsyncIterator[Any]) -> None: - """Consume an async iterator.""" - async for _ in iterator: - pass diff --git a/langgraph/pregel/retry.py b/langgraph/pregel/retry.py new file mode 100644 index 000000000..3c8fc1baf --- /dev/null +++ b/langgraph/pregel/retry.py @@ -0,0 +1,138 @@ +import asyncio +import logging +import random +import time +from typing import Callable, NamedTuple, Union + +import httpx +import requests + +from langgraph.pregel.types import PregelExecutableTask + +logger = logging.getLogger(__name__) + + +def default_retry_on(exc: Exception) -> bool: + if isinstance( + exc, + ( + ValueError, + TypeError, + ArithmeticError, + ImportError, + LookupError, + NameError, + SyntaxError, + RuntimeError, + ), + ): + return False + if isinstance(exc, httpx.HTTPStatusError): + return 500 <= exc.response.status_code < 600 + if isinstance(exc, requests.HTTPError): + return 500 <= exc.response.status_code < 600 if exc.response else True + return True + + +class RetryPolicy(NamedTuple): + initial_interval: float = 0.5 + """Amount of time that must elapse before the first retry occurs. In seconds.""" + backoff_factor: float = 2.0 + """Multiplier by which the interval increases after each retry.""" + max_interval: float = 128.0 + """Maximum amount of time that may elapse between retries. In seconds.""" + max_attempts: int = 10 + """Maximum number of attempts to make before giving up, including the first.""" + jitter: bool = True + """Whether to add random jitter to the interval between retries.""" + retry_on: Union[ + tuple[Exception, ...], Callable[[Exception], bool] + ] = default_retry_on + """List of exceptions that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" + + +def run_with_retry( + task: PregelExecutableTask, + retry_policy: RetryPolicy, +) -> None: + """Run a task with retries.""" + interval = retry_policy.initial_interval + attempts = 0 + while True: + try: + # clear any writes from previous attempts + task.writes.clear() + # run the task + task.proc.invoke(task.input, task.config) + # if successful, end + break + except Exception as exc: + # increment attempts + attempts += 1 + # check if we should retry + if callable(retry_policy.retry_on): + if not retry_policy.retry_on(exc): + raise + elif not isinstance(exc, retry_policy.retry_on): + raise + # check if we should give up + if attempts >= retry_policy.max_attempts: + raise + # sleep before retrying + interval = min( + retry_policy.max_interval, + interval * retry_policy.backoff_factor, + ) + time.sleep( + interval + random.uniform(0, 1) if retry_policy.jitter else interval + ) + # log the retry + logger.info( + f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts})" + ) + + +async def arun_with_retry( + task: PregelExecutableTask, + retry_policy: RetryPolicy, + stream: bool = False, +) -> None: + """Run a task asynchronously with retries.""" + interval = retry_policy.initial_interval + attempts = 0 + while True: + try: + # clear any writes from previous attempts + task.writes.clear() + # run the task + if stream: + async for _ in task.proc.astream(task.input, task.config): + pass + else: + await task.proc.ainvoke(task.input, task.config) + # if successful, end + break + except Exception as exc: + # increment attempts + attempts += 1 + # check if we should retry + if callable(retry_policy.retry_on): + if not retry_policy.retry_on(exc): + raise + elif not isinstance(exc, retry_policy.retry_on): + raise + # check if we should give up + if attempts >= retry_policy.max_attempts: + raise + # sleep before retrying + interval = min( + retry_policy.max_interval, + interval * retry_policy.backoff_factor, + ) + await asyncio.sleep( + interval + random.uniform(0, 1) if retry_policy.jitter else interval + ) + # log the retry + logger.info( + f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts})" + ) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index f3ec83798..d4772b39e 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -729,8 +729,16 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + errored_once = False def raise_if_above_10(input: int) -> int: + nonlocal errored_once + if input > 4: + if errored_once: + pass + else: + errored_once = True + raise OSError("I will be retried") if input > 10: raise ValueError("Input is too large") return input @@ -763,6 +771,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: assert checkpoint["channel_values"].get("total") == 2 # total is now 2, so output is 2+3=5 assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" checkpoint = memory.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 40892621f..6ddf962c2 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -710,8 +710,16 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) async def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) + errored_once = False def raise_if_above_10(input: int) -> int: + nonlocal errored_once + if input > 4: + if errored_once: + pass + else: + errored_once = True + raise OSError("I will be retried") if input > 10: raise ValueError("Input is too large") return input @@ -744,6 +752,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: assert checkpoint["channel_values"].get("total") == 2 # total is now 2, so output is 2+3=5 assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5 + assert errored_once, "errored and retried" checkpoint = await memory.aget({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"].get("total") == 7 From a4e88e930f41c091011ffdc3c760cba43461a74d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 16 May 2024 15:59:45 -0700 Subject: [PATCH 2/3] Disable retries by default --- langgraph/pregel/__init__.py | 2 +- langgraph/pregel/retry.py | 10 +++++++--- tests/test_pregel.py | 2 ++ tests/test_pregel_async.py | 2 ++ 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 145b67fef..cb88512c4 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -214,7 +214,7 @@ class Pregel( checkpointer: Optional[BaseCheckpointSaver] = None """Checkpointer used to save and load graph state. Defaults to None.""" - retry_policy: Optional[RetryPolicy] = RetryPolicy() + retry_policy: Optional[RetryPolicy] = None """Retry policy to use when running tasks. Set to None to disable.""" config_type: Optional[Type[Any]] = None diff --git a/langgraph/pregel/retry.py b/langgraph/pregel/retry.py index 3c8fc1baf..de9699781 100644 --- a/langgraph/pregel/retry.py +++ b/langgraph/pregel/retry.py @@ -2,7 +2,7 @@ import asyncio import logging import random import time -from typing import Callable, NamedTuple, Union +from typing import Callable, NamedTuple, Optional, Union import httpx import requests @@ -53,7 +53,7 @@ class RetryPolicy(NamedTuple): def run_with_retry( task: PregelExecutableTask, - retry_policy: RetryPolicy, + retry_policy: Optional[RetryPolicy], ) -> None: """Run a task with retries.""" interval = retry_policy.initial_interval @@ -67,6 +67,8 @@ def run_with_retry( # if successful, end break except Exception as exc: + if retry_policy is None: + raise # increment attempts attempts += 1 # check if we should retry @@ -94,7 +96,7 @@ def run_with_retry( async def arun_with_retry( task: PregelExecutableTask, - retry_policy: RetryPolicy, + retry_policy: Optional[RetryPolicy], stream: bool = False, ) -> None: """Run a task asynchronously with retries.""" @@ -113,6 +115,8 @@ async def arun_with_retry( # if successful, end break except Exception as exc: + if retry_policy is None: + raise # increment attempts attempts += 1 # check if we should retry diff --git a/tests/test_pregel.py b/tests/test_pregel.py index d4772b39e..0154490f7 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -38,6 +38,7 @@ from langgraph.prebuilt.chat_agent_executor import ( ) from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr from tests.memory_assert import MemorySaverAssertImmutable @@ -762,6 +763,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None: input_channels="input", output_channels="output", checkpointer=memory, + retry_policy=RetryPolicy(), ) # total starts out as 0, so output is 0+2=2 diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 6ddf962c2..94b209c8e 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -38,6 +38,7 @@ from langgraph.prebuilt.chat_agent_executor import ( from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr from tests.memory_assert import MemorySaverAssertImmutable @@ -743,6 +744,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: input_channels="input", output_channels="output", checkpointer=memory, + retry_policy=RetryPolicy(), ) # total starts out as 0, so output is 0+2=2 From 53501f6aaf137d65ccf56975991cd3cea92ea9bb Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 16 May 2024 16:00:19 -0700 Subject: [PATCH 3/3] Lint --- langgraph/pregel/retry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/langgraph/pregel/retry.py b/langgraph/pregel/retry.py index de9699781..24d660e94 100644 --- a/langgraph/pregel/retry.py +++ b/langgraph/pregel/retry.py @@ -56,7 +56,7 @@ def run_with_retry( retry_policy: Optional[RetryPolicy], ) -> None: """Run a task with retries.""" - interval = retry_policy.initial_interval + interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 while True: try: @@ -100,7 +100,7 @@ async def arun_with_retry( stream: bool = False, ) -> None: """Run a task asynchronously with retries.""" - interval = retry_policy.initial_interval + interval = retry_policy.initial_interval if retry_policy else 0 attempts = 0 while True: try: