refactor: unify registration through tokenizer_registration workflow (#17187)

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::{
+    core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService},
-        steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService,
-        UNKNOWN_MODEL_ID,
-    },
     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::{
+    tokenizer::registry::TokenizerRegistry,
-        cache::{CacheConfig, CachedTokenizer},
-        factory as tokenizer_factory,
-        registry::TokenizerRegistry,
-        traits::Tokenizer,
-    },
     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<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);
-            }
-        };
-
-        // 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<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,
-                };
-                Arc::new(CachedTokenizer::new(base_tokenizer, cache_config)) as Arc<dyn Tokenizer>
-            } 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
+    /// Create empty tokenizer registry
-    /// 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).
+    /// Tokenizers are loaded via the tokenizer_registration workflow, which is triggered:
-    fn with_tokenizer_registry(mut self, config: &RouterConfig) -> Result<Self, String> {
+    /// - At startup (if --tokenizer-path or --model-path is provided)
-        // Create empty tokenizer registry
+    /// - When workers connect (registers under model_id)
-        let registry = Arc::new(TokenizerRegistry::new());
+    /// - Via POST /v1/tokenizers API (registers under user-specified name)
-
+    ///
-        // Try to load router-level tokenizer if path is provided
+    /// This unified approach ensures consistent behavior (caching, validation) across all paths.
-        if let Some(tokenizer) = Self::maybe_tokenizer(config)? {
+    fn with_tokenizer_registry(mut self, _config: &RouterConfig) -> Result<Self, String> {
-            // Determine registration key: prefer tokenizer_path, then model_path, finally UNKNOWN_MODEL_ID
+        self.tokenizer_registry = Some(Arc::new(TokenizerRegistry::new()));
-            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);
         Ok(self)
     }
 

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<Self> {
+        if self.enable_l0 || self.enable_l1 {
+            Some(self.clone())
+        } else {
+            None
+        }
+    }
+}
+
 impl Default for TokenizerCacheConfig {
     fn default() -> Self {
         Self {

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;

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<String>,
+    /// Optional cache configuration. If provided, wraps tokenizer with CachedTokenizer.
+    #[serde(default)]
+    pub cache_config: Option<TokenizerCacheConfig>,
 }
 
 /// Configuration for removing a tokenizer
@@ -115,8 +127,15 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
             .clone();
 
         info!(
-            "Loading tokenizer '{}' (id: {}) from source: {}",
+            "Loading tokenizer '{}' (id: {}) from source: {}{}",
-            config.name, config.id, config.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<TokenizerWorkflowData> 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<TokenizerWorkflowData> 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<dyn Tokenizer> = 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]

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",
+                "submit_tokenizer_job",
-                "Register Tokenizer",
+                "Submit Tokenizer Job",
-                Arc::new(RegisterTokenizerStep),
+                Arc::new(SubmitTokenizerJobStep),
             )
-            .with_retry(RetryPolicy {
+            .with_timeout(Duration::from_secs(5))
-                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"]),
         )

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,
+    core::{
-    tokenizer::{factory, TokenizerRegistry},
+        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)
+/// Step: Submit tokenizer registration job for the worker's model
-pub struct RegisterTokenizerStep;
+///
+/// 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<LocalWorkerWorkflowData> for RegisterTokenizerStep {
+impl StepExecutor<LocalWorkerWorkflowData> for SubmitTokenizerJobStep {
     async fn execute(
         &self,
         context: &mut WorkflowContext<LocalWorkerWorkflowData>,
@@ -29,6 +40,16 @@ impl StepExecutor<LocalWorkerWorkflowData> 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<LocalWorkerWorkflowData> 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;
             };
 
+            // Check if tokenizer already exists for this model
+            if app_context.tokenizer_registry.contains(&model_id) {
                 debug!(
-                "Registering tokenizer for model {} from {}",
+                    "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
+            // Create tokenizer config request
-            let tokenizer_id = TokenizerRegistry::generate_id();
+            let config = TokenizerConfigRequest {
-            let source = tokenizer_path.clone();
+                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
+            // Submit job (fire-and-forget, don't wait for completion)
-            let tokenizer_path_owned = tokenizer_path.clone();
+            if let Err(e) = job_queue
-            let template = chat_template.clone();
+                .submit(Job::AddTokenizer {
-            if let Err(e) = app_context
+                    config: Box::new(config),
-                .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())
-                    }
                 })
                 .await
             {
                 warn!(
-                    "Failed to load tokenizer for model {} from {}: {}",
+                    "Failed to submit tokenizer job for model {}: {}",
-                    model_id, source, e
+                    model_id, e
-                );
-            } else {
-                debug!(
-                    "Successfully registered tokenizer for model {} from {}",
-                    model_id, source
                 );
             }
         }
@@ -92,6 +134,6 @@ impl StepExecutor<LocalWorkerWorkflowData> 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
     }
 }

sgl-model-gateway/src/routers/tokenize/handlers.rs

@@ -227,11 +227,14 @@ pub async fn add_tokenizer(context: &Arc<AppContext>, 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 {

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<dyn std::error::Err
         config.router_config.health_check.timeout_secs
     );
 
+    // Submit startup tokenizer job if tokenizer path is configured
+    // This runs before worker initialization to ensure tokenizer is available
+    if let Some(tokenizer_source) = config
+        .router_config
+        .tokenizer_path
+        .as_ref()
+        .or(config.router_config.model_path.as_ref())
+    {
+        info!("Loading startup tokenizer from: {}", tokenizer_source);
+
+        let job_queue = app_context
+            .worker_job_queue
+            .get()
+            .expect("JobQueue should be initialized");
+
+        let tokenizer_config = TokenizerConfigRequest {
+            id: TokenizerRegistry::generate_id(),
+            name: tokenizer_source.clone(),
+            source: tokenizer_source.clone(),
+            chat_template_path: config.router_config.chat_template.clone(),
+            cache_config: config.router_config.tokenizer_cache.to_option(),
+        };
+
+        let job = Job::AddTokenizer {
+            config: Box::new(tokenizer_config),
+        };
+
+        job_queue
+            .submit(job)
+            .await
+            .map_err(|e| format!("Failed to submit startup tokenizer job: {}", e))?;
+
+        info!("Startup tokenizer job submitted (will complete in background)");
+    }
+
     info!(
         "Initializing workers for routing mode: {:?}",
         config.router_config.mode