From d85746db5b396f2d08c75913cb35115526ef100a Mon Sep 17 00:00:00 2001 From: karl Date: Tue, 13 Oct 2026 01:56:03 +0200 Subject: [PATCH] Coordinate foreground search and background inference CPU budgets --- AGENTS.md | 5 + apps/desktop/src-tauri/src/main.rs | 4 +- .../src-tauri/src/semantic_commands.rs | 232 ++++++++++-- apps/desktop/src-tauri/src/settings.rs | 38 +- apps/desktop/src/main.ts | 2 +- crates/searchi-service/src/coordination.rs | 353 ++++++++++++++++++ crates/searchi-service/src/lib.rs | 61 ++- crates/searchi-service/src/pipeline.rs | 5 +- .../tests/budget_google_cpu.rs | 177 +++++++++ docs/PERFORMANCE.md | 29 +- docs/RECOVERY-AND-HANDOFF.md | 7 +- docs/ROADMAP.md | 5 + 12 files changed, 858 insertions(+), 60 deletions(-) create mode 100644 crates/searchi-service/src/coordination.rs create mode 100644 crates/searchi-service/tests/budget_google_cpu.rs diff --git a/AGENTS.md b/AGENTS.md index a4a00f9..5064572 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,6 +46,11 @@ readiness without completed release gates and recorded evidence. graphs before SQL activation; retain the query model/LRU across publication instead of repeating warmup under the foreground mutex. Measure startup, actual first input, and warm p95 separately. +- Desktop CPU threads are a shared inference budget (settings version two), + default two and hardware-clamped. Keep foreground priority and one background + thread; yield between document encodes. Trace admission separately from native + inference. This does not bound decoder, extraction or ANN CPU work. Do not + claim preemption of an already executing media encode on a one-thread budget. - UI should belong to each OS: native typography/chrome/dialogs, themes, keyboard access, reduced motion and high contrast. Brief transitions must not delay requests. Preserve selection/focus as progressive results arrive. diff --git a/apps/desktop/src-tauri/src/main.rs b/apps/desktop/src-tauri/src/main.rs index 7e1b451..98defa8 100644 --- a/apps/desktop/src-tauri/src/main.rs +++ b/apps/desktop/src-tauri/src/main.rs @@ -113,11 +113,9 @@ fn set_cpu_threads(threads: usize, state: State<'_, AppState>) -> Result>, service: Mutex>, prepared_graphs: searchi_service::PreparedGraphs, + inference_budget: InferenceBudget, + lifecycle: Mutex<()>, media_tools: Option, media_setup_error: Option, } @@ -53,12 +55,27 @@ impl SemanticState { error: Mutex::new(None), service: Mutex::new(None), prepared_graphs: searchi_service::PreparedGraphs::default(), + inference_budget: InferenceBudget::new(2), + lifecycle: Mutex::new(()), media_tools, media_setup_error, } } - pub fn reset_session(&self) -> Result<(), String> { + fn ensure_idle(&self) -> Result<(), String> { + if self.embedding.load(Ordering::Acquire) + || self.downloading.load(Ordering::Acquire) + || self.warming.load(Ordering::Acquire) + || self.querying.load(Ordering::Acquire) + { + return Err("Wait for semantic work before changing resource settings.".into()); + } + Ok(()) + } + + #[cfg(test)] + fn reset_session_locked(&self) -> Result<(), String> { + self.ensure_idle()?; *self .service .try_lock() @@ -75,9 +92,19 @@ impl SemanticState { "The installed embedding worker is missing. Repair or reinstall Searchi.".into(), ); } + self.inference_budget.configure(threads); + let threads = self + .inference_budget + .worker_threads(InferenceRole::Foreground); + let admission = self + .inference_budget + .acquire(InferenceRole::Foreground, threads) + .map_err(|error| error.to_string())?; let mut service = HybridSearch::open(database, &self.model_directory, &self.worker, threads) .map_err(|error| error.to_string())?; + drop(admission); + service = service.with_budget(&self.inference_budget, InferenceRole::Foreground, threads); service .share_prepared_graphs(&self.prepared_graphs) .map_err(|error| error.to_string())?; @@ -92,6 +119,12 @@ impl SemanticState { /// First-launch setup is authorized product behavior; inference remains offline. pub fn begin_automatic_setup(self: &std::sync::Arc, database: PathBuf, threads: usize) { + let Ok(admission) = self.lifecycle.lock() else { + return; + }; + if self.embedding.load(Ordering::Acquire) || self.querying.load(Ordering::Acquire) { + return; + } if self .downloading .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) @@ -99,6 +132,11 @@ impl SemanticState { { return; } + drop(admission); + self.spawn_setup(database, threads); + } + + fn spawn_setup(self: &std::sync::Arc, database: PathBuf, threads: usize) { let semantic = self.clone(); tauri::async_runtime::spawn_blocking(move || { let _guard = WorkGuard(&semantic.downloading); @@ -111,6 +149,34 @@ impl SemanticState { }); } + /// Linearize resource reset, persisted settings and setup admission. The + /// lifecycle lock never spans worker startup, inference or a service wait. + pub fn change_resources( + self: &std::sync::Arc, + database: PathBuf, + threads: usize, + save: impl FnOnce() -> Result<(), String>, + ) -> Result<(), String> { + let admission = self + .lifecycle + .lock() + .map_err(|_| "Semantic admission is unavailable.")?; + self.ensure_idle()?; + let mut service = self + .service + .try_lock() + .map_err(|_| "Wait for the current semantic query before changing resources.")?; + // A failed settings write must keep the existing warm session usable. + save()?; + *service = None; + self.warmed.store(false, Ordering::Release); + self.downloading.store(true, Ordering::Release); + drop(service); + drop(admission); + self.spawn_setup(database, threads); + Ok(()) + } + fn record_error(&self, value: Option) { if let Ok(mut error) = self.error.lock() { *error = value; @@ -188,6 +254,9 @@ pub fn semantic_status(state: State<'_, AppState>) -> Result) -> Result<(), String> { let semantic = state.semantic.clone(); + if semantic.embedding.load(Ordering::Acquire) || semantic.querying.load(Ordering::Acquire) { + return Err("Wait for semantic work before restarting model setup.".into()); + } let database = state.database.clone(); let threads = state .settings @@ -195,6 +264,13 @@ pub async fn setup_model(state: State<'_, AppState>) -> Result<(), String> { .map_err(|_| "Settings are unavailable.")? .value .cpu_threads; + let admission = semantic + .lifecycle + .lock() + .map_err(|_| "Semantic admission is unavailable.")?; + if semantic.embedding.load(Ordering::Acquire) || semantic.querying.load(Ordering::Acquire) { + return Err("Wait for semantic work before restarting model setup.".into()); + } if semantic .downloading .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) @@ -202,6 +278,7 @@ pub async fn setup_model(state: State<'_, AppState>) -> Result<(), String> { { return Err("Google model setup is already running.".into()); } + drop(admission); tauri::async_runtime::spawn_blocking(move || { let _guard = WorkGuard(&semantic.downloading); semantic.record_error(None); @@ -219,15 +296,16 @@ pub async fn setup_model(state: State<'_, AppState>) -> Result<(), String> { #[tauri::command] pub async fn index_embeddings(state: State<'_, AppState>) -> Result { let semantic = state.semantic.clone(); + let admission = semantic + .lifecycle + .lock() + .map_err(|_| "Semantic admission is unavailable.")?; + if semantic.downloading.load(Ordering::Acquire) || semantic.warming.load(Ordering::Acquire) { + return Err("Wait for model setup to finish before semantic indexing.".into()); + } if !semantic.verified.load(Ordering::Acquire) || !semantic.worker.is_file() { return Err("Set up the Google model and native worker first.".into()); } - let threads = state - .settings - .lock() - .map_err(|_| "Settings are unavailable.")? - .value - .cpu_threads; if semantic .embedding .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) @@ -236,23 +314,29 @@ pub async fn index_embeddings(state: State<'_, AppState>) -> Result, ) -> Result, String> { let semantic = state.semantic.clone(); + let admission = semantic + .lifecycle + .lock() + .map_err(|_| "Semantic admission is unavailable.")?; if !semantic.warmed.load(Ordering::Acquire) { return Err( "Semantic search is preparing in the background; keyword results are available.".into(), ); } let database = state.database.clone(); - let threads = state - .settings - .lock() - .map_err(|_| "Settings are unavailable.")? - .value - .cpu_threads; if semantic .querying .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) @@ -289,22 +371,36 @@ pub async fn hybrid_search( { return Err("Semantic worker is busy; keyword results remain available.".into()); } + let pending = semantic.inference_budget.pending_foreground(); + drop(admission); tauri::async_runtime::spawn_blocking(move || { let _guard = WorkGuard(&semantic.querying); + let _pending = pending; let mut cached = semantic .service .lock() .map_err(|_| "Semantic session is unavailable.")?; if cached.is_none() { - *cached = Some( - HybridSearch::open( - database, - &semantic.model_directory, - &semantic.worker, - threads, - ) - .map_err(|error| error.to_string())?, - ); + let threads = semantic + .inference_budget + .worker_threads(InferenceRole::Foreground); + let admission = semantic + .inference_budget + .acquire(InferenceRole::Foreground, threads) + .map_err(|error| error.to_string())?; + let session = HybridSearch::open( + database, + &semantic.model_directory, + &semantic.worker, + threads, + ) + .map_err(|error| error.to_string())?; + drop(admission); + *cached = Some(session.with_budget( + &semantic.inference_budget, + InferenceRole::Foreground, + threads, + )); } cached .as_mut() @@ -330,23 +426,18 @@ pub async fn hybrid_search_measured( MetricsMode::Basic }); let semantic = state.semantic.clone(); + let admission = semantic.lifecycle.lock().map_err(|_| { + MeasuredFailure::new( + "Semantic admission is unavailable.", + TraceCollector::new(MetricsMode::Basic), + ) + })?; if !semantic.warmed.load(Ordering::Acquire) { return Err(MeasuredFailure::new( "Semantic search is preparing in the background; keyword results are available.", trace, )); } - let threads = trace.timed("desktop.settings_lock", None, |_, _| { - state - .settings - .lock() - .map(|settings| settings.value.cpu_threads) - .map_err(|_| "Settings are unavailable.") - }); - let threads = match threads { - Ok(threads) => threads, - Err(error) => return Err(MeasuredFailure::new(error, trace)), - }; if semantic .querying .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) @@ -359,8 +450,11 @@ pub async fn hybrid_search_measured( } let database = state.database.clone(); let queue = trace.begin("desktop.worker_queue", None); + let pending = semantic.inference_budget.pending_foreground(); + drop(admission); tauri::async_runtime::spawn_blocking(move || { let _guard = WorkGuard(&semantic.querying); + let _pending = pending; trace.end(queue, Default::default()); let result = (|| { let mut cached = trace.timed("desktop.semantic_lock", None, |_, _| { @@ -370,7 +464,14 @@ pub async fn hybrid_search_measured( .map_err(|_| "Semantic session is unavailable.") })?; if cached.is_none() { - *cached = Some(trace.timed("desktop.session_initialize", None, |_, _| { + let threads = semantic + .inference_budget + .worker_threads(InferenceRole::Foreground); + let admission = semantic + .inference_budget + .acquire(InferenceRole::Foreground, threads) + .map_err(|error| error.to_string())?; + let session = trace.timed("desktop.session_initialize", None, |_, _| { HybridSearch::open( database, &semantic.model_directory, @@ -378,7 +479,13 @@ pub async fn hybrid_search_measured( threads, ) .map_err(|error| error.to_string()) - })?); + })?; + drop(admission); + *cached = Some(session.with_budget( + &semantic.inference_budget, + InferenceRole::Foreground, + threads, + )); } cached .as_mut() @@ -397,3 +504,46 @@ pub async fn hybrid_search_measured( .await .map_err(|error| MeasuredFailure::task(error.to_string(), started))? } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn failed_resource_save_preserves_ready_session_and_admission() { + let state = std::sync::Arc::new(SemanticState::new( + PathBuf::new(), + PathBuf::new(), + Err("not installed".into()), + )); + state.warmed.store(true, Ordering::Release); + state.verified.store(true, Ordering::Release); + let result = state.change_resources(PathBuf::new(), 2, || Err("disk write failed".into())); + assert_eq!(result.unwrap_err(), "disk write failed"); + assert!(state.warmed.load(Ordering::Acquire)); + assert!(state.verified.load(Ordering::Acquire)); + assert!(!state.downloading.load(Ordering::Acquire)); + assert!(!state.warming.load(Ordering::Acquire)); + assert!(state.service.try_lock().is_ok()); + assert!(state.lifecycle.try_lock().is_ok()); + } + + #[test] + fn resources_cannot_reset_active_semantic_work() { + let state = SemanticState::new(PathBuf::new(), PathBuf::new(), Err("not installed".into())); + for busy in [ + &state.embedding, + &state.downloading, + &state.warming, + &state.querying, + ] { + state.warmed.store(true, Ordering::Release); + busy.store(true, Ordering::Release); + assert!(state.reset_session_locked().is_err()); + assert!(state.warmed.load(Ordering::Acquire)); + busy.store(false, Ordering::Release); + } + state.reset_session_locked().unwrap(); + assert!(!state.warmed.load(Ordering::Acquire)); + } +} diff --git a/apps/desktop/src-tauri/src/settings.rs b/apps/desktop/src-tauri/src/settings.rs index e561aba..172b51e 100644 --- a/apps/desktop/src-tauri/src/settings.rs +++ b/apps/desktop/src-tauri/src/settings.rs @@ -7,7 +7,7 @@ use std::{ path::{Path, PathBuf}, }; -const SETTINGS_VERSION: u32 = 1; +const SETTINGS_VERSION: u32 = 2; const MAX_SETTINGS_BYTES: u64 = 128 * 1024; #[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] @@ -27,7 +27,7 @@ impl Default for Settings { Self { version: SETTINGS_VERSION, roots: Vec::new(), - cpu_threads: 1, + cpu_threads: 2, model_directory: None, } } @@ -71,7 +71,7 @@ pub struct SettingsStore { impl SettingsStore { pub fn load(path: PathBuf) -> Result { - let value = match fs::metadata(&path) { + let mut value = match fs::metadata(&path) { Ok(metadata) => { if metadata.len() > MAX_SETTINGS_BYTES { return Err("Settings file exceeds the size limit.".into()); @@ -82,6 +82,15 @@ impl SettingsStore { Err(error) if error.kind() == std::io::ErrorKind::NotFound => Settings::default(), Err(error) => return Err(error.to_string()), }; + // Version 1 applied the value independently to both workers. Reserve + // one background thread while retaining the old query allocation. + if value.version == 1 { + if !(1..=8).contains(&value.cpu_threads) { + return Err("CPU threads must be between 1 and 8.".into()); + } + value.version = SETTINGS_VERSION; + value.cpu_threads = (value.cpu_threads + 1).min(8); + } let store = Self { path, value }; store.value.validate()?; Ok(store) @@ -126,7 +135,7 @@ mod tests { let directory = tempfile::tempdir().unwrap(); let path = directory.path().join("settings.json"); let mut store = SettingsStore::load(path.clone()).unwrap(); - assert_eq!(store.value.cpu_threads, 1); + assert_eq!(store.value.cpu_threads, 2); store.add_root(directory.path()).unwrap(); store.add_root(directory.path()).unwrap(); let mut next = store.value.clone(); @@ -175,6 +184,8 @@ mod tests { fs::write(&path, r#"{"version":1,"roots":[],"cpu_threads":1}"#).unwrap(); let store = SettingsStore::load(path).unwrap(); assert!(store.value.model_directory.is_none()); + assert_eq!(store.value.version, 2); + assert_eq!(store.value.cpu_threads, 2); let invalid = Settings { model_directory: Some(PathBuf::from("relative")), ..Settings::default() @@ -199,3 +210,22 @@ mod tests { ); } } + +#[cfg(test)] +mod budget_migration_tests { + use super::*; + #[test] + fn version_two_single_thread_is_explicit_and_invalid_legacy_budget_is_rejected() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("settings.json"); + fs::write(&path, r#"{"version":2,"roots":[],"cpu_threads":1}"#).unwrap(); + assert_eq!( + SettingsStore::load(path.clone()).unwrap().value.cpu_threads, + 1 + ); + fs::write(&path, r#"{"version":1,"roots":[],"cpu_threads":9}"#).unwrap(); + assert!(SettingsStore::load(path.clone()).is_err()); + fs::write(&path, r#"{"version":1,"roots":[],"cpu_threads":8}"#).unwrap(); + assert_eq!(SettingsStore::load(path).unwrap().value.cpu_threads, 8); + } +} diff --git a/apps/desktop/src/main.ts b/apps/desktop/src/main.ts index b3b93e6..0b328fd 100644 --- a/apps/desktop/src/main.ts +++ b/apps/desktop/src/main.ts @@ -115,7 +115,7 @@ interface SearchRequest { } document.querySelector("#app")!.innerHTML = ` -
searchi
Loading your index…
${icon("filter")}All files${icon("chevron")}

Filter your search

${icon("settings")}

Folders & resources

Search automatically uses supported acceleration with CPU fallback.

    Local timing diagnostics for new searches. Queries and file paths are excluded from copied reports.

    Optional AI answers

    Disabled by default. Only an explicit action sends the shown passage to your configured endpoint.

    Endpoint, model and key stay in this app session. The key authenticates requests to your endpoint.

    +
    searchi
    Loading your index…
    ${icon("filter")}All files${icon("chevron")}

    Filter your search

    ${icon("settings")}

    Folders & resources

    Search automatically uses supported acceleration with CPU fallback. Two or more threads let search run alongside indexing; one thread can delay semantic search during indexing.

      Local timing diagnostics for new searches. Queries and file paths are excluded from copied reports.

      Optional AI answers

      Disabled by default. Only an explicit action sends the shown passage to your configured endpoint.

      Endpoint, model and key stay in this app session. The key authenticates requests to your endpoint.

      ${icon("sparkles")}Preparing local semantic search…Local model${icon("chevron")}
      Google EmbeddingGemma 2

      The verified local model downloads automatically on first launch (865 MB). Your files stay on this device.

      diff --git a/crates/searchi-service/src/coordination.rs b/crates/searchi-service/src/coordination.rs new file mode 100644 index 0000000..a1bf6d5 --- /dev/null +++ b/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, + changed: Condvar, +} +/// Clone this handle across query, setup and indexing sessions. +#[derive(Clone)] +pub struct InferenceBudget(Arc); +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 { + 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, + budget: InferenceBudget, + role: InferenceRole, + threads: usize, + query_started: bool, +} +impl BudgetedProvider { + pub fn new( + inner: Box, + 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 { + 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> { + let _p = self.permit()?; + self.inner.split_windows(text, title) + } + fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result>> { + 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> { + 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> { + 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> { + 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> { + self.query_started = false; + let _p = self.permit()?; + self.query_started = true; + self.inner.embed_code_query_debug(query) + } + fn query_diagnostics(&self) -> BTreeMap { + 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, + ) -> Result> { + 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, + ) -> Result> { + 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>> { + 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> { + let _p = self.permit()?; + self.inner.embed_audio(pcm) + } + fn worker_info(&self) -> Option { + 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); + } +} diff --git a/crates/searchi-service/src/lib.rs b/crates/searchi-service/src/lib.rs index 1b2bc1b..134ada1 100644 --- a/crates/searchi-service/src/lib.rs +++ b/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>; @@ -66,6 +75,32 @@ pub trait EmbeddingProvider: Send { fn embed_query_debug(&mut self, query: &str) -> Result> { self.embed_query(query) } + fn embed_query_traced( + &mut self, + query: &str, + debug: bool, + _trace: &mut TraceCollector, + _parent: Option, + ) -> Result> { + 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, + ) -> Result> { + if debug { + self.embed_code_query_debug(query) + } else { + self.embed_code_query(query) + } + } fn query_diagnostics(&self) -> std::collections::BTreeMap { 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() diff --git a/crates/searchi-service/src/pipeline.rs b/crates/searchi-service/src/pipeline.rs index bb9f718..508b5fc 100644 --- a/crates/searchi-service/src/pipeline.rs +++ b/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, )?; diff --git a/crates/searchi-service/tests/budget_google_cpu.rs b/crates/searchi-service/tests/budget_google_cpu.rs new file mode 100644 index 0000000..1f5aca2 --- /dev/null +++ b/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, + overlap: Arc, + entered: mpsc::Sender<()>, + calls: Arc, +} +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> { + self.inner.split_windows(text, title) + } + fn embed_documents(&mut self, text: &[&str], title: &str) -> Result>> { + 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> { + 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> { + 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 { + self.inner.query_diagnostics() + } + fn worker_info(&self) -> Option { + 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 { + 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(()) +} diff --git a/docs/PERFORMANCE.md b/docs/PERFORMANCE.md index 3217ee0..4cc08a7 100644 --- a/docs/PERFORMANCE.md +++ b/docs/PERFORMANCE.md @@ -181,4 +181,31 @@ The actual CPU regression checks both a repeated cached query and an unseen query after publication; the unseen input must record query_cache_hit zero. See [RECOVERY-AND-HANDOFF.md](RECOVERY-AND-HANDOFF.md). Its single-document request times do not establish representative-corpus p95, native first input -or foreground priority during indexing. Shared CPU admission remains open. +or foreground priority during indexing. + +## Shared inference admission + +The desktop now shares one CPU thread budget across its persistent query and +indexing workers. The default is two threads, clamped to available logical CPUs +and a maximum of eight. Indexing uses one thread; the query worker uses the +remaining threads (at least one). Pending searches have priority over new +background encodes, and document indexing yields between individual encodes. +Cached query vectors do not acquire an inference permit. + +Settings version two stores the total budget. Version-one settings counted +threads per worker; migration adds one background thread, capped at eight. +An explicit version-two budget of one remains one. Resource changes are rejected +during active semantic work; a failed settings save preserves the warm session. + +Debug traces record `semantic.inference.admission_wait` separately from native +inference. Admission waits are bounded at 30 seconds. Already executing native +inference cannot be preempted: a one-thread budget or single-core machine can +still delay a search during a long media encode. This budget covers admitted +inference worker threads, not extraction, decoding, ANN construction, rendering +or total operating-system CPU usage. + +The real Google CPU regression `budget_google_cpu` confirmed overlapping query +and background inference with a two-thread budget: one query thread, one indexing +thread, six background document calls, an unseen query taking 103.407 ms including +0.034 ms admission in the final scoped-cleanup run. These are Windows VM debug measurements on a small fixture, +not representative-corpus p95, native UI latency, GPU or scale evidence. diff --git a/docs/RECOVERY-AND-HANDOFF.md b/docs/RECOVERY-AND-HANDOFF.md index 58b7572..8f9bd44 100644 --- a/docs/RECOVERY-AND-HANDOFF.md +++ b/docs/RECOVERY-AND-HANDOFF.md @@ -77,8 +77,9 @@ are not certified by those smoke tests. ## Remaining performance/release work -Background and query workers still receive independent thread budgets; shared -CPU admission and foreground priority are open. Publication no longer charges +Background and query workers now share inference admission and foreground +priority; see [PERFORMANCE.md](PERFORMANCE.md#shared-inference-admission). +Publication no longer charges foreground graph deserialization, but eligibility refresh, indexing I/O and CPU contention need representative measurements. Native compositor/input timing, large-corpus p95, older CPUs, actual GPU/device-loss, power-loss durability, @@ -97,4 +98,4 @@ installers/upgrades/signatures and macOS packaging remain open. The prepared graph handle is process-local: foreign CLI/MCP index writers cannot share its native Arc. Cold CLI initialization and cross-process reader refresh remain separate costs. Representative concurrent-indexing/foreground measurements -and shared CPU thread admission remain the next performance milestone. +across decoder, extraction and ANN work remain the next performance milestone. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 41a1eda..c7e47d1 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -164,6 +164,11 @@ Audit and measured priorities: [EMBEDDINGGEMMA2-AUDIT.md](EMBEDDINGGEMMA2-AUDIT. See [INCREMENTAL-SEMANTIC.md](INCREMENTAL-SEMANTIC.md). - [ ] Benchmark genuine token-budgeted multi-sequence batching and isolation. - [ ] Implement and measure shared foreground/indexer resource budgets/priority. + - [x] Shared inference admission, foreground priority between background + encodes, total-thread settings migration and retained warm sessions on save + failure; real Google CPU overlap regression passed on Windows. + - [ ] Measure representative mixed-media contention, single-thread latency, + extraction/decoder/ANN work and native UI p95 before closing this gate. - [ ] Attribute startup cost; measure smaller query context and selective encoders. - [ ] Buffered bounded media IPC and linear UTF-8 source-window offset processing. - [ ] Evaluate 256d MRL and content-specific vision budgets on independent data.