perf(langgraph): cache max seen channel version

.changeset/every-windows-lose.md

@@ -0,0 +1,5 @@
+---
+"@langchain/langgraph": patch
+---
+
+Improve performance of scheduling tasks with large graphs

libs/langgraph/src/pregel/algo.ts

@@ -106,6 +106,10 @@ function triggersNextStep(
 }
 
 // Avoids unnecessary double iteration
+const MAX_CHANNEL_MAP_VERSION_CACHE = new WeakMap<
+  Record<string, number | string>,
+  number | string
+>();
 function maxChannelMapVersion(
   channelVersions: Record<string, number | string>
 ): number | string | undefined {
@@ -249,7 +253,7 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   getNextVersion: ((version: any) => any) | undefined,
   triggerToNodes: Record<string, string[]> | undefined
-): void {
+): Set<string> {
   // Sort tasks by first 3 path elements for deterministic order
   // Later path parts (like task IDs) are ignored for sorting
   tasks.sort((a, b) => {
@@ -285,7 +289,9 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
   }
 
   // Find the highest version of all channels
-  let maxVersion = maxChannelMapVersion(checkpoint.channel_versions);
+  let maxVersion =
+    MAX_CHANNEL_MAP_VERSION_CACHE.get(checkpoint.channel_versions) ??
+    maxChannelMapVersion(checkpoint.channel_versions);
 
   // Consume all channels that were read
   const channelsToConsume = new Set(
@@ -294,12 +300,12 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
       .filter((chan) => !RESERVED.includes(chan))
   );
 
-  let usedNewVersion = false;
+  let usedNextVersion = false;
   for (const chan of channelsToConsume) {
     if (chan in onlyChannels && onlyChannels[chan].consume()) {
       if (getNextVersion !== undefined) {
         checkpoint.channel_versions[chan] = getNextVersion(maxVersion);
-        usedNewVersion = true;
+        usedNextVersion = true;
       }
     }
   }
@@ -318,8 +324,11 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
   }
 
   // Find the highest version of all channels
-  if (maxVersion != null && getNextVersion != null) {
-    maxVersion = usedNewVersion ? getNextVersion(maxVersion) : maxVersion;
+  if (maxVersion != null && getNextVersion != null && usedNextVersion) {
+    maxVersion = getNextVersion(maxVersion);
+    usedNextVersion = false;
+
+    MAX_CHANNEL_MAP_VERSION_CACHE.set(checkpoint.channel_versions, maxVersion!);
   }
 
   const updatedChannels: Set<string> = new Set();
@@ -346,6 +355,7 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
       }
       if (updated && getNextVersion !== undefined) {
         checkpoint.channel_versions[chan] = getNextVersion(maxVersion);
+        usedNextVersion = true;
 
         // unavailable channels can't trigger tasks, so don't add them
         if (channel.isAvailable()) updatedChannels.add(chan);
@@ -364,6 +374,7 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
 
         if (updated && getNextVersion !== undefined) {
           checkpoint.channel_versions[chan] = getNextVersion(maxVersion);
+          usedNextVersion = true;
 
           // unavailable channels can't trigger tasks, so don't add them
           if (channel.isAvailable()) updatedChannels.add(chan);
@@ -380,12 +391,63 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
       const channel = onlyChannels[chan];
       if (channel.finish() && getNextVersion !== undefined) {
         checkpoint.channel_versions[chan] = getNextVersion(maxVersion);
+        usedNextVersion = true;
 
         // unavailable channels can't trigger tasks, so don't add them
         if (channel.isAvailable()) updatedChannels.add(chan);
       }
     }
   }
+
+  if (maxVersion != null && getNextVersion != null && usedNextVersion) {
+    maxVersion = getNextVersion(maxVersion);
+    usedNextVersion = false;
+    MAX_CHANNEL_MAP_VERSION_CACHE.set(checkpoint.channel_versions, maxVersion!);
+  }
+
+  return updatedChannels;
+}
+
+function* candidateNodes(
+  checkpoint: ReadonlyCheckpoint,
+  processes: StrRecord<string, PregelNode>,
+  extra: NextTaskExtraFields
+) {
+  // This section is an optimization that allows which
+  // nodes will be active during the next step.
+  // When there's information about:
+  // 1. Which channels were updated in the previous step
+  // 2. Which nodes are triggered by which channels
+  // Then we can determine which nodes should be triggered
+  // in the next step without having to cycle through all nodes.
+  if (extra.updatedChannels != null && extra.triggerToNodes != null) {
+    const triggeredNodes = new Set<string>();
+
+    // Get all nodes that have triggers associated with an updated channel
+    for (const channel of extra.updatedChannels) {
+      const nodeIds = extra.triggerToNodes[channel];
+      for (const id of nodeIds ?? []) triggeredNodes.add(id);
+    }
+
+    // Sort the nodes to ensure deterministic order
+    yield* [...triggeredNodes].sort();
+    return;
+  }
+
+  // If there are no values in checkpoint, no need to run
+  // through all the PULL candidates
+  const isEmptyChannelVersions = (() => {
+    for (const chan in checkpoint.channel_versions) {
+      if (checkpoint.channel_versions[chan] !== null) return false;
+    }
+    return true;
+  })();
+
+  if (isEmptyChannelVersions) return;
+  for (const name in processes) {
+    if (!Object.prototype.hasOwnProperty.call(processes, name)) continue;
+    yield name;
+  }
 }
 
 export type NextTaskExtraFields = {
@@ -395,6 +457,8 @@ export type NextTaskExtraFields = {
   manager?: CallbackManagerForChainRun;
   store?: BaseStore;
   stream?: IterableReadableWritableStream;
+  triggerToNodes?: Record<string, string[]>;
+  updatedChannels?: Set<string>;
 };
 
 export type NextTaskExtraFieldsWithStore = NextTaskExtraFields & {
@@ -478,8 +542,7 @@ export function _prepareNextTasks<
 
   // Check if any processes should be run in next step
   // If so, prepare the values to be passed to them
-  for (const name in processes) {
-    if (!Object.prototype.hasOwnProperty.call(processes, name)) continue;
+  for (const name of candidateNodes(checkpoint, processes, extra)) {
     const task = _prepareSingleTask(
       [PULL, name],
       checkpoint,

libs/langgraph/src/pregel/index.ts

@@ -547,6 +547,7 @@ export class Pregel<
     this.store = fields.store;
     this.cache = fields.cache;
     this.name = fields.name;
+    this.triggerToNodes = fields.triggerToNodes ?? this.triggerToNodes;
 
     if (this.autoValidate) {
       this.validate();

libs/langgraph/src/pregel/loop.ts

@@ -255,6 +255,8 @@ export class PregelLoop {
 
   protected prevCheckpointConfig: RunnableConfig | undefined;
 
+  protected updatedChannels: Set<string> | undefined;
+
   status:
     | "pending"
     | "done"
@@ -701,7 +703,7 @@ export class PregelLoop {
       // finish superstep
       const writes = Object.values(this.tasks).flatMap((t) => t.writes);
       // All tasks have finished
-      _applyWrites(
+      this.updatedChannels = _applyWrites(
         this.checkpoint,
         this.channels,
         Object.values(this.tasks),
@@ -757,6 +759,8 @@ export class PregelLoop {
         manager: this.manager,
         store: this.store,
         stream: this.stream,
+        triggerToNodes: this.triggerToNodes,
+        updatedChannels: this.updatedChannels,
       }
     );
     this.tasks = nextTasks;
@@ -857,7 +861,7 @@ export class PregelLoop {
         this.checkpointPendingWrites.length > 0 &&
         Object.values(this.tasks).some((task) => task.writes.length > 0)
       ) {
-        _applyWrites(
+        this.updatedChannels = _applyWrites(
           this.checkpoint,
           this.channels,
           Object.values(this.tasks),
@@ -1067,7 +1071,7 @@ export class PregelLoop {
           true,
           { step: this.step }
         );
-        _applyWrites(
+        this.updatedChannels = _applyWrites(
           this.checkpoint,
           this.channels,
           (Object.values(discardTasks) as WritesProtocol[]).concat([

libs/langgraph/src/pregel/types.ts

@@ -478,6 +478,12 @@ export type PregelParams<
    * Storage used for node caching.
    */
   cache?: BaseCache;
+
+  /**
+   * The trigger to node mapping for the graph run.
+   * @internal
+   */
+  triggerToNodes?: Record<string, string[]>;
 };
 
 export interface PregelTaskDescription {