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(),