diff --git a/sgl-model-gateway/src/app_context.rs b/sgl-model-gateway/src/app_context.rs index e2e112e5b..1e1c930dc 100644 --- a/sgl-model-gateway/src/app_context.rs +++ b/sgl-model-gateway/src/app_context.rs @@ -4,14 +4,11 @@ use std::{ }; use reqwest::Client; -use tracing::{debug, info}; +use tracing::debug; use crate::{ config::RouterConfig, - core::{ - steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService, - UNKNOWN_MODEL_ID, - }, + core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService}, data_connector::{ create_storage, ConversationItemStorage, ConversationStorage, ResponseStorage, }, @@ -21,12 +18,7 @@ use crate::{ policies::PolicyRegistry, reasoning_parser::ParserFactory as ReasoningParserFactory, routers::router_manager::RouterManager, - tokenizer::{ - cache::{CacheConfig, CachedTokenizer}, - factory as tokenizer_factory, - registry::TokenizerRegistry, - traits::Tokenizer, - }, + tokenizer::registry::TokenizerRegistry, tool_parser::ParserFactory as ToolParserFactory, wasm::{config::WasmRuntimeConfig, module_manager::WasmModuleManager}, }; @@ -396,56 +388,6 @@ impl AppContextBuilder { self } - /// 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); - } - }; - - // 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 - ) - })?; - - // 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, - }; - Arc::new(CachedTokenizer::new(base_tokenizer, cache_config)) as Arc - } else { - // Use base tokenizer directly without caching - base_tokenizer - }; - - Ok(Some(tokenizer)) - } - /// Create reasoning parser factory for gRPC mode or IGW mode fn with_reasoning_parser_factory(mut self) -> Self { // Initialize reasoning parser factory @@ -460,34 +402,16 @@ 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_MODEL_ID 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_MODEL_ID - let source = config - .tokenizer_path - .as_ref() - .or(config.model_path.as_ref()) - .map(|s| s.as_str()) - .unwrap_or(UNKNOWN_MODEL_ID); - - let tokenizer_id = TokenizerRegistry::generate_id(); - registry.register(&tokenizer_id, source, source, tokenizer.clone()); - info!( - "Tokenizer loaded and registered with name '{}' id={} (vocab_size: {})", - source, - tokenizer_id, - tokenizer.vocab_size() - ); - } - - self.tokenizer_registry = Some(registry); + /// Create empty tokenizer registry + /// + /// Tokenizers are loaded via the tokenizer_registration workflow, which is triggered: + /// - At startup (if --tokenizer-path or --model-path is provided) + /// - When workers connect (registers under model_id) + /// - Via POST /v1/tokenizers API (registers under user-specified name) + /// + /// This unified approach ensures consistent behavior (caching, validation) across all paths. + fn with_tokenizer_registry(mut self, _config: &RouterConfig) -> Result { + self.tokenizer_registry = Some(Arc::new(TokenizerRegistry::new())); Ok(self) } diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index 6c22ebfff..95f09bbd9 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -118,6 +118,18 @@ fn default_l1_max_memory() -> usize { 50 * 1024 * 1024 // 50MB } +impl TokenizerCacheConfig { + /// Returns Some(self) if any caching is enabled, None otherwise. + /// Use this when passing cache config to tokenizer registration workflow. + pub fn to_option(&self) -> Option { + if self.enable_l0 || self.enable_l1 { + Some(self.clone()) + } else { + None + } + } +} + impl Default for TokenizerCacheConfig { fn default() -> Self { Self { diff --git a/sgl-model-gateway/src/core/steps/mod.rs b/sgl-model-gateway/src/core/steps/mod.rs index 48a8de8a4..f0ebfe9fa 100644 --- a/sgl-model-gateway/src/core/steps/mod.rs +++ b/sgl-model-gateway/src/core/steps/mod.rs @@ -83,3 +83,5 @@ pub use workflow_data::{ }; // Typed workflow engines pub use workflow_engines::WorkflowEngines; + +pub use crate::config::TokenizerCacheConfig; diff --git a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs index cf4979ab9..9c3319f04 100644 --- a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs +++ b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs @@ -2,6 +2,10 @@ //! //! This module provides a workflow for registering tokenizers asynchronously. //! Tokenizers can be loaded from local paths or downloaded from HuggingFace. +//! +//! This is the **single source of truth** for tokenizer registration. All paths +//! (startup, worker connection, API) should use this workflow to ensure consistent +//! behavior (validation, caching, deduplication). use std::{sync::Arc, time::Duration}; @@ -12,7 +16,12 @@ use tracing::{debug, error, info}; use super::workflow_data::TokenizerWorkflowData; use crate::{ app_context::AppContext, - tokenizer::factory, + config::TokenizerCacheConfig, + tokenizer::{ + cache::{CacheConfig, CachedTokenizer}, + factory, + traits::Tokenizer, + }, workflow::{ BackoffStrategy, FailureAction, RetryPolicy, StepDefinition, StepExecutor, StepId, StepResult, WorkflowContext, WorkflowDefinition, WorkflowError, WorkflowResult, @@ -24,12 +33,15 @@ use crate::{ pub struct TokenizerConfigRequest { /// Pre-generated UUID for this tokenizer pub id: String, - /// User-provided name + /// User-provided name (what to register under in the registry) pub name: String, /// Source: either a local path or HuggingFace model ID pub source: String, /// Optional path to chat template file pub chat_template_path: Option, + /// Optional cache configuration. If provided, wraps tokenizer with CachedTokenizer. + #[serde(default)] + pub cache_config: Option, } /// Configuration for removing a tokenizer @@ -115,8 +127,15 @@ impl StepExecutor for LoadTokenizerStep { .clone(); info!( - "Loading tokenizer '{}' (id: {}) from source: {}", - config.name, config.id, config.source + "Loading tokenizer '{}' (id: {}) from source: {}{}", + config.name, + config.id, + config.source, + if config.cache_config.is_some() { + " with caching" + } else { + "" + } ); // Clone needed values before async move @@ -124,6 +143,7 @@ impl StepExecutor for LoadTokenizerStep { let name = config.name.clone(); let source = config.source.clone(); let chat_template = config.chat_template_path.clone(); + let cache_config = config.cache_config.clone(); // Load the tokenizer using the registry's load method (handles deduplication) let result = app_context @@ -131,13 +151,31 @@ impl StepExecutor for LoadTokenizerStep { .load(&id, &name, &source, || { let source = source.clone(); let chat_template = chat_template.clone(); + let cache_cfg = cache_config.clone(); async move { - factory::create_tokenizer_async_with_chat_template( + // Load base tokenizer + let base_tokenizer = factory::create_tokenizer_async_with_chat_template( &source, chat_template.as_deref(), ) .await - .map_err(|e| format!("Failed to load tokenizer: {}", e)) + .map_err(|e| format!("Failed to load tokenizer: {}", e))?; + + // Wrap with caching layer if configured + let tokenizer: Arc = match cache_cfg { + Some(cfg) if cfg.enable_l0 || cfg.enable_l1 => { + let cache_config = CacheConfig { + enable_l0: cfg.enable_l0, + l0_max_entries: cfg.l0_max_entries, + enable_l1: cfg.enable_l1, + l1_max_memory: cfg.l1_max_memory, + }; + Arc::new(CachedTokenizer::new(base_tokenizer, cache_config)) + } + _ => base_tokenizer, + }; + + Ok(tokenizer) } }) .await; @@ -240,6 +278,7 @@ mod tests { name: "test-model".to_string(), source: "meta-llama/Llama-2-7b-hf".to_string(), chat_template_path: None, + cache_config: None, }; let json = serde_json::to_string(&config).unwrap(); @@ -249,6 +288,32 @@ mod tests { assert_eq!(parsed.name, "test-model"); assert_eq!(parsed.source, "meta-llama/Llama-2-7b-hf"); assert!(parsed.chat_template_path.is_none()); + assert!(parsed.cache_config.is_none()); + } + + #[test] + fn test_tokenizer_config_request_with_cache() { + let config = TokenizerConfigRequest { + id: "test-uuid-1234".to_string(), + name: "test-model".to_string(), + source: "meta-llama/Llama-2-7b-hf".to_string(), + chat_template_path: None, + cache_config: Some(TokenizerCacheConfig { + enable_l0: true, + l0_max_entries: 1000, + enable_l1: false, + l1_max_memory: 0, + }), + }; + + let json = serde_json::to_string(&config).unwrap(); + let parsed: TokenizerConfigRequest = serde_json::from_str(&json).unwrap(); + + assert!(parsed.cache_config.is_some()); + let cache = parsed.cache_config.unwrap(); + assert!(cache.enable_l0); + assert_eq!(cache.l0_max_entries, 1000); + assert!(!cache.enable_l1); } #[test] 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 49ada6bdb..d192cf15d 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -27,7 +27,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 register_tokenizer::SubmitTokenizerJobStep; pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep; pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep; pub use update_policies_for_worker::UpdatePoliciesForWorkerStep; @@ -159,15 +159,11 @@ pub fn create_local_worker_workflow( ) .add_step( StepDefinition::new( - "register_tokenizer", - "Register Tokenizer", - Arc::new(RegisterTokenizerStep), + "submit_tokenizer_job", + "Submit Tokenizer Job", + Arc::new(SubmitTokenizerJobStep), ) - .with_retry(RetryPolicy { - max_attempts: 3, - backoff: BackoffStrategy::Fixed(Duration::from_secs(1)), - }) - .with_timeout(Duration::from_secs(10)) + .with_timeout(Duration::from_secs(5)) .with_failure_action(FailureAction::ContinueNextStep) .depends_on(&["register_workers"]), ) 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 index a1fb5923a..0d86dfbd7 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs @@ -1,19 +1,30 @@ //! Tokenizer registration step for local workers. +//! +//! This step submits a Job::AddTokenizer to the job queue, which triggers the +//! tokenizer_registration workflow. This ensures all tokenizer registrations +//! go through the same workflow with consistent behavior (validation, caching). use async_trait::async_trait; -use tracing::{debug, warn}; +use tracing::{debug, info, warn}; use crate::{ - core::steps::workflow_data::LocalWorkerWorkflowData, - tokenizer::{factory, TokenizerRegistry}, + core::{ + steps::{workflow_data::LocalWorkerWorkflowData, TokenizerConfigRequest}, + Job, + }, + tokenizer::TokenizerRegistry, workflow::{StepExecutor, StepResult, WorkflowContext, WorkflowError, WorkflowResult}, }; -/// Step 6: Register tokenizer for the worker's model (optional, non-blocking) -pub struct RegisterTokenizerStep; +/// Step: Submit tokenizer registration job for the worker's model +/// +/// This step submits a Job::AddTokenizer to the job queue rather than loading +/// the tokenizer directly. This ensures tokenizer registration goes through +/// the unified tokenizer_registration workflow. +pub struct SubmitTokenizerJobStep; #[async_trait] -impl StepExecutor for RegisterTokenizerStep { +impl StepExecutor for SubmitTokenizerJobStep { async fn execute( &self, context: &mut WorkflowContext, @@ -29,6 +40,16 @@ impl StepExecutor for RegisterTokenizerStep { .actual_workers .as_ref() .ok_or_else(|| WorkflowError::ContextValueNotFound("workers".to_string()))?; + + // Get job queue + let job_queue = match app_context.worker_job_queue.get() { + Some(queue) => queue, + None => { + warn!("Job queue not available, skipping tokenizer registration"); + return Ok(StepResult::Success); + } + }; + // Get chat_template: worker config > global router config let chat_template = context .data @@ -37,53 +58,74 @@ impl StepExecutor for RegisterTokenizerStep { .clone() .or_else(|| app_context.router_config.chat_template.clone()); + // Get cache config from router config + let cache_config = app_context.router_config.tokenizer_cache.to_option(); + 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 with fallback chain: + // 1. Worker labels: tokenizer_path + // 2. Worker labels: model_path + // 3. Router config (CLI args): --tokenizer-path + // 4. Router config (CLI args): --model-path + let tokenizer_path: String = if let Some(path) = labels .get("tokenizer_path") .or_else(|| labels.get("model_path")) - else { + { + path.clone() + } else if let Some(path) = app_context + .router_config + .tokenizer_path + .as_ref() + .or(app_context.router_config.model_path.as_ref()) + { + debug!( + "Using router config tokenizer path '{}' for model {}", + path, model_id + ); + path.clone() + } else { warn!( - "No tokenizer_path or model_path found for model {}", + "No tokenizer_path or model_path found for model {} (checked worker labels and router config)", model_id ); - return Ok(StepResult::Success); + continue; }; - debug!( - "Registering tokenizer for model {} from {}", + // Check if tokenizer already exists for this model + if app_context.tokenizer_registry.contains(&model_id) { + debug!( + "Tokenizer already registered for model {}, skipping", + model_id + ); + continue; + } + + info!( + "Submitting tokenizer registration job for model {} from {}", model_id, tokenizer_path ); - // Generate ID for this tokenizer - let tokenizer_id = TokenizerRegistry::generate_id(); - let source = tokenizer_path.clone(); + // Create tokenizer config request + let config = TokenizerConfigRequest { + id: TokenizerRegistry::generate_id(), + name: model_id.clone(), + source: tokenizer_path, + chat_template_path: chat_template.clone(), + cache_config: cache_config.clone(), + }; - // Load tokenizer with thread safe lock - let tokenizer_path_owned = tokenizer_path.clone(); - let template = chat_template.clone(); - if let Err(e) = app_context - .tokenizer_registry - .load(&tokenizer_id, &model_id, &source, move || { - let path = tokenizer_path_owned; - let tmpl = template; - async move { - factory::create_tokenizer_async_with_chat_template(&path, tmpl.as_deref()) - .await - .map_err(|e| e.to_string()) - } + // Submit job (fire-and-forget, don't wait for completion) + if let Err(e) = job_queue + .submit(Job::AddTokenizer { + config: Box::new(config), }) .await { warn!( - "Failed to load tokenizer for model {} from {}: {}", - model_id, source, e - ); - } else { - debug!( - "Successfully registered tokenizer for model {} from {}", - model_id, source + "Failed to submit tokenizer job for model {}: {}", + model_id, e ); } } @@ -92,6 +134,6 @@ impl StepExecutor for RegisterTokenizerStep { } fn is_retryable(&self, _error: &WorkflowError) -> bool { - true // Tokenizer loading failures are retryable (network/IO issues) + false // Job submission failures are not retryable at this level } } diff --git a/sgl-model-gateway/src/routers/tokenize/handlers.rs b/sgl-model-gateway/src/routers/tokenize/handlers.rs index 8e984685d..64d3ac157 100644 --- a/sgl-model-gateway/src/routers/tokenize/handlers.rs +++ b/sgl-model-gateway/src/routers/tokenize/handlers.rs @@ -227,11 +227,14 @@ pub async fn add_tokenizer(context: &Arc, request: AddTokenizerReque let tokenizer_id = TokenizerRegistry::generate_id(); // Create the job with the pre-generated ID + // Note: API-initiated tokenizer loads don't use caching by default + // Caching is applied for startup and worker-initiated loads based on router config let config = TokenizerConfigRequest { id: tokenizer_id.clone(), name: request.name.clone(), source: request.source.clone(), chat_template_path: request.chat_template_path.clone(), + cache_config: None, }; let job = Job::AddTokenizer { diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index cc4749941..4a7dc174a 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -24,7 +24,7 @@ use crate::{ config::{RouterConfig, RoutingMode}, core::{ job_queue::{JobQueue, JobQueueConfig}, - steps::WorkflowEngines, + steps::{TokenizerConfigRequest, WorkflowEngines}, worker::WorkerType, worker_manager::WorkerManager, Job, @@ -60,6 +60,7 @@ use crate::{ }, routers::{conversations, parse, router_manager::RouterManager, tokenize, RouterTrait}, service_discovery::{start_service_discovery, ServiceDiscoveryConfig}, + tokenizer::TokenizerRegistry, wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module}, workflow::LoggingSubscriber, }; @@ -830,6 +831,41 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box