diff --git a/sgl-model-gateway/benches/consistent_hash_bench.rs b/sgl-model-gateway/benches/consistent_hash_bench.rs index 0806bcccd..230e2bf8e 100644 --- a/sgl-model-gateway/benches/consistent_hash_bench.rs +++ b/sgl-model-gateway/benches/consistent_hash_bench.rs @@ -1,5 +1,5 @@ use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion}; -use smg::mesh::consistent_hash::ConsistentHashRing; +use smg_mesh::consistent_hash::ConsistentHashRing; fn setup_ring(node_count: usize) -> ConsistentHashRing { let mut ring = ConsistentHashRing::new(); diff --git a/sgl-model-gateway/bindings/golang/src/client.rs b/sgl-model-gateway/bindings/golang/src/client.rs index 2ac444d20..8be8c23db 100644 --- a/sgl-model-gateway/bindings/golang/src/client.rs +++ b/sgl-model-gateway/bindings/golang/src/client.rs @@ -10,7 +10,7 @@ use uuid::Uuid; use smg::tokenizer::create_tokenizer_from_file; use smg::tokenizer::traits::Tokenizer; -use smg::grpc_client::sglang_scheduler::SglangSchedulerClient; +use smg_grpc_client::sglang_scheduler::SglangSchedulerClient; use smg::protocols::chat::ChatCompletionRequest; use smg::routers::grpc::utils::{process_chat_messages, generate_tool_constraints}; diff --git a/sgl-model-gateway/bindings/golang/src/grpc_converter.rs b/sgl-model-gateway/bindings/golang/src/grpc_converter.rs index 0262fa1b6..8c0195af1 100644 --- a/sgl-model-gateway/bindings/golang/src/grpc_converter.rs +++ b/sgl-model-gateway/bindings/golang/src/grpc_converter.rs @@ -14,7 +14,7 @@ use smg::tokenizer::stream::DecodeStream; use smg::tool_parser::ToolParser; use smg::protocols::common::{Tool, ToolChoice, ToolChoiceValue, ToolCallDelta, FunctionCallDelta, Usage, StringOrArray}; use smg::tokenizer::stop::StopSequenceDecoder; -use smg::grpc_client::sglang_proto as proto; +use smg_grpc_client::sglang_proto as proto; use super::error::{SglErrorCode, set_error_message, clear_error_message}; use super::tokenizer::TokenizerHandle; @@ -390,7 +390,7 @@ pub(crate) async fn convert_proto_chunk_to_openai( created: u64, system_fingerprint: Option<&str>, ) -> Result, String> { - use smg::grpc_client::sglang_proto::generate_response::Response::*; + use smg_grpc_client::sglang_proto::generate_response::Response::*; use smg::protocols::chat::{ChatCompletionStreamResponse, ChatMessageDelta, ChatStreamChoice}; match proto_response.response { diff --git a/sgl-model-gateway/bindings/golang/src/postprocessor.rs b/sgl-model-gateway/bindings/golang/src/postprocessor.rs index 735a5c0ef..995fe5bad 100644 --- a/sgl-model-gateway/bindings/golang/src/postprocessor.rs +++ b/sgl-model-gateway/bindings/golang/src/postprocessor.rs @@ -14,7 +14,7 @@ use std::ptr; use std::sync::Arc; use serde_json::Value; -use smg::grpc_client::sglang_proto as proto; +use smg_grpc_client::sglang_proto as proto; use super::error::{SglErrorCode, set_error_message}; use super::grpc_converter::GrpcResponseConverterHandle; diff --git a/sgl-model-gateway/bindings/golang/src/stream.rs b/sgl-model-gateway/bindings/golang/src/stream.rs index a736e35ba..d6b0d28ef 100644 --- a/sgl-model-gateway/bindings/golang/src/stream.rs +++ b/sgl-model-gateway/bindings/golang/src/stream.rs @@ -23,7 +23,7 @@ use tokio::runtime::Runtime; use once_cell::sync::Lazy; use futures_util::StreamExt; -use smg::grpc_client::{sglang_proto as proto, sglang_scheduler::{SglangSchedulerClient, AbortOnDropStream}}; +use smg_grpc_client::{sglang_proto as proto, sglang_scheduler::{SglangSchedulerClient, AbortOnDropStream}}; use super::error::{SglErrorCode, set_error_message}; use super::grpc_converter::{GrpcResponseConverterHandle, convert_proto_chunk_to_openai}; diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index ba3bf3f6b..cbb47c1e1 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -3,17 +3,17 @@ use std::{ time::Duration, }; +use data_connector::{ + create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage, + StorageFactoryConfig, +}; use reqwest::Client; +use smg_mcp::McpManager; use tracing::debug; use crate::{ config::RouterConfig, core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService}, - data_connector::{ - create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage, - StorageFactoryConfig, - }, - mcp::McpManager, middleware::TokenBucket, observability::inflight_tracker::InFlightRequestTracker, policies::PolicyRegistry, @@ -490,7 +490,7 @@ impl AppContextBuilder { // Always create with empty config and defaults debug!("Initializing MCP manager with empty config and default settings (5 min TTL, 100 max connections)"); - let empty_config = crate::mcp::McpConfig { + let empty_config = smg_mcp::McpConfig { servers: Vec::new(), pool: Default::default(), proxy: None, diff --git a/sgl-model-gateway/src/config/builder.rs b/sgl-model-gateway/src/config/builder.rs index 2e307db40..b103d8a67 100644 --- a/sgl-model-gateway/src/config/builder.rs +++ b/sgl-model-gateway/src/config/builder.rs @@ -1,9 +1,11 @@ +use smg_mcp::McpConfig; + use super::{ CircuitBreakerConfig, ConfigError, ConfigResult, DiscoveryConfig, HealthCheckConfig, HistoryBackend, MetricsConfig, OracleConfig, PolicyConfig, PostgresConfig, RedisConfig, RetryConfig, RouterConfig, RoutingMode, TokenizerCacheConfig, TraceConfig, }; -use crate::{core::ConnectionMode, mcp::McpConfig}; +use crate::core::ConnectionMode; /// Builder for RouterConfig that wraps the config itself /// This eliminates field duplication and stays in sync automatically diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index d1b521875..39e0a1df4 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -1,11 +1,11 @@ use std::collections::HashMap; +// Re-export storage config types from data_connector +pub use data_connector::{HistoryBackend, OracleConfig, PostgresConfig, RedisConfig}; use serde::{Deserialize, Serialize}; use super::ConfigResult; use crate::core::ConnectionMode; -// Re-export storage config types from data_connector -pub use crate::data_connector::{HistoryBackend, OracleConfig, PostgresConfig, RedisConfig}; /// Main router configuration #[derive(Debug, Clone, Serialize, Deserialize)] @@ -82,7 +82,7 @@ pub struct RouterConfig { pub ca_certificates: Vec>, /// Loaded from mcp_config_path during config creation #[serde(skip)] - pub mcp_config: Option, + pub mcp_config: Option, /// Enable WASM support #[serde(default)] pub enable_wasm: bool, diff --git a/sgl-model-gateway/src/core/job_queue.rs b/sgl-model-gateway/src/core/job_queue.rs index e45b1dff6..862cd7165 100644 --- a/sgl-model-gateway/src/core/job_queue.rs +++ b/sgl-model-gateway/src/core/job_queue.rs @@ -10,8 +10,10 @@ use std::{ }; use dashmap::DashMap; +use smg_mcp::McpConfig; use tokio::sync::{mpsc, Semaphore}; use tracing::{debug, error, info, warn}; +use wfaas::WorkflowId; use crate::{ app_context::AppContext, @@ -24,9 +26,7 @@ use crate::{ McpServerConfigRequest, TokenizerConfigRequest, TokenizerRemovalRequest, WasmModuleConfigRequest, WasmModuleRemovalRequest, }, - mcp::McpConfig, protocols::worker_spec::{JobStatus, WorkerConfigRequest, WorkerUpdateRequest}, - workflow::WorkflowId, }; /// Job types for control plane operations diff --git a/sgl-model-gateway/src/core/steps/mcp_registration.rs b/sgl-model-gateway/src/core/steps/mcp_registration.rs index 92a6108c1..641fb222d 100644 --- a/sgl-model-gateway/src/core/steps/mcp_registration.rs +++ b/sgl-model-gateway/src/core/steps/mcp_registration.rs @@ -1,18 +1,15 @@ use std::{sync::Arc, time::Duration}; use async_trait::async_trait; +use smg_mcp::{config::McpServerConfig, manager::McpManager}; use tracing::{debug, error, info, warn}; +use wfaas::{ + BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult, + WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, +}; use super::workflow_data::McpWorkflowData; -use crate::{ - app_context::AppContext, - mcp::{config::McpServerConfig, manager::McpManager}, - observability::metrics::Metrics, - workflow::{ - BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, - StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, - }, -}; +use crate::{app_context::AppContext, observability::metrics::Metrics}; /// MCP server connection configuration #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] diff --git a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs index 6e98673ae..639ae2928 100644 --- a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs +++ b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs @@ -12,6 +12,10 @@ use std::{sync::Arc, time::Duration}; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use tracing::{error, info}; +use wfaas::{ + BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult, + WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, +}; use super::workflow_data::TokenizerWorkflowData; use crate::{ @@ -23,10 +27,6 @@ use crate::{ registry::LoadOutcome, traits::Tokenizer, }, - workflow::{ - BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, - StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, - }, }; /// Configuration for adding a tokenizer diff --git a/sgl-model-gateway/src/core/steps/wasm_module_registration.rs b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs index b188f2ed9..268a39f9b 100644 --- a/sgl-model-gateway/src/core/steps/wasm_module_registration.rs +++ b/sgl-model-gateway/src/core/steps/wasm_module_registration.rs @@ -9,15 +9,15 @@ use sha2::{Digest, Sha256}; use tracing::{debug, info, warn}; use uuid::Uuid; use wasmtime::{component::Component, Config, Engine}; +use wfaas::{ + BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult, + WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, +}; use super::workflow_data::WasmRegistrationWorkflowData; use crate::{ app_context::AppContext, wasm::module::{WasmModule, WasmModuleDescriptor, WasmModuleMeta}, - workflow::{ - BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, - StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, - }, }; /// WASM module registration request diff --git a/sgl-model-gateway/src/core/steps/wasm_module_removal.rs b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs index 5a89a6245..ad524c431 100644 --- a/sgl-model-gateway/src/core/steps/wasm_module_removal.rs +++ b/sgl-model-gateway/src/core/steps/wasm_module_removal.rs @@ -3,15 +3,13 @@ use std::{sync::Arc, time::Duration}; use async_trait::async_trait; use tracing::{debug, info}; use uuid::Uuid; +use wfaas::{ + FailureAction, StepDefinition, StepExecutor, StepId, StepResult, WorkflowContext, + WorkflowDefinition, WorkflowError, WorkflowResult, +}; use super::workflow_data::WasmRemovalWorkflowData; -use crate::{ - app_context::AppContext, - workflow::{ - FailureAction, StepDefinition, StepExecutor, StepId, StepResult, WorkflowContext, - WorkflowDefinition, WorkflowError, WorkflowResult, - }, -}; +use crate::app_context::AppContext; /// WASM module removal request #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] diff --git a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs index cc91a9344..fd6a7d94d 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/create_workers.rs @@ -4,15 +4,13 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use async_trait::async_trait; use tracing::{debug, info}; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; -use crate::{ - core::{ - circuit_breaker::CircuitBreakerConfig, - steps::workflow_data::{ExternalWorkerWorkflowData, WorkerList}, - worker::{HealthConfig, RuntimeType, WorkerType}, - BasicWorkerBuilder, ConnectionMode, Worker, - }, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +use crate::core::{ + circuit_breaker::CircuitBreakerConfig, + steps::workflow_data::{ExternalWorkerWorkflowData, WorkerList}, + worker::{HealthConfig, RuntimeType, WorkerType}, + BasicWorkerBuilder, ConnectionMode, Worker, }; /// Normalize URL for external APIs (ensure https://). diff --git a/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs b/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs index 05edd1a7e..8dd58876b 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/discover_models.rs @@ -8,14 +8,12 @@ use regex::Regex; use reqwest::Client; use serde::Deserialize; use tracing::{debug, info}; +use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; -use crate::{ - core::{ - model_card::{ModelCard, ProviderType}, - model_type::ModelType, - steps::workflow_data::ExternalWorkerWorkflowData, - }, - workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +use crate::core::{ + model_card::{ModelCard, ProviderType}, + model_type::ModelType, + steps::workflow_data::ExternalWorkerWorkflowData, }; // HTTP client for API calls diff --git a/sgl-model-gateway/src/core/steps/worker/external/mod.rs b/sgl-model-gateway/src/core/steps/worker/external/mod.rs index 017b71db7..04f26371f 100644 --- a/sgl-model-gateway/src/core/steps/worker/external/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/external/mod.rs @@ -13,13 +13,12 @@ pub use discover_models::{ group_models_into_cards, infer_model_type_from_id, DiscoverModelsStep, ModelInfo, ModelsResponse, }; +use wfaas::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition}; use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep}; use crate::{ - app_context::AppContext, - core::steps::workflow_data::ExternalWorkerWorkflowData, + app_context::AppContext, core::steps::workflow_data::ExternalWorkerWorkflowData, protocols::worker_spec::WorkerConfigRequest, - workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition}, }; /// Create external worker registration workflow definition. diff --git a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs index 12b577d7a..7acea535e 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/create_worker.rs @@ -4,6 +4,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use async_trait::async_trait; use tracing::debug; +use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use crate::{ app_context::AppContext, @@ -15,7 +16,6 @@ use crate::{ BasicWorkerBuilder, ConnectionMode, DPAwareWorkerBuilder, Worker, UNKNOWN_MODEL_ID, }, protocols::worker_spec::WorkerConfigRequest, - workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, }; /// Step 3: Create worker object(s) with merged configuration + metadata. diff --git a/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs b/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs index 9908de6a6..4f91bb535 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs @@ -5,12 +5,12 @@ use std::time::Duration; use async_trait::async_trait; use reqwest::Client; use tracing::debug; +use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use super::strip_protocol; use crate::{ core::{steps::workflow_data::LocalWorkerWorkflowData, ConnectionMode}, routers::grpc::client::GrpcClient, - workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, }; /// Try HTTP health check. diff --git a/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs b/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs index c2a52b32f..58d1a855b 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs @@ -2,12 +2,10 @@ use async_trait::async_trait; use tracing::debug; +use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use super::discover_metadata::get_server_info; -use crate::{ - core::{steps::workflow_data::LocalWorkerWorkflowData, UNKNOWN_MODEL_ID}, - workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, -}; +use crate::core::{steps::workflow_data::LocalWorkerWorkflowData, UNKNOWN_MODEL_ID}; /// DP (Data Parallel) information for a worker. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] diff --git a/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs index 6524d4295..3ad55e186 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs @@ -8,12 +8,12 @@ use reqwest::Client; use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::{debug, warn}; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use super::strip_protocol; use crate::{ core::{steps::workflow_data::LocalWorkerWorkflowData, ConnectionMode}, routers::grpc::client::GrpcClient, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, }; // HTTP client for metadata fetching diff --git a/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs b/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs index 1dc52f32c..115b985fb 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs @@ -2,12 +2,10 @@ use async_trait::async_trait; use tracing::debug; +use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use super::find_workers_by_url; -use crate::{ - core::steps::workflow_data::WorkerUpdateWorkflowData, - workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, -}; +use crate::core::steps::workflow_data::WorkerUpdateWorkflowData; /// Step to find workers to update based on URL. /// diff --git a/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs b/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs index 22635a1ea..36720c1a5 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs @@ -4,12 +4,10 @@ use std::collections::HashSet; use async_trait::async_trait; use tracing::debug; +use wfaas::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use super::find_workers_by_url; -use crate::{ - core::steps::workflow_data::{WorkerList, WorkerRemovalWorkflowData}, - workflow::{StepExecutor, StepId, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, -}; +use crate::core::steps::workflow_data::{WorkerList, WorkerRemovalWorkflowData}; /// Request structure for worker removal. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] diff --git a/sgl-model-gateway/src/core/steps/worker/local/mod.rs b/sgl-model-gateway/src/core/steps/worker/local/mod.rs index df00fb09a..68f9ea184 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -33,6 +33,7 @@ pub use submit_tokenizer_job::SubmitTokenizerJobStep; pub use update_policies_for_worker::UpdatePoliciesForWorkerStep; pub use update_remaining_policies::UpdateRemainingPoliciesStep; pub use update_worker_properties::UpdateWorkerPropertiesStep; +use wfaas::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition}; use super::shared::{ActivateWorkersStep, RegisterWorkersStep, UpdatePoliciesStep}; use crate::{ @@ -45,7 +46,6 @@ use crate::{ Worker, WorkerRegistry, }, protocols::worker_spec::{WorkerConfigRequest, WorkerUpdateRequest}, - workflow::{BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, WorkflowDefinition}, }; /// Find workers by URL, supporting both DP-aware (prefix match) and regular (exact match) modes. diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs index a68ab1758..c720e0c47 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs @@ -2,11 +2,9 @@ use async_trait::async_trait; use tracing::debug; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; -use crate::{ - core::steps::workflow_data::WorkerRemovalWorkflowData, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, -}; +use crate::core::steps::workflow_data::WorkerRemovalWorkflowData; /// Step to remove workers from the policy registry. /// diff --git a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs index 808d995ed..ae214928b 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs @@ -4,11 +4,10 @@ use std::collections::HashSet; use async_trait::async_trait; use tracing::{debug, warn}; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use crate::{ - core::steps::workflow_data::WorkerRemovalWorkflowData, - observability::metrics::Metrics, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, + core::steps::workflow_data::WorkerRemovalWorkflowData, observability::metrics::Metrics, }; /// Step to remove workers from the worker registry. diff --git a/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs b/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs index a85fe35f9..4996f97a2 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs @@ -6,6 +6,7 @@ use async_trait::async_trait; use tracing::{debug, info, warn}; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; use crate::{ core::{ @@ -13,7 +14,6 @@ use crate::{ Job, }, tokenizer::TokenizerRegistry, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, }; /// Step: Submit tokenizer registration job for the worker's model diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs b/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs index 1ff926389..05215ec25 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs @@ -4,11 +4,9 @@ use std::collections::HashSet; use async_trait::async_trait; use tracing::debug; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; -use crate::{ - core::steps::workflow_data::WorkerUpdateWorkflowData, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, -}; +use crate::core::steps::workflow_data::WorkerUpdateWorkflowData; /// Step to update policies for updated workers. /// diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs b/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs index 0bdb10471..b4da51669 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs @@ -2,11 +2,9 @@ use async_trait::async_trait; use tracing::{debug, info}; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; -use crate::{ - core::steps::workflow_data::WorkerRemovalWorkflowData, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, -}; +use crate::core::steps::workflow_data::WorkerRemovalWorkflowData; /// Step to update cache-aware policies for remaining workers. /// diff --git a/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs b/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs index 48a459876..c326ebaed 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs @@ -4,12 +4,10 @@ use std::sync::Arc; use async_trait::async_trait; use tracing::{debug, info}; +use wfaas::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}; -use crate::{ - core::{ - steps::workflow_data::WorkerUpdateWorkflowData, BasicWorkerBuilder, HealthConfig, Worker, - }, - workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, +use crate::core::{ + steps::workflow_data::WorkerUpdateWorkflowData, BasicWorkerBuilder, HealthConfig, Worker, }; /// Step to update worker properties. diff --git a/sgl-model-gateway/src/core/steps/worker/shared/activate.rs b/sgl-model-gateway/src/core/steps/worker/shared/activate.rs index a56f355a2..877ad6af8 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/activate.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/activate.rs @@ -2,14 +2,12 @@ use async_trait::async_trait; use tracing::info; - -use crate::{ - core::steps::workflow_data::WorkerRegistrationData, - workflow::{ - StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, - }, +use wfaas::{ + StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, }; +use crate::core::steps::workflow_data::WorkerRegistrationData; + /// Unified step to activate workers by marking them as healthy. /// /// This is the final step in any worker registration workflow. diff --git a/sgl-model-gateway/src/core/steps/worker/shared/register.rs b/sgl-model-gateway/src/core/steps/worker/shared/register.rs index fe239bbb2..63be0bd34 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/register.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/register.rs @@ -4,15 +4,12 @@ use std::{collections::HashSet, sync::Arc}; use async_trait::async_trait; use tracing::debug; - -use crate::{ - core::steps::workflow_data::WorkerRegistrationData, - observability::metrics::Metrics, - workflow::{ - StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, - }, +use wfaas::{ + StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, }; +use crate::{core::steps::workflow_data::WorkerRegistrationData, observability::metrics::Metrics}; + /// Unified step to register workers in the registry. /// /// Works with both single workers and batches. Always expects `workers` key diff --git a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs index c9527d73b..bc586d206 100644 --- a/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs +++ b/sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs @@ -4,14 +4,12 @@ use std::sync::Arc; use async_trait::async_trait; use tracing::{debug, warn}; - -use crate::{ - core::{steps::workflow_data::WorkerRegistrationData, Worker}, - workflow::{ - StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, - }, +use wfaas::{ + StepExecutor, StepResult, WorkflowContext, WorkflowData, WorkflowError, WorkflowResult, }; +use crate::core::{steps::workflow_data::WorkerRegistrationData, Worker}; + /// Unified step to update policy registry for registered workers. /// /// Handles both local workers (same model, possibly DP-aware) and diff --git a/sgl-model-gateway/src/core/steps/workflow_data.rs b/sgl-model-gateway/src/core/steps/workflow_data.rs index 570a3904b..9bb90b75e 100644 --- a/sgl-model-gateway/src/core/steps/workflow_data.rs +++ b/sgl-model-gateway/src/core/steps/workflow_data.rs @@ -13,6 +13,7 @@ use std::{collections::HashMap, sync::Arc}; use serde::{Deserialize, Serialize}; +use wfaas::{WorkflowData, WorkflowError}; use super::{ mcp_registration::McpServerConfigRequest, tokenizer_registration::TokenizerConfigRequest, @@ -30,7 +31,6 @@ use crate::{ WorkerConfigRequest as ProtocolWorkerConfigRequest, WorkerUpdateRequest as ProtocolWorkerUpdateRequest, }, - workflow::{WorkflowData, WorkflowError}, }; // ============================================================================ diff --git a/sgl-model-gateway/src/core/steps/workflow_engines.rs b/sgl-model-gateway/src/core/steps/workflow_engines.rs index 3c4088819..e994ef3d2 100644 --- a/sgl-model-gateway/src/core/steps/workflow_engines.rs +++ b/sgl-model-gateway/src/core/steps/workflow_engines.rs @@ -5,6 +5,8 @@ use std::sync::Arc; +use wfaas::{EventSubscriber, InMemoryStore, WorkflowEngine}; + use super::{ create_external_worker_workflow, create_local_worker_workflow, create_mcp_registration_workflow, create_tokenizer_registration_workflow, @@ -13,10 +15,7 @@ use super::{ LocalWorkerWorkflowData, McpWorkflowData, TokenizerWorkflowData, WasmRegistrationWorkflowData, WasmRemovalWorkflowData, WorkerRemovalWorkflowData, WorkerUpdateWorkflowData, }; -use crate::{ - config::RouterConfig, - workflow::{EventSubscriber, InMemoryStore, WorkflowEngine}, -}; +use crate::config::RouterConfig; /// Type alias for local worker workflow engine pub type LocalWorkerEngine = diff --git a/sgl-model-gateway/src/core/worker_registry.rs b/sgl-model-gateway/src/core/worker_registry.rs index f91c9769c..c23dd8719 100644 --- a/sgl-model-gateway/src/core/worker_registry.rs +++ b/sgl-model-gateway/src/core/worker_registry.rs @@ -14,6 +14,7 @@ use std::sync::{Arc, RwLock}; use dashmap::DashMap; +use smg_mesh::OptionalMeshSyncManager; use uuid::Uuid; use crate::{ @@ -22,7 +23,6 @@ use crate::{ worker::{HealthChecker, RuntimeType, WorkerType}, ConnectionMode, Worker, }, - mesh::OptionalMeshSyncManager, observability::metrics::Metrics, }; diff --git a/sgl-model-gateway/src/lib.rs b/sgl-model-gateway/src/lib.rs index 0166819d0..9f92b12ec 100644 --- a/sgl-model-gateway/src/lib.rs +++ b/sgl-model-gateway/src/lib.rs @@ -2,10 +2,6 @@ pub mod app_context; pub use smg_auth as auth; pub mod config; pub mod core; -pub use data_connector; -pub use smg_grpc_client as grpc_client; -pub use smg_mcp as mcp; -pub use smg_mesh as mesh; pub mod middleware; pub mod observability; pub mod policies; @@ -18,4 +14,3 @@ pub use llm_tokenizer as tokenizer; pub use tool_parser; pub mod version; pub mod wasm; -pub use wfaas as workflow; diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index aacea0d9e..a04e10f05 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -11,7 +11,6 @@ use smg::{ TraceConfig, }, core::ConnectionMode, - mesh::service::MeshServerConfig, observability::{ metrics::PrometheusConfig, otel_trace::{is_otel_enabled, shutdown_otel}, @@ -20,6 +19,7 @@ use smg::{ service_discovery::ServiceDiscoveryConfig, version, }; +use smg_mesh::service::MeshServerConfig; fn parse_prefill_args() -> Vec<(String, Option)> { let args: Vec = std::env::args().collect(); let mut prefill_entries = Vec::new(); diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index ae1dae185..aabe70183 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -331,7 +331,7 @@ pub(crate) fn init_metrics() { describe_counter!("smg_db_items_stored", "Total items stored by storage_type"); // Initialize mesh metrics - crate::mesh::metrics::init_mesh_metrics(); + smg_mesh::metrics::init_mesh_metrics(); } pub fn start_prometheus(config: PrometheusConfig) { diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 4fb446106..dfc3ef646 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -64,16 +64,14 @@ use std::sync::Arc; use async_trait::async_trait; use dashmap::DashMap; use rand::Rng; +use smg_mesh::{tree_ops::TreeOperation, OptionalMeshSyncManager}; use tracing::{debug, warn}; use super::{ get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask, CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo, }; -use crate::{ - core::{Worker, UNKNOWN_MODEL_ID}, - mesh::{tree_ops::TreeOperation, OptionalMeshSyncManager}, -}; +use crate::core::{Worker, UNKNOWN_MODEL_ID}; /// Cache-aware routing policy /// @@ -323,7 +321,7 @@ impl CacheAwarePolicy { // Sync insert operation to mesh if enabled (no-op if mesh is not enabled) if let Some(ref mesh_sync) = self.mesh_sync { - use crate::mesh::tree_ops::TreeInsertOp; + use smg_mesh::tree_ops::TreeInsertOp; let op = TreeOperation::Insert(TreeInsertOp { text: text.to_string(), tenant: worker_url.to_string(), @@ -427,7 +425,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Sync insert operation to mesh if enabled (no-op if mesh is not enabled) if let Some(ref mesh_sync) = self.mesh_sync { - use crate::mesh::tree_ops::TreeInsertOp; + use smg_mesh::tree_ops::TreeInsertOp; let op = TreeOperation::Insert(TreeInsertOp { text: text.to_string(), tenant: workers[idx].url().to_string(), @@ -452,7 +450,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { // Sync removal to mesh if enabled (no-op if mesh is not enabled) if let Some(ref mesh_sync) = self.mesh_sync { - use crate::mesh::tree_ops::TreeRemoveOp; + use smg_mesh::tree_ops::TreeRemoveOp; let op = TreeOperation::Remove(TreeRemoveOp { tenant: tenant_url.to_string(), }); @@ -680,7 +678,7 @@ mod tests { async fn test_cache_aware_sync_tree_operation_to_mesh() { use std::sync::Arc; - use crate::mesh::{stores::StateStores, sync::MeshSyncManager}; + use smg_mesh::{stores::StateStores, sync::MeshSyncManager}; let stores = Arc::new(StateStores::with_self_name("node1".to_string())); let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string())); @@ -724,7 +722,7 @@ mod tests { fn test_cache_aware_restore_tree_state_from_mesh() { use std::sync::Arc; - use crate::mesh::{ + use smg_mesh::{ stores::StateStores, sync::MeshSyncManager, tree_ops::{TreeInsertOp, TreeOperation}, @@ -783,7 +781,7 @@ mod tests { fn test_cache_aware_apply_remote_tree_operation() { use std::sync::Arc; - use crate::mesh::{ + use smg_mesh::{ stores::StateStores, sync::MeshSyncManager, tree_ops::{TreeInsertOp, TreeOperation}, @@ -816,7 +814,7 @@ mod tests { fn test_cache_aware_multi_node_consistency() { use std::sync::Arc; - use crate::mesh::{ + use smg_mesh::{ stores::StateStores, sync::MeshSyncManager, tree_ops::{TreeInsertOp, TreeOperation}, diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index fae218a21..debc385a3 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -6,11 +6,9 @@ use std::{fmt::Debug, sync::Arc}; use async_trait::async_trait; +use smg_mesh::OptionalMeshSyncManager; -use crate::{ - core::{HashRing, Worker}, - mesh::OptionalMeshSyncManager, -}; +use crate::core::{HashRing, Worker}; mod bucket; mod cache_aware; diff --git a/sgl-model-gateway/src/policies/registry.rs b/sgl-model-gateway/src/policies/registry.rs index 19f6435b5..2e0e4b9eb 100644 --- a/sgl-model-gateway/src/policies/registry.rs +++ b/sgl-model-gateway/src/policies/registry.rs @@ -2,6 +2,7 @@ use std::sync::{Arc, OnceLock, RwLock}; use dashmap::DashMap; use serde_json; +use smg_mesh::OptionalMeshSyncManager; use tracing::{debug, info, warn}; /// Policy Registry for managing model-to-policy mappings @@ -11,7 +12,7 @@ use tracing::{debug, info, warn}; /// All subsequent workers of the same model use the established policy. /// When the last worker of a model is removed, the policy mapping is cleaned up. use super::{BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PolicyFactory}; -use crate::{config::types::PolicyConfig, core::Worker, mesh::OptionalMeshSyncManager}; +use crate::{config::types::PolicyConfig, core::Worker}; /// Registry for managing model-to-policy mappings #[derive(Clone)] @@ -391,7 +392,7 @@ impl PolicyRegistry { pub fn apply_remote_tree_operation( &self, model_id: &str, - operation: &crate::mesh::tree_ops::TreeOperation, + operation: &smg_mesh::tree_ops::TreeOperation, ) { // Try to find the policy for this model if let Some(policy) = self.get_policy(model_id) { diff --git a/sgl-model-gateway/src/routers/conversations/handlers.rs b/sgl-model-gateway/src/routers/conversations/handlers.rs index 07babe4f8..0dec90467 100644 --- a/sgl-model-gateway/src/routers/conversations/handlers.rs +++ b/sgl-model-gateway/src/routers/conversations/handlers.rs @@ -8,17 +8,14 @@ use axum::{ Json, }; use chrono::Utc; +use data_connector::{ + Conversation, ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, + ConversationStorage, ListParams, NewConversation, NewConversationItem, SortOrder, +}; use serde_json::{json, Value}; use tracing::{info, warn}; -use crate::{ - data_connector::{ - Conversation, ConversationId, ConversationItem, ConversationItemId, - ConversationItemStorage, ConversationStorage, ListParams, NewConversation, - NewConversationItem, SortOrder, - }, - routers::persistence_utils::item_to_json, -}; +use crate::routers::persistence_utils::item_to_json; // ============================================================================ // Constants diff --git a/sgl-model-gateway/src/routers/grpc/client.rs b/sgl-model-gateway/src/routers/grpc/client.rs index 75b520acc..816b3530d 100644 --- a/sgl-model-gateway/src/routers/grpc/client.rs +++ b/sgl-model-gateway/src/routers/grpc/client.rs @@ -2,8 +2,9 @@ use std::sync::Arc; +use smg_grpc_client::{SglangSchedulerClient, VllmEngineClient}; + use crate::{ - grpc_client::{SglangSchedulerClient, VllmEngineClient}, observability::otel_trace::OtelTraceInjector, routers::grpc::proto_wrapper::{ ProtoEmbedRequest, ProtoEmbedResponse, ProtoGenerateRequest, ProtoStream, @@ -159,8 +160,8 @@ impl GrpcClient { /// Unified ModelInfo wrapper pub enum ModelInfo { - Sglang(Box), - Vllm(crate::grpc_client::vllm_proto::GetModelInfoResponse), + Sglang(Box), + Vllm(smg_grpc_client::vllm_proto::GetModelInfoResponse), } impl ModelInfo { diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/context.rs b/sgl-model-gateway/src/routers/grpc/common/responses/context.rs index f2e0336f3..79050286c 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/context.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/context.rs @@ -4,11 +4,10 @@ use std::sync::{Arc, RwLock as StdRwLock}; -use crate::{ - data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, - mcp::McpManager, - routers::grpc::{context::SharedComponents, pipeline::RequestPipeline}, -}; +use data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}; +use smg_mcp::McpManager; + +use crate::routers::grpc::{context::SharedComponents, pipeline::RequestPipeline}; /// Context for /v1/responses endpoint /// diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs b/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs index b28a5a10d..851d0d33a 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs @@ -3,9 +3,10 @@ //! These handlers are used by both pipelines for retrieving and cancelling responses. use axum::response::{IntoResponse, Response}; +use data_connector::ResponseId; use super::ResponsesContext; -use crate::{data_connector::ResponseId, routers::error}; +use crate::routers::error; /// Implementation for GET /v1/responses/{response_id} /// diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/streaming.rs b/sgl-model-gateway/src/routers/grpc/common/responses/streaming.rs index b88225354..2f055e457 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/streaming.rs @@ -5,12 +5,12 @@ use std::collections::HashMap; use axum::{body::Body, http::StatusCode, response::Response}; use bytes::Bytes; use serde_json::json; +use smg_mcp as mcp; use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; use uuid::Uuid; use crate::{ - mcp, protocols::{ chat::ChatCompletionStreamResponse, common::{Usage, UsageInfo}, diff --git a/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs b/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs index 303b9f533..87430a32f 100644 --- a/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs +++ b/sgl-model-gateway/src/routers/grpc/common/responses/utils.rs @@ -3,13 +3,13 @@ use std::sync::Arc; use axum::response::Response; +use data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}; use serde_json::to_value; +use smg_mcp::McpManager; use tracing::{debug, error, warn}; use crate::{ core::WorkerRegistry, - data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, - mcp::McpManager, protocols::{ common::Tool, responses::{ResponseTool, ResponseToolType, ResponsesRequest, ResponsesResponse}, diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/helpers.rs b/sgl-model-gateway/src/routers/grpc/common/stages/helpers.rs index 3dbd7ad78..2fa7a5555 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/helpers.rs +++ b/sgl-model-gateway/src/routers/grpc/common/stages/helpers.rs @@ -3,12 +3,10 @@ use std::sync::Arc; use rand::Rng; +use smg_grpc_client::sglang_proto::DisaggregatedParams; use tracing::debug; -use crate::{ - core::Worker, grpc_client::sglang_proto::DisaggregatedParams, - routers::grpc::proto_wrapper::ProtoGenerateRequest, -}; +use crate::{core::Worker, routers::grpc::proto_wrapper::ProtoGenerateRequest}; /// Inject PD bootstrap metadata into a gRPC request /// diff --git a/sgl-model-gateway/src/routers/grpc/harmony/processor.rs b/sgl-model-gateway/src/routers/grpc/harmony/processor.rs index 908214d71..ae18af954 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/processor.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/processor.rs @@ -3,11 +3,13 @@ use std::sync::Arc; use axum::response::Response; +use smg_grpc_client::sglang_proto::generate_complete::MatchedStop::{ + MatchedStopStr, MatchedTokenId, +}; use tracing::error; use super::HarmonyParserAdapter; use crate::{ - grpc_client::sglang_proto::generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId}, protocols::{ chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse}, common::{CompletionTokensDetails, ToolCall, Usage}, diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs index 704c34d6b..e84bd212d 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs @@ -1,14 +1,14 @@ //! Shared helpers and state tracking for Harmony Responses use axum::response::Response; +use data_connector::ResponseId; use serde_json::{from_value, json, to_string, Value}; +use smg_mcp as mcp; use tracing::{debug, error, warn}; use uuid::Uuid; use super::execution::ToolResult; use crate::{ - data_connector::ResponseId, - mcp, protocols::{ common::{ToolCall, ToolChoice, ToolChoiceValue}, responses::{ diff --git a/sgl-model-gateway/src/routers/grpc/harmony/responses/execution.rs b/sgl-model-gateway/src/routers/grpc/harmony/responses/execution.rs index a7ccfde45..a7f362c43 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/responses/execution.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/responses/execution.rs @@ -4,11 +4,11 @@ use std::{sync::Arc, time::Instant}; use axum::response::Response; use serde_json::{from_str, json, to_string, to_value, Value}; +use smg_mcp::{self as mcp, McpManager}; use tracing::{debug, error, warn}; use super::common::McpCallTracking; use crate::{ - mcp::{self, McpManager}, observability::metrics::{metrics_labels, Metrics}, protocols::{ common::{Function, ToolCall}, diff --git a/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs b/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs index 8b97131b3..d97bbab5d 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs @@ -11,6 +11,9 @@ use axum::{body::Body, http::StatusCode, response::Response}; use bytes::Bytes; use http::header::{HeaderValue, CONTENT_TYPE}; use serde_json::json; +use smg_grpc_client::sglang_proto::generate_complete::MatchedStop::{ + MatchedStopStr, MatchedTokenId, +}; use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::{debug, error}; @@ -19,7 +22,6 @@ use super::{ processor::ResponsesIterationResult, types::HarmonyChannelDelta, HarmonyParserAdapter, }; use crate::{ - grpc_client::sglang_proto::generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId}, observability::metrics::{metrics_labels, Metrics, StreamingMetricsParams}, protocols::{ chat::{ diff --git a/sgl-model-gateway/src/routers/grpc/mod.rs b/sgl-model-gateway/src/routers/grpc/mod.rs index bba0ea8ff..990d7a5f9 100644 --- a/sgl-model-gateway/src/routers/grpc/mod.rs +++ b/sgl-model-gateway/src/routers/grpc/mod.rs @@ -1,6 +1,8 @@ //! gRPC router implementations -use crate::{grpc_client::sglang_proto::MultimodalInputs, protocols::common::StringOrArray}; +use smg_grpc_client::sglang_proto::MultimodalInputs; + +use crate::protocols::common::StringOrArray; pub mod client; // Used by core/ pub(crate) mod common; diff --git a/sgl-model-gateway/src/routers/grpc/proto_wrapper.rs b/sgl-model-gateway/src/routers/grpc/proto_wrapper.rs index 44e42b0b6..d27a9f381 100644 --- a/sgl-model-gateway/src/routers/grpc/proto_wrapper.rs +++ b/sgl-model-gateway/src/routers/grpc/proto_wrapper.rs @@ -4,8 +4,7 @@ //! allowing the router to work with either backend transparently. use futures_util::StreamExt; - -use crate::grpc_client::{ +use smg_grpc_client::{ sglang_proto::{self as sglang, generate_complete::MatchedStop}, sglang_scheduler::AbortOnDropStream as SglangStream, vllm_engine::AbortOnDropStream as VllmStream, diff --git a/sgl-model-gateway/src/routers/grpc/regular/processor.rs b/sgl-model-gateway/src/routers/grpc/regular/processor.rs index aaad3078e..35e2c46eb 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/processor.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/processor.rs @@ -6,10 +6,10 @@ use std::{sync::Arc, time::Instant}; use serde_json::Value; +use smg_grpc_client::sglang_proto::generate_complete::MatchedStop; use tracing::error; use crate::{ - grpc_client::sglang_proto::generate_complete::MatchedStop, protocols::{ chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse}, common::{FunctionCallResponse, ToolCall, ToolChoice, ToolChoiceValue}, diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs index 7b9ab0725..7dee8ef74 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/common.rs @@ -9,13 +9,13 @@ use std::sync::Arc; use axum::response::Response; +use data_connector::{self, ConversationId, ResponseId}; use serde_json::{json, Value}; +use smg_mcp::{self as mcp, McpManager}; use tracing::{debug, warn}; use uuid::Uuid; use crate::{ - data_connector::{self, ConversationId, ResponseId}, - mcp::{self, McpManager}, protocols::{ chat::ChatCompletionRequest, common::{Function, Tool, ToolChoice, ToolChoiceValue}, diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs index 8f4043907..a24683f89 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs @@ -17,6 +17,7 @@ use axum::{ response::Response, }; use bytes::Bytes; +use data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}; use futures_util::StreamExt; use serde_json::{json, Value}; use tokio::sync::mpsc; @@ -32,7 +33,6 @@ use super::{ conversions, }; use crate::{ - data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, observability::metrics::{metrics_labels, Metrics}, protocols::{ chat::{ diff --git a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs index c2673183f..06bd2b97f 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs @@ -8,12 +8,14 @@ use axum::{body::Body, http::StatusCode, response::Response}; use bytes::Bytes; use http::header::{HeaderValue, CONTENT_TYPE}; use serde_json::{json, Value}; +use smg_grpc_client::sglang_proto::generate_complete::MatchedStop::{ + MatchedStopStr, MatchedTokenId, +}; use tokio::sync::{mpsc, mpsc::UnboundedSender}; use tokio_stream::wrappers::UnboundedReceiverStream; use tracing::{debug, error, warn}; use crate::{ - grpc_client::sglang_proto::generate_complete::MatchedStop::{MatchedStopStr, MatchedTokenId}, observability::metrics::{metrics_labels, Metrics, StreamingMetricsParams}, protocols::{ chat::{ChatCompletionRequest, ChatCompletionStreamResponse}, diff --git a/sgl-model-gateway/src/routers/grpc/utils.rs b/sgl-model-gateway/src/routers/grpc/utils.rs index 08cbec032..42dfcbc60 100644 --- a/sgl-model-gateway/src/routers/grpc/utils.rs +++ b/sgl-model-gateway/src/routers/grpc/utils.rs @@ -5,6 +5,7 @@ use std::{collections::HashMap, sync::Arc}; use axum::response::Response; use http::StatusCode; use serde_json::{json, Map, Value}; +use smg_grpc_client::sglang_proto::{InputLogProbs, OutputLogProbs}; use tracing::{error, warn}; use uuid::Uuid; @@ -16,7 +17,6 @@ use super::{ }; use crate::{ core::Worker, - grpc_client::sglang_proto::{InputLogProbs, OutputLogProbs}, observability::metrics::metrics_labels, protocols::{ chat::{ChatCompletionRequest, ChatMessage}, diff --git a/sgl-model-gateway/src/routers/mcp_utils.rs b/sgl-model-gateway/src/routers/mcp_utils.rs index ddba6f40c..4b5625dc6 100644 --- a/sgl-model-gateway/src/routers/mcp_utils.rs +++ b/sgl-model-gateway/src/routers/mcp_utils.rs @@ -5,12 +5,10 @@ use std::sync::Arc; +use smg_mcp::{McpManager, McpServerConfig, McpTransport}; use tracing::warn; -use crate::{ - mcp::{McpManager, McpServerConfig, McpTransport}, - protocols::responses::{ResponseTool, ResponseToolType}, -}; +use crate::protocols::responses::{ResponseTool, ResponseToolType}; // ============================================================================ // Constants diff --git a/sgl-model-gateway/src/routers/openai/context.rs b/sgl-model-gateway/src/routers/openai/context.rs index 6f3dbebe6..7bb3521bb 100644 --- a/sgl-model-gateway/src/routers/openai/context.rs +++ b/sgl-model-gateway/src/routers/openai/context.rs @@ -3,13 +3,13 @@ use std::sync::Arc; use axum::http::HeaderMap; +use data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}; use serde_json::Value; +use smg_mcp::McpManager; use super::provider::Provider; use crate::{ core::Worker, - data_connector::{ConversationItemStorage, ConversationStorage, ResponseStorage}, - mcp::McpManager, protocols::{chat::ChatCompletionRequest, responses::ResponsesRequest}, }; diff --git a/sgl-model-gateway/src/routers/openai/responses/mcp.rs b/sgl-model-gateway/src/routers/openai/responses/mcp.rs index 75425ba90..6f02d4372 100644 --- a/sgl-model-gateway/src/routers/openai/responses/mcp.rs +++ b/sgl-model-gateway/src/routers/openai/responses/mcp.rs @@ -13,11 +13,11 @@ use std::{io, sync::Arc}; use axum::http::HeaderMap; use bytes::Bytes; use serde_json::{json, to_value, Value}; +use smg_mcp as mcp; use tokio::sync::mpsc; use tracing::{debug, info, warn}; use crate::{ - mcp, protocols::{ event_types::{is_function_call_type, ItemType, McpEvent, OutputItemEvent}, responses::{generate_id, ResponseInput, ResponsesRequest}, diff --git a/sgl-model-gateway/src/routers/openai/responses/streaming.rs b/sgl-model-gateway/src/routers/openai/responses/streaming.rs index 2c820517b..2c7da80af 100644 --- a/sgl-model-gateway/src/routers/openai/responses/streaming.rs +++ b/sgl-model-gateway/src/routers/openai/responses/streaming.rs @@ -424,7 +424,7 @@ pub(super) fn send_final_response_event( tx: &mpsc::UnboundedSender>, sequence_number: &mut u64, state: &ToolLoopState, - active_mcp: Option<&Arc>, + active_mcp: Option<&Arc>, ctx: &StreamingEventContext<'_>, ) -> bool { let mut final_response = match handler.snapshot_final_response() { @@ -637,7 +637,7 @@ pub(super) async fn handle_streaming_with_tool_interception( client: &reqwest::Client, headers: Option<&HeaderMap>, req: StreamingRequest, - active_mcp: &Arc, + active_mcp: &Arc, server_keys: Vec, ) -> Response { // Transform MCP tools to function tools in payload diff --git a/sgl-model-gateway/src/routers/openai/router.rs b/sgl-model-gateway/src/routers/openai/router.rs index 336679318..6c920bc14 100644 --- a/sgl-model-gateway/src/routers/openai/router.rs +++ b/sgl-model-gateway/src/routers/openai/router.rs @@ -12,6 +12,7 @@ use axum::{ response::{IntoResponse, Response}, Json, }; +use data_connector::{ConversationId, ListParams, ResponseId, SortOrder}; use futures_util::{future::join_all, StreamExt}; use serde_json::{json, to_value, Value}; use tokio::sync::mpsc; @@ -33,7 +34,6 @@ use crate::{ is_retryable_status, model_type::Endpoint, ModelCard, ProviderType, RetryExecutor, RuntimeType, Worker, WorkerRegistry, }, - data_connector::{ConversationId, ListParams, ResponseId, SortOrder}, observability::metrics::{bool_to_static_str, metrics_labels, Metrics}, protocols::{ chat::ChatCompletionRequest, diff --git a/sgl-model-gateway/src/routers/persistence_utils.rs b/sgl-model-gateway/src/routers/persistence_utils.rs index b3775728b..105afe965 100644 --- a/sgl-model-gateway/src/routers/persistence_utils.rs +++ b/sgl-model-gateway/src/routers/persistence_utils.rs @@ -3,17 +3,15 @@ use std::sync::Arc; use chrono::Utc; +use data_connector::{ + ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, + ConversationStorage, NewConversationItem, ResponseId, ResponseStorage, StoredResponse, +}; use serde_json::{json, Value}; use tracing::{debug, info, warn}; -use crate::{ - data_connector::{ - ConversationId, ConversationItem, ConversationItemId, ConversationItemStorage, - ConversationStorage, NewConversationItem, ResponseId, ResponseStorage, StoredResponse, - }, - protocols::responses::{ - generate_id, ResponseInput, ResponseInputOutputItem, ResponsesRequest, StringOrContentParts, - }, +use crate::protocols::responses::{ + generate_id, ResponseInput, ResponseInputOutputItem, ResponsesRequest, StringOrContentParts, }; // ============================================================================ diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 6a4d5c531..fd762c66e 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -21,6 +21,7 @@ use smg_mesh::{ }; use tokio::{signal, spawn}; use tracing::{debug, error, info, warn, Level}; +use wfaas::LoggingSubscriber; use crate::{ app_context::AppContext, @@ -65,7 +66,6 @@ use crate::{ service_discovery::{start_service_discovery, ServiceDiscoveryConfig}, tokenizer::TokenizerRegistry, wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module}, - workflow::LoggingSubscriber, }; #[derive(Clone)] pub struct AppState { @@ -738,62 +738,61 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box Arc { } // Initialize MCP manager with empty config - use smg::mcp::{McpConfig, McpManager}; + use smg_mcp::{McpConfig, McpManager}; let empty_config = McpConfig { servers: vec![], pool: Default::default(), @@ -510,7 +510,7 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc Arc { - use smg::mcp::{McpConfig, McpManager}; + use smg_mcp::{McpConfig, McpManager}; let client = reqwest::Client::new(); diff --git a/sgl-model-gateway/tests/common/test_app.rs b/sgl-model-gateway/tests/common/test_app.rs index a420e56e7..a9862afaa 100644 --- a/sgl-model-gateway/tests/common/test_app.rs +++ b/sgl-model-gateway/tests/common/test_app.rs @@ -1,6 +1,9 @@ use std::sync::{Arc, OnceLock}; use axum::Router; +use data_connector::{ + MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, +}; use reqwest::Client; use smg::{ app_context::AppContext, @@ -8,16 +11,13 @@ use smg::{ core::{ BasicWorkerBuilder, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry, WorkerType, }, - data_connector::{ - MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, - }, - mcp::{McpConfig, McpManager}, middleware::{AuthConfig, TokenBucket}, policies::PolicyRegistry, routers::RouterTrait, server::{build_app, AppState}, tokenizer::registry::TokenizerRegistry, }; +use smg_mcp::{McpConfig, McpManager}; /// Create a test Axum application using the actual server's build_app function #[allow(dead_code)] diff --git a/sgl-model-gateway/tests/mcp_test.rs b/sgl-model-gateway/tests/mcp_test.rs index 661639ef6..29dcf3e5b 100644 --- a/sgl-model-gateway/tests/mcp_test.rs +++ b/sgl-model-gateway/tests/mcp_test.rs @@ -13,7 +13,7 @@ use std::collections::HashMap; use common::mock_mcp_server::MockMCPServer; use serde_json::json; -use smg::mcp::{error::McpError, McpConfig, McpManager, McpServerConfig, McpTransport}; +use smg_mcp::{error::McpError, McpConfig, McpManager, McpServerConfig, McpTransport}; /// Create a new mock server for testing (each test gets its own) async fn create_mock_server() -> MockMCPServer { diff --git a/sgl-model-gateway/tests/routing/test_openai_routing.rs b/sgl-model-gateway/tests/routing/test_openai_routing.rs index 5ad6325bc..5614f8c75 100644 --- a/sgl-model-gateway/tests/routing/test_openai_routing.rs +++ b/sgl-model-gateway/tests/routing/test_openai_routing.rs @@ -16,10 +16,10 @@ use axum::{ routing::post, Json, Router, }; +use data_connector::{ResponseId, StoredResponse}; use serde_json::json; use smg::{ config::{ConfigError, HistoryBackend, OracleConfig, RouterConfig, RoutingMode}, - data_connector::{ResponseId, StoredResponse}, protocols::{ chat::{ChatCompletionRequest, ChatMessage, MessageContent}, common::StringOrArray, diff --git a/sgl-model-gateway/tests/routing/test_pd_routing.rs b/sgl-model-gateway/tests/routing/test_pd_routing.rs index b5a01f675..185347472 100644 --- a/sgl-model-gateway/tests/routing/test_pd_routing.rs +++ b/sgl-model-gateway/tests/routing/test_pd_routing.rs @@ -215,12 +215,11 @@ mod pd_routing_unit_tests { let app_context = { use std::sync::{Arc, OnceLock}; + use data_connector::{ + MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, + }; use smg::{ core::{LoadMonitor, WorkerRegistry}, - data_connector::{ - MemoryConversationItemStorage, MemoryConversationStorage, - MemoryResponseStorage, - }, middleware::TokenBucket, policies::PolicyRegistry, }; diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs index 008c53791..38b43b134 100644 --- a/sgl-model-gateway/tests/wasm_test.rs +++ b/sgl-model-gateway/tests/wasm_test.rs @@ -15,13 +15,13 @@ use axum::{ extract::Request, http::{header::CONTENT_TYPE, StatusCode}, }; +use data_connector::{ + MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, +}; use smg::{ app_context::AppContext, config::RouterConfig, core::{LoadMonitor, WorkerRegistry}, - data_connector::{ - MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, - }, policies::PolicyRegistry, routers::RouterFactory, server::{build_app, AppState}, @@ -113,7 +113,7 @@ async fn create_test_context_with_wasm() -> Arc { .expect("WorkflowEngines should only be initialized once"); // Initialize MCP manager with empty config - use smg::mcp::{McpConfig, McpManager}; + use smg_mcp::{McpConfig, McpManager}; let empty_config = McpConfig { servers: vec![], pool: Default::default(), @@ -668,10 +668,8 @@ async fn test_wasm_module_execution() { .expect("Workflow engines should be initialized"); // Create workflow context for registration - use smg::{ - core::steps::{WasmModuleConfigRequest, WasmRegistrationWorkflowData}, - workflow::WorkflowId, - }; + use smg::core::steps::{WasmModuleConfigRequest, WasmRegistrationWorkflowData}; + use wfaas::WorkflowId; let descriptor = WasmModuleDescriptor { name: "test_execution_module".to_string(), @@ -716,7 +714,7 @@ async fn test_wasm_module_execution() { .expect("Failed to get workflow status"); match state.status { - smg::workflow::WorkflowStatus::Completed => { + wfaas::WorkflowStatus::Completed => { // Extract module UUID from typed workflow data break state .context @@ -724,7 +722,7 @@ async fn test_wasm_module_execution() { .module_uuid .expect("Module UUID should be in context"); } - smg::workflow::WorkflowStatus::Failed => { + wfaas::WorkflowStatus::Failed => { panic!("Workflow failed: {:?}", state); } _ => {