From f4bd023ab15b2b929f4e082e0b4f0b0a7f54987e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 15 Jan 2025 16:04:14 -0800 Subject: [PATCH] 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