diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index 364a2b0c8..130b61ecc 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -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>; + /// Error type for AppContext builder #[derive(Debug)] pub struct AppContextBuildError(&'static str); @@ -59,13 +65,21 @@ pub struct AppContext { pub configured_reasoning_parser: Option, pub configured_tool_parser: Option, pub worker_job_queue: Arc>>, - pub workflow_engine: Arc>>, + pub workflow_engine: Arc>>, pub mcp_manager: Arc>>, pub wasm_manager: Option>, pub worker_service: Arc, pub inflight_tracker: Arc, } +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, router_config: Option, @@ -81,7 +95,7 @@ pub struct AppContextBuilder { conversation_item_storage: Option>, load_monitor: Option>, worker_job_queue: Option>>>, - workflow_engine: Option>>>, + workflow_engine: Option>>>, mcp_manager: Option>>>, wasm_manager: Option>, } @@ -206,7 +220,10 @@ impl AppContextBuilder { self } - pub fn workflow_engine(mut self, workflow_engine: Arc>>) -> Self { + pub fn workflow_engine( + mut self, + workflow_engine: Arc>>, + ) -> Self { self.workflow_engine = Some(workflow_engine); self } diff --git a/sgl-model-gateway/src/core/job_queue.rs b/sgl-model-gateway/src/core/job_queue.rs index 2a29f3b9f..a657b3723 100644 --- a/sgl-model-gateway/src/core/job_queue.rs +++ b/sgl-model-gateway/src/core/job_queue.rs @@ -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 = 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 = 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 = 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, + engine: &Arc, config: &WorkerConfigRequest, context: &Arc, ) -> Result { - 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, + engine: &Arc, url: &str, context: &Arc, ) -> Result { - 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, + engine: &Arc, url: &str, update: &WorkerUpdateRequest, context: &Arc, ) -> Result { - 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, + engine: &Arc, config: &McpServerConfigRequest, context: &Arc, ) -> Result { - 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, + engine: &Arc, instance_id: WorkflowInstanceId, worker_url: &str, timeout_duration: Duration, diff --git a/sgl-model-gateway/src/core/steps/mcp_registration.rs b/sgl-model-gateway/src/core/steps/mcp_registration.rs index dd49cd6e5..a55398b1d 100644 --- a/sgl-model-gateway/src/core/steps/mcp_registration.rs +++ b/sgl-model-gateway/src/core/steps/mcp_registration.rs @@ -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 { - let config_request: Arc = - context.get_or_err("mcp_server_config")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for ConnectMcpServerStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("mcp_server_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let mcp_client: Arc> = context.get_or_err("mcp_client")?; +impl StepExecutor for DiscoverMcpInventoryStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("mcp_server_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let mcp_client: Arc> = context.get_or_err("mcp_client")?; +impl StepExecutor for RegisterMcpServerStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("mcp_server_config")?; +impl StepExecutor for ValidateRegistrationStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + let data = context.data.as_mcp()?; + let config_request = &data.config; - let client_registered = context - .get::>("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 { 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, +) -> AnyWorkflowData { + AnyWorkflowData::Mcp(McpWorkflowData { + config, + validated: false, + app_context: Some(app_context), + mcp_client: None, + }) +} diff --git a/sgl-model-gateway/src/core/steps/mod.rs b/sgl-model-gateway/src/core/steps/mod.rs index 112ab924b..da1b6e6fa 100644 --- a/sgl-model-gateway/src/core/steps/mod.rs +++ b/sgl-model-gateway/src/core/steps/mod.rs @@ -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, }; diff --git a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs index bbd0218ce..7050fb1af 100644 --- a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs +++ b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs @@ -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 { - let config: Arc = context.get_or_err("tokenizer_config")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for ValidateTokenizerConfigStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config: Arc = context.get_or_err("tokenizer_config")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for LoadTokenizerStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { 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, +) -> 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"); } } diff --git a/sgl-model-gateway/src/core/steps/wasm_module_registration.rs b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs index 5358be6de..5b6b9a9c0 100644 --- a/sgl-model-gateway/src/core/steps/wasm_module_registration.rs +++ b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs @@ -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 { - let config_request: Arc = - context.get_or_err("wasm_module_config")?; - - let descriptor = &config_request.descriptor; +impl StepExecutor for ValidateDescriptorStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("wasm_module_config")?; - - let file_path = &config_request.descriptor.file_path; +impl StepExecutor for CalculateHashStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("wasm_module_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; +impl StepExecutor for CheckDuplicateStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("wasm_module_config")?; - - let file_path = &config_request.descriptor.file_path; +impl StepExecutor for LoadWasmBytesStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("wasm_module_config")?; - let wasm_bytes: Arc> = context.get_or_err("wasm_bytes")?; +impl StepExecutor for ValidateWasmComponentStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let config_request: Arc = - context.get_or_err("wasm_module_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let sha256_hash: Arc<[u8; 32]> = context.get_or_err("sha256_hash")?; - let file_size_bytes: Arc = context.get_or_err("file_size_bytes")?; - let wasm_bytes: Arc> = context.get_or_err("wasm_bytes")?; +impl StepExecutor for RegisterModuleStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { 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, +) -> AnyWorkflowData { + AnyWorkflowData::WasmRegistration(WasmRegistrationWorkflowData { + config, + wasm_bytes: None, + sha256_hash: None, + file_size_bytes: None, + module_uuid: None, + app_context: Some(app_context), + }) +} diff --git a/sgl-model-gateway/src/core/steps/wasm_module_removal.rs b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs index 8405344fb..83692fc5a 100644 --- a/sgl-model-gateway/src/core/steps/wasm_module_removal.rs +++ b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs @@ -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 { - let removal_request: Arc = - context.get_or_err("wasm_module_removal_request")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for FindModuleToRemoveStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let removal_request: Arc = - context.get_or_err("wasm_module_removal_request")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for RemoveModuleStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { 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, +) -> AnyWorkflowData { + AnyWorkflowData::WasmRemoval(WasmRemovalWorkflowData { + config, + module_id: None, + app_context: Some(app_context), + }) +} diff --git a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs index 015784198..ea32494a7 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs @@ -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 { - let config: Arc = context.get_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let model_cards: Arc> = context.get_or_err("model_cards")?; +impl StepExecutor for CreateExternalWorkersStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) } diff --git a/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs b/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs index 066c01f49..eb9820c49 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs @@ -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 pub struct DiscoverModelsStep; #[async_trait] -impl StepExecutor for DiscoverModelsStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; +impl StepExecutor for DiscoverModelsStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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::>("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::>() ); - context.set("model_cards", model_cards); + context.data.as_external_worker_mut()?.model_cards = model_cards; Ok(StepResult::Success) } diff --git a/sgl-model-gateway/src/core/steps/worker/external/mod.rs b/sgl-model-gateway/src/core/steps/worker/external/mod.rs index 99a9787f3..746efbd45 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/mod.rs @@ -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 { 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, +) -> AnyWorkflowData { + AnyWorkflowData::ExternalWorker(ExternalWorkerWorkflowData { + config, + model_cards: Vec::new(), + workers: None, + labels: std::collections::HashMap::new(), + app_context: Some(app_context), + actual_workers: None, + }) +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs index ee1dcbc2c..d7f53d445 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs @@ -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 { - let config: Arc = context.get_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; - let connection_mode: Arc = context.get_or_err("connection_mode")?; - let discovered_labels: Arc> = - context.get_or_err("discovered_labels")?; +impl StepExecutor for CreateLocalWorkerStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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::("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, ) -> Result>, WorkflowError> { - let dp_info: Arc = 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: {})", diff --git a/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs b/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs index bdefd759a..58d31fdd6 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs @@ -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 { - let config: Arc = context.get_or_err("worker_config")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for DetectConnectionModeStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) } diff --git a/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs b/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs index a84f39453..e5c3b7f9e 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs @@ -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 WorkflowResult { - let config: Arc = context.get_or_err("worker_config")?; +impl StepExecutor for DiscoverDPInfoStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) } diff --git a/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs index 635cd309d..501b28b11 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs @@ -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 { - let config: Arc = context.get_or_err("worker_config")?; - let connection_mode: Arc = context.get_or_err("connection_mode")?; +impl StepExecutor for DiscoverMetadataStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) diff --git a/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs b/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs index 2bf9d154d..2a811a944 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs @@ -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 -/// -/// Sets the following context values: -/// - "workers_to_update": Vec> pub struct FindWorkerToUpdateStep; #[async_trait] -impl StepExecutor for FindWorkerToUpdateStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let worker_url: Arc = context.get_or_err("worker_url")?; - let dp_aware: Arc = context.get_or_err("dp_aware")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for FindWorkerToUpdateStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) } diff --git a/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs b/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs index 72441b60a..d627539d3 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs @@ -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 { - let request: Arc = context.get_or_err("removal_request")?; - let app_context: Arc = context.get_or_err("app_context")?; +impl StepExecutor for FindWorkersToRemoveStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) } diff --git a/sgl-model-gateway/src/core/steps/worker/local/mod.rs b/sgl-model-gateway/src/core/steps/worker/local/mod.rs index 7ecdefdd4..411252024 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -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 { 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 { 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 { 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, +) -> 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, +) -> 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, +) -> 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, + }) +} diff --git a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs index b3bfe4318..6b630cc64 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs @@ -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 { - let labels: Arc> = context.get_or_err("labels")?; - let app_context: Arc = context.get_or_err("app_context")?; - let workers: Arc>> = context.get_or_err("workers")?; +impl StepExecutor for RegisterTokenizerStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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()) }) diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs index 2755abefd..216cd5bdd 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs @@ -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 { - let app_context: Arc = context.get_or_err("app_context")?; - let workers_to_remove: Arc>> = - context.get_or_err("workers_to_remove")?; +impl StepExecutor for RemoveFromPolicyRegistryStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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", diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs index f507e5542..57f04953f 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs @@ -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 { - let app_context: Arc = context.get_or_err("app_context")?; - let worker_urls: Arc> = context.get_or_err("worker_urls")?; +impl StepExecutor for RemoveFromWorkerRegistryStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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", diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs index 7dced12a5..511933d22 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs @@ -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 { - let app_context: Arc = context.get_or_err("app_context")?; - let updated_workers: Arc>> = context.get_or_err("updated_workers")?; +impl StepExecutor for UpdatePoliciesForWorkerStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 = updated_workers diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs b/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs index 147529350..3d95c7b21 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs @@ -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 { - let app_context: Arc = context.get_or_err("app_context")?; - let affected_models: Arc> = context.get_or_err("affected_models")?; - let worker_urls: Arc> = context.get_or_err("worker_urls")?; +impl StepExecutor for UpdateRemainingPoliciesStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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)", diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs b/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs index 8bd290d22..b9ea8c123 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs @@ -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 -/// - "workers_to_update": Vec> -/// -/// Sets the following context values: -/// - "updated_workers": Vec> pub struct UpdateWorkerPropertiesStep; #[async_trait] -impl StepExecutor for UpdateWorkerPropertiesStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - let request: Arc = context.get_or_err("update_request")?; - let app_context: Arc = context.get_or_err("app_context")?; - let workers_to_update: Arc>> = - context.get_or_err("workers_to_update")?; +impl StepExecutor for UpdateWorkerPropertiesStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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) } diff --git a/sgl-model-gateway/src/core/steps/worker/mod.rs b/sgl-model-gateway/src/core/steps/worker/mod.rs index 7e3d1a669..fa4baa21c 100644 --- a/sgl-model-gateway/src/core/steps/worker/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/mod.rs @@ -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, }; diff --git a/sgl-model-gateway/src/core/steps/worker/shared/activate.rs b/sgl-model-gateway/src/core/steps/worker/shared/activate.rs index eaee82962..a89c1287b 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/activate.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/activate.rs @@ -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 { - let workers: Arc>> = context.get_or_err("workers")?; +impl StepExecutor for ActivateWorkersStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 } } diff --git a/sgl-model-gateway/src/core/steps/worker/shared/register.rs b/sgl-model-gateway/src/core/steps/worker/shared/register.rs index e0584e927..7fa67ec6b 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/register.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/register.rs @@ -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 { - let app_context: Arc = context.get_or_err("app_context")?; - let workers: Arc>> = context.get_or_err("workers")?; +impl StepExecutor for RegisterWorkersStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 } } diff --git a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs index 62ae9cbd5..4722d0866 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs @@ -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 { - let app_context: Arc = context.get_or_err("app_context")?; - let workers: Arc>> = context.get_or_err("workers")?; - let labels: Arc> = context.get_or_err("labels")?; +impl StepExecutor for UpdatePoliciesStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 } } diff --git a/sgl-model-gateway/src/core/steps/workflow_data.rs b/sgl-model-gateway/src/core/steps/workflow_data.rs new file mode 100644 index 000000000..9874176a0 --- /dev/null +++ b/sgl-model-gateway/src/core/steps/workflow_data.rs @@ -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, so we store URLs) + pub worker_urls: Vec, +} + +impl WorkerList { + pub fn new() -> Self { + Self { + worker_urls: Vec::new(), + } + } + + pub fn from_workers(workers: &[Arc]) -> 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, + /// Application context (transient, must be re-initialized after deserialization) + #[serde(skip, default)] + pub app_context: Option>, +} + +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, + pub discovered_labels: HashMap, + pub dp_info: Option, + pub workers: Option, + pub final_labels: HashMap, + /// Detected runtime type (for gRPC workers) + pub detected_runtime_type: Option, + /// Application context (transient, must be re-initialized after deserialization) + #[serde(skip, default)] + pub app_context: Option>, + /// Actual worker objects (transient, not serialized) + #[serde(skip, default)] + pub actual_workers: Option>>, +} + +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, + pub workers: Option, + /// Labels for policies (derived from config) + pub labels: HashMap, + /// Application context (transient, must be re-initialized after deserialization) + #[serde(skip, default)] + pub app_context: Option>, + /// Actual worker objects (transient, not serialized) + #[serde(skip, default)] + pub actual_workers: Option>>, +} + +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, + /// URLs of workers being removed + pub worker_urls: Vec, + /// Model IDs affected by the removal + pub affected_models: std::collections::HashSet, + /// Application context (transient, must be re-initialized after deserialization) + #[serde(skip, default)] + pub app_context: Option>, + /// Actual worker objects to remove (transient, not serialized) + #[serde(skip, default)] + pub actual_workers_to_remove: Option>>, +} + +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>, + /// Workers to update (transient, not serialized) + #[serde(skip, default)] + pub workers_to_update: Option>>, + /// Updated worker objects (transient, not serialized) + #[serde(skip, default)] + pub updated_workers: Option>>, +} + +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>, + /// Connected MCP client (transient, not serialized) + #[serde(skip, default)] + pub mcp_client: Option>>, +} + +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>, + /// SHA256 hash of the module file (32 bytes) + pub sha256_hash: Option<[u8; 32]>, + /// File size in bytes + pub file_size_bytes: Option, + /// UUID assigned to the registered module + pub module_uuid: Option, + /// Application context (transient, must be re-initialized after deserialization) + #[serde(skip, default)] + pub app_context: Option>, +} + +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, + /// Application context (transient, must be re-initialized after deserialization) + #[serde(skip, default)] + pub app_context: Option>, +} + +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` 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> { + 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>> { + 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>, + ) -> 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> { + 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(), + } + } +} diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index a5491628a..5b855996e 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -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 Result<(), Box { pub id: StepId, pub name: String, - pub executor: Arc, + pub executor: Arc>, pub retry_policy: Option, pub timeout: Option, pub on_failure: FailureAction, pub depends_on: Vec, } -impl StepDefinition { +impl StepDefinition { pub fn new( id: impl Into, name: impl Into, - executor: Arc, + executor: Arc>, ) -> Self { Self { id: StepId::new(id.into()), @@ -64,10 +64,10 @@ impl StepDefinition { } /// Complete workflow definition -pub struct WorkflowDefinition { +pub struct WorkflowDefinition { pub id: WorkflowId, pub name: String, - pub steps: Vec, + pub steps: Vec>, 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, } -impl WorkflowDefinition { +impl WorkflowDefinition { pub fn new(id: impl Into, name: impl Into) -> 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) -> 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) -> &'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) -> 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> = 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>, visited: &mut HashSet<&'a StepId>, rec_stack: &mut HashSet<&'a StepId>, ) -> bool { diff --git a/sgl-model-gateway/src/workflow/engine.rs b/sgl-model-gateway/src/workflow/engine.rs index e3f255dac..61f0b392a 100644 --- a/sgl-model-gateway/src/workflow/engine.rs +++ b/sgl-model-gateway/src/workflow/engine.rs @@ -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`) +/// /// # 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>>>, - state_store: WorkflowStateStore, +pub struct WorkflowEngine = InMemoryStore> { + definitions: Arc>>>>, + state_store: S, event_bus: Arc, /// Shutdown signal sender - when true, engine is shutting down shutdown_tx: Arc>, @@ -125,18 +131,27 @@ pub struct WorkflowEngine { shutdown_rx: watch::Receiver, /// Count of active workflow executions active_workflows: Arc, + _phantom: PhantomData, } -impl WorkflowEngine { +impl WorkflowEngine> { pub fn new() -> Self { + Self::with_store(InMemoryStore::new()) + } +} + +impl + 'static> WorkflowEngine { + /// 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) -> 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 { // 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, + definition: Arc>, ) -> 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, + definition: &WorkflowDefinition, ) -> WorkflowResult { 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 { + pub fn get_status(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { 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 + 'static> { + engine: &'a WorkflowEngine, committed: bool, } -impl<'a> StartGuard<'a> { - fn new(engine: &'a WorkflowEngine) -> Self { +impl<'a, D: WorkflowData, S: StateStore + 'static> StartGuard<'a, D, S> { + fn new(engine: &'a WorkflowEngine) -> Self { engine.active_workflows.fetch_add(1, Ordering::AcqRel); Self { engine, @@ -813,7 +828,7 @@ impl<'a> StartGuard<'a> { } } -impl Drop for StartGuard<'_> { +impl + '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` rather +/// than cloning. +impl + 'static> Clone for WorkflowEngine { fn clone(&self) -> Self { self.clone_for_execution() } } -impl Default for WorkflowEngine { +impl Default for WorkflowEngine> { fn default() -> Self { Self::new() } } -impl std::fmt::Debug for WorkflowEngine { +impl + 'static> std::fmt::Debug for WorkflowEngine { 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() } } diff --git a/sgl-model-gateway/src/workflow/executor.rs b/sgl-model-gateway/src/workflow/executor.rs index 0c9aa9a6a..9dc50468e 100644 --- a/sgl-model-gateway/src/workflow/executor.rs +++ b/sgl-model-gateway/src/workflow/executor.rs @@ -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: Send + Sync { /// Execute the step with the given context - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult; + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult; /// 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) -> 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, _error: &WorkflowError, ) -> WorkflowResult<()> { Ok(()) @@ -40,59 +40,82 @@ pub trait StepExecutor: Send + Sync { } /// Simple function-based step executor -pub struct FunctionStep +pub struct FunctionStep where + D: WorkflowData, F: Fn( - &mut WorkflowContext, + &mut WorkflowContext, ) -> std::pin::Pin< Box> + Send + '_>, > + Send + Sync, { func: F, + _phantom: std::marker::PhantomData, } -impl FunctionStep +impl FunctionStep where + D: WorkflowData, F: Fn( - &mut WorkflowContext, + &mut WorkflowContext, ) -> std::pin::Pin< Box> + Send + '_>, > + Send + Sync, { pub fn new(func: F) -> Self { - Self { func } + Self { + func, + _phantom: std::marker::PhantomData, + } } } #[async_trait] -impl StepExecutor for FunctionStep +impl StepExecutor for FunctionStep where + D: WorkflowData, F: Fn( - &mut WorkflowContext, + &mut WorkflowContext, ) -> std::pin::Pin< Box> + Send + '_>, > + Send + Sync, { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { + async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { (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 { + impl StepExecutor for TestStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { 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()); diff --git a/sgl-model-gateway/src/workflow/mod.rs b/sgl-model-gateway/src/workflow/mod.rs index aa8e625db..291e0a574 100644 --- a/sgl-model-gateway/src/workflow/mod.rs +++ b/sgl-model-gateway/src/workflow/mod.rs @@ -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::*; diff --git a/sgl-model-gateway/src/workflow/state.rs b/sgl-model-gateway/src/workflow/state.rs index 4755e0850..de548bb1e 100644 --- a/sgl-model-gateway/src/workflow/state.rs +++ b/sgl-model-gateway/src/workflow/state.rs @@ -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: Send + Sync + Clone { + /// Save workflow state + fn save(&self, state: WorkflowState) -> WorkflowResult<()>; + + /// Load workflow state by instance ID + fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult>; + + /// Update workflow state using a closure + fn update(&self, instance_id: WorkflowInstanceId, f: F) -> WorkflowResult<()> + where + F: FnOnce(&mut WorkflowState); + + /// Delete workflow state + fn delete(&self, instance_id: WorkflowInstanceId) -> WorkflowResult<()>; + + /// List all active workflows (Running or Pending) + fn list_active(&self) -> WorkflowResult>>; + + /// List all workflows + fn list_all(&self) -> WorkflowResult>>; + + /// Check if workflow is cancelled without loading full state + fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult; + + /// 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>; +} + /// In-memory state storage for workflow instances #[derive(Clone)] -pub struct WorkflowStateStore { - states: Arc>>, +pub struct InMemoryStore { + states: Arc>>>, + _phantom: PhantomData, } -impl WorkflowStateStore { +impl InMemoryStore { 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 { - 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> { - 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> { - 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(&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 { - self.states - .read() - .get(&instance_id) - .map(|s| s.context.clone()) - .ok_or(WorkflowError::NotFound(instance_id)) - } - - /// Check if workflow is cancelled without loading full state - pub fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult { - self.states - .read() - .get(&instance_id) - .map(|s| s.status == WorkflowStatus::Cancelled) - .ok_or(WorkflowError::NotFound(instance_id)) - } - /// Get count of workflows by status pub fn count_by_status(&self, status: WorkflowStatus) -> usize { self.states @@ -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 Default for InMemoryStore { + fn default() -> Self { + Self::new() + } +} + +impl StateStore for InMemoryStore { + fn save(&self, state: WorkflowState) -> WorkflowResult<()> { + self.states.write().insert(state.instance_id, state); + Ok(()) + } + + fn load(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { + self.states + .read() + .get(&instance_id) + .cloned() + .ok_or(WorkflowError::NotFound(instance_id)) + } + + fn list_active(&self) -> WorkflowResult>> { + let states = self.states.read(); + Ok(states + .values() + .filter(|s| matches!(s.status, WorkflowStatus::Running | WorkflowStatus::Pending)) + .cloned() + .collect()) + } + + fn list_all(&self) -> WorkflowResult>> { + 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(&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(()) + } + + fn get_context(&self, instance_id: WorkflowInstanceId) -> WorkflowResult> { + self.states + .read() + .get(&instance_id) + .map(|s| s.context.clone()) + .ok_or(WorkflowError::NotFound(instance_id)) + } + + fn is_cancelled(&self, instance_id: WorkflowInstanceId) -> WorkflowResult { + self.states + .read() + .get(&instance_id) + .map(|s| s.status == WorkflowStatus::Cancelled) + .ok_or(WorkflowError::NotFound(instance_id)) + } + + 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() - } } diff --git a/sgl-model-gateway/src/workflow/types.rs b/sgl-model-gateway/src/workflow/types.rs index a3916b56e..0ba58a8bb 100644 --- a/sgl-model-gateway/src/workflow/types.rs +++ b/sgl-model-gateway/src/workflow/types.rs @@ -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, +/// #[serde(skip, default)] +/// pub app_context: Option>, +/// } +/// +/// 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 { pub instance_id: WorkflowInstanceId, pub definition_id: WorkflowId, pub status: WorkflowStatus, pub current_step: Option, pub step_states: HashMap, - pub context: WorkflowContext, + pub context: WorkflowContext, pub created_at: DateTime, pub updated_at: DateTime, } -impl WorkflowState { - pub fn new(instance_id: WorkflowInstanceId, definition_id: WorkflowId) -> Self { +impl WorkflowState { + 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 { pub instance_id: WorkflowInstanceId, - #[serde(skip)] - data: HashMap>, + 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(&mut self, key: impl Into, value: T) { - self.data.insert(key.into(), Arc::new(value)); - } - - /// Store an Arc directly without double-wrapping - pub fn set_arc(&mut self, key: impl Into, value: Arc) { - self.data.insert(key.into(), value); - } - - /// Retrieve a value from the context - pub fn get(&self, key: &str) -> Option> { - self.data - .get(key) - .and_then(|v| v.clone().downcast::().ok()) - } - - /// Retrieve a value from the context, returning an error if not found - pub fn get_or_err(&self, key: &str) -> Result, WorkflowError> { - self.get::(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 WorkflowContext { + 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, } diff --git a/sgl-model-gateway/tests/common/mod.rs b/sgl-model-gateway/tests/common/mod.rs index af07ffd6b..09f7df33a 100644 --- a/sgl-model-gateway/tests/common/mod.rs +++ b/sgl-model-gateway/tests/common/mod.rs @@ -358,12 +358,12 @@ pub async fn create_test_context(config: RouterConfig) -> Arc { // 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 Arc { // 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::("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; } diff --git a/sgl-model-gateway/tests/workflow_test.rs b/sgl-model-gateway/tests/workflow_test.rs index 6830fbded..0c63251d8 100644 --- a/sgl-model-gateway/tests/workflow_test.rs +++ b/sgl-model-gateway/tests/workflow_test.rs @@ -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, +} + +impl WorkflowData for TestWorkflowData { + fn workflow_type() -> &'static str { + "test_workflow" + } +} + // Test step that counts invocations struct CountingStep { counter: Arc, @@ -18,12 +34,15 @@ struct CountingStep { } #[async_trait::async_trait] -impl StepExecutor for CountingStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { +impl StepExecutor for CountingStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { 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 { +impl StepExecutor for AlwaysSucceedStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { Ok(StepResult::Success) } } #[tokio::test] async fn test_simple_workflow_execution() { - let engine = WorkflowEngine::new(); + let engine: WorkflowEngine = 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 = 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 = 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 = 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 = WorkflowEngine::new(); struct ContextWriterStep { - key: String, value: String, } #[async_trait::async_trait] - impl StepExecutor for ContextWriterStep { - async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult { - context.set(self.key.clone(), self.value.clone()); + impl StepExecutor for ContextWriterStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { - let value: Arc = context - .get(&self.key) - .ok_or_else(|| WorkflowError::ContextValueNotFound(self.key.clone()))?; + impl StepExecutor for ContextReaderStep { + async fn execute( + &self, + context: &mut WorkflowContext, + ) -> WorkflowResult { + 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 { +impl StepExecutor for TimingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { 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 = WorkflowEngine::new(); let start_times: Arc>> = 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 = WorkflowEngine::new(); let start_times: Arc>> = 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 = 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 { + impl StepExecutor for FailingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { Err(WorkflowError::StepFailed { step_id: StepId::new("failing"), message: "Intentional failure".to_string(), @@ -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 { + impl StepExecutor for TrackingStep { + async fn execute( + &self, + _context: &mut WorkflowContext, + ) -> WorkflowResult { 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();