[model-gateway] Optimize workflow engine with pre-computed dependency graph (#15503)

sgl-model-gateway/src/workflow/engine.rs

@@ -5,7 +5,7 @@
 //! wait for all dependencies to complete successfully.
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{HashMap, HashSet, VecDeque},
     sync::Arc,
     time::Duration,
 };
@@ -132,8 +132,8 @@ impl WorkflowEngine {
     }
 
     /// Register a workflow definition
-    pub fn register_workflow(&self, definition: WorkflowDefinition) -> Result<(), String> {
-        // Validate DAG once at registration, not on every execution
+    pub fn register_workflow(&self, mut definition: WorkflowDefinition) -> Result<(), String> {
+        // Validate DAG and build dependency graph once at registration
         definition.validate()?;
 
         let id = definition.id.clone();
@@ -203,6 +203,9 @@ impl WorkflowEngine {
     }
 
     /// Execute a workflow with DAG-based parallel execution
+    ///
+    /// Uses event-driven readiness: instead of scanning all steps each iteration,
+    /// we only check steps whose dependencies just completed.
     async fn execute_workflow(
         &self,
         instance_id: WorkflowInstanceId,
@@ -214,6 +217,13 @@ impl WorkflowEngine {
         let tracker: Arc<RwLock<StepTracker>> = Arc::new(RwLock::new(StepTracker::default()));
         let (tx, mut rx) = mpsc::channel::<(StepId, StepResult)>(step_count.max(1));
 
+        // Initialize with steps that have no dependencies (O(1) lookup)
+        let mut pending_check: VecDeque<usize> = definition
+            .get_initial_step_indices()
+            .iter()
+            .copied()
+            .collect();
+
         loop {
             if self.state_store.is_cancelled(instance_id)? {
                 self.event_bus
@@ -222,32 +232,22 @@ impl WorkflowEngine {
                 return Ok(());
             }
 
-            let (ready_step_indices, total_processed, running_count, blocked_by_failure) = {
+            // Find ready steps from pending_check (not all steps)
+            let (ready_step_indices, total_processed, running_count) = {
                 let t = tracker.read();
 
-                let ready: Vec<usize> = definition
-                    .steps
-                    .iter()
-                    .enumerate()
-                    .filter(|(_, step)| {
+                // Only check steps in pending_check, not all steps
+                let ready: Vec<usize> = pending_check
+                    .drain(..)
+                    .filter(|&idx| {
+                        let step = &definition.steps[idx];
                         t.is_step_processable(&step.id)
                             && t.are_dependencies_satisfied(&step.depends_on)
                             && !t.has_failed_dependency(&step.depends_on)
                     })
-                    .map(|(i, _)| i)
                     .collect();
 
-                let processed = t.total_processed();
-                let running = t.running.len();
-
-                // Check for blocked steps only if needed
-                let blocked = ready.is_empty()
-                    && running == 0
-                    && definition.steps.iter().any(|step| {
-                        t.is_step_processable(&step.id) && t.has_failed_dependency(&step.depends_on)
-                    });
-
-                (ready, processed, running, blocked)
+                (ready, t.total_processed(), t.running.len())
             };
 
             // Check if we're done
@@ -255,9 +255,11 @@ impl WorkflowEngine {
                 break;
             }
 
-            // Handle blocked workflow
-            if ready_step_indices.is_empty() && running_count == 0 {
-                let error_message = if blocked_by_failure {
+            // Handle blocked workflow (no ready steps, none running, but work remains)
+            if ready_step_indices.is_empty() && running_count == 0 && pending_check.is_empty() {
+                // Check if blocked by failure
+                let has_failed = !tracker.read().failed.is_empty();
+                let error_message = if has_failed {
                     "Workflow failed due to step dependency failure".to_string()
                 } else {
                     "Workflow deadlocked: no steps ready and none running. This may indicate a scheduler bug.".to_string()
@@ -345,6 +347,14 @@ impl WorkflowEngine {
                         result = ?result,
                         "Step completed"
                     );
+
+                    // Add dependents of completed step to pending_check (O(1) lookup)
+                    // Only if the step succeeded or was skipped (not failed)
+                    if matches!(result, StepResult::Success | StepResult::Skip) {
+                        for &dep_idx in definition.get_dependent_indices(&completed_step_id) {
+                            pending_check.push_back(dep_idx);
+                        }
+                    }
                 }
             }
         }