[model-gateway] Refine TokenizerRegisty.load() to handle duplication (#17230)
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user