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,
             )?;

crates/searchi-service/tests/budget_google_cpu.rs

@@ -0,0 +1,177 @@
+//! Actual Google CPU inference during background indexing; not a scale benchmark.
+use anyhow::Result;
+use searchi_core::SearchIndex;
+use searchi_semantic::{GemmaEmbedder, TextWindow};
+use searchi_service::{
+    EmbeddingProvider, HybridSearch, InferenceBudget, InferenceRole, PreparedGraphs, SearchOptions,
+};
+use std::{
+    fs,
+    sync::{
+        atomic::{AtomicBool, AtomicUsize, Ordering},
+        mpsc, Arc,
+    },
+    time::Duration,
+};
+
+struct ObservedGoogle {
+    inner: GemmaEmbedder,
+    background: bool,
+    in_flight: Arc<AtomicBool>,
+    overlap: Arc<AtomicBool>,
+    entered: mpsc::Sender<()>,
+    calls: Arc<AtomicUsize>,
+}
+impl EmbeddingProvider for ObservedGoogle {
+    fn fingerprint(&self) -> &str {
+        EmbeddingProvider::fingerprint(&self.inner)
+    }
+    fn dimensions(&self) -> usize {
+        EmbeddingProvider::dimensions(&self.inner)
+    }
+    fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
+        self.inner.split_windows(text, title)
+    }
+    fn embed_documents(&mut self, text: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
+        assert!(self.background);
+        self.calls.fetch_add(1, Ordering::Relaxed);
+        self.in_flight.store(true, Ordering::Release);
+        let _ = self.entered.send(());
+        let result = self.inner.embed_documents(text, title);
+        self.in_flight.store(false, Ordering::Release);
+        result
+    }
+    fn embed_query(&mut self, text: &str) -> Result<Vec<f32>> {
+        if self.in_flight.load(Ordering::Acquire) {
+            self.overlap.store(true, Ordering::Release);
+        }
+        self.inner.embed_query(text)
+    }
+    fn embed_query_debug(&mut self, text: &str) -> Result<Vec<f32>> {
+        if self.in_flight.load(Ordering::Acquire) {
+            self.overlap.store(true, Ordering::Release);
+        }
+        self.inner.embed_query_debug(text)
+    }
+    fn query_diagnostics(&self) -> std::collections::BTreeMap<String, u64> {
+        self.inner.query_diagnostics()
+    }
+    fn worker_info(&self) -> Option<searchi_semantic::WorkerInfo> {
+        Some(self.inner.info.clone())
+    }
+}
+#[test]
+#[ignore = "Requires verified real Google worker, two CPU cores; CPU only"]
+fn foreground_google_query_runs_while_background_encoder_is_admitted() -> Result<()> {
+    assert_eq!(std::env::var("SEARCHI_FORCE_CPU")?, "1");
+    let model = std::env::var_os("SEARCHI_GEMMA_MODEL_DIR").expect("Verified model required");
+    let worker = std::env::var_os("SEARCHI_GEMMA_WORKER").expect("Normal worker required");
+    let budget = InferenceBudget::new(2);
+    assert_eq!(
+        budget.threads(),
+        2,
+        "This concurrent test requires two allocated CPU cores"
+    );
+    let temp = tempfile::tempdir()?;
+    let root = temp.path().join("sources");
+    fs::create_dir(&root)?;
+    fs::write(
+        root.join("garden.txt"),
+        "Tomato plants need regular watering and plenty of sunlight.",
+    )?;
+    fs::write(
+        root.join("supplier.txt"),
+        "The supplier contract ends with sixty days written notice.",
+    )?;
+    let db = temp.path().join("index.db");
+    let mut lexical = SearchIndex::open(&db)?;
+    lexical.index_folder(&root)?;
+    let graphs = PreparedGraphs::default();
+    let mut initial = HybridSearch::open(&db, &model, &worker, 1)?;
+    initial.share_prepared_graphs(&graphs)?;
+    initial.index_embeddings()?;
+    drop(initial);
+    let in_flight = Arc::new(AtomicBool::new(false));
+    let overlap = Arc::new(AtomicBool::new(false));
+    let calls = Arc::new(AtomicUsize::new(0));
+    let (entered, events) = mpsc::channel();
+    let make = |background| -> Result<HybridSearch> {
+        let role = if background {
+            InferenceRole::Background
+        } else {
+            InferenceRole::Foreground
+        };
+        let threads = budget.worker_threads(role);
+        let startup = budget.acquire(role, threads)?;
+        let inner = GemmaEmbedder::load(&model, &worker, threads)?;
+        assert_eq!(inner.info.device, "cpu");
+        drop(startup);
+        let mut service = HybridSearch::with_provider(
+            &db,
+            Box::new(ObservedGoogle {
+                inner,
+                background,
+                in_flight: in_flight.clone(),
+                overlap: overlap.clone(),
+                entered: entered.clone(),
+                calls: calls.clone(),
+            }),
+        )?
+        .with_budget(&budget, role, threads);
+        service.share_prepared_graphs(&graphs)?;
+        Ok(service)
+    };
+    let mut reader = make(false)?;
+    reader.warmup()?;
+    let mut indexer = make(true)?;
+    fs::write(
+        root.join("maintenance.txt"),
+        "Technical maintenance notes for equipment inventory scheduling and routine inspections. "
+            .repeat(800),
+    )?;
+    lexical.index_folder(&root)?;
+    let (report, sample) = std::thread::scope(|scope| -> Result<_> {
+        let indexing = scope.spawn(move || indexer.index_embeddings());
+        let mut sample = None;
+        for query in [
+            "How should we water tomato plants?",
+            "What do tomato plants need for growth?",
+            "How can we keep tomatoes healthy?",
+        ] {
+            events.recv_timeout(Duration::from_secs(90))?;
+            let pending = budget.pending_foreground();
+            let response = reader.search_measured(query, 10, &SearchOptions::default(), true)?;
+            drop(pending);
+            assert_eq!(response.results[0].filename, "garden.txt");
+            let trace = response.metrics.trace.as_ref().expect("Detailed trace");
+            assert!(trace
+                .iter()
+                .any(|span| span.counters.get("query_cache_hit") == Some(&0)));
+            let wait = trace
+                .iter()
+                .find(|span| span.name == "semantic.inference.admission_wait")
+                .expect("Actual admission span")
+                .duration_ms;
+            sample = Some((response.metrics.total_ms, wait));
+            if overlap.load(Ordering::Acquire) {
+                break;
+            }
+        }
+        let report = indexing.join().expect("Index thread panic")?;
+        Ok((report, sample))
+    })?;
+    assert!(
+        overlap.load(Ordering::Acquire),
+        "Must observe genuine query inference while a background call remains active"
+    );
+    assert_eq!(report.reused_sources, 2);
+    assert_eq!(report.encoded_sources, 1);
+    assert_eq!(
+        report.document_inference_calls,
+        calls.load(Ordering::Relaxed)
+    );
+    assert!(report.document_inference_calls > 1);
+    let (query_ms, admission_ms) = sample.unwrap();
+    eprintln!("Concurrent Google CPU fixture: total_budget=2 query_threads=1 indexing_threads=1 query_ms={query_ms:.3} admission_ms={admission_ms:.3} document_calls={}; not scale/p95/UI evidence", report.document_inference_calls);
+    Ok(())
+}