diff --git a/.env b/.env index 2cbcba1..07c1fd6 100644 --- a/.env +++ b/.env @@ -1,3 +1,2 @@ -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..b95d91c 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 = { version = "0.6", features = ["override_allocator_on_supported_platforms"] } diff --git a/Dockerfile b/Dockerfile index 2160251..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 . . @@ -34,7 +35,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..2413639 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 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 \ -H 'content-type: application/json' \ @@ -23,9 +25,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. 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` | _(memory-derived)_ | Number of model instances per pool. | +| `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. | +| `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 @@ -49,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` for 30 seconds after a model load has failed (message in `error` field). 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..87fde29 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -6,10 +6,17 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::{ Arc, Mutex, - atomic::{AtomicUsize, Ordering}, + 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, @@ -26,14 +33,48 @@ 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_pool_size() -> usize { + 1 +} + +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 { + let desired = desired.max(1); + if per_instance_bytes == 0 || budget == 0 { + return 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 { match cgroup_free { Some(cgroup_free) => host_available.min(cgroup_free), @@ -45,6 +86,35 @@ fn next_index(counter: &AtomicUsize, len: usize) -> usize { counter.fetch_add(1, Ordering::Relaxed) % len } +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.saturating_mul(1000) +} + +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,12 +133,19 @@ 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_else(default_pool_size); + let intra_threads = parse_usize_env("EMBEDDING_INTRA_THREADS") + .filter(|&n| n >= 1) + .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, show_download_progress: true, @@ -76,109 +153,319 @@ impl EmbeddingConfig { pool_size, execution_providers: Vec::new(), sub_batch_size: 0, + intra_threads, + idle_unload_secs, } } } pub struct EmbeddingClient { - models: HashMap, + 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, + in_flight: AtomicUsize, + last_failure: Mutex>, + cached_pool_size: AtomicUsize, +} + +struct LoadFailure { + message: String, + at: Instant, +} + +struct ModelSlot { + pool: Option>>>, + 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, } 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(), + Arc::new(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), + in_flight: AtomicUsize::new(0), + last_failure: Mutex::new(None), + cached_pool_size: AtomicUsize::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_blocking(loaded, &self.config)?; + touch_access(loaded); + } + Ok(()) + } + + /// 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 { + self.models + .values() + .find_map(|loaded| Self::cached_load_error(loaded)) + } + + /// 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 = 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) { + continue; + } + let mut slot = match loaded.inner.lock() { + 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; + tracing::info!( + model = loaded.model_name.as_str(), + idle_ms = now.saturating_sub(last), + "unloaded idle embedding model" + ); + } + } + } + + 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 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 + } + } + + 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(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, 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) + } + } + } + + async fn ensure_loaded(&self, loaded: &Arc) -> Result<(), String> { + 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::ensure_loaded_blocking(&loaded, &config)) + .await + .map_err(|e| format!("Failed to join model load: {}", e))? + } + + 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.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, config) + } + fn load_model( model: &EmbeddingModel, model_name: &str, config: &EmbeddingConfig, - ) -> Result { + cached_pool_size: usize, + ) -> 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 nproc = available_cpus(); + + 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( + sys.available_memory(), + sys.cgroup_limits().map(|limits| limits.free_memory), + )) + } else { + None + }; - let mut first_model = Self::init_model(model, config)?; + 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)?); - // 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 { + 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( + 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 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 { + 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), budget_mb = budget / (1024 * 1024), @@ -188,8 +475,6 @@ impl EmbeddingClient { 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), @@ -197,21 +482,35 @@ impl EmbeddingClient { budget_mb = budget / (1024 * 1024), nproc = nproc, desired = desired_pool_size, - max_from_memory = max_memory, capped = capped, "Measured ONNX model memory footprint" ); 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))); + let session_intra = intra_threads_for_pool(config.intra_threads, pool_size, nproc); + 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, + config, + session_intra, + model_name, + )?))); + } for _ in 1..pool_size { - let inst = Self::init_model(model, config)?; + let inst = Self::init_model(model, config, session_intra)?; pool.push(Arc::new(Mutex::new(inst))); } @@ -221,20 +520,33 @@ 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, + session_intra, ep_label, ); - Ok(LoadedModel { - dimension, - model_name: model_name.to_string(), - next: AtomicUsize::new(0), - pool, - tokenizer, - }) + 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, + intra_threads: usize, + model_name: &str, + ) -> Result { + let mut inst = Self::init_model(model, config, intra_threads)?; + Self::warmup_session(&mut inst, model_name)?; + Ok(inst) } fn load_tokenizer( @@ -263,9 +575,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_show_download_progress(config.show_download_progress) + .with_intra_threads(intra_threads); if let Some(cache_dir) = &config.cache_dir { init_options = init_options.with_cache_dir(cache_dir.into()); @@ -279,6 +593,39 @@ impl EmbeddingClient { TextEmbedding::try_new(init_options) .map_err(|e| format!("Failed to initialize embedding model: {}", e)) } + + async fn acquire_instance( + &self, + loaded: &Arc, + ) -> Result>, String> { + for _ in 0..2 { + self.ensure_loaded(loaded).await?; + 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 + )) + } + + async fn tokenizer(&self, loaded: &Arc) -> Result, String> { + self.ensure_loaded(loaded).await?; + 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 +640,12 @@ impl EmbeddingClient { EmbedError::UnknownModel(format!("model not allowed: {}", model_name)) })?; + touch_access(loaded); + 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 } else { @@ -305,9 +658,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(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 || { @@ -327,7 +683,7 @@ impl EmbeddingClient { embeddings.append(&mut batch_result); } - let tokenizer = loaded.tokenizer.clone(); + 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 @@ -338,6 +694,8 @@ impl EmbeddingClient { .await .map_err(|e| format!("Failed to join tokenizer task: {}", e))??; + touch_access(loaded); + Ok(EmbeddingResult { model: loaded.model_name.clone(), embeddings, @@ -346,12 +704,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 +730,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 +793,12 @@ 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, 1); assert!(cfg.show_download_progress); assert!(cfg.execution_providers.is_empty()); assert_eq!(cfg.sub_batch_size, 0); + assert_eq!(cfg.intra_threads, available_cpus()); + assert_eq!(cfg.idle_unload_secs, DEFAULT_IDLE_UNLOAD_SECS); } #[test] @@ -539,7 +895,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, 1); } #[test] @@ -547,7 +903,34 @@ 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, 1); + } + + #[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 should_unload_requires_prior_use_and_timeout() { + assert!(!should_unload(0, 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] @@ -588,9 +971,4 @@ mod tests { assert_eq!(next_index(&counter, 3), 0); assert_eq!(next_index(&counter, 3), 1); } - - #[test] - fn default_pool_size_is_positive() { - assert!(default_pool_size() >= 1); - } } 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 e54e130..15b7c4f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::time::Duration; use axum::{ Json, Router, @@ -38,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( @@ -81,7 +89,27 @@ 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; + 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"); + } + let state = AppState { client }; let app = Router::new() diff --git a/tests/embed_e2e.rs b/tests/embed_e2e.rs index ea92023..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 { @@ -8,6 +10,32 @@ 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, + } +} + +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) } } @@ -54,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}" + ); +}