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