chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Simo Lin
2025-12-19 13:44:57 -10:00
committed by GitHub
parent ba72e759ca
commit 0adfc42b99
3 changed files with 76 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
sgl-model-gateway/src/workflow/definition.rs
View File

@@ -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)
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
}
}

58
sgl-model-gateway/src/workflow/engine.rs
View File

@@ -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);
}
}
}
}
}

6
sgl-model-gateway/tests/workflow_test.rs
View File

@@ -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",