WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-22 15:42:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add retry_policy for graphs

Browse Source 
- default is to retry, fully configurable
- configuration options follow temporal https://docs.temporal.io/retry-policies#properties
This commit is contained in:
Nuno Campos
2024-05-16 15:44:30 -07:00
parent 6f888b6206
commit 2414e0ab5e
4 changed files with 191 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+35 -63
View File
@@ -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
langgraph/pregel/retry.py
+138
View File
@@ -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})"
)
tests/test_pregel.py
+9
View File
@@ -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
tests/test_pregel_async.py
+9
View File
@@ -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