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

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

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,

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();

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;

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)

View File

@@ -235,6 +235,7 @@ pub async fn add_tokenizer(context: &Arc<AppContext>, request: AddTokenizerReque
source: request.source.clone(),
chat_template_path: request.chat_template_path.clone(),
cache_config: None,
fail_on_duplicate: true,
};
let job = Job::AddTokenizer {

View File

@@ -852,6 +852,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
source: tokenizer_source.clone(),
chat_template_path: config.router_config.chat_template.clone(),
cache_config: config.router_config.tokenizer_cache.to_option(),
fail_on_duplicate: false,
};
let job = Job::AddTokenizer {

View File

@@ -27,7 +27,7 @@ mod tests;
use factory::{create_tokenizer_from_file, create_tokenizer_with_chat_template};
// Re-export types used outside this module
pub use huggingface::HuggingFaceTokenizer;
pub use registry::TokenizerRegistry;
pub use registry::{LoadError, LoadOutcome, TokenizerRegistry};
pub use stop::StopSequenceDecoder;
pub use stream::DecodeStream;
pub use traits::{Decoder, Encoder, Encoding, SpecialTokens};

View File

@@ -19,12 +19,51 @@
use std::sync::Arc;
use dashmap::DashMap;
use thiserror::Error;
use tokio::sync::Mutex;
use tracing::{debug, info};
use uuid::Uuid;
use super::traits::Tokenizer;
/// Outcome of a tokenizer load operation
#[derive(Debug, Clone)]
pub enum LoadOutcome {
/// Tokenizer was newly loaded and registered
Loaded { id: String },
/// Tokenizer already existed, returning existing ID
AlreadyExists { id: String },
}
impl LoadOutcome {
/// Get the ID regardless of outcome
pub fn id(&self) -> &str {
match self {
LoadOutcome::Loaded { id } => id,
LoadOutcome::AlreadyExists { id } => id,
}
}
/// Returns true if the tokenizer was newly loaded
pub fn is_newly_loaded(&self) -> bool {
matches!(self, LoadOutcome::Loaded { .. })
}
}
/// Error type for tokenizer loading operations
#[derive(Debug, Error)]
pub enum LoadError {
/// Name cannot be empty
#[error("tokenizer name cannot be empty")]
EmptyName,
/// Source cannot be empty
#[error("tokenizer source cannot be empty")]
EmptySource,
/// Loading failed
#[error("{0}")]
LoadFailed(String),
}
/// Metadata and tokenizer instance for a registered tokenizer
#[derive(Clone)]
pub struct TokenizerEntry {
@@ -53,6 +92,19 @@ pub struct TokenizerRegistry {
loading_locks: DashMap<String, Arc<Mutex<()>>>,
}
/// RAII guard that removes the loading lock entry on drop.
/// Ensures cleanup happens on normal completion, early return, or panic.
struct LoadingLockGuard<'a> {
locks: &'a DashMap<String, Arc<Mutex<()>>>,
key: String,
}
impl Drop for LoadingLockGuard<'_> {
fn drop(&mut self) {
self.locks.remove(&self.key);
}
}
impl TokenizerRegistry {
/// Create a new empty registry
pub fn new() -> Self {
@@ -68,36 +120,46 @@ impl TokenizerRegistry {
Uuid::new_v4().to_string()
}
/// Load and register a tokenizer with a pre-generated ID
/// Load and register a tokenizer
///
/// If the tokenizer is already loaded (by name), returns the existing ID.
/// Otherwise, uses the provided loader function to load it.
/// Validates inputs, handles deduplication, and loads the tokenizer if needed.
/// Per-key locking ensures only one load happens per name, preventing race conditions.
///
/// # Arguments
/// * `id` - Pre-generated UUID for this tokenizer
/// * `name` - User-provided name
/// * `source` - Source path or HuggingFace model ID
/// * `id` - Pre-generated UUID (use `generate_id()` to create one)
/// * `name` - User-provided name (used for deduplication, must not be empty)
/// * `source` - Source path or HuggingFace model ID (must not be empty)
/// * `loader` - Async function that loads the tokenizer
///
/// # Returns
/// * `Ok(id)` - Successfully loaded and registered (returns the ID)
/// * `Err(message)` - Error message if loading fails
/// * `Ok(LoadOutcome::Loaded { id })` - Tokenizer was newly loaded
/// * `Ok(LoadOutcome::AlreadyExists { id })` - Tokenizer already existed
/// * `Err(LoadError)` - Validation failed or loading failed
pub async fn load<F, Fut>(
&self,
id: &str,
name: &str,
source: &str,
loader: F,
) -> Result<String, String>
) -> Result<LoadOutcome, LoadError>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<Arc<dyn Tokenizer>, String>>,
{
// Validate inputs
if name.is_empty() {
return Err(LoadError::EmptyName);
}
if source.is_empty() {
return Err(LoadError::EmptySource);
}
// Fast path: already loaded by name
if let Some(existing_id) = self.name_to_id.get(name) {
debug!("Tokenizer already registered for name: {}", name);
return Ok(existing_id.clone());
return Ok(LoadOutcome::AlreadyExists {
id: existing_id.clone(),
});
}
debug!("Tokenizer cache miss for name: {}", name);
@@ -109,24 +171,27 @@ impl TokenizerRegistry {
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone();
let _guard = lock.lock().await;
let _mutex_guard = lock.lock().await;
let _lock_cleanup = LoadingLockGuard {
locks: &self.loading_locks,
key: name.to_string(),
};
// Double-check after acquiring lock (another thread may have loaded it)
if let Some(existing_id) = self.name_to_id.get(name) {
debug!("Tokenizer loaded by another thread for name: {}", name);
return Ok(existing_id.clone());
return Ok(LoadOutcome::AlreadyExists {
id: existing_id.clone(),
});
}
// Load tokenizer
info!("Loading tokenizer '{}' from source: {}", name, source);
let result = loader().await;
// Always clean up the lock, whether loading succeeded or failed
self.loading_locks.remove(name);
let tokenizer = result.map_err(LoadError::LoadFailed)?;
let tokenizer = result?;
// Create entry
// Create entry with provided ID
let entry = TokenizerEntry {
id: id.to_string(),
name: name.to_string(),
@@ -143,18 +208,21 @@ impl TokenizerRegistry {
name, id
);
Ok(id.to_string())
Ok(LoadOutcome::Loaded { id: id.to_string() })
}
/// Register a pre-loaded tokenizer with a pre-generated ID
/// Register a preloaded tokenizer with a pre-generated ID
///
/// Atomically inserts a tokenizer into the registry only if no tokenizer
/// with the same name exists. Returns the ID if successful.
///
/// This is primarily used for testing. Production code should use `load()`.
///
/// # Returns
/// * `Some(id)` - If the tokenizer was successfully registered
/// * `None` - If a tokenizer with this name already existed
pub fn register(
#[cfg(test)]
pub(crate) fn register(
&self,
id: &str,
name: &str,
@@ -292,13 +360,17 @@ mod tests {
// Load and register a tokenizer
let id = TokenizerRegistry::generate_id();
registry
let outcome = registry
.load(&id, "model1", "path/to/model", || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
.unwrap();
// Verify LoadOutcome::Loaded
assert!(outcome.is_newly_loaded());
assert_eq!(outcome.id(), id);
// Verify it's loaded
assert!(!registry.is_empty());
assert_eq!(registry.len(), 1);
@@ -317,6 +389,65 @@ mod tests {
assert!(registry.is_empty());
}
#[tokio::test]
async fn test_load_returns_already_exists() {
let registry = TokenizerRegistry::new();
let id1 = TokenizerRegistry::generate_id();
let id2 = TokenizerRegistry::generate_id();
// First load should return Loaded
let outcome1 = registry
.load(&id1, "model1", "source1", || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
.unwrap();
assert!(outcome1.is_newly_loaded());
assert_eq!(outcome1.id(), id1);
// Second load with same name should return AlreadyExists with ORIGINAL id
let outcome2 = registry
.load(&id2, "model1", "source2", || async {
panic!("Loader should not be called for duplicate name");
})
.await
.unwrap();
assert!(!outcome2.is_newly_loaded());
assert_eq!(outcome2.id(), id1); // Returns the original ID, not id2
// Registry still has only one entry
assert_eq!(registry.len(), 1);
// Original source is preserved
let entry = registry.get_by_name("model1").unwrap();
assert_eq!(entry.source, "source1");
}
#[tokio::test]
async fn test_load_validation() {
let registry = TokenizerRegistry::new();
let id = TokenizerRegistry::generate_id();
// Empty name should fail
let result = registry
.load(&id, "", "source", || async {
panic!("Loader should not be called for invalid input");
})
.await;
assert!(matches!(result, Err(LoadError::EmptyName)));
// Empty source should fail
let result = registry
.load(&id, "model", "", || async {
panic!("Loader should not be called for invalid input");
})
.await;
assert!(matches!(result, Err(LoadError::EmptySource)));
// Registry should be empty (nothing was loaded)
assert!(registry.is_empty());
}
#[tokio::test]
async fn test_load_prevents_duplicate_loading() {
let registry = Arc::new(TokenizerRegistry::new());
@@ -459,4 +590,88 @@ mod tests {
assert!(result3.is_some());
assert_eq!(registry.len(), 2);
}
#[tokio::test]
async fn test_loading_lock_cleanup_on_panic() {
let registry = Arc::new(TokenizerRegistry::new());
// Spawn a task that will panic during loading
let registry_clone = registry.clone();
let handle = tokio::spawn(async move {
registry_clone
.load(
&TokenizerRegistry::generate_id(),
"panic-model",
"source",
|| async {
panic!("Simulated panic during tokenizer loading");
},
)
.await
});
// Wait for the task - it should panic
let result = handle.await;
assert!(result.is_err(), "Task should have panicked");
// The RAII guard should have cleaned up the loading lock.
// Verify by attempting another load with the same name - it should work,
// not deadlock or fail due to stale lock.
let id = TokenizerRegistry::generate_id();
let outcome = registry
.load(&id, "panic-model", "source", || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await;
// Should succeed - the lock was properly cleaned up
assert!(outcome.is_ok(), "Load should succeed after panic cleanup");
assert!(outcome.unwrap().is_newly_loaded());
assert_eq!(registry.len(), 1);
assert!(registry.contains("panic-model"));
}
#[tokio::test]
async fn test_loading_lock_cleanup_on_early_return() {
let registry = Arc::new(TokenizerRegistry::new());
// Load a tokenizer
let id1 = TokenizerRegistry::generate_id();
registry
.load(&id1, "model1", "source1", || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
.unwrap();
// Now simulate concurrent load attempts where one thread wins
// and another thread sees "already exists" in the double-check.
// The RAII guard should clean up the lock on early return.
// First, verify loading_locks is empty after successful load
// by checking that we can load a different model without issues
let id2 = TokenizerRegistry::generate_id();
let outcome = registry
.load(&id2, "model2", "source2", || async {
Ok(Arc::new(MockTokenizer::default()) as Arc<dyn Tokenizer>)
})
.await
.unwrap();
assert!(outcome.is_newly_loaded());
assert_eq!(registry.len(), 2);
// Try to load model1 again - should return AlreadyExists
// and the lock should be cleaned up (not leak)
let id3 = TokenizerRegistry::generate_id();
let outcome = registry
.load(&id3, "model1", "source1", || async {
panic!("Loader should not be called for existing model");
})
.await
.unwrap();
assert!(!outcome.is_newly_loaded());
assert_eq!(outcome.id(), id1); // Returns original ID
}
}