diff --git a/libs/scheduler-kafka/README.md b/libs/scheduler-kafka/README.md index cedfcc53f..c279f5e38 100644 --- a/libs/scheduler-kafka/README.md +++ b/libs/scheduler-kafka/README.md @@ -1,3 +1,134 @@ # LangGraph Scheduler for Kafka -... +This library implements a distributed scheduler for LangGraph using Kafka as the message broker. + +## Architecture + +![](./langgraph-distributed.png) + +- Combination of Kafka (at least once) with a LangGraph Checkpointer provides exactly once semantics both for orchestrator and executor messages +- Checkpointer ensures writes for a given task are saved only once, even if the task is re-executed +- Checkpointer is used to record whether each task in each step has been successfully published to Kafka, to ensure tasks aren't lost, or published more than once +- Orchestrator and Executor manage commit of offsets manually to ensure tasks are marked as done only after finished processing +- Orchestrator and Executor pick up from the earliest message not yet consumed when restarted, to ensure no message is lost, and avoid processing messages more than once +- Orchestrator messages are keyed by thread ID and checkpoint NS, to ensure that no two consumers can process updates for same step of same thread concurrently +- Executor messages are not keyed, as they can be processed concurrently +- Orchestrator and Executor execute messages in configurable batches (up to N messages within space of X seconds), and dedupe messages intra-batch where appropriate (this is purely a performance optimization, with no impact on correctness whether applied or not) + +## Basic Usage + +Launch orchestrator and executor processes: + +`orchestrator.py` + +```python +import asyncio +import logging +import os + +from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator +from langgraph.scheduler.kafka.types import Topics + +from your_lib import graph # graph expected to be a compiled LangGraph graph + +logger = logging.getLogger(__name__) + +topics = Topics( + orchestrator: os.environ['KAFKA_TOPIC_ORCHESTRATOR'], + executor: os.environ['KAFKA_TOPIC_EXECUTOR'], + error: os.environ['KAFKA_TOPIC_ERROR'], +) + +async def main(): + async with KafkaOrchestrator(graph, topics) as orch: + async for msgs in orch: + logger.info('Procesed %d messages', len(msgs)) + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO) + asyncio.run(main()) +``` + +`executor.py` + +```python +import asyncio +import logging +import os + +from langgraph.scheduler.kafka.executor import KafkaExecutor +from langgraph.scheduler.kafka.types import Topics + +from your_lib import graph # graph expected to be a compiled LangGraph graph + +logger = logging.getLogger(__name__) + +topics = Topics( + orchestrator: os.environ['KAFKA_TOPIC_ORCHESTRATOR'], + executor: os.environ['KAFKA_TOPIC_EXECUTOR'], + error: os.environ['KAFKA_TOPIC_ERROR'], +) + +async def main(): + async with KafkaExecutor(graph, topics) as orch: + async for msgs in orch: + logger.info('Procesed %d messages', len(msgs)) + +if __name__ == '__main__': + logging.basicConfig(level=logging.INFO) + asyncio.run(main()) +``` + +```bash +export KAFKA_TOPIC_ORCHESTRATOR='orchestrator' +export KAFKA_TOPIC_EXECUTOR='executor' +export KAFKA_TOPIC_ERROR='error' +python orchestrator.py & +python executor.py & +``` + +## Configuration + +You can pass any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor` to configure the consumer: + +- group_id (str): a name for the consumer group. Defaults to 'orchestrator' or 'executor', respectively. +- batch_max_n (int): Maximum number of messages to include in a single batch. Default: 10. +- batch_max_ms (int): Maximum time in milliseconds to wait for messages to include in a batch. Default: 1000. +- retry_policy (langgraph.pregel.types.RetryPolicy): Controls which graph-level errors will be retried when processing messages. A good use for this is to retry database errors thrown by the checkpointer. Defaults to None. + +### Connection settings + +By default the orchestrator and executor will attempt to connect to a Kafka broker running on `localhost:9092`. You can change connection settings by passing any of the following values as `kwargs` to either `KafkaOrchestrator` or `KafkaExecutor`: + +- bootstrap_servers: 'host[:port]' string (or list of 'host[:port]' + strings) that the consumer should contact to bootstrap initial + cluster metadata. This does not have to be the full node list. + It just needs to have at least one broker that will respond to + Metadata API Request. Default port is 9092. If no servers are + specified, will default to localhost:9092. +- client_id (str): a name for this client. This string is passed in + each request to servers and can be used to identify specific + server-side log entries that correspond to this client. Also + submitted to GroupCoordinator for logging with respect to + consumer group administration. Default: 'aiokafka-{ver}' +- request_timeout_ms (int): Client request timeout in milliseconds. + Default: 40000. +- metadata_max_age_ms (int): The period of time in milliseconds after + which we force a refresh of metadata even if we haven't seen + any partition leadership changes to proactively discover any + new brokers or partitions. Default: 300000 +- retry_backoff_ms (int): Milliseconds to backoff when retrying on + errors. Default: 100. +- api_version (str): specify which kafka API version to use. + AIOKafka supports Kafka API versions >=0.9 only. + If set to 'auto', will attempt to infer the broker version by + probing various APIs. Default: auto +- security_protocol (str): Protocol used to communicate with brokers. + Valid values are: PLAINTEXT, SSL, SASL_PLAINTEXT, SASL_SSL. + Default: PLAINTEXT. +- ssl_context (ssl.SSLContext): pre-configured SSLContext for wrapping + socket connections. For more information see :ref:`ssl_auth`. + Default: None. +- connections_max_idle_ms (int): Close idle connections after the number + of milliseconds specified by this config. Specifying `None` will + disable idle checks. Default: 540000 (9 minutes). diff --git a/libs/scheduler-kafka/langgraph-distributed.png b/libs/scheduler-kafka/langgraph-distributed.png new file mode 100644 index 000000000..4315a01b5 Binary files /dev/null and b/libs/scheduler-kafka/langgraph-distributed.png differ diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index f76c66052..8aec3ef9e 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -36,12 +36,16 @@ class KafkaExecutor(AbstractAsyncContextManager): batch_max_n: int = 10, batch_max_ms: int = 1000, retry_policy: Optional[RetryPolicy] = None, + consumer_kwargs: Optional[dict[str, Any]] = None, + producer_kwargs: Optional[dict[str, Any]] = None, **kwargs: Any, ) -> None: self.graph = graph self.topics = topics self.stack = AsyncExitStack() self.kwargs = kwargs + self.consumer_kwargs = consumer_kwargs or {} + self.producer_kwargs = producer_kwargs or {} self.group_id = group_id self.batch_max_n = batch_max_n self.batch_max_ms = batch_max_ms diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 9d548280d..a3da188e8 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -38,12 +38,16 @@ class KafkaOrchestrator(AbstractAsyncContextManager): batch_max_n: int = 10, batch_max_ms: int = 1000, retry_policy: Optional[RetryPolicy] = None, + consumer_kwargs: Optional[dict[str, Any]] = None, + producer_kwargs: Optional[dict[str, Any]] = None, **kwargs: Any, ) -> None: self.graph = graph self.topics = topics self.stack = AsyncExitStack() self.kwargs = kwargs + self.consumer_kwargs = consumer_kwargs or {} + self.producer_kwargs = producer_kwargs or {} self.group_id = group_id self.batch_max_n = batch_max_n self.batch_max_ms = batch_max_ms @@ -57,10 +61,15 @@ class KafkaOrchestrator(AbstractAsyncContextManager): group_id=self.group_id, enable_auto_commit=False, **self.kwargs, + **self.consumer_kwargs, ) ) self.producer = await self.stack.enter_async_context( - aiokafka.AIOKafkaProducer(value_serializer=serde.dumps, **self.kwargs) + aiokafka.AIOKafkaProducer( + value_serializer=serde.dumps, + **self.kwargs, + **self.producer_kwargs, + ) ) self.subgraphs = { k: v async for k, v in self.graph.aget_subgraphs(recurse=True) diff --git a/libs/scheduler-kafka/tests/drain.py b/libs/scheduler-kafka/tests/drain.py index b58906362..b6199adbb 100644 --- a/libs/scheduler-kafka/tests/drain.py +++ b/libs/scheduler-kafka/tests/drain.py @@ -1,10 +1,11 @@ import asyncio import functools -from typing import Callable, Optional, ParamSpec, TypeVar +from typing import Callable, Optional, TypeVar import anyio from aiokafka import AIOKafkaConsumer from langchain_core.runnables import RunnableConfig +from typing_extensions import ParamSpec from langgraph.pregel import Pregel from langgraph.pregel.types import StateSnapshot diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index a9c0062a8..8c8a0f530 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -1,4 +1,4 @@ -from typing import Literal, ParamSpec, TypeVar, cast +from typing import Literal, cast import pytest from aiokafka import AIOKafkaProducer @@ -19,8 +19,6 @@ from tests.any import AnyDict, AnyStr from tests.drain import drain_topics pytestmark = pytest.mark.anyio -C = ParamSpec("C") -R = TypeVar("R") def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel: