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

refactor: unify registration through tokenizer_registration workflow (#17187)

Browse Source
This commit is contained in:
Chang Su
2026-01-16 10:53:36 -08:00
committed by GitHub
parent 3e968ab369
commit a7f5677abe
8 changed files with 222 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/app_context.rs
+13 -89
View File
@@ -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<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
/// 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<Self, String> {
// Create empty tokenizer registry
let registry = Arc::new(TokenizerRegistry::new());
// Try to load router-level tokenizer if path is provided
if let Some(tokenizer) = Self::maybe_tokenizer(config)? {
// Determine registration key: prefer tokenizer_path, then model_path, finally UNKNOWN_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, String> {
self.tokenizer_registry = Some(Arc::new(TokenizerRegistry::new()));
Ok(self)
}
sgl-model-gateway/src/config/types.rs
+12
View File
@@ -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
+2
View File
@@ -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
+71 -6
View File
@@ -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: {}",
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<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
+5 -9
View File
@@ -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"]),
)
sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs
+79 -37
View File
@@ -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<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;
};
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<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
+3
View File
@@ -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
+37 -1
View File
@@ -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