[model-gateway] refactor: extract workflow engine to src/workflow module (#14996)

This commit is contained in:
Simo Lin
2025-12-12 06:29:12 -08:00
committed by GitHub
parent c7c837cd1d
commit 10c68f6236
23 changed files with 41 additions and 48 deletions
+2 -1
View File
@@ -8,7 +8,7 @@ use tracing::info;
use crate::{ use crate::{
config::RouterConfig, config::RouterConfig,
core::{workflow::WorkflowEngine, ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry}, core::{ConnectionMode, JobQueue, LoadMonitor, WorkerRegistry},
data_connector::{ data_connector::{
create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage, create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage,
}, },
@@ -24,6 +24,7 @@ use crate::{
}, },
tool_parser::ParserFactory as ToolParserFactory, tool_parser::ParserFactory as ToolParserFactory,
wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager}, wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
workflow::WorkflowEngine,
}; };
/// Error type for AppContext builder /// Error type for AppContext builder
+2 -4
View File
@@ -16,16 +16,14 @@ use tracing::{debug, error, info, warn};
use crate::{ use crate::{
app_context::AppContext, app_context::AppContext,
config::{RouterConfig, RoutingMode}, config::{RouterConfig, RoutingMode},
core::workflow::{ core::steps::{
steps::{
McpServerConfigRequest, WasmModuleConfigRequest, WasmModuleRemovalRequest, McpServerConfigRequest, WasmModuleConfigRequest, WasmModuleRemovalRequest,
WorkerRemovalRequest, WorkerRemovalRequest,
}, },
WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus,
},
mcp::McpConfig, mcp::McpConfig,
observability::metrics::RouterMetrics, observability::metrics::RouterMetrics,
protocols::worker_spec::{JobStatus, WorkerConfigRequest}, protocols::worker_spec::{JobStatus, WorkerConfigRequest},
workflow::{WorkflowContext, WorkflowEngine, WorkflowId, WorkflowInstanceId, WorkflowStatus},
}; };
/// Job types for control plane operations /// Job types for control plane operations
+2 -2
View File
@@ -6,7 +6,7 @@
//! - Error types //! - Error types
//! - Circuit breaker for reliability //! - Circuit breaker for reliability
//! - Token buckets for rate limiting //! - Token buckets for rate limiting
//! - Workflow engine for multi-step operations //! - Workflow steps for multi-step operations
//! - Common utilities //! - Common utilities
pub mod circuit_breaker; pub mod circuit_breaker;
@@ -16,12 +16,12 @@ pub mod metrics_aggregator;
pub mod model_card; pub mod model_card;
pub mod model_type; pub mod model_type;
pub mod retry; pub mod retry;
pub mod steps;
pub mod token_bucket; pub mod token_bucket;
pub mod worker; pub mod worker;
pub mod worker_builder; pub mod worker_builder;
pub mod worker_manager; pub mod worker_manager;
pub mod worker_registry; pub mod worker_registry;
pub mod workflow;
pub use circuit_breaker::{ pub use circuit_breaker::{
CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState, CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState,
@@ -28,11 +28,11 @@ use crate::{
core::{ core::{
model_card::{ModelCard, ProviderType}, model_card::{ModelCard, ProviderType},
model_type::ModelType, model_type::ModelType,
workflow::*,
BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, HealthConfig, RuntimeType, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, HealthConfig, RuntimeType,
Worker, WorkerType, Worker, WorkerType,
}, },
protocols::worker_spec::WorkerConfigRequest, protocols::worker_spec::WorkerConfigRequest,
workflow::*,
}; };
// HTTP client for API calls // HTTP client for API calls
@@ -16,8 +16,8 @@ use tracing::{debug, error, info, warn};
use crate::{ use crate::{
app_context::AppContext, app_context::AppContext,
core::workflow::*,
mcp::{config::McpServerConfig, manager::McpManager}, mcp::{config::McpServerConfig, manager::McpManager},
workflow::*,
}; };
/// MCP server connection configuration /// MCP server connection configuration
@@ -24,8 +24,8 @@ use wasmtime::{component::Component, Config, Engine};
use crate::{ use crate::{
app_context::AppContext, app_context::AppContext,
core::workflow::*,
wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta}, wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta},
workflow::*,
}; };
/// WASM module registration request /// WASM module registration request
@@ -12,7 +12,7 @@ use async_trait::async_trait;
use tracing::{debug, info}; use tracing::{debug, info};
use uuid::Uuid; use uuid::Uuid;
use crate::{app_context::AppContext, core::workflow::*}; use crate::{app_context::AppContext, workflow::*};
/// WASM module removal request /// WASM module removal request
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -24,11 +24,12 @@ use tracing::{debug, info, warn};
use crate::{ use crate::{
app_context::AppContext, app_context::AppContext,
core::{ core::{
workflow::*, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, BasicWorkerBuilder, CircuitBreakerConfig, ConnectionMode, DPAwareWorkerBuilder,
DPAwareWorkerBuilder, HealthConfig, ModelCard, RuntimeType, Worker, WorkerType, HealthConfig, ModelCard, RuntimeType, Worker, WorkerType,
}, },
protocols::worker_spec::WorkerConfigRequest, protocols::worker_spec::WorkerConfigRequest,
routers::grpc::client::GrpcClient, routers::grpc::client::GrpcClient,
workflow::*,
}; };
// HTTP client for metadata fetching // HTTP client for metadata fetching
@@ -14,10 +14,7 @@ use std::{collections::HashSet, sync::Arc, time::Duration};
use async_trait::async_trait; use async_trait::async_trait;
use tracing::{debug, info}; use tracing::{debug, info};
use crate::{ use crate::{app_context::AppContext, core::Worker, workflow::*};
app_context::AppContext,
core::{workflow::*, Worker},
};
/// Request structure for worker removal /// Request structure for worker removal
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
+1
View File
@@ -17,3 +17,4 @@ pub mod tokenizer;
pub mod tool_parser; pub mod tool_parser;
pub mod version; pub mod version;
pub mod wasm; pub mod wasm;
pub mod workflow;
+4 -5
View File
@@ -23,14 +23,12 @@ use crate::{
app_context::AppContext, app_context::AppContext,
config::{RouterConfig, RoutingMode}, config::{RouterConfig, RoutingMode},
core::{ core::{
worker_to_info, steps::{
workflow::{
create_external_worker_registration_workflow, create_mcp_registration_workflow, create_external_worker_registration_workflow, create_mcp_registration_workflow,
create_wasm_module_registration_workflow, create_wasm_module_removal_workflow, create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
create_worker_registration_workflow, create_worker_removal_workflow, LoggingSubscriber, create_worker_registration_workflow, create_worker_removal_workflow,
WorkflowEngine,
}, },
Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType, worker_to_info, Job, JobQueue, JobQueueConfig, WorkerManager, WorkerType,
}, },
middleware::{self, AuthConfig, QueuedRequest}, middleware::{self, AuthConfig, QueuedRequest},
observability::{ observability::{
@@ -52,6 +50,7 @@ use crate::{
routers::{conversations, router_manager::RouterManager, RouterTrait}, routers::{conversations, router_manager::RouterManager, RouterTrait},
service_discovery::{start_service_discovery, ServiceDiscoveryConfig}, service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module}, wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},
workflow::{LoggingSubscriber, WorkflowEngine},
}; };
#[derive(Clone)] #[derive(Clone)]
+2 -2
View File
@@ -15,7 +15,7 @@ use axum::{
use uuid::Uuid; use uuid::Uuid;
use crate::{ use crate::{
core::{job_queue::Job, workflow::steps::WasmModuleConfigRequest}, core::{job_queue::Job, steps::WasmModuleConfigRequest},
server::AppState, server::AppState,
wasm::module::{ wasm::module::{
WasmMetrics, WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAddResult, WasmMetrics, WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAddResult,
@@ -172,7 +172,7 @@ pub async fn remove_wasm_module(
return StatusCode::INTERNAL_SERVER_ERROR.into_response(); return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}; };
use crate::core::workflow::steps::WasmModuleRemovalRequest; use crate::core::steps::WasmModuleRemovalRequest;
let removal_request = WasmModuleRemovalRequest::new(module_uuid); let removal_request = WasmModuleRemovalRequest::new(module_uuid);
@@ -84,7 +84,7 @@ where
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::core::workflow::types::WorkflowInstanceId; use crate::workflow::types::WorkflowInstanceId;
struct TestStep { struct TestStep {
should_succeed: bool, should_succeed: bool,
@@ -97,7 +97,7 @@ mod tests {
Ok(StepResult::Success) Ok(StepResult::Success)
} else { } else {
Err(WorkflowError::StepFailed { Err(WorkflowError::StepFailed {
step_id: crate::core::workflow::types::StepId::new("test"), step_id: crate::workflow::types::StepId::new("test"),
message: "test error".to_string(), message: "test error".to_string(),
}) })
} }
@@ -5,7 +5,6 @@ mod engine;
mod event; mod event;
mod executor; mod executor;
mod state; mod state;
pub mod steps;
pub mod types; pub mod types;
// Re-export main types // Re-export main types
@@ -14,9 +13,4 @@ pub use engine::WorkflowEngine;
pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent}; pub use event::{EventBus, EventSubscriber, LoggingSubscriber, WorkflowEvent};
pub use executor::{FunctionStep, StepExecutor}; pub use executor::{FunctionStep, StepExecutor};
pub use state::WorkflowStateStore; pub use state::WorkflowStateStore;
pub use steps::{
create_external_worker_registration_workflow, create_mcp_registration_workflow,
create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
create_worker_registration_workflow, create_worker_removal_workflow,
};
pub use types::*; pub use types::*;
+6 -4
View File
@@ -102,8 +102,9 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
.expect("JobQueue should only be initialized once"); .expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows // Initialize WorkflowEngine and register workflows
use sgl_model_gateway::core::workflow::{ use sgl_model_gateway::{
create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine, core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
}; };
let engine = Arc::new(WorkflowEngine::new()); let engine = Arc::new(WorkflowEngine::new());
engine.register_workflow(create_worker_registration_workflow(&config)); engine.register_workflow(create_worker_registration_workflow(&config));
@@ -233,8 +234,9 @@ pub async fn create_test_context_with_mcp_config(
.expect("JobQueue should only be initialized once"); .expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows // Initialize WorkflowEngine and register workflows
use sgl_model_gateway::core::workflow::{ use sgl_model_gateway::{
create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine, core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
}; };
let engine = Arc::new(WorkflowEngine::new()); let engine = Arc::new(WorkflowEngine::new());
engine.register_workflow(create_worker_registration_workflow(&config)); engine.register_workflow(create_worker_registration_workflow(&config));
+9 -9
View File
@@ -18,9 +18,7 @@ use axum::{
use sgl_model_gateway::{ use sgl_model_gateway::{
app_context::AppContext, app_context::AppContext,
config::RouterConfig, config::RouterConfig,
core::workflow::{ core::steps::{create_wasm_module_registration_workflow, create_wasm_module_removal_workflow},
create_wasm_module_registration_workflow, create_wasm_module_removal_workflow,
},
routers::RouterFactory, routers::RouterFactory,
server::{build_app, AppState}, server::{build_app, AppState},
wasm::{ wasm::{
@@ -113,8 +111,9 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
.expect("JobQueue should only be initialized once"); .expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows // Initialize WorkflowEngine and register workflows
use sgl_model_gateway::core::workflow::{ use sgl_model_gateway::{
create_worker_registration_workflow, create_worker_removal_workflow, WorkflowEngine, core::steps::{create_worker_registration_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
}; };
let engine = Arc::new(WorkflowEngine::new()); let engine = Arc::new(WorkflowEngine::new());
engine.register_workflow(create_worker_registration_workflow(&config)); engine.register_workflow(create_worker_registration_workflow(&config));
@@ -684,8 +683,9 @@ async fn test_wasm_module_execution() {
.expect("Workflow engine should be initialized"); .expect("Workflow engine should be initialized");
// Create workflow context for registration // Create workflow context for registration
use sgl_model_gateway::core::workflow::{ use sgl_model_gateway::{
steps::WasmModuleConfigRequest, WorkflowContext, WorkflowId, WorkflowInstanceId, core::steps::WasmModuleConfigRequest,
workflow::{WorkflowContext, WorkflowId, WorkflowInstanceId},
}; };
let descriptor = WasmModuleDescriptor { let descriptor = WasmModuleDescriptor {
@@ -727,14 +727,14 @@ async fn test_wasm_module_execution() {
.expect("Failed to get workflow status"); .expect("Failed to get workflow status");
match state.status { match state.status {
sgl_model_gateway::core::workflow::WorkflowStatus::Completed => { sgl_model_gateway::workflow::WorkflowStatus::Completed => {
// Extract module UUID from context // Extract module UUID from context
if let Some(uuid_arc) = state.context.get::<Uuid>("module_uuid") { if let Some(uuid_arc) = state.context.get::<Uuid>("module_uuid") {
module_uuid = Some(*uuid_arc.as_ref()); module_uuid = Some(*uuid_arc.as_ref());
} }
break; break;
} }
sgl_model_gateway::core::workflow::WorkflowStatus::Failed => { sgl_model_gateway::workflow::WorkflowStatus::Failed => {
panic!("Workflow failed: {:?}", state); panic!("Workflow failed: {:?}", state);
} }
_ => { _ => {
+1 -1
View File
@@ -8,7 +8,7 @@ use std::{
time::Duration, time::Duration,
}; };
use sgl_model_gateway::core::workflow::*; use sgl_model_gateway::workflow::*;
use tokio::time::sleep; use tokio::time::sleep;
// Test step that counts invocations // Test step that counts invocations