From 18b82cb8e295ab3f05afd13ee18b03c0bafa937d Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 19 Mar 2025 15:29:50 -0400 Subject: [PATCH] x --- libs/langgraph/langgraph/pregel/__init__.py | 30 +++++++++++++------ libs/langgraph/langgraph/pregel/algo.py | 32 ++++++++++++++++----- libs/langgraph/langgraph/pregel/loop.py | 26 ++++++----------- 3 files changed, 56 insertions(+), 32 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d93f8ab11..77f682fa3 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -543,12 +543,6 @@ class Pregel(PregelProtocol): self.config_type = config_type self.input_model = input_model self.config = config - # Index from a trigger to nodes that depend on it - trigger_to_node = {} - for name, node in self.nodes.items(): - for trigger in node.triggers: - trigger_to_node.setdefault(trigger, []).append(name) - self.trigger_to_node = trigger_to_node self.name = name if auto_validate: @@ -2006,7 +2000,12 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - triggers_to_nodes=self.trigger_to_node, + # `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), ) as loop: # create runner runner = PregelRunner( @@ -2285,6 +2284,7 @@ class Pregel(PregelProtocol): stream.put_nowait, ((), "custom", c) ) ) + async with AsyncPregelLoop( input, input_model=self.input_model, @@ -2300,7 +2300,12 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - triggers_to_nodes=self.trigger_to_node, + # `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), ) as loop: # create runner runner = PregelRunner( @@ -2469,3 +2474,12 @@ class Pregel(PregelProtocol): return latest else: return chunks + + +def _trigger_to_nodes(nodes: dict[str, PregelNode]) -> dict[str, list[str]]: + """Index from a trigger to nodes that depend on it.""" + trigger_to_nodes = {} + for name, node in nodes.items(): + for trigger in node.triggers: + trigger_to_nodes.setdefault(trigger, []).append(name) + return trigger_to_nodes diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 9bc0f59b4..01d8c1ff8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -329,7 +329,6 @@ def apply_writes( max_version, channels[chan], ) - # Return managed values writes to be applied externally return pending_writes_by_managed, updated_channels @@ -348,6 +347,8 @@ def prepare_next_tasks( store: Literal[None] = None, checkpointer: Literal[None] = None, manager: Literal[None] = None, + trigger_to_nodes: Optional[Mapping[str, set[str]]] = None, + updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelTask]: ... @@ -365,6 +366,8 @@ def prepare_next_tasks( store: Optional[BaseStore], checkpointer: Optional[BaseCheckpointSaver], manager: Union[None, ParentRunManager, AsyncParentRunManager], + trigger_to_nodes: Optional[Mapping[str, set[str]]] = None, + updated_channels: Optional[set[str]] = None, ) -> dict[str, PregelExecutableTask]: ... @@ -381,8 +384,8 @@ def prepare_next_tasks( store: Optional[BaseStore] = None, checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, - # Nodes that are known to have been triggered in the previous step - triggered_nodes: Optional[set[str]] = None, + trigger_to_nodes: Optional[Mapping[str, set[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. @@ -398,7 +401,12 @@ def prepare_next_tasks( store: An instance of BaseStore to make it available for usage within tasks. checkpointer: Checkpointer instance used for saving checkpoints. manager: The parent run manager to use for the tasks. - triggered_nodes: The set of nodes that were triggered in the previous step + trigger_to_nodes: Optional: Mapping of channel names to the set of nodes + that are can be triggered by that channel. + updated_channels: Optional. Set of channel names that have been updated during + the previous step. Using in conjunction with trigger_to_nodes to speed + up the process of determining which nodes should be triggered in the next + step. Returns: A dictionary of tasks to be executed. The keys are the task ids and the values @@ -428,11 +436,21 @@ def prepare_next_tasks( manager=manager, ): tasks.append(task) + # If updated channels is available, we project only a subset of the nodes. + # since only those nodes have been updated? + if updated_channels and trigger_to_nodes: + triggered_nodes: set[str] = set() + # Get all nodes that have triggers associated with an updated channel + for channel in updated_channels: + if node_ids := trigger_to_nodes.get(channel): + triggered_nodes.update(node_ids) + # Sort the nodes to ensure deterministic order + candidate_nodes: Iterable[str] = sorted(triggered_nodes) + else: + candidate_nodes = processes.keys() + # Check if any processes should be run in next step # If so, prepare the values to be passed to them - - candidate_nodes: Iterable[str] = processes.keys() if triggered_nodes is None else sorted(triggered_nodes) - for name in candidate_nodes: # Check if we know which channels have been updated. if task := prepare_single_task( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index ef1ae00d9..518372834 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -209,7 +209,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - triggers_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( step=0, @@ -233,6 +233,7 @@ class PregelLoop(LoopProtocol): CONFIG_KEY_CHECKPOINT_ID not in config[CONF] or CONFIG_KEY_DEDUPE_TASKS in config[CONF] ) + self.trigger_to_nodes = trigger_to_nodes 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]) @@ -279,7 +280,6 @@ class PregelLoop(LoopProtocol): if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) else () ) - self.triggers_to_nodes = triggers_to_nodes self.prev_checkpoint_config = None def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: @@ -431,7 +431,7 @@ class PregelLoop(LoopProtocol): ), ) # all tasks have finished - mv_writes, _ = apply_writes( + mv_writes, updated_channels = apply_writes( self.checkpoint, self.channels, self.tasks.values(), @@ -484,15 +484,6 @@ class PregelLoop(LoopProtocol): self.status = "out_of_steps" return False - # If updated channels is available, we project only a subset of the nodes. - # since only those nodes have been updated? - triggered_nodes: set[str] = set() - if updated_channels is not None: - # Get all nodes that have triggers associated with an updated channel - for channel in updated_channels: - if node_ids := self.triggers_to_nodes.get(channel): - triggered_nodes.update(node_ids) - # prepare next tasks self.tasks = prepare_next_tasks( self.checkpoint, @@ -506,7 +497,8 @@ class PregelLoop(LoopProtocol): manager=self.manager, store=self.store, checkpointer=self.checkpointer, - triggered_nodes=triggered_nodes or None, + trigger_to_nodes=self.trigger_to_nodes, + updated_channels=updated_channels, ) self.to_interrupt = [] @@ -897,7 +889,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - triggers_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, @@ -914,7 +906,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): interrupt_before=interrupt_before, manager=manager, debug=debug, - triggers_to_nodes=triggers_to_nodes, + trigger_to_nodes=trigger_to_nodes, ) self.stack = ExitStack() if checkpointer: @@ -1040,7 +1032,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, - triggers_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( input, @@ -1057,7 +1049,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): interrupt_before=interrupt_before, manager=manager, debug=debug, - triggers_to_nodes=triggers_to_nodes, + trigger_to_nodes=trigger_to_nodes, ) self.stack = AsyncExitStack() if checkpointer: