Coordinate foreground search and background inference CPU budgets

2026-10-13 01:56:03 +02:00
1 parent 6b78c8bb7f
commit d85746db5b
@@ -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(())
+}