perf(langgraph): avoid iterating on all nodes when scheduling new tasks (#1628)

.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

@@ -249,7 +249,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) => {
@@ -386,6 +386,50 @@ export function _applyWrites<Cc extends Record<string, BaseChannel>>(
       }
     }
   }
+
+  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 +439,8 @@ export type NextTaskExtraFields = {
   manager?: CallbackManagerForChainRun;
   store?: BaseStore;
   stream?: IterableReadableWritableStream;
+  triggerToNodes?: Record<string, string[]>;
+  updatedChannels?: Set<string>;
 };
 
 export type NextTaskExtraFieldsWithStore = NextTaskExtraFields & {
@@ -478,8 +524,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 {