From e6fa98d6993add0a4fd74332a35de3dbff46aa33 Mon Sep 17 00:00:00 2001 From: Chirag Aggarwal Date: Fri, 4 Sep 2026 16:44:25 +0100 Subject: [PATCH 1/6] fix: size the model pool from the cgroup memory limit, not the host Inside a container sysinfo's available_memory() reports the host, so a service started with --memory=1g on a large machine sized its pool for the whole machine, loaded several instances, was OOM-killed and restart-looped at full CPU. Take the smaller of host-available and cgroup-free memory for both the pool and sub-batch budgets, and warn when even a single instance is estimated to exceed the budget. --- src/embedding.rs | 54 +++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 51 insertions(+), 3 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index 22f0208..4049ee8 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -34,6 +34,27 @@ fn default_pool_size() -> usize { .unwrap_or(2) } +/// Memory this process can still allocate before the kernel OOM-kills it. +/// +/// `System::available_memory()` reports the host, which inside a container is +/// the wrong number: a service started with `--memory=1g` on a 32 GB machine +/// would size its pool for 32 GB, load several model instances, get killed, +/// and restart in a loop. When a cgroup limit applies, the cgroup's free +/// memory is the real ceiling, so take the smaller of the two. +fn available_memory(sys: &sysinfo::System) -> u64 { + memory_budget( + sys.available_memory(), + sys.cgroup_limits().map(|limits| limits.free_memory), + ) +} + +fn memory_budget(host_available: u64, cgroup_free: Option) -> u64 { + match cgroup_free { + Some(cgroup_free) => host_available.min(cgroup_free), + None => host_available, + } +} + fn next_index(counter: &AtomicUsize, len: usize) -> usize { counter.fetch_add(1, Ordering::Relaxed) % len } @@ -129,7 +150,7 @@ impl EmbeddingClient { // loading instance and measuring memory footprint let mut sys = sysinfo::System::new(); sys.refresh_memory(); - let mem_before_loading_model = sys.available_memory(); + let mem_before_loading_model = available_memory(&sys); let has_gpu_providers = !config.execution_providers.is_empty(); @@ -149,7 +170,7 @@ impl EmbeddingClient { .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; sys.refresh_memory(); - let memory_after_loading_model = sys.available_memory(); + let memory_after_loading_model = available_memory(&sys); let per_instance_loaded = mem_before_loading_model.saturating_sub(memory_after_loading_model); @@ -165,6 +186,16 @@ impl EmbeddingClient { // 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!( + 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!( @@ -333,7 +364,7 @@ impl EmbeddingClient { fn compute_sub_batch(dimension: usize, gpu: bool) -> usize { let mut sys = sysinfo::System::new(); sys.refresh_memory(); - let available_mb = sys.available_memory() / (1024 * 1024); + let available_mb = available_memory(&sys) / (1024 * 1024); if available_mb == 0 { return 32; @@ -405,6 +436,23 @@ mod tests { } } + const GIB: u64 = 1024 * 1024 * 1024; + + #[test] + fn memory_budget_uses_host_without_cgroup() { + assert_eq!(memory_budget(32 * GIB, None), 32 * GIB); + } + + #[test] + fn memory_budget_is_capped_by_cgroup_limit() { + assert_eq!(memory_budget(32 * GIB, Some(GIB)), GIB); + } + + #[test] + fn memory_budget_is_capped_by_host_when_cgroup_is_unlimited() { + assert_eq!(memory_budget(2 * GIB, Some(32 * GIB)), 2 * GIB); + } + #[test] fn from_env_uses_nomic_when_unset() { let _g = isolate_env(); From 07a632a66769b3505872908316637ebee93dc74d Mon Sep 17 00:00:00 2001 From: Chirag Aggarwal Date: Fri, 4 Sep 2026 16:48:36 +0100 Subject: [PATCH 2/6] refactor: inline the memory budget at its call sites --- src/embedding.rs | 29 ++++++++++++----------------- 1 file changed, 12 insertions(+), 17 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index 4049ee8..f98b21f 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -34,20 +34,6 @@ fn default_pool_size() -> usize { .unwrap_or(2) } -/// Memory this process can still allocate before the kernel OOM-kills it. -/// -/// `System::available_memory()` reports the host, which inside a container is -/// the wrong number: a service started with `--memory=1g` on a 32 GB machine -/// would size its pool for 32 GB, load several model instances, get killed, -/// and restart in a loop. When a cgroup limit applies, the cgroup's free -/// memory is the real ceiling, so take the smaller of the two. -fn available_memory(sys: &sysinfo::System) -> u64 { - memory_budget( - sys.available_memory(), - sys.cgroup_limits().map(|limits| limits.free_memory), - ) -} - fn memory_budget(host_available: u64, cgroup_free: Option) -> u64 { match cgroup_free { Some(cgroup_free) => host_available.min(cgroup_free), @@ -150,7 +136,10 @@ impl EmbeddingClient { // loading instance and measuring memory footprint let mut sys = sysinfo::System::new(); sys.refresh_memory(); - let mem_before_loading_model = available_memory(&sys); + 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(); @@ -170,7 +159,10 @@ impl EmbeddingClient { .map_err(|e| format!("warmup inference failed for {}: {}", model_name, e))?; sys.refresh_memory(); - let memory_after_loading_model = available_memory(&sys); + 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); @@ -364,7 +356,10 @@ impl EmbeddingClient { fn compute_sub_batch(dimension: usize, gpu: bool) -> usize { let mut sys = sysinfo::System::new(); sys.refresh_memory(); - let available_mb = available_memory(&sys) / (1024 * 1024); + let available_mb = memory_budget( + sys.available_memory(), + sys.cgroup_limits().map(|limits| limits.free_memory), + ) / (1024 * 1024); if available_mb == 0 { return 32; From 91f5f045436e762225da05e693aa25f3eed79563 Mon Sep 17 00:00:00 2001 From: Chirag Aggarwal Date: Fri, 4 Sep 2026 16:50:41 +0100 Subject: [PATCH 3/6] fix: let the sub-batch shrink to one text on small memory budgets --- src/embedding.rs | 47 +++++++++++++++++++++++++++++++++++------------ 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index f98b21f..dd9271c 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -361,20 +361,24 @@ impl EmbeddingClient { sys.cgroup_limits().map(|limits| limits.free_memory), ) / (1024 * 1024); - if available_mb == 0 { - return 32; - } + sub_batch_for(available_mb, dimension, gpu) + } +} - // Per-text memory estimate for ONNX inference. Attention matrices - // dominate: heads × seq² × 4 bytes. For 768-dim BERT-like models - // (12 heads) processing ~1000-2000 token code chunks, attention alone - // is 50-200 MB per text. The estimate below is conservative so the - // sub-batch stays small enough to prevent arena over-allocation. - let mb_per_text: u64 = if dimension >= 768 { 100 } else { 40 }; - let budget_mb = available_mb / 2; - let max_batch = if gpu { 256 } else { 16 }; - (budget_mb / mb_per_text).clamp(4, max_batch) as usize +fn sub_batch_for(available_mb: u64, dimension: usize, gpu: bool) -> usize { + if available_mb == 0 { + return 32; } + + // Per-text memory estimate for ONNX inference. Attention matrices + // dominate: heads × seq² × 4 bytes. For 768-dim BERT-like models + // (12 heads) processing ~1000-2000 token code chunks, attention alone + // is 50-200 MB per text. The estimate below is conservative so the + // sub-batch stays small enough to prevent arena over-allocation. + let mb_per_text: u64 = if dimension >= 768 { 100 } else { 40 }; + let budget_mb = available_mb / 2; + let max_batch = if gpu { 256 } else { 16 }; + (budget_mb / mb_per_text).clamp(1, max_batch) as usize } #[cfg(test)] @@ -448,6 +452,25 @@ mod tests { assert_eq!(memory_budget(2 * GIB, Some(32 * GIB)), 2 * GIB); } + #[test] + fn sub_batch_falls_back_when_memory_is_unknown() { + assert_eq!(sub_batch_for(0, 768, false), 32); + } + + #[test] + fn sub_batch_never_exceeds_a_small_budget() { + // 200 MiB free -> 100 MiB budget -> one 100 MiB text, not four. + assert_eq!(sub_batch_for(200, 768, false), 1); + assert_eq!(sub_batch_for(1, 768, false), 1); + assert_eq!(sub_batch_for(200, 384, false), 2); + } + + #[test] + fn sub_batch_is_capped_by_device() { + assert_eq!(sub_batch_for(64 * 1024, 768, false), 16); + assert_eq!(sub_batch_for(64 * 1024, 768, true), 256); + } + #[test] fn from_env_uses_nomic_when_unset() { let _g = isolate_env(); From 89573d386eb157b0d24a6eae5a788cd037142f86 Mon Sep 17 00:00:00 2001 From: Chirag Aggarwal Date: Fri, 4 Sep 2026 16:53:27 +0100 Subject: [PATCH 4/6] refactor: pass available memory into compute_sub_batch --- src/embedding.rs | 100 +++++++++++++++++------------------------------ 1 file changed, 35 insertions(+), 65 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index dd9271c..03f6855 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -296,7 +296,13 @@ impl EmbeddingClient { let sub_batch = if self.sub_batch_override > 0 { self.sub_batch_override } else { - Self::compute_sub_batch(loaded.dimension, self.gpu) + let mut sys = sysinfo::System::new(); + sys.refresh_memory(); + let available_mb = memory_budget( + sys.available_memory(), + sys.cgroup_limits().map(|limits| limits.free_memory), + ) / (1024 * 1024); + Self::compute_sub_batch(available_mb, loaded.dimension, self.gpu) }; let mut handles = Vec::new(); @@ -353,32 +359,21 @@ impl EmbeddingClient { /// When GPU is enabled, the upper clamp is raised to 256 (GPU VRAM can /// handle much larger batches than CPU). /// TODO: add here the config batch size - fn compute_sub_batch(dimension: usize, gpu: bool) -> usize { - let mut sys = sysinfo::System::new(); - sys.refresh_memory(); - let available_mb = memory_budget( - sys.available_memory(), - sys.cgroup_limits().map(|limits| limits.free_memory), - ) / (1024 * 1024); - - sub_batch_for(available_mb, dimension, gpu) - } -} + fn compute_sub_batch(available_mb: u64, dimension: usize, gpu: bool) -> usize { + if available_mb == 0 { + return 32; + } -fn sub_batch_for(available_mb: u64, dimension: usize, gpu: bool) -> usize { - if available_mb == 0 { - return 32; + // Per-text memory estimate for ONNX inference. Attention matrices + // dominate: heads × seq² × 4 bytes. For 768-dim BERT-like models + // (12 heads) processing ~1000-2000 token code chunks, attention alone + // is 50-200 MB per text. The estimate below is conservative so the + // sub-batch stays small enough to prevent arena over-allocation. + let mb_per_text: u64 = if dimension >= 768 { 100 } else { 40 }; + let budget_mb = available_mb / 2; + let max_batch = if gpu { 256 } else { 16 }; + (budget_mb / mb_per_text).clamp(1, max_batch) as usize } - - // Per-text memory estimate for ONNX inference. Attention matrices - // dominate: heads × seq² × 4 bytes. For 768-dim BERT-like models - // (12 heads) processing ~1000-2000 token code chunks, attention alone - // is 50-200 MB per text. The estimate below is conservative so the - // sub-batch stays small enough to prevent arena over-allocation. - let mb_per_text: u64 = if dimension >= 768 { 100 } else { 40 }; - let budget_mb = available_mb / 2; - let max_batch = if gpu { 256 } else { 16 }; - (budget_mb / mb_per_text).clamp(1, max_batch) as usize } #[cfg(test)] @@ -452,25 +447,6 @@ mod tests { assert_eq!(memory_budget(2 * GIB, Some(32 * GIB)), 2 * GIB); } - #[test] - fn sub_batch_falls_back_when_memory_is_unknown() { - assert_eq!(sub_batch_for(0, 768, false), 32); - } - - #[test] - fn sub_batch_never_exceeds_a_small_budget() { - // 200 MiB free -> 100 MiB budget -> one 100 MiB text, not four. - assert_eq!(sub_batch_for(200, 768, false), 1); - assert_eq!(sub_batch_for(1, 768, false), 1); - assert_eq!(sub_batch_for(200, 384, false), 2); - } - - #[test] - fn sub_batch_is_capped_by_device() { - assert_eq!(sub_batch_for(64 * 1024, 768, false), 16); - assert_eq!(sub_batch_for(64 * 1024, 768, true), 256); - } - #[test] fn from_env_uses_nomic_when_unset() { let _g = isolate_env(); @@ -589,33 +565,27 @@ mod tests { } #[test] - fn compute_sub_batch_cpu_768_within_clamp() { - let result = EmbeddingClient::compute_sub_batch(768, false); - // 32 is the sysinfo-zero fallback; otherwise the CPU clamp is [4, 16]. - assert!( - result == 32 || (4..=16).contains(&result), - "got {} for cpu/768", - result - ); + fn compute_sub_batch_falls_back_when_memory_is_unknown() { + assert_eq!(EmbeddingClient::compute_sub_batch(0, 768, false), 32); } #[test] - fn compute_sub_batch_cpu_384_within_clamp() { - let result = EmbeddingClient::compute_sub_batch(384, false); - assert!( - result == 32 || (4..=16).contains(&result), - "got {} for cpu/384", - result - ); + fn compute_sub_batch_never_exceeds_a_small_budget() { + // 200 MiB free -> 100 MiB budget -> one 100 MiB text, not four. + assert_eq!(EmbeddingClient::compute_sub_batch(200, 768, false), 1); + assert_eq!(EmbeddingClient::compute_sub_batch(1, 768, false), 1); + assert_eq!(EmbeddingClient::compute_sub_batch(200, 384, false), 2); } #[test] - fn compute_sub_batch_gpu_raises_ceiling() { - let result = EmbeddingClient::compute_sub_batch(768, true); - assert!( - result == 32 || (4..=256).contains(&result), - "got {} for gpu/768", - result + fn compute_sub_batch_is_capped_by_device() { + assert_eq!( + EmbeddingClient::compute_sub_batch(64 * 1024, 768, false), + 16 + ); + assert_eq!( + EmbeddingClient::compute_sub_batch(64 * 1024, 768, true), + 256 ); } From 9587705d571b97e3f34ea9f904fcc2417f88e25c Mon Sep 17 00:00:00 2001 From: Chirag Aggarwal Date: Fri, 4 Sep 2026 16:54:50 +0100 Subject: [PATCH 5/6] test: collapse memory sizing coverage into two cases --- src/embedding.rs | 25 +++---------------------- 1 file changed, 3 insertions(+), 22 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index 03f6855..b234ac5 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -430,20 +430,11 @@ mod tests { } } - const GIB: u64 = 1024 * 1024 * 1024; - - #[test] - fn memory_budget_uses_host_without_cgroup() { - assert_eq!(memory_budget(32 * GIB, None), 32 * GIB); - } - #[test] fn memory_budget_is_capped_by_cgroup_limit() { + const GIB: u64 = 1024 * 1024 * 1024; + assert_eq!(memory_budget(32 * GIB, None), 32 * GIB); assert_eq!(memory_budget(32 * GIB, Some(GIB)), GIB); - } - - #[test] - fn memory_budget_is_capped_by_host_when_cgroup_is_unlimited() { assert_eq!(memory_budget(2 * GIB, Some(32 * GIB)), 2 * GIB); } @@ -565,20 +556,10 @@ mod tests { } #[test] - fn compute_sub_batch_falls_back_when_memory_is_unknown() { + fn compute_sub_batch_shrinks_to_fit_small_budgets() { assert_eq!(EmbeddingClient::compute_sub_batch(0, 768, false), 32); - } - - #[test] - fn compute_sub_batch_never_exceeds_a_small_budget() { // 200 MiB free -> 100 MiB budget -> one 100 MiB text, not four. assert_eq!(EmbeddingClient::compute_sub_batch(200, 768, false), 1); - assert_eq!(EmbeddingClient::compute_sub_batch(1, 768, false), 1); - assert_eq!(EmbeddingClient::compute_sub_batch(200, 384, false), 2); - } - - #[test] - fn compute_sub_batch_is_capped_by_device() { assert_eq!( EmbeddingClient::compute_sub_batch(64 * 1024, 768, false), 16 From f49e1a0b260fa378470758008809ded59671187a Mon Sep 17 00:00:00 2001 From: Chirag Aggarwal Date: Fri, 4 Sep 2026 17:01:04 +0100 Subject: [PATCH 6/6] fix: treat zero available memory as exhausted, not unknown --- src/embedding.rs | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/embedding.rs b/src/embedding.rs index b234ac5..19c8140 100644 --- a/src/embedding.rs +++ b/src/embedding.rs @@ -354,16 +354,11 @@ impl EmbeddingClient { /// Compute sub-batch size based on available system memory. /// - /// Uses 50% of available RAM as a budget. Falls back to 32 if sysinfo - /// reports 0. + /// Uses 50% of available RAM as a budget. /// When GPU is enabled, the upper clamp is raised to 256 (GPU VRAM can /// handle much larger batches than CPU). /// TODO: add here the config batch size fn compute_sub_batch(available_mb: u64, dimension: usize, gpu: bool) -> usize { - if available_mb == 0 { - return 32; - } - // Per-text memory estimate for ONNX inference. Attention matrices // dominate: heads × seq² × 4 bytes. For 768-dim BERT-like models // (12 heads) processing ~1000-2000 token code chunks, attention alone @@ -557,7 +552,7 @@ mod tests { #[test] fn compute_sub_batch_shrinks_to_fit_small_budgets() { - assert_eq!(EmbeddingClient::compute_sub_batch(0, 768, false), 32); + assert_eq!(EmbeddingClient::compute_sub_batch(0, 768, false), 1); // 200 MiB free -> 100 MiB budget -> one 100 MiB text, not four. assert_eq!(EmbeddingClient::compute_sub_batch(200, 768, false), 1); assert_eq!(