Add error topic, add retries (eg for checkpointer database exceptions)

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:

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"],

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,
+            )

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]

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()