Reduce to 4.1s

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,

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,

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(