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

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

@@ -70,6 +70,10 @@ pub struct WorkflowDefinition {
     pub steps: Vec<StepDefinition>,
     pub default_retry_policy: RetryPolicy,
     pub default_timeout: Duration,
+    /// Pre-computed reverse dependencies: step_id -> indices of steps that depend on it
+    reverse_deps: HashMap<StepId, Vec<usize>>,
+    /// Pre-computed indices of steps with no dependencies (can start immediately)
+    initial_step_indices: Vec<usize>,
 }
 
 impl WorkflowDefinition {
@@ -80,6 +84,8 @@ impl WorkflowDefinition {
             steps: Vec::new(),
             default_retry_policy: RetryPolicy::default(),
             default_timeout: Duration::from_secs(300), // 5 minutes
+            reverse_deps: HashMap::new(),
+            initial_step_indices: Vec::new(),
         }
     }
 
@@ -110,11 +116,13 @@ impl WorkflowDefinition {
         step.timeout.unwrap_or(self.default_timeout)
     }
 
-    /// Validate the workflow DAG structure.
+    /// Validate the workflow DAG structure and build dependency graph.
     /// Returns an error if:
     /// - A step depends on a non-existent step
     /// - There's a cycle in the dependencies
-    pub fn validate(&self) -> Result<(), String> {
+    ///
+    /// On success, pre-computes reverse dependencies for O(1) dependent lookup.
+    pub fn validate(&mut self) -> Result<(), String> {
         // Build HashMap for O(1) lookup instead of O(n) linear search
         let steps_map: HashMap<&StepId, &StepDefinition> =
             self.steps.iter().map(|s| (&s.id, s)).collect();
@@ -143,6 +151,26 @@ impl WorkflowDefinition {
             }
         }
 
+        // Build reverse dependency map: for each step, which steps depend on it?
+        self.reverse_deps.clear();
+        for (idx, step) in self.steps.iter().enumerate() {
+            for dep_id in &step.depends_on {
+                self.reverse_deps
+                    .entry(dep_id.clone())
+                    .or_default()
+                    .push(idx);
+            }
+        }
+
+        // Cache indices of steps with no dependencies (can start immediately)
+        self.initial_step_indices = self
+            .steps
+            .iter()
+            .enumerate()
+            .filter(|(_, s)| s.depends_on.is_empty())
+            .map(|(i, _)| i)
+            .collect();
+
         Ok(())
     }
 
@@ -176,19 +204,16 @@ impl WorkflowDefinition {
         false
     }
 
-    /// Get steps that have no dependencies (can run immediately)
+    /// Get indices of steps that depend on the given step
-    pub fn get_initial_steps(&self) -> Vec<&StepDefinition> {
+    pub fn get_dependent_indices(&self, step_id: &StepId) -> &[usize] {
-        self.steps
+        self.reverse_deps
-            .iter()
+            .get(step_id)
-            .filter(|s| s.depends_on.is_empty())
+            .map(|v| v.as_slice())
-            .collect()
+            .unwrap_or(&[])
     }
 
-    /// Get steps that depend on the given step
+    /// Get indices of steps with no dependencies
-    pub fn get_dependents(&self, step_id: &StepId) -> Vec<&StepDefinition> {
+    pub fn get_initial_step_indices(&self) -> &[usize] {
-        self.steps
+        &self.initial_step_indices
-            .iter()
-            .filter(|s| s.depends_on.contains(step_id))
-            .collect()
     }
 }

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> {
+    pub fn register_workflow(&self, mut definition: WorkflowDefinition) -> Result<(), String> {
-        // Validate DAG once at registration, not on every execution
+        // 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
+                // Only check steps in pending_check, not all steps
-                    .steps
+                let ready: Vec<usize> = pending_check
-                    .iter()
+                    .drain(..)
-                    .enumerate()
+                    .filter(|&idx| {
-                    .filter(|(_, step)| {
+                        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();
+                (ready, t.total_processed(), t.running.len())
-                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)
             };
 
             // Check if we're done
@@ -255,9 +255,11 @@ impl WorkflowEngine {
                 break;
             }
 
-            // Handle blocked workflow
+            // Handle blocked workflow (no ready steps, none running, but work remains)
-            if ready_step_indices.is_empty() && running_count == 0 {
+            if ready_step_indices.is_empty() && running_count == 0 && pending_check.is_empty() {
-                let error_message = if blocked_by_failure {
+                // 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);
+                        }
+                    }
                 }
             }
         }

sgl-model-gateway/tests/workflow_test.rs

@@ -598,7 +598,7 @@ async fn test_dag_dependency_failure_blocks_dependents() {
 #[test]
 fn test_dag_validation_cycle_detection() {
     // Create a workflow with a cycle: A -> B -> C -> A
-    let workflow = WorkflowDefinition::new("cyclic_workflow", "Cyclic Test")
+    let mut workflow = WorkflowDefinition::new("cyclic_workflow", "Cyclic Test")
         .add_step(
             StepDefinition::new("step_a", "Step A", Arc::new(AlwaysSucceedStep))
                 .depends_on(&["step_c"]),
@@ -620,7 +620,7 @@ fn test_dag_validation_cycle_detection() {
 #[test]
 fn test_dag_validation_missing_dependency() {
     // Create a workflow with a missing dependency
-    let workflow = WorkflowDefinition::new("missing_dep_workflow", "Missing Dep Test")
+    let mut workflow = WorkflowDefinition::new("missing_dep_workflow", "Missing Dep Test")
         .add_step(StepDefinition::new(
             "step_a",
             "Step A",
@@ -639,7 +639,7 @@ fn test_dag_validation_missing_dependency() {
 #[test]
 fn test_dag_validation_valid_workflow() {
     // Create a valid DAG workflow
-    let workflow = WorkflowDefinition::new("valid_workflow", "Valid Test")
+    let mut workflow = WorkflowDefinition::new("valid_workflow", "Valid Test")
         .add_step(StepDefinition::new(
             "step_a",
             "Step A",