diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7772827c4..7d4df7092 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -336,6 +336,13 @@ class PregelLoop: self.status = "done" return False + print( + self.step, + self.skip_done_tasks, + [(t.id, t.name) for t in self.tasks.values()], + self.checkpoint_pending_writes, + ) + # if there are pending writes from a previous loop, apply them if self.skip_done_tasks and self.checkpoint_pending_writes: for tid, k, v in self.checkpoint_pending_writes: @@ -343,7 +350,7 @@ class PregelLoop: continue if task := self.tasks.get(tid): if k == SCHEDULED: - if v == max( + if True or v == max( self.checkpoint["versions_seen"] .get(INTERRUPT, {}) .values(), diff --git a/libs/scheduler-kafka/tests/any.py b/libs/scheduler-kafka/tests/any.py new file mode 100644 index 000000000..73744a1e8 --- /dev/null +++ b/libs/scheduler-kafka/tests/any.py @@ -0,0 +1,37 @@ +import re +from typing import Union + + +class AnyStr(str): + def __init__(self, prefix: Union[str, re.Pattern] = "") -> None: + super().__init__() + self.prefix = prefix + + def __eq__(self, other: object) -> bool: + return isinstance(other, str) and ( + other.startswith(self.prefix) + if isinstance(self.prefix, str) + else self.prefix.match(other) + ) + + def __hash__(self) -> int: + return hash((str(self), self.prefix)) + + +class AnyDict(dict): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + + def __eq__(self, other: object) -> bool: + if not self and isinstance(other, dict): + return True + if not isinstance(other, dict) or len(self) != len(other): + return False + for k, v in self.items(): + if kk := next((kk for kk in other if kk == k), None): + if v == other[kk]: + continue + else: + return False + else: + return True diff --git a/libs/scheduler-kafka/tests/conftest.py b/libs/scheduler-kafka/tests/conftest.py index 7d855a113..74b8a7800 100644 --- a/libs/scheduler-kafka/tests/conftest.py +++ b/libs/scheduler-kafka/tests/conftest.py @@ -19,9 +19,9 @@ def anyio_backend(): @pytest.fixture def topics() -> Iterator[Topics]: - o = f"test_{uuid4().hex[:16]}" - e = f"test_{uuid4().hex[:16]}" - z = f"test_{uuid4().hex[:16]}" + o = f"test_o_{uuid4().hex[:16]}" + e = f"test_e_{uuid4().hex[:16]}" + z = f"test_z_{uuid4().hex[:16]}" admin = kafka.admin.KafkaAdminClient() # create topics admin.create_topics( diff --git a/libs/scheduler-kafka/tests/run.py b/libs/scheduler-kafka/tests/run.py new file mode 100644 index 000000000..03b3108c4 --- /dev/null +++ b/libs/scheduler-kafka/tests/run.py @@ -0,0 +1,97 @@ +import asyncio +import functools +from typing import Callable, Optional, ParamSpec, TypeVar + +import anyio +from aiokafka import AIOKafkaConsumer +from langchain_core.runnables import RunnableConfig + +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 + +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 + + +@timeout(20) +async def drain_topics( + topics: Topics, + graph: Pregel, + config: RunnableConfig, + *, + until: Callable[[StateSnapshot], bool], + debug: bool = False, +) -> tuple[list[MessageToOrchestrator], list[MessageToOrchestrator]]: + scope: Optional[anyio.CancelScope] = None + orch_msgs = [] + exec_msgs = [] + errors = [] + + async def orchestrator() -> None: + async with KafkaOrchestrator(graph, topics) as orch: + async for msgs in orch: + orch_msgs.extend(msgs) + if debug: + print("orch", len(msgs)) + + async def executor() -> None: + async with KafkaExecutor(graph, topics) as exec: + async for msgs in exec: + exec_msgs.extend(msgs) + if debug: + print("exec", len(msgs)) + + async def error_consumer() -> None: + async with AIOKafkaConsumer(topics.error) as consumer: + async for msg in consumer: + errors.append(msg) + 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 + 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: + scope = tg.cancel_scope + tg.start_soon(orchestrator, name="orchestrator") + tg.start_soon(executor, name="executor") + + # cancel error consumer and poller + error_task.cancel() + poller_task.cancel() + + try: + await asyncio.gather(error_task, poller_task) + except asyncio.CancelledError: + pass + + # check no errors + assert not errors + + return orch_msgs, exec_msgs diff --git a/libs/scheduler-kafka/tests/test_fanout.py b/libs/scheduler-kafka/tests/test_fanout.py index ac8889045..e00787aee 100644 --- a/libs/scheduler-kafka/tests/test_fanout.py +++ b/libs/scheduler-kafka/tests/test_fanout.py @@ -1,35 +1,24 @@ import asyncio -import functools import operator -from typing import Annotated, Callable, ParamSpec, Sequence, TypedDict, TypeVar, Union +from typing import ( + Annotated, + Sequence, + TypedDict, + Union, +) -import anyio import pytest -from aiokafka import AIOKafkaConsumer, AIOKafkaProducer +from aiokafka import AIOKafkaProducer from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.graph.state import StateGraph from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde -from langgraph.scheduler.kafka.executor import KafkaExecutor -from langgraph.scheduler.kafka.orchestrator import KafkaOrchestrator from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics +from tests.any import AnyDict +from tests.run import drain_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( @@ -97,31 +86,10 @@ def mk_fanout_graph( return builder.compile(checkpointer, interrupt_before=interrupt_before) -@timeout(10) async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) -> None: input = {"query": "what is weather in sf"} config = {"configurable": {"thread_id": "1"}} graph = mk_fanout_graph(checkpointer) - n_orch_msgs = 0 - n_exec_msgs = 0 - - async def orchestrator(expected: int) -> None: - nonlocal n_orch_msgs - async with KafkaOrchestrator(graph, topics) as orch: - async for msgs in orch: - print("orch", msgs) - n_orch_msgs += len(msgs) - if n_orch_msgs == expected: - break - - async def executor(expected: int) -> None: - nonlocal n_exec_msgs - async with KafkaExecutor(graph, topics) as exec: - async for msgs in exec: - print("exec", msgs) - n_exec_msgs += len(msgs) - if n_exec_msgs == expected: - break # start a new run async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: @@ -130,20 +98,14 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - MessageToOrchestrator(input=input, config=config), ) - # run the orchestrator and executor - async with anyio.create_task_group() as tg: - tg.start_soon(orchestrator, 13, name="orchestrator") - tg.start_soon(executor, 12, name="executor") - - # check no errors - async with AIOKafkaConsumer(topics.error) as consumer: - assert len(consumer.assignment()) > 0 - for tp in consumer.assignment(): - assert await consumer.position(tp) == 0 + # drain topics + orch_msgs, exec_msgs = await drain_topics( + topics, graph, config, until=lambda s: s.values and s.next == () + ) + # check state state = await graph.aget_state(config) - assert n_orch_msgs == 13 - assert n_exec_msgs == 12 + assert state.next == () assert ( state.values == await graph.ainvoke(input, {"configurable": {"thread_id": "2"}}) @@ -154,34 +116,62 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - } ) + # check history + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 11 + + # check messages + assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + } + for c in reversed(history) + for _ in c.tasks + ] + assert exec_msgs == [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + } + for c in reversed(history) + for t in c.tasks + ] + -@timeout(10) async def test_fanout_graph_w_interrupt( topics: Topics, checkpointer: BaseCheckpointSaver ) -> None: input = {"query": "what is weather in sf"} config = {"configurable": {"thread_id": "1"}} graph = mk_fanout_graph(checkpointer, interrupt_before=["qa"]) - n_orch_msgs = 0 - n_exec_msgs = 0 - - async def orchestrator(expected: int) -> None: - nonlocal n_orch_msgs - async with KafkaOrchestrator(graph, topics) as orch: - async for msgs in orch: - print("orch", msgs) - n_orch_msgs += len(msgs) - if n_orch_msgs == expected: - break - - async def executor(expected: int) -> None: - nonlocal n_exec_msgs - async with KafkaExecutor(graph, topics) as exec: - async for msgs in exec: - print("exec", msgs) - n_exec_msgs += len(msgs) - if n_exec_msgs == expected: - break # start a new run async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: @@ -190,21 +180,12 @@ async def test_fanout_graph_w_interrupt( MessageToOrchestrator(input=input, config=config), ) - # run the orchestrator and executor - async with anyio.create_task_group() as tg: - tg.start_soon(orchestrator, 12, name="orchestrator") - tg.start_soon(executor, 11, name="executor") - - # check no errors - async with AIOKafkaConsumer(topics.error) as consumer: - assert len(consumer.assignment()) > 0 - for tp in consumer.assignment(): - assert await consumer.position(tp) == 0 + orch_msgs, exec_msgs = await drain_topics( + topics, graph, config, until=lambda s: s.values and s.next == ("qa",) + ) # check interrupted state state = await graph.aget_state(config) - assert n_orch_msgs == 12 - assert n_exec_msgs == 11 assert state.next == ("qa",) assert ( state.values @@ -215,6 +196,55 @@ async def test_fanout_graph_w_interrupt( } ) + # check history + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 10 + + # check messages + assert orch_msgs == [MessageToOrchestrator(input=input, config=config)] + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + } + for c in reversed(history[1:]) # the last one wasn't executed + for _ in c.tasks + ] + assert exec_msgs == [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + } + for c in reversed(history[1:]) # the last one wasn't executed + for t in c.tasks + ] + # resume the thread async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: await producer.send_and_wait( @@ -222,21 +252,12 @@ async def test_fanout_graph_w_interrupt( MessageToOrchestrator(input=None, config=config), ) - # run the orchestrator and executor - async with anyio.create_task_group() as tg: - tg.start_soon(orchestrator, 14, name="orchestrator") - tg.start_soon(executor, 12, name="executor") - - # check no errors - async with AIOKafkaConsumer(topics.error) as consumer: - assert len(consumer.assignment()) > 0 - for tp in consumer.assignment(): - assert await consumer.position(tp) == 0 + orch_msgs, exec_msgs = await drain_topics( + topics, graph, config, until=lambda s: s.values and s.next == () + ) # check final state state = await graph.aget_state(config) - assert n_orch_msgs == 14 - assert n_exec_msgs == 12 assert state.next == () assert ( state.values