From 5de46f7472de8e7b4854b7f0add5d2d13f681b17 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 10 Sep 2024 16:51:49 -0700 Subject: [PATCH] Update stop condition for drain --- libs/scheduler-kafka/tests/drain.py | 46 +++++++++------------ libs/scheduler-kafka/tests/test_fanout.py | 12 ++---- libs/scheduler-kafka/tests/test_subgraph.py | 11 +---- 3 files changed, 25 insertions(+), 44 deletions(-) diff --git a/libs/scheduler-kafka/tests/drain.py b/libs/scheduler-kafka/tests/drain.py index 2a6ed426b..21fcd4d4c 100644 --- a/libs/scheduler-kafka/tests/drain.py +++ b/libs/scheduler-kafka/tests/drain.py @@ -1,13 +1,11 @@ import asyncio -from typing import Callable, Optional, TypeVar +from typing import 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 from langgraph.scheduler.kafka.executor import KafkaExecutor from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics @@ -17,31 +15,38 @@ R = TypeVar("R") async def drain_topics( - topics: Topics, - graph: Pregel, - config: RunnableConfig, - *, - until: Callable[[StateSnapshot], bool], - debug: bool = False, + topics: Topics, graph: Pregel, *, debug: bool = False ) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]: scope: Optional[anyio.CancelScope] = None orch_msgs = [] exec_msgs = [] errors = [] + def done() -> bool: + return ( + len(orch_msgs) > 0 + and len(exec_msgs) > 0 + and not orch_msgs[-1] + and not exec_msgs[-1] + ) + async def orchestrator() -> None: async with KafkaOrchestrator(graph, topics) as orch: async for msgs in orch: - orch_msgs.extend(msgs) + orch_msgs.append(msgs) if debug: print("\n---\norch", len(msgs), msgs) + if done(): + scope.cancel() async def executor() -> None: async with KafkaExecutor(graph, topics) as exec: async for msgs in exec: - exec_msgs.extend(msgs) + exec_msgs.append(msgs) if debug: print("\n---\nexec", len(msgs), msgs) + if done(): + scope.cancel() async def error_consumer() -> None: async with AIOKafkaConsumer(topics.error) as consumer: @@ -50,18 +55,8 @@ async def drain_topics( if scope: scope.cancel() - async def poller(expected_next: tuple[str, ...]) -> None: - while True: - await asyncio.sleep(0.5) - state = await graph.aget_state(config) - if until(state): - break - if scope: - scope.cancel() - - # start error consumer and poller + # start error consumer error_task = asyncio.create_task(error_consumer(), name="error_consumer") - poller_task = asyncio.create_task(poller(()), name="poller") # run the orchestrator and executor until break_when async with anyio.create_task_group() as tg: @@ -70,16 +65,15 @@ async def drain_topics( tg.start_soon(orchestrator, name="orchestrator") tg.start_soon(executor, name="executor") - # cancel error consumer and poller + # cancel error consumer error_task.cancel() - poller_task.cancel() try: - await asyncio.gather(error_task, poller_task) + await error_task except asyncio.CancelledError: pass # check no errors assert not errors, errors - return orch_msgs, exec_msgs + return [m for mm in orch_msgs for m in mm], [m for mm in exec_msgs for m in mm] diff --git a/libs/scheduler-kafka/tests/test_fanout.py b/libs/scheduler-kafka/tests/test_fanout.py index 0ecc3d80d..9f3726648 100644 --- a/libs/scheduler-kafka/tests/test_fanout.py +++ b/libs/scheduler-kafka/tests/test_fanout.py @@ -99,9 +99,7 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - ) # drain topics - orch_msgs, exec_msgs = await drain_topics( - topics, graph, config, until=lambda s: s.values and s.next == () - ) + orch_msgs, exec_msgs = await drain_topics(topics, graph) # check state state = await graph.aget_state(config) @@ -184,9 +182,7 @@ async def test_fanout_graph_w_interrupt( MessageToOrchestrator(input=input, config=config), ) - orch_msgs, exec_msgs = await drain_topics( - topics, graph, config, until=lambda s: s.values and s.next == ("qa",) - ) + orch_msgs, exec_msgs = await drain_topics(topics, graph) # check interrupted state state = await graph.aget_state(config) @@ -262,9 +258,7 @@ async def test_fanout_graph_w_interrupt( MessageToOrchestrator(input=None, config=config), ) - orch_msgs, exec_msgs = await drain_topics( - topics, graph, config, until=lambda s: s.values and s.next == () - ) + orch_msgs, exec_msgs = await drain_topics(topics, graph) # check final state state = await graph.aget_state(config) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index fa4ec50d3..ab3b4d20c 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -129,12 +129,7 @@ async def test_subgraph_w_interrupt( MessageToOrchestrator(input=input, config=config), ) - orch_msgs, exec_msgs = await drain_topics( - topics, - graph, - config, - until=lambda state: state.next == ("weather_graph",), - ) + orch_msgs, exec_msgs = await drain_topics(topics, graph) # check interrupted state state = await graph.aget_state(config) @@ -419,9 +414,7 @@ async def test_subgraph_w_interrupt( MessageToOrchestrator(input=None, config=config), ) - orch_msgs, exec_msgs = await drain_topics( - topics, graph, config, until=lambda state: state.next == () - ) + orch_msgs, exec_msgs = await drain_topics(topics, graph) # check final state state = await graph.aget_state(config)