From f7dffa023c7129239499cf44f0511e4ee9f3a136 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 10 Sep 2024 15:16:35 -0700 Subject: [PATCH] Implement subgraph delegation for distributed arch --- libs/langgraph/langgraph/constants.py | 2 + libs/langgraph/langgraph/errors.py | 7 + libs/langgraph/langgraph/pregel/loop.py | 30 +- libs/langgraph/langgraph/pregel/runner.py | 6 +- .../langgraph/scheduler/kafka/executor.py | 58 +- .../langgraph/scheduler/kafka/orchestrator.py | 49 +- .../langgraph/scheduler/kafka/types.py | 4 +- libs/scheduler-kafka/tests/drain.py | 4 +- libs/scheduler-kafka/tests/test_fanout.py | 8 +- libs/scheduler-kafka/tests/test_subgraph.py | 554 +++++++++++++++++- 10 files changed, 691 insertions(+), 31 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 15dc47fd9..dd7efd6f7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -11,6 +11,7 @@ CONFIG_KEY_RESUMING = "__pregel_resuming" CONFIG_KEY_TASK_ID = "__pregel_task_id" CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks" CONFIG_KEY_ENSURE_LATEST = "__pregel_ensure_latest" +CONFIG_KEY_DELEGATE = "__pregel_delegate" # this one part of public API so more readable CONFIG_KEY_CHECKPOINT_MAP = "checkpoint_map" INTERRUPT = "__interrupt__" @@ -38,6 +39,7 @@ RESERVED = { CONFIG_KEY_TASK_ID, CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, + CONFIG_KEY_DELEGATE, INPUT, RUNTIME_PLACEHOLDER, } diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 2fcd64473..ec84e0b28 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -43,6 +43,13 @@ class NodeInterrupt(GraphInterrupt): super().__init__([Interrupt(value)]) +class GraphDelegate(Exception): + """Raised when a graph is delegated.""" + + def __init__(self, *args: dict[str, Any]) -> None: + super().__init__(*args) + + class EmptyInputError(Exception): """Raised when graph receives an empty input.""" diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index c0551c14b..9d96fbc3e 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -40,6 +40,7 @@ from langgraph.checkpoint.base import ( from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_DEDUPE_TASKS, + CONFIG_KEY_DELEGATE, CONFIG_KEY_ENSURE_LATEST, CONFIG_KEY_RESUMING, CONFIG_KEY_STREAM, @@ -50,7 +51,12 @@ from langgraph.constants import ( SCHEDULED, TAG_HIDDEN, ) -from langgraph.errors import CheckpointNotLatest, EmptyInputError, GraphInterrupt +from langgraph.errors import ( + CheckpointNotLatest, + EmptyInputError, + GraphDelegate, + GraphInterrupt, +) from langgraph.managed.base import ( ManagedValueMapping, ManagedValueSpec, @@ -337,6 +343,18 @@ class PregelLoop: self.status = "done" return False + # check if we should delegate (used by subgraphs in distributed mode) + if self.config["configurable"].get(CONFIG_KEY_DELEGATE): + assert self.input is INPUT_RESUMING + raise GraphDelegate( + { + "config": patch_configurable( + self.config, {CONFIG_KEY_DELEGATE: False} + ), + "input": None, + } + ) + # 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: @@ -412,6 +430,16 @@ class PregelLoop: ) # map inputs to channel updates elif input_writes := deque(map_input(input_keys, self.input)): + # check if we should delegate (used by subgraphs in distributed mode) + if self.config["configurable"].get(CONFIG_KEY_DELEGATE): + raise GraphDelegate( + { + "config": patch_configurable( + self.config, {CONFIG_KEY_DELEGATE: False} + ), + "input": self.input, + } + ) # discard any unfinished tasks from previous checkpoint discard_tasks = prepare_next_tasks( self.checkpoint, diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 086afa9da..7a09afde3 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -13,7 +13,7 @@ from typing import ( ) from langgraph.constants import ERROR, INTERRUPT, NO_WRITES -from langgraph.errors import GraphInterrupt +from langgraph.errors import GraphDelegate, GraphInterrupt from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry from langgraph.pregel.types import PregelExecutableTask, RetryPolicy @@ -71,6 +71,8 @@ class PregelRunner: # save interrupt to checkpointer if interrupts := [(INTERRUPT, i) for i in exc.args[0]]: self.put_writes(task.id, interrupts) + elif isinstance(exc, GraphDelegate): + raise exc else: # save error to checkpointer self.put_writes(task.id, [(ERROR, exc)]) @@ -135,6 +137,8 @@ class PregelRunner: # save interrupt to checkpointer if interrupts := [(INTERRUPT, i) for i in exc.args[0]]: self.put_writes(task.id, interrupts) + elif isinstance(exc, GraphDelegate): + raise exc else: # save error to checkpointer self.put_writes(task.id, [(ERROR, exc)]) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 714fbb647..f76c66052 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -4,11 +4,12 @@ from functools import partial from typing import Any, Optional, Self, Sequence import aiokafka +import orjson from langchain_core.runnables import RunnableConfig import langgraph.scheduler.kafka.serde as serde -from langgraph.constants import ERROR -from langgraph.errors import CheckpointNotLatest, TaskNotFound +from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP +from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel from langgraph.pregel.algo import prepare_single_task from langgraph.pregel.executor import AsyncBackgroundExecutor, Submit @@ -59,10 +60,14 @@ class KafkaExecutor(AbstractAsyncContextManager): ) self.producer = await self.stack.enter_async_context( aiokafka.AIOKafkaProducer( + key_serializer=serde.dumps, value_serializer=serde.dumps, **self.kwargs, ) ) + self.subgraphs = { + k: v async for k, v in self.graph.aget_subgraphs(recurse=True) + } return self async def __aexit__(self, *args: Any) -> None: @@ -94,6 +99,23 @@ class KafkaExecutor(AbstractAsyncContextManager): await aretry(self.retry_policy, self.attempt, msg) except CheckpointNotLatest: pass + except GraphDelegate as exc: + for arg in exc.args: + await self.producer.send_and_wait( + self.topics.orchestrator, + value=MessageToOrchestrator( + config=arg["config"], + input=orjson.Fragment( + self.graph.checkpointer.serde.dumps(arg["input"]) + ), + finally_executor=[msg], + ), + # use thread_id, checkpoint_ns as partition key + key=( + arg["config"]["configurable"]["thread_id"], + arg["config"]["configurable"].get("checkpoint_ns"), + ), + ) except Exception as exc: await self.producer.send_and_wait( self.topics.error, @@ -105,6 +127,19 @@ class KafkaExecutor(AbstractAsyncContextManager): ) async def attempt(self, msg: MessageToExecutor) -> None: + # find graph + if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): + # remove task_ids from checkpoint_ns + recast_checkpoint_ns = NS_SEP.join( + part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) + ) + # find the subgraph with the matching name + if recast_checkpoint_ns in self.subgraphs: + graph = self.subgraphs[recast_checkpoint_ns] + else: + raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + else: + graph = self.graph # process message saved = await self.graph.checkpointer.aget_tuple( patch_configurable(msg["config"], {"checkpoint_id": None}) @@ -114,17 +149,17 @@ class KafkaExecutor(AbstractAsyncContextManager): if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() async with AsyncChannelsManager( - self.graph.channels, saved.checkpoint, msg["config"], self.graph.store + graph.channels, saved.checkpoint, msg["config"], self.graph.store ) as (channels, managed), AsyncBackgroundExecutor() as submit: if task := await asyncio.to_thread( prepare_single_task, msg["task"]["path"], msg["task"]["id"], checkpoint=saved.checkpoint, - processes=self.graph.nodes, + processes=graph.nodes, channels=channels, managed=managed, - config=msg["config"], + config=patch_configurable(msg["config"], {CONFIG_KEY_DELEGATE: True}), step=saved.metadata["step"] + 1, for_execution=True, checkpointer=self.graph.checkpointer, @@ -144,9 +179,16 @@ class KafkaExecutor(AbstractAsyncContextManager): # notify orchestrator await self.producer.send_and_wait( self.topics.orchestrator, - value=MessageToOrchestrator(input=None, config=msg["config"]), - # use thread_id as partition key - key=msg["config"]["configurable"]["thread_id"].encode(), + value=MessageToOrchestrator( + input=None, + config=msg["config"], + finally_executor=msg.get("finally_executor"), + ), + # use thread_id, checkpoint_ns as partition key + key=( + msg["config"]["configurable"]["thread_id"], + msg["config"]["configurable"].get("checkpoint_ns"), + ), ) def _put_writes( diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 9abadf743..9d548280d 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -10,9 +10,11 @@ from langgraph.constants import ( CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, INTERRUPT, + NS_END, + NS_SEP, SCHEDULED, ) -from langgraph.errors import CheckpointNotLatest +from langgraph.errors import CheckpointNotLatest, GraphInterrupt from langgraph.pregel import Pregel from langgraph.pregel.loop import AsyncPregelLoop from langgraph.pregel.types import RetryPolicy @@ -60,6 +62,9 @@ class KafkaOrchestrator(AbstractAsyncContextManager): self.producer = await self.stack.enter_async_context( aiokafka.AIOKafkaProducer(value_serializer=serde.dumps, **self.kwargs) ) + self.subgraphs = { + k: v async for k, v in self.graph.aget_subgraphs(recurse=True) + } return self async def __aexit__(self, *args: Any) -> None: @@ -91,6 +96,8 @@ class KafkaOrchestrator(AbstractAsyncContextManager): await aretry(self.retry_policy, self.attempt, msg) except CheckpointNotLatest: pass + except GraphInterrupt: + pass except Exception as exc: await self.producer.send_and_wait( self.topics.error, @@ -102,6 +109,19 @@ class KafkaOrchestrator(AbstractAsyncContextManager): ) async def attempt(self, msg: MessageToOrchestrator) -> None: + # find graph + if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): + # remove task_ids from checkpoint_ns + recast_checkpoint_ns = NS_SEP.join( + part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) + ) + # find the subgraph with the matching name + if recast_checkpoint_ns in self.subgraphs: + graph = self.subgraphs[recast_checkpoint_ns] + else: + raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + else: + graph = self.graph # process message async with AsyncPregelLoop( msg["input"], @@ -109,15 +129,15 @@ class KafkaOrchestrator(AbstractAsyncContextManager): stream=None, store=self.graph.store, checkpointer=self.graph.checkpointer, - nodes=self.graph.nodes, - specs=self.graph.channels, - output_keys=self.graph.output_channels, - stream_keys=self.graph.stream_channels, + nodes=graph.nodes, + specs=graph.channels, + output_keys=graph.output_channels, + stream_keys=graph.stream_channels, ) as loop: if loop.tick( - input_keys=self.graph.input_channels, - interrupt_after=self.graph.interrupt_after_nodes, - interrupt_before=self.graph.interrupt_before_nodes, + input_keys=graph.input_channels, + interrupt_after=graph.interrupt_after_nodes, + interrupt_before=graph.interrupt_before_nodes, ): # wait for checkpoint to be saved if hasattr(loop, "_put_checkpoint_fut"): @@ -139,6 +159,7 @@ class KafkaOrchestrator(AbstractAsyncContextManager): }, ), task=ExecutorTask(id=task.id, path=task.path), + finally_executor=msg.get("finally_executor"), ), ) for task in new_tasks @@ -162,5 +183,13 @@ class KafkaOrchestrator(AbstractAsyncContextManager): ) ], ) - else: - pass + elif loop.status == "done" and msg.get("finally_executor"): + # schedule any finally_executor tasks + futs = await asyncio.gather( + *( + self.producer.send(self.topics.executor, value=m) + for m in msg["finally_executor"] + ) + ) + # wait for messages to be sent + await asyncio.gather(*futs) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py index a63057319..3bc298b93 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/types.py @@ -1,4 +1,4 @@ -from typing import Any, NamedTuple, Optional, TypedDict, Union +from typing import Any, NamedTuple, Optional, Sequence, TypedDict, Union from langchain_core.runnables import RunnableConfig @@ -12,6 +12,7 @@ class Topics(NamedTuple): class MessageToOrchestrator(TypedDict): input: Optional[dict[str, Any]] config: RunnableConfig + finally_executor: Optional[Sequence["MessageToExecutor"]] class ExecutorTask(TypedDict): @@ -22,6 +23,7 @@ class ExecutorTask(TypedDict): class MessageToExecutor(TypedDict): config: RunnableConfig task: ExecutorTask + finally_executor: Optional[Sequence["MessageToExecutor"]] class ErrorMessage(TypedDict): diff --git a/libs/scheduler-kafka/tests/drain.py b/libs/scheduler-kafka/tests/drain.py index 6543dd0ee..b58906362 100644 --- a/libs/scheduler-kafka/tests/drain.py +++ b/libs/scheduler-kafka/tests/drain.py @@ -47,14 +47,14 @@ async def drain_topics( async for msgs in orch: orch_msgs.extend(msgs) if debug: - print("orch", len(msgs)) + print("\n---\norch", len(msgs), 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)) + print("\n---\nexec", len(msgs), msgs) async def error_consumer() -> None: async with AIOKafkaConsumer(topics.error) as consumer: diff --git a/libs/scheduler-kafka/tests/test_fanout.py b/libs/scheduler-kafka/tests/test_fanout.py index 4e5e9586a..0ecc3d80d 100644 --- a/libs/scheduler-kafka/tests/test_fanout.py +++ b/libs/scheduler-kafka/tests/test_fanout.py @@ -16,7 +16,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics from tests.any import AnyDict -from tests.run import drain_topics +from tests.drain import drain_topics pytestmark = pytest.mark.anyio @@ -138,6 +138,7 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - "tags": [], }, "input": None, + "finally_executor": None, } for c in reversed(history) for _ in c.tasks @@ -162,6 +163,7 @@ async def test_fanout_graph(topics: Topics, checkpointer: BaseCheckpointSaver) - "id": t.id, "path": list(t.path), }, + "finally_executor": None, } for c in reversed(history) for t in c.tasks @@ -220,8 +222,11 @@ async def test_fanout_graph_w_interrupt( "tags": [], }, "input": None, + "finally_executor": None, } for c in reversed(history[1:]) # the last one wasn't executed + # orchestrator messages appear only after tasks for that checkpoint + # finish executing, ie. after executor sends message to resume checkpoint for _ in c.tasks ] assert exec_msgs == [ @@ -244,6 +249,7 @@ async def test_fanout_graph_w_interrupt( "id": t.id, "path": list(t.path), }, + "finally_executor": None, } for c in reversed(history[1:]) # the last one wasn't executed for t in c.tasks diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index f051d1369..a9c0062a8 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -15,8 +15,8 @@ from langgraph.graph.state import StateGraph from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyStr -from tests.run import drain_topics +from tests.any import AnyDict, AnyStr +from tests.drain import drain_topics pytestmark = pytest.mark.anyio C = ParamSpec("C") @@ -140,14 +140,282 @@ async def test_subgraph_w_interrupt( # check interrupted state state = await graph.aget_state(config) - assert len(orch_msgs) == 4 - assert len(exec_msgs) == 3 + assert len(orch_msgs) == 6 + assert len(exec_msgs) == 5 assert state.next == ("weather_graph",) assert state.values == { "messages": [HumanMessage(id=AnyStr(), content="what's the weather in sf")], "route": "weather", } + # check outer history + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 3 + + # check child history + child_history = [ + c async for c in graph.aget_state_history(history[0].tasks[0].state) + ] + assert len(child_history) == 3 + + # check messages + assert ( + orch_msgs + == ( + # initial message to outer graph + [MessageToOrchestrator(input=input, config=config)] + # outer graph messages, until interrupted + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__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, + "finally_executor": None, + } + for c in reversed(history[1:]) # the last one wasn't executed + # orchestrator messages appear only after tasks for that checkpoint + # finish executing, ie. after executor sends message to resume checkpoint + for _ in c.tasks + ] + # initial message to child graph + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "__pregel_task_id": history[0].tasks[0].id, + "checkpoint_id": None, + "checkpoint_map": { + "": history[0].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[0] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": { + "messages": [ + { + "id": [ + "langchain", + "schema", + "messages", + "HumanMessage", + ], + "kwargs": { + "content": "what's the weather in sf", + "id": AnyStr(), + "type": "human", + }, + "lc": 1, + "type": "constructor", + } + ], + "route": "weather", + }, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": False, + "checkpoint_id": history[0].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[0].tasks[0].id, + "path": list(history[0].tasks[0].path), + }, + } + ], + } + ] + # child graph messages, until interrupted + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "__pregel_task_id": history[0].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[0].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[0] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": False, + "checkpoint_id": history[0].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[0].tasks[0].id, + "path": list(history[0].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[1:]) # the last one wasn't executed + # orchestrator messages appear only after tasks for that checkpoint + # finish executing, ie. after executor sends message to resume checkpoint + for _ in c.tasks + ] + ) + ) + assert ( + exec_msgs + == ( + # outer graph tasks + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__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), + }, + "finally_executor": None, + } + for c in reversed(history) + for t in c.tasks + ] + # child graph tasks + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": False, + "__pregel_task_id": history[0].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[0].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[0] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": False, + "checkpoint_id": history[0].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[0].tasks[0].id, + "path": list(history[0].tasks[0].path), + }, + } + ], + } + for c in reversed(child_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( @@ -156,13 +424,13 @@ async def test_subgraph_w_interrupt( ) orch_msgs, exec_msgs = await drain_topics( - topics, graph, config, until=lambda state: state.next == (), debug=True + topics, graph, config, until=lambda state: state.next == () ) # check final state state = await graph.aget_state(config) - assert len(orch_msgs) == 2 - assert len(exec_msgs) == 1 + assert len(orch_msgs) == 4 + assert len(exec_msgs) == 3 assert state.next == () assert state.values == { "messages": [ @@ -171,3 +439,275 @@ async def test_subgraph_w_interrupt( ], "route": "weather", } + + # check outer history + history = [c async for c in graph.aget_state_history(config)] + assert len(history) == 4 + + # check child history + # accessing second to last checkpoint, since that's the one w/ subgraph task + child_history = [ + c async for c in graph.aget_state_history(history[1].tasks[0].state) + ] + assert len(child_history) == 4 + + # check messages + assert ( + orch_msgs + == ( + # initial message to outer graph + [MessageToOrchestrator(input=None, config=config)] + # initial message to child graph + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "__pregel_task_id": history[1].tasks[0].id, + "checkpoint_id": None, + "checkpoint_map": { + "": history[1].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[1] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ], + } + ] + # child graph messages, from previous last checkpoint onwards + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "__pregel_task_id": history[1].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[1].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[1] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[:2]) + for _ in c.tasks + ] + # outer graph messages, from previous last checkpoint onwards + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "input": None, + "finally_executor": None, + } + for c in reversed(history[:2]) + for _ in c.tasks + ] + ) + ) + assert ( + exec_msgs + == ( + # outer graph tasks + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "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), + }, + "finally_executor": None, + } + for c in reversed(history[:2]) + for t in c.tasks + ] + # child graph tasks + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_checkpointer": None, + "__pregel_delegate": False, + "__pregel_read": None, + "__pregel_send": None, + "__pregel_ensure_latest": True, + "__pregel_dedupe_tasks": True, + "__pregel_resuming": True, + "__pregel_task_id": history[1].tasks[0].id, + "checkpoint_id": c.config["configurable"]["checkpoint_id"], + "checkpoint_map": { + "": history[1].config["configurable"]["checkpoint_id"] + }, + "checkpoint_ns": history[1] + .tasks[0] + .state["configurable"]["checkpoint_ns"], + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "task": { + "id": t.id, + "path": list(t.path), + }, + "finally_executor": [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ], + } + for c in reversed(child_history[:2]) + for t in c.tasks + ] + # "finally" tasks + + [ + { + "config": { + "callbacks": None, + "configurable": { + "__pregel_dedupe_tasks": True, + "__pregel_ensure_latest": True, + "__pregel_resuming": True, + "checkpoint_id": history[1].config["configurable"][ + "checkpoint_id" + ], + "checkpoint_ns": "", + "thread_id": "1", + }, + "metadata": AnyDict(), + "recursion_limit": 25, + "tags": [], + }, + "finally_executor": None, + "task": { + "id": history[1].tasks[0].id, + "path": list(history[1].tasks[0].path), + }, + } + ] + ) + )