From eaa18cc2dd5d9e0785648eedaf4baaaf7b625275 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 19 Mar 2025 18:04:19 -0700 Subject: [PATCH 1/6] Use fast path for prepare_next_tasks on input - When there are no values in checkpoint no need to run through all the PULL candidates - When there are input writes save updated_channels to use on the next call to prepare_next_tasks --- libs/langgraph/langgraph/pregel/__init__.py | 10 ++++------ libs/langgraph/langgraph/pregel/algo.py | 2 ++ libs/langgraph/langgraph/pregel/loop.py | 9 ++++++--- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 4dd69806d..79470680b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -504,6 +504,8 @@ class Pregel(PregelProtocol): name: str = "LangGraph" + trigger_to_nodes: dict[str, set[str]] = None + def __init__( self, *, @@ -577,6 +579,7 @@ class Pregel(PregelProtocol): self.interrupt_after_nodes, self.interrupt_before_nodes, ) + self.trigger_to_nodes = _trigger_to_nodes(self.nodes) return self @property @@ -2276,12 +2279,7 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - # `self.nodes` can be modified after creation of `Pregel`. For example, - # that's how StateGraph compilation currently works. - # For now, we recompute the trigger_to_nodes mapping every time the - # loop is created. We could potentially memoize this if it becomes a - # performance issue. - trigger_to_nodes=_trigger_to_nodes(self.nodes), + trigger_to_nodes=self.trigger_to_nodes, ) as loop: # create runner runner = PregelRunner( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c37c95373..4546ab300 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -452,6 +452,8 @@ def prepare_next_tasks( triggered_nodes.update(node_ids) # Sort the nodes to ensure deterministic order candidate_nodes: Iterable[str] = sorted(triggered_nodes) + elif not checkpoint["channel_versions"]: + candidate_nodes = () else: candidate_nodes = processes.keys() diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 1dd604f02..a55941a30 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -412,7 +412,7 @@ class PregelLoop(LoopProtocol): updated_channels: set[str] | None = None if self.input not in (INPUT_DONE, INPUT_RESUMING, INPUT_SHOULD_VALIDATE): - self._first(input_keys=input_keys) + updated_channels = self._first(input_keys=input_keys) elif self.to_interrupt: # if we need to interrupt, do so self.status = "interrupt_before" @@ -582,7 +582,7 @@ class PregelLoop(LoopProtocol): else: task.writes.append((k, v)) - def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None: + def _first(self, *, input_keys: Union[str, Sequence[str]]) -> set[str] | None: # resuming from previous checkpoint requires # - finding a previous checkpoint # - receiving None input (outer graph) or RESUMING flag (subgraph) @@ -599,6 +599,8 @@ class PregelLoop(LoopProtocol): ), ) ) + # this can be set only when there are input_writes + updated_channels: set[str] | None = None # map command to writes if isinstance(self.input, Command): @@ -668,7 +670,7 @@ class PregelLoop(LoopProtocol): manager=None, ) # apply input writes - mv_writes, _ = apply_writes( + mv_writes, updated_channels = apply_writes( self.checkpoint, self.channels, [ @@ -698,6 +700,7 @@ class PregelLoop(LoopProtocol): self.config = patch_configurable( self.config, {CONFIG_KEY_RESUMING: is_resuming} ) + return updated_channels def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: for k, v in self.config["metadata"].items(): From d3bb2b9aa0424e603a7c2b2e06fb637beeb85393 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 20 Mar 2025 08:18:17 -0700 Subject: [PATCH 2/6] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 8 ++++---- libs/langgraph/langgraph/pregel/algo.py | 6 +++--- libs/langgraph/langgraph/pregel/loop.py | 6 +++--- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 79470680b..ac1f2781d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -504,7 +504,7 @@ class Pregel(PregelProtocol): name: str = "LangGraph" - trigger_to_nodes: dict[str, set[str]] = None + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None def __init__( self, @@ -2749,10 +2749,10 @@ class Pregel(PregelProtocol): return chunks -def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, list[str]]: +def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> Mapping[str, Sequence[str]]: """Index from a trigger to nodes that depend on it.""" trigger_to_nodes: defaultdict[str, list[str]] = defaultdict(list) for name, node in nodes.items(): for trigger in node.triggers: - trigger_to_nodes.setdefault(trigger, []).append(name) - return cast(Mapping[str, list[str]], trigger_to_nodes) + trigger_to_nodes[trigger].append(name) + return dict(trigger_to_nodes) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 4546ab300..01d724c9f 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -347,7 +347,7 @@ def prepare_next_tasks( store: Literal[None] = None, checkpointer: Literal[None] = None, manager: Literal[None] = None, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelTask]: ... @@ -366,7 +366,7 @@ def prepare_next_tasks( store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], manager: Union[None, ParentRunManager, AsyncParentRunManager], - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelExecutableTask]: ... @@ -384,7 +384,7 @@ def prepare_next_tasks( store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, updated_channels: Optional[set[str]] = None, ) -> Union[dict[str, PregelTask], dict[str, PregelExecutableTask]]: """Prepare the set of tasks that will make up the next Pregel step. diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index a55941a30..f21d61442 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -210,7 +210,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( step=0, @@ -893,7 +893,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, @@ -1036,7 +1036,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - trigger_to_nodes: Optional[Mapping[str, list[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, From b82d70a66a50134461b0337e98aaa7e14b6c90f1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 20 Mar 2025 08:22:16 -0700 Subject: [PATCH 3/6] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index ac1f2781d..e8ddb8489 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -527,6 +527,7 @@ class Pregel(PregelProtocol): config_type: Optional[Type[Any]] = None, input_model: Optional[Type[BaseModel]] = None, config: Optional[RunnableConfig] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, name: str = "LangGraph", ) -> None: self.nodes = nodes @@ -546,6 +547,7 @@ class Pregel(PregelProtocol): self.config_type = config_type self.input_model = input_model self.config = config + self.trigger_to_nodes = trigger_to_nodes self.name = name if auto_validate: self.validate() From 7d0857f263851e81bfcd1cb82423af2cd40955ff Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 20 Mar 2025 08:24:31 -0700 Subject: [PATCH 4/6] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 10 ++++++---- libs/langgraph/langgraph/pregel/loop.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index e8ddb8489..45efe5c17 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -553,20 +553,22 @@ class Pregel(PregelProtocol): self.validate() def get_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False + self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False ) -> Graph: raise NotImplementedError async def aget_graph( - self, config: RunnableConfig | None = None, *, xray: int | bool = False + self, config: Optional[RunnableConfig] = None, *, xray: Union[int, bool] = False ) -> Graph: raise NotImplementedError - def copy(self, update: dict[str, Any] | None = None) -> Self: + def copy(self, update: Optional[dict[str, Any]] = None) -> Self: attrs = {**self.__dict__, **(update or {})} return self.__class__(**attrs) - def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: + def with_config( + self, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Self: return self.copy( {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))} ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index f21d61442..ff196a778 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -582,7 +582,7 @@ class PregelLoop(LoopProtocol): else: task.writes.append((k, v)) - def _first(self, *, input_keys: Union[str, Sequence[str]]) -> set[str] | None: + def _first(self, *, input_keys: Union[str, Sequence[str]]) -> Union[set[str]]: # resuming from previous checkpoint requires # - finding a previous checkpoint # - receiving None input (outer graph) or RESUMING flag (subgraph) From 43f5a174163491d123e5af6ce504d0f207fe6fd7 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 20 Mar 2025 08:24:51 -0700 Subject: [PATCH 5/6] Lint --- libs/langgraph/langgraph/pregel/loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index ff196a778..8f3331349 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -600,7 +600,7 @@ class PregelLoop(LoopProtocol): ) ) # this can be set only when there are input_writes - updated_channels: set[str] | None = None + updated_channels: Union[set[str]] = None # map command to writes if isinstance(self.input, Command): From 0ac29434a7e28351e7164a22988dd0dfa2f273ec Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 20 Mar 2025 08:40:05 -0700 Subject: [PATCH 6/6] Lint --- libs/langgraph/langgraph/pregel/loop.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 8f3331349..94faff87f 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -582,7 +582,7 @@ class PregelLoop(LoopProtocol): else: task.writes.append((k, v)) - def _first(self, *, input_keys: Union[str, Sequence[str]]) -> Union[set[str]]: + def _first(self, *, input_keys: Union[str, Sequence[str]]) -> Optional[set[str]]: # resuming from previous checkpoint requires # - finding a previous checkpoint # - receiving None input (outer graph) or RESUMING flag (subgraph) @@ -600,7 +600,7 @@ class PregelLoop(LoopProtocol): ) ) # this can be set only when there are input_writes - updated_channels: Union[set[str]] = None + updated_channels: Optional[set[str]] = None # map command to writes if isinstance(self.input, Command):