WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-28 18:59:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

Improvements to consumer and producer patterns

Browse Source 
- 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
This commit is contained in:
Nuno Campos
2024-09-10 16:17:22 -07:00
parent dc5be3fd58
commit e6d6f565a5
4 changed files with 179 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+69 -31
View File
@@ -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)
libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py
+74 -42
View File
@@ -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
libs/scheduler-kafka/pyproject.toml
+1 -1
View File
@@ -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"]
libs/scheduler-kafka/tests/test_scheduler.py
+35 -21
View File
@@ -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