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

[model-gateway] refactor workflow engine from type erasure to typed engines (#16973)

Browse Source
This commit is contained in:
Simo Lin
2026-01-12 10:47:00 -08:00
committed by GitHub
parent fa51b85466
commit ed729d22b3
36 changed files with 803 additions and 917 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/tests/common/mod.rs
+27 -54
View File
@@ -322,9 +322,9 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
// Create empty OnceLock for worker job queue, workflow engines, and mcp manager
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
let mcp_manager_lock = Arc::new(OnceLock::new());
let app_context = Arc::new(
@@ -342,7 +342,7 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
@@ -356,22 +356,13 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
// Initialize typed workflow engines
use smg::core::steps::WorkflowEngines;
let engines = WorkflowEngines::new(&config);
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
.workflow_engines
.set(engines)
.expect("WorkflowEngines should only be initialized once");
// Register external workers for OpenAI mode
if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
@@ -451,9 +442,9 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
// Create empty OnceLock for worker job queue, workflow engines, and mcp manager
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
let mcp_manager_lock = Arc::new(OnceLock::new());
// Initialize parser factories
@@ -475,7 +466,7 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
@@ -489,22 +480,13 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
// Initialize typed workflow engines
use smg::core::steps::WorkflowEngines;
let engines = WorkflowEngines::new(&config);
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
.workflow_engines
.set(engines)
.expect("WorkflowEngines should only be initialized once");
// Register external workers for OpenAI mode
if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
@@ -588,9 +570,9 @@ pub async fn create_test_context_with_mcp_config(
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
// Create empty OnceLock for worker job queue, workflow engines, and mcp manager
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
let mcp_manager_lock = Arc::new(OnceLock::new());
let app_context = Arc::new(
@@ -608,7 +590,7 @@ pub async fn create_test_context_with_mcp_config(
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
@@ -622,22 +604,13 @@ pub async fn create_test_context_with_mcp_config(
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
// Initialize typed workflow engines
use smg::core::steps::WorkflowEngines;
let engines = WorkflowEngines::new(&config);
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
.workflow_engines
.set(engines)
.expect("WorkflowEngines should only be initialized once");
// Register external workers for OpenAI mode
if let RoutingMode::OpenAI { worker_urls, .. } = &config.mode {
sgl-model-gateway/tests/common/test_app.rs
+5 -5
View File
@@ -58,9 +58,9 @@ pub fn create_test_app(
router_config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue and workflow engine
// Create empty OnceLock for worker job queue and workflow engines
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
// Create AppContext using builder pattern
let app_context = Arc::new(
@@ -78,7 +78,7 @@ pub fn create_test_app(
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.build()
.unwrap(),
);
@@ -168,7 +168,7 @@ pub async fn create_test_app_context() -> Arc<AppContext> {
// Initialize empty OnceLocks
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
// Initialize MCP manager with empty config
let mcp_manager_lock = Arc::new(OnceLock::new());
@@ -208,7 +208,7 @@ pub async fn create_test_app_context() -> Arc<AppContext> {
.conversation_item_storage(conversation_item_storage)
.load_monitor(None)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager_lock)
.build()
.unwrap(),
sgl-model-gateway/tests/routing/test_pd_routing.rs
+3 -3
View File
@@ -247,9 +247,9 @@ mod pd_routing_unit_tests {
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
// Create empty OnceLock for worker job queue, workflow engines, and mcp manager
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
let mcp_manager = Arc::new(OnceLock::new());
Arc::new(
@@ -267,7 +267,7 @@ mod pd_routing_unit_tests {
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager)
.build()
.unwrap(),
sgl-model-gateway/tests/wasm_test.rs
+27 -48
View File
@@ -18,10 +18,7 @@ use axum::{
use smg::{
app_context::AppContext,
config::RouterConfig,
core::{
steps::{create_wasm_module_registration_workflow, create_wasm_module_removal_workflow},
LoadMonitor, WorkerRegistry,
},
core::{LoadMonitor, WorkerRegistry},
data_connector::{
MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
},
@@ -71,10 +68,10 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
config.worker_startup_check_interval_secs,
)));
// Create empty OnceLock for worker job queue, workflow engine, and mcp manager
// Create empty OnceLock for worker job queue, workflow engines, and mcp manager
use std::sync::OnceLock;
let worker_job_queue = Arc::new(OnceLock::new());
let workflow_engine = Arc::new(OnceLock::new());
let workflow_engines = Arc::new(OnceLock::new());
let mcp_manager_lock = Arc::new(OnceLock::new());
let app_context = Arc::new(
@@ -92,7 +89,7 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
.conversation_item_storage(conversation_item_storage)
.load_monitor(load_monitor)
.worker_job_queue(worker_job_queue)
.workflow_engine(workflow_engine)
.workflow_engines(workflow_engines)
.mcp_manager(mcp_manager_lock)
.wasm_manager(Some(wasm_manager))
.build()
@@ -107,28 +104,13 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
.set(job_queue)
.expect("JobQueue should only be initialized once");
// Initialize WorkflowEngine and register workflows
use smg::{
core::steps::{create_local_worker_workflow, create_worker_removal_workflow},
workflow::WorkflowEngine,
};
let engine = Arc::new(WorkflowEngine::new());
engine
.register_workflow(create_local_worker_workflow(&config))
.expect("worker_registration workflow should be valid");
engine
.register_workflow(create_worker_removal_workflow())
.expect("worker_removal workflow should be valid");
engine
.register_workflow(create_wasm_module_registration_workflow())
.expect("wasm_module_registration workflow should be valid");
engine
.register_workflow(create_wasm_module_removal_workflow())
.expect("wasm_module_removal workflow should be valid");
// Initialize WorkflowEngines
use smg::core::steps::WorkflowEngines;
let engines = WorkflowEngines::new(&config);
app_context
.workflow_engine
.set(engine)
.expect("WorkflowEngine should only be initialized once");
.workflow_engines
.set(engines)
.expect("WorkflowEngines should only be initialized once");
// Initialize MCP manager with empty config
use smg::mcp::{McpConfig, McpManager};
@@ -678,17 +660,14 @@ async fn test_wasm_module_execution() {
.as_ref()
.expect("WASM manager should be initialized");
let engine = app_context
.workflow_engine
let engines = app_context
.workflow_engines
.get()
.expect("Workflow engine should be initialized");
.expect("Workflow engines should be initialized");
// Create workflow context for registration
use smg::{
core::steps::{
workflow_data::{AnyWorkflowData, WasmRegistrationWorkflowData},
WasmModuleConfigRequest,
},
core::steps::{WasmModuleConfigRequest, WasmRegistrationWorkflowData},
workflow::WorkflowId,
};
@@ -703,17 +682,18 @@ async fn test_wasm_module_execution() {
};
let config_request = WasmModuleConfigRequest { descriptor };
let workflow_data = AnyWorkflowData::WasmRegistration(WasmRegistrationWorkflowData {
let workflow_data = 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
let instance_id = engines
.wasm_registration
.start_workflow(WorkflowId::new("wasm_module_registration"), workflow_data)
.await
.expect("Failed to start workflow");
@@ -721,24 +701,25 @@ async fn test_wasm_module_execution() {
// Wait for workflow to complete
let timeout = Duration::from_secs(30);
let start = std::time::Instant::now();
let mut module_uuid: Option<Uuid> = None;
loop {
let module_uuid = loop {
if start.elapsed() > timeout {
panic!("Workflow timeout");
}
let state = engine
let state = engines
.wasm_registration
.get_status(instance_id)
.expect("Failed to get workflow status");
match state.status {
smg::workflow::WorkflowStatus::Completed => {
// Extract module UUID from typed workflow data
if let AnyWorkflowData::WasmRegistration(ref data) = state.context.data {
module_uuid = data.module_uuid;
}
break;
break state
.context
.data
.module_uuid
.expect("Module UUID should be in context");
}
smg::workflow::WorkflowStatus::Failed => {
panic!("Workflow failed: {:?}", state);
@@ -747,9 +728,7 @@ async fn test_wasm_module_execution() {
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
}
let module_uuid = module_uuid.expect("Module UUID should be in context");
};
// Verify module is registered
let module = wasm_manager