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

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

Browse Source
This commit is contained in:
Simo Lin
2025-12-12 08:31:12 -08:00
committed by GitHub
parent 10c68f6236
commit 306e5b8d0b
13 changed files with 768 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
sgl-model-gateway/src/core/steps/external_worker_registration.rs
View File

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

19
sgl-model-gateway/src/core/steps/mcp_registration.rs
View File

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

27
sgl-model-gateway/src/core/steps/wasm_module_registration.rs
View File

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

11
sgl-model-gateway/src/core/steps/wasm_module_removal.rs
View File

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

32
sgl-model-gateway/src/core/steps/worker_registration.rs
View File

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

20
sgl-model-gateway/src/core/steps/worker_removal.rs
View File

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

24
sgl-model-gateway/src/server.rs
View File

@@ -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.register_workflow(create_external_worker_registration_workflow());
engine.register_workflow(create_worker_removal_workflow());
engine.register_workflow(create_mcp_registration_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.router_config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_external_worker_registration_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)

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

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

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

@@ -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,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<usize> = 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 {

22
sgl-model-gateway/src/workflow/state.rs
View File

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

16
sgl-model-gateway/tests/common/mod.rs
View File

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

34
sgl-model-gateway/tests/wasm_test.rs
View File

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

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

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