chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[model-gateway] Replace tokenizer with tokenizer registry for dynamic tokenizer loading in gRPC router (#12968)

Browse Source
This commit is contained in:
Arthur Cheng
2025-12-23 07:58:47 -08:00
committed by GitHub
parent 53f974b973
commit dd620987d1
27 changed files with 793 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

113
sgl-model-gateway/src/app_context.rs
View File

@@ -4,7 +4,7 @@ use std::{
};
use reqwest::Client;
use tracing::debug;
use tracing::{debug, info};
use crate::{
config::RouterConfig,
@@ -21,6 +21,7 @@ use crate::{
cache::{CacheConfig, CachedTokenizer},
factory as tokenizer_factory,
traits::Tokenizer,
TokenizerRegistry,
},
tool_parser::ParserFactory as ToolParserFactory,
wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager},
@@ -44,7 +45,7 @@ pub struct AppContext {
pub client: Client,
pub router_config: RouterConfig,
pub rate_limiter: Option<Arc<TokenBucket>>,
pub tokenizer: Option<Arc<dyn Tokenizer>>,
pub tokenizer_registry: Arc<TokenizerRegistry>,
pub reasoning_parser_factory: Option<ReasoningParserFactory>,
pub tool_parser_factory: Option<ToolParserFactory>,
pub worker_registry: Arc<WorkerRegistry>,
@@ -67,7 +68,7 @@ pub struct AppContextBuilder {
client: Option<Client>,
router_config: Option<RouterConfig>,
rate_limiter: Option<Arc<TokenBucket>>,
tokenizer: Option<Arc<dyn Tokenizer>>,
tokenizer_registry: Option<Arc<TokenizerRegistry>>,
reasoning_parser_factory: Option<ReasoningParserFactory>,
tool_parser_factory: Option<ToolParserFactory>,
worker_registry: Option<Arc<WorkerRegistry>>,
@@ -107,7 +108,7 @@ impl AppContextBuilder {
client: None,
router_config: None,
rate_limiter: None,
tokenizer: None,
tokenizer_registry: None,
reasoning_parser_factory: None,
tool_parser_factory: None,
worker_registry: None,
@@ -139,8 +140,8 @@ impl AppContextBuilder {
self
}
pub fn tokenizer(mut self, tokenizer: Option<Arc<dyn Tokenizer>>) -> Self {
self.tokenizer = tokenizer;
pub fn tokenizer_registry(mut self, tokenizer_registry: Arc<TokenizerRegistry>) -> Self {
self.tokenizer_registry = Some(tokenizer_registry);
self
}
@@ -243,7 +244,9 @@ impl AppContextBuilder {
client: self.client.ok_or(AppContextBuildError("client"))?,
router_config,
rate_limiter: self.rate_limiter,
tokenizer: self.tokenizer,
tokenizer_registry: self
.tokenizer_registry
.ok_or(AppContextBuildError("tokenizer_registry"))?,
reasoning_parser_factory: self.reasoning_parser_factory,
tool_parser_factory: self.tool_parser_factory,
worker_registry,
@@ -284,7 +287,7 @@ impl AppContextBuilder {
Ok(Self::new()
.with_client(&router_config, request_timeout_secs)?
.maybe_rate_limiter(&router_config)
.maybe_tokenizer(&router_config)?
.with_tokenizer_registry(&router_config)?
.maybe_reasoning_parser_factory(&router_config)
.maybe_tool_parser_factory(&router_config)
.with_worker_registry()
@@ -380,49 +383,54 @@ impl AppContextBuilder {
self
}
/// Create tokenizer for gRPC mode
fn maybe_tokenizer(mut self, config: &RouterConfig) -> Result<Self, String> {
if matches!(config.connection_mode, ConnectionMode::Grpc { .. }) {
let tokenizer_path = config
.tokenizer_path
.clone()
.or_else(|| config.model_path.clone())
.ok_or_else(|| {
"gRPC mode requires either --tokenizer-path or --model-path to be specified"
.to_string()
})?;
/// Load tokenizer if tokenizer_path is provided
///
/// This is a pure function that loads the tokenizer from the provided path
/// and applies caching configuration. Returns None if no tokenizer path is configured.
fn maybe_tokenizer(config: &RouterConfig) -> Result<Option<Arc<dyn Tokenizer>>, String> {
// Check if tokenizer path is provided
let tokenizer_path = match config
.tokenizer_path
.clone()
.or_else(|| config.model_path.clone())
{
Some(path) => path,
None => {
info!("Tokenizer path is not provided, will load from worker on the fly");
return Ok(None);
}
};
let base_tokenizer = tokenizer_factory::create_tokenizer_with_chat_template_blocking(
&tokenizer_path,
config.chat_template.as_deref(),
// Load base tokenizer
let base_tokenizer = tokenizer_factory::create_tokenizer_with_chat_template_blocking(
&tokenizer_path,
config.chat_template.as_deref(),
)
.map_err(|e| {
format!(
"Failed to create tokenizer from '{}': {}. \
Ensure the path is valid and points to a tokenizer file (tokenizer.json) \
or a HuggingFace model ID. For directories, ensure they contain tokenizer files.",
tokenizer_path, e
)
.map_err(|e| {
format!(
"Failed to create tokenizer from '{}': {}. \
Ensure the path is valid and points to a tokenizer file (tokenizer.json) \
or a HuggingFace model ID. For directories, ensure they contain tokenizer files.",
tokenizer_path, e
)
})?;
})?;
// Conditionally wrap with caching layer if at least one cache is enabled
self.tokenizer = if config.tokenizer_cache.enable_l0 || config.tokenizer_cache.enable_l1
{
// Conditionally wrap with caching layer if at least one cache is enabled
let tokenizer: Arc<dyn Tokenizer> =
if config.tokenizer_cache.enable_l0 || config.tokenizer_cache.enable_l1 {
let cache_config = CacheConfig {
enable_l0: config.tokenizer_cache.enable_l0,
l0_max_entries: config.tokenizer_cache.l0_max_entries,
enable_l1: config.tokenizer_cache.enable_l1,
l1_max_memory: config.tokenizer_cache.l1_max_memory,
};
Some(Arc::new(CachedTokenizer::new(base_tokenizer, cache_config))
as Arc<dyn Tokenizer>)
Arc::new(CachedTokenizer::new(base_tokenizer, cache_config)) as Arc<dyn Tokenizer>
} else {
// Use base tokenizer directly without caching
Some(base_tokenizer)
base_tokenizer
};
}
Ok(self)
Ok(Some(tokenizer))
}
/// Create reasoning parser factory for gRPC mode
@@ -441,6 +449,35 @@ impl AppContextBuilder {
self
}
/// Create tokenizer registry and optionally load tokenizer
/// If a tokenizer is successfully loaded, it is registered with a key derived from
/// tokenizer_path or model_path (falling back to "unknown" if neither exists).
fn with_tokenizer_registry(mut self, config: &RouterConfig) -> Result<Self, String> {
// Create empty tokenizer registry
let registry = Arc::new(TokenizerRegistry::new());
// Try to load router-level tokenizer if path is provided
if let Some(tokenizer) = Self::maybe_tokenizer(config)? {
// Determine registration key: prefer tokenizer_path, then model_path, finally "unknown"
let tokenizer_key = config
.tokenizer_path
.as_ref()
.or(config.model_path.as_ref())
.map(|s| s.as_str())
.unwrap_or("unknown");
registry.register(tokenizer_key, tokenizer.clone());
info!(
"Tokenizer loaded and registered with key '{}' (vocab_size: {})",
tokenizer_key,
tokenizer.vocab_size()
);
}
self.tokenizer_registry = Some(registry);
Ok(self)
}
/// Create worker registry
fn with_worker_registry(mut self) -> Self {
self.worker_registry = Some(Arc::new(WorkerRegistry::new()));

32
sgl-model-gateway/src/config/validation.rs
View File

@@ -1,5 +1,4 @@
use super::*;
use crate::core::ConnectionMode;
/// Configuration validator
pub struct ConfigValidator;
@@ -517,15 +516,6 @@ impl ConfigValidator {
return Ok(());
}
if matches!(config.connection_mode, ConnectionMode::Grpc { .. })
&& config.tokenizer_path.is_none()
&& config.model_path.is_none()
{
return Err(ConfigError::ValidationFailed {
reason: "gRPC connection mode requires either --tokenizer-path or --model-path to be specified".to_string(),
});
}
Self::validate_mtls(config)?;
let has_service_discovery = config.discovery.as_ref().is_some_and(|d| d.enabled);
@@ -624,6 +614,7 @@ impl ConfigValidator {
#[cfg(test)]
mod tests {
use super::*;
use crate::core::ConnectionMode;
#[test]
fn test_validate_regular_mode() {
@@ -952,27 +943,6 @@ mod tests {
assert!(ConfigValidator::validate(&config).is_ok());
}
#[test]
fn test_validate_grpc_requires_tokenizer() {
let mut config = RouterConfig::new(
RoutingMode::Regular {
worker_urls: vec!["grpc://worker:50051".to_string()],
},
PolicyConfig::Random,
);
// Set connection mode to gRPC without tokenizer config
config.connection_mode = ConnectionMode::Grpc { port: None };
config.tokenizer_path = None;
config.model_path = None;
let result = ConfigValidator::validate(&config);
assert!(result.is_err());
if let Err(e) = result {
assert!(e.to_string().contains("gRPC connection mode requires"));
}
}
#[test]
fn test_validate_grpc_with_model_path() {
let mut config = RouterConfig::new(

2
sgl-model-gateway/src/core/model_card.rs
View File

@@ -335,7 +335,7 @@ impl ModelCard {
impl Default for ModelCard {
fn default() -> Self {
Self::new("default")
Self::new("unknown")
}
}

16
sgl-model-gateway/src/core/steps/worker/local/mod.rs
View File

@@ -4,6 +4,7 @@ mod discover_dp;
mod discover_metadata;
mod find_worker_to_update;
mod find_workers_to_remove;
mod register_tokenizer;
mod remove_from_policy_registry;
mod remove_from_worker_registry;
mod update_policies_for_worker;
@@ -18,6 +19,7 @@ pub use discover_dp::{get_dp_info, DiscoverDPInfoStep, DpInfo};
pub use discover_metadata::DiscoverMetadataStep;
pub use find_worker_to_update::FindWorkerToUpdateStep;
pub use find_workers_to_remove::{FindWorkersToRemoveStep, WorkerRemovalRequest};
pub use register_tokenizer::RegisterTokenizerStep;
pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep;
pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep;
pub use update_policies_for_worker::UpdatePoliciesForWorkerStep;
@@ -138,6 +140,20 @@ pub fn create_local_worker_workflow(router_config: &RouterConfig) -> WorkflowDef
.with_failure_action(FailureAction::FailWorkflow)
.depends_on(&["create_worker"]),
)
.add_step(
StepDefinition::new(
"register_tokenizer",
"Register Tokenizer",
Arc::new(RegisterTokenizerStep),
)
.with_retry(RetryPolicy {
max_attempts: 3,
backoff: BackoffStrategy::Fixed(Duration::from_secs(1)),
})
.with_timeout(Duration::from_secs(10))
.with_failure_action(FailureAction::ContinueNextStep)
.depends_on(&["register_workers"]),
)
// Step 5a: Update policies (parallel with activation)
.add_step(
StepDefinition::new(

78
sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs Normal file
View File

@@ -0,0 +1,78 @@
//! Connection mode detection step.
use std::{collections::HashMap, sync::Arc};
use async_trait::async_trait;
use tracing::{debug, warn};
use crate::{
app_context::AppContext,
core::Worker,
tokenizer::factory,
workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult},
};
/// Step 6: Register tokenizer for the worker's model (optional, non-blocking)
pub struct RegisterTokenizerStep;
#[async_trait]
impl StepExecutor for RegisterTokenizerStep {
async fn execute(&self, context: &mut WorkflowContext) -> WorkflowResult<StepResult> {
let labels: Arc<HashMap<String, String>> = context
.get("labels")
.ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?;
let app_context: Arc<AppContext> = context
.get("app_context")
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
let workers: Arc<Vec<Arc<dyn Worker>>> = context
.get("workers")
.ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?;
for worker in workers.iter() {
let model_id = worker.model_id().to_string();
// Get tokenizer path (prefer tokenizer_path, fallback to model_path)
let Some(tokenizer_path) = labels
.get("tokenizer_path")
.or_else(|| labels.get("model_path"))
else {
warn!(
"No tokenizer_path or model_path found for model {}",
model_id
);
return Ok(StepResult::Success);
};
debug!(
"Registering tokenizer for model {} from {}",
model_id, tokenizer_path
);
// Load tokenizer with thread safe lock
if let Err(e) = app_context
.tokenizer_registry
.load(&model_id, || async move {
factory::create_tokenizer_async(&tokenizer_path.to_string())
.await
.map_err(|e| e.to_string())
})
.await
{
warn!(
"Failed to load tokenizer for model {} from {}: {}",
model_id, tokenizer_path, e
);
} else {
debug!(
"Successfully registered tokenizer for model {} from {}",
model_id, tokenizer_path
);
}
}
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
true // Tokenizer loading failures are retryable (network/IO issues)
}
}

18
sgl-model-gateway/src/policies/bucket.rs
View File

@@ -79,10 +79,10 @@ impl BucketPolicy {
// Group workers by model
let mut model_workers: HashMap<String, Vec<&Arc<dyn Worker>>> = HashMap::new();
for worker in prefill_workers {
// Use "default" for unknown/empty model_ids for backward compatibility
// Use "unknown" for empty model_ids
let model_id = worker.model_id();
let model_key = if model_id.is_empty() || model_id == "unknown" {
"default"
let model_key = if model_id.is_empty() {
"unknown"
} else {
model_id
};
@@ -119,8 +119,8 @@ impl BucketPolicy {
pub fn add_prefill_url(&self, worker: &dyn Worker) {
let model_id = worker.model_id();
let model_key = if model_id.is_empty() || model_id == "unknown" {
"default"
let model_key = if model_id.is_empty() {
"unknown"
} else {
model_id
};
@@ -167,8 +167,8 @@ impl BucketPolicy {
pub fn remove_prefill_url(&self, worker: &dyn Worker) {
let model_id = worker.model_id();
let model_key = if model_id.is_empty() || model_id == "unknown" {
"default"
let model_key = if model_id.is_empty() {
"unknown"
} else {
model_id
};
@@ -236,8 +236,8 @@ impl LoadBalancingPolicy for BucketPolicy {
// Determine the model for this set of workers (router pre-filters by model)
// All workers should be from the same model
let first_model = workers[healthy_indices[0]].model_id();
let model_key = if first_model.is_empty() || first_model == "unknown" {
"default"
let model_key = if first_model.is_empty() {
"unknown"
} else {
first_model
};

18
sgl-model-gateway/src/policies/cache_aware.rs
View File

@@ -160,10 +160,10 @@ impl CacheAwarePolicy {
let mut model_workers: std::collections::HashMap<String, Vec<&Arc<dyn Worker>>> =
std::collections::HashMap::new();
for worker in workers {
// Use "default" for unknown/empty model_ids for backward compatibility
// Use "unknown" for empty model_ids
let model_id = worker.model_id();
let tree_key = if model_id.is_empty() || model_id == "unknown" {
"default"
let tree_key = if model_id.is_empty() {
"unknown"
} else {
model_id
};
@@ -190,8 +190,8 @@ impl CacheAwarePolicy {
// For backward compatibility: if model_id is "unknown" or empty,
// use a default tree. This preserves existing behavior for single-model routers.
let model_id = worker.model_id();
let tree_key = if model_id.is_empty() || model_id == "unknown" {
"default"
let tree_key = if model_id.is_empty() {
"unknown"
} else {
model_id
};
@@ -215,8 +215,8 @@ impl CacheAwarePolicy {
pub fn remove_worker(&self, worker: &dyn Worker) {
// Use same logic as add_worker for consistency
let model_id = worker.model_id();
let tree_key = if model_id.is_empty() || model_id == "unknown" {
"default"
let tree_key = if model_id.is_empty() {
"unknown"
} else {
model_id
};
@@ -314,8 +314,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
// Determine the model for this set of workers (router pre-filters by model)
// All workers should be from the same model
let first_model = workers[healthy_indices[0]].model_id();
let model_id = if first_model.is_empty() || first_model == "unknown" {
"default"
let model_id = if first_model.is_empty() {
"unknown"
} else {
first_model
};

4
sgl-model-gateway/src/routers/grpc/common/stages/dispatch_metadata.rs
View File

@@ -31,11 +31,11 @@ impl PipelineStage for DispatchMetadataStage {
RequestType::Chat(req) => req.model.clone(),
RequestType::Generate(_req) => {
// Generate requests don't have a model field
// Use model_id from input or default
// Use model_id from input or unknown
ctx.input
.model_id
.clone()
.unwrap_or_else(|| "default".to_string())
.unwrap_or_else(|| "unknown".to_string())
}
RequestType::Responses(req) => req.model.clone(),
};

4
sgl-model-gateway/src/routers/grpc/context.rs
View File

@@ -21,7 +21,7 @@ use crate::{
responses::ResponsesRequest,
},
reasoning_parser::ParserFactory as ReasoningParserFactory,
tokenizer::{stop::StopSequenceDecoder, traits::Tokenizer},
tokenizer::{stop::StopSequenceDecoder, TokenizerRegistry},
tool_parser::ParserFactory as ToolParserFactory,
};
@@ -53,7 +53,7 @@ pub enum RequestType {
/// Shared components (injected once at creation)
pub struct SharedComponents {
pub tokenizer: Arc<dyn Tokenizer>,
pub tokenizer_registry: Arc<TokenizerRegistry>,
pub tool_parser_factory: ToolParserFactory,
pub reasoning_parser_factory: ReasoningParserFactory,
}

12
sgl-model-gateway/src/routers/grpc/pd_router.rs
View File

@@ -30,12 +30,9 @@ impl GrpcPDRouter {
let worker_registry = ctx.worker_registry.clone();
let policy_registry = ctx.policy_registry.clone();
// Extract necessary components from context
let tokenizer = ctx
.tokenizer
.as_ref()
.ok_or_else(|| "gRPC PD router requires tokenizer".to_string())?
.clone();
// Get tokenizer registry (no longer requires pre-loaded tokenizer)
let tokenizer_registry = ctx.tokenizer_registry.clone();
let reasoning_parser_factory = ctx
.reasoning_parser_factory
.as_ref()
@@ -49,7 +46,7 @@ impl GrpcPDRouter {
// Create shared components for pipeline
let shared_components = Arc::new(SharedComponents {
tokenizer: tokenizer.clone(),
tokenizer_registry: tokenizer_registry.clone(),
tool_parser_factory: tool_parser_factory.clone(),
reasoning_parser_factory: reasoning_parser_factory.clone(),
});
@@ -58,7 +55,6 @@ impl GrpcPDRouter {
let pipeline = RequestPipeline::new_pd(
worker_registry.clone(),
policy_registry.clone(),
tokenizer.clone(),
tool_parser_factory.clone(),
reasoning_parser_factory.clone(),
ctx.configured_tool_parser.clone(),

9
sgl-model-gateway/src/routers/grpc/pipeline.rs
View File

@@ -25,7 +25,6 @@ use crate::{
},
reasoning_parser::ParserFactory as ReasoningParserFactory,
routers::error,
tokenizer::traits::Tokenizer,
tool_parser::ParserFactory as ToolParserFactory,
};
@@ -45,14 +44,12 @@ impl RequestPipeline {
pub fn new_regular(
worker_registry: Arc<WorkerRegistry>,
policy_registry: Arc<PolicyRegistry>,
tokenizer: Arc<dyn Tokenizer>,
tool_parser_factory: ToolParserFactory,
reasoning_parser_factory: ReasoningParserFactory,
configured_tool_parser: Option<String>,
configured_reasoning_parser: Option<String>,
) -> Self {
let processor = processor::ResponseProcessor::new(
tokenizer.clone(),
tool_parser_factory.clone(),
reasoning_parser_factory.clone(),
configured_tool_parser.clone(),
@@ -60,7 +57,6 @@ impl RequestPipeline {
);
let streaming_processor = Arc::new(streaming::StreamingProcessor::new(
tokenizer,
tool_parser_factory,
reasoning_parser_factory,
configured_tool_parser,
@@ -92,7 +88,6 @@ impl RequestPipeline {
pub fn new_harmony(
worker_registry: Arc<WorkerRegistry>,
policy_registry: Arc<PolicyRegistry>,
_tokenizer: Arc<dyn Tokenizer>,
_tool_parser_factory: ToolParserFactory,
_reasoning_parser_factory: ReasoningParserFactory,
_configured_tool_parser: Option<String>,
@@ -122,7 +117,6 @@ impl RequestPipeline {
pub fn new_harmony_pd(
worker_registry: Arc<WorkerRegistry>,
policy_registry: Arc<PolicyRegistry>,
_tokenizer: Arc<dyn Tokenizer>,
_tool_parser_factory: ToolParserFactory,
_reasoning_parser_factory: ReasoningParserFactory,
_configured_tool_parser: Option<String>,
@@ -152,14 +146,12 @@ impl RequestPipeline {
pub fn new_pd(
worker_registry: Arc<WorkerRegistry>,
policy_registry: Arc<PolicyRegistry>,
tokenizer: Arc<dyn Tokenizer>,
tool_parser_factory: ToolParserFactory,
reasoning_parser_factory: ReasoningParserFactory,
configured_tool_parser: Option<String>,
configured_reasoning_parser: Option<String>,
) -> Self {
let processor = processor::ResponseProcessor::new(
tokenizer.clone(),
tool_parser_factory.clone(),
reasoning_parser_factory.clone(),
configured_tool_parser.clone(),
@@ -167,7 +159,6 @@ impl RequestPipeline {
);
let streaming_processor = Arc::new(streaming::StreamingProcessor::new(
tokenizer,
tool_parser_factory,
reasoning_parser_factory,
configured_tool_parser,

8
sgl-model-gateway/src/routers/grpc/regular/processor.rs
View File

@@ -35,7 +35,6 @@ use crate::{
/// Unified response processor for both routers
#[derive(Clone)]
pub struct ResponseProcessor {
pub tokenizer: Arc<dyn Tokenizer>,
pub tool_parser_factory: ToolParserFactory,
pub reasoning_parser_factory: ReasoningParserFactory,
pub configured_tool_parser: Option<String>,
@@ -44,14 +43,12 @@ pub struct ResponseProcessor {
impl ResponseProcessor {
pub fn new(
tokenizer: Arc<dyn Tokenizer>,
tool_parser_factory: ToolParserFactory,
reasoning_parser_factory: ReasoningParserFactory,
configured_tool_parser: Option<String>,
configured_reasoning_parser: Option<String>,
) -> Self {
Self {
tokenizer,
tool_parser_factory,
reasoning_parser_factory,
configured_tool_parser,
@@ -66,6 +63,7 @@ impl ResponseProcessor {
complete: &ProtoGenerateComplete,
index: usize,
original_request: &ChatCompletionRequest,
tokenizer: &Arc<dyn Tokenizer>,
stop_decoder: &mut StopSequenceDecoder,
history_tool_calls_count: usize,
reasoning_parser_available: bool,
@@ -176,7 +174,7 @@ impl ResponseProcessor {
// Step 4: Convert output logprobs if present
let logprobs = if let Some(proto_logprobs) = complete.output_logprobs() {
match utils::convert_proto_to_openai_logprobs(proto_logprobs, &self.tokenizer) {
match utils::convert_proto_to_openai_logprobs(proto_logprobs, tokenizer) {
Ok(logprobs) => Some(logprobs),
Err(e) => {
error!("Failed to convert logprobs: {}", e);
@@ -216,6 +214,7 @@ impl ResponseProcessor {
execution_result: ExecutionResult,
chat_request: Arc<ChatCompletionRequest>,
dispatch: DispatchMetadata,
tokenizer: Arc<dyn Tokenizer>,
stop_decoder: &mut StopSequenceDecoder,
request_logprobs: bool,
) -> Result<ChatCompletionResponse, axum::response::Response> {
@@ -269,6 +268,7 @@ impl ResponseProcessor {
complete,
index,
&chat_request,
&tokenizer,
stop_decoder,
history_tool_calls_count,
reasoning_parser_available,

2
sgl-model-gateway/src/routers/grpc/regular/responses/conversions.rs
View File

@@ -171,7 +171,7 @@ pub fn responses_to_chat(req: &ResponsesRequest) -> Result<ChatCompletionRequest
Ok(ChatCompletionRequest {
messages,
model: if req.model.is_empty() {
"default".to_string()
"unknown".to_string()
} else {
req.model.clone()
},

27
sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs
View File

@@ -43,14 +43,29 @@ impl ChatPreparationStage {
ctx: &mut RequestContext,
request: &ChatCompletionRequest,
) -> Result<(), Response> {
// Step 0: Resolve tokenizer from registry
let model_id = ctx.input.model_id.as_deref().unwrap();
let tokenizer = ctx
.components
.tokenizer_registry
.get(model_id)
.ok_or_else(|| {
error!(
function = "ChatPreparationStage::prepare_chat",
model = %model_id,
"Tokenizer not found for model"
);
error::internal_error(
"tokenizer_not_found",
format!("Tokenizer not found for model: {}", model_id),
)
})?;
// Step 1: Filter tools if needed
let body_ref = utils::filter_chat_request_by_tool_choice(request);
// Step 2: Process messages and apply chat template
let processed_messages = match utils::process_chat_messages(
&body_ref,
&*ctx.components.tokenizer,
) {
let processed_messages = match utils::process_chat_messages(&body_ref, &*tokenizer) {
Ok(msgs) => msgs,
Err(e) => {
error!(function = "ChatPreparationStage::execute", error = %e, "Failed to process chat messages");
@@ -59,7 +74,7 @@ impl ChatPreparationStage {
};
// Step 3: Tokenize the processed text
let encoding = match ctx.components.tokenizer.encode(&processed_messages.text) {
let encoding = match tokenizer.encode(&processed_messages.text) {
Ok(encoding) => encoding,
Err(e) => {
error!(function = "ChatPreparationStage::execute", error = %e, "Tokenization failed");
@@ -85,7 +100,7 @@ impl ChatPreparationStage {
// Step 5: Create stop sequence decoder (build once, reuse in non-stream)
let stop_decoder = utils::create_stop_decoder(
&ctx.components.tokenizer,
&tokenizer,
request.stop.as_ref(),
request.stop_token_ids.as_ref(),
request.skip_special_tokens,

20
sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
View File

@@ -79,12 +79,31 @@ impl ChatResponseProcessingStage {
})?
.clone();
// Get tokenizer in real time
let model_id = ctx.input.model_id.as_deref().unwrap();
let tokenizer = ctx
.components
.tokenizer_registry
.get(model_id)
.ok_or_else(|| {
error!(
function = "ChatPreparationStage::prepare_chat",
model = model_id,
"Tokenizer not found for model"
);
error::internal_error(
"tokenizer_not_found",
format!("Tokenizer not found for model: {}", model_id),
)
})?;
if is_streaming {
// Streaming: Use StreamingProcessor and return SSE response
let response = self.streaming_processor.clone().process_streaming_response(
execution_result,
ctx.chat_request_arc(), // Cheap Arc clone (8 bytes)
dispatch,
tokenizer,
);
// Attach load guards to response body for proper RAII lifecycle
@@ -118,6 +137,7 @@ impl ChatResponseProcessingStage {
execution_result,
chat_request,
dispatch,
tokenizer,
stop_decoder,
request_logprobs,
)

28
sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs
View File

@@ -44,8 +44,26 @@ impl GeneratePreparationStage {
ctx: &mut RequestContext,
request: &GenerateRequest,
) -> Result<(), Response> {
// Resolve input (text, prompt, or input_ids)
let (original_text, token_ids) = match self.resolve_generate_input(ctx, request) {
// Get model_id from context (normalized by router_manager)
let model_id = ctx.input.model_id.as_deref().unwrap();
let tokenizer = ctx
.components
.tokenizer_registry
.get(model_id)
.ok_or_else(|| {
error!(
function = "GeneratePreparationStage::execute",
model = %model_id,
"Tokenizer not found for model"
);
error::internal_error(
"tokenizer_not_found",
format!("Tokenizer not found for model: {}", model_id),
)
})?;
let (original_text, token_ids) = match self.resolve_generate_input(request, &tokenizer) {
Ok(res) => res,
Err(msg) => {
error!(function = "GeneratePreparationStage::execute", error = %msg, "Failed to resolve generate input");
@@ -56,7 +74,7 @@ impl GeneratePreparationStage {
// Create stop sequence decoder for generate requests
let params = request.sampling_params.as_ref();
let stop_decoder = utils::create_stop_decoder(
&ctx.components.tokenizer,
&tokenizer,
params.and_then(|p| p.stop.as_ref()),
params.and_then(|p| p.stop_token_ids.as_ref()),
params.and_then(|p| p.skip_special_tokens).unwrap_or(true),
@@ -84,12 +102,12 @@ impl GeneratePreparationStage {
fn resolve_generate_input(
&self,
ctx: &RequestContext,
request: &GenerateRequest,
tokenizer: &Arc<dyn Tokenizer>,
) -> Result<(Option<String>, Vec<u32>), String> {
if let Some(text) = &request.text {
return self
.tokenize_single_text(&ctx.components.tokenizer, text)
.tokenize_single_text(tokenizer, text)
.map(|(original, ids)| (Some(original), ids));
}

20
sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
View File

@@ -77,12 +77,32 @@ impl GenerateResponseProcessingStage {
})?
.clone();
// Get model_id from context
let model_id = ctx.input.model_id.as_deref().unwrap();
let tokenizer = ctx
.components
.tokenizer_registry
.get(model_id)
.ok_or_else(|| {
error!(
function = "GeneratePreparationStage::execute",
model = %model_id,
"Tokenizer not found for model"
);
error::internal_error(
"tokenizer_not_found",
format!("Tokenizer not found for model: {}", model_id),
)
})?;
if is_streaming {
// Streaming: Use StreamingProcessor and return SSE response
let response = self.streaming_processor.clone().process_streaming_generate(
execution_result,
ctx.generate_request_arc(), // Cheap Arc clone (8 bytes)
dispatch,
tokenizer,
);
// Attach load guards to response body for proper RAII lifecycle

33
sgl-model-gateway/src/routers/grpc/regular/streaming.rs
View File

@@ -39,7 +39,6 @@ use crate::{
/// Shared streaming processor for both single and dual dispatch modes
#[derive(Clone)]
pub struct StreamingProcessor {
tokenizer: Arc<dyn Tokenizer>,
tool_parser_factory: ToolParserFactory,
reasoning_parser_factory: ReasoningParserFactory,
configured_tool_parser: Option<String>,
@@ -58,7 +57,6 @@ struct GenerateStreamContext {
impl StreamingProcessor {
pub fn new(
tokenizer: Arc<dyn Tokenizer>,
tool_parser_factory: ToolParserFactory,
reasoning_parser_factory: ReasoningParserFactory,
configured_tool_parser: Option<String>,
@@ -66,7 +64,6 @@ impl StreamingProcessor {
backend_type: &'static str,
) -> Self {
Self {
tokenizer,
tool_parser_factory,
reasoning_parser_factory,
configured_tool_parser,
@@ -89,6 +86,7 @@ impl StreamingProcessor {
execution_result: context::ExecutionResult,
chat_request: Arc<ChatCompletionRequest>,
dispatch: context::DispatchMetadata,
tokenizer: Arc<dyn Tokenizer>,
) -> Response {
use bytes::Bytes;
use tokio::sync::mpsc;
@@ -108,11 +106,13 @@ impl StreamingProcessor {
context::ExecutionResult::Single { stream } => {
let processor = self.clone();
let dispatch_clone = dispatch.clone();
let tokenizer_clone = tokenizer.clone();
tokio::spawn(async move {
let result = processor
.process_streaming_chunks(
stream,
dispatch_clone,
tokenizer_clone,
stop_params,
chat_request,
&tx,
@@ -137,12 +137,14 @@ impl StreamingProcessor {
}
context::ExecutionResult::Dual { prefill, decode } => {
let processor = self.clone();
let tokenizer_clone = tokenizer.clone();
tokio::spawn(async move {
let result = processor
.process_dual_streaming_chunks(
prefill,
*decode,
dispatch,
tokenizer_clone,
stop_params,
chat_request,
&tx,
@@ -176,6 +178,7 @@ impl StreamingProcessor {
&self,
mut grpc_stream: ProtoStream,
dispatch: context::DispatchMetadata,
tokenizer: Arc<dyn Tokenizer>,
stop_params: (Option<StringOrArray>, Option<Vec<u32>>, bool, bool),
original_request: Arc<ChatCompletionRequest>,
tx: &UnboundedSender<Result<Bytes, io::Error>>,
@@ -285,7 +288,7 @@ impl StreamingProcessor {
let (ref stop, ref stop_token_ids, skip_special_tokens, no_stop_trim) =
stop_params;
utils::create_stop_decoder(
&self.tokenizer,
&tokenizer,
stop.as_ref(),
stop_token_ids.as_ref(),
skip_special_tokens,
@@ -303,10 +306,7 @@ impl StreamingProcessor {
// Process logprobs if present
let choice_logprobs = if let Some(proto_logprobs) = chunk.output_logprobs() {
match utils::convert_proto_to_openai_logprobs(
proto_logprobs,
&self.tokenizer,
) {
match utils::convert_proto_to_openai_logprobs(proto_logprobs, &tokenizer) {
Ok(logprobs) => Some(logprobs),
Err(e) => {
warn!("Failed to process logprobs: {}", e);
@@ -591,11 +591,13 @@ impl StreamingProcessor {
}
/// Process dual streaming chunks (prefill + decode) - PD mode
#[allow(clippy::too_many_arguments)]
pub async fn process_dual_streaming_chunks(
&self,
mut prefill_stream: ProtoStream,
decode_stream: ProtoStream,
dispatch: context::DispatchMetadata,
tokenizer: Arc<dyn Tokenizer>,
stop_params: (Option<StringOrArray>, Option<Vec<u32>>, bool, bool),
original_request: Arc<ChatCompletionRequest>,
tx: &UnboundedSender<Result<Bytes, io::Error>>,
@@ -621,7 +623,14 @@ impl StreamingProcessor {
// Phase 2-5: Process decode stream (same as single mode)
// Note: decode_stream will be marked completed inside process_streaming_chunks
let result = self
.process_streaming_chunks(decode_stream, dispatch, stop_params, original_request, tx)
.process_streaming_chunks(
decode_stream,
dispatch,
tokenizer,
stop_params,
original_request,
tx,
)
.await;
// Mark prefill stream as completed AFTER decode completes successfully
@@ -644,6 +653,7 @@ impl StreamingProcessor {
execution_result: context::ExecutionResult,
generate_request: Arc<GenerateRequest>,
dispatch: context::DispatchMetadata,
tokenizer: Arc<dyn Tokenizer>,
) -> Response {
// Create SSE channel
let (tx, rx) = mpsc::unbounded_channel::<Result<Bytes, io::Error>>();
@@ -663,7 +673,7 @@ impl StreamingProcessor {
// Spawn background task based on execution mode
match execution_result {
context::ExecutionResult::Single { stream } => {
let tokenizer = self.tokenizer.clone();
let tokenizer = tokenizer.clone();
tokio::spawn(async move {
let result =
Self::process_generate_streaming(tokenizer, stream, ctx, &tx).await;
@@ -677,7 +687,8 @@ impl StreamingProcessor {
});
}
context::ExecutionResult::Dual { prefill, decode } => {
let tokenizer = self.tokenizer.clone();
// For PD mode, need to handle prefill stream for input_logprobs
let tokenizer = tokenizer.clone();
tokio::spawn(async move {
let result = Self::process_generate_streaming_dual(
tokenizer, prefill, *decode, ctx, &tx,

13
sgl-model-gateway/src/routers/grpc/router.rs
View File

@@ -49,12 +49,9 @@ pub struct GrpcRouter {
impl GrpcRouter {
/// Create a new gRPC router
pub async fn new(ctx: &Arc<AppContext>) -> Result<Self, String> {
// Extract necessary components from context
let tokenizer = ctx
.tokenizer
.as_ref()
.ok_or_else(|| "gRPC router requires tokenizer".to_string())?
.clone();
// Get tokenizer registry (no longer requires pre-loaded tokenizer)
let tokenizer_registry = ctx.tokenizer_registry.clone();
let reasoning_parser_factory = ctx
.reasoning_parser_factory
.as_ref()
@@ -71,7 +68,7 @@ impl GrpcRouter {
// Create shared components for pipeline
let shared_components = Arc::new(SharedComponents {
tokenizer: tokenizer.clone(),
tokenizer_registry: tokenizer_registry.clone(),
tool_parser_factory: tool_parser_factory.clone(),
reasoning_parser_factory: reasoning_parser_factory.clone(),
});
@@ -80,7 +77,6 @@ impl GrpcRouter {
let pipeline = RequestPipeline::new_regular(
worker_registry.clone(),
_policy_registry.clone(),
tokenizer.clone(),
tool_parser_factory.clone(),
reasoning_parser_factory.clone(),
ctx.configured_tool_parser.clone(),
@@ -91,7 +87,6 @@ impl GrpcRouter {
let harmony_pipeline = RequestPipeline::new_harmony(
worker_registry.clone(),
_policy_registry.clone(),
tokenizer.clone(),
tool_parser_factory.clone(),
reasoning_parser_factory.clone(),
ctx.configured_tool_parser.clone(),

58
sgl-model-gateway/src/routers/router_manager.rs
View File

@@ -180,6 +180,52 @@ impl RouterManager {
self.routers.len()
}
/// Resolve model_id for a request, inferring from available workers if not specified
/// - If model_id is provided, use it directly
/// - If not provided and only one model exists, use it as implicit default
/// - If not provided and multiple models exist, return error requiring specification
/// - If no models exist, return service unavailable error
fn resolve_model_id(&self, model_id: Option<&str>) -> Result<String, Box<Response>> {
// If model_id is provided, use it
if let Some(id) = model_id {
return Ok(id.to_string());
}
// Get all available models from worker registry
let available_models = self.worker_registry.get_models();
match available_models.len() {
0 => Err(Box::new(
(
StatusCode::SERVICE_UNAVAILABLE,
"No models available - no workers registered",
)
.into_response(),
)),
1 => {
// Single model: use it as implicit default
debug!(
"Model not specified, using implicit default: {}",
available_models[0]
);
Ok(available_models[0].clone())
}
_ => {
// Multiple models: require explicit model specification
Err(Box::new(
(
StatusCode::BAD_REQUEST,
format!(
"Model must be specified. Available models: {}",
available_models.join(", ")
),
)
.into_response(),
))
}
}
}
pub fn get_router_for_model(&self, model_id: &str) -> Option<Arc<dyn RouterTrait>> {
let workers = self.worker_registry.get_by_model(model_id);
@@ -385,10 +431,18 @@ impl RouterTrait for RouterManager {
body: &GenerateRequest,
model_id: Option<&str>,
) -> Response {
let router = self.select_router_for_request(headers, model_id);
// Resolve model_id intelligently instead of falling back to "unknown"
let resolved_model_id = match self.resolve_model_id(model_id) {
Ok(id) => id,
Err(err_response) => return *err_response,
};
let router = self.select_router_for_request(headers, Some(&resolved_model_id));
if let Some(router) = router {
router.route_generate(headers, body, model_id).await
router
.route_generate(headers, body, Some(&resolved_model_id))
.await
} else {
(
StatusCode::NOT_FOUND,

2
sgl-model-gateway/src/service_discovery.rs
View File

@@ -628,7 +628,6 @@ mod tests {
policy_registry: Arc::new(crate::policies::PolicyRegistry::new(
router_config.policy.clone(),
)),
tokenizer: None,
reasoning_parser_factory: None,
tool_parser_factory: None,
router_manager: None,
@@ -643,6 +642,7 @@ mod tests {
worker_job_queue: worker_job_queue.clone(),
workflow_engine: Arc::new(std::sync::OnceLock::new()),
mcp_manager: Arc::new(std::sync::OnceLock::new()),
tokenizer_registry: Arc::new(crate::tokenizer::registry::TokenizerRegistry::new()),
wasm_manager: None,
worker_service: Arc::new(WorkerService::new(
worker_registry,

2
sgl-model-gateway/src/tokenizer/mod.rs
View File

@@ -6,6 +6,7 @@ pub mod cache;
pub mod factory;
pub mod hub;
pub mod mock;
pub mod registry;
pub mod sequence;
pub mod stop;
pub mod stream;
@@ -30,6 +31,7 @@ pub use factory::{
create_tokenizer_with_chat_template_blocking, TokenizerType,
};
pub use huggingface::HuggingFaceTokenizer;
pub use registry::TokenizerRegistry;
pub use sequence::Sequence;
pub use stop::{SequenceDecoderOutput, StopSequenceConfig, StopSequenceDecoder};
pub use stream::DecodeStream;

409
sgl-model-gateway/src/tokenizer/registry.rs Normal file
View File

@@ -0,0 +1,409 @@
//! Tokenizer Registry for dynamic tokenizer loading
//!
//! Provides thread-safe, deduplicated tokenizer loading for IGW mode where
//! multiple routers (HTTP and gRPC) need to share tokenizers across workers.
use std::sync::Arc;
use dashmap::DashMap;
use tokio::sync::Mutex;
use tracing::{debug, info};
use super::traits::Tokenizer;
/// Registry for managing tokenizers keyed by served_model_name
///
/// Features:
/// - Thread-safe concurrent access using DashMap
/// - Per-key locking to prevent duplicate loading
/// - Simple key scheme: served_model_name
pub struct TokenizerRegistry {
/// Storage for loaded tokenizers
tokenizers: DashMap<String, Arc<dyn Tokenizer>>,
/// Per-key locks to prevent duplicate loading
loading_locks: DashMap<String, Arc<Mutex<()>>>,
}
impl TokenizerRegistry {
/// Create a new empty registry
pub fn new() -> Self {
Self {
tokenizers: DashMap::new(),
loading_locks: DashMap::new(),
}
}
/// Load and register a tokenizer by model ID
///
/// If the tokenizer is already loaded, returns true immediately.
/// Otherwise, uses the provided loader function to load it.
/// Per-key locking ensures only one load happens per model, preventing race conditions.
///
/// # Arguments
/// * `model_id` - The model identifier to use as key
/// * `loader` - Async function that loads the tokenizer
///
/// # Returns
/// * `Ok(true)` - Successfully loaded and registered (or already registered)
/// * `Err(message)` - Error message if loading fails
///
/// # Example
/// ```ignore
/// registry.load("meta-llama/Llama-2-7b", || async {
/// create_tokenizer_async("/path/to/tokenizer").await
/// }).await?;
/// ```
pub async fn load<F, Fut>(&self, model_id: &str, loader: F) -> Result<bool, String>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<Arc<dyn Tokenizer>, String>>,
{
// Fast path: already loaded
if self.tokenizers.contains_key(model_id) {
debug!("Tokenizer already registered for model: {}", model_id);
return Ok(true);
}
debug!("Tokenizer cache miss for model: {}", model_id);
// Acquire per-key lock to prevent duplicate loading
let lock = self
.loading_locks
.entry(model_id.to_string())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone();
let _guard = lock.lock().await;
// Double-check after acquiring lock (another thread may have loaded it)
if self.tokenizers.contains_key(model_id) {
debug!("Tokenizer loaded by another thread for model: {}", model_id);
return Ok(true);
}
// Load tokenizer
info!("Loading tokenizer for model: {}", model_id);
let tokenizer = loader().await?;
// Store in registry
self.tokenizers.insert(model_id.to_string(), tokenizer);
// Remove the lock since it's no longer needed for this model.
self.loading_locks.remove(model_id);
info!(
"Successfully loaded and registered tokenizer for model: {}",
model_id
);
Ok(true)
}
/// Register a pre-loaded tokenizer
///
/// Atomically inserts a tokenizer into the registry only if no tokenizer
/// with the same model_name exists. Returns true if the tokenizer was inserted,
/// false if one already existed.
///
/// This method is thread-safe and uses atomic operations to prevent race conditions.
/// If you need to replace an existing tokenizer, first use `remove()` then `register()`.
///
/// # Arguments
/// * `model_name` - The served_model_name to use as key
/// * `tokenizer` - The tokenizer to register
///
/// # Returns
/// * `true` - If the tokenizer was successfully registered (didn't exist before)
/// * `false` - If a tokenizer with this model_name already existed
///
/// # Example
/// ```ignore
/// let tokenizer = create_tokenizer_blocking("/path/to/tokenizer")?;
/// if registry.register("meta-llama/Llama-2-7b", tokenizer) {
/// info!("Tokenizer registered successfully");
/// } else {
/// info!("Tokenizer already exists");
/// }
/// ```
pub fn register(&self, model_name: &str, tokenizer: Arc<dyn Tokenizer>) -> bool {
use dashmap::mapref::entry::Entry;
match self.tokenizers.entry(model_name.to_string()) {
Entry::Occupied(_) => {
debug!(
"Tokenizer already exists for model: {}, skipping registration",
model_name
);
false
}
Entry::Vacant(entry) => {
info!("Registering tokenizer for model: {}", model_name);
entry.insert(tokenizer);
true
}
}
}
/// Get a tokenizer if it's already loaded
///
/// Returns None if the tokenizer hasn't been loaded yet.
pub fn get(&self, model_name: &str) -> Option<Arc<dyn Tokenizer>> {
self.tokenizers.get(model_name).map(|t| t.clone())
}
/// Check if a tokenizer is loaded for the given model
pub fn contains(&self, model_name: &str) -> bool {
self.tokenizers.contains_key(model_name)
}
/// Get the number of loaded tokenizers
pub fn len(&self) -> usize {
self.tokenizers.len()
}
/// Check if the registry is empty
pub fn is_empty(&self) -> bool {
self.tokenizers.is_empty()
}
/// List all registered tokenizer keys (model names)
///
/// Returns a sorted vector of model names that have registered tokenizers.
/// Returns an empty vector if no tokenizers are registered.
pub fn list(&self) -> Vec<String> {
let mut keys: Vec<String> = self
.tokenizers
.iter()
.map(|entry| entry.key().clone())
.collect();
keys.sort();
keys
}
/// Remove a tokenizer from the registry
///
/// Returns the tokenizer if it was present.
pub fn remove(&self, model_name: &str) -> Option<Arc<dyn Tokenizer>> {
self.tokenizers.remove(model_name).map(|(_, v)| v)
}
/// Clear all tokenizers from the registry
pub fn clear(&self) {
self.tokenizers.clear();
self.loading_locks.clear();
}
}
impl Default for TokenizerRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use tokio::time::sleep;
use super::*;
use crate::tokenizer::mock::MockTokenizer;
#[tokio::test]
async fn test_basic_operations() {
let registry = TokenizerRegistry::new();
// Registry starts empty
assert!(registry.is_empty());
assert_eq!(registry.len(), 0);
assert!(!registry.contains("model1"));
// Load and register a tokenizer
registry
.load("model1", || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
.unwrap();
// Verify it's loaded
assert!(!registry.is_empty());
assert_eq!(registry.len(), 1);
assert!(registry.contains("model1"));
// Get returns the tokenizer
let tokenizer = registry.get("model1").unwrap();
assert_eq!(
tokenizer.vocab_size(),
MockTokenizer::default().vocab_size()
);
// Remove works
let removed = registry.remove("model1");
assert!(removed.is_some());
assert!(registry.is_empty());
}
#[tokio::test]
async fn test_load_prevents_duplicate_loading() {
let registry = Arc::new(TokenizerRegistry::new());
let load_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
// Spawn multiple tasks trying to load the same tokenizer
let mut handles = vec![];
for _ in 0..10 {
let registry = registry.clone();
let load_count = load_count.clone();
let handle = tokio::spawn(async move {
registry
.load("model1", || async {
// Simulate slow loading
sleep(Duration::from_millis(10)).await;
load_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
});
handles.push(handle);
}
// Wait for all tasks
for handle in handles {
handle.await.unwrap().unwrap();
}
// Verify tokenizer was loaded only once
assert_eq!(
load_count.load(std::sync::atomic::Ordering::SeqCst),
1,
"Tokenizer should be loaded exactly once despite concurrent requests"
);
assert_eq!(registry.len(), 1);
}
#[tokio::test]
async fn test_multiple_models() {
let registry = TokenizerRegistry::new();
// Load multiple tokenizers
for i in 1..=5 {
let model_name = format!("model{}", i);
registry
.load(&model_name, || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
.unwrap();
}
assert_eq!(registry.len(), 5);
assert!(registry.contains("model1"));
assert!(registry.contains("model5"));
assert!(!registry.contains("model6"));
// Clear all
registry.clear();
assert!(registry.is_empty());
}
#[tokio::test]
async fn test_load_failure() {
let registry = TokenizerRegistry::new();
// Try to load with a failing loader
let result = registry
.load("failing_model", || async { Err("Load failed".to_string()) })
.await;
assert!(result.is_err());
assert!(!registry.contains("failing_model"));
assert!(registry.is_empty());
}
#[tokio::test]
async fn test_concurrent_different_models() {
let registry = Arc::new(TokenizerRegistry::new());
let mut handles = vec![];
// Load different models concurrently
for i in 1..=10 {
let registry = registry.clone();
let handle = tokio::spawn(async move {
let model_name = format!("model{}", i);
registry
.load(&model_name, || async {
sleep(Duration::from_millis(5)).await;
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap().unwrap();
}
assert_eq!(registry.len(), 10);
}
#[tokio::test]
async fn test_register_only_if_absent() {
let registry = TokenizerRegistry::new();
let tokenizer1 = Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>;
let tokenizer2 = Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>;
// First registration should succeed
assert!(registry.register("model1", tokenizer1.clone()));
assert_eq!(registry.len(), 1);
assert!(registry.contains("model1"));
// Second registration with same key should fail
assert!(!registry.register("model1", tokenizer2.clone()));
assert_eq!(registry.len(), 1);
// Verify the original tokenizer is still there (not replaced)
let retrieved = registry.get("model1").unwrap();
assert_eq!(
Arc::as_ptr(&retrieved),
Arc::as_ptr(&tokenizer1),
"Original tokenizer should not be replaced"
);
// Registration with different key should succeed
assert!(registry.register("model2", tokenizer2));
assert_eq!(registry.len(), 2);
}
#[tokio::test]
async fn test_concurrent_register_same_model() {
let registry = Arc::new(TokenizerRegistry::new());
let success_count = Arc::new(std::sync::atomic::AtomicUsize::new(0));
// Spawn multiple tasks trying to register the same model
let mut handles = vec![];
for _ in 0..10 {
let registry = registry.clone();
let success_count = success_count.clone();
let handle = tokio::spawn(async move {
let tokenizer = Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>;
if registry.register("model1", tokenizer) {
success_count.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
});
handles.push(handle);
}
// Wait for all tasks
for handle in handles {
handle.await.unwrap();
}
// Verify only one registration succeeded
assert_eq!(
success_count.load(std::sync::atomic::Ordering::SeqCst),
1,
"Only one concurrent registration should succeed"
);
assert_eq!(registry.len(), 1);
}
}

8
sgl-model-gateway/tests/common/mod.rs
View File

@@ -27,6 +27,7 @@ use sgl_model_gateway::{
policies::PolicyRegistry,
protocols::common::{Function, Tool},
reasoning_parser::ParserFactory as ReasoningParserFactory,
tokenizer::registry::TokenizerRegistry,
tool_parser::ParserFactory as ToolParserFactory,
};
@@ -76,7 +77,7 @@ pub async fn create_test_context(config: RouterConfig) -> Arc<AppContext> {
.router_config(config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
.reasoning_parser_factory(None) // reasoning_parser_factory
.tool_parser_factory(None) // tool_parser_factory
.worker_registry(worker_registry)
@@ -181,6 +182,7 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
};
// Initialize registries
let tokenizer_registry = Arc::new(TokenizerRegistry::new());
let worker_registry = Arc::new(WorkerRegistry::new());
let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
@@ -211,7 +213,7 @@ pub async fn create_test_context_with_parsers(config: RouterConfig) -> Arc<AppCo
.router_config(config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.tokenizer_registry(tokenizer_registry)
.reasoning_parser_factory(reasoning_parser_factory)
.tool_parser_factory(tool_parser_factory)
.worker_registry(worker_registry)
@@ -347,7 +349,7 @@ pub async fn create_test_context_with_mcp_config(
.router_config(config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
.reasoning_parser_factory(None) // reasoning_parser_factory
.tool_parser_factory(None) // tool_parser_factory
.worker_registry(worker_registry)

5
sgl-model-gateway/tests/common/test_app.rs
View File

@@ -16,6 +16,7 @@ use sgl_model_gateway::{
policies::PolicyRegistry,
routers::RouterTrait,
server::{build_app, AppState},
tokenizer::registry::TokenizerRegistry,
};
/// Create a test Axum application using the actual server's build_app function
@@ -67,7 +68,7 @@ pub fn create_test_app(
.router_config(router_config.clone())
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
.reasoning_parser_factory(None) // reasoning_parser_factory
.tool_parser_factory(None) // tool_parser_factory
.worker_registry(worker_registry)
@@ -195,7 +196,7 @@ pub async fn create_test_app_context() -> Arc<AppContext> {
.router_config(router_config)
.client(client)
.rate_limiter(None)
.tokenizer(None)
.tokenizer_registry(Arc::new(TokenizerRegistry::new()))
.reasoning_parser_factory(None)
.tool_parser_factory(None)
.worker_registry(worker_registry)

3
sgl-model-gateway/tests/test_pd_routing.rs
View File

@@ -6,6 +6,7 @@ mod test_pd_routing {
config::{PolicyConfig, RouterConfig, RoutingMode},
core::{BasicWorkerBuilder, Worker, WorkerType},
routers::{http::pd_types::PDSelectionPolicy, RouterFactory},
tokenizer::registry::TokenizerRegistry,
};
#[derive(Debug)]
@@ -256,7 +257,7 @@ mod test_pd_routing {
.router_config(config)
.client(client)
.rate_limiter(rate_limiter)
.tokenizer(None) // tokenizer
.tokenizer_registry(Arc::new(TokenizerRegistry::new())) // tokenizer
.reasoning_parser_factory(None) // reasoning_parser_factory
.tool_parser_factory(None) // tool_parser_factory
.worker_registry(worker_registry)

4
sgl-model-gateway/tests/wasm_test.rs
View File

@@ -28,6 +28,7 @@ use sgl_model_gateway::{
policies::PolicyRegistry,
routers::RouterFactory,
server::{build_app, AppState},
tokenizer::TokenizerRegistry,
wasm::{
module::{
WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAttachPoint,
@@ -53,6 +54,7 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
// Create AppContext with wasm_manager from the start
let client = reqwest::Client::new();
let tokenizer_registry = Arc::new(TokenizerRegistry::new());
let worker_registry = Arc::new(WorkerRegistry::new());
let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone()));
@@ -80,7 +82,7 @@ async fn create_test_context_with_wasm() -> Arc<AppContext> {
.router_config(config.clone())
.client(client)
.rate_limiter(None)
.tokenizer(None)
.tokenizer_registry(tokenizer_registry)
.reasoning_parser_factory(None)
.tool_parser_factory(None)
.worker_registry(worker_registry)