diff --git a/sgl-model-gateway/src/workflow/definition.rs b/sgl-model-gateway/src/workflow/definition.rs index 5cdb38425..5059be00e 100644 --- a/sgl-model-gateway/src/workflow/definition.rs +++ b/sgl-model-gateway/src/workflow/definition.rs @@ -1,6 +1,10 @@ //! Workflow definition types -use std::{collections::HashSet, sync::Arc, time::Duration}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; use super::{ executor::StepExecutor, @@ -111,12 +115,14 @@ impl WorkflowDefinition { /// - A step depends on a non-existent step /// - There's a cycle in the dependencies pub fn validate(&self) -> Result<(), String> { - let step_ids: HashSet<_> = self.steps.iter().map(|s| &s.id).collect(); + // 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(); // Check all dependencies exist for step in &self.steps { for dep in &step.depends_on { - if !step_ids.contains(dep) { + if !steps_map.contains_key(dep) { return Err(format!( "Step '{}' depends on non-existent step '{}'", step.id, dep @@ -130,7 +136,9 @@ impl WorkflowDefinition { let mut rec_stack = HashSet::new(); for step in &self.steps { - if self.has_cycle(&step.id, &mut visited, &mut rec_stack) { + if !visited.contains(&step.id) + && Self::has_cycle(&step.id, &steps_map, &mut visited, &mut rec_stack) + { return Err(format!("Cycle detected involving step '{}'", step.id)); } } @@ -138,12 +146,12 @@ impl WorkflowDefinition { Ok(()) } - /// DFS helper for cycle detection - fn has_cycle( - &self, - step_id: &StepId, - visited: &mut HashSet, - rec_stack: &mut HashSet, + /// DFS helper for cycle detection with O(1) HashMap lookup + fn has_cycle<'a>( + step_id: &'a StepId, + steps_map: &HashMap<&'a StepId, &'a StepDefinition>, + visited: &mut HashSet<&'a StepId>, + rec_stack: &mut HashSet<&'a StepId>, ) -> bool { if rec_stack.contains(step_id) { return true; // Back edge found - cycle! @@ -152,13 +160,13 @@ impl WorkflowDefinition { return false; // Already fully processed } - visited.insert(step_id.clone()); - rec_stack.insert(step_id.clone()); + visited.insert(step_id); + rec_stack.insert(step_id); - // Find the step and check its dependencies - if let Some(step) = self.steps.iter().find(|s| &s.id == step_id) { + // O(1) lookup instead of linear search + if let Some(step) = steps_map.get(step_id) { for dep in &step.depends_on { - if self.has_cycle(dep, visited, rec_stack) { + if Self::has_cycle(dep, steps_map, visited, rec_stack) { return true; } } diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs index 7576ddf9b..56540ef95 100644 --- a/sgl-model-gateway/src/workflow/engine.rs +++ b/sgl-model-gateway/src/workflow/engine.rs @@ -266,22 +266,26 @@ impl WorkflowEngine { // Handle blocked workflow if ready_step_indices.is_empty() && running_count == 0 { - if blocked_by_failure { - self.state_store.update(instance_id, |s| { - s.status = WorkflowStatus::Failed; - })?; + let error_message = if blocked_by_failure { + "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() + }; - let failed_step = tracker.read().failed.iter().next().cloned(); - self.event_bus - .publish(WorkflowEvent::WorkflowFailed { - instance_id, - failed_step: failed_step.unwrap_or_else(|| StepId::new("unknown")), - error: "Workflow failed due to step dependency failure".to_string(), - }) - .await; - return Ok(()); - } - break; + self.state_store.update(instance_id, |s| { + s.status = WorkflowStatus::Failed; + })?; + + let failed_step = tracker.read().failed.iter().next().cloned(); + self.event_bus + .publish(WorkflowEvent::WorkflowFailed { + instance_id, + failed_step: failed_step + .unwrap_or_else(|| StepId::new("internal_scheduler")), + error: error_message, + }) + .await; + return Ok(()); } // Launch ready steps in parallel