From 283fb3953e03064e7e2ef21b3190e11e5cb0e78b Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 15 Sep 2026 15:23:45 +0000 Subject: [PATCH 01/11] fix: cut idle embedding RSS without changing nomic quality Keep full-precision NomicEmbedTextV15. Default pool size to 1, lazy-load ONNX sessions so /health does not pin weights, unload after idle, cap ORT intra-op threads, and use jemalloc so dropped sessions return RSS. Co-authored-by: Torsten Dittmann --- .env | 2 +- Cargo.lock | 39 ++++- Cargo.toml | 7 +- Dockerfile | 2 +- README.md | 6 +- docker-compose.yml | 2 +- src/bin/warmup.rs | 3 +- src/embedding.rs | 424 ++++++++++++++++++++++++++++++++++----------- src/main.rs | 22 +++ tests/embed_e2e.rs | 2 + 10 files changed, 388 insertions(+), 121 deletions(-) diff --git a/.env b/.env index 2cbcba1..2ae8efe 100644 --- a/.env +++ b/.env @@ -1,3 +1,3 @@ -EMBEDDING_MODELS=nomic,bge-small +EMBEDDING_MODELS=nomic EMBEDDING_CACHE_DIR=./models EMBEDDING_POOL_SIZE=1 diff --git a/Cargo.lock b/Cargo.lock index 0605437..c3da6e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -592,6 +592,7 @@ dependencies = [ "serde", "serde_json", "sysinfo", + "tikv-jemallocator", "tokenizers", "tokio", "tracing", @@ -675,9 +676,9 @@ dependencies = [ [[package]] name = "fastembed" -version = "5.13.4" +version = "5.17.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0112bd54a5d1903b19c85609c282949523bb8bb39f1614d4db0017e0ef3b0ff" +checksum = "4539f4a2c4472269adc227587b935c0a973e6b5fc4a03e14bbe62608e06c2298" dependencies = [ "anyhow", "hf-hub", @@ -1865,9 +1866,9 @@ checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" [[package]] name = "ort" -version = "2.0.0-rc.12" +version = "2.0.0-rc.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7de3af33d24a745ffb8fab904b13478438d1cd52868e6f17735ef6e1f8bf133" +checksum = "4336a1e2b38848325241c72889086886004e589b7c74f335e60a8e8db5138a0b" dependencies = [ "ndarray", "ort-sys", @@ -1878,9 +1879,9 @@ dependencies = [ [[package]] name = "ort-sys" -version = "2.0.0-rc.12" +version = "2.0.0-rc.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90" +checksum = "cf211e3776eea6aec988552fa118dd746d70e1b1e5e244058d1c98015f3e5872" dependencies = [ "hmac-sha256", "lzma-rust2", @@ -2339,13 +2340,15 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "safetensors" -version = "0.7.0" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5" +checksum = "79b079b829cb27a1c3c374341345ed2e8b2c0c839034522cee576c140bd7f846" dependencies = [ "hashbrown 0.16.1", + "libc", "serde", "serde_json", + "tempfile", ] [[package]] @@ -2684,6 +2687,26 @@ dependencies = [ "zune-jpeg", ] +[[package]] +name = "tikv-jemalloc-sys" +version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd8aa5b2ab86a2cefa406d889139c162cbb230092f7d1d7cbc1716405d852a3b" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.47" diff --git a/Cargo.toml b/Cargo.toml index 60efb3b..3a43710 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,13 +5,16 @@ edition = "2024" [dependencies] axum = "0.8.7" -fastembed = "5.13.4" +fastembed = "5.15.0" futures = "0.3.32" hf-hub = "0.5.0" serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.145" sysinfo = "0.39.2" tokenizers = "0.22.2" -tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal"] } +tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal", "time"] } tracing = "0.1.44" tracing-subscriber = { version = "0.3.20", features = ["env-filter"] } + +[target.'cfg(not(target_env = "msvc"))'.dependencies] +tikv-jemallocator = "0.6" diff --git a/Dockerfile b/Dockerfile index 2160251..6c1500d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -34,7 +34,7 @@ COPY --from=builder /usr/local/bin/warmup /usr/local/bin/warmup # the model set with `--build-arg EMBEDDING_MODELS=...` (docker compose passes # this from .env). Pool size is forced to 1 to keep the build's memory low — # it only affects the warmup, not the runtime pool. -ARG EMBEDDING_MODELS=nomic,bge-small +ARG EMBEDDING_MODELS=nomic RUN EMBEDDING_MODELS="${EMBEDDING_MODELS}" EMBEDDING_POOL_SIZE=1 /usr/local/bin/warmup # EXPOSE is build-time metadata only; the actual port is controlled by the diff --git a/README.md b/README.md index a78b556..3fd5bf6 100644 --- a/README.md +++ b/README.md @@ -23,9 +23,11 @@ Configured via environment variables (set them in `.env`): | Variable | Default | Description | | --- | --- | --- | | `EMBEDDING_PORT` | `3000` | Port the service listens on. | -| `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models to load. | +| `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. | | `EMBEDDING_CACHE_DIR` | _(default cache)_ | Directory for downloaded model files. | -| `EMBEDDING_POOL_SIZE` | _(memory-derived)_ | Number of model instances per pool. | +| `EMBEDDING_POOL_SIZE` | `1` | Number of ONNX sessions per model. Raise this for concurrent throughput; each extra session keeps another copy of the weights resident. | +| `EMBEDDING_INTRA_THREADS` | `min(4, nproc)` | ONNX Runtime intra-op threads per session. | +| `EMBEDDING_IDLE_UNLOAD_SECS` | `300` | Drop a model's sessions this many seconds after last use (`0` disables). The next `/embed` reloads the same checkpoint from `EMBEDDING_CACHE_DIR`. | ## API diff --git a/docker-compose.yml b/docker-compose.yml index 3e3f08c..a5a8d02 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -4,7 +4,7 @@ services: context: . dockerfile: Dockerfile args: - EMBEDDING_MODELS: ${EMBEDDING_MODELS:-nomic,bge-small} + EMBEDDING_MODELS: ${EMBEDDING_MODELS:-nomic} image: embedding:latest container_name: embedding ports: diff --git a/src/bin/warmup.rs b/src/bin/warmup.rs index 4f3a348..cd3a17f 100644 --- a/src/bin/warmup.rs +++ b/src/bin/warmup.rs @@ -13,7 +13,8 @@ fn main() -> Result<(), Box> { "warmup: downloading and initializing {} model(s)", config.models.len() ); - let _ = EmbeddingClient::new(config)?; + let client = EmbeddingClient::new(config)?; + client.preload()?; tracing::info!("warmup: models cached and ready"); Ok(()) diff --git a/src/embedding.rs b/src/embedding.rs index 19c8140..c2f5dd5 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -6,10 +6,15 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::{ Arc, Mutex, - atomic::{AtomicUsize, Ordering}, + atomic::{AtomicU64, AtomicUsize, Ordering}, }; use tokenizers::Tokenizer; +/// Default number of ONNX sessions per model when `EMBEDDING_POOL_SIZE` is unset. +pub const DEFAULT_POOL_SIZE: usize = 1; +/// Unload a model this many seconds after last use. `0` disables unloading. +pub const DEFAULT_IDLE_UNLOAD_SECS: u64 = 300; + #[derive(Debug, Clone)] pub struct EmbeddingResult { pub model: String, @@ -26,14 +31,20 @@ pub struct EmbeddingConfig { pub pool_size: usize, pub execution_providers: Vec, pub sub_batch_size: usize, + pub intra_threads: usize, + pub idle_unload_secs: u64, } -fn default_pool_size() -> usize { +fn available_cpus() -> usize { std::thread::available_parallelism() .map(|n| n.get()) .unwrap_or(2) } +fn default_intra_threads() -> usize { + available_cpus().min(4) +} + fn memory_budget(host_available: u64, cgroup_free: Option) -> u64 { match cgroup_free { Some(cgroup_free) => host_available.min(cgroup_free), @@ -45,6 +56,24 @@ fn next_index(counter: &AtomicUsize, len: usize) -> usize { counter.fetch_add(1, Ordering::Relaxed) % len } +fn unix_now() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +/// Whether an idle model slot should drop its ONNX sessions. +pub(crate) fn should_unload(last_access: u64, now: u64, idle_unload_secs: u64) -> bool { + idle_unload_secs > 0 && last_access > 0 && now.saturating_sub(last_access) >= idle_unload_secs +} + +fn parse_usize_env(key: &str) -> Option { + std::env::var(key) + .ok() + .and_then(|size| size.parse::().ok()) +} + impl EmbeddingConfig { pub fn from_env() -> Self { let models: Vec = std::env::var("EMBEDDING_MODELS") @@ -63,11 +92,18 @@ impl EmbeddingConfig { let cache_dir = std::env::var("EMBEDDING_CACHE_DIR").ok(); - let pool_size = std::env::var("EMBEDDING_POOL_SIZE") - .ok() - .and_then(|size| size.parse::().ok()) + let pool_size = parse_usize_env("EMBEDDING_POOL_SIZE") + .filter(|&n| n >= 1) + .unwrap_or(DEFAULT_POOL_SIZE); + + let intra_threads = parse_usize_env("EMBEDDING_INTRA_THREADS") .filter(|&n| n >= 1) - .unwrap_or_else(default_pool_size); + .unwrap_or_else(default_intra_threads); + + let idle_unload_secs = std::env::var("EMBEDDING_IDLE_UNLOAD_SECS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(DEFAULT_IDLE_UNLOAD_SECS); Self { models, @@ -76,137 +112,240 @@ impl EmbeddingConfig { pool_size, execution_providers: Vec::new(), sub_batch_size: 0, + intra_threads, + idle_unload_secs, } } } pub struct EmbeddingClient { models: HashMap, + config: EmbeddingConfig, sub_batch_override: usize, gpu: bool, } struct LoadedModel { + spec: EmbeddingModel, model_name: String, next: AtomicUsize, - pool: Vec>>, + load: Mutex<()>, + inner: Mutex, dimension: usize, + last_access: AtomicU64, +} + +struct ModelSlot { + pool: Option>>>, + tokenizer: Option>, +} + +struct BuiltModel { + pool: Vec>>, tokenizer: Arc, } impl EmbeddingClient { + /// Register configured models without loading ONNX weights. Sessions are + /// created on the first `/embed` (or [`Self::preload`]). pub fn new(config: EmbeddingConfig) -> Result { - // ONNX model loading memory != actual inference memory. - // - // `TextEmbedding::try_new()` mainly loads: - // - model weights - // - tokenizer - // - ONNX graph/session - // - // However, ONNX Runtime lazily allocates most execution memory - // (attention buffers, tensor arenas, activations, kernel workspaces) - // only during the first real inference call. - // - // We therefore run a warmup inference before measuring memory usage, - // otherwise pool sizing would severely underestimate the true runtime - // footprint and may cause OOMs under load. let mut models = HashMap::new(); for model in &config.models { let model_name = format!("{:?}", model); - let loaded_model = Self::load_model(model, &model_name, &config)?; - models.insert(model_name, loaded_model); + models.insert( + model_name.clone(), + LoadedModel { + spec: model.clone(), + model_name, + next: AtomicUsize::new(0), + load: Mutex::new(()), + inner: Mutex::new(ModelSlot { + pool: None, + tokenizer: None, + }), + dimension: model::dimension(model), + last_access: AtomicU64::new(0), + }, + ); } let sub_batch_override = config.sub_batch_size; let gpu = !config.execution_providers.is_empty(); Ok(Self { models, + config, sub_batch_override, gpu, }) } + /// Load every configured model now. Used by the warmup binary so image + /// builds still populate the ONNX cache. + pub fn preload(&self) -> Result<(), String> { + for loaded in self.models.values() { + self.ensure_loaded(loaded)?; + loaded.last_access.store(unix_now(), Ordering::Relaxed); + } + Ok(()) + } + + /// Drop ONNX sessions that have been unused for `idle_unload_secs`. + pub fn unload_idle(&self) { + if self.config.idle_unload_secs == 0 { + return; + } + let now = unix_now(); + for loaded in self.models.values() { + let last = loaded.last_access.load(Ordering::Relaxed); + if !should_unload(last, now, self.config.idle_unload_secs) { + continue; + } + let mut slot = match loaded.inner.lock() { + Ok(slot) => slot, + Err(poisoned) => poisoned.into_inner(), + }; + if slot.pool.is_some() || slot.tokenizer.is_some() { + slot.pool = None; + slot.tokenizer = None; + tracing::info!( + model = loaded.model_name.as_str(), + idle_secs = now.saturating_sub(last), + "unloaded idle embedding model" + ); + } + } + } + + pub fn idle_unload_secs(&self) -> u64 { + self.config.idle_unload_secs + } + + fn ensure_loaded(&self, loaded: &LoadedModel) -> Result<(), String> { + { + let slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + if slot.pool.is_some() && slot.tokenizer.is_some() { + return Ok(()); + } + } + + let _load = loaded + .load + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + { + let slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + if slot.pool.is_some() && slot.tokenizer.is_some() { + return Ok(()); + } + } + + let built = Self::load_model(&loaded.spec, &loaded.model_name, &self.config)?; + + let mut slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + slot.pool = Some(built.pool); + slot.tokenizer = Some(built.tokenizer); + Ok(()) + } + fn load_model( model: &EmbeddingModel, model_name: &str, config: &EmbeddingConfig, - ) -> Result { + ) -> Result { let desired_pool_size = config.pool_size.max(1); let dimension = model::dimension(model); - // loading instance and measuring memory footprint - let mut sys = sysinfo::System::new(); - sys.refresh_memory(); - let mem_before_loading_model = memory_budget( - sys.available_memory(), - sys.cgroup_limits().map(|limits| limits.free_memory), - ); - let has_gpu_providers = !config.execution_providers.is_empty(); + let mem_before_loading_model = if desired_pool_size > 1 { + let mut sys = sysinfo::System::new(); + sys.refresh_memory(); + Some(memory_budget( + sys.available_memory(), + sys.cgroup_limits().map(|limits| limits.free_memory), + )) + } else { + None + }; + let mut first_model = Self::init_model(model, config)?; // Tokenizer is fetched from the same cache dir fastembed just populated, // so this is a cache hit (no network) after the first model load. let tokenizer = Arc::new(Self::load_tokenizer(model, config)?); - // Run a warmup inference so the ONNX Runtime arena is allocated before - // we measure memory. Without this, per_instance only captures model - // weights and misses the arena buffers. A failure here means the model - // can't serve requests at all, so fail loudly — swallowing it would let - // the memory delta read ~0 and silently mis-size the pool (no OOM guard). - first_model - .embed(vec!["warmup"], None) - .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; - - sys.refresh_memory(); - let memory_after_loading_model = memory_budget( - sys.available_memory(), - sys.cgroup_limits().map(|limits| limits.free_memory), - ); - let per_instance_loaded = - mem_before_loading_model.saturating_sub(memory_after_loading_model); - - // ONNX Runtime uses arena allocation that grows with - // batch_size × sequence_length² (attention matrices) and is never - // released. The warmup above only allocates a minimal arena for a - // single short text. Apply a 3× multiplier to account for realistic - // inference workloads (batch=8-32 texts of 1000-2000 tokens each). - let per_instance_bytes = per_instance_loaded.saturating_mul(3); - - // determining pool size based on ram and capacity provided - let nproc = default_pool_size(); - // 60% of memory that was available before loading first model - let budget = mem_before_loading_model * 6 / 10; - let pool_size = if let Some(max_memory) = budget.checked_div(per_instance_bytes) { - if max_memory == 0 { - tracing::warn!( + let pool_size = if let Some(mem_before_loading_model) = mem_before_loading_model { + // ONNX model loading memory != actual inference memory. + // + // `TextEmbedding::try_new()` mainly loads weights / graph / session. + // ONNX Runtime lazily allocates execution memory on first inference. + // Warmup before measuring so extra pool slots are not oversized. + first_model + .embed(vec!["warmup"], None) + .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; + + let mut sys = sysinfo::System::new(); + sys.refresh_memory(); + let memory_after_loading_model = memory_budget( + sys.available_memory(), + sys.cgroup_limits().map(|limits| limits.free_memory), + ); + let per_instance_loaded = + mem_before_loading_model.saturating_sub(memory_after_loading_model); + + // ONNX Runtime uses arena allocation that grows with + // batch_size × sequence_length² (attention matrices) and is never + // released. The warmup above only allocates a minimal arena for a + // single short text. Apply a 3× multiplier to account for realistic + // inference workloads (batch=8-32 texts of 1000-2000 tokens each). + let per_instance_bytes = per_instance_loaded.saturating_mul(3); + + let nproc = available_cpus(); + let budget = mem_before_loading_model * 6 / 10; + if let Some(max_memory) = budget.checked_div(per_instance_bytes) { + if max_memory == 0 { + tracing::warn!( + estimated_with_arena_mb = per_instance_bytes / (1024 * 1024), + budget_mb = budget / (1024 * 1024), + "A single {} instance is estimated to exceed the memory budget; \ + running with pool_size=1 but the process may be OOM-killed under load. \ + Raise the container memory limit or pick a smaller model.", + model_name + ); + } + let max_memory = (max_memory as usize).max(1); + let capped = max_memory.min(desired_pool_size); + tracing::info!( + per_instance_mb = per_instance_loaded / (1024 * 1024), estimated_with_arena_mb = per_instance_bytes / (1024 * 1024), + available_mb = mem_before_loading_model / (1024 * 1024), budget_mb = budget / (1024 * 1024), - "A single {} instance is estimated to exceed the memory budget; \ - running with pool_size=1 but the process may be OOM-killed under load. \ - Raise the container memory limit or pick a smaller model.", - model_name + nproc = nproc, + desired = desired_pool_size, + max_from_memory = max_memory, + capped = capped, + "Measured ONNX model memory footprint" ); + capped + } else { + desired_pool_size } - let max_memory = (max_memory as usize).max(1); - let capped = max_memory.min(desired_pool_size); + } else { tracing::info!( - per_instance_mb = per_instance_loaded / (1024 * 1024), - estimated_with_arena_mb = per_instance_bytes / (1024 * 1024), - available_mb = mem_before_loading_model / (1024 * 1024), - budget_mb = budget / (1024 * 1024), - nproc = nproc, - desired = desired_pool_size, - max_from_memory = max_memory, - capped = capped, - "Measured ONNX model memory footprint" + model = model_name, + "Skipping warmup inference; pool_size=1 does not need a memory-based cap" ); - capped - } else { - desired_pool_size + 1 }; - // loading remaining instances in the pool along with the first model let mut pool = Vec::with_capacity(pool_size); pool.push(Arc::new(Mutex::new(first_model))); @@ -221,20 +360,15 @@ impl EmbeddingClient { "CPU" }; tracing::info!( - "Initialized embedding model: {} ({}d, pool_size={}, execution_provider={})", + "Initialized embedding model: {} ({}d, pool_size={}, intra_threads={}, execution_provider={})", model_name, dimension, pool_size, + config.intra_threads, ep_label, ); - Ok(LoadedModel { - dimension, - model_name: model_name.to_string(), - next: AtomicUsize::new(0), - pool, - tokenizer, - }) + Ok(BuiltModel { pool, tokenizer }) } fn load_tokenizer( @@ -265,7 +399,8 @@ impl EmbeddingClient { config: &EmbeddingConfig, ) -> Result { let mut init_options = InitOptions::new(model.clone()) - .with_show_download_progress(config.show_download_progress); + .with_show_download_progress(config.show_download_progress) + .with_intra_threads(config.intra_threads); if let Some(cache_dir) = &config.cache_dir { init_options = init_options.with_cache_dir(cache_dir.into()); @@ -279,6 +414,36 @@ impl EmbeddingClient { TextEmbedding::try_new(init_options) .map_err(|e| format!("Failed to initialize embedding model: {}", e)) } + + fn acquire_instance(&self, loaded: &LoadedModel) -> Result>, String> { + for _ in 0..2 { + self.ensure_loaded(loaded)?; + let slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + if let Some(pool) = slot.pool.as_ref() { + let idx = next_index(&loaded.next, pool.len()); + return Ok(pool[idx].clone()); + } + } + Err(format!( + "embedding model {} unloaded during acquire", + loaded.model_name + )) + } + + fn tokenizer(&self, loaded: &LoadedModel) -> Result, String> { + self.ensure_loaded(loaded)?; + let slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + slot.tokenizer + .clone() + .ok_or_else(|| format!("tokenizer missing for {}", loaded.model_name)) + } + pub async fn embed( &self, model_name: &str, @@ -293,6 +458,8 @@ impl EmbeddingClient { EmbedError::UnknownModel(format!("model not allowed: {}", model_name)) })?; + loaded.last_access.store(unix_now(), Ordering::Relaxed); + let sub_batch = if self.sub_batch_override > 0 { self.sub_batch_override } else { @@ -307,7 +474,7 @@ impl EmbeddingClient { let mut handles = Vec::new(); for chunk in texts.chunks(sub_batch) { - let inst = Self::acquire(loaded); + let inst = self.acquire_instance(loaded)?; let chunked_texts: Vec = chunk.iter().map(|t| (*t).to_owned()).collect(); handles.push(tokio::task::spawn_blocking(move || { @@ -327,7 +494,7 @@ impl EmbeddingClient { embeddings.append(&mut batch_result); } - let tokenizer = loaded.tokenizer.clone(); + let tokenizer = self.tokenizer(loaded)?; let owned_texts: Vec = texts.iter().map(|t| t.to_string()).collect(); let tokens = tokio::task::spawn_blocking(move || -> Result { let encodings = tokenizer @@ -338,6 +505,8 @@ impl EmbeddingClient { .await .map_err(|e| format!("Failed to join tokenizer task: {}", e))??; + loaded.last_access.store(unix_now(), Ordering::Relaxed); + Ok(EmbeddingResult { model: loaded.model_name.clone(), embeddings, @@ -346,12 +515,6 @@ impl EmbeddingClient { }) } - /// Round-robin acquire one instance from a model's pool. - fn acquire(loaded: &LoadedModel) -> Arc> { - let idx = next_index(&loaded.next, loaded.pool.len()); - loaded.pool[idx].clone() - } - /// Compute sub-batch size based on available system memory. /// /// Uses 50% of available RAM as a budget. @@ -378,10 +541,12 @@ mod tests { static ENV_LOCK: Mutex<()> = Mutex::new(()); - const ENV_KEYS: [&str; 3] = [ + const ENV_KEYS: [&str; 5] = [ "EMBEDDING_MODELS", "EMBEDDING_CACHE_DIR", "EMBEDDING_POOL_SIZE", + "EMBEDDING_INTRA_THREADS", + "EMBEDDING_IDLE_UNLOAD_SECS", ]; /// Holds the env mutex and restores the original values on drop. Tests that @@ -439,10 +604,13 @@ mod tests { let cfg = EmbeddingConfig::from_env(); assert!(matches!(cfg.models[0], EmbeddingModel::NomicEmbedTextV15)); assert_eq!(cfg.cache_dir, None); - assert!(cfg.pool_size >= 1); + assert_eq!(cfg.pool_size, DEFAULT_POOL_SIZE); assert!(cfg.show_download_progress); assert!(cfg.execution_providers.is_empty()); assert_eq!(cfg.sub_batch_size, 0); + assert!(cfg.intra_threads >= 1); + assert!(cfg.intra_threads <= 4); + assert_eq!(cfg.idle_unload_secs, DEFAULT_IDLE_UNLOAD_SECS); } #[test] @@ -539,7 +707,7 @@ mod tests { let _g = isolate_env(); set("EMBEDDING_POOL_SIZE", "0"); let cfg = EmbeddingConfig::from_env(); - assert_eq!(cfg.pool_size, default_pool_size()); + assert_eq!(cfg.pool_size, DEFAULT_POOL_SIZE); } #[test] @@ -547,7 +715,51 @@ mod tests { let _g = isolate_env(); set("EMBEDDING_POOL_SIZE", "not-a-number"); let cfg = EmbeddingConfig::from_env(); - assert_eq!(cfg.pool_size, default_pool_size()); + assert_eq!(cfg.pool_size, DEFAULT_POOL_SIZE); + } + + #[test] + fn from_env_parses_intra_threads() { + let _g = isolate_env(); + set("EMBEDDING_INTRA_THREADS", "2"); + let cfg = EmbeddingConfig::from_env(); + assert_eq!(cfg.intra_threads, 2); + } + + #[test] + fn from_env_parses_idle_unload_secs() { + let _g = isolate_env(); + set("EMBEDDING_IDLE_UNLOAD_SECS", "0"); + let cfg = EmbeddingConfig::from_env(); + assert_eq!(cfg.idle_unload_secs, 0); + set("EMBEDDING_IDLE_UNLOAD_SECS", "60"); + let cfg = EmbeddingConfig::from_env(); + assert_eq!(cfg.idle_unload_secs, 60); + } + + #[test] + fn new_does_not_load_onnx_sessions() { + let cfg = EmbeddingConfig { + models: vec![EmbeddingModel::AllMiniLML6V2], + show_download_progress: false, + cache_dir: None, + pool_size: 1, + execution_providers: Vec::new(), + sub_batch_size: 0, + intra_threads: 1, + idle_unload_secs: 0, + }; + let client = EmbeddingClient::new(cfg).expect("lazy construct"); + assert_eq!(client.idle_unload_secs(), 0); + assert_eq!(client.models.len(), 1); + } + + #[test] + fn should_unload_requires_prior_use_and_timeout() { + assert!(!should_unload(0, 1_000, 300)); + assert!(!should_unload(900, 1_000, 0)); + assert!(!should_unload(800, 1_000, 300)); + assert!(should_unload(700, 1_000, 300)); } #[test] @@ -590,7 +802,9 @@ mod tests { } #[test] - fn default_pool_size_is_positive() { - assert!(default_pool_size() >= 1); + fn default_intra_threads_is_capped() { + assert!(default_intra_threads() >= 1); + assert!(default_intra_threads() <= 4); + assert!(available_cpus() >= 1); } } diff --git a/src/main.rs b/src/main.rs index e54e130..28c49ff 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,9 @@ use std::sync::Arc; +use std::time::Duration; + +#[cfg(not(target_env = "msvc"))] +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; use axum::{ Json, Router, @@ -81,7 +86,24 @@ async fn main() -> Result<(), Box> { .init(); let config = EmbeddingConfig::from_env(); + let idle_unload_secs = config.idle_unload_secs; let client = Arc::new(EmbeddingClient::new(config)?); + + if idle_unload_secs > 0 { + let client_bg = client.clone(); + let tick_secs = (idle_unload_secs / 6).clamp(10, 30); + tokio::spawn(async move { + let mut ticker = tokio::time::interval(Duration::from_secs(tick_secs)); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + ticker.tick().await; // skip the immediate first tick + loop { + ticker.tick().await; + client_bg.unload_idle(); + } + }); + tracing::info!(idle_unload_secs, tick_secs, "idle model unload enabled"); + } + let state = AppState { client }; let app = Router::new() diff --git a/tests/embed_e2e.rs b/tests/embed_e2e.rs index ea92023..7a076d7 100644 --- a/tests/embed_e2e.rs +++ b/tests/embed_e2e.rs @@ -8,6 +8,8 @@ fn small_model_config(pool_size: usize) -> EmbeddingConfig { pool_size, execution_providers: Vec::new(), sub_batch_size: 0, + intra_threads: 1, + idle_unload_secs: 0, } } From 616f526af95c2cb641f1fd7fe830fd949e20b790 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 15 Sep 2026 15:30:34 +0000 Subject: [PATCH 02/11] fix: route ONNX allocations through jemalloc so idle unload returns RSS The Rust global allocator alone did not capture ONNX Runtime's C++ heap. Override libc malloc on supported platforms so dropped sessions can be returned to the OS after idle unload. Co-authored-by: Torsten Dittmann --- Cargo.lock | 1 + Cargo.toml | 3 ++- src/main.rs | 5 +++++ 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index c3da6e9..7e5014c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -592,6 +592,7 @@ dependencies = [ "serde", "serde_json", "sysinfo", + "tikv-jemalloc-sys", "tikv-jemallocator", "tokenizers", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 3a43710..4f745b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,4 +17,5 @@ tracing = "0.1.44" tracing-subscriber = { version = "0.3.20", features = ["env-filter"] } [target.'cfg(not(target_env = "msvc"))'.dependencies] -tikv-jemallocator = "0.6" +tikv-jemallocator = { version = "0.6", features = ["override_allocator_on_supported_platforms"] } +tikv-jemalloc-sys = "0.6" diff --git a/src/main.rs b/src/main.rs index 28c49ff..fcd47bb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,11 @@ use std::time::Duration; #[global_allocator] static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; +// Make jemalloc override libc malloc so ONNX Runtime's C++ allocations +// (the bulk of idle RSS) go through jemalloc and can be returned to the OS. +#[cfg(not(target_env = "msvc"))] +use tikv_jemalloc_sys as _; + use axum::{ Json, Router, extract::State, From 2e2099d949b10f1901608e9a83b80a9431046b6f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 15 Sep 2026 15:38:12 +0000 Subject: [PATCH 03/11] fix: install make in the Docker builder so jemalloc can compile tikv-jemalloc-sys runs configure then make. The slim image had g++ but not make, so the image build failed with ENOENT after jemalloc configured. Co-authored-by: Torsten Dittmann --- Dockerfile | 1 + 1 file changed, 1 insertion(+) diff --git a/Dockerfile b/Dockerfile index 6c1500d..c3fe1e7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,6 +5,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ pkg-config \ libssl-dev \ g++ \ + make \ && rm -rf /var/lib/apt/lists/* WORKDIR /app COPY . . From 7de58518606b84efb63057cf592d6b5da7c63c4c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 15 Sep 2026 17:56:14 +0000 Subject: [PATCH 04/11] fix: keep a RAM-capped session pool while the model is loaded Idle unload still drops every ONNX session. Desired pool size is CPU count again (capped by available memory at load) so concurrent /embed requests are not serialized on a single mutex. In-flight embeds block unload so a slow request cannot reload a second full pool. Co-authored-by: Torsten Dittmann --- .env | 1 - README.md | 2 +- src/embedding.rs | 66 +++++++++++++++++++++++++++++++++++------------- 3 files changed, 50 insertions(+), 19 deletions(-) diff --git a/.env b/.env index 2ae8efe..07c1fd6 100644 --- a/.env +++ b/.env @@ -1,3 +1,2 @@ EMBEDDING_MODELS=nomic EMBEDDING_CACHE_DIR=./models -EMBEDDING_POOL_SIZE=1 diff --git a/README.md b/README.md index 3fd5bf6..a7310d8 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,7 @@ Configured via environment variables (set them in `.env`): | `EMBEDDING_PORT` | `3000` | Port the service listens on. | | `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. | | `EMBEDDING_CACHE_DIR` | _(default cache)_ | Directory for downloaded model files. | -| `EMBEDDING_POOL_SIZE` | `1` | Number of ONNX sessions per model. Raise this for concurrent throughput; each extra session keeps another copy of the weights resident. | +| `EMBEDDING_POOL_SIZE` | CPU count | Number of ONNX sessions per model while it is loaded, then capped by available RAM. Concurrent `/embed` calls round-robin across sessions. Idle unload drops every session. | | `EMBEDDING_INTRA_THREADS` | `min(4, nproc)` | ONNX Runtime intra-op threads per session. | | `EMBEDDING_IDLE_UNLOAD_SECS` | `300` | Drop a model's sessions this many seconds after last use (`0` disables). The next `/embed` reloads the same checkpoint from `EMBEDDING_CACHE_DIR`. | diff --git a/src/embedding.rs b/src/embedding.rs index c2f5dd5..38deb25 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -10,8 +10,6 @@ use std::sync::{ }; use tokenizers::Tokenizer; -/// Default number of ONNX sessions per model when `EMBEDDING_POOL_SIZE` is unset. -pub const DEFAULT_POOL_SIZE: usize = 1; /// Unload a model this many seconds after last use. `0` disables unloading. pub const DEFAULT_IDLE_UNLOAD_SECS: u64 = 300; @@ -41,6 +39,10 @@ fn available_cpus() -> usize { .unwrap_or(2) } +fn default_pool_size() -> usize { + available_cpus() +} + fn default_intra_threads() -> usize { available_cpus().min(4) } @@ -94,7 +96,7 @@ impl EmbeddingConfig { let pool_size = parse_usize_env("EMBEDDING_POOL_SIZE") .filter(|&n| n >= 1) - .unwrap_or(DEFAULT_POOL_SIZE); + .unwrap_or_else(default_pool_size); let intra_threads = parse_usize_env("EMBEDDING_INTRA_THREADS") .filter(|&n| n >= 1) @@ -133,6 +135,7 @@ struct LoadedModel { inner: Mutex, dimension: usize, last_access: AtomicU64, + in_flight: AtomicUsize, } struct ModelSlot { @@ -140,6 +143,16 @@ struct ModelSlot { tokenizer: Option>, } +struct InFlightGuard<'a> { + counter: &'a AtomicUsize, +} + +impl Drop for InFlightGuard<'_> { + fn drop(&mut self) { + self.counter.fetch_sub(1, Ordering::SeqCst); + } +} + struct BuiltModel { pool: Vec>>, tokenizer: Arc, @@ -165,6 +178,7 @@ impl EmbeddingClient { }), dimension: model::dimension(model), last_access: AtomicU64::new(0), + in_flight: AtomicUsize::new(0), }, ); } @@ -204,6 +218,11 @@ impl EmbeddingClient { Ok(slot) => slot, Err(poisoned) => poisoned.into_inner(), }; + // Recheck under the pool lock so we never drop sessions that an + // in-flight embed already acquired (which would reload a second pool). + if loaded.in_flight.load(Ordering::SeqCst) > 0 { + continue; + } if slot.pool.is_some() || slot.tokenizer.is_some() { slot.pool = None; slot.tokenizer = None; @@ -282,16 +301,14 @@ impl EmbeddingClient { // so this is a cache hit (no network) after the first model load. let tokenizer = Arc::new(Self::load_tokenizer(model, config)?); - let pool_size = if let Some(mem_before_loading_model) = mem_before_loading_model { - // ONNX model loading memory != actual inference memory. - // - // `TextEmbedding::try_new()` mainly loads weights / graph / session. - // ONNX Runtime lazily allocates execution memory on first inference. - // Warmup before measuring so extra pool slots are not oversized. - first_model - .embed(vec!["warmup"], None) - .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; + // Always run one inference so Docker warmup (pool_size=1) still proves + // the session can execute, and so extra pool slots are sized from a + // post-arena RSS delta when desired_pool_size > 1. + first_model + .embed(vec!["warmup"], None) + .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; + let pool_size = if let Some(mem_before_loading_model) = mem_before_loading_model { let mut sys = sysinfo::System::new(); sys.refresh_memory(); let memory_after_loading_model = memory_budget( @@ -341,16 +358,21 @@ impl EmbeddingClient { } else { tracing::info!( model = model_name, - "Skipping warmup inference; pool_size=1 does not need a memory-based cap" + "Using pool_size=1; extra sessions are not created" ); 1 }; + let mut extra_config = config.clone(); + extra_config.intra_threads = (available_cpus() / pool_size) + .max(1) + .min(config.intra_threads); + let mut pool = Vec::with_capacity(pool_size); pool.push(Arc::new(Mutex::new(first_model))); for _ in 1..pool_size { - let inst = Self::init_model(model, config)?; + let inst = Self::init_model(model, &extra_config)?; pool.push(Arc::new(Mutex::new(inst))); } @@ -459,6 +481,10 @@ impl EmbeddingClient { })?; loaded.last_access.store(unix_now(), Ordering::Relaxed); + loaded.in_flight.fetch_add(1, Ordering::SeqCst); + let _in_flight = InFlightGuard { + counter: &loaded.in_flight, + }; let sub_batch = if self.sub_batch_override > 0 { self.sub_batch_override @@ -604,7 +630,7 @@ mod tests { let cfg = EmbeddingConfig::from_env(); assert!(matches!(cfg.models[0], EmbeddingModel::NomicEmbedTextV15)); assert_eq!(cfg.cache_dir, None); - assert_eq!(cfg.pool_size, DEFAULT_POOL_SIZE); + assert_eq!(cfg.pool_size, available_cpus()); assert!(cfg.show_download_progress); assert!(cfg.execution_providers.is_empty()); assert_eq!(cfg.sub_batch_size, 0); @@ -707,7 +733,7 @@ mod tests { let _g = isolate_env(); set("EMBEDDING_POOL_SIZE", "0"); let cfg = EmbeddingConfig::from_env(); - assert_eq!(cfg.pool_size, DEFAULT_POOL_SIZE); + assert_eq!(cfg.pool_size, available_cpus()); } #[test] @@ -715,7 +741,7 @@ mod tests { let _g = isolate_env(); set("EMBEDDING_POOL_SIZE", "not-a-number"); let cfg = EmbeddingConfig::from_env(); - assert_eq!(cfg.pool_size, DEFAULT_POOL_SIZE); + assert_eq!(cfg.pool_size, available_cpus()); } #[test] @@ -801,6 +827,12 @@ mod tests { assert_eq!(next_index(&counter, 3), 1); } + #[test] + fn default_pool_size_follows_cpu_count() { + assert_eq!(default_pool_size(), available_cpus()); + assert!(default_pool_size() >= 1); + } + #[test] fn default_intra_threads_is_capped() { assert!(default_intra_threads() >= 1); From 926c43749435e3f8b62350ce981b81f455ead2e9 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 11:14:50 +0000 Subject: [PATCH 05/11] fix: default to one ONNX session and cover idle/nomic e2e Keep EMBEDDING_POOL_SIZE at 1 unless operators opt in. Cap extra sessions from RAM and fall back to one session when the RSS delta is unusable. Default intra-op threads to the CPU count so the single session can still use the machine. Add ignored MiniLM unload lifecycle tests and a Nomic 768-d cosine ranking test. Co-authored-by: Torsten Dittmann --- README.md | 4 +- src/embedding.rs | 111 +++++++++++++++++++------------------- tests/embed_e2e.rs | 131 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 190 insertions(+), 56 deletions(-) diff --git a/README.md b/README.md index a7310d8..1f8bd2b 100644 --- a/README.md +++ b/README.md @@ -25,8 +25,8 @@ Configured via environment variables (set them in `.env`): | `EMBEDDING_PORT` | `3000` | Port the service listens on. | | `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. | | `EMBEDDING_CACHE_DIR` | _(default cache)_ | Directory for downloaded model files. | -| `EMBEDDING_POOL_SIZE` | CPU count | Number of ONNX sessions per model while it is loaded, then capped by available RAM. Concurrent `/embed` calls round-robin across sessions. Idle unload drops every session. | -| `EMBEDDING_INTRA_THREADS` | `min(4, nproc)` | ONNX Runtime intra-op threads per session. | +| `EMBEDDING_POOL_SIZE` | `1` | Number of ONNX sessions per model while it is loaded, then capped by available RAM. Concurrent `/embed` calls round-robin across sessions. Raise this for parallel HTTP throughput; each extra session keeps another copy of the weights resident until idle unload. | +| `EMBEDDING_INTRA_THREADS` | CPU count | ONNX Runtime intra-op threads per session. The default uses the whole machine on the single default session. | | `EMBEDDING_IDLE_UNLOAD_SECS` | `300` | Drop a model's sessions this many seconds after last use (`0` disables). The next `/embed` reloads the same checkpoint from `EMBEDDING_CACHE_DIR`. | ## API diff --git a/src/embedding.rs b/src/embedding.rs index 38deb25..dc34377 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -40,11 +40,25 @@ fn available_cpus() -> usize { } fn default_pool_size() -> usize { - available_cpus() + 1 } fn default_intra_threads() -> usize { - available_cpus().min(4) + available_cpus() +} + +/// Cap extra ONNX sessions so a failed RSS delta cannot fall back to an +/// uncapped CPU-count pool. +fn cap_pool_from_memory(desired: usize, budget: u64, per_instance_bytes: u64) -> usize { + let desired = desired.max(1); + if per_instance_bytes == 0 { + return 1; + } + match budget.checked_div(per_instance_bytes) { + Some(0) => 1, + Some(max_from_memory) => desired.min(max_from_memory as usize).max(1), + None => 1, + } } fn memory_budget(host_available: u64, cgroup_free: Option) -> u64 { @@ -327,34 +341,34 @@ impl EmbeddingClient { let nproc = available_cpus(); let budget = mem_before_loading_model * 6 / 10; - if let Some(max_memory) = budget.checked_div(per_instance_bytes) { - if max_memory == 0 { - tracing::warn!( - estimated_with_arena_mb = per_instance_bytes / (1024 * 1024), - budget_mb = budget / (1024 * 1024), - "A single {} instance is estimated to exceed the memory budget; \ - running with pool_size=1 but the process may be OOM-killed under load. \ - Raise the container memory limit or pick a smaller model.", - model_name - ); - } - let max_memory = (max_memory as usize).max(1); - let capped = max_memory.min(desired_pool_size); - tracing::info!( - per_instance_mb = per_instance_loaded / (1024 * 1024), + let capped = cap_pool_from_memory(desired_pool_size, budget, per_instance_bytes); + if per_instance_bytes == 0 { + tracing::warn!( + model = model_name, + desired = desired_pool_size, + "Could not measure ONNX instance size; using pool_size=1" + ); + } else if capped == 1 && desired_pool_size > 1 { + tracing::warn!( estimated_with_arena_mb = per_instance_bytes / (1024 * 1024), - available_mb = mem_before_loading_model / (1024 * 1024), budget_mb = budget / (1024 * 1024), - nproc = nproc, - desired = desired_pool_size, - max_from_memory = max_memory, - capped = capped, - "Measured ONNX model memory footprint" + "A single {} instance is estimated to exceed the memory budget; \ + running with pool_size=1 but the process may be OOM-killed under load. \ + Raise the container memory limit or pick a smaller model.", + model_name ); - capped - } else { - desired_pool_size } + tracing::info!( + per_instance_mb = per_instance_loaded / (1024 * 1024), + estimated_with_arena_mb = per_instance_bytes / (1024 * 1024), + available_mb = mem_before_loading_model / (1024 * 1024), + budget_mb = budget / (1024 * 1024), + nproc = nproc, + desired = desired_pool_size, + capped = capped, + "Measured ONNX model memory footprint" + ); + capped } else { tracing::info!( model = model_name, @@ -630,12 +644,11 @@ mod tests { let cfg = EmbeddingConfig::from_env(); assert!(matches!(cfg.models[0], EmbeddingModel::NomicEmbedTextV15)); assert_eq!(cfg.cache_dir, None); - assert_eq!(cfg.pool_size, available_cpus()); + assert_eq!(cfg.pool_size, 1); assert!(cfg.show_download_progress); assert!(cfg.execution_providers.is_empty()); assert_eq!(cfg.sub_batch_size, 0); - assert!(cfg.intra_threads >= 1); - assert!(cfg.intra_threads <= 4); + assert_eq!(cfg.intra_threads, available_cpus()); assert_eq!(cfg.idle_unload_secs, DEFAULT_IDLE_UNLOAD_SECS); } @@ -733,7 +746,7 @@ mod tests { let _g = isolate_env(); set("EMBEDDING_POOL_SIZE", "0"); let cfg = EmbeddingConfig::from_env(); - assert_eq!(cfg.pool_size, available_cpus()); + assert_eq!(cfg.pool_size, 1); } #[test] @@ -741,7 +754,7 @@ mod tests { let _g = isolate_env(); set("EMBEDDING_POOL_SIZE", "not-a-number"); let cfg = EmbeddingConfig::from_env(); - assert_eq!(cfg.pool_size, available_cpus()); + assert_eq!(cfg.pool_size, 1); } #[test] @@ -763,23 +776,6 @@ mod tests { assert_eq!(cfg.idle_unload_secs, 60); } - #[test] - fn new_does_not_load_onnx_sessions() { - let cfg = EmbeddingConfig { - models: vec![EmbeddingModel::AllMiniLML6V2], - show_download_progress: false, - cache_dir: None, - pool_size: 1, - execution_providers: Vec::new(), - sub_batch_size: 0, - intra_threads: 1, - idle_unload_secs: 0, - }; - let client = EmbeddingClient::new(cfg).expect("lazy construct"); - assert_eq!(client.idle_unload_secs(), 0); - assert_eq!(client.models.len(), 1); - } - #[test] fn should_unload_requires_prior_use_and_timeout() { assert!(!should_unload(0, 1_000, 300)); @@ -828,15 +824,22 @@ mod tests { } #[test] - fn default_pool_size_follows_cpu_count() { - assert_eq!(default_pool_size(), available_cpus()); - assert!(default_pool_size() >= 1); + fn cap_pool_from_memory_never_exceeds_budget_or_desired() { + const MIB: u64 = 1024 * 1024; + assert_eq!(cap_pool_from_memory(8, 600 * MIB, 0), 1); + assert_eq!(cap_pool_from_memory(8, 0, 100 * MIB), 1); + assert_eq!(cap_pool_from_memory(8, 250 * MIB, 100 * MIB), 2); + assert_eq!(cap_pool_from_memory(1, 10_000 * MIB, 100 * MIB), 1); + } + + #[test] + fn default_pool_size_is_one() { + assert_eq!(default_pool_size(), 1); } #[test] - fn default_intra_threads_is_capped() { + fn default_intra_threads_uses_available_cpus() { + assert_eq!(default_intra_threads(), available_cpus()); assert!(default_intra_threads() >= 1); - assert!(default_intra_threads() <= 4); - assert!(available_cpus() >= 1); } } diff --git a/tests/embed_e2e.rs b/tests/embed_e2e.rs index 7a076d7..741af35 100644 --- a/tests/embed_e2e.rs +++ b/tests/embed_e2e.rs @@ -1,4 +1,6 @@ use embedding::{EmbeddingClient, EmbeddingConfig, EmbeddingModel}; +use std::sync::Arc; +use std::time::Duration; fn small_model_config(pool_size: usize) -> EmbeddingConfig { EmbeddingConfig { @@ -13,6 +15,30 @@ fn small_model_config(pool_size: usize) -> EmbeddingConfig { } } +fn nomic_config() -> EmbeddingConfig { + EmbeddingConfig { + models: vec![EmbeddingModel::NomicEmbedTextV15], + show_download_progress: false, + cache_dir: None, + pool_size: 1, + execution_providers: Vec::new(), + sub_batch_size: 0, + intra_threads: 1, + idle_unload_secs: 0, + } +} + +fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); + let na = a.iter().map(|x| x * x).sum::().sqrt(); + let nb = b.iter().map(|x| x * x).sum::().sqrt(); + if na == 0.0 || nb == 0.0 { + 0.0 + } else { + dot / (na * nb) + } +} + #[tokio::test(flavor = "multi_thread")] #[ignore = "downloads the AllMiniLML6V2 ONNX model on first run"] async fn embed_single_text_returns_expected_dimension() { @@ -56,3 +82,108 @@ async fn embed_distinct_inputs_produce_distinct_vectors() { assert_eq!(result.embeddings.len(), 2); assert_ne!(result.embeddings[0], result.embeddings[1]); } + +#[tokio::test(flavor = "multi_thread")] +#[ignore = "downloads the AllMiniLML6V2 ONNX model on first run"] +async fn embed_after_idle_unload_reloads() { + let mut cfg = small_model_config(1); + cfg.idle_unload_secs = 1; + let client = EmbeddingClient::new(cfg).expect("client init"); + + let first = client + .embed("minilm", &["hello world"]) + .await + .expect("first embed should succeed"); + assert_eq!(first.embeddings.len(), 1); + assert_eq!(first.embeddings[0].len(), 384); + + tokio::time::sleep(Duration::from_secs(2)).await; + client.unload_idle(); + + let second = client + .embed("minilm", &["hello world"]) + .await + .expect("embed after idle unload should reload"); + assert_eq!(second.embeddings.len(), 1); + assert_eq!(second.embeddings[0].len(), 384); + assert!(second.tokens > 0); +} + +#[tokio::test(flavor = "multi_thread")] +#[ignore = "downloads the AllMiniLML6V2 ONNX model on first run"] +async fn in_flight_embed_survives_idle_unload() { + let mut cfg = small_model_config(1); + cfg.idle_unload_secs = 1; + let client = Arc::new(EmbeddingClient::new(cfg).expect("client init")); + + let embed_client = client.clone(); + let embed_task = tokio::spawn(async move { + embed_client + .embed( + "minilm", + &[ + "the cat sat on the mat", + "rust is a systems language", + "idle unload must not drop in-flight sessions", + ], + ) + .await + }); + + for _ in 0..50 { + client.unload_idle(); + tokio::time::sleep(Duration::from_millis(2)).await; + } + + let result = embed_task + .await + .expect("embed task should join") + .expect("in-flight embed should succeed while unload runs"); + assert_eq!(result.embeddings.len(), 3); + for embedding in &result.embeddings { + assert_eq!(embedding.len(), 384); + } + + tokio::time::sleep(Duration::from_secs(2)).await; + client.unload_idle(); + let after = client + .embed("minilm", &["hello"]) + .await + .expect("embed after in-flight request should succeed"); + assert_eq!(after.embeddings[0].len(), 384); +} + +#[tokio::test(flavor = "multi_thread")] +#[ignore = "downloads the NomicEmbedTextV15 ONNX model on first run"] +async fn nomic_embed_is_768d_and_ranks_similar_text_higher() { + let client = EmbeddingClient::new(nomic_config()).expect("client init"); + let result = client + .embed( + "nomic", + &[ + "The cat sat on the mat", + "A kitten rested on the rug", + "Rust is a systems programming language", + ], + ) + .await + .expect("nomic embed should succeed"); + + assert!( + result.model.to_lowercase().contains("nomic"), + "public model identity should be nomic, got {}", + result.model + ); + assert_eq!(result.embeddings.len(), 3); + for embedding in &result.embeddings { + assert_eq!(embedding.len(), 768); + } + assert!(result.tokens > 0); + + let similar = cosine_similarity(&result.embeddings[0], &result.embeddings[1]); + let dissimilar = cosine_similarity(&result.embeddings[0], &result.embeddings[2]); + assert!( + similar > dissimilar, + "related sentences should rank above an unrelated one: similar={similar} dissimilar={dissimilar}" + ); +} From 66c141d427bbdce06b9f50e7956a564134a25a83 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 11:21:35 +0000 Subject: [PATCH 06/11] fix: split intra-op threads across every pool session When EMBEDDING_POOL_SIZE is greater than one, the first ONNX session kept the full CPU-count thread setting while extras used nproc / pool_size. Rebuild the first session when the per-session count changes so concurrent embeds cannot oversubscribe the host. Co-authored-by: Torsten Dittmann --- README.md | 2 +- src/embedding.rs | 51 ++++++++++++++++++++++++++++++++++++++---------- 2 files changed, 42 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 1f8bd2b..e50f1dc 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ Configured via environment variables (set them in `.env`): | `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. | | `EMBEDDING_CACHE_DIR` | _(default cache)_ | Directory for downloaded model files. | | `EMBEDDING_POOL_SIZE` | `1` | Number of ONNX sessions per model while it is loaded, then capped by available RAM. Concurrent `/embed` calls round-robin across sessions. Raise this for parallel HTTP throughput; each extra session keeps another copy of the weights resident until idle unload. | -| `EMBEDDING_INTRA_THREADS` | CPU count | ONNX Runtime intra-op threads per session. The default uses the whole machine on the single default session. | +| `EMBEDDING_INTRA_THREADS` | CPU count | ONNX Runtime intra-op threads per session. The default uses the whole machine on the single default session. When `EMBEDDING_POOL_SIZE` is greater than one, threads are split across sessions (`nproc / pool_size`, still capped by this value) so concurrent embeds do not oversubscribe the host. | | `EMBEDDING_IDLE_UNLOAD_SECS` | `300` | Drop a model's sessions this many seconds after last use (`0` disables). The next `/embed` reloads the same checkpoint from `EMBEDDING_CACHE_DIR`. | ## API diff --git a/src/embedding.rs b/src/embedding.rs index dc34377..d322b09 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -47,6 +47,19 @@ fn default_intra_threads() -> usize { available_cpus() } +/// Intra-op threads for each ONNX session in a pool. A single session keeps +/// the configured value (CPU count by default). Extra sessions split CPUs +/// so concurrent embeds cannot oversubscribe the host. +fn intra_threads_for_pool(configured: usize, pool_size: usize, nproc: usize) -> usize { + let configured = configured.max(1); + let pool_size = pool_size.max(1); + if pool_size == 1 { + configured + } else { + (nproc / pool_size).max(1).min(configured) + } +} + /// Cap extra ONNX sessions so a failed RSS delta cannot fall back to an /// uncapped CPU-count pool. fn cap_pool_from_memory(desired: usize, budget: u64, per_instance_bytes: u64) -> usize { @@ -297,6 +310,10 @@ impl EmbeddingClient { let desired_pool_size = config.pool_size.max(1); let dimension = model::dimension(model); let has_gpu_providers = !config.execution_providers.is_empty(); + let nproc = available_cpus(); + let probe_intra = intra_threads_for_pool(config.intra_threads, desired_pool_size, nproc); + let mut probe_config = config.clone(); + probe_config.intra_threads = probe_intra; let mem_before_loading_model = if desired_pool_size > 1 { let mut sys = sysinfo::System::new(); @@ -309,7 +326,7 @@ impl EmbeddingClient { None }; - let mut first_model = Self::init_model(model, config)?; + let mut first_model = Self::init_model(model, &probe_config)?; // Tokenizer is fetched from the same cache dir fastembed just populated, // so this is a cache hit (no network) after the first model load. @@ -339,7 +356,6 @@ impl EmbeddingClient { // inference workloads (batch=8-32 texts of 1000-2000 tokens each). let per_instance_bytes = per_instance_loaded.saturating_mul(3); - let nproc = available_cpus(); let budget = mem_before_loading_model * 6 / 10; let capped = cap_pool_from_memory(desired_pool_size, budget, per_instance_bytes); if per_instance_bytes == 0 { @@ -377,16 +393,22 @@ impl EmbeddingClient { 1 }; - let mut extra_config = config.clone(); - extra_config.intra_threads = (available_cpus() / pool_size) - .max(1) - .min(config.intra_threads); + let session_intra = intra_threads_for_pool(config.intra_threads, pool_size, nproc); + let mut session_config = config.clone(); + session_config.intra_threads = session_intra; let mut pool = Vec::with_capacity(pool_size); - pool.push(Arc::new(Mutex::new(first_model))); - + if session_intra == probe_intra { + pool.push(Arc::new(Mutex::new(first_model))); + } else { + drop(first_model); + pool.push(Arc::new(Mutex::new(Self::init_model( + model, + &session_config, + )?))); + } for _ in 1..pool_size { - let inst = Self::init_model(model, &extra_config)?; + let inst = Self::init_model(model, &session_config)?; pool.push(Arc::new(Mutex::new(inst))); } @@ -400,7 +422,7 @@ impl EmbeddingClient { model_name, dimension, pool_size, - config.intra_threads, + session_intra, ep_label, ); @@ -842,4 +864,13 @@ mod tests { assert_eq!(default_intra_threads(), available_cpus()); assert!(default_intra_threads() >= 1); } + + #[test] + fn intra_threads_for_pool_splits_cpus_across_sessions() { + assert_eq!(intra_threads_for_pool(64, 1, 64), 64); + assert_eq!(intra_threads_for_pool(64, 4, 64), 16); + assert_eq!(intra_threads_for_pool(4, 4, 64), 4); + assert_eq!(intra_threads_for_pool(64, 8, 4), 1); + assert_eq!(intra_threads_for_pool(128, 1, 8), 128); + } } From 66412322943ef803a7dedd5b27bb2ff8c5c2f6a8 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 11:28:53 +0000 Subject: [PATCH 07/11] fix: warmup the rebuilt serving session after a RAM cap When the pool shrinks and intra-op threads change, drop the probe session only after initializing its replacement, then run the same warmup inference so preload cannot report ready on a session that has never executed. Co-authored-by: Torsten Dittmann --- src/embedding.rs | 38 ++++++++++++++++++++------------------ 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index d322b09..981c7c9 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -326,19 +326,12 @@ impl EmbeddingClient { None }; - let mut first_model = Self::init_model(model, &probe_config)?; + let first_model = Self::init_and_warmup(model, &probe_config, model_name)?; // Tokenizer is fetched from the same cache dir fastembed just populated, // so this is a cache hit (no network) after the first model load. let tokenizer = Arc::new(Self::load_tokenizer(model, config)?); - // Always run one inference so Docker warmup (pool_size=1) still proves - // the session can execute, and so extra pool slots are sized from a - // post-arena RSS delta when desired_pool_size > 1. - first_model - .embed(vec!["warmup"], None) - .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; - let pool_size = if let Some(mem_before_loading_model) = mem_before_loading_model { let mut sys = sysinfo::System::new(); sys.refresh_memory(); @@ -402,9 +395,10 @@ impl EmbeddingClient { pool.push(Arc::new(Mutex::new(first_model))); } else { drop(first_model); - pool.push(Arc::new(Mutex::new(Self::init_model( + pool.push(Arc::new(Mutex::new(Self::init_and_warmup( model, &session_config, + model_name, )?))); } for _ in 1..pool_size { @@ -429,6 +423,23 @@ impl EmbeddingClient { Ok(BuiltModel { pool, tokenizer }) } + fn warmup_session(model: &mut TextEmbedding, model_name: &str) -> Result<(), String> { + model + .embed(vec!["warmup"], None) + .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; + Ok(()) + } + + fn init_and_warmup( + model: &EmbeddingModel, + config: &EmbeddingConfig, + model_name: &str, + ) -> Result { + let mut inst = Self::init_model(model, config)?; + Self::warmup_session(&mut inst, model_name)?; + Ok(inst) + } + fn load_tokenizer( model: &EmbeddingModel, config: &EmbeddingConfig, @@ -864,13 +875,4 @@ mod tests { assert_eq!(default_intra_threads(), available_cpus()); assert!(default_intra_threads() >= 1); } - - #[test] - fn intra_threads_for_pool_splits_cpus_across_sessions() { - assert_eq!(intra_threads_for_pool(64, 1, 64), 64); - assert_eq!(intra_threads_for_pool(64, 4, 64), 16); - assert_eq!(intra_threads_for_pool(4, 4, 64), 4); - assert_eq!(intra_threads_for_pool(64, 8, 4), 1); - assert_eq!(intra_threads_for_pool(128, 1, 8), 128); - } } From 3e5806f1c6405a92770731f5456e52b6a1b2a110 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 11:34:26 +0000 Subject: [PATCH 08/11] test: drop helper-only pool default assertions Public env parsing already checks pool_size=1 and intra_threads=nproc. Idle unload and Nomic quality are covered by ignored e2e, not by reasserting private helper return values. Co-authored-by: Torsten Dittmann --- src/embedding.rs | 20 -------------------- 1 file changed, 20 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index 981c7c9..d69f417 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -855,24 +855,4 @@ mod tests { assert_eq!(next_index(&counter, 3), 0); assert_eq!(next_index(&counter, 3), 1); } - - #[test] - fn cap_pool_from_memory_never_exceeds_budget_or_desired() { - const MIB: u64 = 1024 * 1024; - assert_eq!(cap_pool_from_memory(8, 600 * MIB, 0), 1); - assert_eq!(cap_pool_from_memory(8, 0, 100 * MIB), 1); - assert_eq!(cap_pool_from_memory(8, 250 * MIB, 100 * MIB), 2); - assert_eq!(cap_pool_from_memory(1, 10_000 * MIB, 100 * MIB), 1); - } - - #[test] - fn default_pool_size_is_one() { - assert_eq!(default_pool_size(), 1); - } - - #[test] - fn default_intra_threads_uses_available_cpus() { - assert_eq!(default_intra_threads(), available_cpus()); - assert!(default_intra_threads() >= 1); - } } From 05a38d7f015462a52419f7b187a39eeca09ee702 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 16:25:45 +0000 Subject: [PATCH 09/11] fix: keep tokio workers free during model load Load ONNX sessions with block_in_place behind a tokio mutex so cold /embed cannot stall /health. Surface a failed load as 503 from /health and retry at most every 30s. Idle unload runs on a blocking thread; last-access uses a monotonic clock; the RAM cap is reused across reloads; jemalloc is the crate allocator so tests match the binary. Co-authored-by: Torsten Dittmann --- Cargo.lock | 1 - Cargo.toml | 3 +- README.md | 7 +- src/embedding.rs | 273 ++++++++++++++++++++++++++++++++++------------- src/lib.rs | 4 + src/main.rs | 25 ++--- 6 files changed, 221 insertions(+), 92 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7e5014c..c3da6e9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -592,7 +592,6 @@ dependencies = [ "serde", "serde_json", "sysinfo", - "tikv-jemalloc-sys", "tikv-jemallocator", "tokenizers", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 4f745b7..a4fae90 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,10 +12,9 @@ serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.145" sysinfo = "0.39.2" tokenizers = "0.22.2" -tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal", "time"] } +tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal", "sync", "time"] } tracing = "0.1.44" tracing-subscriber = { version = "0.3.20", features = ["env-filter"] } [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemallocator = { version = "0.6", features = ["override_allocator_on_supported_platforms"] } -tikv-jemalloc-sys = "0.6" diff --git a/README.md b/README.md index e50f1dc..a228a57 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,8 @@ docker compose up --build First request triggers the model download into `./models` (bind-mounted into the container); subsequent restarts reuse it. +`GET /health` is liveness: it stays `200` while unused so idle RSS is not pinned. After a failed model load it returns `503` with the error so a broken cache is visible without calling `/embed`. Failed loads are retried at most once every 30 seconds. + ```bash curl -X POST http://localhost:3000/embed \ -H 'content-type: application/json' \ @@ -23,7 +25,7 @@ Configured via environment variables (set them in `.env`): | Variable | Default | Description | | --- | --- | --- | | `EMBEDDING_PORT` | `3000` | Port the service listens on. | -| `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. | +| `EMBEDDING_MODELS` | `nomic` | Comma-separated list of models allowed to load. ONNX sessions are created on first `/embed`, not at process start. Image/compose defaults are nomic only; `bge-small` and other aliases still work if you add them here. | | `EMBEDDING_CACHE_DIR` | _(default cache)_ | Directory for downloaded model files. | | `EMBEDDING_POOL_SIZE` | `1` | Number of ONNX sessions per model while it is loaded, then capped by available RAM. Concurrent `/embed` calls round-robin across sessions. Raise this for parallel HTTP throughput; each extra session keeps another copy of the weights resident until idle unload. | | `EMBEDDING_INTRA_THREADS` | CPU count | ONNX Runtime intra-op threads per session. The default uses the whole machine on the single default session. When `EMBEDDING_POOL_SIZE` is greater than one, threads are split across sessions (`nproc / pool_size`, still capped by this value) so concurrent embeds do not oversubscribe the host. | @@ -51,5 +53,6 @@ Response: Errors: -- `400 Bad Request` — `texts` is empty. +- `400 Bad Request` — `texts` is empty or the model alias is not in `EMBEDDING_MODELS`. - `500 Internal Server Error` — embedding or tokenizer failure (message in `error` field). +- `503 Service Unavailable` — `GET /health` after a model load has failed (message in `error` field). diff --git a/src/embedding.rs b/src/embedding.rs index d69f417..605cb5b 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -8,11 +8,15 @@ use std::sync::{ Arc, Mutex, atomic::{AtomicU64, AtomicUsize, Ordering}, }; +use std::time::{Duration, Instant}; use tokenizers::Tokenizer; /// Unload a model this many seconds after last use. `0` disables unloading. pub const DEFAULT_IDLE_UNLOAD_SECS: u64 = 300; +/// After a failed load, skip another download/init for this long. +const LOAD_RETRY_SECS: u64 = 30; + #[derive(Debug, Clone)] pub struct EmbeddingResult { pub model: String, @@ -64,14 +68,11 @@ fn intra_threads_for_pool(configured: usize, pool_size: usize, nproc: usize) -> /// uncapped CPU-count pool. fn cap_pool_from_memory(desired: usize, budget: u64, per_instance_bytes: u64) -> usize { let desired = desired.max(1); - if per_instance_bytes == 0 { + if per_instance_bytes == 0 || budget == 0 { return 1; } - match budget.checked_div(per_instance_bytes) { - Some(0) => 1, - Some(max_from_memory) => desired.min(max_from_memory as usize).max(1), - None => 1, - } + let max_from_memory = (budget / per_instance_bytes) as usize; + desired.min(max_from_memory.max(1)) } fn memory_budget(host_available: u64, cgroup_free: Option) -> u64 { @@ -85,16 +86,27 @@ fn next_index(counter: &AtomicUsize, len: usize) -> usize { counter.fetch_add(1, Ordering::Relaxed) % len } -fn unix_now() -> u64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_secs()) - .unwrap_or(0) +fn process_origin() -> Instant { + static ORIGIN: std::sync::OnceLock = std::sync::OnceLock::new(); + *ORIGIN.get_or_init(Instant::now) +} + +fn mono_now_ms() -> u64 { + process_origin().elapsed().as_millis() as u64 +} + +fn touch_access(loaded: &LoadedModel) { + loaded + .last_access + .store(mono_now_ms().max(1), Ordering::Relaxed); } /// Whether an idle model slot should drop its ONNX sessions. +/// `last_access` and `now` are monotonic milliseconds; `idle_unload_secs` is seconds. pub(crate) fn should_unload(last_access: u64, now: u64, idle_unload_secs: u64) -> bool { - idle_unload_secs > 0 && last_access > 0 && now.saturating_sub(last_access) >= idle_unload_secs + idle_unload_secs > 0 + && last_access > 0 + && now.saturating_sub(last_access) >= idle_unload_secs.saturating_mul(1000) } fn parse_usize_env(key: &str) -> Option { @@ -158,11 +170,18 @@ struct LoadedModel { spec: EmbeddingModel, model_name: String, next: AtomicUsize, - load: Mutex<()>, + load: tokio::sync::Mutex<()>, inner: Mutex, dimension: usize, last_access: AtomicU64, in_flight: AtomicUsize, + last_failure: Mutex>, + cached_pool_size: AtomicUsize, +} + +struct LoadFailure { + message: String, + at: Instant, } struct ModelSlot { @@ -198,7 +217,7 @@ impl EmbeddingClient { spec: model.clone(), model_name, next: AtomicUsize::new(0), - load: Mutex::new(()), + load: tokio::sync::Mutex::new(()), inner: Mutex::new(ModelSlot { pool: None, tokenizer: None, @@ -206,6 +225,8 @@ impl EmbeddingClient { dimension: model::dimension(model), last_access: AtomicU64::new(0), in_flight: AtomicUsize::new(0), + last_failure: Mutex::new(None), + cached_pool_size: AtomicUsize::new(0), }, ); } @@ -224,18 +245,31 @@ impl EmbeddingClient { /// builds still populate the ONNX cache. pub fn preload(&self) -> Result<(), String> { for loaded in self.models.values() { - self.ensure_loaded(loaded)?; - loaded.last_access.store(unix_now(), Ordering::Relaxed); + self.ensure_loaded_blocking(loaded)?; + touch_access(loaded); } Ok(()) } + /// Most recent model-load failure, if any. Used by `/health` so a stuck + /// cache or unreadable `EMBEDDING_CACHE_DIR` is visible without `/embed`. + pub fn last_load_error(&self) -> Option { + for loaded in self.models.values() { + if let Ok(guard) = loaded.last_failure.lock() + && let Some(fail) = guard.as_ref() + { + return Some(fail.message.clone()); + } + } + None + } + /// Drop ONNX sessions that have been unused for `idle_unload_secs`. pub fn unload_idle(&self) { if self.config.idle_unload_secs == 0 { return; } - let now = unix_now(); + let now = mono_now_ms(); for loaded in self.models.values() { let last = loaded.last_access.load(Ordering::Relaxed); if !should_unload(last, now, self.config.idle_unload_secs) { @@ -255,67 +289,140 @@ impl EmbeddingClient { slot.tokenizer = None; tracing::info!( model = loaded.model_name.as_str(), - idle_secs = now.saturating_sub(last), + idle_ms = now.saturating_sub(last), "unloaded idle embedding model" ); } } } - pub fn idle_unload_secs(&self) -> u64 { - self.config.idle_unload_secs + fn slot_is_loaded(loaded: &LoadedModel) -> Result { + let slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + Ok(slot.pool.is_some() && slot.tokenizer.is_some()) } - fn ensure_loaded(&self, loaded: &LoadedModel) -> Result<(), String> { - { - let slot = loaded - .inner - .lock() - .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; - if slot.pool.is_some() && slot.tokenizer.is_some() { - return Ok(()); - } + fn cached_load_error(loaded: &LoadedModel) -> Option { + let guard = loaded.last_failure.lock().ok()?; + let fail = guard.as_ref()?; + if fail.at.elapsed() < Duration::from_secs(LOAD_RETRY_SECS) { + Some(fail.message.clone()) + } else { + None } + } - let _load = loaded - .load - .lock() - .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; - { - let slot = loaded - .inner - .lock() - .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; - if slot.pool.is_some() && slot.tokenizer.is_some() { - return Ok(()); + fn record_load_failure(loaded: &LoadedModel, message: String) { + if let Ok(mut guard) = loaded.last_failure.lock() { + *guard = Some(LoadFailure { + message, + at: Instant::now(), + }); + } + } + + fn clear_load_failure(loaded: &LoadedModel) { + if let Ok(mut guard) = loaded.last_failure.lock() { + *guard = None; + } + } + + fn finish_load(&self, loaded: &LoadedModel) -> Result<(), String> { + let cached_pool_size = loaded.cached_pool_size.load(Ordering::Relaxed); + match Self::load_model( + &loaded.spec, + &loaded.model_name, + &self.config, + cached_pool_size, + ) { + Ok(built) => { + loaded + .cached_pool_size + .store(built.pool.len(), Ordering::Relaxed); + let mut slot = loaded + .inner + .lock() + .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; + slot.pool = Some(built.pool); + slot.tokenizer = Some(built.tokenizer); + Self::clear_load_failure(loaded); + Ok(()) + } + Err(err) => { + Self::record_load_failure(loaded, err.clone()); + Err(err) } } + } - let built = Self::load_model(&loaded.spec, &loaded.model_name, &self.config)?; + async fn ensure_loaded(&self, loaded: &LoadedModel) -> Result<(), String> { + if Self::slot_is_loaded(loaded)? { + return Ok(()); + } + if let Some(err) = Self::cached_load_error(loaded) { + return Err(err); + } + let _load = loaded.load.lock().await; + if Self::slot_is_loaded(loaded)? { + return Ok(()); + } + if let Some(err) = Self::cached_load_error(loaded) { + return Err(err); + } + tokio::task::block_in_place(|| self.finish_load(loaded)) + } - let mut slot = loaded - .inner - .lock() - .map_err(|e| format!("Embedding model lock poisoned: {}", e))?; - slot.pool = Some(built.pool); - slot.tokenizer = Some(built.tokenizer); - Ok(()) + fn ensure_loaded_blocking(&self, loaded: &LoadedModel) -> Result<(), String> { + if Self::slot_is_loaded(loaded)? { + return Ok(()); + } + if let Some(err) = Self::cached_load_error(loaded) { + return Err(err); + } + let _load = loaded.load.blocking_lock(); + if Self::slot_is_loaded(loaded)? { + return Ok(()); + } + if let Some(err) = Self::cached_load_error(loaded) { + return Err(err); + } + self.finish_load(loaded) } fn load_model( model: &EmbeddingModel, model_name: &str, config: &EmbeddingConfig, + cached_pool_size: usize, ) -> Result { let desired_pool_size = config.pool_size.max(1); let dimension = model::dimension(model); let has_gpu_providers = !config.execution_providers.is_empty(); let nproc = available_cpus(); - let probe_intra = intra_threads_for_pool(config.intra_threads, desired_pool_size, nproc); - let mut probe_config = config.clone(); - probe_config.intra_threads = probe_intra; - let mem_before_loading_model = if desired_pool_size > 1 { + let known_pool_size = if desired_pool_size == 1 { + Some(1) + } else if cached_pool_size > 0 { + let reused = cached_pool_size.min(desired_pool_size).max(1); + if reused != cached_pool_size { + tracing::warn!( + model = model_name, + previous = cached_pool_size, + new = reused, + "reusing capped pool size from the first load" + ); + } + Some(reused) + } else { + None + }; + + let probe_size = known_pool_size.unwrap_or(desired_pool_size); + let probe_intra = intra_threads_for_pool(config.intra_threads, probe_size, nproc); + + let mem_before_loading_model = if known_pool_size.is_none() { let mut sys = sysinfo::System::new(); sys.refresh_memory(); Some(memory_budget( @@ -326,13 +433,21 @@ impl EmbeddingClient { None }; - let first_model = Self::init_and_warmup(model, &probe_config, model_name)?; + let first_model = Self::init_and_warmup(model, config, probe_intra, model_name)?; // Tokenizer is fetched from the same cache dir fastembed just populated, // so this is a cache hit (no network) after the first model load. let tokenizer = Arc::new(Self::load_tokenizer(model, config)?); - let pool_size = if let Some(mem_before_loading_model) = mem_before_loading_model { + let pool_size = if let Some(known) = known_pool_size { + if known == 1 { + tracing::info!( + model = model_name, + "Using pool_size=1; extra sessions are not created" + ); + } + known + } else if let Some(mem_before_loading_model) = mem_before_loading_model { let mut sys = sysinfo::System::new(); sys.refresh_memory(); let memory_after_loading_model = memory_budget( @@ -379,30 +494,30 @@ impl EmbeddingClient { ); capped } else { - tracing::info!( - model = model_name, - "Using pool_size=1; extra sessions are not created" - ); 1 }; let session_intra = intra_threads_for_pool(config.intra_threads, pool_size, nproc); - let mut session_config = config.clone(); - session_config.intra_threads = session_intra; let mut pool = Vec::with_capacity(pool_size); if session_intra == probe_intra { pool.push(Arc::new(Mutex::new(first_model))); } else { + // RAM cap changed the pool size, so the probe session was built + // with the wrong intra-op thread count. Drop it and rebuild the + // serving session (with warmup) at the final count. The probe + // measured a session with fewer threads than serving will use, + // so the cap can underestimate RSS. drop(first_model); pool.push(Arc::new(Mutex::new(Self::init_and_warmup( model, - &session_config, + config, + session_intra, model_name, )?))); } for _ in 1..pool_size { - let inst = Self::init_model(model, &session_config)?; + let inst = Self::init_model(model, config, session_intra)?; pool.push(Arc::new(Mutex::new(inst))); } @@ -433,9 +548,10 @@ impl EmbeddingClient { fn init_and_warmup( model: &EmbeddingModel, config: &EmbeddingConfig, + intra_threads: usize, model_name: &str, ) -> Result { - let mut inst = Self::init_model(model, config)?; + let mut inst = Self::init_model(model, config, intra_threads)?; Self::warmup_session(&mut inst, model_name)?; Ok(inst) } @@ -466,10 +582,11 @@ impl EmbeddingClient { fn init_model( model: &EmbeddingModel, config: &EmbeddingConfig, + intra_threads: usize, ) -> Result { let mut init_options = InitOptions::new(model.clone()) .with_show_download_progress(config.show_download_progress) - .with_intra_threads(config.intra_threads); + .with_intra_threads(intra_threads); if let Some(cache_dir) = &config.cache_dir { init_options = init_options.with_cache_dir(cache_dir.into()); @@ -484,9 +601,12 @@ impl EmbeddingClient { .map_err(|e| format!("Failed to initialize embedding model: {}", e)) } - fn acquire_instance(&self, loaded: &LoadedModel) -> Result>, String> { + async fn acquire_instance( + &self, + loaded: &LoadedModel, + ) -> Result>, String> { for _ in 0..2 { - self.ensure_loaded(loaded)?; + self.ensure_loaded(loaded).await?; let slot = loaded .inner .lock() @@ -502,8 +622,8 @@ impl EmbeddingClient { )) } - fn tokenizer(&self, loaded: &LoadedModel) -> Result, String> { - self.ensure_loaded(loaded)?; + async fn tokenizer(&self, loaded: &LoadedModel) -> Result, String> { + self.ensure_loaded(loaded).await?; let slot = loaded .inner .lock() @@ -527,7 +647,7 @@ impl EmbeddingClient { EmbedError::UnknownModel(format!("model not allowed: {}", model_name)) })?; - loaded.last_access.store(unix_now(), Ordering::Relaxed); + touch_access(loaded); loaded.in_flight.fetch_add(1, Ordering::SeqCst); let _in_flight = InFlightGuard { counter: &loaded.in_flight, @@ -545,9 +665,12 @@ impl EmbeddingClient { Self::compute_sub_batch(available_mb, loaded.dimension, self.gpu) }; + // Each chunk runs in spawn_blocking so ORT inference does not occupy + // a tokio worker. At pool_size=1 every chunk still serializes on the + // same session mutex; extra tasks are then just overhead. let mut handles = Vec::new(); for chunk in texts.chunks(sub_batch) { - let inst = self.acquire_instance(loaded)?; + let inst = self.acquire_instance(loaded).await?; let chunked_texts: Vec = chunk.iter().map(|t| (*t).to_owned()).collect(); handles.push(tokio::task::spawn_blocking(move || { @@ -567,7 +690,7 @@ impl EmbeddingClient { embeddings.append(&mut batch_result); } - let tokenizer = self.tokenizer(loaded)?; + let tokenizer = self.tokenizer(loaded).await?; let owned_texts: Vec = texts.iter().map(|t| t.to_string()).collect(); let tokens = tokio::task::spawn_blocking(move || -> Result { let encodings = tokenizer @@ -578,7 +701,7 @@ impl EmbeddingClient { .await .map_err(|e| format!("Failed to join tokenizer task: {}", e))??; - loaded.last_access.store(unix_now(), Ordering::Relaxed); + touch_access(loaded); Ok(EmbeddingResult { model: loaded.model_name.clone(), @@ -812,9 +935,9 @@ mod tests { #[test] fn should_unload_requires_prior_use_and_timeout() { assert!(!should_unload(0, 1_000, 300)); - assert!(!should_unload(900, 1_000, 0)); - assert!(!should_unload(800, 1_000, 300)); - assert!(should_unload(700, 1_000, 300)); + assert!(!should_unload(900_000, 1_000_000, 0)); + assert!(!should_unload(800_000, 1_000_000, 300)); + assert!(should_unload(700_000, 1_000_000, 300)); } #[test] diff --git a/src/lib.rs b/src/lib.rs index 6dba6f3..fb5386b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,7 @@ +#[cfg(not(target_env = "msvc"))] +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + mod embedding; mod error; mod model; diff --git a/src/main.rs b/src/main.rs index fcd47bb..15b7c4f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,6 @@ use std::sync::Arc; use std::time::Duration; -#[cfg(not(target_env = "msvc"))] -#[global_allocator] -static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; - -// Make jemalloc override libc malloc so ONNX Runtime's C++ allocations -// (the bulk of idle RSS) go through jemalloc and can be returned to the OS. -#[cfg(not(target_env = "msvc"))] -use tikv_jemalloc_sys as _; - use axum::{ Json, Router, extract::State, @@ -48,8 +39,15 @@ impl IntoResponse for AppError { } } -async fn health() -> impl IntoResponse { - StatusCode::OK +async fn health(State(state): State) -> Response { + if let Some(err) = state.client.last_load_error() { + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ "error": err })), + ) + .into_response(); + } + StatusCode::OK.into_response() } async fn embed( @@ -103,7 +101,10 @@ async fn main() -> Result<(), Box> { ticker.tick().await; // skip the immediate first tick loop { ticker.tick().await; - client_bg.unload_idle(); + let client = client_bg.clone(); + if let Err(err) = tokio::task::spawn_blocking(move || client.unload_idle()).await { + tracing::warn!(error = %err, "idle unload task failed"); + } } }); tracing::info!(idle_unload_secs, tick_secs, "idle model unload enabled"); From dad9f7a2fe8a6ee53130a5014dbfe787eb8ee2c0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 16:31:59 +0000 Subject: [PATCH 10/11] fix: expire health 503 and load models on spawn_blocking /health 503 now lasts only for the 30s retry window so a liveness probe can recover; the next /embed retries the load. Model init uses spawn_blocking instead of block_in_place so current-thread runtimes do not panic. Co-authored-by: Torsten Dittmann --- README.md | 4 ++-- src/embedding.rs | 45 ++++++++++++++++++++------------------------- 2 files changed, 22 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index a228a57..2413639 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ docker compose up --build First request triggers the model download into `./models` (bind-mounted into the container); subsequent restarts reuse it. -`GET /health` is liveness: it stays `200` while unused so idle RSS is not pinned. After a failed model load it returns `503` with the error so a broken cache is visible without calling `/embed`. Failed loads are retried at most once every 30 seconds. +`GET /health` is liveness: it stays `200` while unused so idle RSS is not pinned. After a failed model load it returns `503` with the error for 30 seconds, then `200` again so a probe can recover; the next `/embed` retries the load. ```bash curl -X POST http://localhost:3000/embed \ @@ -55,4 +55,4 @@ Errors: - `400 Bad Request` — `texts` is empty or the model alias is not in `EMBEDDING_MODELS`. - `500 Internal Server Error` — embedding or tokenizer failure (message in `error` field). -- `503 Service Unavailable` — `GET /health` after a model load has failed (message in `error` field). +- `503 Service Unavailable` — `GET /health` for 30 seconds after a model load has failed (message in `error` field). diff --git a/src/embedding.rs b/src/embedding.rs index 605cb5b..e1571d3 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -160,7 +160,7 @@ impl EmbeddingConfig { } pub struct EmbeddingClient { - models: HashMap, + models: HashMap>, config: EmbeddingConfig, sub_batch_override: usize, gpu: bool, @@ -213,7 +213,7 @@ impl EmbeddingClient { let model_name = format!("{:?}", model); models.insert( model_name.clone(), - LoadedModel { + Arc::new(LoadedModel { spec: model.clone(), model_name, next: AtomicUsize::new(0), @@ -227,7 +227,7 @@ impl EmbeddingClient { in_flight: AtomicUsize::new(0), last_failure: Mutex::new(None), cached_pool_size: AtomicUsize::new(0), - }, + }), ); } @@ -251,17 +251,13 @@ impl EmbeddingClient { Ok(()) } - /// Most recent model-load failure, if any. Used by `/health` so a stuck - /// cache or unreadable `EMBEDDING_CACHE_DIR` is visible without `/embed`. + /// Most recent in-window model-load failure, if any. Used by `/health`. + /// The error expires with the 30s retry window so a liveness probe can + /// recover; the next `/embed` then retries the load. pub fn last_load_error(&self) -> Option { - for loaded in self.models.values() { - if let Ok(guard) = loaded.last_failure.lock() - && let Some(fail) = guard.as_ref() - { - return Some(fail.message.clone()); - } - } - None + self.models + .values() + .find_map(|loaded| Self::cached_load_error(loaded)) } /// Drop ONNX sessions that have been unused for `idle_unload_secs`. @@ -329,14 +325,9 @@ impl EmbeddingClient { } } - fn finish_load(&self, loaded: &LoadedModel) -> Result<(), String> { + fn finish_load(loaded: &LoadedModel, config: &EmbeddingConfig) -> Result<(), String> { let cached_pool_size = loaded.cached_pool_size.load(Ordering::Relaxed); - match Self::load_model( - &loaded.spec, - &loaded.model_name, - &self.config, - cached_pool_size, - ) { + match Self::load_model(&loaded.spec, &loaded.model_name, config, cached_pool_size) { Ok(built) => { loaded .cached_pool_size @@ -357,7 +348,7 @@ impl EmbeddingClient { } } - async fn ensure_loaded(&self, loaded: &LoadedModel) -> Result<(), String> { + async fn ensure_loaded(&self, loaded: &Arc) -> Result<(), String> { if Self::slot_is_loaded(loaded)? { return Ok(()); } @@ -371,7 +362,11 @@ impl EmbeddingClient { if let Some(err) = Self::cached_load_error(loaded) { return Err(err); } - tokio::task::block_in_place(|| self.finish_load(loaded)) + let loaded = Arc::clone(loaded); + let config = self.config.clone(); + tokio::task::spawn_blocking(move || Self::finish_load(&loaded, &config)) + .await + .map_err(|e| format!("Failed to join model load: {}", e))? } fn ensure_loaded_blocking(&self, loaded: &LoadedModel) -> Result<(), String> { @@ -388,7 +383,7 @@ impl EmbeddingClient { if let Some(err) = Self::cached_load_error(loaded) { return Err(err); } - self.finish_load(loaded) + Self::finish_load(loaded, &self.config) } fn load_model( @@ -603,7 +598,7 @@ impl EmbeddingClient { async fn acquire_instance( &self, - loaded: &LoadedModel, + loaded: &Arc, ) -> Result>, String> { for _ in 0..2 { self.ensure_loaded(loaded).await?; @@ -622,7 +617,7 @@ impl EmbeddingClient { )) } - async fn tokenizer(&self, loaded: &LoadedModel) -> Result, String> { + async fn tokenizer(&self, loaded: &Arc) -> Result, String> { self.ensure_loaded(loaded).await?; let slot = loaded .inner From 8ac60069948dd04c8c5126d053f6cc7087205fa3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 16 Sep 2026 16:39:02 +0000 Subject: [PATCH 11/11] fix: keep model-load ownership if /embed is cancelled Acquire the load mutex inside spawn_blocking so a cancelled cold /embed cannot drop ownership while finish_load is still running. Co-authored-by: Torsten Dittmann --- Cargo.toml | 2 +- src/embedding.rs | 26 ++++++++++++-------------- 2 files changed, 13 insertions(+), 15 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index a4fae90..b95d91c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.145" sysinfo = "0.39.2" tokenizers = "0.22.2" -tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal", "sync", "time"] } +tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread", "net", "signal", "time"] } tracing = "0.1.44" tracing-subscriber = { version = "0.3.20", features = ["env-filter"] } diff --git a/src/embedding.rs b/src/embedding.rs index e1571d3..87fde29 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -170,7 +170,7 @@ struct LoadedModel { spec: EmbeddingModel, model_name: String, next: AtomicUsize, - load: tokio::sync::Mutex<()>, + load: Mutex<()>, inner: Mutex, dimension: usize, last_access: AtomicU64, @@ -217,7 +217,7 @@ impl EmbeddingClient { spec: model.clone(), model_name, next: AtomicUsize::new(0), - load: tokio::sync::Mutex::new(()), + load: Mutex::new(()), inner: Mutex::new(ModelSlot { pool: None, tokenizer: None, @@ -245,7 +245,7 @@ impl EmbeddingClient { /// builds still populate the ONNX cache. pub fn preload(&self) -> Result<(), String> { for loaded in self.models.values() { - self.ensure_loaded_blocking(loaded)?; + Self::ensure_loaded_blocking(loaded, &self.config)?; touch_access(loaded); } Ok(()) @@ -355,35 +355,33 @@ impl EmbeddingClient { if let Some(err) = Self::cached_load_error(loaded) { return Err(err); } - let _load = loaded.load.lock().await; - if Self::slot_is_loaded(loaded)? { - return Ok(()); - } - if let Some(err) = Self::cached_load_error(loaded) { - return Err(err); - } + // The std mutex is acquired inside spawn_blocking so cancelling this + // request cannot drop load ownership while finish_load is still running. let loaded = Arc::clone(loaded); let config = self.config.clone(); - tokio::task::spawn_blocking(move || Self::finish_load(&loaded, &config)) + tokio::task::spawn_blocking(move || Self::ensure_loaded_blocking(&loaded, &config)) .await .map_err(|e| format!("Failed to join model load: {}", e))? } - fn ensure_loaded_blocking(&self, loaded: &LoadedModel) -> Result<(), String> { + fn ensure_loaded_blocking( + loaded: &LoadedModel, + config: &EmbeddingConfig, + ) -> Result<(), String> { if Self::slot_is_loaded(loaded)? { return Ok(()); } if let Some(err) = Self::cached_load_error(loaded) { return Err(err); } - let _load = loaded.load.blocking_lock(); + let _load = loaded.load.lock().unwrap_or_else(|e| e.into_inner()); if Self::slot_is_loaded(loaded)? { return Ok(()); } if let Some(err) = Self::cached_load_error(loaded) { return Err(err); } - Self::finish_load(loaded, &self.config) + Self::finish_load(loaded, config) } fn load_model(