From 306e5b8d0b21a9dcea255d6e3245356f93fd5d60 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Fri, 12 Dec 2025 08:31:12 -0800 Subject: [PATCH] [model-gateway] feat: add DAG parallel execution support and workflow optimization (#14999) --- .../steps/external_worker_registration.rs | 28 +- .../src/core/steps/mcp_registration.rs | 19 +- .../core/steps/wasm_module_registration.rs | 27 +- .../src/core/steps/wasm_module_removal.rs | 11 +- .../src/core/steps/worker_registration.rs | 32 +- .../src/core/steps/worker_removal.rs | 20 +- sgl-model-gateway/src/server.rs | 24 +- sgl-model-gateway/src/workflow/definition.rs | 90 ++++- sgl-model-gateway/src/workflow/engine.rs | 284 ++++++++++---- sgl-model-gateway/src/workflow/state.rs | 22 ++ sgl-model-gateway/tests/common/mod.rs | 16 +- sgl-model-gateway/tests/wasm_test.rs | 34 +- sgl-model-gateway/tests/workflow_test.rs | 354 +++++++++++++++++- 13 files changed, 768 insertions(+), 193 deletions(-) diff --git a/sgl-model-gateway/src/core/steps/external_worker_registration.rs b/sgl-model-gateway/src/core/steps/external_worker_registration.rs index 101b1abef..640205843 100644 --- a/sgl-model-gateway/src/core/steps/external_worker_registration.rs +++ b/sgl-model-gateway/src/core/steps/external_worker_registration.rs @@ -1,19 +1,3 @@ -//! External worker registration workflow steps -//! -//! This workflow handles registration of external API endpoints (OpenAI, xAI, Anthropic, etc.) -//! -//! Key features: -//! - Fetches models from /v1/models endpoint -//! - Groups dated model variants under base model names (e.g., gpt-4o, gpt-4o-2024-08-06) -//! - Infers ModelType from model ID patterns (LLM, embedding, image gen, audio, etc.) -//! -//! Workflow order: -//! 1. DiscoverModels - Fetch available models from /v1/models endpoint -//! 2. CreateExternalWorkers - Build worker objects for each discovered model -//! 3. RegisterWorkers - Register workers in registry -//! 4. UpdatePolicies - Update policy registry -//! 5. ActivateWorkers - Mark workers as healthy - use std::{collections::HashMap, sync::Arc, time::Duration}; use async_trait::async_trait; @@ -571,7 +555,8 @@ pub fn create_external_worker_registration_workflow() -> WorkflowDefinition { Arc::new(CreateExternalWorkersStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["discover_models"]), ) .add_step( StepDefinition::new( @@ -580,7 +565,8 @@ pub fn create_external_worker_registration_workflow() -> WorkflowDefinition { Arc::new(RegisterExternalWorkersStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["create_workers"]), ) .add_step( StepDefinition::new( @@ -589,7 +575,8 @@ pub fn create_external_worker_registration_workflow() -> WorkflowDefinition { Arc::new(UpdateExternalPoliciesStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::ContinueNextStep), + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["register_workers"]), ) .add_step( StepDefinition::new( @@ -598,6 +585,7 @@ pub fn create_external_worker_registration_workflow() -> WorkflowDefinition { Arc::new(ActivateExternalWorkersStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["update_policies"]), ) } diff --git a/sgl-model-gateway/src/core/steps/mcp_registration.rs b/sgl-model-gateway/src/core/steps/mcp_registration.rs index f236a80a2..e496fdab1 100644 --- a/sgl-model-gateway/src/core/steps/mcp_registration.rs +++ b/sgl-model-gateway/src/core/steps/mcp_registration.rs @@ -1,13 +1,3 @@ -//! MCP server registration workflow steps -//! -//! Each step is atomic and performs a single operation in the MCP server registration process. -//! Updated for flat manager architecture - single McpManager manages all clients directly. -//! -//! Workflow order: -//! 1. ConnectMcpServer - Establish connection to MCP server using McpManager::connect_server() -//! 2. DiscoverMcpInventory - Discover and cache inventory using McpManager::load_server_inventory() -//! 3. RegisterMcpServer - Register McpClient in McpManager's client map - use std::{sync::Arc, time::Duration}; use async_trait::async_trait; @@ -263,7 +253,8 @@ pub fn create_mcp_registration_workflow() -> WorkflowDefinition { backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::ContinueNextStep), + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["connect_mcp_server"]), ) .add_step( StepDefinition::new( @@ -272,7 +263,8 @@ pub fn create_mcp_registration_workflow() -> WorkflowDefinition { Arc::new(RegisterMcpServerStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::ContinueNextStep), + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["discover_mcp_inventory"]), ) .add_step( StepDefinition::new( @@ -281,6 +273,7 @@ pub fn create_mcp_registration_workflow() -> WorkflowDefinition { Arc::new(ValidateRegistrationStep), ) .with_timeout(Duration::from_secs(1)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["register_mcp_server"]), ) } diff --git a/sgl-model-gateway/src/core/steps/wasm_module_registration.rs b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs index 417d5bad3..5358be6de 100644 --- a/sgl-model-gateway/src/core/steps/wasm_module_registration.rs +++ b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs @@ -1,15 +1,3 @@ -//! WASM Module Registration Workflow Steps -//! -//! Each step is atomic and performs a single operation in the WASM module registration process. -//! -//! Workflow order: -//! 1. ValidateDescriptor - Validate module descriptor (name, file_path, file existence) -//! 2. CalculateHash - Calculate SHA256 hash of the module file -//! 3. CheckDuplicate - Check for duplicate SHA256 hash -//! 4. LoadWasmBytes - Load WASM bytes into memory -//! 5. ValidateWasmComponent - Validate WASM component format -//! 6. RegisterModule - Register module in WasmModuleManager - use std::{ path::{Component as PathComponent, Path}, sync::Arc, @@ -538,7 +526,8 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition { backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(60)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["validate_descriptor"]), ) .add_step( StepDefinition::new( @@ -547,7 +536,8 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition { Arc::new(CheckDuplicateStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["calculate_hash"]), ) .add_step( StepDefinition::new( @@ -560,7 +550,8 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition { backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(60)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["check_duplicate"]), ) .add_step( StepDefinition::new( @@ -569,7 +560,8 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition { Arc::new(ValidateWasmComponentStep), ) .with_timeout(Duration::from_secs(30)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["load_wasm_bytes"]), ) .add_step( StepDefinition::new( @@ -578,6 +570,7 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition { Arc::new(RegisterModuleStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["validate_wasm_component"]), ) } diff --git a/sgl-model-gateway/src/core/steps/wasm_module_removal.rs b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs index 3a0efc023..8405344fb 100644 --- a/sgl-model-gateway/src/core/steps/wasm_module_removal.rs +++ b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs @@ -1,11 +1,3 @@ -//! WASM Module Removal Workflow Steps -//! -//! Each step is atomic and performs a single operation in the WASM module removal process. -//! -//! Workflow order: -//! 1. FindModuleToRemove - Find the module to remove by UUID -//! 2. RemoveModule - Remove module from WasmModuleManager - use std::{sync::Arc, time::Duration}; use async_trait::async_trait; @@ -147,6 +139,7 @@ pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition { .add_step( StepDefinition::new("remove_module", "Remove Module", Arc::new(RemoveModuleStep)) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["find_module_to_remove"]), ) } diff --git a/sgl-model-gateway/src/core/steps/worker_registration.rs b/sgl-model-gateway/src/core/steps/worker_registration.rs index d4d4dee8d..c936615a8 100644 --- a/sgl-model-gateway/src/core/steps/worker_registration.rs +++ b/sgl-model-gateway/src/core/steps/worker_registration.rs @@ -1,17 +1,3 @@ -//! Local worker registration workflow steps -//! -//! This workflow handles registration of local inference workers (SGLang, vLLM). -//! For external API endpoints (OpenAI, xAI, etc.), see external_worker_registration.rs. -//! -//! Workflow order: -//! 1. DetectConnectionMode - Probe HTTP and gRPC to determine connection mode -//! 2. DiscoverMetadata - Fetch metadata from /server_info or gRPC -//! 3. DiscoverDPInfo - Fetch DP (Data Parallel) information (only for DP-aware workers) -//! 4. CreateWorker - Build worker object(s) with merged config + metadata -//! 5. RegisterWorker - Register worker(s) in registry -//! 6. UpdatePolicies - Update policy registry with worker information -//! 7. ActivateWorker - Mark worker(s) as healthy - use std::{collections::HashMap, sync::Arc, time::Duration}; use async_trait::async_trait; @@ -901,7 +887,8 @@ pub fn create_worker_registration_workflow( backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::ContinueNextStep), + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["detect_connection_mode"]), ) .add_step( StepDefinition::new( @@ -914,12 +901,14 @@ pub fn create_worker_registration_workflow( backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), }) .with_timeout(Duration::from_secs(10)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["discover_metadata"]), ) .add_step( StepDefinition::new("create_worker", "Create Worker", Arc::new(CreateWorkerStep)) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["discover_dp_info"]), ) .add_step( StepDefinition::new( @@ -928,7 +917,8 @@ pub fn create_worker_registration_workflow( Arc::new(RegisterWorkerStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["create_worker"]), ) .add_step( StepDefinition::new( @@ -937,7 +927,8 @@ pub fn create_worker_registration_workflow( Arc::new(UpdatePoliciesStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::ContinueNextStep), + .with_failure_action(FailureAction::ContinueNextStep) + .depends_on(&["register_worker"]), ) .add_step( StepDefinition::new( @@ -946,6 +937,7 @@ pub fn create_worker_registration_workflow( Arc::new(ActivateWorkerStep), ) .with_timeout(Duration::from_secs(5)) - .with_failure_action(FailureAction::FailWorkflow), + .with_failure_action(FailureAction::FailWorkflow) + .depends_on(&["update_policies"]), ) } diff --git a/sgl-model-gateway/src/core/steps/worker_removal.rs b/sgl-model-gateway/src/core/steps/worker_removal.rs index f8770ab7e..dad96fa90 100644 --- a/sgl-model-gateway/src/core/steps/worker_removal.rs +++ b/sgl-model-gateway/src/core/steps/worker_removal.rs @@ -1,14 +1,3 @@ -//! Worker Removal Workflow Steps -//! -//! This module implements the workflow steps for removing workers from the router. -//! Handles both single worker removal and DP-aware worker removal with prefix matching. -//! -//! Steps: -//! 1. FindWorkersToRemove - Identify workers to remove based on URL (handles DP-aware prefix matching) -//! 2. RemoveFromPolicyRegistry - Remove workers from policy registry and cache-aware policies -//! 3. RemoveFromWorkerRegistry - Remove workers from worker registry -//! 4. UpdateRemainingPolicies - Update cache-aware policies for remaining workers - use std::{collections::HashSet, sync::Arc, time::Duration}; use async_trait::async_trait; @@ -259,7 +248,8 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition { .with_retry(RetryPolicy { max_attempts: 1, backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }), + }) + .depends_on(&["find_workers_to_remove"]), ) .add_step( StepDefinition::new( @@ -271,7 +261,8 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition { .with_retry(RetryPolicy { max_attempts: 1, backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }), + }) + .depends_on(&["remove_from_policy_registry"]), ) .add_step( StepDefinition::new( @@ -283,6 +274,7 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition { .with_retry(RetryPolicy { max_attempts: 1, backoff: BackoffStrategy::Fixed(Duration::from_secs(0)), - }), + }) + .depends_on(&["remove_from_worker_registry"]), ) } diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 8cc03b85a..5f992d756 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -770,12 +770,24 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box, pub timeout: Option, pub on_failure: FailureAction, + pub depends_on: Vec, } impl StepDefinition { @@ -30,6 +31,7 @@ impl StepDefinition { retry_policy: None, timeout: None, on_failure: FailureAction::FailWorkflow, + depends_on: Vec::new(), } } @@ -47,6 +49,14 @@ impl StepDefinition { self.on_failure = action; self } + + /// Set dependencies for this step. + /// The step will only run after all specified dependencies have completed successfully. + /// Empty slice means no dependencies - step can run immediately in parallel with others. + pub fn depends_on(mut self, deps: &[&str]) -> Self { + self.depends_on = deps.iter().map(|s| StepId::new(*s)).collect(); + self + } } /// Complete workflow definition @@ -95,4 +105,82 @@ impl WorkflowDefinition { pub fn get_timeout(&self, step: &StepDefinition) -> Duration { step.timeout.unwrap_or(self.default_timeout) } + + /// Validate the workflow DAG structure. + /// 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> { + let step_ids: HashSet<_> = self.steps.iter().map(|s| &s.id).collect(); + + // Check all dependencies exist + for step in &self.steps { + for dep in &step.depends_on { + if !step_ids.contains(dep) { + return Err(format!( + "Step '{}' depends on non-existent step '{}'", + step.id, dep + )); + } + } + } + + // Check for cycles using DFS + let mut visited = HashSet::new(); + let mut rec_stack = HashSet::new(); + + for step in &self.steps { + if self.has_cycle(&step.id, &mut visited, &mut rec_stack) { + return Err(format!("Cycle detected involving step '{}'", step.id)); + } + } + + Ok(()) + } + + /// DFS helper for cycle detection + fn has_cycle( + &self, + step_id: &StepId, + visited: &mut HashSet, + rec_stack: &mut HashSet, + ) -> bool { + if rec_stack.contains(step_id) { + return true; // Back edge found - cycle! + } + if visited.contains(step_id) { + return false; // Already fully processed + } + + visited.insert(step_id.clone()); + rec_stack.insert(step_id.clone()); + + // Find the step and check its dependencies + if let Some(step) = self.steps.iter().find(|s| &s.id == step_id) { + for dep in &step.depends_on { + if self.has_cycle(dep, visited, rec_stack) { + return true; + } + } + } + + rec_stack.remove(step_id); + 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 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() + } } diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs index 630b36304..7576ddf9b 100644 --- a/sgl-model-gateway/src/workflow/engine.rs +++ b/sgl-model-gateway/src/workflow/engine.rs @@ -1,11 +1,19 @@ //! Workflow execution engine +//! +//! Supports DAG-based parallel execution of workflow steps. +//! Steps with no dependencies run in parallel, steps with dependencies +//! wait for all dependencies to complete successfully. -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{ + collections::{HashMap, HashSet}, + sync::Arc, + time::Duration, +}; use backoff::{backoff::Backoff, ExponentialBackoffBuilder}; use chrono::Utc; use parking_lot::RwLock; -use tokio::time::timeout; +use tokio::{sync::mpsc, time::timeout}; use super::{ definition::{StepDefinition, WorkflowDefinition}, @@ -14,6 +22,39 @@ use super::{ types::*, }; +/// Consolidated step execution tracking to minimize lock contention. +/// Single lock instead of 4 separate locks. +#[derive(Default)] +struct StepTracker { + completed: HashSet, + failed: HashSet, + skipped: HashSet, + running: HashSet, +} + +impl StepTracker { + fn total_processed(&self) -> usize { + self.completed.len() + self.failed.len() + self.skipped.len() + } + + fn is_step_processable(&self, step_id: &StepId) -> bool { + !self.completed.contains(step_id) + && !self.failed.contains(step_id) + && !self.skipped.contains(step_id) + && !self.running.contains(step_id) + } + + fn are_dependencies_satisfied(&self, depends_on: &[StepId]) -> bool { + depends_on + .iter() + .all(|dep| self.completed.contains(dep) || self.skipped.contains(dep)) + } + + fn has_failed_dependency(&self, depends_on: &[StepId]) -> bool { + depends_on.iter().any(|dep| self.failed.contains(dep)) + } +} + /// Linear backoff implementation that increases delay by a fixed amount each retry struct LinearBackoff { current: Duration, @@ -32,15 +73,15 @@ impl LinearBackoff { } impl Backoff for LinearBackoff { + fn reset(&mut self) { + self.current = self.increment; + } + fn next_backoff(&mut self) -> Option { let next = self.current; self.current = (self.current + self.increment).min(self.max); Some(next) } - - fn reset(&mut self) { - self.current = self.increment; - } } /// Main workflow execution engine @@ -93,9 +134,13 @@ impl WorkflowEngine { } /// Register a workflow definition - pub fn register_workflow(&self, definition: WorkflowDefinition) { + pub fn register_workflow(&self, definition: WorkflowDefinition) -> Result<(), String> { + // Validate DAG once at registration, not on every execution + definition.validate()?; + let id = definition.id.clone(); self.definitions.write().insert(id, Arc::new(definition)); + Ok(()) } /// Get the event bus for subscribing to workflow events @@ -159,15 +204,23 @@ impl WorkflowEngine { Ok(instance_id) } - /// Execute a workflow (internal) + /// Execute a workflow with DAG-based parallel execution async fn execute_workflow( &self, instance_id: WorkflowInstanceId, definition: Arc, ) -> WorkflowResult<()> { let start_time = std::time::Instant::now(); + let step_count = definition.steps.len(); - for step in &definition.steps { + // Consolidated tracking state - single lock instead of 4 + let tracker: Arc> = Arc::new(RwLock::new(StepTracker::default())); + + // Channel for step completion notifications + // Capacity equals step count to prevent blocking on send + let (tx, mut rx) = mpsc::channel::<(StepId, StepResult)>(step_count.max(1)); + + loop { // Check if workflow was cancelled let state = self.state_store.load(instance_id)?; if state.status == WorkflowStatus::Cancelled { @@ -177,67 +230,161 @@ impl WorkflowEngine { return Ok(()); } - // Execute step with retry - match self - .execute_step_with_retry(instance_id, step, &definition) - .await - { - Ok(StepResult::Success) => { - // Continue to next step - } - Ok(StepResult::Skip) => { - // Step was skipped, continue to next - continue; - } - Ok(StepResult::Failure) | Err(_) => { - // Handle failure based on failure action - match step.on_failure { - FailureAction::FailWorkflow => { - let error_msg = format!("Step {} failed", step.id); - self.state_store.update(instance_id, |s| { - s.status = WorkflowStatus::Failed; - })?; + // Single lock acquisition for all tracking state reads + let (ready_step_indices, total_processed, running_count, blocked_by_failure) = { + let t = tracker.read(); - self.event_bus - .publish(WorkflowEvent::WorkflowFailed { - instance_id, - failed_step: step.id.clone(), - error: error_msg, - }) - .await; + let ready: Vec = definition + .steps + .iter() + .enumerate() + .filter(|(_, step)| { + t.is_step_processable(&step.id) + && t.are_dependencies_satisfied(&step.depends_on) + && !t.has_failed_dependency(&step.depends_on) + }) + .map(|(i, _)| i) + .collect(); - return Ok(()); - } - FailureAction::ContinueNextStep => { - // Mark step as skipped and continue - self.state_store.update(instance_id, |s| { - if let Some(step_state) = s.step_states.get_mut(&step.id) { - step_state.status = StepStatus::Skipped; + 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) + }; + + // Check if we're done + if total_processed == step_count { + break; + } + + // 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 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; + } + + // Launch ready steps in parallel + for step_idx in ready_step_indices { + let step = &definition.steps[step_idx]; + tracker.write().running.insert(step.id.clone()); + + let engine = self.clone_for_execution(); + let def = Arc::clone(&definition); + let step_id = step.id.clone(); + let tx = tx.clone(); + let tracker = Arc::clone(&tracker); + + tokio::spawn(async move { + let step = &def.steps[step_idx]; + let result = engine + .execute_step_with_retry(instance_id, step, &def) + .await; + + // Single lock acquisition to update tracking state + { + let mut t = tracker.write(); + t.running.remove(&step_id); + + match result { + Ok(StepResult::Success) => { + t.completed.insert(step_id.clone()); + } + Ok(StepResult::Skip) => { + t.skipped.insert(step_id.clone()); + } + Ok(StepResult::Failure) | Err(_) => match step.on_failure { + FailureAction::FailWorkflow | FailureAction::RetryIndefinitely => { + t.failed.insert(step_id.clone()); } - })?; - continue; - } - FailureAction::RetryIndefinitely => { - // This should not happen as execute_step_with_retry handles it - unreachable!("RetryIndefinitely should be handled in retry logic"); + FailureAction::ContinueNextStep => { + if let Err(e) = engine.state_store.update(instance_id, |s| { + if let Some(step_state) = s.step_states.get_mut(&step_id) { + step_state.status = StepStatus::Skipped; + } + }) { + tracing::warn!( + step_id = %step_id, + error = ?e, + "Failed to update step state to Skipped" + ); + } + t.skipped.insert(step_id.clone()); + } + }, } } + + // Send completion signal (result determines the message) + let signal = match result { + Ok(r) => r, + Err(_) => StepResult::Failure, + }; + let _ = tx.send((step_id, signal)).await; + }); + } + + // Wait for at least one step to complete (if any running) + if !tracker.read().running.is_empty() { + if let Some((completed_step_id, result)) = rx.recv().await { + tracing::debug!( + step_id = %completed_step_id, + result = ?result, + "Step completed" + ); } } } - // Workflow completed successfully - self.state_store.update(instance_id, |s| { - s.status = WorkflowStatus::Completed; - })?; + // Check final status - single lock read + let has_failures = !tracker.read().failed.is_empty(); - let duration = start_time.elapsed(); - self.event_bus - .publish(WorkflowEvent::WorkflowCompleted { - instance_id, - duration, - }) - .await; + if has_failures { + 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("unknown")), + error: "One or more steps failed".to_string(), + }) + .await; + } else { + self.state_store.update(instance_id, |s| { + s.status = WorkflowStatus::Completed; + })?; + + let duration = start_time.elapsed(); + self.event_bus + .publish(WorkflowEvent::WorkflowCompleted { + instance_id, + duration, + }) + .await; + } Ok(()) } @@ -262,12 +409,9 @@ impl WorkflowEngine { let mut backoff = Self::create_backoff(&retry_policy.backoff); loop { - // Check for cancellation before starting/retrying step - { - let state = self.state_store.load(instance_id)?; - if state.status == WorkflowStatus::Cancelled { - return Err(WorkflowError::Cancelled(instance_id)); - } + // Check for cancellation before starting/retrying step (optimized: no full state load) + if self.state_store.is_cancelled(instance_id)? { + return Err(WorkflowError::Cancelled(instance_id)); } // Update step state @@ -293,8 +437,8 @@ impl WorkflowEngine { }) .await; - // Get current context - let mut context = self.state_store.load(instance_id)?.context; + // Get current context (optimized: only clone context, not full state) + let mut context = self.state_store.get_context(instance_id)?; // Execute step with timeout let step_start = std::time::Instant::now(); @@ -302,9 +446,9 @@ impl WorkflowEngine { let step_duration = step_start.elapsed(); - // Save updated context + // Save updated context (use std::mem::replace to avoid extra clone) self.state_store.update(instance_id, |s| { - s.context = context.clone(); + s.context = std::mem::replace(&mut context, WorkflowContext::new(instance_id)); })?; match result { diff --git a/sgl-model-gateway/src/workflow/state.rs b/sgl-model-gateway/src/workflow/state.rs index 1c0530e0b..4755e0850 100644 --- a/sgl-model-gateway/src/workflow/state.rs +++ b/sgl-model-gateway/src/workflow/state.rs @@ -86,6 +86,28 @@ impl WorkflowStateStore { Ok(()) } + /// Get just the workflow context without cloning the entire state. + /// More efficient when you only need the context for step execution. + pub fn get_context( + &self, + instance_id: WorkflowInstanceId, + ) -> WorkflowResult { + self.states + .read() + .get(&instance_id) + .map(|s| s.context.clone()) + .ok_or(WorkflowError::NotFound(instance_id)) + } + + /// Check if workflow is cancelled without loading full state + pub fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult { + self.states + .read() + .get(&instance_id) + .map(|s| s.status == WorkflowStatus::Cancelled) + .ok_or(WorkflowError::NotFound(instance_id)) + } + /// Get count of workflows by status pub fn count_by_status(&self, status: WorkflowStatus) -> usize { self.states diff --git a/sgl-model-gateway/tests/common/mod.rs b/sgl-model-gateway/tests/common/mod.rs index 9f8a8d9c5..de8254048 100644 --- a/sgl-model-gateway/tests/common/mod.rs +++ b/sgl-model-gateway/tests/common/mod.rs @@ -107,8 +107,12 @@ pub async fn create_test_context(config: RouterConfig) -> Arc { workflow::WorkflowEngine, }; let engine = Arc::new(WorkflowEngine::new()); - engine.register_workflow(create_worker_registration_workflow(&config)); - engine.register_workflow(create_worker_removal_workflow()); + engine + .register_workflow(create_worker_registration_workflow(&config)) + .expect("worker_registration workflow should be valid"); + engine + .register_workflow(create_worker_removal_workflow()) + .expect("worker_removal workflow should be valid"); app_context .workflow_engine .set(engine) @@ -239,8 +243,12 @@ pub async fn create_test_context_with_mcp_config( workflow::WorkflowEngine, }; let engine = Arc::new(WorkflowEngine::new()); - engine.register_workflow(create_worker_registration_workflow(&config)); - engine.register_workflow(create_worker_removal_workflow()); + engine + .register_workflow(create_worker_registration_workflow(&config)) + .expect("worker_registration workflow should be valid"); + engine + .register_workflow(create_worker_removal_workflow()) + .expect("worker_removal workflow should be valid"); app_context .workflow_engine .set(engine) diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs index caccd812b..3b9621260 100644 --- a/sgl-model-gateway/tests/wasm_test.rs +++ b/sgl-model-gateway/tests/wasm_test.rs @@ -18,7 +18,14 @@ use axum::{ use sgl_model_gateway::{ app_context::AppContext, config::RouterConfig, - core::steps::{create_wasm_module_registration_workflow, create_wasm_module_removal_workflow}, + core::{ + steps::{create_wasm_module_registration_workflow, create_wasm_module_removal_workflow}, + LoadMonitor, WorkerRegistry, + }, + data_connector::{ + MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, + }, + policies::PolicyRegistry, routers::RouterFactory, server::{build_app, AppState}, wasm::{ @@ -46,15 +53,6 @@ async fn create_test_context_with_wasm() -> Arc { // Create AppContext with wasm_manager from the start let client = reqwest::Client::new(); - // Initialize registries - use sgl_model_gateway::{ - core::{LoadMonitor, WorkerRegistry}, - data_connector::{ - MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, - }, - policies::PolicyRegistry, - }; - let worker_registry = Arc::new(WorkerRegistry::new()); let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone())); @@ -116,10 +114,18 @@ async fn create_test_context_with_wasm() -> Arc { workflow::WorkflowEngine, }; let engine = Arc::new(WorkflowEngine::new()); - engine.register_workflow(create_worker_registration_workflow(&config)); - engine.register_workflow(create_worker_removal_workflow()); - engine.register_workflow(create_wasm_module_registration_workflow()); - engine.register_workflow(create_wasm_module_removal_workflow()); + engine + .register_workflow(create_worker_registration_workflow(&config)) + .expect("worker_registration workflow should be valid"); + engine + .register_workflow(create_worker_removal_workflow()) + .expect("worker_removal workflow should be valid"); + engine + .register_workflow(create_wasm_module_registration_workflow()) + .expect("wasm_module_registration workflow should be valid"); + engine + .register_workflow(create_wasm_module_removal_workflow()) + .expect("wasm_module_removal workflow should be valid"); app_context .workflow_engine .set(engine) diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs index de5c6ca97..f14fa82dc 100644 --- a/sgl-model-gateway/tests/workflow_test.rs +++ b/sgl-model-gateway/tests/workflow_test.rs @@ -70,7 +70,7 @@ async fn test_simple_workflow_execution() { )); let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow); + engine.register_workflow(workflow).unwrap(); // Start workflow let instance_id = engine @@ -115,7 +115,7 @@ async fn test_workflow_with_retry() { ); let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow); + engine.register_workflow(workflow).unwrap(); // Start workflow let instance_id = engine @@ -166,7 +166,7 @@ async fn test_workflow_failure_after_max_retries() { ); let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow); + engine.register_workflow(workflow).unwrap(); // Start workflow let instance_id = engine @@ -223,7 +223,7 @@ async fn test_workflow_continue_on_failure() { )); let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow); + engine.register_workflow(workflow).unwrap(); // Start workflow let instance_id = engine @@ -306,7 +306,7 @@ async fn test_workflow_context_sharing() { )); let workflow_id = workflow.id.clone(); - engine.register_workflow(workflow); + engine.register_workflow(workflow).unwrap(); let instance_id = engine .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new())) @@ -318,3 +318,347 @@ async fn test_workflow_context_sharing() { let state = engine.get_status(instance_id).unwrap(); assert_eq!(state.status, WorkflowStatus::Completed); } + +// ============================================================================ +// DAG / Parallel Execution Tests +// ============================================================================ + +// Step that records when it starts and ends (for testing parallel execution) +struct TimingStep { + step_name: String, + duration_ms: u64, + start_times: Arc>>, + end_times: Arc>>, +} + +#[async_trait::async_trait] +impl StepExecutor for TimingStep { + async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult { + let start = std::time::Instant::now(); + self.start_times + .write() + .push((self.step_name.clone(), start)); + + sleep(Duration::from_millis(self.duration_ms)).await; + + let end = std::time::Instant::now(); + self.end_times.write().push((self.step_name.clone(), end)); + + Ok(StepResult::Success) + } +} + +#[tokio::test] +async fn test_parallel_execution_no_dependencies() { + // Steps without dependencies should run in parallel + let engine = WorkflowEngine::new(); + + let start_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + let end_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + + // Three steps, each taking 100ms, no dependencies + // If parallel: ~100ms total + // If sequential: ~300ms total + let workflow = WorkflowDefinition::new("parallel_workflow", "Parallel Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(TimingStep { + step_name: "step_a".to_string(), + duration_ms: 100, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step(StepDefinition::new( + "step_b", + "Step B", + Arc::new(TimingStep { + step_name: "step_b".to_string(), + duration_ms: 100, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step(StepDefinition::new( + "step_c", + "Step C", + Arc::new(TimingStep { + step_name: "step_c".to_string(), + duration_ms: 100, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let overall_start = std::time::Instant::now(); + let instance_id = engine + .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new())) + .await + .unwrap(); + + // Wait for completion - give enough time for async scheduling + for _ in 0..50 { + sleep(Duration::from_millis(50)).await; + let state = engine.get_status(instance_id).unwrap(); + if state.status != WorkflowStatus::Running { + break; + } + } + + let overall_duration = overall_start.elapsed(); + + let state = engine.get_status(instance_id).unwrap(); + assert_eq!(state.status, WorkflowStatus::Completed); + + // Check that all steps completed + assert_eq!(end_times.read().len(), 3); + + // Verify parallel execution: all steps should start around the same time + let starts = start_times.read(); + let first_start = starts.iter().map(|(_, t)| t).min().unwrap(); + let last_start = starts.iter().map(|(_, t)| t).max().unwrap(); + + // All starts should be within 100ms of each other (allowing for scheduling variance) + let start_spread = last_start.duration_since(*first_start); + assert!( + start_spread < Duration::from_millis(100), + "Steps did not start in parallel, spread: {:?}", + start_spread + ); + + // Total duration should be less than sequential (300ms) - use generous threshold + assert!( + overall_duration < Duration::from_millis(500), + "Parallel execution took too long: {:?}", + overall_duration + ); +} + +#[tokio::test] +async fn test_dag_with_dependencies() { + // DAG: A and B run in parallel, C waits for both + // A ──┐ + // ├──> C + // B ──┘ + let engine = WorkflowEngine::new(); + + let start_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + let end_times: Arc>> = + Arc::new(parking_lot::RwLock::new(Vec::new())); + + let workflow = WorkflowDefinition::new("dag_workflow", "DAG Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(TimingStep { + step_name: "step_a".to_string(), + duration_ms: 50, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step(StepDefinition::new( + "step_b", + "Step B", + Arc::new(TimingStep { + step_name: "step_b".to_string(), + duration_ms: 100, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + )) + .add_step( + StepDefinition::new( + "step_c", + "Step C", + Arc::new(TimingStep { + step_name: "step_c".to_string(), + duration_ms: 50, + start_times: Arc::clone(&start_times), + end_times: Arc::clone(&end_times), + }), + ) + .depends_on(&["step_a", "step_b"]), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new())) + .await + .unwrap(); + + // Poll until workflow completes (or timeout) + for _ in 0..50 { + sleep(Duration::from_millis(50)).await; + let state = engine.get_status(instance_id).unwrap(); + if state.status != WorkflowStatus::Running { + break; + } + } + + let state = engine.get_status(instance_id).unwrap(); + assert_eq!(state.status, WorkflowStatus::Completed); + + // Verify step C started after both A and B finished + let starts = start_times.read(); + let ends = end_times.read(); + + let c_start = starts.iter().find(|(n, _)| n == "step_c").unwrap().1; + let a_end = ends.iter().find(|(n, _)| n == "step_a").unwrap().1; + let b_end = ends.iter().find(|(n, _)| n == "step_b").unwrap().1; + + assert!(c_start >= a_end, "Step C started before Step A finished"); + assert!(c_start >= b_end, "Step C started before Step B finished"); +} + +#[tokio::test] +async fn test_dag_dependency_failure_blocks_dependents() { + // If step A fails with FailWorkflow, step B (depends on A) should not run + let engine = WorkflowEngine::new(); + + let b_executed = Arc::new(AtomicU32::new(0)); + + struct FailingStep; + + #[async_trait::async_trait] + impl StepExecutor for FailingStep { + async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult { + Err(WorkflowError::StepFailed { + step_id: StepId::new("failing"), + message: "Intentional failure".to_string(), + }) + } + + fn is_retryable(&self, _error: &WorkflowError) -> bool { + false // Disable retries for this test + } + } + + struct TrackingStep { + counter: Arc, + } + + #[async_trait::async_trait] + impl StepExecutor for TrackingStep { + async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult { + self.counter.fetch_add(1, Ordering::SeqCst); + Ok(StepResult::Success) + } + } + + let workflow = WorkflowDefinition::new("blocked_workflow", "Blocked Test") + .add_step( + StepDefinition::new("step_a", "Step A", Arc::new(FailingStep)) + .with_failure_action(FailureAction::FailWorkflow), + ) + .add_step( + StepDefinition::new( + "step_b", + "Step B", + Arc::new(TrackingStep { + counter: Arc::clone(&b_executed), + }), + ) + .depends_on(&["step_a"]), + ); + + let workflow_id = workflow.id.clone(); + engine.register_workflow(workflow).unwrap(); + + let instance_id = engine + .start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new())) + .await + .unwrap(); + + // Poll until workflow completes (or timeout) + for _ in 0..50 { + sleep(Duration::from_millis(50)).await; + let state = engine.get_status(instance_id).unwrap(); + if state.status != WorkflowStatus::Running { + break; + } + } + + let state = engine.get_status(instance_id).unwrap(); + assert_eq!(state.status, WorkflowStatus::Failed); + + // Step B should not have executed + assert_eq!(b_executed.load(Ordering::SeqCst), 0); +} + +#[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") + .add_step( + StepDefinition::new("step_a", "Step A", Arc::new(AlwaysSucceedStep)) + .depends_on(&["step_c"]), + ) + .add_step( + StepDefinition::new("step_b", "Step B", Arc::new(AlwaysSucceedStep)) + .depends_on(&["step_a"]), + ) + .add_step( + StepDefinition::new("step_c", "Step C", Arc::new(AlwaysSucceedStep)) + .depends_on(&["step_b"]), + ); + + let result = workflow.validate(); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("Cycle detected")); +} + +#[test] +fn test_dag_validation_missing_dependency() { + // Create a workflow with a missing dependency + let workflow = WorkflowDefinition::new("missing_dep_workflow", "Missing Dep Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(AlwaysSucceedStep), + )) + .add_step( + StepDefinition::new("step_b", "Step B", Arc::new(AlwaysSucceedStep)) + .depends_on(&["nonexistent_step"]), + ); + + let result = workflow.validate(); + assert!(result.is_err()); + assert!(result.unwrap_err().contains("non-existent step")); +} + +#[test] +fn test_dag_validation_valid_workflow() { + // Create a valid DAG workflow + let workflow = WorkflowDefinition::new("valid_workflow", "Valid Test") + .add_step(StepDefinition::new( + "step_a", + "Step A", + Arc::new(AlwaysSucceedStep), + )) + .add_step(StepDefinition::new( + "step_b", + "Step B", + Arc::new(AlwaysSucceedStep), + )) + .add_step( + StepDefinition::new("step_c", "Step C", Arc::new(AlwaysSucceedStep)) + .depends_on(&["step_a", "step_b"]), + ) + .add_step( + StepDefinition::new("step_d", "Step D", Arc::new(AlwaysSucceedStep)) + .depends_on(&["step_c"]), + ); + + let result = workflow.validate(); + assert!(result.is_ok()); +}