Coordinate foreground search and background inference CPU budgets

crates/searchi-service/src/coordination.rs

@@ -0,0 +1,353 @@
+//! Shared admission for persistent inference sessions. This bounds admitted
+//! native worker threads, not all extraction/ANN/OS work. Already executing
+//! inference cannot be preempted; long audio windows remain a latency limit.
+use crate::EmbeddingProvider;
+use anyhow::{bail, Result};
+use searchi_core::metrics::{SpanToken, TraceCollector};
+use searchi_semantic::TextWindow;
+use std::{
+    collections::BTreeMap,
+    path::Path,
+    sync::{Arc, Condvar, Mutex},
+    time::{Duration, Instant},
+};
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum InferenceRole {
+    Foreground,
+    Background,
+}
+#[derive(Default)]
+struct State {
+    limit: usize,
+    used: usize,
+    foreground: usize,
+}
+struct Shared {
+    state: Mutex<State>,
+    changed: Condvar,
+}
+/// Clone this handle across query, setup and indexing sessions.
+#[derive(Clone)]
+pub struct InferenceBudget(Arc<Shared>);
+impl InferenceBudget {
+    pub fn new(requested: usize) -> Self {
+        Self(Arc::new(Shared {
+            state: Mutex::new(State {
+                limit: clamp(requested),
+                ..State::default()
+            }),
+            changed: Condvar::new(),
+        }))
+    }
+    pub fn threads(&self) -> usize {
+        self.0.state.lock().unwrap_or_else(|e| e.into_inner()).limit
+    }
+    /// Leave one admitted worker thread for background indexing. On a
+    /// single-thread budget, already executing inference still cannot yield.
+    pub fn worker_threads(&self, role: InferenceRole) -> usize {
+        match role {
+            InferenceRole::Foreground => self.threads().saturating_sub(1).max(1),
+            InferenceRole::Background => 1,
+        }
+    }
+    /// Update admission immediately; already admitted work is not cancelled.
+    /// Recreate sessions whose worker thread count exceeds a reduced limit.
+    pub fn configure(&self, requested: usize) -> usize {
+        let limit = clamp(requested);
+        self.0.state.lock().unwrap_or_else(|e| e.into_inner()).limit = limit;
+        self.0.changed.notify_all();
+        limit
+    }
+    /// Register before waiting for a service mutex. This reserves priority, not
+    /// threads, so the foreground does not deadlock with its indexing session.
+    pub fn pending_foreground(&self) -> PendingForeground {
+        self.0
+            .state
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .foreground += 1;
+        self.0.changed.notify_all();
+        PendingForeground(self.clone())
+    }
+    pub fn acquire(
+        &self,
+        role: InferenceRole,
+        actual_worker_threads: usize,
+    ) -> Result<InferencePermit> {
+        let weight = actual_worker_threads.max(1);
+        let _pending = (role == InferenceRole::Foreground).then(|| self.pending_foreground());
+        let deadline = Instant::now() + Duration::from_secs(30);
+        let mut state = self.0.state.lock().unwrap_or_else(|e| e.into_inner());
+        loop {
+            if weight > state.limit {
+                bail!("Inference session uses {weight} threads but shared budget is {}; recreate the session after changing its thread count", state.limit);
+            }
+            if state.used <= state.limit - weight
+                && (role == InferenceRole::Foreground || state.foreground == 0)
+            {
+                state.used += weight;
+                return Ok(InferencePermit {
+                    budget: self.clone(),
+                    weight,
+                });
+            }
+            let remaining = deadline.saturating_duration_since(Instant::now());
+            if remaining.is_zero() {
+                bail!("Timed out waiting for shared inference budget");
+            }
+            state = self
+                .0
+                .changed
+                .wait_timeout(state, remaining)
+                .unwrap_or_else(|e| e.into_inner())
+                .0;
+        }
+    }
+}
+fn clamp(requested: usize) -> usize {
+    requested
+        .clamp(1, 8)
+        .min(std::thread::available_parallelism().map_or(1, usize::from))
+}
+pub struct PendingForeground(InferenceBudget);
+impl Drop for PendingForeground {
+    fn drop(&mut self) {
+        let mut s = self.0 .0.state.lock().unwrap_or_else(|e| e.into_inner());
+        s.foreground -= 1;
+        self.0 .0.changed.notify_all();
+    }
+}
+pub struct InferencePermit {
+    budget: InferenceBudget,
+    weight: usize,
+}
+impl Drop for InferencePermit {
+    fn drop(&mut self) {
+        let mut s = self
+            .budget
+            .0
+            .state
+            .lock()
+            .unwrap_or_else(|e| e.into_inner());
+        s.used -= self.weight;
+        self.budget.0.changed.notify_all();
+    }
+}
+
+/// Delegates genuine inference unchanged. Background documents yield between
+/// individual encodes even when the caller supplies a larger batch.
+pub struct BudgetedProvider {
+    inner: Box<dyn EmbeddingProvider>,
+    budget: InferenceBudget,
+    role: InferenceRole,
+    threads: usize,
+    query_started: bool,
+}
+impl BudgetedProvider {
+    pub fn new(
+        inner: Box<dyn EmbeddingProvider>,
+        budget: &InferenceBudget,
+        role: InferenceRole,
+        actual_worker_threads: usize,
+    ) -> Self {
+        Self {
+            inner,
+            budget: budget.clone(),
+            role,
+            threads: actual_worker_threads.max(1),
+            query_started: false,
+        }
+    }
+    fn permit(&self) -> Result<InferencePermit> {
+        self.budget.acquire(self.role, self.threads)
+    }
+}
+impl EmbeddingProvider for BudgetedProvider {
+    fn document_batch_limit(&self) -> usize {
+        if self.role == InferenceRole::Background {
+            1
+        } else {
+            self.inner.document_batch_limit()
+        }
+    }
+    fn fingerprint(&self) -> &str {
+        self.inner.fingerprint()
+    }
+    fn dimensions(&self) -> usize {
+        self.inner.dimensions()
+    }
+    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
+        let _p = self.permit()?;
+        self.inner.split_windows(text, title)
+    }
+    fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
+        if self.role == InferenceRole::Background {
+            let mut output = Vec::with_capacity(texts.len());
+            for text in texts {
+                let _p = self.permit()?;
+                output.extend(self.inner.embed_documents(&[*text], title)?);
+            }
+            Ok(output)
+        } else {
+            let _p = self.permit()?;
+            self.inner.embed_documents(texts, title)
+        }
+    }
+    fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_started = false;
+        let _p = self.permit()?;
+        self.query_started = true;
+        self.inner.embed_query(query)
+    }
+    fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_started = false;
+        let _p = self.permit()?;
+        self.query_started = true;
+        self.inner.embed_query_debug(query)
+    }
+    fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_started = false;
+        let _p = self.permit()?;
+        self.query_started = true;
+        self.inner.embed_code_query(query)
+    }
+    fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
+        self.query_started = false;
+        let _p = self.permit()?;
+        self.query_started = true;
+        self.inner.embed_code_query_debug(query)
+    }
+    fn query_diagnostics(&self) -> BTreeMap<String, u64> {
+        if self.query_started {
+            self.inner.query_diagnostics()
+        } else {
+            Default::default()
+        }
+    }
+    fn embed_query_traced(
+        &mut self,
+        query: &str,
+        debug: bool,
+        trace: &mut TraceCollector,
+        parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        self.query_started = false;
+        let span = trace.begin("semantic.inference.admission_wait", parent);
+        let permit = self.permit();
+        trace.end(
+            span,
+            if trace.is_debug() {
+                BTreeMap::from([
+                    ("worker_threads".into(), self.threads as u64),
+                    ("failed".into(), u64::from(permit.is_err())),
+                ])
+            } else {
+                Default::default()
+            },
+        );
+        let _permit = permit?;
+        self.query_started = true;
+        self.inner.embed_query_traced(query, debug, trace, parent)
+    }
+    fn embed_code_query_traced(
+        &mut self,
+        query: &str,
+        debug: bool,
+        trace: &mut TraceCollector,
+        parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        self.query_started = false;
+        let span = trace.begin("semantic.inference.admission_wait", parent);
+        let permit = self.permit();
+        trace.end(
+            span,
+            if trace.is_debug() {
+                BTreeMap::from([
+                    ("worker_threads".into(), self.threads as u64),
+                    ("failed".into(), u64::from(permit.is_err())),
+                ])
+            } else {
+                Default::default()
+            },
+        );
+        let _permit = permit?;
+        self.query_started = true;
+        self.inner
+            .embed_code_query_traced(query, debug, trace, parent)
+    }
+    fn embed_images(&mut self, paths: &[&Path], budget: u32) -> Result<Vec<Vec<f32>>> {
+        if self.role == InferenceRole::Background {
+            let mut output = Vec::with_capacity(paths.len());
+            for path in paths {
+                let _p = self.permit()?;
+                output.extend(self.inner.embed_images(&[*path], budget)?);
+            }
+            return Ok(output);
+        }
+        let _p = self.permit()?;
+        self.inner.embed_images(paths, budget)
+    }
+    fn embed_audio(&mut self, pcm: &[f32]) -> Result<Vec<f32>> {
+        let _p = self.permit()?;
+        self.inner.embed_audio(pcm)
+    }
+    fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
+        self.inner.worker_info()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn reserved_query_capacity_admits_foreground_during_background() {
+        let budget = InferenceBudget::new(1);
+        // Deliberately exercise a two-thread scheduler independent of CI hardware.
+        budget.0.state.lock().unwrap().limit = 2;
+        assert_eq!(budget.worker_threads(InferenceRole::Foreground), 1);
+        assert_eq!(budget.worker_threads(InferenceRole::Background), 1);
+        let background = budget.acquire(InferenceRole::Background, 1).unwrap();
+        let foreground = budget.acquire(InferenceRole::Foreground, 1).unwrap();
+        assert_eq!(budget.0.state.lock().unwrap().used, 2);
+        drop(foreground);
+        drop(background);
+    }
+    #[test]
+    fn pending_query_blocks_background_until_released() {
+        let budget = InferenceBudget::new(1);
+        let pending = budget.pending_foreground();
+        let other = budget.clone();
+        let (tx, rx) = std::sync::mpsc::channel();
+        let worker = std::thread::spawn(move || {
+            let _permit = other.acquire(InferenceRole::Background, 1).unwrap();
+            tx.send(()).unwrap();
+        });
+        assert!(rx.recv_timeout(Duration::from_millis(50)).is_err());
+        let foreground = budget.acquire(InferenceRole::Foreground, 1).unwrap();
+        drop(pending);
+        assert!(rx.recv_timeout(Duration::from_millis(50)).is_err());
+        drop(foreground);
+        rx.recv_timeout(Duration::from_secs(2)).unwrap();
+        worker.join().unwrap();
+    }
+    #[test]
+    fn panic_releases_threads_and_priority() {
+        let budget = InferenceBudget::new(1);
+        let shared = budget.clone();
+        assert!(std::panic::catch_unwind(move || {
+            let _pending = shared.pending_foreground();
+            let _permit = shared.acquire(InferenceRole::Foreground, 1).unwrap();
+            panic!("owned inference failure");
+        })
+        .is_err());
+        let _permit = budget.acquire(InferenceRole::Background, 1).unwrap();
+    }
+    #[test]
+    fn invalid_session_weight_fails_instead_of_oversubscribing() {
+        let budget = InferenceBudget::new(1);
+        assert!(budget.acquire(InferenceRole::Foreground, 2).is_err());
+        assert_eq!(budget.configure(0), 1);
+        assert!(InferenceBudget::new(usize::MAX).threads() <= 8);
+    }
+}

crates/searchi-service/src/lib.rs

@@ -1,6 +1,10 @@
 //! Shared offline hybrid search using real Google embeddings and persisted,
 //! filtered HNSW retrieval. Production-scale validation remains a release gate.
+mod coordination;
 mod media;
+pub use coordination::{
+    BudgetedProvider, InferenceBudget, InferencePermit, InferenceRole, PendingForeground,
+};
 mod prepared;
 pub use prepared::PreparedGraphs;
 mod pipeline;
@@ -52,6 +56,11 @@ fn reusable_policy(expected: &str, active: &str) -> bool {
 /// Model-specific query/document prompts and actual tokenizer windows belong
 /// to the provider. Storage and ranking never assume a particular dimension.
 pub trait EmbeddingProvider: Send {
+    /// Maximum document batch admitted as one inference call. Coordinated
+    /// background sessions yield at every document boundary.
+    fn document_batch_limit(&self) -> usize {
+        MAX_BATCH
+    }
     fn fingerprint(&self) -> &str;
     fn dimensions(&self) -> usize;
     fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>>;
@@ -66,6 +75,32 @@ pub trait EmbeddingProvider: Send {
     fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
         self.embed_query(query)
     }
+    fn embed_query_traced(
+        &mut self,
+        query: &str,
+        debug: bool,
+        _trace: &mut TraceCollector,
+        _parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        if debug {
+            self.embed_query_debug(query)
+        } else {
+            self.embed_query(query)
+        }
+    }
+    fn embed_code_query_traced(
+        &mut self,
+        query: &str,
+        debug: bool,
+        _trace: &mut TraceCollector,
+        _parent: Option<SpanToken>,
+    ) -> Result<Vec<f32>> {
+        if debug {
+            self.embed_code_query_debug(query)
+        } else {
+            self.embed_code_query(query)
+        }
+    }
     fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
         Default::default()
     }
@@ -174,6 +209,21 @@ pub struct HybridSearch {
 }
 
 impl HybridSearch {
+    /// Attach shared admission without replacing the warmed model or query LRU.
+    pub fn with_budget(
+        mut self,
+        budget: &InferenceBudget,
+        role: InferenceRole,
+        actual_worker_threads: usize,
+    ) -> Self {
+        self.embedder = Box::new(BudgetedProvider::new(
+            self.embedder,
+            budget,
+            role,
+            actual_worker_threads,
+        ));
+        self
+    }
     /// Connect background publication to an existing warmed reader.
     pub fn share_prepared_graphs(&mut self, graphs: &PreparedGraphs) -> Result<()> {
         self.store.share_prepared_graphs(graphs)
@@ -442,11 +492,12 @@ impl HybridSearch {
         trace.counter(span, "query_cache_hit", 0);
         trace.counter(span, "query_cache_miss", 1);
         let vector = (|| {
-            let vector = match (code, trace.is_debug()) {
-                (true, true) => self.embedder.embed_code_query_debug(query)?,
-                (true, false) => self.embedder.embed_code_query(query)?,
-                (false, true) => self.embedder.embed_query_debug(query)?,
-                (false, false) => self.embedder.embed_query(query)?,
+            let vector = if code {
+                self.embedder
+                    .embed_code_query_traced(query, trace.is_debug(), trace, span)?
+            } else {
+                self.embedder
+                    .embed_query_traced(query, trace.is_debug(), trace, span)?
             };
             ensure!(
                 vector.len() == self.embedder.dimensions()

crates/searchi-service/src/pipeline.rs

@@ -258,14 +258,15 @@ impl HybridSearch {
             work.report.chunks + windows.len() <= MAX_SCAN_CHUNKS,
             "Corpus exceeds ANN generation resource budget of {MAX_SCAN_CHUNKS} chunks"
         );
-        for (batch_index, batch) in windows.chunks(MAX_BATCH).enumerate() {
+        let batch_limit = self.embedder.document_batch_limit().clamp(1, MAX_BATCH);
+        for (batch_index, batch) in windows.chunks(batch_limit).enumerate() {
             let texts: Vec<_> = batch.iter().map(|window| window.text.as_str()).collect();
             work.report.document_inference_calls += 1;
             let vectors = self.embedder.embed_documents(&texts, title)?;
             self.store.insert(
                 generation,
                 passage,
-                batch_index * MAX_BATCH,
+                batch_index * batch_limit,
                 batch,
                 &vectors,
             )?;