From dd620987d17591ae8688a5f4630de0cb24024a30 Mon Sep 17 00:00:00 2001 From: Arthur Cheng Date: Tue, 23 Dec 2025 07:58:47 -0800 Subject: [PATCH] [model-gateway] Replace tokenizer with tokenizer registry for dynamic tokenizer loading in gRPC router (#12968) --- sgl-model-gateway/src/app_context.rs | 113 +++-- sgl-model-gateway/src/config/validation.rs | 32 +- sgl-model-gateway/src/core/model_card.rs | 2 +- .../src/core/steps/worker/local/mod.rs | 16 + .../steps/worker/local/register_tokenizer.rs | 78 ++++ sgl-model-gateway/src/policies/bucket.rs | 18 +- sgl-model-gateway/src/policies/cache_aware.rs | 18 +- .../grpc/common/stages/dispatch_metadata.rs | 4 +- sgl-model-gateway/src/routers/grpc/context.rs | 4 +- .../src/routers/grpc/pd_router.rs | 12 +- .../src/routers/grpc/pipeline.rs | 9 - .../src/routers/grpc/regular/processor.rs | 8 +- .../grpc/regular/responses/conversions.rs | 2 +- .../grpc/regular/stages/chat/preparation.rs | 27 +- .../stages/chat/response_processing.rs | 20 + .../regular/stages/generate/preparation.rs | 28 +- .../stages/generate/response_processing.rs | 20 + .../src/routers/grpc/regular/streaming.rs | 33 +- sgl-model-gateway/src/routers/grpc/router.rs | 13 +- .../src/routers/router_manager.rs | 58 ++- sgl-model-gateway/src/service_discovery.rs | 2 +- sgl-model-gateway/src/tokenizer/mod.rs | 2 + sgl-model-gateway/src/tokenizer/registry.rs | 409 ++++++++++++++++++ sgl-model-gateway/tests/common/mod.rs | 8 +- sgl-model-gateway/tests/common/test_app.rs | 5 +- sgl-model-gateway/tests/test_pd_routing.rs | 3 +- sgl-model-gateway/tests/wasm_test.rs | 4 +- 27 files changed, 793 insertions(+), 155 deletions(-) create mode 100644 sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs create mode 100644 sgl-model-gateway/src/tokenizer/registry.rs diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index 7488991dc..ce42d7749 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -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>, - pub tokenizer: Option>, + pub tokenizer_registry: Arc, pub reasoning_parser_factory: Option, pub tool_parser_factory: Option, pub worker_registry: Arc, @@ -67,7 +68,7 @@ pub struct AppContextBuilder { client: Option, router_config: Option, rate_limiter: Option>, - tokenizer: Option>, + tokenizer_registry: Option>, reasoning_parser_factory: Option, tool_parser_factory: Option, worker_registry: Option>, @@ -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>) -> Self { - self.tokenizer = tokenizer; + pub fn tokenizer_registry(mut self, tokenizer_registry: Arc) -> 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 { - 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>, 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 = + 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) + Arc::new(CachedTokenizer::new(base_tokenizer, cache_config)) as Arc } 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 { + // 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())); diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index f8a564849..44d46bd9a 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -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( diff --git a/sgl-model-gateway/src/core/model_card.rs b/sgl-model-gateway/src/core/model_card.rs index eee56a54c..b0da07e82 100644 --- a/sgl-model-gateway/src/core/model_card.rs +++ b/sgl-model-gateway/src/core/model_card.rs @@ -335,7 +335,7 @@ impl ModelCard { impl Default for ModelCard { fn default() -> Self { - Self::new("default") + Self::new("unknown") } } diff --git a/sgl-model-gateway/src/core/steps/worker/local/mod.rs b/sgl-model-gateway/src/core/steps/worker/local/mod.rs index 47a8a0e48..6c12d599f 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -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( diff --git a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs new file mode 100644 index 000000000..521caecab --- /dev/null +++ b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs @@ -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 { + let labels: Arc> = context + .get("labels") + .ok_or_else(|| WorkflowError::ContextValueNotFound("labels".to_string()))?; + let app_context: Arc = context + .get("app_context") + .ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?; + let workers: Arc>> = 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) + } +} diff --git a/sgl-model-gateway/src/policies/bucket.rs b/sgl-model-gateway/src/policies/bucket.rs index c86b2fbc7..473a20623 100644 --- a/sgl-model-gateway/src/policies/bucket.rs +++ b/sgl-model-gateway/src/policies/bucket.rs @@ -79,10 +79,10 @@ impl BucketPolicy { // Group workers by model let mut model_workers: HashMap>> = 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 }; diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 6a6545367..01b07b063 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -160,10 +160,10 @@ impl CacheAwarePolicy { let mut model_workers: std::collections::HashMap>> = 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 }; diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/dispatch_metadata.rs b/sgl-model-gateway/src/routers/grpc/common/stages/dispatch_metadata.rs index 3ec211dcb..51960fb1e 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/dispatch_metadata.rs +++ b/sgl-model-gateway/src/routers/grpc/common/stages/dispatch_metadata.rs @@ -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(), }; diff --git a/sgl-model-gateway/src/routers/grpc/context.rs b/sgl-model-gateway/src/routers/grpc/context.rs index c90339943..09636f2b3 100644 --- a/sgl-model-gateway/src/routers/grpc/context.rs +++ b/sgl-model-gateway/src/routers/grpc/context.rs @@ -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, + pub tokenizer_registry: Arc, pub tool_parser_factory: ToolParserFactory, pub reasoning_parser_factory: ReasoningParserFactory, } diff --git a/sgl-model-gateway/src/routers/grpc/pd_router.rs b/sgl-model-gateway/src/routers/grpc/pd_router.rs index 5f9b081ad..fd9ffcd38 100644 --- a/sgl-model-gateway/src/routers/grpc/pd_router.rs +++ b/sgl-model-gateway/src/routers/grpc/pd_router.rs @@ -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(), diff --git a/sgl-model-gateway/src/routers/grpc/pipeline.rs b/sgl-model-gateway/src/routers/grpc/pipeline.rs index 183ab2d91..b596ec7a2 100644 --- a/sgl-model-gateway/src/routers/grpc/pipeline.rs +++ b/sgl-model-gateway/src/routers/grpc/pipeline.rs @@ -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, policy_registry: Arc, - tokenizer: Arc, tool_parser_factory: ToolParserFactory, reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, ) -> 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, policy_registry: Arc, - _tokenizer: Arc, _tool_parser_factory: ToolParserFactory, _reasoning_parser_factory: ReasoningParserFactory, _configured_tool_parser: Option, @@ -122,7 +117,6 @@ impl RequestPipeline { pub fn new_harmony_pd( worker_registry: Arc, policy_registry: Arc, - _tokenizer: Arc, _tool_parser_factory: ToolParserFactory, _reasoning_parser_factory: ReasoningParserFactory, _configured_tool_parser: Option, @@ -152,14 +146,12 @@ impl RequestPipeline { pub fn new_pd( worker_registry: Arc, policy_registry: Arc, - tokenizer: Arc, tool_parser_factory: ToolParserFactory, reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, ) -> 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, diff --git a/sgl-model-gateway/src/routers/grpc/regular/processor.rs b/sgl-model-gateway/src/routers/grpc/regular/processor.rs index 1499603bf..e68fe9286 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/processor.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/processor.rs @@ -35,7 +35,6 @@ use crate::{ /// Unified response processor for both routers #[derive(Clone)] pub struct ResponseProcessor { - pub tokenizer: Arc, pub tool_parser_factory: ToolParserFactory, pub reasoning_parser_factory: ReasoningParserFactory, pub configured_tool_parser: Option, @@ -44,14 +43,12 @@ pub struct ResponseProcessor { impl ResponseProcessor { pub fn new( - tokenizer: Arc, tool_parser_factory: ToolParserFactory, reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, configured_reasoning_parser: Option, ) -> 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, 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, dispatch: DispatchMetadata, + tokenizer: Arc, stop_decoder: &mut StopSequenceDecoder, request_logprobs: bool, ) -> Result { @@ -269,6 +268,7 @@ impl ResponseProcessor { complete, index, &chat_request, + &tokenizer, stop_decoder, history_tool_calls_count, reasoning_parser_available, diff --git a/sgl-model-gateway/src/routers/grpc/regular/responses/conversions.rs b/sgl-model-gateway/src/routers/grpc/regular/responses/conversions.rs index e89d47f6b..2faf41712 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/responses/conversions.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/responses/conversions.rs @@ -171,7 +171,7 @@ pub fn responses_to_chat(req: &ResponsesRequest) -> Result 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, diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs index 5c970c2ad..fc17f0a6b 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs @@ -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, ) diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs index da365c5d8..a9d46d140 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs @@ -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, ) -> Result<(Option, Vec), 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)); } diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs index e386b4d9a..e6b6012e3 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs @@ -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 diff --git a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs index e62bdee53..926178ce0 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs @@ -39,7 +39,6 @@ use crate::{ /// Shared streaming processor for both single and dual dispatch modes #[derive(Clone)] pub struct StreamingProcessor { - tokenizer: Arc, tool_parser_factory: ToolParserFactory, reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, @@ -58,7 +57,6 @@ struct GenerateStreamContext { impl StreamingProcessor { pub fn new( - tokenizer: Arc, tool_parser_factory: ToolParserFactory, reasoning_parser_factory: ReasoningParserFactory, configured_tool_parser: Option, @@ -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, dispatch: context::DispatchMetadata, + tokenizer: Arc, ) -> 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, stop_params: (Option, Option>, bool, bool), original_request: Arc, tx: &UnboundedSender>, @@ -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, stop_params: (Option, Option>, bool, bool), original_request: Arc, tx: &UnboundedSender>, @@ -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, dispatch: context::DispatchMetadata, + tokenizer: Arc, ) -> Response { // Create SSE channel let (tx, rx) = mpsc::unbounded_channel::>(); @@ -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, diff --git a/sgl-model-gateway/src/routers/grpc/router.rs b/sgl-model-gateway/src/routers/grpc/router.rs index 7e8b5db2a..c1570270e 100644 --- a/sgl-model-gateway/src/routers/grpc/router.rs +++ b/sgl-model-gateway/src/routers/grpc/router.rs @@ -49,12 +49,9 @@ pub struct GrpcRouter { impl GrpcRouter { /// Create a new gRPC router pub async fn new(ctx: &Arc) -> Result { - // 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(), diff --git a/sgl-model-gateway/src/routers/router_manager.rs b/sgl-model-gateway/src/routers/router_manager.rs index fb46c2602..9bcf01e52 100644 --- a/sgl-model-gateway/src/routers/router_manager.rs +++ b/sgl-model-gateway/src/routers/router_manager.rs @@ -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> { + // 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> { 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, diff --git a/sgl-model-gateway/src/service_discovery.rs b/sgl-model-gateway/src/service_discovery.rs index 7c6fdd271..c89e15462 100644 --- a/sgl-model-gateway/src/service_discovery.rs +++ b/sgl-model-gateway/src/service_discovery.rs @@ -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, diff --git a/sgl-model-gateway/src/tokenizer/mod.rs b/sgl-model-gateway/src/tokenizer/mod.rs index 78fe39158..ff494790a 100644 --- a/sgl-model-gateway/src/tokenizer/mod.rs +++ b/sgl-model-gateway/src/tokenizer/mod.rs @@ -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; diff --git a/sgl-model-gateway/src/tokenizer/registry.rs b/sgl-model-gateway/src/tokenizer/registry.rs new file mode 100644 index 000000000..d8a9f1e2a --- /dev/null +++ b/sgl-model-gateway/src/tokenizer/registry.rs @@ -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>, + /// Per-key locks to prevent duplicate loading + loading_locks: DashMap>>, +} + +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(&self, model_id: &str, loader: F) -> Result + where + F: FnOnce() -> Fut, + Fut: std::future::Future, 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) -> 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> { + 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 { + let mut keys: Vec = 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> { + 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) + }) + .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) + }) + .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) + }) + .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) + }) + .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; + let tokenizer2 = Arc::new(MockTokenizer::default()) as Arc; + + // 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; + 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); + } +} diff --git a/sgl-model-gateway/tests/common/mod.rs b/sgl-model-gateway/tests/common/mod.rs index c49a4fd3e..8bd223081 100644 --- a/sgl-model-gateway/tests/common/mod.rs +++ b/sgl-model-gateway/tests/common/mod.rs @@ -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 { .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 Arc Arc { .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) diff --git a/sgl-model-gateway/tests/test_pd_routing.rs b/sgl-model-gateway/tests/test_pd_routing.rs index f6a2e069a..5c0787fae 100644 --- a/sgl-model-gateway/tests/test_pd_routing.rs +++ b/sgl-model-gateway/tests/test_pd_routing.rs @@ -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) diff --git a/sgl-model-gateway/tests/wasm_test.rs b/sgl-model-gateway/tests/wasm_test.rs index 3b9621260..cc965ceb1 100644 --- a/sgl-model-gateway/tests/wasm_test.rs +++ b/sgl-model-gateway/tests/wasm_test.rs @@ -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 { // 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 { .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)