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

Speed up prepare_single_task (#3893)

Browse Source 
- sequential(2000) goes from 8.4s to 4.1s
- replace UUID(str).bytes with faster binascii.unhexlify, and do it only
once per step
- find only the first active trigger, instead of the full list
- use a dedicated function for checking active trigger
This commit is contained in:
Nuno Campos
2025-03-18 10:15:17 -07:00
committed by GitHub
parent e7fbdeeb13 7a959f62cc
commit 4836f8b18b
8 changed files with 128 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/algo.py
+54 -17
View File
@@ -1,3 +1,4 @@
import binascii
import itertools
import sys
from collections import defaultdict, deque
@@ -18,7 +19,6 @@ from typing import (
cast,
overload,
)
from uuid import UUID
from langchain_core.callbacks import Callbacks
from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
@@ -27,6 +27,7 @@ from langchain_core.runnables.config import RunnableConfig
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import (
BaseCheckpointSaver,
ChannelVersions,
Checkpoint,
PendingWrite,
V,
@@ -373,6 +374,8 @@ def prepare_next_tasks(
"""Prepare the set of tasks that will make up the next Pregel step.
This is the union of all PUSH tasks (Sends) and PULL tasks (nodes triggered
by edges)."""
checkpoint_id_bytes = binascii.unhexlify(checkpoint["id"].replace("-", ""))
null_version = checkpoint_null_version(checkpoint)
tasks: list[Union[PregelTask, PregelExecutableTask]] = []
# Consume pending_sends from previous step
for idx, _ in enumerate(checkpoint["pending_sends"]):
@@ -380,6 +383,8 @@ def prepare_next_tasks(
(PUSH, idx),
None,
checkpoint=checkpoint,
checkpoint_id_bytes=checkpoint_id_bytes,
checkpoint_null_version=null_version,
pending_writes=pending_writes,
processes=processes,
channels=channels,
@@ -399,6 +404,8 @@ def prepare_next_tasks(
(PULL, name),
None,
checkpoint=checkpoint,
checkpoint_id_bytes=checkpoint_id_bytes,
checkpoint_null_version=null_version,
pending_writes=pending_writes,
processes=processes,
channels=channels,
@@ -414,11 +421,16 @@ def prepare_next_tasks(
return {t.id: t for t in tasks}
PUSH_TRIGGER = (PUSH,)
def prepare_single_task(
task_path: tuple[Any, ...],
task_id_checksum: Optional[str],
*,
checkpoint: Checkpoint,
checkpoint_id_bytes: bytes,
checkpoint_null_version: Optional[V],
pending_writes: list[PendingWrite],
processes: Mapping[str, PregelNode],
channels: Mapping[str, BaseChannel],
@@ -432,7 +444,6 @@ def prepare_single_task(
) -> Union[None, PregelTask, PregelExecutableTask]:
"""Prepares a single task for the next Pregel step, given a task path, which
uniquely identifies a PUSH or PULL task within the graph."""
checkpoint_id = UUID(checkpoint["id"]).bytes
configurable = config.get(CONF, {})
parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "")
@@ -445,10 +456,10 @@ def prepare_single_task(
if name is None:
raise ValueError("`call` functions must have a `__name__` attribute")
# create task id
triggers = [PUSH]
triggers: Sequence[str] = PUSH_TRIGGER
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = _uuid5_str(
checkpoint_id,
checkpoint_id_bytes,
checkpoint_ns,
str(step),
name,
@@ -539,12 +550,12 @@ def prepare_single_task(
)
return
# create task id
triggers = [PUSH]
triggers = PUSH_TRIGGER
checkpoint_ns = (
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = _uuid5_str(
checkpoint_id,
checkpoint_id_bytes,
checkpoint_ns,
str(step),
packet.node,
@@ -641,18 +652,15 @@ def prepare_single_task(
if name not in processes:
return
proc = processes[name]
version_type = type(next(iter(checkpoint["channel_versions"].values()), None))
null_version = version_type() # type: ignore[misc]
if null_version is None:
if checkpoint_null_version is None:
return
seen = checkpoint["versions_seen"].get(name, {})
# If any of the channels read by this process were updated
if triggers := sorted(
chan
for chan in proc.triggers
if channels[chan].is_available()
and checkpoint["channel_versions"].get(chan, null_version) # type: ignore[operator]
> seen.get(chan, null_version)
if triggers := _triggers(
channels,
checkpoint["channel_versions"],
checkpoint["versions_seen"].get(name),
checkpoint_null_version,
proc,
):
try:
val = next(
@@ -670,7 +678,7 @@ def prepare_single_task(
# create task id
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = _uuid5_str(
checkpoint_id,
checkpoint_id_bytes,
checkpoint_ns,
str(step),
name,
@@ -761,6 +769,35 @@ def prepare_single_task(
return PregelTask(task_id, name, task_path[:3])
def checkpoint_null_version(
checkpoint: Checkpoint,
) -> Optional[V]:
"""Get the null version for the checkpoint, if available."""
for version in checkpoint["channel_versions"].values():
return type(version)()
return None
def _triggers(
channels: Mapping[str, BaseChannel],
versions: ChannelVersions,
seen: Optional[ChannelVersions],
null_version: V,
proc: PregelNode,
) -> Sequence[str]:
if seen is None:
for chan in proc.triggers:
if channels[chan].is_available():
return (chan,)
else:
for chan in proc.triggers:
if channels[chan].is_available() and versions.get( # type: ignore[operator]
chan, null_version
) > seen.get(chan, null_version):
return (chan,)
return EMPTY_SEQ
def _scratchpad(
config: RunnableConfig,
pending_writes: list[PendingWrite],
libs/langgraph/langgraph/pregel/loop.py
+6
View File
@@ -1,4 +1,5 @@
import asyncio
import binascii
import concurrent.futures
from collections import defaultdict, deque
from contextlib import AsyncExitStack, ExitStack
@@ -79,6 +80,7 @@ from langgraph.pregel.algo import (
GetNextVersion,
PregelTaskWrites,
apply_writes,
checkpoint_null_version,
increment,
prepare_next_tasks,
prepare_single_task,
@@ -347,12 +349,16 @@ class PregelLoop(LoopProtocol):
):
self.to_interrupt.append(task)
return
checkpoint_id_bytes = binascii.unhexlify(self.checkpoint["id"].replace("-", ""))
null_version = checkpoint_null_version(self.checkpoint)
if pushed := cast(
Optional[PregelExecutableTask],
prepare_single_task(
(PUSH, task.path, write_idx, task.id, call),
None,
checkpoint=self.checkpoint,
checkpoint_id_bytes=checkpoint_id_bytes,
checkpoint_null_version=null_version,
pending_writes=self.checkpoint_pending_writes,
processes=self.nodes,
channels=self.channels,
libs/langgraph/langgraph/types.py
+1 -1
View File
@@ -149,7 +149,7 @@ class PregelExecutableTask(NamedTuple):
proc: Runnable
writes: deque[tuple[str, Any]]
config: RunnableConfig
triggers: list[str]
triggers: Sequence[str]
retry_policy: Optional[RetryPolicy]
cache_policy: Optional[CachePolicy]
id: str
libs/langgraph/tests/test_large_cases.py
+13 -13
View File
@@ -2483,7 +2483,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
{
"langgraph_step": 1,
"langgraph_node": "agent",
"langgraph_triggers": ["start:agent"],
"langgraph_triggers": ("start:agent",),
"langgraph_path": (PULL, "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
@@ -2500,7 +2500,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
{
"langgraph_step": 2,
"langgraph_node": "tools",
"langgraph_triggers": ["branch:to:tools"],
"langgraph_triggers": ("branch:to:tools",),
"langgraph_path": (PULL, "tools"),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
@@ -2542,7 +2542,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
{
"langgraph_step": 3,
"langgraph_node": "agent",
"langgraph_triggers": ["tools"],
"langgraph_triggers": ("tools",),
"langgraph_path": (PULL, "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
@@ -2559,7 +2559,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
{
"langgraph_step": 4,
"langgraph_node": "tools",
"langgraph_triggers": ["branch:to:tools"],
"langgraph_triggers": ("branch:to:tools",),
"langgraph_path": (PULL, "tools"),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
@@ -2573,7 +2573,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
{
"langgraph_step": 4,
"langgraph_node": "tools",
"langgraph_triggers": ["branch:to:tools"],
"langgraph_triggers": ("branch:to:tools",),
"langgraph_path": (PULL, "tools"),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
@@ -2585,7 +2585,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
{
"langgraph_step": 5,
"langgraph_node": "agent",
"langgraph_triggers": ["tools"],
"langgraph_triggers": ("tools",),
"langgraph_path": (PULL, "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
@@ -5501,7 +5501,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": AnyStr(),
"name": "rewrite_query",
"input": {"query": "what is weather in sf", "docs": []},
"triggers": ["start:rewrite_query"],
"triggers": ("start:rewrite_query",),
},
},
),
@@ -5532,7 +5532,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": AnyStr(),
"name": "retriever_one",
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
"triggers": ("rewrite_query",),
},
},
),
@@ -5546,7 +5546,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"id": AnyStr(),
"name": "retriever_two",
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
"triggers": ("rewrite_query",),
},
},
),
@@ -5608,7 +5608,7 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
"triggers": ["retriever_one", "retriever_two"],
"triggers": (AnyStr("retriever_"),),
},
},
),
@@ -6634,7 +6634,7 @@ def test_branch_then(
"id": AnyStr(),
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
"triggers": ("start:prepare",),
},
},
{
@@ -6706,7 +6706,7 @@ def test_branch_then(
"id": AnyStr(),
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:to:tool_two_slow"],
"triggers": ("branch:to:tool_two_slow",),
},
},
{
@@ -6773,7 +6773,7 @@ def test_branch_then(
"id": AnyStr(),
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition::then"],
"triggers": ("branch:prepare:condition::then",),
},
},
{
libs/langgraph/tests/test_large_cases_async.py
+14 -14
View File
@@ -2300,7 +2300,7 @@ async def test_prebuilt_tool_chat() -> None:
{
"langgraph_step": 1,
"langgraph_node": "agent",
"langgraph_triggers": ["start:agent"],
"langgraph_triggers": ("start:agent",),
"langgraph_path": ("__pregel_pull", "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
@@ -2317,7 +2317,7 @@ async def test_prebuilt_tool_chat() -> None:
{
"langgraph_step": 2,
"langgraph_node": "tools",
"langgraph_triggers": ["branch:to:tools"],
"langgraph_triggers": ("branch:to:tools",),
"langgraph_path": ("__pregel_pull", "tools"),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
@@ -2359,7 +2359,7 @@ async def test_prebuilt_tool_chat() -> None:
{
"langgraph_step": 3,
"langgraph_node": "agent",
"langgraph_triggers": ["tools"],
"langgraph_triggers": ("tools",),
"langgraph_path": ("__pregel_pull", "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
@@ -2376,7 +2376,7 @@ async def test_prebuilt_tool_chat() -> None:
{
"langgraph_step": 4,
"langgraph_node": "tools",
"langgraph_triggers": ["branch:to:tools"],
"langgraph_triggers": ("branch:to:tools",),
"langgraph_path": ("__pregel_pull", "tools"),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
@@ -2390,7 +2390,7 @@ async def test_prebuilt_tool_chat() -> None:
{
"langgraph_step": 4,
"langgraph_node": "tools",
"langgraph_triggers": ["branch:to:tools"],
"langgraph_triggers": ("branch:to:tools",),
"langgraph_path": ("__pregel_pull", "tools"),
"langgraph_checkpoint_ns": AnyStr("tools:"),
},
@@ -2402,7 +2402,7 @@ async def test_prebuilt_tool_chat() -> None:
{
"langgraph_step": 5,
"langgraph_node": "agent",
"langgraph_triggers": ["tools"],
"langgraph_triggers": ("tools",),
"langgraph_path": ("__pregel_pull", "agent"),
"langgraph_checkpoint_ns": AnyStr("agent:"),
"checkpoint_ns": AnyStr("agent:"),
@@ -3883,7 +3883,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": AnyStr(),
"name": "rewrite_query",
"input": {"query": "what is weather in sf", "docs": []},
"triggers": ["start:rewrite_query"],
"triggers": ("start:rewrite_query",),
},
},
),
@@ -3914,7 +3914,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": AnyStr(),
"name": "retriever_one",
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
"triggers": ("rewrite_query",),
},
},
),
@@ -3928,7 +3928,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"id": AnyStr(),
"name": "retriever_two",
"input": {"query": "query: what is weather in sf", "docs": []},
"triggers": ["rewrite_query"],
"triggers": ("rewrite_query",),
},
},
),
@@ -3990,7 +3990,7 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
"triggers": ["retriever_one", "retriever_two"],
"triggers": (AnyStr("retriever_"),),
},
},
),
@@ -4465,7 +4465,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
"id": AnyStr(),
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
"triggers": ("start:prepare",),
},
},
{
@@ -4537,7 +4537,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
"id": AnyStr(),
"name": "tool_two_slow",
"input": {"my_key": "value prepared", "market": "DE"},
"triggers": ["branch:to:tool_two_slow"],
"triggers": ("branch:to:tool_two_slow",),
},
},
{
@@ -4609,7 +4609,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
"id": AnyStr(),
"name": "finish",
"input": {"my_key": "value prepared slow", "market": "DE"},
"triggers": ["branch:prepare:condition::then"],
"triggers": ("branch:prepare:condition::then",),
},
},
{
@@ -4778,7 +4778,7 @@ async def test_branch_then(checkpointer_name: str) -> None:
"id": AnyStr(),
"name": "prepare",
"input": {"my_key": "value", "market": "DE"},
"triggers": ["start:prepare"],
"triggers": ("start:prepare",),
},
},
{
libs/langgraph/tests/test_pregel.py
+26 -24
View File
@@ -817,7 +817,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": AnyStr(),
"name": "one",
"input": 2,
"triggers": ["input"],
"triggers": ("input",),
},
},
{
@@ -828,7 +828,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": AnyStr(),
"name": "two",
"input": [12],
"triggers": ["inbox"],
"triggers": ("inbox",),
},
},
{
@@ -863,7 +863,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": AnyStr(),
"name": "two",
"input": [3],
"triggers": ["inbox"],
"triggers": ("inbox",),
},
},
{
@@ -3247,14 +3247,24 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
assert [
c for c in app_w_interrupt.stream({"query": "what is weather in sf"}, config)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"qa": {"answer": ""}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"__interrupt__": ()},
]
] in (
[
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"qa": {"answer": ""}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"__interrupt__": ()},
],
[
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"qa": {"answer": ""}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
{"retriever_one": {"docs": ["doc1", "doc2"]}},
{"__interrupt__": ()},
],
)
assert [c for c in app_w_interrupt.stream(None, config)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
@@ -5969,9 +5979,7 @@ def test_falsy_return_from_task(
"a": 5,
},
"name": "graph",
"triggers": [
"__start__",
],
"triggers": ("__start__",),
},
"step": 0,
"timestamp": AnyStr(),
@@ -5985,9 +5993,7 @@ def test_falsy_return_from_task(
{},
),
"name": "falsy_task",
"triggers": [
"__pregel_push",
],
"triggers": ("__pregel_push",),
},
"step": 0,
"timestamp": AnyStr(),
@@ -6094,9 +6100,7 @@ def test_falsy_return_from_task(
"a": 5,
},
"name": "graph",
"triggers": [
"__start__",
],
"triggers": ("__start__",),
},
"step": 0,
"timestamp": AnyStr(),
@@ -6110,9 +6114,7 @@ def test_falsy_return_from_task(
{},
),
"name": "falsy_task",
"triggers": [
"__pregel_push",
],
"triggers": ("__pregel_push",),
},
"step": 0,
"timestamp": AnyStr(),
@@ -6923,7 +6925,7 @@ def test_tags_stream_mode_messages() -> None:
{
"langgraph_step": 1,
"langgraph_node": "call_model",
"langgraph_triggers": ["start:call_model"],
"langgraph_triggers": ("start:call_model",),
"langgraph_path": ("__pregel_pull", "call_model"),
"langgraph_checkpoint_ns": AnyStr("call_model:"),
"checkpoint_ns": AnyStr("call_model:"),
libs/langgraph/tests/test_pregel_async.py
+4 -4
View File
@@ -1672,7 +1672,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": AnyStr(),
"name": "one",
"input": 2,
"triggers": ["input"],
"triggers": ("input",),
},
},
{
@@ -1683,7 +1683,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": AnyStr(),
"name": "two",
"input": [12],
"triggers": ["inbox"],
"triggers": ("inbox",),
},
},
{
@@ -1718,7 +1718,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
"id": AnyStr(),
"name": "two",
"input": [3],
"triggers": ["inbox"],
"triggers": ("inbox",),
},
},
{
@@ -7571,7 +7571,7 @@ async def test_tags_stream_mode_messages() -> None:
{
"langgraph_step": 1,
"langgraph_node": "call_model",
"langgraph_triggers": ["start:call_model"],
"langgraph_triggers": ("start:call_model",),
"langgraph_path": ("__pregel_pull", "call_model"),
"langgraph_checkpoint_ns": AnyStr("call_model:"),
"checkpoint_ns": AnyStr("call_model:"),
libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py
+10 -1
View File
@@ -1,4 +1,5 @@
import asyncio
import binascii
import concurrent.futures
from collections.abc import Sequence
from contextlib import (
@@ -19,7 +20,7 @@ import langgraph.scheduler.kafka.serde as serde
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
from langgraph.pregel.algo import checkpoint_null_version, prepare_single_task
from langgraph.pregel.executor import (
AsyncBackgroundExecutor,
BackgroundExecutor,
@@ -209,6 +210,10 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager):
for_execution=True,
checkpointer=self.graph.checkpointer,
store=self.graph.store,
checkpoint_id_bytes=binascii.unhexlify(
saved.checkpoint["id"].replace("-", "")
),
checkpoint_null_version=checkpoint_null_version(saved.checkpoint),
):
# execute task, saving writes
runner = PregelRunner(
@@ -421,6 +426,10 @@ class KafkaExecutor(AbstractContextManager):
step=saved.metadata["step"] + 1,
for_execution=True,
checkpointer=self.graph.checkpointer,
checkpoint_id_bytes=binascii.unhexlify(
saved.checkpoint["id"].replace("-", "")
),
checkpoint_null_version=checkpoint_null_version(saved.checkpoint),
):
# execute task, saving writes
runner = PregelRunner(