mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-11 04:07:52 +02:00
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
This commit is contained in:
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user