From 2d02e5e0c5d2f5ef286b4ee90d1c02b4f9cdc0c5 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 6 Sep 2024 15:18:13 -0700 Subject: [PATCH] Add error topic, add retries (eg for checkpointer database exceptions) --- .../langgraph/scheduler/kafka/executor.py | 20 ++++++- .../langgraph/scheduler/kafka/orchestrator.py | 20 ++++++- .../langgraph/scheduler/kafka/retry.py | 52 +++++++++++++++++++ .../langgraph/scheduler/kafka/types.py | 9 +++- libs/scheduler-kafka/tests/conftest.py | 6 ++- 5 files changed, 102 insertions(+), 5 deletions(-) create mode 100644 libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 1cf55266f..c81455c50 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,7 +1,7 @@ import asyncio from contextlib import AbstractAsyncContextManager, AsyncExitStack from functools import partial -from typing import Any, Self, Sequence +from typing import Any, Optional, Self, Sequence import aiokafka from langchain_core.runnables import RunnableConfig @@ -14,7 +14,10 @@ from langgraph.pregel.algo import prepare_single_task from langgraph.pregel.executor import AsyncBackgroundExecutor, Submit from langgraph.pregel.manager import AsyncChannelsManager from langgraph.pregel.runner import PregelRunner +from langgraph.pregel.types import RetryPolicy +from langgraph.scheduler.kafka.retry import aretry from langgraph.scheduler.kafka.types import ( + ErrorMessage, MessageToExecutor, MessageToOrchestrator, Topics, @@ -30,6 +33,7 @@ class KafkaExecutor(AbstractAsyncContextManager): group_id: str = "executor", batch_max_n: int = 10, batch_max_ms: int = 1000, + retry_policy: Optional[RetryPolicy] = None, **kwargs: Any, ) -> None: self.graph = graph @@ -39,6 +43,7 @@ class KafkaExecutor(AbstractAsyncContextManager): self.group_id = group_id self.batch_max_n = batch_max_n self.batch_max_ms = batch_max_ms + self.retry_policy = retry_policy async def __aenter__(self) -> Self: self.consumer = await self.stack.enter_async_context( @@ -84,6 +89,19 @@ class KafkaExecutor(AbstractAsyncContextManager): return msgs async def each(self, msg: MessageToExecutor) -> None: + try: + await aretry(self.retry_policy, self.attempt, msg) + except Exception as exc: + await self.producer.send_and_wait( + self.topics.error, + value=ErrorMessage( + topic=self.topics.executor, + msg=msg, + error=repr(exc), + ), + ) + + async def attempt(self, msg: MessageToExecutor) -> None: # process message saved = await self.graph.checkpointer.aget_tuple(msg["config"]) if saved is None: diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 408499f70..a8f41d99e 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -1,6 +1,6 @@ import asyncio from contextlib import AbstractAsyncContextManager, AsyncExitStack -from typing import Any, Self +from typing import Any, Optional, Self import aiokafka from langchain_core.runnables import ensure_config @@ -9,7 +9,10 @@ import langgraph.scheduler.kafka.serde as serde from langgraph.constants import CONFIG_KEY_DEDUPE_TASKS, SCHEDULED from langgraph.pregel import Pregel from langgraph.pregel.loop import INPUT_RESUMING, AsyncPregelLoop +from langgraph.pregel.types import RetryPolicy +from langgraph.scheduler.kafka.retry import aretry from langgraph.scheduler.kafka.types import ( + ErrorMessage, ExecutorTask, MessageToExecutor, MessageToOrchestrator, @@ -26,6 +29,7 @@ class KafkaOrchestrator(AbstractAsyncContextManager): group_id: str = "orchestrator", batch_max_n: int = 10, batch_max_ms: int = 1000, + retry_policy: Optional[RetryPolicy] = None, **kwargs: Any, ) -> None: self.graph = graph @@ -35,6 +39,7 @@ class KafkaOrchestrator(AbstractAsyncContextManager): self.group_id = group_id self.batch_max_n = batch_max_n self.batch_max_ms = batch_max_ms + self.retry_policy = retry_policy async def __aenter__(self) -> Self: self.consumer = await self.stack.enter_async_context( @@ -77,6 +82,19 @@ class KafkaOrchestrator(AbstractAsyncContextManager): return msgs async def each(self, msg: MessageToOrchestrator) -> None: + try: + await aretry(self.retry_policy, self.attempt, msg) + except Exception as exc: + await self.producer.send_and_wait( + self.topics.error, + value=ErrorMessage( + topic=self.topics.orchestrator, + msg=msg, + error=repr(exc), + ), + ) + + async def attempt(self, msg: MessageToOrchestrator) -> None: # process message async with AsyncPregelLoop( msg["input"], diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py new file mode 100644 index 000000000..52ad044b7 --- /dev/null +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py @@ -0,0 +1,52 @@ +import asyncio +import logging +import random +from typing import Awaitable, Callable, Optional, ParamSpec + +from langgraph.pregel.types import RetryPolicy + +logger = logging.getLogger(__name__) +P = ParamSpec("P") + + +async def aretry( + retry_policy: Optional[RetryPolicy], + func: Callable[P, Awaitable[None]], + *args: P.args, + **kwargs: P.kwargs, +) -> None: + """Run a task asynchronously with retries.""" + interval = retry_policy.initial_interval if retry_policy else 0 + attempts = 0 + while True: + try: + await func(*args, **kwargs) + # if successful, end + break + except Exception as exc: + if retry_policy is None: + raise + # 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 function {func} with {args} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}", + exc_info=exc, + ) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py index ce2a01d9b..78396c88e 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py @@ -1,4 +1,4 @@ -from typing import Any, NamedTuple, Optional, TypedDict +from typing import Any, NamedTuple, Optional, TypedDict, Union from langchain_core.runnables import RunnableConfig @@ -6,6 +6,7 @@ from langchain_core.runnables import RunnableConfig class Topics(NamedTuple): orchestrator: str executor: str + error: str class MessageToOrchestrator(TypedDict): @@ -23,3 +24,9 @@ class ExecutorTask(TypedDict): class MessageToExecutor(TypedDict): config: RunnableConfig task: ExecutorTask + + +class ErrorMessage(TypedDict): + topic: str + error: str + msg: Union[MessageToExecutor, MessageToOrchestrator] diff --git a/libs/scheduler-kafka/tests/conftest.py b/libs/scheduler-kafka/tests/conftest.py index 1d50ecc94..7d855a113 100644 --- a/libs/scheduler-kafka/tests/conftest.py +++ b/libs/scheduler-kafka/tests/conftest.py @@ -21,18 +21,20 @@ def anyio_backend(): def topics() -> Iterator[Topics]: o = f"test_{uuid4().hex[:16]}" e = f"test_{uuid4().hex[:16]}" + z = f"test_{uuid4().hex[:16]}" admin = kafka.admin.KafkaAdminClient() # create topics admin.create_topics( [ kafka.admin.NewTopic(name=o, num_partitions=1, replication_factor=1), kafka.admin.NewTopic(name=e, num_partitions=1, replication_factor=1), + kafka.admin.NewTopic(name=z, num_partitions=1, replication_factor=1), ] ) # yield topics - yield Topics(orchestrator=o, executor=e) + yield Topics(orchestrator=o, executor=e, error=z) # delete topics - admin.delete_topics([o, e]) + admin.delete_topics([o, e, z]) admin.close()