@ai-setting/roy-agent-core 1.0.0

package/src/env/workflow/engine/scheduler.ts

@@ -0,0 +1,318 @@
+import { DAGManager } from './dag-manager';
+
+/**
+ * Scheduler Options
+ */
+export interface SchedulerOptions {
+  /** Maximum number of nodes that can run in parallel. null = unlimited */
+  parallelLimit: number | null;
+}
+
+/**
+ * Scheduler State
+ */
+export interface SchedulerState {
+  /** Nodes that have not been started yet */
+  pending: string[];
+  /** Nodes that are ready to execute (dependencies met, not started) */
+  ready: string[];
+  /** Nodes that are currently running */
+  running: string[];
+  /** Nodes that have completed successfully */
+  completed: string[];
+  /** Nodes that have failed */
+  failed: string[];
+  /** Nodes that were skipped (not executed) */
+  skipped: string[];
+}
+
+/**
+ * Scheduler
+ * 
+ * Determines which nodes are ready to execute based on:
+ * - DAG dependencies
+ * - Parallel execution limit
+ * - Node completion status
+ * 
+ * The scheduler tracks the state of all nodes and manages the execution queue.
+ */
+export class Scheduler {
+  private dagManager: DAGManager;
+  private options: SchedulerOptions;
+  
+  // Node state tracking
+  private pending: Set<string> = new Set();
+  private ready: Set<string> = new Set();
+  private running: Set<string> = new Set();
+  private completed: Set<string> = new Set();
+  private failed: Set<string> = new Set();
+  private skipped: Set<string> = new Set();
+
+  /**
+   * Create a new Scheduler
+   * @param dagManager - The DAG manager for the workflow
+   * @param options - Scheduler options
+   */
+  constructor(dagManager: DAGManager, options?: Partial<SchedulerOptions>) {
+    this.dagManager = dagManager;
+    // Use explicit undefined check to preserve null as valid value
+    const parallelLimit = options?.parallelLimit;
+    this.options = {
+      parallelLimit: parallelLimit === undefined ? 1 : parallelLimit,
+    };
+    
+    // Initialize pending set with all nodes
+    const analysis = dagManager.analyze();
+    for (const nodeId of analysis.dependencies.keys()) {
+      this.pending.add(nodeId);
+    }
+    
+    // Calculate initial ready nodes
+    this.updateReadyNodes(new Set());
+  }
+
+  /**
+   * Get nodes that are ready to execute
+   * Ready nodes are those whose dependencies are met and are not already running/completed/failed
+   * @param completedNodes - Set of completed node IDs (external tracking)
+   */
+  getReadyNodes(completedNodes: Set<string>): string[] {
+    // Calculate all completed nodes (including those tracked internally)
+    const allCompleted = new Set(completedNodes);
+    for (const nodeId of this.completed) {
+      allCompleted.add(nodeId);
+    }
+    
+    // Get ready nodes from DAG manager
+    const readyFromDag = this.dagManager.getReadyNodes(allCompleted);
+    
+    // Filter out nodes that are already running, completed, failed, or skipped
+    // Also ensure the node is in pending or ready state (not already processed)
+    const ready: string[] = [];
+    for (const nodeId of readyFromDag) {
+      const isProcessed = this.running.has(nodeId) || 
+                          this.completed.has(nodeId) || 
+                          this.failed.has(nodeId) ||
+                          this.skipped.has(nodeId);
+      
+      if (!isProcessed) {
+        ready.push(nodeId);
+      }
+    }
+    
+    // Sort for deterministic order
+    return ready.sort();
+  }
+
+  /**
+   * Get the number of currently running nodes
+   */
+  getRunningCount(): number {
+    return this.running.size;
+  }
+
+  /**
+   * Check if more nodes can be started based on the parallel limit
+   */
+  canStartMore(): boolean {
+    // null means unlimited
+    if (this.options.parallelLimit === null) {
+      return true;
+    }
+    
+    if (this.options.parallelLimit <= 0) {
+      return false;
+    }
+    
+    return this.running.size < this.options.parallelLimit;
+  }
+
+  /**
+   * Mark a node as started
+   * @param nodeId - The node ID
+   */
+  markStarted(nodeId: string): void {
+    // Only track if not already running
+    if (!this.running.has(nodeId)) {
+      this.running.add(nodeId);
+      this.ready.delete(nodeId);
+      this.pending.delete(nodeId);
+    }
+  }
+  
+  /**
+   * Check if a node is currently running
+   * @param nodeId - The node ID
+   */
+  isRunning(nodeId: string): boolean {
+    return this.running.has(nodeId);
+  }
+  
+  /**
+   * Check if a node has completed
+   * @param nodeId - The node ID
+   */
+  isCompleted(nodeId: string): boolean {
+    return this.completed.has(nodeId);
+  }
+  
+  /**
+   * Check if a node has failed
+   * @param nodeId - The node ID
+   */
+  isFailed(nodeId: string): boolean {
+    return this.failed.has(nodeId);
+  }
+  
+  /**
+   * Check if a node is in pending state
+   * @param nodeId - The node ID
+   */
+  isPending(nodeId: string): boolean {
+    return this.pending.has(nodeId);
+  }
+  
+  /**
+   * Check if a node is ready to execute
+   * @param nodeId - The node ID
+   */
+  isReady(nodeId: string): boolean {
+    return this.ready.has(nodeId);
+  }
+
+  /**
+   * Mark a node as completed successfully
+   * @param nodeId - The node ID
+   */
+  markCompleted(nodeId: string): void {
+    // Remove from other states
+    this.running.delete(nodeId);
+    this.ready.delete(nodeId);
+    this.pending.delete(nodeId);
+    this.failed.delete(nodeId);
+    this.skipped.delete(nodeId);
+    
+    // Add to completed
+    this.completed.add(nodeId);
+  }
+
+  /**
+   * Mark a node as failed
+   * @param nodeId - The node ID
+   */
+  markFailed(nodeId: string): void {
+    // Remove from any other state
+    this.running.delete(nodeId);
+    this.ready.delete(nodeId);
+    this.pending.delete(nodeId);
+    this.completed.delete(nodeId);
+    this.skipped.delete(nodeId);
+    
+    // Add to failed
+    this.failed.add(nodeId);
+  }
+
+  /**
+   * Mark a node as skipped (not executed)
+   * @param nodeId - The node ID
+   */
+  markSkipped(nodeId: string): void {
+    if (this.running.has(nodeId)) {
+      this.running.delete(nodeId);
+    }
+    if (this.ready.has(nodeId)) {
+      this.ready.delete(nodeId);
+    }
+    if (this.pending.has(nodeId)) {
+      this.pending.delete(nodeId);
+    }
+    this.skipped.add(nodeId);
+  }
+
+  /**
+   * Reset the scheduler state
+   * Clears all running, completed, failed, and skipped nodes
+   */
+  reset(): void {
+    this.running.clear();
+    this.completed.clear();
+    this.failed.clear();
+    this.skipped.clear();
+    
+    // Rebuild pending and ready from DAG analysis
+    const analysis = this.dagManager.analyze();
+    
+    this.pending.clear();
+    this.ready.clear();
+    
+    for (const nodeId of analysis.dependencies.keys()) {
+      this.pending.add(nodeId);
+    }
+    
+    this.updateReadyNodes(new Set());
+  }
+
+  /**
+   * Get the current state of the scheduler
+   */
+  getState(): SchedulerState {
+    // Rebuild ready set based on current state
+    this.updateReadyNodes(new Set());
+    
+    return {
+      pending: Array.from(this.pending).sort(),
+      ready: Array.from(this.ready).sort(),
+      running: Array.from(this.running).sort(),
+      completed: Array.from(this.completed).sort(),
+      failed: Array.from(this.failed).sort(),
+      skipped: Array.from(this.skipped).sort(),
+    };
+  }
+
+  /**
+   * Update the ready set based on completed nodes
+   */
+  private updateReadyNodes(completedNodes: Set<string>): void {
+    const allCompleted = new Set(completedNodes);
+    for (const nodeId of this.completed) {
+      allCompleted.add(nodeId);
+    }
+    
+    // Get all node IDs
+    const analysis = this.dagManager.analyze();
+    const allNodeIds = Array.from(analysis.dependencies.keys());
+    
+    // Clear and rebuild ready set
+    this.ready.clear();
+    
+    for (const nodeId of allNodeIds) {
+      // Skip if already in some terminal or transitional state
+      if (this.running.has(nodeId) || 
+          this.completed.has(nodeId) || 
+          this.failed.has(nodeId) ||
+          this.skipped.has(nodeId)) {
+        continue;
+      }
+      
+      // Check if dependencies are met
+      if (this.dagManager.areDependenciesMet(nodeId, allCompleted)) {
+        this.ready.add(nodeId);
+      }
+    }
+  }
+
+  /**
+   * Get the parallel limit
+   */
+  getParallelLimit(): number | null {
+    return this.options.parallelLimit;
+  }
+
+  /**
+   * Update the parallel limit
+   * @param limit - New parallel limit (null for unlimited)
+   */
+  setParallelLimit(limit: number | null): void {
+    this.options.parallelLimit = limit;
+  }
+}