From 8bcdba822e4146286590203e546b2e316342991a Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 18 Mar 2025 07:17:50 -0700 Subject: [PATCH] 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(