WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-11 12:17:53 +02:00
Code Issues Packages Projects Releases Wiki Activity

Use fast path for prepare_next_tasks on input (#3931)

Browse Source 
- 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
This commit is contained in:
Nuno Campos
2025-03-20 08:46:48 -07:00
committed by GitHub
parent 5fb037171d 0ac29434a7
commit 77c833e1e5
3 changed files with 29 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+15 -13
View File
@@ -504,6 +504,8 @@ class Pregel(PregelProtocol):
name: str = "LangGraph"
trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None
def __init__(
self,
*,
@@ -525,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
@@ -544,25 +547,28 @@ 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()
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))}
)
@@ -577,6 +583,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 +2283,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(
@@ -2751,10 +2753,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)
libs/langgraph/langgraph/pregel/algo.py
+5 -3
View File
@@ -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.
@@ -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()
libs/langgraph/langgraph/pregel/loop.py
+9 -6
View File
@@ -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,
@@ -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]]) -> Optional[set[str]]:
# 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: Optional[set[str]] = 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():
@@ -890,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,
@@ -1033,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,