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

[model-gateway] Refine TokenizerRegisty.load() to handle duplication (#17230)

Browse Source
This commit is contained in:
Chang Su
2026-01-16 13:53:21 -08:00
committed by GitHub
parent 2e14407983
commit 4229de3b13
8 changed files with 343 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/core/steps/mod.rs
+1 -1
View File
@@ -22,7 +22,7 @@ pub use mcp_registration::{
};
pub use tokenizer_registration::{
create_tokenizer_registration_workflow, create_tokenizer_workflow_data, LoadTokenizerStep,
TokenizerConfigRequest, TokenizerRemovalRequest, ValidateTokenizerConfigStep,
TokenizerConfigRequest, TokenizerRemovalRequest,
};
pub use wasm_module_registration::{
create_wasm_module_registration_workflow, create_wasm_registration_workflow_data,
sgl-model-gateway/src/core/steps/tokenizer_registration.rs
+95 -93
View File
@@ -11,7 +11,7 @@ use std::{sync::Arc, time::Duration};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tracing::{debug, error, info};
use tracing::{error, info};
use super::workflow_data::TokenizerWorkflowData;
use crate::{
@@ -20,6 +20,7 @@ use crate::{
tokenizer::{
cache::{CacheConfig, CachedTokenizer},
factory,
registry::LoadOutcome,
traits::Tokenizer,
},
workflow::{
@@ -42,6 +43,11 @@ pub struct TokenizerConfigRequest {
/// Optional cache configuration. If provided, wraps tokenizer with CachedTokenizer.
#[serde(default)]
pub cache_config: Option<TokenizerCacheConfig>,
/// If true, the workflow fails when a tokenizer with the same name already exists.
/// If false (default), the workflow succeeds and returns the existing tokenizer's ID.
/// API callers should set this to true.
#[serde(default)]
pub fail_on_duplicate: bool,
}
/// Configuration for removing a tokenizer
@@ -55,61 +61,13 @@ pub struct TokenizerRemovalRequest {
// Workflow Steps
// ============================================================================
/// Step 1: Validate the tokenizer configuration
pub struct ValidateTokenizerConfigStep;
#[async_trait]
impl StepExecutor<TokenizerWorkflowData> for ValidateTokenizerConfigStep {
async fn execute(
&self,
context: &mut WorkflowContext<TokenizerWorkflowData>,
) -> WorkflowResult<StepResult> {
let config = &context.data.config;
let app_context = context
.data
.app_context
.as_ref()
.ok_or_else(|| WorkflowError::ContextValueNotFound("app_context".to_string()))?;
debug!(
"Validating tokenizer config: name={}, source={}",
config.name, config.source
);
// Validate name is not empty
if config.name.is_empty() {
return Err(WorkflowError::StepFailed {
step_id: StepId::new("validate_config"),
message: "Tokenizer name cannot be empty".to_string(),
});
}
// Validate source is not empty
if config.source.is_empty() {
return Err(WorkflowError::StepFailed {
step_id: StepId::new("validate_config"),
message: "Tokenizer source cannot be empty".to_string(),
});
}
// Check if tokenizer already exists
if app_context.tokenizer_registry.contains(&config.name) {
return Err(WorkflowError::StepFailed {
step_id: StepId::new("validate_config"),
message: format!("Tokenizer '{}' already exists", config.name),
});
}
debug!("Tokenizer config validated successfully");
Ok(StepResult::Success)
}
fn is_retryable(&self, _error: &WorkflowError) -> bool {
false // Validation errors are not retryable
}
}
/// Step 2: Load the tokenizer from source (local path or HuggingFace)
/// Load the tokenizer from source (local path or HuggingFace)
///
/// This step handles:
/// - Input validation (via registry.load())
/// - Deduplication (returns success if already exists)
/// - Loading from local path or HuggingFace
/// - Optional caching layer wrapping
pub struct LoadTokenizerStep;
#[async_trait]
@@ -145,7 +103,8 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
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)
// Load the tokenizer using the registry's load method
// This handles: validation, deduplication, and loading
let result = app_context
.tokenizer_registry
.load(&id, &name, &source, || {
@@ -181,17 +140,38 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
.await;
match result {
Ok(loaded_id) => {
Ok(outcome) => {
let loaded_id = outcome.id();
// Get vocab size for logging
let vocab_size = app_context
.tokenizer_registry
.get_by_id(&loaded_id)
.get_by_id(loaded_id)
.map(|e| e.tokenizer.vocab_size());
info!(
"Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
name, loaded_id, vocab_size
);
match &outcome {
LoadOutcome::Loaded { id } => {
info!(
"Successfully loaded tokenizer '{}' (id: {}) with vocab_size: {:?}",
name, id, vocab_size
);
}
LoadOutcome::AlreadyExists { id } => {
if config.fail_on_duplicate {
return Err(WorkflowError::StepFailed {
step_id: StepId::new("load_tokenizer"),
message: format!(
"Tokenizer '{}' already exists (id: {})",
name, id
),
});
}
info!(
"Tokenizer '{}' already exists (id: {}), skipping load",
name, id
);
}
}
// Store vocab size in typed data
if let Some(size) = vocab_size {
@@ -204,7 +184,7 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
error!("Failed to load tokenizer '{}': {}", name, e);
Err(WorkflowError::StepFailed {
step_id: StepId::new("load_tokenizer"),
message: e,
message: e.to_string(),
})
}
}
@@ -221,38 +201,29 @@ impl StepExecutor<TokenizerWorkflowData> for LoadTokenizerStep {
/// Create the tokenizer registration workflow
///
/// This workflow:
/// - Validates the tokenizer configuration
/// - Loads the tokenizer from local path or HuggingFace
/// This workflow loads and registers a tokenizer. The single LoadTokenizerStep handles:
/// - Input validation (empty name/source)
/// - Deduplication (returns success if already exists)
/// - Loading from local path or HuggingFace
/// - Optional caching layer wrapping
///
/// Workflow configuration:
/// - ValidateConfig: No retry, 5s timeout (fast validation)
/// - LoadTokenizer: 3 retries, 5min timeout (may need to download from HuggingFace)
/// Configuration:
/// - 3 retries with 2s backoff (for network issues)
/// - 5 minute timeout (HuggingFace downloads can be slow)
pub fn create_tokenizer_registration_workflow() -> WorkflowDefinition<TokenizerWorkflowData> {
WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration")
.add_step(
StepDefinition::new(
"validate_config",
"Validate Configuration",
Arc::new(ValidateTokenizerConfigStep),
)
.with_timeout(Duration::from_secs(5))
.with_failure_action(FailureAction::FailWorkflow),
)
.add_step(
StepDefinition::new(
"load_tokenizer",
"Load Tokenizer",
Arc::new(LoadTokenizerStep),
)
.with_retry(RetryPolicy {
max_attempts: 3,
backoff: BackoffStrategy::Fixed(Duration::from_secs(2)),
})
.with_timeout(Duration::from_secs(300)) // 5 min for HuggingFace downloads
.with_failure_action(FailureAction::FailWorkflow)
.depends_on(&["validate_config"]),
WorkflowDefinition::new("tokenizer_registration", "Tokenizer Registration").add_step(
StepDefinition::new(
"load_tokenizer",
"Load Tokenizer",
Arc::new(LoadTokenizerStep),
)
.with_retry(RetryPolicy {
max_attempts: 3,
backoff: BackoffStrategy::Fixed(Duration::from_secs(2)),
})
.with_timeout(Duration::from_secs(300)) // 5 min for HuggingFace downloads
.with_failure_action(FailureAction::FailWorkflow),
)
}
/// Helper to create initial workflow data for tokenizer registration
@@ -279,6 +250,7 @@ mod tests {
source: "meta-llama/Llama-2-7b-hf".to_string(),
chat_template_path: None,
cache_config: None,
fail_on_duplicate: false,
};
let json = serde_json::to_string(&config).unwrap();
@@ -289,6 +261,35 @@ mod tests {
assert_eq!(parsed.source, "meta-llama/Llama-2-7b-hf");
assert!(parsed.chat_template_path.is_none());
assert!(parsed.cache_config.is_none());
assert!(!parsed.fail_on_duplicate);
}
#[test]
fn test_tokenizer_config_request_fail_on_duplicate_defaults_to_false() {
// Test that fail_on_duplicate defaults to false when not specified in JSON
let json = r#"{
"id": "test-uuid",
"name": "test-model",
"source": "/path/to/tokenizer"
}"#;
let parsed: TokenizerConfigRequest = serde_json::from_str(json).unwrap();
assert!(!parsed.fail_on_duplicate);
}
#[test]
fn test_tokenizer_config_request_fail_on_duplicate_true() {
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: None,
fail_on_duplicate: true,
};
let json = serde_json::to_string(&config).unwrap();
let parsed: TokenizerConfigRequest = serde_json::from_str(&json).unwrap();
assert!(parsed.fail_on_duplicate);
}
#[test]
@@ -304,6 +305,7 @@ mod tests {
enable_l1: false,
l1_max_memory: 0,
}),
fail_on_duplicate: false,
};
let json = serde_json::to_string(&config).unwrap();
sgl-model-gateway/src/core/steps/worker/local/mod.rs
+2 -2
View File
@@ -4,9 +4,9 @@ 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 submit_tokenizer_job;
mod update_policies_for_worker;
mod update_remaining_policies;
mod update_worker_properties;
@@ -27,9 +27,9 @@ 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::SubmitTokenizerJobStep;
pub use remove_from_policy_registry::RemoveFromPolicyRegistryStep;
pub use remove_from_worker_registry::RemoveFromWorkerRegistryStep;
pub use submit_tokenizer_job::SubmitTokenizerJobStep;
pub use update_policies_for_worker::UpdatePoliciesForWorkerStep;
pub use update_remaining_policies::UpdateRemainingPoliciesStep;
pub use update_worker_properties::UpdateWorkerPropertiesStep;
sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs → sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs
+6 -10
View File
@@ -1,8 +1,8 @@
//! 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).
//! tokenizer_registration workflow. The workflow handles validation, deduplication,
//! and caching - this step just submits the job.
use async_trait::async_trait;
use tracing::{debug, info, warn};
@@ -93,14 +93,9 @@ impl StepExecutor<LocalWorkerWorkflowData> for SubmitTokenizerJobStep {
continue;
};
// 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;
}
// Note: We don't check if tokenizer already exists here.
// The registry.load() handles deduplication gracefully (returns AlreadyExists).
// This simplifies the code and ensures consistent behavior.
info!(
"Submitting tokenizer registration job for model {} from {}",
@@ -114,6 +109,7 @@ impl StepExecutor<LocalWorkerWorkflowData> for SubmitTokenizerJobStep {
source: tokenizer_path,
chat_template_path: chat_template.clone(),
cache_config: cache_config.clone(),
fail_on_duplicate: false,
};
// Submit job (fire-and-forget, don't wait for completion)