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

refactor: remove crate re-export aliases from lib.rs (#18737)

Browse Source
This commit is contained in:
Simo Lin
2026-02-12 09:17:24 -08:00
committed by GitHub
parent e730c728d3
commit 0236522214
74 changed files with 240 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/benches/consistent_hash_bench.rs
+1 -1
View File
@@ -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();
sgl-model-gateway/bindings/golang/src/client.rs
+1 -1
View File
@@ -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};
sgl-model-gateway/bindings/golang/src/grpc_converter.rs
+2 -2
View File
@@ -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<Option<smg::protocols::chat::ChatCompletionStreamResponse>, 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 {
sgl-model-gateway/bindings/golang/src/postprocessor.rs
+1 -1
View File
@@ -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;
sgl-model-gateway/bindings/golang/src/stream.rs
+1 -1
View File
@@ -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};
sgl-model-gateway/src/app_context.rs
+6 -6
View File
@@ -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,
sgl-model-gateway/src/config/builder.rs
+3 -1
View File
@@ -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
sgl-model-gateway/src/config/types.rs
+3 -3
View File
@@ -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<Vec<u8>>,
/// Loaded from mcp_config_path during config creation
#[serde(skip)]
pub mcp_config: Option<crate::mcp::McpConfig>,
pub mcp_config: Option<smg_mcp::McpConfig>,
/// Enable WASM support
#[serde(default)]
pub enable_wasm: bool,
sgl-model-gateway/src/core/job_queue.rs
+2 -2
View File
@@ -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
sgl-model-gateway/src/core/steps/mcp_registration.rs
+6 -9
View File
@@ -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)]
sgl-model-gateway/src/core/steps/tokenizer_registration.rs
+4 -4
View File
@@ -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
sgl-model-gateway/src/core/steps/wasm_module_registration.rs
+4 -4
View File
@@ -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
sgl-model-gateway/src/core/steps/wasm_module_removal.rs
+5 -7
View File
@@ -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)]
sgl-model-gateway/src/core/steps/worker/external/create_workers.rs
Vendored
+6 -8
View File
@@ -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://).
sgl-model-gateway/src/core/steps/worker/external/discover_models.rs
Vendored
+5 -7
View File
@@ -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
sgl-model-gateway/src/core/steps/worker/external/mod.rs
Vendored
+2 -3
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/local/create_worker.rs
+1 -1
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/local/detect_connection.rs
+1 -1
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/local/discover_dp.rs
+2 -4
View File
@@ -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)]
sgl-model-gateway/src/core/steps/worker/local/discover_metadata.rs
+1 -1
View File
@@ -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
sgl-model-gateway/src/core/steps/worker/local/find_worker_to_update.rs
+2 -4
View File
@@ -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.
///
sgl-model-gateway/src/core/steps/worker/local/find_workers_to_remove.rs
+2 -4
View File
@@ -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)]
sgl-model-gateway/src/core/steps/worker/local/mod.rs
+1 -1
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/local/remove_from_policy_registry.rs
+2 -4
View File
@@ -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.
///
sgl-model-gateway/src/core/steps/worker/local/remove_from_worker_registry.rs
+2 -3
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs
+1 -1
View File
@@ -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
sgl-model-gateway/src/core/steps/worker/local/update_policies_for_worker.rs
+2 -4
View File
@@ -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.
///
sgl-model-gateway/src/core/steps/worker/local/update_remaining_policies.rs
+2 -4
View File
@@ -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.
///
sgl-model-gateway/src/core/steps/worker/local/update_worker_properties.rs
+3 -5
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/shared/activate.rs
+4 -6
View File
@@ -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.
sgl-model-gateway/src/core/steps/worker/shared/register.rs
+4 -7
View File
@@ -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
sgl-model-gateway/src/core/steps/worker/shared/update_policies.rs
+4 -6
View File
@@ -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
sgl-model-gateway/src/core/steps/workflow_data.rs
+1 -1
View File
@@ -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},
};
// ============================================================================
sgl-model-gateway/src/core/steps/workflow_engines.rs
+3 -4
View File
@@ -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 =
sgl-model-gateway/src/core/worker_registry.rs
+1 -1
View File
@@ -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,
};
sgl-model-gateway/src/lib.rs
-5
View File
@@ -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;
sgl-model-gateway/src/main.rs
+1 -1
View File
@@ -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<u16>)> {
let args: Vec<String> = std::env::args().collect();
let mut prefill_entries = Vec::new();
sgl-model-gateway/src/observability/metrics.rs
+1 -1
View File
@@ -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) {
sgl-model-gateway/src/policies/cache_aware.rs
+9 -11
View File
@@ -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},
sgl-model-gateway/src/policies/mod.rs
+2 -4
View File
@@ -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;
sgl-model-gateway/src/policies/registry.rs
+3 -2
View File
@@ -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) {
sgl-model-gateway/src/routers/conversations/handlers.rs
+5 -8
View File
@@ -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
sgl-model-gateway/src/routers/grpc/client.rs
+4 -3
View File
@@ -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<crate::grpc_client::sglang_proto::GetModelInfoResponse>),
Vllm(crate::grpc_client::vllm_proto::GetModelInfoResponse),
Sglang(Box<smg_grpc_client::sglang_proto::GetModelInfoResponse>),
Vllm(smg_grpc_client::vllm_proto::GetModelInfoResponse),
}
impl ModelInfo {
sgl-model-gateway/src/routers/grpc/common/responses/context.rs
+4 -5
View File
@@ -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
///
sgl-model-gateway/src/routers/grpc/common/responses/handlers.rs
+2 -1
View File
@@ -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}
///
sgl-model-gateway/src/routers/grpc/common/responses/streaming.rs
+1 -1
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/common/responses/utils.rs
+2 -2
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/common/stages/helpers.rs
+2 -4
View File
@@ -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
///
sgl-model-gateway/src/routers/grpc/harmony/processor.rs
+3 -1
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/harmony/responses/common.rs
+2 -2
View File
@@ -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::{
sgl-model-gateway/src/routers/grpc/harmony/responses/execution.rs
+1 -1
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/harmony/streaming.rs
+3 -1
View File
@@ -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::{
sgl-model-gateway/src/routers/grpc/mod.rs
+3 -1
View File
@@ -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;
sgl-model-gateway/src/routers/grpc/proto_wrapper.rs
+1 -2
View File
@@ -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,
sgl-model-gateway/src/routers/grpc/regular/processor.rs
+1 -1
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/regular/responses/common.rs
+2 -2
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/regular/responses/streaming.rs
+1 -1
View File
@@ -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::{
sgl-model-gateway/src/routers/grpc/regular/streaming.rs
+3 -1
View File
@@ -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},
sgl-model-gateway/src/routers/grpc/utils.rs
+1 -1
View File
@@ -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},
sgl-model-gateway/src/routers/mcp_utils.rs
+2 -4
View File
@@ -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
sgl-model-gateway/src/routers/openai/context.rs
+2 -2
View File
@@ -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},
};
sgl-model-gateway/src/routers/openai/responses/mcp.rs
+1 -1
View File
@@ -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},
sgl-model-gateway/src/routers/openai/responses/streaming.rs
+2 -2
View File
@@ -424,7 +424,7 @@ pub(super) fn send_final_response_event(
tx: &mpsc::UnboundedSender<Result<Bytes, io::Error>>,
sequence_number: &mut u64,
state: &ToolLoopState,
active_mcp: Option<&Arc<crate::mcp::McpManager>>,
active_mcp: Option<&Arc<smg_mcp::McpManager>>,
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<crate::mcp::McpManager>,
active_mcp: &Arc<smg_mcp::McpManager>,
server_keys: Vec<String>,
) -> Response {
// Transform MCP tools to function tools in payload
sgl-model-gateway/src/routers/openai/router.rs
+1 -1
View File
@@ -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,
sgl-model-gateway/src/routers/persistence_utils.rs
+6 -8
View File
@@ -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,
};
// ============================================================================
sgl-model-gateway/src/server.rs
+50 -51
View File
@@ -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<dyn std::error::Err
metrics::start_prometheus(prometheus_config.clone());
}
let (mesh_handler, mesh_sync_manager) =
if let Some(mesh_server_config) = &config.mesh_server_config {
// Create HA sync manager with stores first
use crate::mesh::{
partition::PartitionDetector, stores::StateStores, sync::MeshSyncManager,
};
let stores = Arc::new(StateStores::with_self_name(
mesh_server_config.self_name.clone(),
));
let sync_manager = Arc::new(MeshSyncManager::new(
stores.clone(),
mesh_server_config.self_name.clone(),
));
let (mesh_handler, mesh_sync_manager) = if let Some(mesh_server_config) =
&config.mesh_server_config
{
// Create HA sync manager with stores first
use smg_mesh::{partition::PartitionDetector, stores::StateStores, sync::MeshSyncManager};
let stores = Arc::new(StateStores::with_self_name(
mesh_server_config.self_name.clone(),
));
let sync_manager = Arc::new(MeshSyncManager::new(
stores.clone(),
mesh_server_config.self_name.clone(),
));
// Create partition detector
let partition_detector = Arc::new(PartitionDetector::default());
// Create partition detector
let partition_detector = Arc::new(PartitionDetector::default());
// Initialize rate-limit hash ring with current membership
sync_manager.update_rate_limit_membership();
// Initialize rate-limit hash ring with current membership
sync_manager.update_rate_limit_membership();
// Start rate limit window reset task
let window_manager = RateLimitWindow::new(sync_manager.clone(), 1); // Reset every 1 second
spawn(async move {
window_manager.start_reset_task().await;
});
// Start rate limit window reset task
let window_manager = RateLimitWindow::new(sync_manager.clone(), 1); // Reset every 1 second
spawn(async move {
window_manager.start_reset_task().await;
});
// Create mesh server builder and build with stores
use crate::mesh::service::MeshServerBuilder;
let builder = MeshServerBuilder::new(
mesh_server_config.self_name.clone(),
mesh_server_config.self_addr,
mesh_server_config.init_peer,
);
let (mesh_server, handler) = builder.build_with_stores(Some(stores.clone()));
// Create mesh server builder and build with stores
use smg_mesh::service::MeshServerBuilder;
let builder = MeshServerBuilder::new(
mesh_server_config.self_name.clone(),
mesh_server_config.self_addr,
mesh_server_config.init_peer,
);
let (mesh_server, handler) = builder.build_with_stores(Some(stores.clone()));
// Spawn the mesh server with stores and partition detector
let stores_for_server = stores.clone();
let sync_manager_for_server = sync_manager.clone();
let partition_detector_for_server = partition_detector.clone();
spawn(async move {
if let Err(e) = mesh_server
.start_serve_with_stores(
Some(stores_for_server),
Some(sync_manager_for_server),
Some(partition_detector_for_server),
)
.await
{
tracing::error!("Mesh server failed: {}", e);
}
});
// Spawn the mesh server with stores and partition detector
let stores_for_server = stores.clone();
let sync_manager_for_server = sync_manager.clone();
let partition_detector_for_server = partition_detector.clone();
spawn(async move {
if let Err(e) = mesh_server
.start_serve_with_stores(
Some(stores_for_server),
Some(sync_manager_for_server),
Some(partition_detector_for_server),
)
.await
{
tracing::error!("Mesh server failed: {}", e);
}
});
(Some(Arc::new(handler)), Some(sync_manager))
} else {
(None, None)
};
(Some(Arc::new(handler)), Some(sync_manager))
} else {
(None, None)
};
info!(
"Starting router on {}:{} | mode: {:?} | policy: {:?} | max_payload: {}MB",
sgl-model-gateway/src/service_discovery.rs
+4 -4
View File
@@ -15,16 +15,16 @@ use kube::{
Client,
};
use rustls;
use smg_mesh::service::{
gossip::{NodeState, NodeStatus},
ClusterState,
};
use tokio::{task, time};
use tracing::{debug, error, info, warn};
use crate::{
app_context::AppContext,
core::Job,
mesh::service::{
gossip::{NodeState, NodeStatus},
ClusterState,
},
observability::metrics::{metrics_labels, Metrics},
protocols::worker_spec::WorkerConfigRequest,
};
sgl-model-gateway/tests/api/api_endpoints_test.rs
+1 -1
View File
@@ -880,7 +880,7 @@ mod responses_endpoint_tests {
let app = ctx.create_app().await;
// Directly store a response in the storage to test the retrieval endpoint
use smg::data_connector::{ResponseId, StoredResponse};
use data_connector::{ResponseId, StoredResponse};
let mut stored_response = StoredResponse::new(None);
stored_response.id = ResponseId::from("resp_test_input_items");
stored_response.input = json!([
sgl-model-gateway/tests/common/mod.rs
+6 -6
View File
@@ -18,6 +18,9 @@ use std::{
sync::{Arc, Mutex, OnceLock},
};
use data_connector::{
MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
};
use mock_worker::{MockWorker, MockWorkerConfig};
use serde_json::json;
use smg::{
@@ -27,9 +30,6 @@ use smg::{
BasicWorkerBuilder, Job, LoadMonitor, ModelCard, RuntimeType, Worker, WorkerRegistry,
WorkerType,
},
data_connector::{
MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage,
},
middleware::TokenBucket,
policies::PolicyRegistry,
protocols::common::{Function, Tool},
@@ -386,7 +386,7 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
}
// 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<AppCo
}
// 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(),
@@ -535,7 +535,7 @@ pub async fn create_test_context_with_mcp_config(
config: RouterConfig,
mcp_config_path: &str,
) -> Arc<AppContext> {
use smg::mcp::{McpConfig, McpManager};
use smg_mcp::{McpConfig, McpManager};
let client = reqwest::Client::new();
sgl-model-gateway/tests/common/test_app.rs
+4 -4
View File
@@ -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)]
sgl-model-gateway/tests/mcp_test.rs
+1 -1
View File
@@ -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 {
sgl-model-gateway/tests/routing/test_openai_routing.rs
+1 -1
View File
@@ -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,
sgl-model-gateway/tests/routing/test_pd_routing.rs
+3 -4
View File
@@ -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,
};
sgl-model-gateway/tests/wasm_test.rs
+8 -10
View File
@@ -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<AppContext> {
.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);
}
_ => {