karl/searchi
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases 10 Wiki Activity

Coordinate foreground search and background inference CPU budgets

Browse Source
This commit is contained in:
karl committed 2026-10-13 01:56:03 +02:00
1 parent 6b78c8bb7f
commit d85746db5b
12 files changed
+858 -60
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

apps/desktop/src-tauri/src/main.rs
+1 -3
View File
@@ -113,11 +113,9 @@ fn set_cpu_threads(threads: usize, state: State<'_, AppState>) -> Result<Setting
let mut next = store.value.clone();
next.cpu_threads = threads;
next.validate()?;
state.semantic.reset_session()?;
store.save(next)?;
state
.semantic
.begin_automatic_setup(state.database.clone(), threads);
.change_resources(state.database.clone(), threads, || store.save(next))?;
Ok(store.value.clone())
}
apps/desktop/src-tauri/src/semantic_commands.rs
+191 -41
View File
@@ -2,7 +2,7 @@
use crate::{measured_outcome, model_download, AppState, Filters, MeasuredFailure};
use searchi_core::metrics::{MetricsMode, SearchResponse, TraceCollector};
use searchi_core::SearchResult;
use searchi_service::{EmbeddingReport, HybridSearch};
use searchi_service::{EmbeddingReport, HybridSearch, InferenceBudget, InferenceRole};
use serde::Serialize;
use std::{
path::PathBuf,
@@ -26,6 +26,8 @@ pub struct SemanticState {
error: Mutex<Option<String>>,
service: Mutex<Option<HybridSearch>>,
prepared_graphs: searchi_service::PreparedGraphs,
inference_budget: InferenceBudget,
lifecycle: Mutex<()>,
media_tools: Option<searchi_media::FfmpegTools>,
media_setup_error: Option<String>,
}
@@ -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<Self>, 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<Self>, 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<Self>,
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<String>) {
if let Ok(mut error) = self.error.lock() {
*error = value;
@@ -188,6 +254,9 @@ pub fn semantic_status(state: State<'_, AppState>) -> Result<SemanticStatus, Str
#[tauri::command]
pub async fn setup_model(state: State<'_, AppState>) -> 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<EmbeddingReport, String> {
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<EmbeddingRep
return Err("Semantic indexing is already running.".into());
}
let database = state.database.clone();
drop(admission);
tauri::async_runtime::spawn_blocking(move || {
let _guard = WorkGuard(&semantic.embedding);
semantic.record_error(None);
// Indexing uses its own isolated worker; interactive queries keep their own session.
let result = HybridSearch::open(
// Both workers share one clamped budget. Admission is per encoder call,
// so waiting queries stop the next background batch without a service lock.
let threads = semantic.inference_budget.worker_threads(InferenceRole::Background);
let result = (|| {
let admission = semantic.inference_budget.acquire(InferenceRole::Background, threads)?;
let service = HybridSearch::open(
&database,
&semantic.model_directory,
&semantic.worker,
threads,
)
.and_then(|mut service| {
)?;
drop(admission);
let mut service = service.with_budget(&semantic.inference_budget, InferenceRole::Background, threads);
service.share_prepared_graphs(&semantic.prepared_graphs)?;
let tools = semantic.media_tools.as_ref().ok_or_else(|| {
anyhow::anyhow!("The media decoder is missing from this installation. Repair or reinstall searchi.")
})?;
service.index_embeddings_with_media(tools, &searchi_media::MediaLimits::default())
})
})()
.map_err(|error| error.to_string());
// The publisher shares its prepared immutable ANN before activation.
// Keep the query model and its cached vectors warm throughout indexing.
@@ -270,18 +354,16 @@ pub async fn hybrid_search(
state: State<'_, AppState>,
) -> Result<Vec<SearchResult>, 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));
}
}
apps/desktop/src-tauri/src/settings.rs
+34 -4
View File
@@ -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<Self, String> {
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);
}
}