diff --git a/sgl-model-gateway/src/core/steps/mod.rs b/sgl-model-gateway/src/core/steps/mod.rs index f0ebfe9fa..406c09b7e 100644 --- a/sgl-model-gateway/src/core/steps/mod.rs +++ b/sgl-model-gateway/src/core/steps/mod.rs @@ -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, diff --git a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs index 9c3319f04..6e98673ae 100644 --- a/sgl-model-gateway/src/core/steps/tokenizer_registration.rs +++ b/sgl-model-gateway/src/core/steps/tokenizer_registration.rs @@ -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, + /// 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 for ValidateTokenizerConfigStep { - async fn execute( - &self, - context: &mut WorkflowContext, - ) -> WorkflowResult { - 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 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 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 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 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 { - 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(); 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 d192cf15d..df00fb09a 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/mod.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/mod.rs @@ -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; diff --git a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs b/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs similarity index 90% rename from sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs rename to sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs index 0d86dfbd7..a85fe35f9 100644 --- a/sgl-model-gateway/src/core/steps/worker/local/register_tokenizer.rs +++ b/sgl-model-gateway/src/core/steps/worker/local/submit_tokenizer_job.rs @@ -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 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 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) diff --git a/sgl-model-gateway/src/routers/tokenize/handlers.rs b/sgl-model-gateway/src/routers/tokenize/handlers.rs index 64d3ac157..8698565b2 100644 --- a/sgl-model-gateway/src/routers/tokenize/handlers.rs +++ b/sgl-model-gateway/src/routers/tokenize/handlers.rs @@ -235,6 +235,7 @@ pub async fn add_tokenizer(context: &Arc, 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 { diff --git a/sgl-model-gateway/src/server.rs b/sgl-model-gateway/src/server.rs index 4a7dc174a..cfbfc9755 100644 --- a/sgl-model-gateway/src/server.rs +++ b/sgl-model-gateway/src/server.rs @@ -852,6 +852,7 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box &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>>, } +/// 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>>, + 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( &self, id: &str, name: &str, source: &str, loader: F, - ) -> Result + ) -> Result where F: FnOnce() -> Fut, Fut: std::future::Future, 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) }) .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) + }) + .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) + }) + .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) + }) + .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) + }) + .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 + } }