Add readme, lint

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

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

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)

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

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: