From 0adfc42b99fad2d5c01ee9644e8b503e5ccf73d1 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Fri, 19 Dec 2025 13:44:57 -1000 Subject: [PATCH] [model-gateway] Optimize workflow engine with pre-computed dependency graph (#15503) --- sgl-model-gateway/src/workflow/definition.rs | 53 +++++++++++++----- sgl-model-gateway/src/workflow/engine.rs | 58 ++++++++++++-------- sgl-model-gateway/tests/workflow_test.rs | 6 +- 3 files changed, 76 insertions(+), 41 deletions(-) diff --git a/sgl-model-gateway/src/workflow/definition.rs b/sgl-model-gateway/src/workflow/definition.rs index 5059be00e..71729863b 100644 --- a/sgl-model-gateway/src/workflow/definition.rs +++ b/sgl-model-gateway/src/workflow/definition.rs @@ -70,6 +70,10 @@ pub struct WorkflowDefinition { pub steps: Vec, 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>, + /// Pre-computed indices of steps with no dependencies (can start immediately) + initial_step_indices: Vec, } 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) - pub fn get_initial_steps(&self) -> Vec<&StepDefinition> { - self.steps - .iter() - .filter(|s| s.depends_on.is_empty()) - .collect() + /// Get indices of steps that depend on the given step + pub fn get_dependent_indices(&self, step_id: &StepId) -> &[usize] { + self.reverse_deps + .get(step_id) + .map(|v| v.as_slice()) + .unwrap_or(&[]) } - /// Get steps that depend on the given step - pub fn get_dependents(&self, step_id: &StepId) -> Vec<&StepDefinition> { - self.steps - .iter() - .filter(|s| s.depends_on.contains(step_id)) - .collect() + /// Get indices of steps with no dependencies + pub fn get_initial_step_indices(&self) -> &[usize] { + &self.initial_step_indices } } diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs index c5d8e6998..a68e7a741 100644 --- a/sgl-model-gateway/src/workflow/engine.rs +++ b/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> = 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 = 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 = definition - .steps - .iter() - .enumerate() - .filter(|(_, step)| { + // Only check steps in pending_check, not all steps + let ready: Vec = 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); + } + } } } } diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs index f14fa82dc..183b29436 100644 --- a/sgl-model-gateway/tests/workflow_test.rs +++ b/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",