From e6d6f565a52f5bf981f317db7f729c330cd9735c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 6 Sep 2024 12:46:10 -0700 Subject: [PATCH] Improvements to consumer and producer patterns - await future returned by send() instead of flush() - use consumer groups by default - process tasks in batches by default, configurable - manually commit offsets when batch is processed --- .../langgraph/scheduler/kafka/executor.py | 100 ++++++++++----- .../langgraph/scheduler/kafka/orchestrator.py | 116 +++++++++++------- libs/scheduler-kafka/pyproject.toml | 2 +- libs/scheduler-kafka/tests/test_scheduler.py | 56 +++++---- 4 files changed, 179 insertions(+), 95 deletions(-) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index de1e61e09..82549cad9 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,15 +1,17 @@ import asyncio -from contextlib import AbstractAsyncContextManager -from typing import Any +from contextlib import AbstractAsyncContextManager, AsyncExitStack +from functools import partial +from typing import Any, Self, Sequence import aiokafka +from langchain_core.runnables import RunnableConfig import langgraph.scheduler.kafka.serde as serde from langgraph.constants import ERROR from langgraph.errors import TaskNotFound from langgraph.pregel import Pregel from langgraph.pregel.algo import prepare_single_task -from langgraph.pregel.executor import AsyncBackgroundExecutor +from langgraph.pregel.executor import AsyncBackgroundExecutor, Submit from langgraph.pregel.manager import AsyncChannelsManager from langgraph.pregel.runner import PregelRunner from langgraph.scheduler.kafka.types import ( @@ -20,35 +22,68 @@ from langgraph.scheduler.kafka.types import ( class KafkaExecutor(AbstractAsyncContextManager): - def __init__(self, graph: Pregel, topics: Topics, **kwargs: Any) -> None: + def __init__( + self, + graph: Pregel, + topics: Topics, + *, + group_id: str = "executor", + batch_max_n: int = 10, + batch_max_ms: int = 1000, + **kwargs: Any, + ) -> None: self.graph = graph self.topics = topics - self.consumer = aiokafka.AIOKafkaConsumer( - topics.executor, value_deserializer=serde.loads, **kwargs - ) - self.producer = aiokafka.AIOKafkaProducer( - value_serializer=serde.dumps, **kwargs - ) + self.stack = AsyncExitStack() + self.kwargs = kwargs + self.group_id = group_id + self.batch_max_n = batch_max_n + self.batch_max_ms = batch_max_ms - async def __aenter__(self) -> "KafkaExecutor": - await self.consumer.start() - await self.producer.start() + async def __aenter__(self) -> Self: + self.consumer = await self.stack.enter_async_context( + aiokafka.AIOKafkaConsumer( + self.topics.executor, + value_deserializer=serde.loads, + auto_offset_reset="earliest", + group_id=self.group_id, + enable_auto_commit=False, + **self.kwargs, + ) + ) + self.producer = await self.stack.enter_async_context( + aiokafka.AIOKafkaProducer( + value_serializer=serde.dumps, + **self.kwargs, + ) + ) return self async def __aexit__(self, *args: Any) -> None: - await self.consumer.stop() - await self.producer.stop() + await self.stack.__aexit__(*args) - def __aiter__(self) -> "KafkaExecutor": + def __aiter__(self) -> Self: return self - async def __anext__(self) -> Any: - # wait for next message + async def __anext__(self) -> Sequence[MessageToExecutor]: + # wait for next batch try: - rec = await self.consumer.getone() - msg: MessageToExecutor = rec.value + recs = await self.consumer.getmany( + timeout_ms=self.batch_max_ms, max_records=self.batch_max_n + ) + msgs: list[MessageToExecutor] = [ + msg.value for msgs in recs.values() for msg in msgs + ] except aiokafka.ConsumerStoppedError: raise StopAsyncIteration from None + # process batch + await asyncio.gather(*(self.each(msg) for msg in msgs)) + # commit offsets + await self.consumer.commit() + # return message + return msgs + + async def each(self, msg: MessageToExecutor) -> None: # process message saved = await self.graph.checkpointer.aget_tuple(msg["config"]) if saved is None: @@ -56,13 +91,6 @@ class KafkaExecutor(AbstractAsyncContextManager): async with AsyncChannelsManager( self.graph.channels, saved.checkpoint, msg["config"], self.graph.store ) as (channels, managed), AsyncBackgroundExecutor() as submit: - - def put_writes(task_id: str, writes: list[tuple[str, Any]]) -> None: - print("put_writes", task_id, writes) - return submit( - self.graph.checkpointer.aput_writes, msg["config"], writes, task_id - ) - if task := await asyncio.to_thread( prepare_single_task, msg["task"]["path"], @@ -77,7 +105,10 @@ class KafkaExecutor(AbstractAsyncContextManager): is_resuming=msg["task"]["resuming"], ): # execute task, saving writes - runner = PregelRunner(submit=submit, put_writes=put_writes) + runner = PregelRunner( + submit=submit, + put_writes=partial(self._put_writes, submit, msg["config"]), + ) async for _ in runner.atick([task]): pass else: @@ -86,9 +117,16 @@ class KafkaExecutor(AbstractAsyncContextManager): msg["config"], [(ERROR, TaskNotFound())] ) # notify orchestrator - await self.producer.send( + await self.producer.send_and_wait( self.topics.orchestrator, value=MessageToOrchestrator(input=None, config=msg["config"]), ) - # return message - return msg + + def _put_writes( + self, + submit: Submit, + config: RunnableConfig, + task_id: str, + writes: list[tuple[str, Any]], + ) -> None: + return submit(self.graph.checkpointer.aput_writes, config, writes, task_id) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 5187f0435..b8fe1f003 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -1,5 +1,6 @@ -from contextlib import AbstractAsyncContextManager -from typing import Any +import asyncio +from contextlib import AbstractAsyncContextManager, AsyncExitStack +from typing import Any, Self import aiokafka from langchain_core.runnables import ensure_config @@ -18,35 +19,64 @@ from langgraph.utils.config import patch_configurable class KafkaOrchestrator(AbstractAsyncContextManager): - def __init__(self, graph: Pregel, topics: Topics, **kwargs: Any) -> None: + def __init__( + self, + graph: Pregel, + topics: Topics, + group_id: str = "orchestrator", + batch_max_n: int = 10, + batch_max_ms: int = 1000, + **kwargs: Any, + ) -> None: self.graph = graph self.topics = topics - self.consumer = aiokafka.AIOKafkaConsumer( - topics.orchestrator, value_deserializer=serde.loads, **kwargs - ) - self.producer = aiokafka.AIOKafkaProducer( - value_serializer=serde.dumps, **kwargs - ) + self.stack = AsyncExitStack() + self.kwargs = kwargs + self.group_id = group_id + self.batch_max_n = batch_max_n + self.batch_max_ms = batch_max_ms - async def __aenter__(self) -> "KafkaOrchestrator": - await self.consumer.start() - await self.producer.start() + async def __aenter__(self) -> Self: + self.consumer = await self.stack.enter_async_context( + aiokafka.AIOKafkaConsumer( + self.topics.orchestrator, + value_deserializer=serde.loads, + auto_offset_reset="earliest", + group_id=self.group_id, + enable_auto_commit=False, + **self.kwargs, + ) + ) + self.producer = await self.stack.enter_async_context( + aiokafka.AIOKafkaProducer(value_serializer=serde.dumps, **self.kwargs) + ) return self async def __aexit__(self, *args: Any) -> None: - await self.consumer.stop() - await self.producer.stop() + await self.stack.__aexit__(*args) - def __aiter__(self) -> "KafkaOrchestrator": + def __aiter__(self) -> Self: return self - async def __anext__(self) -> Any: - # wait for next message + async def __anext__(self) -> list[MessageToOrchestrator]: + # wait for next batch try: - rec = await self.consumer.getone() - msg: MessageToOrchestrator = rec.value + recs = await self.consumer.getmany( + timeout_ms=self.batch_max_ms, max_records=self.batch_max_n + ) + msgs: list[MessageToOrchestrator] = [ + msg.value for msgs in recs.values() for msg in msgs + ] except aiokafka.ConsumerStoppedError: raise StopAsyncIteration from None + # process batch + await asyncio.gather(*(self.each(msg) for msg in msgs)) + # commit offsets + await self.consumer.commit() + # return message + return msgs + + async def each(self, msg: MessageToOrchestrator) -> None: # process message async with AsyncPregelLoop( msg["input"], @@ -60,35 +90,37 @@ class KafkaOrchestrator(AbstractAsyncContextManager): stream_keys=self.graph.stream_channels, ) as loop: if loop.tick(input_keys=self.graph.input_channels): + # wait for checkpoint to be saved if hasattr(loop, "_put_checkpoint_fut"): await loop._put_checkpoint_fut + # schedule any new tasks if new_tasks := [t for t in loop.tasks.values() if not t.scheduled]: # send messages to executor - for task in new_tasks: - if task.scheduled: - continue - await self.producer.send( - self.topics.executor, - value=MessageToExecutor( - config=patch_configurable( - loop.config, - { - **loop.checkpoint_config["configurable"], - CONFIG_KEY_DEDUPE_TASKS: True, - }, + futures: list[asyncio.Future] = await asyncio.gather( + *( + self.producer.send( + self.topics.executor, + value=MessageToExecutor( + config=patch_configurable( + loop.config, + { + **loop.checkpoint_config["configurable"], + CONFIG_KEY_DEDUPE_TASKS: True, + }, + ), + task=ExecutorTask( + id=task.id, + path=task.path, + step=loop.step, + resuming=loop.input is INPUT_RESUMING, + ), ), - task=ExecutorTask( - id=task.id, - path=task.path, - step=loop.step, - resuming=loop.input is INPUT_RESUMING, - ), - ), + ) + for task in new_tasks ) - # flush producer - await self.producer.flush() + ) + # wait for messages to be sent + await asyncio.gather(*futures) # mark as scheduled for task in new_tasks: loop.put_writes(task.id, [(SCHEDULED, None)]) - # return message - return msg diff --git a/libs/scheduler-kafka/pyproject.toml b/libs/scheduler-kafka/pyproject.toml index d107afcc8..1799acb44 100644 --- a/libs/scheduler-kafka/pyproject.toml +++ b/libs/scheduler-kafka/pyproject.toml @@ -53,5 +53,5 @@ lint.ignore = ["E501", "B008", "UP007", "UP006"] [tool.pytest-watcher] now = true delay = 0.1 -runner_args = ["-s", "--ff", "-v", "--tb", "short"] +runner_args = ["--ff", "-v", "--tb", "short"] patterns = ["*.py"] diff --git a/libs/scheduler-kafka/tests/test_scheduler.py b/libs/scheduler-kafka/tests/test_scheduler.py index b53abd361..d4cda6210 100644 --- a/libs/scheduler-kafka/tests/test_scheduler.py +++ b/libs/scheduler-kafka/tests/test_scheduler.py @@ -1,6 +1,7 @@ import asyncio +import functools import operator -from typing import Annotated, TypedDict, Union +from typing import Annotated, Callable, ParamSpec, TypedDict, TypeVar, Union import pytest from aiokafka import AIOKafkaProducer @@ -14,6 +15,20 @@ from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics pytestmark = pytest.mark.anyio +C = ParamSpec("C") +R = TypeVar("R") + + +def timeout(delay: int): + def decorator(func: Callable[C, R]) -> Callable[C, R]: + @functools.wraps(func) + async def new_func(*args: C.args, **kwargs: C.kwargs) -> R: + async with asyncio.timeout(delay): + return await func(*args, **kwargs) + + return new_func + + return decorator def mk_fanout_graph(checkpointer: BaseCheckpointSaver) -> Pregel: @@ -80,32 +95,32 @@ def mk_fanout_graph(checkpointer: BaseCheckpointSaver) -> Pregel: return builder.compile(checkpointer) +@timeout(5) async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None: graph = mk_fanout_graph(checkpointer) n_orch_msgs = 0 n_exec_msgs = 0 - async def orchestrator() -> None: + async def orchestrator(expected: int) -> None: nonlocal n_orch_msgs async with KafkaOrchestrator(graph, topics) as orch: - async for msg in orch: - n_orch_msgs += 1 - print("orch", msg) + async for msgs in orch: + n_orch_msgs += len(msgs) + print("orch", msgs) + if n_orch_msgs == expected: + break - async def executor() -> None: + async def executor(expected: int) -> None: nonlocal n_exec_msgs async with KafkaExecutor(graph, topics) as exec: - async for msg in exec: - n_exec_msgs += 1 - print("exec", msg) + async for msgs in exec: + n_exec_msgs += len(msgs) + print("exec", msgs) + if n_exec_msgs == expected: + break - async with asyncio.TaskGroup() as tg: - o = tg.create_task(orchestrator(), name="orchestrator") - e = tg.create_task(executor(), name="executor") - - # start a new run - producer = AIOKafkaProducer(value_serializer=serde.dumps) - await producer.start() + # start a new run + async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: await producer.send_and_wait( topics.orchestrator, MessageToOrchestrator( @@ -113,12 +128,11 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - config={"configurable": {"thread_id": "1"}}, ), ) - await producer.stop() - # wait for the run to finish - await asyncio.sleep(5) - o.cancel() - e.cancel() + # run the orchestrator and executor + async with asyncio.TaskGroup() as tg: + tg.create_task(orchestrator(13), name="orchestrator") + tg.create_task(executor(12), name="executor") assert n_orch_msgs == 13 assert n_exec_msgs == 12