From 60fc49b448c25847e3119cd0e3f9a549a01923af Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Mar 2025 21:56:58 -0700 Subject: [PATCH 1/7] Speed up prepare_single_task - sequential(2000) goes from 8.4s to 4.7s - replace UUID(str).bytes with simpler str.encode() - find only the first active trigger, instead of the full list - use a dedicated function for checking active trigger --- libs/langgraph/langgraph/pregel/algo.py | 40 ++++++++++++++++++------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index ba028d7bd..614f61dc8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -18,7 +18,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 +26,7 @@ from langchain_core.runnables.config import RunnableConfig from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( BaseCheckpointSaver, + ChannelVersions, Checkpoint, PendingWrite, V, @@ -432,7 +432,7 @@ 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 + checkpoint_id = checkpoint["id"].encode() configurable = config.get(CONF, {}) parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -641,18 +641,18 @@ def prepare_single_task( if name not in processes: return proc = processes[name] - version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) + versions = checkpoint["channel_versions"] + version_type = type(next(iter(versions.values()), None)) null_version = version_type() # type: ignore[misc] if 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, + versions, + checkpoint["versions_seen"].get(name), + null_version, + proc, ): try: val = next( @@ -761,6 +761,26 @@ def prepare_single_task( return PregelTask(task_id, name, task_path[:3]) +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( + chan, null_version + ) > seen.get(chan, null_version): # type: ignore[operator] + return (chan,) + return EMPTY_SEQ + + def _scratchpad( config: RunnableConfig, pending_writes: list[PendingWrite], From 8bcdba822e4146286590203e546b2e316342991a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 07:17:50 -0700 Subject: [PATCH 2/7] Reduce to 4.1s --- libs/langgraph/langgraph/pregel/algo.py | 34 +++++++++++++------ libs/langgraph/langgraph/pregel/loop.py | 6 ++++ .../langgraph/scheduler/kafka/executor.py | 7 +++- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 614f61dc8..1044cd1c4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -1,3 +1,4 @@ +import binascii import itertools import sys from collections import defaultdict, deque @@ -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, @@ -419,6 +426,8 @@ def prepare_single_task( 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 +441,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 = checkpoint["id"].encode() configurable = config.get(CONF, {}) parent_ns = configurable.get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -448,7 +456,7 @@ def prepare_single_task( triggers = [PUSH] 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, @@ -544,7 +552,7 @@ def prepare_single_task( 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,17 +649,14 @@ def prepare_single_task( if name not in processes: return proc = processes[name] - versions = checkpoint["channel_versions"] - version_type = type(next(iter(versions.values()), None)) - null_version = version_type() # type: ignore[misc] - if null_version is None: + if checkpoint_null_version is None: return # If any of the channels read by this process were updated if triggers := _triggers( channels, - versions, + checkpoint["channel_versions"], checkpoint["versions_seen"].get(name), - null_version, + checkpoint_null_version, proc, ): try: @@ -670,7 +675,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 +766,15 @@ 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, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 7ab3431b7..edd69db01 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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, diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index fa9a221d0..b8ab27674 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -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, @@ -421,6 +422,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( From 951131c8ecc40277165038459ebd07e3c64ab14d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 08:42:12 -0700 Subject: [PATCH 3/7] Lint --- libs/langgraph/langgraph/pregel/algo.py | 11 +++++++---- libs/langgraph/langgraph/types.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 1044cd1c4..027f51ee7 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -421,6 +421,9 @@ 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], @@ -453,7 +456,7 @@ 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_bytes, @@ -547,7 +550,7 @@ 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 ) @@ -788,9 +791,9 @@ def _triggers( return (chan,) else: for chan in proc.triggers: - if channels[chan].is_available() and versions.get( + if channels[chan].is_available() and versions.get( # type: ignore[operator] chan, null_version - ) > seen.get(chan, null_version): # type: ignore[operator] + ) > seen.get(chan, null_version): return (chan,) return EMPTY_SEQ diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 5efb89385..4c8cf6da4 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -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 From 98b8ff904ceaa45e629f290d629af250ba49ba6c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 09:30:29 -0700 Subject: [PATCH 4/7] Update test assertions for triggers --- libs/langgraph/tests/test_large_cases.py | 26 ++++++++--------- .../langgraph/tests/test_large_cases_async.py | 28 +++++++++---------- libs/langgraph/tests/test_pregel.py | 24 ++++++---------- libs/langgraph/tests/test_pregel_async.py | 8 +++--- 4 files changed, 39 insertions(+), 47 deletions(-) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 57c9b4071..a9d54b98d 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -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": ("retriever_one",), }, }, ), @@ -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",), }, }, { diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 95fcb9646..f7f11da0c 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -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": ("retriever_one",), }, }, ), @@ -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",), }, }, { diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ccd899ce9..f4af98b52 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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",), }, }, { @@ -5969,9 +5969,7 @@ def test_falsy_return_from_task( "a": 5, }, "name": "graph", - "triggers": [ - "__start__", - ], + "triggers": ("__start__",), }, "step": 0, "timestamp": AnyStr(), @@ -5985,9 +5983,7 @@ def test_falsy_return_from_task( {}, ), "name": "falsy_task", - "triggers": [ - "__pregel_push", - ], + "triggers": ("__pregel_push",), }, "step": 0, "timestamp": AnyStr(), @@ -6094,9 +6090,7 @@ def test_falsy_return_from_task( "a": 5, }, "name": "graph", - "triggers": [ - "__start__", - ], + "triggers": ("__start__",), }, "step": 0, "timestamp": AnyStr(), @@ -6110,9 +6104,7 @@ def test_falsy_return_from_task( {}, ), "name": "falsy_task", - "triggers": [ - "__pregel_push", - ], + "triggers": ("__pregel_push",), }, "step": 0, "timestamp": AnyStr(), @@ -6923,7 +6915,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:"), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a42a84a28..f1bb4d8ab 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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:"), From fa96c0ac761350fbd8c0da5f508732bfe06e24d0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 09:34:52 -0700 Subject: [PATCH 5/7] One more --- libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index b8ab27674..2d442e131 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -210,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( From 9b5549f759aa3c4d10bc9211741d91e329e4d078 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 09:43:56 -0700 Subject: [PATCH 6/7] Fix flaky assertion --- libs/langgraph/tests/test_pregel.py | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f4af98b52..d8d51a8a0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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"}}, From 7a959f62ccaf5fd978623ed6f5e20baaaf42aad9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 09:54:20 -0700 Subject: [PATCH 7/7] Fix assertion --- libs/langgraph/tests/test_large_cases.py | 2 +- libs/langgraph/tests/test_large_cases_async.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index a9d54b98d..d6637f000 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -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",), + "triggers": (AnyStr("retriever_"),), }, }, ), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index f7f11da0c..157e0e080 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -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",), + "triggers": (AnyStr("retriever_"),), }, }, ),