[model-gateway] fix: handle workflow deadlock and optimize cycle detection (#15000)

This commit is contained in:
Simo Lin
2025-12-12 08:53:45 -08:00
committed by GitHub
parent 306e5b8d0b
commit 56d0ad47f4
2 changed files with 42 additions and 30 deletions
+23 -15
View File
@@ -1,6 +1,10 @@
//! Workflow definition types //! Workflow definition types
use std::{collections::HashSet, sync::Arc, time::Duration}; use std::{
collections::{HashMap, HashSet},
sync::Arc,
time::Duration,
};
use super::{ use super::{
executor::StepExecutor, executor::StepExecutor,
@@ -111,12 +115,14 @@ impl WorkflowDefinition {
/// - A step depends on a non-existent step /// - A step depends on a non-existent step
/// - There's a cycle in the dependencies /// - There's a cycle in the dependencies
pub fn validate(&self) -> Result<(), String> { 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 // Check all dependencies exist
for step in &self.steps { for step in &self.steps {
for dep in &step.depends_on { for dep in &step.depends_on {
if !step_ids.contains(dep) { if !steps_map.contains_key(dep) {
return Err(format!( return Err(format!(
"Step '{}' depends on non-existent step '{}'", "Step '{}' depends on non-existent step '{}'",
step.id, dep step.id, dep
@@ -130,7 +136,9 @@ impl WorkflowDefinition {
let mut rec_stack = HashSet::new(); let mut rec_stack = HashSet::new();
for step in &self.steps { 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)); return Err(format!("Cycle detected involving step '{}'", step.id));
} }
} }
@@ -138,12 +146,12 @@ impl WorkflowDefinition {
Ok(()) Ok(())
} }
/// DFS helper for cycle detection /// DFS helper for cycle detection with O(1) HashMap lookup
fn has_cycle( fn has_cycle<'a>(
&self, step_id: &'a StepId,
step_id: &StepId, steps_map: &HashMap<&'a StepId, &'a StepDefinition>,
visited: &mut HashSet<StepId>, visited: &mut HashSet<&'a StepId>,
rec_stack: &mut HashSet<StepId>, rec_stack: &mut HashSet<&'a StepId>,
) -> bool { ) -> bool {
if rec_stack.contains(step_id) { if rec_stack.contains(step_id) {
return true; // Back edge found - cycle! return true; // Back edge found - cycle!
@@ -152,13 +160,13 @@ impl WorkflowDefinition {
return false; // Already fully processed return false; // Already fully processed
} }
visited.insert(step_id.clone()); visited.insert(step_id);
rec_stack.insert(step_id.clone()); rec_stack.insert(step_id);
// Find the step and check its dependencies // O(1) lookup instead of linear search
if let Some(step) = self.steps.iter().find(|s| &s.id == step_id) { if let Some(step) = steps_map.get(step_id) {
for dep in &step.depends_on { 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; return true;
} }
} }
+9 -5
View File
@@ -266,7 +266,12 @@ impl WorkflowEngine {
// Handle blocked workflow // Handle blocked workflow
if ready_step_indices.is_empty() && running_count == 0 { if ready_step_indices.is_empty() && running_count == 0 {
if blocked_by_failure { 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()
};
self.state_store.update(instance_id, |s| { self.state_store.update(instance_id, |s| {
s.status = WorkflowStatus::Failed; s.status = WorkflowStatus::Failed;
})?; })?;
@@ -275,14 +280,13 @@ impl WorkflowEngine {
self.event_bus self.event_bus
.publish(WorkflowEvent::WorkflowFailed { .publish(WorkflowEvent::WorkflowFailed {
instance_id, instance_id,
failed_step: failed_step.unwrap_or_else(|| StepId::new("unknown")), failed_step: failed_step
error: "Workflow failed due to step dependency failure".to_string(), .unwrap_or_else(|| StepId::new("internal_scheduler")),
error: error_message,
}) })
.await; .await;
return Ok(()); return Ok(());
} }
break;
}
// Launch ready steps in parallel // Launch ready steps in parallel
for step_idx in ready_step_indices { for step_idx in ready_step_indices {