From d6492ef048618bbd9713aba2968041fcfa1140ed Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 15:32:48 -0800 Subject: [PATCH 1/6] Add support for multiple subgraphs called in a single node --- libs/langgraph/langgraph/errors.py | 15 ------------- libs/langgraph/langgraph/pregel/algo.py | 2 ++ libs/langgraph/langgraph/pregel/loop.py | 27 ++++++++++++++--------- libs/langgraph/langgraph/pregel/retry.py | 16 +------------- libs/langgraph/langgraph/types.py | 2 ++ libs/langgraph/tests/test_pregel.py | 7 +++--- libs/langgraph/tests/test_pregel_async.py | 7 +++--- 7 files changed, 27 insertions(+), 49 deletions(-) diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 0737a31d0..8e78a8784 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -107,18 +107,3 @@ class CheckpointNotLatest(Exception): """Raised when the checkpoint is not the latest version (for distributed mode).""" pass - - -class MultipleSubgraphsError(Exception): - """Raised when multiple subgraphs are called inside the same node. - - Troubleshooting guides: - - - [MULTIPLE_SUBGRAPHS](https://python.langchain.com/docs/troubleshooting/errors/MULTIPLE_SUBGRAPHS) - """ - - pass - - -_SEEN_CHECKPOINT_NS: set[str] = set() -"""Used for subgraph detection.""" diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c46138550..205793ab4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -766,6 +766,8 @@ def _scratchpad( (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), MISSING, ), + # subgraph + subgraph_counter=0, ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 9a745a2e7..ff244b7de 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -63,12 +63,10 @@ from langgraph.constants import ( TAG_HIDDEN, ) from langgraph.errors import ( - _SEEN_CHECKPOINT_NS, CheckpointNotLatest, EmptyInputError, GraphDelegate, GraphInterrupt, - MultipleSubgraphsError, ) from langgraph.managed.base import ( ManagedValueMapping, @@ -116,6 +114,7 @@ from langgraph.types import ( Command, LoopProtocol, PregelExecutableTask, + PregelScratchpad, StreamChunk, StreamProtocol, ) @@ -230,20 +229,26 @@ class PregelLoop(LoopProtocol): self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) + scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) + if scratchpad is not None: + if scratchpad["subgraph_counter"]: + self.config = patch_configurable( + self.config, + { + CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join( + ( + config[CONF][CONFIG_KEY_CHECKPOINT_NS], + str(scratchpad["subgraph_counter"]), + ) + ) + }, + ) + scratchpad["subgraph_counter"] += 1 if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): self.config = patch_configurable( self.config, {CONFIG_KEY_CHECKPOINT_NS: "", CONFIG_KEY_CHECKPOINT_ID: None}, ) - if check_subgraphs and self.is_nested and self.checkpointer is not None: - if self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] in _SEEN_CHECKPOINT_NS: - raise MultipleSubgraphsError( - "Multiple subgraphs called inside the same node\n\n" - "Troubleshooting URL: https://python.langchain.com/docs" - "/troubleshooting/errors/MULTIPLE_SUBGRAPHS/" - ) - else: - _SEEN_CHECKPOINT_NS.add(self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]) if ( CONFIG_KEY_CHECKPOINT_MAP in self.config[CONF] and self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 29faaab21..43e7e8d9e 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -12,7 +12,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, NS_SEP, ) -from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand +from langgraph.errors import GraphBubbleUp, ParentCommand from langgraph.types import Command, PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable @@ -96,13 +96,6 @@ def run_with_retry( ) # signal subgraphs to resume (if available) config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) - finally: - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) async def arun_with_retry( @@ -188,10 +181,3 @@ async def arun_with_retry( ) # signal subgraphs to resume (if available) config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) - finally: - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 3b1bcd213..09f777d01 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -346,6 +346,8 @@ class PregelScratchpad(TypedDict): interrupt_counter: int resume: list[Any] null_resume: Any + # subgraph + subgraph_counter: int def interrupt(value: Any) -> Any: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 67536deed..76a067533 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -51,7 +51,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import MemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START -from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError +from langgraph.errors import InvalidUpdateError from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages @@ -1745,9 +1745,8 @@ def test_invoke_join_then_call_other_pregel( # add checkpointer app.checkpointer = checkpointer - # subgraph is called twice in the same node, through .map(), so raises - with pytest.raises(MultipleSubgraphsError): - app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) + # subgraph is called twice in the same node, but that works + assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 # set inner graph checkpointer NeverCheckpoint inner_app.checkpointer = False diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 2f7d7a1c4..a6fab066f 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -48,7 +48,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import MemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START -from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt +from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessagesState, add_messages @@ -4068,9 +4068,8 @@ async def test_invoke_join_then_call_other_pregel( async with awith_checkpointer(checkpointer_name) as checkpointer: # add checkpointer app.checkpointer = checkpointer - # subgraph is called twice in the same node, through .map(), so raises - with pytest.raises(MultipleSubgraphsError): - await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) + # subgraph is called twice, and that works + assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 # set inner graph checkpointer NeverCheckpoint inner_app.checkpointer = False From e8a73e1505502b9c8a7e2984a59cd42d788f2648 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 15:35:06 -0800 Subject: [PATCH 2/6] Remove flag --- libs/langgraph/langgraph/pregel/__init__.py | 1 - libs/langgraph/langgraph/pregel/loop.py | 5 ----- .../langgraph/scheduler/kafka/orchestrator.py | 2 -- 3 files changed, 8 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 102e68be8..0a90aa10f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1642,7 +1642,6 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - check_subgraphs=self.checkpointer is not True, ) as loop: # create runner runner = PregelRunner( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index ff244b7de..0afe62658 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -202,7 +202,6 @@ class PregelLoop(LoopProtocol): interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, - check_subgraphs: bool = True, debug: bool = False, ) -> None: super().__init__( @@ -822,7 +821,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - check_subgraphs: bool = True, debug: bool = False, ) -> None: super().__init__( @@ -837,7 +835,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys=stream_keys, interrupt_after=interrupt_after, interrupt_before=interrupt_before, - check_subgraphs=check_subgraphs, manager=manager, debug=debug, ) @@ -959,7 +956,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - check_subgraphs: bool = True, debug: bool = False, ) -> None: super().__init__( @@ -974,7 +970,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys=stream_keys, interrupt_after=interrupt_after, interrupt_before=interrupt_before, - check_subgraphs=check_subgraphs, manager=manager, debug=debug, ) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 4e5be8470..e3701b529 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -163,7 +163,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): stream_keys=graph.stream_channels, interrupt_after=graph.interrupt_after_nodes, interrupt_before=graph.interrupt_before_nodes, - check_subgraphs=False, ) as loop: if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved @@ -353,7 +352,6 @@ class KafkaOrchestrator(AbstractContextManager): stream_keys=graph.stream_channels, interrupt_after=graph.interrupt_after_nodes, interrupt_before=graph.interrupt_before_nodes, - check_subgraphs=False, ) as loop: if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved From d402bf73799c33c91659d0d3ede58a79d9d80754 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 15:36:16 -0800 Subject: [PATCH 3/6] Remove flag --- libs/langgraph/langgraph/pregel/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 0a90aa10f..abf51f5be 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1869,7 +1869,6 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - check_subgraphs=self.checkpointer is not True, ) as loop: # create runner runner = PregelRunner( From f4bd023ab15b2b929f4e082e0b4f0b0a7f54987e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 16:04:14 -0800 Subject: [PATCH 4/6] Fix --- libs/langgraph/langgraph/pregel/__init__.py | 61 +++++----------- libs/langgraph/langgraph/pregel/retry.py | 10 ++- libs/langgraph/langgraph/utils/config.py | 16 ++++ .../langgraph/scheduler/kafka/executor.py | 73 ++++++++++--------- .../langgraph/scheduler/kafka/orchestrator.py | 24 +++--- 5 files changed, 90 insertions(+), 94 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index abf51f5be..e9ac31502 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -115,6 +115,7 @@ from langgraph.utils.config import ( patch_checkpoint_map, patch_config, patch_configurable, + recast_checkpoint_ns, ) from langgraph.utils.fields import get_enhanced_type_hints from langgraph.utils.pydantic import create_model @@ -694,19 +695,15 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - for _, pregel in self.get_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): return pregel.get_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), subgraphs=subgraphs, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs(self.config, config) if self.config else config saved = checkpointer.get_tuple(config) @@ -731,19 +728,15 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): return await pregel.aget_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), subgraphs=subgraphs, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs(self.config, config) if self.config else config saved = await checkpointer.aget_tuple(config) @@ -774,13 +767,9 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - for _, pregel in self.get_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): yield from pregel.get_state_history( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), filter=filter, @@ -789,7 +778,7 @@ class Pregel(PregelProtocol): ) return else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs( self.config, @@ -824,13 +813,9 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): async for state in pregel.aget_state_history( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), filter=filter, @@ -840,7 +825,7 @@ class Pregel(PregelProtocol): yield state return else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs( self.config, @@ -879,20 +864,16 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - for _, pregel in self.get_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): return pregel.update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), values, as_node, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") # get last checkpoint config = ensure_config(self.config, config) @@ -1163,20 +1144,16 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): return await pregel.aupdate_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), values, as_node, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") # get last checkpoint config = ensure_config(self.config, config) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 43e7e8d9e..6d0e43b54 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -48,7 +48,10 @@ def run_with_retry( break elif cmd.graph == Command.PARENT: # this command is for the parent graph, assign it to the parent - parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1]) + parts = ns.split(NS_SEP) + if parts[-1].isdigit(): + parts.pop() + parent_ns = NS_SEP.join(parts[:-1]) exc.args = (replace(cmd, graph=parent_ns),) # bubble up raise @@ -133,7 +136,10 @@ async def arun_with_retry( break elif cmd.graph == Command.PARENT: # this command is for the parent graph, assign it to the parent - parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1]) + parts = ns.split(NS_SEP) + if parts[-1].isdigit(): + parts.pop() + parent_ns = NS_SEP.join(parts[:-1]) exc.args = (replace(cmd, graph=parent_ns),) # bubble up raise diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index ac803cc35..309c6d6be 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -23,9 +23,25 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, + NS_END, + NS_SEP, ) +def recast_checkpoint_ns(ns: str) -> str: + """Remove task IDs from checkpoint namespace. + + Args: + ns (str): The checkpoint namespace with task IDs. + + Returns: + str: The checkpoint namespace without task IDs. + """ + return NS_SEP.join( + part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit() + ) + + def patch_configurable( config: Optional[RunnableConfig], patch: dict[str, Any] ) -> RunnableConfig: diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 970d55be8..fa9a221d0 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,5 +1,6 @@ import asyncio import concurrent.futures +from collections.abc import Sequence from contextlib import ( AbstractAsyncContextManager, AbstractContextManager, @@ -7,7 +8,7 @@ from contextlib import ( ExitStack, ) from functools import partial -from typing import Any, Optional, Sequence +from typing import Any, Optional from uuid import UUID import orjson @@ -15,7 +16,7 @@ from langchain_core.runnables import RunnableConfig from typing_extensions import Self import langgraph.scheduler.kafka.serde as serde -from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP +from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel from langgraph.pregel.algo import prepare_single_task @@ -39,7 +40,7 @@ from langgraph.scheduler.kafka.types import ( Topics, ) from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy -from langgraph.utils.config import patch_configurable +from langgraph.utils.config import patch_configurable, recast_checkpoint_ns class AsyncKafkaExecutor(AbstractAsyncContextManager): @@ -165,14 +166,12 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -183,16 +182,19 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): raise RuntimeError("Checkpoint not found") if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() - async with AsyncChannelsManager( - graph.channels, - saved.checkpoint, - LoopProtocol( - config=msg["config"], - store=self.graph.store, - step=saved.metadata["step"] + 1, - stop=saved.metadata["step"] + 2, - ), - ) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit: + async with ( + AsyncChannelsManager( + graph.channels, + saved.checkpoint, + LoopProtocol( + config=msg["config"], + store=self.graph.store, + step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, + ), + ) as (channels, managed), + AsyncBackgroundExecutor(msg["config"]) as submit, + ): if task := await asyncio.to_thread( prepare_single_task, msg["task"]["path"], @@ -378,14 +380,12 @@ class KafkaExecutor(AbstractContextManager): # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -396,16 +396,19 @@ class KafkaExecutor(AbstractContextManager): raise RuntimeError("Checkpoint not found") if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() - with ChannelsManager( - graph.channels, - saved.checkpoint, - LoopProtocol( - config=msg["config"], - store=self.graph.store, - step=saved.metadata["step"] + 1, - stop=saved.metadata["step"] + 2, - ), - ) as (channels, managed), BackgroundExecutor({}) as submit: + with ( + ChannelsManager( + graph.channels, + saved.checkpoint, + LoopProtocol( + config=msg["config"], + store=self.graph.store, + step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, + ), + ) as (channels, managed), + BackgroundExecutor({}) as submit, + ): if task := prepare_single_task( msg["task"]["path"], msg["task"]["id"], diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index e3701b529..5527ec964 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -16,8 +16,6 @@ from langgraph.constants import ( CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, INTERRUPT, - NS_END, - NS_SEP, SCHEDULED, ) from langgraph.errors import CheckpointNotLatest, GraphInterrupt @@ -37,7 +35,7 @@ from langgraph.scheduler.kafka.types import ( Topics, ) from langgraph.types import RetryPolicy -from langgraph.utils.config import patch_configurable +from langgraph.utils.config import patch_configurable, recast_checkpoint_ns class AsyncKafkaOrchestrator(AbstractAsyncContextManager): @@ -140,14 +138,12 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -329,14 +325,12 @@ class KafkaOrchestrator(AbstractContextManager): # 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) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message From c26b0e78b6766e59b2c616692a4e447b46fab870 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 17:15:58 -0800 Subject: [PATCH 5/6] Fix kafka lib --- libs/langgraph/langgraph/pregel/algo.py | 2 +- libs/langgraph/langgraph/pregel/loop.py | 2 +- .../langgraph/scheduler/kafka/orchestrator.py | 23 +++++++++++-------- libs/scheduler-kafka/tests/any.py | 8 +++++++ libs/scheduler-kafka/tests/test_subgraph.py | 8 ++++++- .../tests/test_subgraph_sync.py | 8 ++++++- 6 files changed, 37 insertions(+), 14 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 205793ab4..39bc9462d 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -680,7 +680,7 @@ def prepare_single_task( "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: - assert task_id == task_id_checksum + assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: if node := proc.node: if proc.metadata: diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 0afe62658..1f2fe91a3 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -229,7 +229,7 @@ class PregelLoop(LoopProtocol): if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) - if scratchpad is not None: + if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None: if scratchpad["subgraph_counter"]: self.config = patch_configurable( self.config, diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 5527ec964..57971e175 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -13,8 +13,10 @@ from typing_extensions import Self import langgraph.scheduler.kafka.serde as serde from langgraph.constants import ( + CONF, CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, + CONFIG_KEY_SCRATCHPAD, INTERRUPT, SCHEDULED, ) @@ -168,6 +170,16 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): if new_tasks := [ t for t in loop.tasks.values() if not t.scheduled and not t.writes ]: + config = patch_configurable( + loop.config, + { + **loop.checkpoint_config["configurable"], + CONFIG_KEY_DEDUPE_TASKS: True, + CONFIG_KEY_ENSURE_LATEST: True, + }, + ) + if CONFIG_KEY_SCRATCHPAD in config[CONF]: + config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = await asyncio.gather( *( @@ -175,16 +187,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): self.topics.executor, value=serde.dumps( MessageToExecutor( - config=patch_configurable( - loop.config, - { - **loop.checkpoint_config[ - "configurable" - ], - CONFIG_KEY_DEDUPE_TASKS: True, - CONFIG_KEY_ENSURE_LATEST: True, - }, - ), + config=config, task=ExecutorTask(id=task.id, path=task.path), finally_send=msg.get("finally_send"), ) diff --git a/libs/scheduler-kafka/tests/any.py b/libs/scheduler-kafka/tests/any.py index 3ea224173..0336d8506 100644 --- a/libs/scheduler-kafka/tests/any.py +++ b/libs/scheduler-kafka/tests/any.py @@ -53,3 +53,11 @@ class AnyList(list): return False else: return True + + +class AnyInt(int): + def __init__(self) -> None: + super().__init__() + + def __eq__(self, other: object) -> bool: + return isinstance(other, int) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 55303cba0..053eaedd0 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -15,7 +15,7 @@ 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 AnyDict +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics_async from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -198,6 +198,7 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -269,6 +270,7 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -370,6 +372,7 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -481,6 +484,7 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -547,6 +551,7 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -669,6 +674,7 @@ async def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index a67919dda..7d5de920c 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -15,7 +15,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.default_sync import DefaultProducer from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -197,6 +197,7 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -268,6 +269,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -369,6 +371,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -479,6 +482,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": True, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -545,6 +549,7 @@ def test_subgraph_w_interrupt( "__pregel_resuming": True, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, @@ -667,6 +672,7 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), "call_counter": 0, "interrupt_counter": -1, "null_resume": None, From 6b9369876fc199ee77e2055a6ba96dd2b29f5446 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 17:52:19 -0800 Subject: [PATCH 6/6] Fix sync --- .../langgraph/scheduler/kafka/orchestrator.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 57971e175..9ed72cd02 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -358,20 +358,23 @@ class KafkaOrchestrator(AbstractContextManager): if new_tasks := [ t for t in loop.tasks.values() if not t.scheduled and not t.writes ]: + config = patch_configurable( + loop.config, + { + **loop.checkpoint_config["configurable"], + CONFIG_KEY_DEDUPE_TASKS: True, + CONFIG_KEY_ENSURE_LATEST: True, + }, + ) + if CONFIG_KEY_SCRATCHPAD in config[CONF]: + config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = [ self.producer.send( self.topics.executor, value=serde.dumps( MessageToExecutor( - config=patch_configurable( - loop.config, - { - **loop.checkpoint_config["configurable"], - CONFIG_KEY_DEDUPE_TASKS: True, - CONFIG_KEY_ENSURE_LATEST: True, - }, - ), + config=config, task=ExecutorTask(id=task.id, path=task.path), finally_send=msg.get("finally_send"), )