diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 2f947be22..d0f8e712e 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -518,6 +518,7 @@ def prepare_single_task( **configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}), parent_ns: checkpoint["id"], }, + "checkpoint_id": None, "checkpoint_ns": task_checkpoint_ns, }, ), diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7d4df7092..7772827c4 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -336,13 +336,6 @@ 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: @@ -350,7 +343,7 @@ class PregelLoop: continue if task := self.tasks.get(tid): if k == SCHEDULED: - if True or v == max( + if v == max( self.checkpoint["versions_seen"] .get(INTERRUPT, {}) .values(), diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 5b06180e8..91b53f31d 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -120,6 +120,7 @@ class KafkaExecutor(AbstractAsyncContextManager): config=msg["config"], step=saved.metadata["step"] + 1, for_execution=True, + checkpointer=self.graph.checkpointer, ): # execute task, saving writes runner = PregelRunner( @@ -146,5 +147,4 @@ class KafkaExecutor(AbstractAsyncContextManager): task_id: str, writes: list[tuple[str, Any]], ) -> None: - print("put_writes", task_id, writes) return submit(self.graph.checkpointer.aput_writes, config, writes, task_id) diff --git a/libs/scheduler-kafka/tests/run.py b/libs/scheduler-kafka/tests/run.py index 03b3108c4..6543dd0ee 100644 --- a/libs/scheduler-kafka/tests/run.py +++ b/libs/scheduler-kafka/tests/run.py @@ -92,6 +92,6 @@ async def drain_topics( pass # check no errors - assert not errors + assert not errors, errors return orch_msgs, exec_msgs diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 15bb52ac2..f051d1369 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -1,11 +1,7 @@ -import asyncio -import functools -import re -from typing import Callable, Literal, Optional, ParamSpec, TypeVar, Union, cast +from typing import Literal, ParamSpec, TypeVar, cast -import anyio import pytest -from aiokafka import AIOKafkaConsumer, AIOKafkaProducer +from aiokafka import AIOKafkaProducer from langchain_core.language_models.fake_chat_models import ( FakeMessagesListChatModel, ) @@ -18,43 +14,15 @@ from langgraph.graph import MessagesState 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 AnyStr +from tests.run import drain_topics pytestmark = pytest.mark.anyio C = ParamSpec("C") R = TypeVar("R") -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)) - - -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_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel: # copied from test_weather_subgraph @@ -149,45 +117,12 @@ def mk_weather_graph(checkpointer: BaseCheckpointSaver) -> Pregel: return graph.compile(checkpointer=checkpointer) -@timeout(10) async def test_subgraph_w_interrupt( topics: Topics, checkpointer: BaseCheckpointSaver ) -> None: input = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} config = {"configurable": {"thread_id": "1"}} graph = mk_weather_graph(checkpointer) - n_orch_msgs = 0 - n_exec_msgs = 0 - errors = [] - scope: Optional[anyio.CancelScope] = None - - async def orchestrator(expected: int) -> None: - nonlocal n_orch_msgs - async with KafkaOrchestrator(graph, topics) as orch: - async for msgs in orch: - n_orch_msgs += len(msgs) - print("orch", n_orch_msgs, 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: - n_exec_msgs += len(msgs) - print("exec", n_exec_msgs, msgs) - if n_exec_msgs == expected: - break - - async def error_consumer() -> None: - async with AIOKafkaConsumer(topics.error) as consumer: - async for msg in consumer: - errors.append(msg) - if scope: - scope.cancel() - - # start error consumer - error_task = asyncio.create_task(error_consumer(), name="error_consumer") # start a new run async with AIOKafkaProducer(value_serializer=serde.dumps) as producer: @@ -196,19 +131,17 @@ async def test_subgraph_w_interrupt( MessageToOrchestrator(input=input, config=config), ) - # run the orchestrator and executor - async with anyio.create_task_group() as tg: - scope = tg.cancel_scope - tg.start_soon(orchestrator, 4, name="orchestrator") - tg.start_soon(executor, 3, name="executor") - - # check no errors - assert not errors + orch_msgs, exec_msgs = await drain_topics( + topics, + graph, + config, + until=lambda state: state.next == ("weather_graph",), + ) # check interrupted state state = await graph.aget_state(config) - assert n_orch_msgs == 4 - assert n_exec_msgs == 3 + assert len(orch_msgs) == 4 + assert len(exec_msgs) == 3 assert state.next == ("weather_graph",) assert state.values == { "messages": [HumanMessage(id=AnyStr(), content="what's the weather in sf")], @@ -222,24 +155,19 @@ async def test_subgraph_w_interrupt( MessageToOrchestrator(input=None, config=config), ) - # run the orchestrator and executor - async with anyio.create_task_group() as tg: - scope = tg.cancel_scope - tg.start_soon(orchestrator, 6, name="orchestrator") - tg.start_soon(executor, 4, name="executor") - - # check no errors - assert not errors + orch_msgs, exec_msgs = await drain_topics( + topics, graph, config, until=lambda state: state.next == (), debug=True + ) # check final state state = await graph.aget_state(config) - assert n_orch_msgs == 6 - assert n_exec_msgs == 4 + assert len(orch_msgs) == 2 + assert len(exec_msgs) == 1 assert state.next == () assert state.values == { - "answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4", - "docs": ["doc1", "doc1", "doc2", "doc2", "doc3", "doc3", "doc4", "doc4"], - "query": "analyzed: query: analyzed: query: what is weather in sf", + "messages": [ + HumanMessage(id=AnyStr(), content="what's the weather in sf"), + AIMessage(content="I'ts sunny in San Francisco!", id=AnyStr()), + ], + "route": "weather", } - - error_task.cancel()