[model-gateway] refactor: extract workflow engine to src/workflow module (#14996)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
+1
-1
@@ -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
|
||||
+1
-1
@@ -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
|
||||
+1
-1
@@ -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
|
||||
+1
-1
@@ -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)]
|
||||
+3
-2
@@ -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
|
||||
+1
-4
@@ -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)]
|
||||
@@ -17,3 +17,4 @@ pub mod tokenizer;
|
||||
pub mod tool_parser;
|
||||
pub mod version;
|
||||
pub mod wasm;
|
||||
pub mod workflow;
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
+2
-2
@@ -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::*;
|
||||
@@ -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));
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
_ => {
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user