[model-gateway] feat: add DAG parallel execution support and workflow optimization (#14999)

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"]),
     )
 }

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"]),
         )
 }

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"]),
         )
 }

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"]),
         )
 }

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"]),
         )
 }

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"]),
         )
 }

sgl-model-gateway/src/server.rs

@@ -770,12 +770,24 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
         .subscribe(Arc::new(LoggingSubscriber))
         .await;
 
-    engine.register_workflow(create_worker_registration_workflow(&config.router_config));
+    engine
-    engine.register_workflow(create_external_worker_registration_workflow());
+        .register_workflow(create_worker_registration_workflow(&config.router_config))
-    engine.register_workflow(create_worker_removal_workflow());
+        .expect("worker_registration workflow should be valid");
-    engine.register_workflow(create_mcp_registration_workflow());
+    engine
-    engine.register_workflow(create_wasm_module_registration_workflow());
+        .register_workflow(create_external_worker_registration_workflow())
-    engine.register_workflow(create_wasm_module_removal_workflow());
+        .expect("external_worker_registration workflow should be valid");
+    engine
+        .register_workflow(create_worker_removal_workflow())
+        .expect("worker_removal workflow should be valid");
+    engine
+        .register_workflow(create_mcp_registration_workflow())
+        .expect("mcp_registration 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)

sgl-model-gateway/src/workflow/definition.rs

@@ -1,6 +1,6 @@
 //! Workflow definition types
 
-use std::{sync::Arc, time::Duration};
+use std::{collections::HashSet, sync::Arc, time::Duration};
 
 use super::{
     executor::StepExecutor,
@@ -15,6 +15,7 @@ pub struct StepDefinition {
     pub retry_policy: Option<RetryPolicy>,
     pub timeout: Option<Duration>,
     pub on_failure: FailureAction,
+    pub depends_on: Vec<StepId>,
 }
 
 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<StepId>,
+        rec_stack: &mut HashSet<StepId>,
+    ) -> 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()
+    }
 }

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<StepId>,
+    failed: HashSet<StepId>,
+    skipped: HashSet<StepId>,
+    running: HashSet<StepId>,
+}
+
+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<Duration> {
         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<WorkflowDefinition>,
     ) -> 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<RwLock<StepTracker>> = 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,56 +230,149 @@ impl WorkflowEngine {
                 return Ok(());
             }
 
-            // Execute step with retry
+            // Single lock acquisition for all tracking state reads
-            match self
+            let (ready_step_indices, total_processed, running_count, blocked_by_failure) = {
-                .execute_step_with_retry(instance_id, step, &definition)
+                let t = tracker.read();
-                .await
+
-            {
+                let ready: Vec<usize> = definition
-                Ok(StepResult::Success) => {
+                    .steps
-                    // Continue to next step
+                    .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();
+
+                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;
             }
-                Ok(StepResult::Skip) => {
+
-                    // Step was skipped, continue to next
+            // Handle blocked workflow
-                    continue;
+            if ready_step_indices.is_empty() && running_count == 0 {
-                }
+                if blocked_by_failure {
-                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;
                     })?;
 
+                    let failed_step = tracker.read().failed.iter().next().cloned();
                     self.event_bus
                         .publish(WorkflowEvent::WorkflowFailed {
                             instance_id,
-                                    failed_step: step.id.clone(),
+                            failed_step: failed_step.unwrap_or_else(|| StepId::new("unknown")),
-                                    error: error_msg,
+                            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());
+                                }
                                 FailureAction::ContinueNextStep => {
-                            // Mark step as skipped and continue
+                                    if let Err(e) = engine.state_store.update(instance_id, |s| {
-                            self.state_store.update(instance_id, |s| {
+                                        if let Some(step_state) = s.step_states.get_mut(&step_id) {
-                                if let Some(step_state) = s.step_states.get_mut(&step.id) {
                                             step_state.status = StepStatus::Skipped;
                                         }
-                            })?;
+                                    }) {
-                            continue;
+                                        tracing::warn!(
+                                            step_id = %step_id,
+                                            error = ?e,
+                                            "Failed to update step state to Skipped"
+                                        );
                                     }
-                        FailureAction::RetryIndefinitely => {
+                                    t.skipped.insert(step_id.clone());
-                            // This should not happen as execute_step_with_retry handles it
+                                }
-                            unreachable!("RetryIndefinitely should be handled in retry logic");
+                            },
                         }
                     }
+
+                    // 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
+        // Check final status - single lock read
+        let has_failures = !tracker.read().failed.is_empty();
+
+        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;
             })?;
@@ -238,6 +384,7 @@ impl WorkflowEngine {
                     duration,
                 })
                 .await;
+        }
 
         Ok(())
     }
@@ -262,13 +409,10 @@ impl WorkflowEngine {
         let mut backoff = Self::create_backoff(&retry_policy.backoff);
 
         loop {
-            // Check for cancellation before starting/retrying step
+            // Check for cancellation before starting/retrying step (optimized: no full state load)
-            {
+            if self.state_store.is_cancelled(instance_id)? {
-                let state = self.state_store.load(instance_id)?;
-                if state.status == WorkflowStatus::Cancelled {
                 return Err(WorkflowError::Cancelled(instance_id));
             }
-            }
 
             // Update step state
             self.state_store.update(instance_id, |s| {
@@ -293,8 +437,8 @@ impl WorkflowEngine {
                 })
                 .await;
 
-            // Get current context
+            // Get current context (optimized: only clone context, not full state)
-            let mut context = self.state_store.load(instance_id)?.context;
+            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 {

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<super::types::WorkflowContext> {
+        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<bool> {
+        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

sgl-model-gateway/tests/common/mod.rs

@@ -107,8 +107,12 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
         workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
-    engine.register_workflow(create_worker_registration_workflow(&config));
+    engine
-    engine.register_workflow(create_worker_removal_workflow());
+        .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
-    engine.register_workflow(create_worker_removal_workflow());
+        .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)

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<AppContext> {
     // 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<AppContext> {
         workflow::WorkflowEngine,
     };
     let engine = Arc::new(WorkflowEngine::new());
-    engine.register_workflow(create_worker_registration_workflow(&config));
+    engine
-    engine.register_workflow(create_worker_removal_workflow());
+        .register_workflow(create_worker_registration_workflow(&config))
-    engine.register_workflow(create_wasm_module_registration_workflow());
+        .expect("worker_registration workflow should be valid");
-    engine.register_workflow(create_wasm_module_removal_workflow());
+    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)

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<parking_lot::RwLock<Vec<(String, std::time::Instant)>>>,
+    end_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>>,
+}
+
+#[async_trait::async_trait]
+impl StepExecutor for TimingStep {
+    async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+        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<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
+        Arc::new(parking_lot::RwLock::new(Vec::new()));
+    let end_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
+        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<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
+        Arc::new(parking_lot::RwLock::new(Vec::new()));
+    let end_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
+        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<StepResult> {
+            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<AtomicU32>,
+    }
+
+    #[async_trait::async_trait]
+    impl StepExecutor for TrackingStep {
+        async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
+            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());
+}