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

[model-gateway] convert workflow system to type-safe workflow data (#16970)

Browse Source
This commit is contained in:
Simo Lin
2026-01-12 10:06:36 -08:00
committed by GitHub
parent 559ff9ecaf
commit fa51b85466
38 changed files with 1789 additions and 746 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/app_context.rs
+22 -5
View File
@@ -8,7 +8,10 @@ use tracing::{debug, info};
use crate::{
config::RouterConfig,
core::{JobQueue, LoadMonitor, WorkerRegistry, WorkerService, UNKNOWN_MODEL_ID},
core::{
steps::workflow_data::AnyWorkflowData, JobQueue, LoadMonitor, WorkerRegistry,
WorkerService, UNKNOWN_MODEL_ID,
},
data_connector::{
create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage,
},
@@ -26,9 +29,12 @@ use crate::{
},
tool_parser::ParserFactory as ToolParserFactory,
wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
workflow::WorkflowEngine,
workflow::{InMemoryStore, WorkflowEngine},
};
/// Type alias for the concrete workflow engine used in the application
pub type AppWorkflowEngine = WorkflowEngine<AnyWorkflowData, InMemoryStore<AnyWorkflowData>>;
/// Error type for AppContext builder
#[derive(Debug)]
pub struct AppContextBuildError(&'static str);
@@ -59,13 +65,21 @@ pub struct AppContext {
pub configured_reasoning_parser: Option<String>,
pub configured_tool_parser: Option<String>,
pub worker_job_queue: Arc<OnceLock<Arc<JobQueue>>>,
pub workflow_engine: Arc<OnceLock<Arc<WorkflowEngine>>>,
pub workflow_engine: Arc<OnceLock<Arc<AppWorkflowEngine>>>,
pub mcp_manager: Arc<OnceLock<Arc<McpManager>>>,
pub wasm_manager: Option<Arc<WasmModuleManager>>,
pub worker_service: Arc<WorkerService>,
pub inflight_tracker: Arc<InFlightRequestTracker>,
}
impl std::fmt::Debug for AppContext {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AppContext")
.field("router_config", &self.router_config)
.finish_non_exhaustive()
}
}
pub struct AppContextBuilder {
client: Option<Client>,
router_config: Option<RouterConfig>,
@@ -81,7 +95,7 @@ pub struct AppContextBuilder {
conversation_item_storage: Option<Arc<dyn ConversationItemStorage>>,
load_monitor: Option<Arc<LoadMonitor>>,
worker_job_queue: Option<Arc<OnceLock<Arc<JobQueue>>>>,
workflow_engine: Option<Arc<OnceLock<Arc<WorkflowEngine>>>>,
workflow_engine: Option<Arc<OnceLock<Arc<AppWorkflowEngine>>>>,
mcp_manager: Option<Arc<OnceLock<Arc<McpManager>>>>,
wasm_manager: Option<Arc<WasmModuleManager>>,
}
@@ -206,7 +220,10 @@ impl AppContextBuilder {
self
}
pub fn workflow_engine(mut self, workflow_engine: Arc<OnceLock<Arc<WorkflowEngine>>>) -> Self {
pub fn workflow_engine(
mut self,
workflow_engine: Arc<OnceLock<Arc<AppWorkflowEngine>>>,
) -> Self {
self.workflow_engine = Some(workflow_engine);
self
}
sgl-model-gateway/src/core/job_queue.rs
+45 -56
View File
@@ -14,15 +14,19 @@ use tokio::sync::{mpsc, Semaphore};
use tracing::{debug, error, info, warn};
use crate::{
app_context::AppContext,
app_context::{AppContext, AppWorkflowEngine},
config::{RouterConfig, RoutingMode},
core::steps::{
create_external_worker_workflow_data, create_local_worker_workflow_data,
create_mcp_workflow_data, create_tokenizer_workflow_data,
create_wasm_registration_workflow_data, create_wasm_removal_workflow_data,
create_worker_removal_workflow_data, create_worker_update_workflow_data,
McpServerConfigRequest, TokenizerConfigRequest, TokenizerRemovalRequest,
WasmModuleConfigRequest, WasmModuleRemovalRequest, WorkerRemovalRequest,
WasmModuleConfigRequest, WasmModuleRemovalRequest,
},
mcp::McpConfig,
protocols::worker_spec::{JobStatus, WorkerConfigRequest, WorkerUpdateRequest},
workflow::{WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus},
workflow::{WorkflowId, WorkflowInstanceId, WorkflowStatus},
};
/// Job types for control plane operations
@@ -404,17 +408,11 @@ impl JobQueue {
.get()
.ok_or_else(|| "Workflow engine not initialized".to_string())?;
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
// Convert Box to Arc for context storage
let config_arc: Arc<WasmModuleConfigRequest> = Arc::new(*config.clone());
workflow_context.set_arc("wasm_module_config", config_arc);
workflow_context.set_arc("app_context", Arc::clone(context));
let workflow_data =
create_wasm_registration_workflow_data(*config.clone(), Arc::clone(context));
let instance_id = engine
.start_workflow(
WorkflowId::new("wasm_module_registration"),
workflow_context,
)
.start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
.await
.map_err(|e| {
format!("Failed to start WASM module registration workflow: {:?}", e)
@@ -441,14 +439,11 @@ impl JobQueue {
.get()
.ok_or_else(|| "Workflow engine not initialized".to_string())?;
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
// Convert Box to Arc for context storage
let request_arc: Arc<WasmModuleRemovalRequest> = Arc::new(*request.clone());
workflow_context.set_arc("wasm_module_removal_request", request_arc);
workflow_context.set_arc("app_context", Arc::clone(context));
let workflow_data =
create_wasm_removal_workflow_data(*request.clone(), Arc::clone(context));
let instance_id = engine
.start_workflow(WorkflowId::new("wasm_module_removal"), workflow_context)
.start_workflow(WorkflowId::new("wasm_module_removal"), workflow_data)
.await
.map_err(|e| {
format!("Failed to start WASM module removal workflow: {:?}", e)
@@ -674,13 +669,11 @@ impl JobQueue {
.get()
.ok_or_else(|| "Workflow engine not initialized".to_string())?;
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
let config_arc: Arc<TokenizerConfigRequest> = Arc::new(*config.clone());
workflow_context.set_arc("tokenizer_config", config_arc);
workflow_context.set_arc("app_context", Arc::clone(context));
let workflow_data =
create_tokenizer_workflow_data(*config.clone(), Arc::clone(context));
let instance_id = engine
.start_workflow(WorkflowId::new("tokenizer_registration"), workflow_context)
.start_workflow(WorkflowId::new("tokenizer_registration"), workflow_data)
.await
.map_err(|e| {
format!("Failed to start tokenizer registration workflow: {:?}", e)
@@ -719,86 +712,82 @@ impl JobQueue {
/// Start a workflow and return its instance ID
async fn start_worker_workflow(
engine: &Arc<WorkflowEngine>,
engine: &Arc<AppWorkflowEngine>,
config: &WorkerConfigRequest,
context: &Arc<AppContext>,
) -> Result<WorkflowInstanceId, String> {
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
workflow_context.set("worker_config", config.clone());
workflow_context.set_arc("app_context", Arc::clone(context));
// Select workflow based on runtime field
let workflow_id = match config.runtime.as_deref() {
Some("external") => WorkflowId::new("external_worker_registration"),
_ => WorkflowId::new("local_worker_registration"),
let (workflow_id, workflow_data) = match config.runtime.as_deref() {
Some("external") => (
WorkflowId::new("external_worker_registration"),
create_external_worker_workflow_data(config.clone(), Arc::clone(context)),
),
_ => (
WorkflowId::new("local_worker_registration"),
create_local_worker_workflow_data(config.clone(), Arc::clone(context)),
),
};
engine
.start_workflow(workflow_id, workflow_context)
.start_workflow(workflow_id, workflow_data)
.await
.map_err(|e| format!("Failed to start worker registration workflow: {:?}", e))
}
/// Start worker removal workflow
async fn start_worker_removal_workflow(
engine: &Arc<WorkflowEngine>,
engine: &Arc<AppWorkflowEngine>,
url: &str,
context: &Arc<AppContext>,
) -> Result<WorkflowInstanceId, String> {
let removal_request = WorkerRemovalRequest {
url: url.to_string(),
dp_aware: context.router_config.dp_aware,
};
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
workflow_context.set("removal_request", removal_request);
workflow_context.set_arc("app_context", Arc::clone(context));
let workflow_data = create_worker_removal_workflow_data(
url.to_string(),
context.router_config.dp_aware,
Arc::clone(context),
);
engine
.start_workflow(WorkflowId::new("worker_removal"), workflow_context)
.start_workflow(WorkflowId::new("worker_removal"), workflow_data)
.await
.map_err(|e| format!("Failed to start worker removal workflow: {:?}", e))
}
/// Start worker update workflow
async fn start_worker_update_workflow(
engine: &Arc<WorkflowEngine>,
engine: &Arc<AppWorkflowEngine>,
url: &str,
update: &WorkerUpdateRequest,
context: &Arc<AppContext>,
) -> Result<WorkflowInstanceId, String> {
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
// Pass URL and dp_aware separately, workflow step handles the rest
workflow_context.set("worker_url", url.to_string());
workflow_context.set("dp_aware", context.router_config.dp_aware);
workflow_context.set("update_request", update.clone());
workflow_context.set_arc("app_context", Arc::clone(context));
let workflow_data = create_worker_update_workflow_data(
url.to_string(),
update.clone(),
Arc::clone(context),
);
engine
.start_workflow(WorkflowId::new("worker_update"), workflow_context)
.start_workflow(WorkflowId::new("worker_update"), workflow_data)
.await
.map_err(|e| format!("Failed to start worker update workflow: {:?}", e))
}
/// Start MCP server registration workflow
async fn start_mcp_registration_workflow(
engine: &Arc<WorkflowEngine>,
engine: &Arc<AppWorkflowEngine>,
config: &McpServerConfigRequest,
context: &Arc<AppContext>,
) -> Result<WorkflowInstanceId, String> {
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
workflow_context.set("mcp_server_config", config.clone());
workflow_context.set_arc("app_context", Arc::clone(context));
let workflow_data = create_mcp_workflow_data(config.clone(), Arc::clone(context));
engine
.start_workflow(WorkflowId::new("mcp_registration"), workflow_context)
.start_workflow(WorkflowId::new("mcp_registration"), workflow_data)
.await
.map_err(|e| format!("Failed to start MCP registration workflow: {:?}", e))
}
/// Wait for workflow completion with adaptive polling
async fn wait_for_workflow_completion(
engine: &Arc<WorkflowEngine>,
engine: &Arc<AppWorkflowEngine>,
instance_id: WorkflowInstanceId,
worker_url: &str,
timeout_duration: Duration,
sgl-model-gateway/src/core/steps/mcp_registration.rs
+79 -31
View File
@@ -1,18 +1,21 @@
use std::{sync::Arc, time::Duration};
use async_trait::async_trait;
use rmcp::{service::RunningService, RoleClient};
use tracing::{debug, error, info, warn};
use super::workflow_data::{AnyWorkflowData, McpWorkflowData};
use crate::{
app_context::AppContext,
mcp::{config::McpServerConfig, manager::McpManager},
observability::metrics::Metrics,
workflow::*,
workflow::{
BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId,
StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
},
};
/// MCP server connection configuration
#[derive(Debug, Clone)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct McpServerConfigRequest {
/// Server name (unique identifier)
pub name: String,
@@ -35,11 +38,17 @@ impl McpServerConfigRequest {
pub struct ConnectMcpServerStep;
#[async_trait]
impl StepExecutor for ConnectMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for ConnectMcpServerStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_mcp()?;
let config_request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!("Connecting to MCP server: {}", config_request.name);
@@ -66,8 +75,9 @@ impl StepExecutor for ConnectMcpServerStep {
config_request.name
);
// Store client in context (context.set() will wrap in Arc)
context.set("mcp_client", client);
// Store client in typed data
let data_mut = context.data.as_mcp_mut()?;
data_mut.mcp_client = Some(Arc::new(client));
Ok(StepResult::Success)
}
@@ -86,12 +96,21 @@ impl StepExecutor for ConnectMcpServerStep {
pub struct DiscoverMcpInventoryStep;
#[async_trait]
impl StepExecutor for DiscoverMcpInventoryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
impl StepExecutor<AnyWorkflowData> for DiscoverMcpInventoryStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_mcp()?;
let config_request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let mcp_client = data
.mcp_client
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?;
debug!(
"Discovering inventory for MCP server: {}",
@@ -111,7 +130,7 @@ impl StepExecutor for DiscoverMcpInventoryStep {
let inventory = mcp_manager.inventory();
// Use the public load_server_inventory method
McpManager::load_server_inventory(&inventory, &config_request.name, &mcp_client).await;
McpManager::load_server_inventory(&inventory, &config_request.name, mcp_client).await;
info!("Completed inventory discovery for {}", config_request.name);
@@ -130,12 +149,22 @@ impl StepExecutor for DiscoverMcpInventoryStep {
pub struct RegisterMcpServerStep;
#[async_trait]
impl StepExecutor for RegisterMcpServerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let mcp_client: Arc<RunningService<RoleClient, ()>> = context.get_or_err("mcp_client")?;
impl StepExecutor<AnyWorkflowData> for RegisterMcpServerStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_mcp()?;
let config_request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let mcp_client = data
.mcp_client
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("mcp_client".to_string()))?
.clone();
debug!("Registering MCP server: {}", config_request.name);
@@ -174,20 +203,26 @@ impl StepExecutor for RegisterMcpServerStep {
pub struct ValidateRegistrationStep;
#[async_trait]
impl StepExecutor for ValidateRegistrationStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<McpServerConfigRequest> =
context.get_or_err("mcp_server_config")?;
impl StepExecutor<AnyWorkflowData> for ValidateRegistrationStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_mcp()?;
let config_request = &data.config;
let client_registered = context
.get::<RunningService<RoleClient, ()>>("mcp_client")
.is_some();
let client_registered = data.mcp_client.is_some();
if client_registered {
info!(
"MCP server '{}' registered successfully",
config_request.name
);
// Mark as validated
let data_mut = context.data.as_mcp_mut()?;
data_mut.validated = true;
return Ok(StepResult::Success);
}
@@ -228,7 +263,7 @@ impl StepExecutor for ValidateRegistrationStep {
/// - DiscoverMcpInventory: 3 retries, 10s timeout (discovery + caching)
/// - RegisterMcpServer: No retry, 5s timeout (fast registration)
/// - ValidateRegistration: Final validation step
pub fn create_mcp_registration_workflow() -> WorkflowDefinition {
pub fn create_mcp_registration_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new("mcp_registration", "MCP Server Registration")
.add_step(
StepDefinition::new(
@@ -281,3 +316,16 @@ pub fn create_mcp_registration_workflow() -> WorkflowDefinition {
.depends_on(&["register_mcp_server"]),
)
}
/// Helper to create initial workflow data for MCP registration
pub fn create_mcp_workflow_data(
config: McpServerConfigRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::Mcp(McpWorkflowData {
config,
validated: false,
app_context: Some(app_context),
mcp_client: None,
})
}
sgl-model-gateway/src/core/steps/mod.rs
+30 -30
View File
@@ -11,19 +11,38 @@ pub mod tokenizer_registration;
pub mod wasm_module_registration;
pub mod wasm_module_removal;
pub mod worker;
pub mod workflow_data;
// Worker management (registration, removal)
#[allow(deprecated)]
pub use worker::create_external_worker_registration_workflow;
// Backward compatibility aliases
#[allow(deprecated)]
pub use worker::create_worker_registration_workflow;
pub use mcp_registration::{
create_mcp_registration_workflow, create_mcp_workflow_data, ConnectMcpServerStep,
DiscoverMcpInventoryStep, McpServerConfigRequest, RegisterMcpServerStep,
ValidateRegistrationStep,
};
pub use tokenizer_registration::{
create_tokenizer_registration_workflow, create_tokenizer_workflow_data, LoadTokenizerStep,
TokenizerConfigRequest, TokenizerRemovalRequest, ValidateTokenizerConfigStep,
};
pub use wasm_module_registration::{
create_wasm_module_registration_workflow, create_wasm_registration_workflow_data,
CalculateHashStep, CheckDuplicateStep, LoadWasmBytesStep, RegisterModuleStep,
ValidateDescriptorStep, ValidateWasmComponentStep, WasmModuleConfigRequest,
};
pub use wasm_module_removal::{
create_wasm_module_removal_workflow, create_wasm_removal_workflow_data, FindModuleToRemoveStep,
RemoveModuleStep, WasmModuleRemovalRequest,
};
pub use worker::{
// Workflow builders
create_external_worker_workflow,
// Workflow data helpers
create_external_worker_workflow_data,
create_local_worker_workflow,
create_local_worker_workflow_data,
create_worker_removal_workflow,
create_worker_removal_workflow_data,
create_worker_update_workflow,
create_worker_update_workflow_data,
// Utility functions
group_models_into_cards,
infer_model_type_from_id,
@@ -54,29 +73,10 @@ pub use worker::{
WorkerList,
WorkerRemovalRequest,
};
// Legacy type aliases for backward compatibility
pub type ActivateWorkerStep = ActivateWorkersStep;
pub type RegisterWorkerStep = RegisterWorkersStep;
pub type CreateWorkerStep = CreateLocalWorkerStep;
pub type ActivateExternalWorkersStep = ActivateWorkersStep;
pub type RegisterExternalWorkersStep = RegisterWorkersStep;
pub type UpdateExternalPoliciesStep = UpdatePoliciesStep;
pub use mcp_registration::{
create_mcp_registration_workflow, ConnectMcpServerStep, DiscoverMcpInventoryStep,
McpServerConfigRequest, RegisterMcpServerStep, ValidateRegistrationStep,
};
pub use tokenizer_registration::{
create_tokenizer_registration_workflow, LoadTokenizerStep, TokenizerConfigRequest,
TokenizerRemovalRequest, ValidateTokenizerConfigStep,
};
pub use wasm_module_registration::{
create_wasm_module_registration_workflow, CalculateHashStep, CheckDuplicateStep,
LoadWasmBytesStep, RegisterModuleStep, ValidateDescriptorStep, ValidateWasmComponentStep,
WasmModuleConfigRequest,
};
pub use wasm_module_removal::{
create_wasm_module_removal_workflow, FindModuleToRemoveStep, RemoveModuleStep,
WasmModuleRemovalRequest,
// Typed workflow data structures
pub use workflow_data::{
AnyWorkflowData, ExternalWorkerWorkflowData, LocalWorkerWorkflowData, McpWorkflowData,
ProtocolUpdateRequest, TokenizerWorkflowData, WasmRegistrationWorkflowData,
WasmRemovalWorkflowData, WorkerConfigRequest, WorkerList as WorkflowWorkerList,
WorkerRemovalWorkflowData, WorkerUpdateWorkflowData,
};
sgl-model-gateway/src/core/steps/tokenizer_registration.rs
+64 -18
View File
@@ -9,7 +9,15 @@ use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tracing::{debug, error, info};
use crate::{app_context::AppContext, tokenizer::factory, workflow::*};
use super::workflow_data::{AnyWorkflowData, TokenizerWorkflowData};
use crate::{
app_context::AppContext,
tokenizer::factory,
workflow::{
BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId,
StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
},
};
/// Configuration for adding a tokenizer
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -39,10 +47,17 @@ pub struct TokenizerRemovalRequest {
pub struct ValidateTokenizerConfigStep;
#[async_trait]
impl StepExecutor for ValidateTokenizerConfigStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for ValidateTokenizerConfigStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_tokenizer()?;
let config = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!(
"Validating tokenizer config: name={}, source={}",
@@ -86,22 +101,36 @@ impl StepExecutor for ValidateTokenizerConfigStep {
pub struct LoadTokenizerStep;
#[async_trait]
impl StepExecutor for LoadTokenizerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<TokenizerConfigRequest> = context.get_or_err("tokenizer_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for LoadTokenizerStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_tokenizer()?;
let config = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
.clone();
info!(
"Loading tokenizer '{}' (id: {}) from source: {}",
config.name, config.id, config.source
);
// Clone needed values before async move
let id = config.id.clone();
let name = config.name.clone();
let source = config.source.clone();
let chat_template = config.chat_template_path.clone();
// Load the tokenizer using the registry's load method (handles deduplication)
let result = app_context
.tokenizer_registry
.load(&config.id, &config.name, &config.source, || {
let source = config.source.clone();
let chat_template = config.chat_template_path.clone();
.load(&id, &name, &source, || {
let source = source.clone();
let chat_template = chat_template.clone();
async move {
factory::create_tokenizer_async_with_chat_template(
&source,
@@ -123,18 +152,19 @@ impl StepExecutor for LoadTokenizerStep {
info!(
"Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
config.name, loaded_id, vocab_size
name, loaded_id, vocab_size
);
// Store vocab size in context for later use
// Store vocab size in typed data
if let Some(size) = vocab_size {
context.set("vocab_size", size);
let data_mut = context.data.as_tokenizer_mut()?;
data_mut.vocab_size = Some(size);
}
Ok(StepResult::Success)
}
Err(e) => {
error!("Failed to load tokenizer '{}': {}", config.name, e);
error!("Failed to load tokenizer '{}': {}", name, e);
Err(WorkflowError::StepFailed {
step_id: StepId::new("load_tokenizer"),
message: e,
@@ -161,7 +191,7 @@ impl StepExecutor for LoadTokenizerStep {
/// Workflow configuration:
/// - ValidateConfig: No retry, 5s timeout (fast validation)
/// - LoadTokenizer: 3 retries, 5min timeout (may need to download from HuggingFace)
pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition {
pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration")
.add_step(
StepDefinition::new(
@@ -188,6 +218,18 @@ pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition {
)
}
/// Helper to create initial workflow data for tokenizer registration
pub fn create_tokenizer_workflow_data(
config: TokenizerConfigRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::Tokenizer(TokenizerWorkflowData {
config,
vocab_size: None,
app_context: Some(app_context),
})
}
#[cfg(test)]
mod tests {
use super::*;
@@ -212,7 +254,11 @@ mod tests {
#[test]
fn test_workflow_creation() {
let workflow = create_tokenizer_registration_workflow();
let mut workflow = create_tokenizer_registration_workflow();
assert_eq!(workflow.id.to_string(), "tokenizer_registration");
// Validate the workflow DAG
workflow
.validate()
.expect("Workflow validation should pass");
}
}
sgl-model-gateway/src/core/steps/wasm_module_registration.rs
+131 -68
View File
@@ -10,14 +10,18 @@ use tracing::{debug, info, warn};
use uuid::Uuid;
use wasmtime::{component::Component, Config, Engine};
use super::workflow_data::{AnyWorkflowData, WasmRegistrationWorkflowData};
use crate::{
app_context::AppContext,
wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta},
workflow::*,
workflow::{
BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId,
StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult,
},
};
/// WASM module registration request
#[derive(Debug, Clone)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WasmModuleConfigRequest {
/// Module descriptor containing name, file_path, attach_points, etc.
pub descriptor: WasmModuleDescriptor,
@@ -61,12 +65,13 @@ fn has_wasm_extension(path: &Path) -> bool {
pub struct ValidateDescriptorStep;
#[async_trait]
impl StepExecutor for ValidateDescriptorStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let descriptor = &config_request.descriptor;
impl StepExecutor<AnyWorkflowData> for ValidateDescriptorStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_registration()?;
let descriptor = &data.config.descriptor;
debug!("Validating WASM module descriptor: {}", descriptor.name);
@@ -198,12 +203,16 @@ impl StepExecutor for ValidateDescriptorStep {
});
}
// Store file size in context for later steps
context.set("file_size_bytes", metadata.len());
// Clone name for logging before mutable borrow
let module_name = descriptor.name.clone();
// Store file size in typed data
let data_mut = context.data.as_wasm_registration_mut()?;
data_mut.file_size_bytes = Some(metadata.len());
info!(
"Descriptor validated successfully for module: {}",
descriptor.name
module_name
);
Ok(StepResult::Success)
}
@@ -220,12 +229,13 @@ impl StepExecutor for ValidateDescriptorStep {
pub struct CalculateHashStep;
#[async_trait]
impl StepExecutor for CalculateHashStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let file_path = &config_request.descriptor.file_path;
impl StepExecutor<AnyWorkflowData> for CalculateHashStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_registration()?;
let file_path = &data.config.descriptor.file_path;
debug!("Calculating SHA256 hash for: {}", file_path);
@@ -260,10 +270,14 @@ impl StepExecutor for CalculateHashStep {
let hash: [u8; 32] = hasher.finalize().into();
// Store hash in context
context.set("sha256_hash", hash);
// Clone path for logging before mutable borrow
let path_for_log = file_path.clone();
info!("SHA256 hash calculated for: {}", file_path);
// Store hash in typed data
let data_mut = context.data.as_wasm_registration_mut()?;
data_mut.sha256_hash = Some(hash);
info!("SHA256 hash calculated for: {}", path_for_log);
Ok(StepResult::Success)
}
@@ -279,16 +293,24 @@ impl StepExecutor for CalculateHashStep {
pub struct CheckDuplicateStep;
#[async_trait]
impl StepExecutor for CheckDuplicateStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
impl StepExecutor<AnyWorkflowData> for CheckDuplicateStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_registration()?;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let sha256_hash = data
.sha256_hash
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
debug!(
"Checking for duplicate SHA256 hash for module: {}",
config_request.descriptor.name
data.config.descriptor.name
);
// Get WASM module manager from app context
@@ -303,7 +325,7 @@ impl StepExecutor for CheckDuplicateStep {
// Check for duplicate hash using manager's internal method
wasm_manager
.check_duplicate_sha256_hash(sha256_hash.as_ref())
.check_duplicate_sha256_hash(sha256_hash)
.map_err(|e| WorkflowError::StepFailed {
step_id: StepId::new("check_duplicate"),
message: format!("Duplicate SHA256 hash detected: {}", e),
@@ -311,7 +333,7 @@ impl StepExecutor for CheckDuplicateStep {
info!(
"No duplicate found for module: {}",
config_request.descriptor.name
data.config.descriptor.name
);
Ok(StepResult::Success)
}
@@ -328,15 +350,19 @@ impl StepExecutor for CheckDuplicateStep {
pub struct LoadWasmBytesStep;
#[async_trait]
impl StepExecutor for LoadWasmBytesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let file_path = &config_request.descriptor.file_path;
impl StepExecutor<AnyWorkflowData> for LoadWasmBytesStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_registration()?;
let file_path = &data.config.descriptor.file_path;
debug!("Loading WASM bytes from: {}", file_path);
// Clone path for logging before mutable borrow
let path_for_log = file_path.clone();
let wasm_bytes =
tokio::fs::read(file_path)
.await
@@ -345,10 +371,11 @@ impl StepExecutor for LoadWasmBytesStep {
message: format!("Failed to read WASM file {}: {}", file_path, e),
})?;
// Store WASM bytes in context
context.set("wasm_bytes", wasm_bytes);
// Store WASM bytes in typed data
let data_mut = context.data.as_wasm_registration_mut()?;
data_mut.wasm_bytes = Some(wasm_bytes);
info!("WASM bytes loaded from: {}", file_path);
info!("WASM bytes loaded from: {}", path_for_log);
Ok(StepResult::Success)
}
@@ -364,15 +391,20 @@ impl StepExecutor for LoadWasmBytesStep {
pub struct ValidateWasmComponentStep;
#[async_trait]
impl StepExecutor for ValidateWasmComponentStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
impl StepExecutor<AnyWorkflowData> for ValidateWasmComponentStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_registration()?;
let wasm_bytes = data
.wasm_bytes
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?;
debug!(
"Validating WASM component format for module: {}",
config_request.descriptor.name
data.config.descriptor.name
);
// Create a temporary engine to validate the component
@@ -386,7 +418,7 @@ impl StepExecutor for ValidateWasmComponentStep {
})?;
// Attempt to compile the component to validate it
Component::new(&engine, wasm_bytes.as_ref())
Component::new(&engine, wasm_bytes)
.map_err(|e| WorkflowError::StepFailed {
step_id: StepId::new("validate_wasm_component"),
message: format!(
@@ -399,7 +431,7 @@ impl StepExecutor for ValidateWasmComponentStep {
info!(
"WASM component validated successfully for module: {}",
config_request.descriptor.name
data.config.descriptor.name
);
Ok(StepResult::Success)
}
@@ -416,19 +448,31 @@ impl StepExecutor for ValidateWasmComponentStep {
pub struct RegisterModuleStep;
#[async_trait]
impl StepExecutor for RegisterModuleStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config_request: Arc<WasmModuleConfigRequest> =
context.get_or_err("wasm_module_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?;
let file_size_bytes: Arc<u64> = context.get_or_err("file_size_bytes")?;
let wasm_bytes: Arc<Vec<u8>> = context.get_or_err("wasm_bytes")?;
impl StepExecutor<AnyWorkflowData> for RegisterModuleStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_registration()?;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let sha256_hash = data
.sha256_hash
.ok_or_else(|| WorkflowError::ContextValueNotFound("sha256_hash".to_string()))?;
let file_size_bytes = data
.file_size_bytes
.ok_or_else(|| WorkflowError::ContextValueNotFound("file_size_bytes".to_string()))?;
let wasm_bytes = data
.wasm_bytes
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("wasm_bytes".to_string()))?
.clone();
debug!(
"Registering WASM module in manager: {}",
config_request.descriptor.name
);
let descriptor = &data.config.descriptor;
debug!("Registering WASM module in manager: {}", descriptor.name);
// Get WASM module manager from app context
let wasm_manager =
@@ -451,18 +495,21 @@ impl StepExecutor for RegisterModuleStep {
let module = WasmModule {
module_uuid,
module_meta: WasmModuleMeta {
name: config_request.descriptor.name.clone(),
file_path: config_request.descriptor.file_path.clone(),
sha256_hash: *sha256_hash.as_ref(),
size_bytes: *file_size_bytes.as_ref(),
name: descriptor.name.clone(),
file_path: descriptor.file_path.clone(),
sha256_hash,
size_bytes: file_size_bytes,
created_at: now,
last_accessed_at: now,
access_count: 0,
attach_points: config_request.descriptor.attach_points.clone(),
wasm_bytes: wasm_bytes.as_ref().clone(),
attach_points: descriptor.attach_points.clone(),
wasm_bytes,
},
};
// Clone name for logging before mutable borrow
let module_name = descriptor.name.clone();
// Register module in manager
wasm_manager
.register_module_internal(module)
@@ -471,12 +518,13 @@ impl StepExecutor for RegisterModuleStep {
message: format!("Failed to register module: {}", e),
})?;
// Store module UUID in context for return value
context.set("module_uuid", module_uuid);
// Store module UUID in typed data
let data_mut = context.data.as_wasm_registration_mut()?;
data_mut.module_uuid = Some(module_uuid);
info!(
"WASM module registered successfully: {} (UUID: {})",
config_request.descriptor.name, module_uuid
module_name, module_uuid
);
Ok(StepResult::Success)
@@ -504,7 +552,7 @@ impl StepExecutor for RegisterModuleStep {
/// - LoadWasmBytes: 3 retries, 60s timeout (I/O intensive)
/// - ValidateWasmComponent: No retry, 30s timeout (CPU intensive validation)
/// - RegisterModule: No retry, 5s timeout (fast registration)
pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition {
pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new("wasm_module_registration", "WASM Module Registration")
.add_step(
StepDefinition::new(
@@ -574,3 +622,18 @@ pub fn create_wasm_module_registration_workflow() -> WorkflowDefinition {
.depends_on(&["validate_wasm_component"]),
)
}
/// Helper to create initial workflow data for WASM module registration
pub fn create_wasm_registration_workflow_data(
config: WasmModuleConfigRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::WasmRegistration(WasmRegistrationWorkflowData {
config,
wasm_bytes: None,
sha256_hash: None,
file_size_bytes: None,
module_uuid: None,
app_context: Some(app_context),
})
}
sgl-model-gateway/src/core/steps/wasm_module_removal.rs
+52 -14
View File
@@ -4,10 +4,17 @@ use async_trait::async_trait;
use tracing::{debug, info};
use uuid::Uuid;
use crate::{app_context::AppContext, workflow::*};
use super::workflow_data::{AnyWorkflowData, WasmRemovalWorkflowData};
use crate::{
app_context::AppContext,
workflow::{
FailureAction, StepDefinition, StepExecutor, StepId, StepResult, WorkflowContext,
WorkflowDefinition, WorkflowError, WorkflowResult,
},
};
/// WASM module removal request
#[derive(Debug, Clone)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WasmModuleRemovalRequest {
/// Module UUID to remove
pub module_uuid: Uuid,
@@ -30,11 +37,17 @@ impl WasmModuleRemovalRequest {
pub struct FindModuleToRemoveStep;
#[async_trait]
impl StepExecutor for FindModuleToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let removal_request: Arc<WasmModuleRemovalRequest> =
context.get_or_err("wasm_module_removal_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for FindModuleToRemoveStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_removal()?;
let removal_request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!("Finding module to remove: {}", removal_request.module_uuid);
@@ -63,7 +76,14 @@ impl StepExecutor for FindModuleToRemoveStep {
});
}
info!("Module found for removal: {}", removal_request.module_uuid);
// Clone uuid for logging before mutable borrow
let module_uuid = removal_request.module_uuid;
// Store the module ID in typed data
let data_mut = context.data.as_wasm_removal_mut()?;
data_mut.module_id = Some(module_uuid.to_string());
info!("Module found for removal: {}", module_uuid);
Ok(StepResult::Success)
}
@@ -78,11 +98,17 @@ impl StepExecutor for FindModuleToRemoveStep {
pub struct RemoveModuleStep;
#[async_trait]
impl StepExecutor for RemoveModuleStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let removal_request: Arc<WasmModuleRemovalRequest> =
context.get_or_err("wasm_module_removal_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for RemoveModuleStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_wasm_removal()?;
let removal_request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!("Removing WASM module: {}", removal_request.module_uuid);
@@ -125,7 +151,7 @@ impl StepExecutor for RemoveModuleStep {
/// Workflow configuration:
/// - FindModuleToRemove: No retry, 5s timeout (fast lookup)
/// - RemoveModule: No retry, 5s timeout (fast removal)
pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition {
pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new("wasm_module_removal", "WASM Module Removal")
.add_step(
StepDefinition::new(
@@ -143,3 +169,15 @@ pub fn create_wasm_module_removal_workflow() -> WorkflowDefinition {
.depends_on(&["find_module_to_remove"]),
)
}
/// Helper to create initial workflow data for WASM module removal
pub fn create_wasm_removal_workflow_data(
config: WasmModuleRemovalRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::WasmRemoval(WasmRemovalWorkflowData {
config,
module_id: None,
app_context: Some(app_context),
})
}
sgl-model-gateway/src/core/steps/worker/external/create_workers.rs
Vendored
+18 -10
View File
@@ -6,14 +6,12 @@ use async_trait::async_trait;
use tracing::{debug, info};
use crate::{
app_context::AppContext,
core::{
circuit_breaker::CircuitBreakerConfig,
model_card::ModelCard,
steps::workflow_data::{AnyWorkflowData, WorkerList},
worker::{HealthConfig, RuntimeType, WorkerType},
BasicWorkerBuilder, ConnectionMode, Worker,
},
protocols::worker_spec::WorkerConfigRequest,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -30,11 +28,18 @@ fn normalize_external_url(url: &str) -> String {
pub struct CreateExternalWorkersStep;
#[async_trait]
impl StepExecutor for CreateExternalWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let model_cards: Arc<Vec<ModelCard>> = context.get_or_err("model_cards")?;
impl StepExecutor<AnyWorkflowData> for CreateExternalWorkersStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_external_worker()?;
let config = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let model_cards = &data.model_cards;
// Build configs from router settings
let circuit_breaker_config = {
@@ -144,8 +149,11 @@ impl StepExecutor for CreateExternalWorkersStep {
);
}
context.set("workers", workers);
context.set("labels", labels);
// Store results in workflow data
let data_mut = context.data.as_external_worker_mut()?;
data_mut.workers = Some(WorkerList::from_workers(&workers));
data_mut.actual_workers = Some(workers);
data_mut.labels = labels;
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/external/discover_models.rs
Vendored
+11 -7
View File
@@ -1,6 +1,6 @@
//! Model discovery step for external API endpoints.
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, time::Duration};
use async_trait::async_trait;
use once_cell::sync::Lazy;
@@ -13,8 +13,8 @@ use crate::{
core::{
model_card::{ModelCard, ProviderType},
model_type::ModelType,
steps::workflow_data::AnyWorkflowData,
},
protocols::worker_spec::WorkerConfigRequest,
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -225,9 +225,13 @@ async fn fetch_models(url: &str, api_key: Option<&str>) -> Result<Vec<ModelCard>
pub struct DiscoverModelsStep;
#[async_trait]
impl StepExecutor for DiscoverModelsStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
impl StepExecutor<AnyWorkflowData> for DiscoverModelsStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_external_worker()?;
let config = &data.config;
// If no API key is provided, skip model discovery and use wildcard mode.
if config.api_key.as_ref().is_none_or(|k| k.is_empty()) {
@@ -236,7 +240,7 @@ impl StepExecutor for DiscoverModelsStep {
User's Authorization header will be forwarded to backend.",
config.url
);
context.set::<Vec<ModelCard>>("model_cards", vec![]);
// Leave model_cards empty for wildcard mode
return Ok(StepResult::Success);
}
@@ -263,7 +267,7 @@ impl StepExecutor for DiscoverModelsStep {
model_cards.iter().map(|c| &c.id).collect::<Vec<_>>()
);
context.set("model_cards", model_cards);
context.data.as_external_worker_mut()?.model_cards = model_cards;
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/external/mod.rs
Vendored
+21 -3
View File
@@ -15,8 +15,11 @@ pub use discover_models::{
};
use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
use crate::workflow::{
BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition,
use crate::{
app_context::AppContext,
core::steps::workflow_data::{AnyWorkflowData, ExternalWorkerWorkflowData},
protocols::worker_spec::WorkerConfigRequest,
workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
};
/// Create external worker registration workflow definition.
@@ -35,7 +38,7 @@ use crate::workflow::{
/// │ │
/// └────────────┴────────────┘
/// ```
pub fn create_external_worker_workflow() -> WorkflowDefinition {
pub fn create_external_worker_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new(
"external_worker_registration",
"External Worker Registration",
@@ -102,3 +105,18 @@ pub fn create_external_worker_workflow() -> WorkflowDefinition {
.depends_on(&["register_workers"]),
)
}
/// Helper to create initial workflow data for external worker registration
pub fn create_external_worker_workflow_data(
config: WorkerConfigRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::ExternalWorker(ExternalWorkerWorkflowData {
config,
model_cards: Vec::new(),
workers: None,
labels: std::collections::HashMap::new(),
app_context: Some(app_context),
actual_workers: None,
})
}
sgl-model-gateway/src/core/steps/worker/local/create_worker.rs
+41 -29
View File
@@ -5,12 +5,12 @@ use std::{collections::HashMap, sync::Arc, time::Duration};
use async_trait::async_trait;
use tracing::debug;
use super::discover_dp::DpInfo;
use crate::{
app_context::AppContext,
core::{
circuit_breaker::CircuitBreakerConfig,
model_card::ModelCard,
steps::workflow_data::{AnyWorkflowData, LocalWorkerWorkflowData},
worker::{HealthConfig, RuntimeType, WorkerType},
BasicWorkerBuilder, ConnectionMode, DPAwareWorkerBuilder, Worker, UNKNOWN_MODEL_ID,
},
@@ -29,13 +29,22 @@ use crate::{
pub struct CreateLocalWorkerStep;
#[async_trait]
impl StepExecutor for CreateLocalWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
let discovered_labels: Arc<HashMap<String, String>> =
context.get_or_err("discovered_labels")?;
impl StepExecutor<AnyWorkflowData> for CreateLocalWorkerStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_local_worker()?;
let config = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let connection_mode = data
.connection_mode
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?;
let discovered_labels = &data.discovered_labels;
// Check if worker already exists
if app_context
@@ -59,7 +68,7 @@ impl StepExecutor for CreateLocalWorkerStep {
}
// Merge: discovered labels first, then config labels (config takes precedence)
let mut final_labels = discovered_labels.as_ref().clone();
let mut final_labels = discovered_labels.clone();
for (key, value) in &config_labels {
final_labels.insert(key.clone(), value.clone());
}
@@ -77,7 +86,7 @@ impl StepExecutor for CreateLocalWorkerStep {
}
// Create ModelCard
let model_card = build_model_card(&model_id, &config, &final_labels);
let model_card = build_model_card(&model_id, config, &final_labels);
debug!(
"Creating worker {} with {} discovered + {} config = {} final labels",
@@ -88,41 +97,39 @@ impl StepExecutor for CreateLocalWorkerStep {
);
// Parse worker type
let worker_type = parse_worker_type(&config);
let worker_type = parse_worker_type(config);
// Get runtime type (for gRPC workers)
let runtime_type = determine_runtime_type(&connection_mode, context, &config);
let runtime_type = determine_runtime_type(connection_mode, data, config);
// Build circuit breaker config
let circuit_breaker_config = build_circuit_breaker_config(&app_context);
let circuit_breaker_config = build_circuit_breaker_config(app_context);
// Build health config
let health_config = build_health_config(&app_context);
let health_config = build_health_config(app_context);
// Normalize URL
let normalized_url = normalize_url(&config.url, &connection_mode);
let normalized_url = normalize_url(&config.url, connection_mode);
if normalized_url != config.url {
debug!(
"Normalized worker URL: {} -> {} ({:?})",
config.url,
normalized_url,
connection_mode.as_ref()
config.url, normalized_url, connection_mode
);
}
// Create workers - always output as Vec for unified downstream handling
let workers = if config.dp_aware {
create_dp_aware_workers(
context,
data,
&normalized_url,
model_card,
worker_type,
&connection_mode,
connection_mode,
runtime_type,
circuit_breaker_config,
health_config,
&config,
config,
&final_labels,
)?
} else {
@@ -130,17 +137,19 @@ impl StepExecutor for CreateLocalWorkerStep {
&normalized_url,
model_card,
worker_type,
&connection_mode,
connection_mode,
runtime_type,
circuit_breaker_config,
health_config,
&config,
config,
&final_labels,
)
};
context.set("workers", workers);
context.set("labels", final_labels);
// Update workflow data
let data_mut = context.data.as_local_worker_mut()?;
data_mut.actual_workers = Some(workers);
data_mut.final_labels = final_labels;
Ok(StepResult::Success)
}
@@ -230,14 +239,14 @@ fn parse_worker_type(config: &WorkerConfigRequest) -> WorkerType {
fn determine_runtime_type(
connection_mode: &ConnectionMode,
context: &WorkflowContext,
data: &LocalWorkerWorkflowData,
config: &WorkerConfigRequest,
) -> RuntimeType {
if !matches!(connection_mode, ConnectionMode::Grpc { .. }) {
return RuntimeType::Sglang;
}
if let Some(detected_runtime) = context.get::<String>("detected_runtime_type") {
if let Some(ref detected_runtime) = data.detected_runtime_type {
match detected_runtime.as_str() {
"vllm" => RuntimeType::Vllm,
_ => RuntimeType::Sglang,
@@ -286,7 +295,7 @@ fn normalize_url(url: &str, connection_mode: &ConnectionMode) -> String {
#[allow(clippy::too_many_arguments)]
fn create_dp_aware_workers(
context: &WorkflowContext,
data: &LocalWorkerWorkflowData,
normalized_url: &str,
model_card: ModelCard,
worker_type: WorkerType,
@@ -297,7 +306,10 @@ fn create_dp_aware_workers(
config: &WorkerConfigRequest,
final_labels: &HashMap<String, String>,
) -> Result<Vec<Arc<dyn Worker>>, WorkflowError> {
let dp_info: Arc<DpInfo> = context.get_or_err("dp_info")?;
let dp_info = data
.dp_info
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("dp_info".to_string()))?;
debug!(
"Creating {} DP-aware workers for {} (dp_size: {})",
sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs
+14 -9
View File
@@ -1,6 +1,6 @@
//! Connection mode detection step.
use std::{sync::Arc, time::Duration};
use std::time::Duration;
use async_trait::async_trait;
use reqwest::Client;
@@ -8,9 +8,7 @@ use tracing::debug;
use super::strip_protocol;
use crate::{
app_context::AppContext,
core::ConnectionMode,
protocols::worker_spec::WorkerConfigRequest,
core::{steps::workflow_data::AnyWorkflowData, ConnectionMode},
routers::grpc::client::GrpcClient,
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -88,10 +86,17 @@ async fn try_grpc_health_check(
pub struct DetectConnectionModeStep;
#[async_trait]
impl StepExecutor for DetectConnectionModeStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for DetectConnectionModeStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_local_worker()?;
let config = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!(
"Detecting connection mode for {} (timeout: {}s, max_attempts: {})",
@@ -129,7 +134,7 @@ impl StepExecutor for DetectConnectionModeStep {
}
};
context.set("connection_mode", connection_mode);
context.data.as_local_worker_mut()?.connection_mode = Some(connection_mode);
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs
+10 -9
View File
@@ -1,19 +1,16 @@
//! Data Parallel (DP) information discovery step.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::debug;
use super::discover_metadata::get_server_info;
use crate::{
core::UNKNOWN_MODEL_ID,
protocols::worker_spec::WorkerConfigRequest,
core::{steps::workflow_data::AnyWorkflowData, UNKNOWN_MODEL_ID},
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// DP (Data Parallel) information for a worker.
#[derive(Debug, Clone)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct DpInfo {
pub dp_size: usize,
pub model_id: String,
@@ -44,9 +41,13 @@ pub async fn get_dp_info(url: &str, api_key: Option<&str>) -> Result<DpInfo, Str
pub struct DiscoverDPInfoStep;
#[async_trait]
impl StepExecutor for DiscoverDPInfoStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
impl StepExecutor<AnyWorkflowData> for DiscoverDPInfoStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_local_worker()?;
let config = &data.config;
if !config.dp_aware {
debug!(
@@ -70,7 +71,7 @@ impl StepExecutor for DiscoverDPInfoStep {
dp_info.dp_size, config.url, dp_info.model_id
);
context.set("dp_info", dp_info);
context.data.as_local_worker_mut()?.dp_info = Some(dp_info);
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs
+21 -12
View File
@@ -1,6 +1,6 @@
//! Metadata discovery step for local workers.
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{collections::HashMap, time::Duration};
use async_trait::async_trait;
use once_cell::sync::Lazy;
@@ -11,8 +11,7 @@ use tracing::{debug, warn};
use super::strip_protocol;
use crate::{
core::ConnectionMode,
protocols::worker_spec::WorkerConfigRequest,
core::{steps::workflow_data::AnyWorkflowData, ConnectionMode},
routers::grpc::client::GrpcClient,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -219,17 +218,24 @@ async fn fetch_grpc_metadata(
pub struct DiscoverMetadataStep;
#[async_trait]
impl StepExecutor for DiscoverMetadataStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let config: Arc<WorkerConfigRequest> = context.get_or_err("worker_config")?;
let connection_mode: Arc<ConnectionMode> = context.get_or_err("connection_mode")?;
impl StepExecutor<AnyWorkflowData> for DiscoverMetadataStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_local_worker()?;
let config = &data.config;
let connection_mode = data
.connection_mode
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("connection_mode".to_string()))?;
debug!(
"Discovering metadata for {} ({:?})",
config.url, *connection_mode
config.url, connection_mode
);
let (discovered_labels, detected_runtime) = match connection_mode.as_ref() {
let (discovered_labels, detected_runtime) = match connection_mode {
ConnectionMode::Http => {
let mut labels = HashMap::new();
@@ -287,16 +293,19 @@ impl StepExecutor for DiscoverMetadataStep {
(HashMap::new(), None)
});
let url = config.url.clone();
debug!(
"Discovered {} metadata labels for {}",
discovered_labels.len(),
config.url
url
);
context.set("discovered_labels", discovered_labels);
// Update workflow data
let data_mut = context.data.as_local_worker_mut()?;
data_mut.discovered_labels = discovered_labels;
if let Some(runtime) = detected_runtime {
debug!("Detected runtime type: {}", runtime);
context.set("detected_runtime_type", runtime);
data_mut.detected_runtime_type = Some(runtime);
}
Ok(StepResult::Success)
sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs
+19 -22
View File
@@ -1,13 +1,11 @@
//! Step to find a worker to update based on URL.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::debug;
use super::find_workers_by_url;
use crate::{
app_context::AppContext,
core::steps::workflow_data::AnyWorkflowData,
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -15,31 +13,30 @@ use crate::{
///
/// For DP-aware workers, finds all workers with matching URL prefix.
/// For regular workers, finds the single worker with exact URL match.
///
/// Expects the following context values:
/// - "worker_url": String - the URL of the worker to update
/// - "dp_aware": bool - whether to find all DP-aware workers with matching prefix
/// - "app_context": Arc<AppContext>
///
/// Sets the following context values:
/// - "workers_to_update": Vec<Arc<dyn Worker>>
pub struct FindWorkerToUpdateStep;
#[async_trait]
impl StepExecutor for FindWorkerToUpdateStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let worker_url: Arc<String> = context.get_or_err("worker_url")?;
let dp_aware: Arc<bool> = context.get_or_err("dp_aware")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for FindWorkerToUpdateStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_update()?;
let worker_url = &data.worker_url;
let dp_aware = data.dp_aware;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers_to_update =
find_workers_by_url(&app_context.worker_registry, &worker_url, *dp_aware);
find_workers_by_url(&app_context.worker_registry, worker_url, dp_aware);
if workers_to_update.is_empty() {
let error_msg = if *dp_aware {
format!("No workers found with prefix {}@", *worker_url)
let error_msg = if dp_aware {
format!("No workers found with prefix {}@", worker_url)
} else {
format!("Worker {} not found", *worker_url)
format!("Worker {} not found", worker_url)
};
return Err(WorkflowError::StepFailed {
step_id: StepId::new("find_worker_to_update"),
@@ -50,10 +47,10 @@ impl StepExecutor for FindWorkerToUpdateStep {
debug!(
"Found {} worker(s) to update for {}",
workers_to_update.len(),
*worker_url
worker_url
);
context.set("workers_to_update", workers_to_update);
context.data.as_worker_update_mut()?.workers_to_update = Some(workers_to_update);
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs
+20 -10
View File
@@ -1,18 +1,18 @@
//! Step to find workers to remove based on URL.
use std::{collections::HashSet, sync::Arc};
use std::collections::HashSet;
use async_trait::async_trait;
use tracing::debug;
use super::find_workers_by_url;
use crate::{
app_context::AppContext,
core::steps::workflow_data::{AnyWorkflowData, WorkerList},
workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Request structure for worker removal.
#[derive(Debug, Clone)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct WorkerRemovalRequest {
pub url: String,
pub dp_aware: bool,
@@ -25,10 +25,17 @@ pub struct WorkerRemovalRequest {
pub struct FindWorkersToRemoveStep;
#[async_trait]
impl StepExecutor for FindWorkersToRemoveStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerRemovalRequest> = context.get_or_err("removal_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
impl StepExecutor<AnyWorkflowData> for FindWorkersToRemoveStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_removal()?;
let request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers_to_remove =
find_workers_by_url(&app_context.worker_registry, &request.url, request.dp_aware);
@@ -62,9 +69,12 @@ impl StepExecutor for FindWorkersToRemoveStep {
.map(|w| w.model_id().to_string())
.collect();
context.set("workers_to_remove", workers_to_remove);
context.set("worker_urls", worker_urls);
context.set("affected_models", affected_models);
// Update workflow data
let data_mut = context.data.as_worker_removal_mut()?;
data_mut.workers_to_remove = Some(WorkerList::from_workers(&workers_to_remove));
data_mut.actual_workers_to_remove = Some(workers_to_remove);
data_mut.worker_urls = worker_urls;
data_mut.affected_models = affected_models;
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/local/mod.rs
+66 -4
View File
@@ -36,8 +36,16 @@ pub use update_worker_properties::UpdateWorkerPropertiesStep;
use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep};
use crate::{
app_context::AppContext,
config::RouterConfig,
core::{Worker, WorkerRegistry},
core::{
steps::workflow_data::{
AnyWorkflowData, LocalWorkerWorkflowData, WorkerRemovalWorkflowData,
WorkerUpdateWorkflowData,
},
Worker, WorkerRegistry,
},
protocols::worker_spec::{WorkerConfigRequest, WorkerUpdateRequest},
workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition},
};
@@ -66,7 +74,9 @@ pub(crate) fn find_workers_by_url(
}
}
pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDefinition {
pub fn create_local_worker_workflow(
router_config: &RouterConfig,
) -> WorkflowDefinition<AnyWorkflowData> {
let detect_timeout = Duration::from_secs(router_config.worker_startup_timeout_secs);
// Calculate max_attempts based on timeout
@@ -198,7 +208,7 @@ pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDef
/// │
/// update_remaining_policies
/// ```
pub fn create_worker_removal_workflow() -> WorkflowDefinition {
pub fn create_worker_removal_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new("worker_removal", "Remove worker from router")
.add_step(
StepDefinition::new(
@@ -263,7 +273,7 @@ pub fn create_worker_removal_workflow() -> WorkflowDefinition {
/// │
/// update_policies_for_worker
/// ```
pub fn create_worker_update_workflow() -> WorkflowDefinition {
pub fn create_worker_update_workflow() -> WorkflowDefinition<AnyWorkflowData> {
WorkflowDefinition::new("worker_update", "Update worker properties")
.add_step(
StepDefinition::new(
@@ -304,3 +314,55 @@ pub fn create_worker_update_workflow() -> WorkflowDefinition {
.depends_on(&["update_worker_properties"]),
)
}
/// Helper to create initial workflow data for local worker registration
pub fn create_local_worker_workflow_data(
config: WorkerConfigRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::LocalWorker(LocalWorkerWorkflowData {
config,
connection_mode: None,
discovered_labels: std::collections::HashMap::new(),
dp_info: None,
workers: None,
final_labels: std::collections::HashMap::new(),
detected_runtime_type: None,
app_context: Some(app_context),
actual_workers: None,
})
}
/// Helper to create initial workflow data for worker removal
pub fn create_worker_removal_workflow_data(
url: String,
dp_aware: bool,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
AnyWorkflowData::WorkerRemoval(WorkerRemovalWorkflowData {
config: WorkerRemovalRequest { url, dp_aware },
workers_to_remove: None,
worker_urls: Vec::new(),
affected_models: std::collections::HashSet::new(),
app_context: Some(app_context),
actual_workers_to_remove: None,
})
}
/// Helper to create initial workflow data for worker update
pub fn create_worker_update_workflow_data(
worker_url: String,
update_config: WorkerUpdateRequest,
app_context: Arc<AppContext>,
) -> AnyWorkflowData {
// Determine if this is a DP-aware update based on URL pattern
let dp_aware = worker_url.contains('@');
AnyWorkflowData::WorkerUpdate(WorkerUpdateWorkflowData {
config: update_config,
worker_url,
dp_aware,
app_context: Some(app_context),
workers_to_update: None,
updated_workers: None,
})
}
sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs
+18 -10
View File
@@ -1,13 +1,10 @@
//! Tokenizer registration step for local workers.
use std::{collections::HashMap, sync::Arc};
use async_trait::async_trait;
use tracing::{debug, warn};
use crate::{
app_context::AppContext,
core::Worker,
core::steps::workflow_data::AnyWorkflowData,
tokenizer::{factory, TokenizerRegistry},
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -16,11 +13,21 @@ use crate::{
pub struct RegisterTokenizerStep;
#[async_trait]
impl StepExecutor for RegisterTokenizerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
impl StepExecutor<AnyWorkflowData> for RegisterTokenizerStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_local_worker()?;
let labels = &data.final_labels;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers = data
.actual_workers
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
for worker in workers.iter() {
let model_id = worker.model_id().to_string();
@@ -46,10 +53,11 @@ impl StepExecutor for RegisterTokenizerStep {
let source = tokenizer_path.clone();
// Load tokenizer with thread safe lock
let tokenizer_path_owned = tokenizer_path.clone();
if let Err(e) = app_context
.tokenizer_registry
.load(&tokenizer_id, &model_id, &source, || async move {
factory::create_tokenizer_async(&tokenizer_path.to_string())
factory::create_tokenizer_async(&tokenizer_path_owned)
.await
.map_err(|e| e.to_string())
})
sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs
+15 -9
View File
@@ -1,13 +1,10 @@
//! Step to remove workers from policy registry.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::debug;
use crate::{
app_context::AppContext,
core::Worker,
core::steps::workflow_data::AnyWorkflowData,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -18,11 +15,20 @@ use crate::{
pub struct RemoveFromPolicyRegistryStep;
#[async_trait]
impl StepExecutor for RemoveFromPolicyRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_remove: Arc<Vec<Arc<dyn Worker>>> =
context.get_or_err("workers_to_remove")?;
impl StepExecutor<AnyWorkflowData> for RemoveFromPolicyRegistryStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_removal()?;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers_to_remove = data
.actual_workers_to_remove
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_remove".to_string()))?;
debug!(
"Removing {} worker(s) from policy registry",
sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs
+13 -6
View File
@@ -1,12 +1,12 @@
//! Step to remove workers from worker registry.
use std::{collections::HashSet, sync::Arc};
use std::collections::HashSet;
use async_trait::async_trait;
use tracing::{debug, warn};
use crate::{
app_context::AppContext,
core::steps::workflow_data::AnyWorkflowData,
observability::metrics::Metrics,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -17,10 +17,17 @@ use crate::{
pub struct RemoveFromWorkerRegistryStep;
#[async_trait]
impl StepExecutor for RemoveFromWorkerRegistryStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
impl StepExecutor<AnyWorkflowData> for RemoveFromWorkerRegistryStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_removal()?;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let worker_urls = &data.worker_urls;
debug!(
"Removing {} worker(s) from worker registry",
sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs
+16 -7
View File
@@ -1,13 +1,12 @@
//! Step to update policies for updated workers.
use std::{collections::HashSet, sync::Arc};
use std::collections::HashSet;
use async_trait::async_trait;
use tracing::debug;
use crate::{
app_context::AppContext,
core::Worker,
core::steps::workflow_data::AnyWorkflowData,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -18,10 +17,20 @@ use crate::{
pub struct UpdatePoliciesForWorkerStep;
#[async_trait]
impl StepExecutor for UpdatePoliciesForWorkerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let updated_workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("updated_workers")?;
impl StepExecutor<AnyWorkflowData> for UpdatePoliciesForWorkerStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_update()?;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let updated_workers = data
.updated_workers
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("updated_workers".to_string()))?;
// Collect affected models
let affected_models: HashSet<String> = updated_workers
sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs
+13 -8
View File
@@ -1,12 +1,10 @@
//! Step to update cache-aware policies for remaining workers after removal.
use std::{collections::HashSet, sync::Arc};
use async_trait::async_trait;
use tracing::{debug, info};
use crate::{
app_context::AppContext,
core::steps::workflow_data::AnyWorkflowData,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -17,11 +15,18 @@ use crate::{
pub struct UpdateRemainingPoliciesStep;
#[async_trait]
impl StepExecutor for UpdateRemainingPoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let affected_models: Arc<HashSet<String>> = context.get_or_err("affected_models")?;
let worker_urls: Arc<Vec<String>> = context.get_or_err("worker_urls")?;
impl StepExecutor<AnyWorkflowData> for UpdateRemainingPoliciesStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_removal()?;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let affected_models = &data.affected_models;
let worker_urls = &data.worker_urls;
debug!(
"Updating cache-aware policies for {} affected model(s)",
sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs
+19 -18
View File
@@ -6,9 +6,7 @@ use async_trait::async_trait;
use tracing::{debug, info};
use crate::{
app_context::AppContext,
core::{BasicWorkerBuilder, HealthConfig, Worker},
protocols::worker_spec::WorkerUpdateRequest,
core::{steps::workflow_data::AnyWorkflowData, BasicWorkerBuilder, HealthConfig, Worker},
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
@@ -16,23 +14,26 @@ use crate::{
///
/// This step creates new worker instances with updated properties and
/// re-registers them to replace the old workers in the registry.
///
/// Expects the following context values:
/// - "update_request": WorkerUpdateRequest (from protocols::worker_spec)
/// - "app_context": Arc<AppContext>
/// - "workers_to_update": Vec<Arc<dyn Worker>>
///
/// Sets the following context values:
/// - "updated_workers": Vec<Arc<dyn Worker>>
pub struct UpdateWorkerPropertiesStep;
#[async_trait]
impl StepExecutor for UpdateWorkerPropertiesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let request: Arc<WorkerUpdateRequest> = context.get_or_err("update_request")?;
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers_to_update: Arc<Vec<Arc<dyn Worker>>> =
context.get_or_err("workers_to_update")?;
impl StepExecutor<AnyWorkflowData> for UpdateWorkerPropertiesStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let data = context.data.as_worker_update()?;
let request = &data.config;
let app_context = data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
.clone();
let workers_to_update = data
.workers_to_update
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers_to_update".to_string()))?
.clone();
debug!(
"Updating properties for {} worker(s)",
@@ -136,7 +137,7 @@ impl StepExecutor for UpdateWorkerPropertiesStep {
}
// Store updated workers for subsequent steps
context.set("updated_workers", updated_workers);
context.data.as_worker_update_mut()?.updated_workers = Some(updated_workers);
Ok(StepResult::Success)
}
sgl-model-gateway/src/core/steps/worker/mod.rs
+8 -7
View File
@@ -3,15 +3,16 @@ pub mod local;
pub mod shared;
pub use external::{
create_external_worker_workflow as create_external_worker_registration_workflow,
create_external_worker_workflow, group_models_into_cards, infer_model_type_from_id,
CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo, ModelsResponse,
create_external_worker_workflow, create_external_worker_workflow_data, group_models_into_cards,
infer_model_type_from_id, CreateExternalWorkersStep, DiscoverModelsStep, ModelInfo,
ModelsResponse,
};
pub use local::{
create_local_worker_workflow as create_worker_registration_workflow,
create_local_worker_workflow, create_worker_removal_workflow, create_worker_update_workflow,
CreateLocalWorkerStep, DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep,
DpInfo, FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
create_local_worker_workflow, create_local_worker_workflow_data,
create_worker_removal_workflow, create_worker_removal_workflow_data,
create_worker_update_workflow, create_worker_update_workflow_data, CreateLocalWorkerStep,
DetectConnectionModeStep, DiscoverDPInfoStep, DiscoverMetadataStep, DpInfo,
FindWorkerToUpdateStep, FindWorkersToRemoveStep, RemoveFromPolicyRegistryStep,
RemoveFromWorkerRegistryStep, UpdatePoliciesForWorkerStep, UpdateRemainingPoliciesStep,
UpdateWorkerPropertiesStep, WorkerRemovalRequest,
};
sgl-model-gateway/src/core/steps/worker/shared/activate.rs
+12 -8
View File
@@ -1,13 +1,11 @@
//! Unified worker activation step.
use std::sync::Arc;
use async_trait::async_trait;
use tracing::info;
use crate::{
core::Worker,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
core::steps::workflow_data::AnyWorkflowData,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Unified step to activate workers by marking them as healthy.
@@ -16,9 +14,15 @@ use crate::{
pub struct ActivateWorkersStep;
#[async_trait]
impl StepExecutor for ActivateWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
impl StepExecutor<AnyWorkflowData> for ActivateWorkersStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let workers = context
.data
.get_actual_workers()
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
for worker in workers.iter() {
worker.set_healthy(true);
@@ -29,7 +33,7 @@ impl StepExecutor for ActivateWorkersStep {
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}
sgl-model-gateway/src/core/steps/worker/shared/register.rs
+26 -9
View File
@@ -6,10 +6,9 @@ use async_trait::async_trait;
use tracing::debug;
use crate::{
app_context::AppContext,
core::Worker,
core::steps::workflow_data::AnyWorkflowData,
observability::metrics::Metrics,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Unified step to register workers in the registry.
@@ -19,10 +18,21 @@ use crate::{
pub struct RegisterWorkersStep;
#[async_trait]
impl StepExecutor for RegisterWorkersStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
impl StepExecutor<AnyWorkflowData> for RegisterWorkersStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let app_context = context
.data
.get_app_context()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
.clone();
let workers = context
.data
.get_actual_workers()
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let mut worker_ids = Vec::with_capacity(workers.len());
@@ -75,11 +85,18 @@ impl StepExecutor for RegisterWorkersStep {
);
}
context.set("worker_ids", worker_ids);
// Note: worker_ids are stored for potential future use but not persisted
// as they are internal registry identifiers
debug!(
"Registered {} workers with IDs: {:?}",
worker_ids.len(),
worker_ids
);
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}
sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs
+24 -10
View File
@@ -1,14 +1,13 @@
//! Unified policy update step.
use std::{collections::HashMap, sync::Arc};
use std::sync::Arc;
use async_trait::async_trait;
use tracing::{debug, warn};
use crate::{
app_context::AppContext,
core::Worker,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowResult},
core::{steps::workflow_data::AnyWorkflowData, Worker},
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Unified step to update policy registry for registered workers.
@@ -82,11 +81,26 @@ impl UpdatePoliciesStep {
}
#[async_trait]
impl StepExecutor for UpdatePoliciesStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let app_context: Arc<AppContext> = context.get_or_err("app_context")?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context.get_or_err("workers")?;
let labels: Arc<HashMap<String, String>> = context.get_or_err("labels")?;
impl StepExecutor<AnyWorkflowData> for UpdatePoliciesStep {
async fn execute(
&self,
context: &mut WorkflowContext<AnyWorkflowData>,
) -> WorkflowResult<StepResult> {
let app_context = context
.data
.get_app_context()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?
.clone();
let workers = context
.data
.get_actual_workers()
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
let labels = context
.data
.get_labels()
.ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?;
let policy_hint = labels.get("policy").map(|s| s.as_str());
@@ -139,7 +153,7 @@ impl StepExecutor for UpdatePoliciesStep {
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &crate::workflow::WorkflowError) -> bool {
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false
}
}
sgl-model-gateway/src/core/steps/workflow_data.rs
+559
View File
@@ -0,0 +1,559 @@
//! Typed workflow data structures
//!
//! This module defines the typed data structures for all workflows, enabling
//! compile-time type safety and state persistence.
use std::{collections::HashMap, sync::Arc};
use serde::{Deserialize, Serialize};
use super::{
mcp_registration::McpServerConfigRequest, tokenizer_registration::TokenizerConfigRequest,
wasm_module_registration::WasmModuleConfigRequest,
wasm_module_removal::WasmModuleRemovalRequest, worker::local::WorkerRemovalRequest,
};
/// Re-export the protocol types for convenience
pub use crate::protocols::worker_spec::{
WorkerConfigRequest, WorkerUpdateRequest as ProtocolUpdateRequest,
};
use crate::{
app_context::AppContext,
core::{model_card::ModelCard, Worker},
protocols::worker_spec::{
WorkerConfigRequest as ProtocolWorkerConfigRequest,
WorkerUpdateRequest as ProtocolWorkerUpdateRequest,
},
workflow::{WorkflowData, WorkflowError},
};
/// Wrapper for worker list that can be serialized
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct WorkerList {
/// Worker URLs (we can't serialize Arc<dyn Worker>, so we store URLs)
pub worker_urls: Vec<String>,
}
impl WorkerList {
pub fn new() -> Self {
Self {
worker_urls: Vec::new(),
}
}
pub fn from_workers(workers: &[Arc<dyn Worker>]) -> Self {
Self {
worker_urls: workers.iter().map(|w| w.url().to_string()).collect(),
}
}
}
// ============================================================================
// Workflow-specific data types
// ============================================================================
/// Data for tokenizer registration workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TokenizerWorkflowData {
pub config: TokenizerConfigRequest,
pub vocab_size: Option<usize>,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
}
impl WorkflowData for TokenizerWorkflowData {
fn workflow_type() -> &'static str {
"tokenizer_registration"
}
}
impl TokenizerWorkflowData {
/// Validate that all transient fields are properly initialized.
///
/// Call this after deserializing workflow state to ensure runtime fields
/// have been repopulated.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for local worker registration workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LocalWorkerWorkflowData {
pub config: ProtocolWorkerConfigRequest,
pub connection_mode: Option<crate::core::ConnectionMode>,
pub discovered_labels: HashMap<String, String>,
pub dp_info: Option<super::worker::local::DpInfo>,
pub workers: Option<WorkerList>,
pub final_labels: HashMap<String, String>,
/// Detected runtime type (for gRPC workers)
pub detected_runtime_type: Option<String>,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
/// Actual worker objects (transient, not serialized)
#[serde(skip, default)]
pub actual_workers: Option<Vec<Arc<dyn Worker>>>,
}
impl WorkflowData for LocalWorkerWorkflowData {
fn workflow_type() -> &'static str {
"local_worker_registration"
}
}
impl LocalWorkerWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for external worker registration workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExternalWorkerWorkflowData {
pub config: ProtocolWorkerConfigRequest,
/// Discovered model cards from /v1/models endpoint
pub model_cards: Vec<ModelCard>,
pub workers: Option<WorkerList>,
/// Labels for policies (derived from config)
pub labels: HashMap<String, String>,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
/// Actual worker objects (transient, not serialized)
#[serde(skip, default)]
pub actual_workers: Option<Vec<Arc<dyn Worker>>>,
}
impl WorkflowData for ExternalWorkerWorkflowData {
fn workflow_type() -> &'static str {
"external_worker_registration"
}
}
impl ExternalWorkerWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for worker removal workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerRemovalWorkflowData {
pub config: WorkerRemovalRequest,
pub workers_to_remove: Option<WorkerList>,
/// URLs of workers being removed
pub worker_urls: Vec<String>,
/// Model IDs affected by the removal
pub affected_models: std::collections::HashSet<String>,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
/// Actual worker objects to remove (transient, not serialized)
#[serde(skip, default)]
pub actual_workers_to_remove: Option<Vec<Arc<dyn Worker>>>,
}
impl WorkflowData for WorkerRemovalWorkflowData {
fn workflow_type() -> &'static str {
"worker_removal"
}
}
impl WorkerRemovalWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for worker update workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkerUpdateWorkflowData {
pub config: ProtocolWorkerUpdateRequest,
/// URL of worker(s) to update
pub worker_url: String,
/// Whether to update all DP-aware workers with matching prefix
pub dp_aware: bool,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
/// Workers to update (transient, not serialized)
#[serde(skip, default)]
pub workers_to_update: Option<Vec<Arc<dyn Worker>>>,
/// Updated worker objects (transient, not serialized)
#[serde(skip, default)]
pub updated_workers: Option<Vec<Arc<dyn Worker>>>,
}
impl WorkflowData for WorkerUpdateWorkflowData {
fn workflow_type() -> &'static str {
"worker_update"
}
}
impl WorkerUpdateWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for MCP server registration workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpWorkflowData {
pub config: McpServerConfigRequest,
pub validated: bool,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
/// Connected MCP client (transient, not serialized)
#[serde(skip, default)]
pub mcp_client: Option<Arc<rmcp::service::RunningService<rmcp::RoleClient, ()>>>,
}
impl WorkflowData for McpWorkflowData {
fn workflow_type() -> &'static str {
"mcp_registration"
}
}
impl McpWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for WASM module registration workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmRegistrationWorkflowData {
pub config: WasmModuleConfigRequest,
pub wasm_bytes: Option<Vec<u8>>,
/// SHA256 hash of the module file (32 bytes)
pub sha256_hash: Option<[u8; 32]>,
/// File size in bytes
pub file_size_bytes: Option<u64>,
/// UUID assigned to the registered module
pub module_uuid: Option<uuid::Uuid>,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
}
impl WorkflowData for WasmRegistrationWorkflowData {
fn workflow_type() -> &'static str {
"wasm_module_registration"
}
}
impl WasmRegistrationWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
/// Data for WASM module removal workflow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WasmRemovalWorkflowData {
pub config: WasmModuleRemovalRequest,
pub module_id: Option<String>,
/// Application context (transient, must be re-initialized after deserialization)
#[serde(skip, default)]
pub app_context: Option<Arc<AppContext>>,
}
impl WorkflowData for WasmRemovalWorkflowData {
fn workflow_type() -> &'static str {
"wasm_module_removal"
}
}
impl WasmRemovalWorkflowData {
/// Validate that all transient fields are properly initialized.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
if self.app_context.is_none() {
return Err(WorkflowError::ContextValueNotFound(
"app_context not initialized after deserialization".into(),
));
}
Ok(())
}
}
// ============================================================================
// Unified enum for all workflow types
// ============================================================================
/// Macro to generate type-safe accessor methods for AnyWorkflowData variants.
///
/// This reduces boilerplate and ensures consistent error handling across all accessors.
macro_rules! impl_workflow_accessor {
($fn_name:ident, $fn_name_mut:ident, $variant:ident, $ty:ty, $type_name:expr) => {
/// Extract the inner data, returning an error if this is a different variant.
#[must_use = "this returns the result of the operation, without modifying the original"]
pub fn $fn_name(&self) -> Result<&$ty, WorkflowError> {
match self {
AnyWorkflowData::$variant(data) => Ok(data),
_ => Err(WorkflowError::TypeMismatch {
expected: $type_name,
actual: self.concrete_type(),
}),
}
}
/// Extract the inner data mutably, returning an error if this is a different variant.
pub fn $fn_name_mut(&mut self) -> Result<&mut $ty, WorkflowError> {
// Store the type name before the mutable borrow
let actual = self.concrete_type();
match self {
AnyWorkflowData::$variant(data) => Ok(data),
_ => Err(WorkflowError::TypeMismatch {
expected: $type_name,
actual,
}),
}
}
};
}
/// Macro to generate From implementations for AnyWorkflowData variants.
macro_rules! impl_from_workflow_data {
($variant:ident, $ty:ty) => {
impl From<$ty> for AnyWorkflowData {
fn from(data: $ty) -> Self {
AnyWorkflowData::$variant(data)
}
}
};
}
/// Unified workflow data enum covering all workflow types.
///
/// This allows a single `WorkflowEngine<AnyWorkflowData>` to handle all workflows
/// while maintaining type safety at the step level.
///
/// # Type Erasure
///
/// `AnyWorkflowData` implements `WorkflowData` with `workflow_type()` returning `"any"`.
/// This is intentional: the static method cannot know the runtime variant. Use
/// [`concrete_type()`](Self::concrete_type) to get the actual workflow type at runtime.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AnyWorkflowData {
Tokenizer(TokenizerWorkflowData),
LocalWorker(LocalWorkerWorkflowData),
ExternalWorker(ExternalWorkerWorkflowData),
WorkerRemoval(WorkerRemovalWorkflowData),
WorkerUpdate(WorkerUpdateWorkflowData),
Mcp(McpWorkflowData),
WasmRegistration(WasmRegistrationWorkflowData),
WasmRemoval(WasmRemovalWorkflowData),
}
impl WorkflowData for AnyWorkflowData {
/// Returns `"any"` as this is a type-erased container.
///
/// Use [`concrete_type()`](Self::concrete_type) to get the actual workflow type at runtime.
fn workflow_type() -> &'static str {
"any"
}
}
// Generate From implementations for ergonomic construction
impl_from_workflow_data!(Tokenizer, TokenizerWorkflowData);
impl_from_workflow_data!(LocalWorker, LocalWorkerWorkflowData);
impl_from_workflow_data!(ExternalWorker, ExternalWorkerWorkflowData);
impl_from_workflow_data!(WorkerRemoval, WorkerRemovalWorkflowData);
impl_from_workflow_data!(WorkerUpdate, WorkerUpdateWorkflowData);
impl_from_workflow_data!(Mcp, McpWorkflowData);
impl_from_workflow_data!(WasmRegistration, WasmRegistrationWorkflowData);
impl_from_workflow_data!(WasmRemoval, WasmRemovalWorkflowData);
impl AnyWorkflowData {
/// Get the concrete workflow type name at runtime.
///
/// Unlike the static `workflow_type()` method, this returns the actual
/// type of the contained workflow data.
#[must_use]
pub fn concrete_type(&self) -> &'static str {
match self {
AnyWorkflowData::Tokenizer(_) => TokenizerWorkflowData::workflow_type(),
AnyWorkflowData::LocalWorker(_) => LocalWorkerWorkflowData::workflow_type(),
AnyWorkflowData::ExternalWorker(_) => ExternalWorkerWorkflowData::workflow_type(),
AnyWorkflowData::WorkerRemoval(_) => WorkerRemovalWorkflowData::workflow_type(),
AnyWorkflowData::WorkerUpdate(_) => WorkerUpdateWorkflowData::workflow_type(),
AnyWorkflowData::Mcp(_) => McpWorkflowData::workflow_type(),
AnyWorkflowData::WasmRegistration(_) => WasmRegistrationWorkflowData::workflow_type(),
AnyWorkflowData::WasmRemoval(_) => WasmRemovalWorkflowData::workflow_type(),
}
}
// Generate all accessor methods using the macro
impl_workflow_accessor!(
as_tokenizer,
as_tokenizer_mut,
Tokenizer,
TokenizerWorkflowData,
"tokenizer_registration"
);
impl_workflow_accessor!(
as_local_worker,
as_local_worker_mut,
LocalWorker,
LocalWorkerWorkflowData,
"local_worker_registration"
);
impl_workflow_accessor!(
as_external_worker,
as_external_worker_mut,
ExternalWorker,
ExternalWorkerWorkflowData,
"external_worker_registration"
);
impl_workflow_accessor!(
as_worker_removal,
as_worker_removal_mut,
WorkerRemoval,
WorkerRemovalWorkflowData,
"worker_removal"
);
impl_workflow_accessor!(
as_worker_update,
as_worker_update_mut,
WorkerUpdate,
WorkerUpdateWorkflowData,
"worker_update"
);
impl_workflow_accessor!(as_mcp, as_mcp_mut, Mcp, McpWorkflowData, "mcp_registration");
impl_workflow_accessor!(
as_wasm_registration,
as_wasm_registration_mut,
WasmRegistration,
WasmRegistrationWorkflowData,
"wasm_module_registration"
);
impl_workflow_accessor!(
as_wasm_removal,
as_wasm_removal_mut,
WasmRemoval,
WasmRemovalWorkflowData,
"wasm_module_removal"
);
// ========================================================================
// Helper methods for shared worker steps
// ========================================================================
/// Get app_context from any workflow data type that has it.
#[must_use]
pub fn get_app_context(&self) -> Option<&Arc<AppContext>> {
match self {
AnyWorkflowData::Tokenizer(d) => d.app_context.as_ref(),
AnyWorkflowData::LocalWorker(d) => d.app_context.as_ref(),
AnyWorkflowData::ExternalWorker(d) => d.app_context.as_ref(),
AnyWorkflowData::WorkerRemoval(d) => d.app_context.as_ref(),
AnyWorkflowData::WorkerUpdate(d) => d.app_context.as_ref(),
AnyWorkflowData::Mcp(d) => d.app_context.as_ref(),
AnyWorkflowData::WasmRegistration(d) => d.app_context.as_ref(),
AnyWorkflowData::WasmRemoval(d) => d.app_context.as_ref(),
}
}
/// Get actual workers from local or external worker workflows.
#[must_use]
pub fn get_actual_workers(&self) -> Option<&Vec<Arc<dyn Worker>>> {
match self {
AnyWorkflowData::LocalWorker(d) => d.actual_workers.as_ref(),
AnyWorkflowData::ExternalWorker(d) => d.actual_workers.as_ref(),
_ => None,
}
}
/// Set actual workers for local or external worker workflows.
pub fn set_actual_workers(
&mut self,
workers: Vec<Arc<dyn Worker>>,
) -> Result<(), WorkflowError> {
match self {
AnyWorkflowData::LocalWorker(d) => {
d.workers = Some(WorkerList::from_workers(&workers));
d.actual_workers = Some(workers);
Ok(())
}
AnyWorkflowData::ExternalWorker(d) => {
d.workers = Some(WorkerList::from_workers(&workers));
d.actual_workers = Some(workers);
Ok(())
}
_ => Err(WorkflowError::TypeMismatch {
expected: "LocalWorker or ExternalWorker",
actual: self.concrete_type(),
}),
}
}
/// Get labels for policy configuration (from local or external worker workflows).
#[must_use]
pub fn get_labels(&self) -> Option<&HashMap<String, String>> {
match self {
AnyWorkflowData::LocalWorker(d) => Some(&d.final_labels),
AnyWorkflowData::ExternalWorker(d) => Some(&d.labels),
_ => None,
}
}
/// Validate that all transient fields are properly initialized.
///
/// Call this after deserializing workflow state to ensure runtime fields
/// have been repopulated.
pub fn validate_initialized(&self) -> Result<(), WorkflowError> {
match self {
AnyWorkflowData::Tokenizer(d) => d.validate_initialized(),
AnyWorkflowData::LocalWorker(d) => d.validate_initialized(),
AnyWorkflowData::ExternalWorker(d) => d.validate_initialized(),
AnyWorkflowData::WorkerRemoval(d) => d.validate_initialized(),
AnyWorkflowData::WorkerUpdate(d) => d.validate_initialized(),
AnyWorkflowData::Mcp(d) => d.validate_initialized(),
AnyWorkflowData::WasmRegistration(d) => d.validate_initialized(),
AnyWorkflowData::WasmRemoval(d) => d.validate_initialized(),
}
}
}
sgl-model-gateway/src/server.rs
+9 -9
View File
@@ -20,14 +20,14 @@ use tokio::{signal, spawn};
use tracing::{debug, error, info, warn, Level};
use crate::{
app_context::AppContext,
app_context::{AppContext, AppWorkflowEngine},
config::{RouterConfig, RoutingMode},
core::{
job_queue::{JobQueue, JobQueueConfig},
steps::{
create_external_worker_registration_workflow, create_mcp_registration_workflow,
create_tokenizer_registration_workflow, create_wasm_module_registration_workflow,
create_wasm_module_removal_workflow, create_worker_registration_workflow,
create_external_worker_workflow, create_local_worker_workflow,
create_mcp_registration_workflow, create_tokenizer_registration_workflow,
create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
create_worker_removal_workflow, create_worker_update_workflow,
},
worker::WorkerType,
@@ -56,7 +56,7 @@ use crate::{
routers::{conversations, parse, router_manager::RouterManager, tokenize, RouterTrait},
service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},
workflow::{LoggingSubscriber, WorkflowEngine},
workflow::LoggingSubscriber,
};
#[derive(Clone)]
pub struct AppState {
@@ -731,7 +731,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
.expect("JobQueue should only be initialized once");
// Initialize workflow engine and register workflows
let engine = Arc::new(WorkflowEngine::new());
let engine = Arc::new(AppWorkflowEngine::new());
engine
.event_bus()
@@ -739,10 +739,10 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
.await;
engine
.register_workflow(create_worker_registration_workflow(&config.router_config))
.expect("worker_registration workflow should be valid");
.register_workflow(create_local_worker_workflow(&config.router_config))
.expect("local_worker_registration workflow should be valid");
engine
.register_workflow(create_external_worker_registration_workflow())
.register_workflow(create_external_worker_workflow())
.expect("external_worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
sgl-model-gateway/src/workflow/definition.rs
+13 -13
View File
@@ -8,25 +8,25 @@ use std::{
use super::{
executor::StepExecutor,
types::{FailureAction, RetryPolicy, StepId, WorkflowId},
types::{FailureAction, RetryPolicy, StepId, WorkflowData, WorkflowId},
};
/// Definition of a single step within a workflow
pub struct StepDefinition {
pub struct StepDefinition<D: WorkflowData> {
pub id: StepId,
pub name: String,
pub executor: Arc<dyn StepExecutor>,
pub executor: Arc<dyn StepExecutor<D>>,
pub retry_policy: Option<RetryPolicy>,
pub timeout: Option<Duration>,
pub on_failure: FailureAction,
pub depends_on: Vec<StepId>,
}
impl StepDefinition {
impl<D: WorkflowData> StepDefinition<D> {
pub fn new(
id: impl Into<String>,
name: impl Into<String>,
executor: Arc<dyn StepExecutor>,
executor: Arc<dyn StepExecutor<D>>,
) -> Self {
Self {
id: StepId::new(id.into()),
@@ -64,10 +64,10 @@ impl StepDefinition {
}
/// Complete workflow definition
pub struct WorkflowDefinition {
pub struct WorkflowDefinition<D: WorkflowData> {
pub id: WorkflowId,
pub name: String,
pub steps: Vec<StepDefinition>,
pub steps: Vec<StepDefinition<D>>,
pub default_retry_policy: RetryPolicy,
pub default_timeout: Duration,
/// Pre-computed reverse dependencies: step_id -> indices of steps that depend on it
@@ -76,7 +76,7 @@ pub struct WorkflowDefinition {
initial_step_indices: Vec<usize>,
}
impl WorkflowDefinition {
impl<D: WorkflowData> WorkflowDefinition<D> {
pub fn new(id: impl Into<String>, name: impl Into<String>) -> Self {
Self {
id: WorkflowId::new(id.into()),
@@ -89,7 +89,7 @@ impl WorkflowDefinition {
}
}
pub fn add_step(mut self, step: StepDefinition) -> Self {
pub fn add_step(mut self, step: StepDefinition<D>) -> Self {
self.steps.push(step);
self
}
@@ -105,14 +105,14 @@ impl WorkflowDefinition {
}
/// Get the retry policy for a step (step-specific or default)
pub fn get_retry_policy<'a>(&'a self, step: &'a StepDefinition) -> &'a RetryPolicy {
pub fn get_retry_policy<'a>(&'a self, step: &'a StepDefinition<D>) -> &'a RetryPolicy {
step.retry_policy
.as_ref()
.unwrap_or(&self.default_retry_policy)
}
/// Get the timeout for a step (step-specific or default)
pub fn get_timeout(&self, step: &StepDefinition) -> Duration {
pub fn get_timeout(&self, step: &StepDefinition<D>) -> Duration {
step.timeout.unwrap_or(self.default_timeout)
}
@@ -124,7 +124,7 @@ impl WorkflowDefinition {
/// On success, pre-computes reverse dependencies for O(1) dependent lookup.
pub fn validate(&mut self) -> Result<(), String> {
// Build HashMap for O(1) lookup instead of O(n) linear search
let steps_map: HashMap<&StepId, &StepDefinition> =
let steps_map: HashMap<&StepId, &StepDefinition<D>> =
self.steps.iter().map(|s| (&s.id, s)).collect();
// Check all dependencies exist
@@ -177,7 +177,7 @@ impl WorkflowDefinition {
/// DFS helper for cycle detection with O(1) HashMap lookup
fn has_cycle<'a>(
step_id: &'a StepId,
steps_map: &HashMap<&'a StepId, &'a StepDefinition>,
steps_map: &HashMap<&'a StepId, &'a StepDefinition<D>>,
visited: &mut HashSet<&'a StepId>,
rec_stack: &mut HashSet<&'a StepId>,
) -> bool {
sgl-model-gateway/src/workflow/engine.rs
+53 -26
View File
@@ -6,6 +6,7 @@
use std::{
collections::{HashMap, HashSet, VecDeque},
marker::PhantomData,
sync::{
atomic::{AtomicUsize, Ordering},
Arc,
@@ -24,7 +25,7 @@ use tokio::{
use super::{
definition::{StepDefinition, WorkflowDefinition},
event::{EventBus, WorkflowEvent},
state::WorkflowStateStore,
state::{InMemoryStore, StateStore},
types::*,
};
@@ -101,6 +102,11 @@ impl Backoff for LinearBackoff {
/// Main workflow execution engine
///
/// # Type Parameters
///
/// * `D` - The workflow data type that implements `WorkflowData`
/// * `S` - The state store implementation (defaults to `InMemoryStore<D>`)
///
/// # Graceful Shutdown
///
/// The engine supports graceful shutdown via [`shutdown()`](Self::shutdown):
@@ -115,9 +121,9 @@ impl Backoff for LinearBackoff {
/// engine.force_cancel_all().await;
/// }
/// ```
pub struct WorkflowEngine {
definitions: Arc<RwLock<HashMap<WorkflowId, Arc<WorkflowDefinition>>>>,
state_store: WorkflowStateStore,
pub struct WorkflowEngine<D: WorkflowData, S: StateStore<D> = InMemoryStore<D>> {
definitions: Arc<RwLock<HashMap<WorkflowId, Arc<WorkflowDefinition<D>>>>>,
state_store: S,
event_bus: Arc<EventBus>,
/// Shutdown signal sender - when true, engine is shutting down
shutdown_tx: Arc<watch::Sender<bool>>,
@@ -125,18 +131,27 @@ pub struct WorkflowEngine {
shutdown_rx: watch::Receiver<bool>,
/// Count of active workflow executions
active_workflows: Arc<AtomicUsize>,
_phantom: PhantomData<D>,
}
impl WorkflowEngine {
impl<D: WorkflowData> WorkflowEngine<D, InMemoryStore<D>> {
pub fn new() -> Self {
Self::with_store(InMemoryStore::new())
}
}
impl<D: WorkflowData, S: StateStore<D> + 'static> WorkflowEngine<D, S> {
/// Create a new workflow engine with a custom state store
pub fn with_store(state_store: S) -> Self {
let (shutdown_tx, shutdown_rx) = watch::channel(false);
Self {
definitions: Arc::new(RwLock::new(HashMap::new())),
state_store: WorkflowStateStore::new(),
state_store,
event_bus: Arc::new(EventBus::new()),
shutdown_tx: Arc::new(shutdown_tx),
shutdown_rx,
active_workflows: Arc::new(AtomicUsize::new(0)),
_phantom: PhantomData,
}
}
@@ -284,7 +299,7 @@ impl WorkflowEngine {
}
/// Register a workflow definition
pub fn register_workflow(&self, mut definition: WorkflowDefinition) -> Result<(), String> {
pub fn register_workflow(&self, mut definition: WorkflowDefinition<D>) -> Result<(), String> {
// Validate DAG and build dependency graph once at registration
definition.validate()?;
@@ -299,7 +314,7 @@ impl WorkflowEngine {
}
/// Get the state store
pub fn state_store(&self) -> &WorkflowStateStore {
pub fn state_store(&self) -> &S {
&self.state_store
}
@@ -309,7 +324,7 @@ impl WorkflowEngine {
pub async fn start_workflow(
&self,
definition_id: WorkflowId,
context: WorkflowContext,
data: D,
) -> WorkflowResult<WorkflowInstanceId> {
// Guard increments counter and decrements on drop unless committed.
// This handles all error paths automatically.
@@ -326,10 +341,9 @@ impl WorkflowEngine {
.cloned()
.ok_or_else(|| WorkflowError::DefinitionNotFound(definition_id.clone()))?;
let instance_id = context.instance_id;
let mut state = WorkflowState::new(instance_id, definition_id.clone());
let instance_id = WorkflowInstanceId::new();
let mut state = WorkflowState::new(instance_id, definition_id.clone(), data);
state.status = WorkflowStatus::Running;
state.context = context;
for step in &definition.steps {
state
@@ -369,7 +383,7 @@ impl WorkflowEngine {
async fn execute_workflow(
&self,
instance_id: WorkflowInstanceId,
definition: Arc<WorkflowDefinition>,
definition: Arc<WorkflowDefinition<D>>,
) -> WorkflowResult<()> {
let start_time = std::time::Instant::now();
let step_count = definition.steps.len();
@@ -572,8 +586,8 @@ impl WorkflowEngine {
async fn execute_step_with_retry(
&self,
instance_id: WorkflowInstanceId,
step: &StepDefinition,
definition: &WorkflowDefinition,
step: &StepDefinition<D>,
definition: &WorkflowDefinition<D>,
) -> WorkflowResult<StepResult> {
let retry_policy = definition.get_retry_policy(step);
let step_timeout = definition.get_timeout(step);
@@ -624,7 +638,7 @@ impl WorkflowEngine {
let step_duration = step_start.elapsed();
self.state_store.update(instance_id, |s| {
s.context = std::mem::replace(&mut context, WorkflowContext::new(instance_id));
s.context = context.clone();
})?;
match result {
@@ -763,7 +777,7 @@ impl WorkflowEngine {
}
/// Get workflow status
pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState> {
pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
self.state_store.load(instance_id)
}
@@ -776,6 +790,7 @@ impl WorkflowEngine {
shutdown_tx: Arc::clone(&self.shutdown_tx),
shutdown_rx: self.shutdown_rx.clone(),
active_workflows: Arc::clone(&self.active_workflows),
_phantom: PhantomData,
}
}
}
@@ -794,13 +809,13 @@ impl Drop for ActiveWorkflowGuard {
/// RAII guard for start_workflow that increments on creation and decrements on drop
/// unless commit() is called. Handles all error paths automatically.
struct StartGuard<'a> {
engine: &'a WorkflowEngine,
struct StartGuard<'a, D: WorkflowData, S: StateStore<D> + 'static> {
engine: &'a WorkflowEngine<D, S>,
committed: bool,
}
impl<'a> StartGuard<'a> {
fn new(engine: &'a WorkflowEngine) -> Self {
impl<'a, D: WorkflowData, S: StateStore<D> + 'static> StartGuard<'a, D, S> {
fn new(engine: &'a WorkflowEngine<D, S>) -> Self {
engine.active_workflows.fetch_add(1, Ordering::AcqRel);
Self {
engine,
@@ -813,7 +828,7 @@ impl<'a> StartGuard<'a> {
}
}
impl Drop for StartGuard<'_> {
impl<D: WorkflowData, S: StateStore<D> + 'static> Drop for StartGuard<'_, D, S> {
fn drop(&mut self) {
if !self.committed {
self.engine.workflow_finished();
@@ -821,23 +836,35 @@ impl Drop for StartGuard<'_> {
}
}
impl Clone for WorkflowEngine {
/// Clone implementation for internal use.
///
/// **Note**: This creates a shallow clone that shares state with the original engine.
/// Both engines will share the same:
/// - Workflow definitions
/// - State store
/// - Event bus
/// - Shutdown signal
/// - Active workflow counter
///
/// This is intentional for spawning async tasks that need access to the engine.
/// For most use cases, prefer sharing the engine via `Arc<WorkflowEngine>` rather
/// than cloning.
impl<D: WorkflowData, S: StateStore<D> + 'static> Clone for WorkflowEngine<D, S> {
fn clone(&self) -> Self {
self.clone_for_execution()
}
}
impl Default for WorkflowEngine {
impl<D: WorkflowData> Default for WorkflowEngine<D, InMemoryStore<D>> {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for WorkflowEngine {
impl<D: WorkflowData, S: StateStore<D> + 'static> std::fmt::Debug for WorkflowEngine<D, S> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("WorkflowEngine")
.field("definitions_count", &self.definitions.read().len())
.field("state_count", &self.state_store.count())
.finish()
}
}
sgl-model-gateway/src/workflow/executor.rs
+40 -17
View File
@@ -2,13 +2,13 @@
use async_trait::async_trait;
use super::types::{StepResult, WorkflowContext, WorkflowError, WorkflowResult};
use super::types::{StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult};
/// Trait for executing individual workflow steps
#[async_trait]
pub trait StepExecutor: Send + Sync {
pub trait StepExecutor<D: WorkflowData>: Send + Sync {
/// Execute the step with the given context
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult>;
async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult>;
/// Check if an error is retry-able
///
@@ -22,7 +22,7 @@ pub trait StepExecutor: Send + Sync {
///
/// This hook allows steps to perform cleanup or additional actions
/// after successful execution.
async fn on_success(&self, _context: &WorkflowContext) -> WorkflowResult<()> {
async fn on_success(&self, _context: &WorkflowContext<D>) -> WorkflowResult<()> {
Ok(())
}
@@ -32,7 +32,7 @@ pub trait StepExecutor: Send + Sync {
/// when the step cannot complete successfully.
async fn on_failure(
&self,
_context: &WorkflowContext,
_context: &WorkflowContext<D>,
_error: &WorkflowError,
) -> WorkflowResult<()> {
Ok(())
@@ -40,59 +40,82 @@ pub trait StepExecutor: Send + Sync {
}
/// Simple function-based step executor
pub struct FunctionStep<F>
pub struct FunctionStep<D, F>
where
D: WorkflowData,
F: Fn(
&mut WorkflowContext,
&mut WorkflowContext<D>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = WorkflowResult<StepResult>> + Send + '_>,
> + Send
+ Sync,
{
func: F,
_phantom: std::marker::PhantomData<D>,
}
impl<F> FunctionStep<F>
impl<D, F> FunctionStep<D, F>
where
D: WorkflowData,
F: Fn(
&mut WorkflowContext,
&mut WorkflowContext<D>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = WorkflowResult<StepResult>> + Send + '_>,
> + Send
+ Sync,
{
pub fn new(func: F) -> Self {
Self { func }
Self {
func,
_phantom: std::marker::PhantomData,
}
}
}
#[async_trait]
impl<F> StepExecutor for FunctionStep<F>
impl<D, F> StepExecutor<D> for FunctionStep<D, F>
where
D: WorkflowData,
F: Fn(
&mut WorkflowContext,
&mut WorkflowContext<D>,
) -> std::pin::Pin<
Box<dyn std::future::Future<Output = WorkflowResult<StepResult>> + Send + '_>,
> + Send
+ Sync,
{
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
async fn execute(&self, context: &mut WorkflowContext<D>) -> WorkflowResult<StepResult> {
(self.func)(context).await
}
}
#[cfg(test)]
mod tests {
use serde::{Deserialize, Serialize};
use super::*;
use crate::workflow::types::WorkflowInstanceId;
#[derive(Debug, Clone, Serialize, Deserialize)]
struct TestData {
value: i32,
}
impl WorkflowData for TestData {
fn workflow_type() -> &'static str {
"test"
}
}
struct TestStep {
should_succeed: bool,
}
#[async_trait]
impl StepExecutor for TestStep {
async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
impl StepExecutor<TestData> for TestStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestData>,
) -> WorkflowResult<StepResult> {
if self.should_succeed {
Ok(StepResult::Success)
} else {
@@ -109,7 +132,7 @@ mod tests {
let step = TestStep {
should_succeed: true,
};
let mut context = WorkflowContext::new(WorkflowInstanceId::new());
let mut context = WorkflowContext::new(WorkflowInstanceId::new(), TestData { value: 42 });
let result = step.execute(&mut context).await;
assert!(result.is_ok());
@@ -121,7 +144,7 @@ mod tests {
let step = TestStep {
should_succeed: false,
};
let mut context = WorkflowContext::new(WorkflowInstanceId::new());
let mut context = WorkflowContext::new(WorkflowInstanceId::new(), TestData { value: 42 });
let result = step.execute(&mut context).await;
assert!(result.is_err());
sgl-model-gateway/src/workflow/mod.rs
+1 -1
View File
@@ -11,5 +11,5 @@ pub use definition::{StepDefinition, WorkflowDefinition};
pub use engine::WorkflowEngine;
pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent};
pub use executor::{FunctionStep, StepExecutor};
pub use state::WorkflowStateStore;
pub use state::{InMemoryStore, StateStore};
pub use types::*;
sgl-model-gateway/src/workflow/state.rs
+128 -133
View File
@@ -1,113 +1,64 @@
//! Workflow state management
use std::{collections::HashMap, sync::Arc};
use std::{collections::HashMap, marker::PhantomData, sync::Arc, time::Duration};
use parking_lot::RwLock;
use super::types::{
WorkflowError, WorkflowInstanceId, WorkflowResult, WorkflowState, WorkflowStatus,
WorkflowContext, WorkflowData, WorkflowError, WorkflowInstanceId, WorkflowResult,
WorkflowState, WorkflowStatus,
};
/// Trait for workflow state persistence.
///
/// Implement this trait to provide custom storage backends (e.g., PostgreSQL, Redis).
/// The default implementation is `InMemoryStore` which keeps state in memory.
pub trait StateStore<D: WorkflowData>: Send + Sync + Clone {
/// Save workflow state
fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()>;
/// Load workflow state by instance ID
fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>>;
/// Update workflow state using a closure
fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where
F: FnOnce(&mut WorkflowState<D>);
/// Delete workflow state
fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>;
/// List all active workflows (Running or Pending)
fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
/// List all workflows
fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>>;
/// Check if workflow is cancelled without loading full state
fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<bool>;
/// Clean up old completed/failed/cancelled workflows beyond a time threshold
fn cleanup_old_workflows(&self, ttl: Duration) -> usize;
/// Get just the workflow context without cloning the entire state
fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>>;
}
/// In-memory state storage for workflow instances
#[derive(Clone)]
pub struct WorkflowStateStore {
states: Arc<RwLock<HashMap<WorkflowInstanceId, WorkflowState>>>,
pub struct InMemoryStore<D: WorkflowData> {
states: Arc<RwLock<HashMap<WorkflowInstanceId, WorkflowState<D>>>>,
_phantom: PhantomData<D>,
}
impl WorkflowStateStore {
impl<D: WorkflowData> InMemoryStore<D> {
pub fn new() -> Self {
Self {
states: Arc::new(RwLock::new(HashMap::new())),
_phantom: PhantomData,
}
}
/// Save workflow state
///
/// # Note
///
/// This emits a debug log if the workflow context contains unserializable data,
/// which would be lost if state persistence is later implemented.
pub fn save(&self, state: WorkflowState) -> WorkflowResult<()> {
if state.context.has_unserializable_data() {
tracing::debug!(
instance_id = %state.instance_id,
data_count = state.context.data_len(),
"Saving workflow state with {} unserializable context entries. \
This data cannot be persisted and will be lost on restart.",
state.context.data_len()
);
}
self.states.write().insert(state.instance_id, state);
Ok(())
}
/// Load workflow state by instance ID
pub fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState> {
self.states
.read()
.get(&instance_id)
.cloned()
.ok_or(WorkflowError::NotFound(instance_id))
}
/// List all active workflows (Running or Pending)
pub fn list_active(&self) -> WorkflowResult<Vec<WorkflowState>> {
let states = self.states.read();
Ok(states
.values()
.filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
.cloned()
.collect())
}
/// List all workflows
pub fn list_all(&self) -> WorkflowResult<Vec<WorkflowState>> {
let states = self.states.read();
Ok(states.values().cloned().collect())
}
/// Delete workflow state
pub fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.states.write().remove(&instance_id);
Ok(())
}
/// Update workflow state using a closure
pub fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where
F: FnOnce(&mut WorkflowState),
{
let mut states = self.states.write();
let state = states
.get_mut(&instance_id)
.ok_or(WorkflowError::NotFound(instance_id))?;
f(state);
state.updated_at = chrono::Utc::now();
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
@@ -122,23 +73,91 @@ impl WorkflowStateStore {
self.states.read().len()
}
/// Clean up old completed/failed/cancelled workflows beyond a time threshold
///
/// This prevents unbounded memory growth by removing workflow states that
/// have been in a terminal state (Completed, Failed, Cancelled) for longer
/// than the specified TTL (time-to-live).
///
/// Active workflows (Running, Pending, Paused) are never cleaned up.
///
/// # Arguments
///
/// * `ttl` - Time-to-live for terminal workflows. Workflows in terminal states
/// older than this will be removed.
///
/// # Returns
///
/// The number of workflow states removed.
pub fn cleanup_old_workflows(&self, ttl: std::time::Duration) -> usize {
/// Clean up a specific completed workflow immediately
pub fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
let mut states = self.states.write();
if let Some(state) = states.get(&instance_id) {
if matches!(
state.status,
WorkflowStatus::Completed | WorkflowStatus::Failed | WorkflowStatus::Cancelled
) {
states.remove(&instance_id);
return true;
}
}
false
}
}
impl<D: WorkflowData> Default for InMemoryStore<D> {
fn default() -> Self {
Self::new()
}
}
impl<D: WorkflowData> StateStore<D> for InMemoryStore<D> {
fn save(&self, state: WorkflowState<D>) -> WorkflowResult<()> {
self.states.write().insert(state.instance_id, state);
Ok(())
}
fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowState<D>> {
self.states
.read()
.get(&instance_id)
.cloned()
.ok_or(WorkflowError::NotFound(instance_id))
}
fn list_active(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states
.values()
.filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending))
.cloned()
.collect())
}
fn list_all(&self) -> WorkflowResult<Vec<WorkflowState<D>>> {
let states = self.states.read();
Ok(states.values().cloned().collect())
}
fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()> {
self.states.write().remove(&instance_id);
Ok(())
}
fn update<F>(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()>
where
F: FnOnce(&mut WorkflowState<D>),
{
let mut states = self.states.write();
let state = states
.get_mut(&instance_id)
.ok_or(WorkflowError::NotFound(instance_id))?;
f(state);
state.updated_at = chrono::Utc::now();
Ok(())
}
fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<WorkflowContext<D>> {
self.states
.read()
.get(&instance_id)
.map(|s| s.context.clone())
.ok_or(WorkflowError::NotFound(instance_id))
}
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))
}
fn cleanup_old_workflows(&self, ttl: Duration) -> usize {
let now = chrono::Utc::now();
let mut states = self.states.write();
let initial_count = states.len();
@@ -170,28 +189,4 @@ impl WorkflowStateStore {
}
removed_count
}
/// Clean up a specific completed workflow immediately
///
/// This is useful for cleaning up workflows right after they complete
/// when you know they won't be queried again.
pub fn cleanup_if_terminal(&self, instance_id: WorkflowInstanceId) -> bool {
let mut states = self.states.write();
if let Some(state) = states.get(&instance_id) {
if matches!(
state.status,
WorkflowStatus::Completed | WorkflowStatus::Failed | WorkflowStatus::Cancelled
) {
states.remove(&instance_id);
return true;
}
}
false
}
}
impl Default for WorkflowStateStore {
fn default() -> Self {
Self::new()
}
}
sgl-model-gateway/src/workflow/types.rs
+58 -58
View File
@@ -1,11 +1,36 @@
//! Core workflow types and definitions
use std::{collections::HashMap, fmt, sync::Arc, time::Duration};
use std::{collections::HashMap, fmt, time::Duration};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use uuid::Uuid;
/// Trait for workflow data that can be passed through workflow steps.
///
/// Implementing this trait allows your data type to be used as the typed
/// context for a workflow. The data must be serializable for state persistence.
///
/// # Example
///
/// ```ignore
/// #[derive(Debug, Clone, Serialize, Deserialize)]
/// pub struct MyWorkflowData {
/// pub config: MyConfig,
/// pub result: Option<MyResult>,
/// #[serde(skip, default)]
/// pub app_context: Option<Arc<AppContext>>,
/// }
///
/// impl WorkflowData for MyWorkflowData {
/// fn workflow_type() -> &'static str { "my_workflow" }
/// }
/// ```
pub trait WorkflowData: Serialize + DeserializeOwned + Send + Sync + Clone + 'static {
/// Human-readable name for logging and identification
fn workflow_type() -> &'static str;
}
/// Unique identifier for a workflow definition
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct WorkflowId(String);
@@ -147,19 +172,23 @@ impl Default for StepState {
/// Workflow instance state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowState {
#[serde(bound(
serialize = "D: Serialize",
deserialize = "D: serde::de::DeserializeOwned"
))]
pub struct WorkflowState<D: WorkflowData> {
pub instance_id: WorkflowInstanceId,
pub definition_id: WorkflowId,
pub status: WorkflowStatus,
pub current_step: Option<StepId>,
pub step_states: HashMap<StepId, StepState>,
pub context: WorkflowContext,
pub context: WorkflowContext<D>,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
}
impl WorkflowState {
pub fn new(instance_id: WorkflowInstanceId, definition_id: WorkflowId) -> Self {
impl<D: WorkflowData> WorkflowState<D> {
pub fn new(instance_id: WorkflowInstanceId, definition_id: WorkflowId, data: D) -> Self {
let now = Utc::now();
Self {
instance_id,
@@ -167,71 +196,36 @@ impl WorkflowState {
status: WorkflowStatus::Pending,
current_step: None,
step_states: HashMap::new(),
context: WorkflowContext::new(instance_id),
context: WorkflowContext::new(instance_id, data),
created_at: now,
updated_at: now,
}
}
}
/// Shared context passed between workflow steps
/// Shared context passed between workflow steps.
///
/// # Serialization Warning
/// The context contains typed workflow data that is fully serializable,
/// enabling state persistence and workflow recovery.
///
/// The `data` field contains type-erased values that cannot be serialized.
/// This means workflow context is **not preserved** across:
/// - Process restarts
/// - State persistence to disk
/// - Network serialization
/// # Type Parameter
///
/// The workflow engine only supports **in-memory execution**. If you need
/// durable workflows, consider implementing a custom serializable context type.
/// `D` - The workflow-specific data type implementing `WorkflowData`.
/// This type holds all the state needed by workflow steps and must be
/// serializable (except for fields marked with `#[serde(skip)]`).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WorkflowContext {
#[serde(bound(
serialize = "D: Serialize",
deserialize = "D: serde::de::DeserializeOwned"
))]
pub struct WorkflowContext<D: WorkflowData> {
pub instance_id: WorkflowInstanceId,
#[serde(skip)]
data: HashMap<String, Arc<dyn std::any::Any + Send + Sync>>,
pub data: D,
}
impl WorkflowContext {
pub fn new(instance_id: WorkflowInstanceId) -> Self {
Self {
instance_id,
data: HashMap::new(),
}
}
/// Store a value in the context (will be wrapped in Arc)
pub fn set<T: Send + Sync + 'static>(&mut self, key: impl Into<String>, value: T) {
self.data.insert(key.into(), Arc::new(value));
}
/// Store an Arc directly without double-wrapping
pub fn set_arc<T: Send + Sync + 'static>(&mut self, key: impl Into<String>, value: Arc<T>) {
self.data.insert(key.into(), value);
}
/// Retrieve a value from the context
pub fn get<T: Send + Sync + 'static>(&self, key: &str) -> Option<Arc<T>> {
self.data
.get(key)
.and_then(|v| v.clone().downcast::<T>().ok())
}
/// Retrieve a value from the context, returning an error if not found
pub fn get_or_err<T: Send + Sync + 'static>(&self, key: &str) -> Result<Arc<T>, WorkflowError> {
self.get::<T>(key)
.ok_or_else(|| WorkflowError::ContextValueNotFound(key.to_string()))
}
/// Check if the context has any data that would be lost during serialization
pub fn has_unserializable_data(&self) -> bool {
!self.data.is_empty()
}
/// Get the number of context entries (useful for debugging)
pub fn data_len(&self) -> usize {
self.data.len()
impl<D: WorkflowData> WorkflowContext<D> {
pub fn new(instance_id: WorkflowInstanceId, data: D) -> Self {
Self { instance_id, data }
}
}
@@ -270,6 +264,12 @@ pub enum WorkflowError {
#[error("Context value not found: {0}")]
ContextValueNotFound(String),
#[error("Type mismatch: expected {expected}, got {actual}")]
TypeMismatch {
expected: &'static str,
actual: &'static str,
},
#[error("Engine is shutting down, not accepting new workflows")]
ShuttingDown,
}
sgl-model-gateway/tests/common/mod.rs
+6 -6
View File
@@ -358,12 +358,12 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_worker_registration_workflow(&config))
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
@@ -491,12 +491,12 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_worker_registration_workflow(&config))
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
@@ -624,12 +624,12 @@ pub async fn create_test_context_with_mcp_config(
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_worker_registration_workflow(&config))
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
sgl-model-gateway/tests/wasm_test.rs
+19 -14
View File
@@ -109,12 +109,12 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_worker_registration_workflow(&config))
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
@@ -685,8 +685,11 @@ async fn test_wasm_module_execution() {
// Create workflow context for registration
use smg::{
core::steps::WasmModuleConfigRequest,
workflow::{WorkflowContext, WorkflowId, WorkflowInstanceId},
core::steps::{
workflow_data::{AnyWorkflowData, WasmRegistrationWorkflowData},
WasmModuleConfigRequest,
},
workflow::WorkflowId,
};
let descriptor = WasmModuleDescriptor {
@@ -700,16 +703,18 @@ async fn test_wasm_module_execution() {
};
let config_request = WasmModuleConfigRequest { descriptor };
let mut workflow_context = WorkflowContext::new(WorkflowInstanceId::new());
workflow_context.set_arc("wasm_module_config", Arc::new(config_request));
workflow_context.set_arc("app_context", app_context.clone());
let workflow_data = AnyWorkflowData::WasmRegistration(WasmRegistrationWorkflowData {
config: config_request,
wasm_bytes: None,
sha256_hash: None,
file_size_bytes: None,
module_uuid: None,
app_context: Some(app_context.clone()),
});
// Start workflow
let instance_id = engine
.start_workflow(
WorkflowId::new("wasm_module_registration"),
workflow_context,
)
.start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
.await
.expect("Failed to start workflow");
@@ -729,9 +734,9 @@ async fn test_wasm_module_execution() {
match state.status {
smg::workflow::WorkflowStatus::Completed => {
// Extract module UUID from context
if let Some(uuid_arc) = state.context.get::<Uuid>("module_uuid") {
module_uuid = Some(*uuid_arc.as_ref());
// Extract module UUID from typed workflow data
if let AnyWorkflowData::WasmRegistration(ref data) = state.context.data {
module_uuid = data.module_uuid;
}
break;
}
sgl-model-gateway/tests/workflow_test.rs
+75 -40
View File
@@ -8,9 +8,25 @@ use std::{
time::Duration,
};
use serde::{Deserialize, Serialize};
use smg::workflow::*;
use tokio::time::sleep;
/// Test workflow data type for integration tests.
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
struct TestWorkflowData {
/// Execution count for tracking step invocations
pub execution_count: u32,
/// Test key for context sharing tests
pub test_key: Option<String>,
}
impl WorkflowData for TestWorkflowData {
fn workflow_type() -> &'static str {
"test_workflow"
}
}
// Test step that counts invocations
struct CountingStep {
counter: Arc<AtomicU32>,
@@ -18,12 +34,15 @@ struct CountingStep {
}
#[async_trait::async_trait]
impl StepExecutor for CountingStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
impl StepExecutor<TestWorkflowData> for CountingStep {
async fn execute(
&self,
context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
let count = self.counter.fetch_add(1, Ordering::SeqCst) + 1;
// Store count in context
context.set("execution_count", count);
context.data.execution_count = count;
if count >= self.should_succeed_after {
Ok(StepResult::Success)
@@ -40,15 +59,18 @@ impl StepExecutor for CountingStep {
struct AlwaysSucceedStep;
#[async_trait::async_trait]
impl StepExecutor for AlwaysSucceedStep {
async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
impl StepExecutor<TestWorkflowData> for AlwaysSucceedStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
Ok(StepResult::Success)
}
}
#[tokio::test]
async fn test_simple_workflow_execution() {
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
// Subscribe to events for logging
engine
@@ -74,7 +96,7 @@ async fn test_simple_workflow_execution() {
// Start workflow
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -89,7 +111,7 @@ async fn test_simple_workflow_execution() {
#[tokio::test]
async fn test_workflow_with_retry() {
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
engine
.event_bus()
.subscribe(Arc::new(LoggingSubscriber))
@@ -119,7 +141,7 @@ async fn test_workflow_with_retry() {
// Start workflow
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -140,7 +162,7 @@ async fn test_workflow_with_retry() {
#[tokio::test]
async fn test_workflow_failure_after_max_retries() {
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
engine
.event_bus()
.subscribe(Arc::new(LoggingSubscriber))
@@ -170,7 +192,7 @@ async fn test_workflow_failure_after_max_retries() {
// Start workflow
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -191,7 +213,7 @@ async fn test_workflow_failure_after_max_retries() {
#[tokio::test]
async fn test_workflow_continue_on_failure() {
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
engine
.event_bus()
.subscribe(Arc::new(LoggingSubscriber))
@@ -227,7 +249,7 @@ async fn test_workflow_continue_on_failure() {
// Start workflow
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -249,34 +271,40 @@ async fn test_workflow_continue_on_failure() {
#[tokio::test]
async fn test_workflow_context_sharing() {
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
struct ContextWriterStep {
key: String,
value: String,
}
#[async_trait::async_trait]
impl StepExecutor for ContextWriterStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
context.set(self.key.clone(), self.value.clone());
impl StepExecutor<TestWorkflowData> for ContextWriterStep {
async fn execute(
&self,
context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
context.data.test_key = Some(self.value.clone());
Ok(StepResult::Success)
}
}
struct ContextReaderStep {
key: String,
expected_value: String,
}
#[async_trait::async_trait]
impl StepExecutor for ContextReaderStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let value: Arc<String> = context
.get(&self.key)
.ok_or_else(|| WorkflowError::ContextValueNotFound(self.key.clone()))?;
impl StepExecutor<TestWorkflowData> for ContextReaderStep {
async fn execute(
&self,
context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
let value = context
.data
.test_key
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("test_key".to_string()))?;
if *value == self.expected_value {
if value == &self.expected_value {
Ok(StepResult::Success)
} else {
Err(WorkflowError::StepFailed {
@@ -292,7 +320,6 @@ async fn test_workflow_context_sharing() {
"writer",
"Write to context",
Arc::new(ContextWriterStep {
key: "test_key".to_string(),
value: "test_value".to_string(),
}),
))
@@ -300,7 +327,6 @@ async fn test_workflow_context_sharing() {
"reader",
"Read from context",
Arc::new(ContextReaderStep {
key: "test_key".to_string(),
expected_value: "test_value".to_string(),
}),
));
@@ -309,7 +335,7 @@ async fn test_workflow_context_sharing() {
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -332,8 +358,11 @@ struct TimingStep {
}
#[async_trait::async_trait]
impl StepExecutor for TimingStep {
async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
impl StepExecutor<TestWorkflowData> for TimingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
let start = std::time::Instant::now();
self.start_times
.write()
@@ -351,7 +380,7 @@ impl StepExecutor for TimingStep {
#[tokio::test]
async fn test_parallel_execution_no_dependencies() {
// Steps without dependencies should run in parallel
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let start_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
Arc::new(parking_lot::RwLock::new(Vec::new()));
@@ -398,7 +427,7 @@ async fn test_parallel_execution_no_dependencies() {
let overall_start = std::time::Instant::now();
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -446,7 +475,7 @@ async fn test_dag_with_dependencies() {
// A ──┐
// ├──> C
// B ──┘
let engine = WorkflowEngine::new();
let engine: WorkflowEngine<TestWorkflowData> = WorkflowEngine::new();
let start_times: Arc<parking_lot::RwLock<Vec<(String, std::time::Instant)>>> =
Arc::new(parking_lot::RwLock::new(Vec::new()));
@@ -492,7 +521,7 @@ async fn test_dag_with_dependencies() {
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();
@@ -523,15 +552,18 @@ async fn test_dag_with_dependencies() {
#[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 engine: WorkflowEngine<TestWorkflowData> = 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> {
impl StepExecutor<TestWorkflowData> for FailingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
Err(WorkflowError::StepFailed {
step_id: StepId::new("failing"),
message: "Intentional failure".to_string(),
@@ -548,8 +580,11 @@ async fn test_dag_dependency_failure_blocks_dependents() {
}
#[async_trait::async_trait]
impl StepExecutor for TrackingStep {
async fn execute(&self, _context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
impl StepExecutor<TestWorkflowData> for TrackingStep {
async fn execute(
&self,
_context: &mut WorkflowContext<TestWorkflowData>,
) -> WorkflowResult<StepResult> {
self.counter.fetch_add(1, Ordering::SeqCst);
Ok(StepResult::Success)
}
@@ -575,7 +610,7 @@ async fn test_dag_dependency_failure_blocks_dependents() {
engine.register_workflow(workflow).unwrap();
let instance_id = engine
.start_workflow(workflow_id, WorkflowContext::new(WorkflowInstanceId::new()))
.start_workflow(workflow_id, TestWorkflowData::default())
.await
.unwrap();