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