build: package universal Windows runtime and verify automatic CPU fallback
This commit is contained in:
1 parent
0fcc8ec21c
commit
278b09f82c
16 files changed
+695
-34
No files matched your search
@@ -78,6 +78,27 @@ impl Drop for GemmaEmbedder {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct StartupTransportFailure(String);
|
||||
impl std::fmt::Display for StartupTransportFailure {
|
||||
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
formatter.write_str(&self.0)
|
||||
}
|
||||
}
|
||||
impl std::error::Error for StartupTransportFailure {}
|
||||
|
||||
/// Retry only process/transport failure, never an incompatible protocol or model.
|
||||
fn start_with_cpu_fallback<T>(cpu: bool, mut start: impl FnMut(bool) -> Result<T>) -> Result<T> {
|
||||
match start(cpu) {
|
||||
Err(error) if !cpu && error.downcast_ref::<StartupTransportFailure>().is_some() => {
|
||||
start(true).with_context(|| {
|
||||
format!("CPU initialization failed after automatic worker startup failure: {error}")
|
||||
})
|
||||
}
|
||||
result => result,
|
||||
}
|
||||
}
|
||||
|
||||
impl GemmaEmbedder {
|
||||
/// Load the verified Google model and explicitly supplied local worker.
|
||||
/// Production selects available acceleration with CPU fallback. Tests set
|
||||
@@ -113,7 +134,16 @@ impl GemmaEmbedder {
|
||||
);
|
||||
let force_cpu = std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1");
|
||||
let cpu = force_cpu || matches!(device, Device::Cpu);
|
||||
let mut command = Command::new(&worker);
|
||||
start_with_cpu_fallback(cpu, |cpu| {
|
||||
Self::start_verified(&model_dir, &worker, threads, cpu)
|
||||
})
|
||||
}
|
||||
|
||||
// Both attempts share the already validated model/worker paths. CPU fallback
|
||||
// does not repeat the expensive hash check or weaken the handshake contract.
|
||||
fn start_verified(model_dir: &Path, worker: &Path, threads: usize, cpu: bool) -> Result<Self> {
|
||||
let model = model_dir.join(MODEL_FILE);
|
||||
let mut command = Command::new(worker);
|
||||
command
|
||||
.args(["--model"])
|
||||
.arg(model)
|
||||
@@ -132,9 +162,14 @@ impl GemmaEmbedder {
|
||||
use std::os::windows::process::CommandExt;
|
||||
command.creation_flags(0x08000000); // CREATE_NO_WINDOW
|
||||
}
|
||||
let mut child = command.spawn().context("Cannot start local Gemma worker")?;
|
||||
let stdin = child.stdin.take().context("Worker stdin missing")?;
|
||||
let stdout = child.stdout.take().context("Worker stdout missing")?;
|
||||
let mut child = command.spawn().map_err(|error| {
|
||||
StartupTransportFailure(format!("Cannot start local Gemma worker: {error}"))
|
||||
})?;
|
||||
let (Some(stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
anyhow::bail!("Gemma worker protocol pipes missing");
|
||||
};
|
||||
let (sender, responses) = mpsc::sync_channel(1);
|
||||
thread::spawn(move || {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
@@ -151,7 +186,16 @@ impl GemmaEmbedder {
|
||||
other => {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
anyhow::bail!("Gemma worker did not initialize: {other:?}");
|
||||
// Malformed/oversized JSON is a protocol violation, not evidence
|
||||
// that hardware acceleration failed. Do not mask it with retry.
|
||||
if matches!(&other, Ok(Err(error)) if error.starts_with("Invalid Gemma worker JSON:") || error == "Gemma worker response exceeds protocol byte budget")
|
||||
{
|
||||
anyhow::bail!("Invalid Gemma worker startup protocol: {other:?}");
|
||||
}
|
||||
return Err(StartupTransportFailure(format!(
|
||||
"Gemma worker did not initialize: {other:?}"
|
||||
))
|
||||
.into());
|
||||
}
|
||||
};
|
||||
let info: WorkerInfo = match serde_json::from_value(ready) {
|
||||
@@ -181,8 +225,8 @@ impl GemmaEmbedder {
|
||||
info,
|
||||
healthy: true,
|
||||
force_cpu: cpu,
|
||||
model_dir,
|
||||
worker_path: worker,
|
||||
model_dir: model_dir.to_path_buf(),
|
||||
worker_path: worker.to_path_buf(),
|
||||
projector_verified: false,
|
||||
query_profile: None,
|
||||
})
|
||||
@@ -666,6 +710,73 @@ fn verify_artifact(path: &Path, bytes: u64, sha: &str) -> Result<()> {
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn automatic_startup_transport_failure_retries_cpu_once() {
|
||||
let mut attempts = Vec::new();
|
||||
let result = start_with_cpu_fallback(false, |cpu| {
|
||||
attempts.push(cpu);
|
||||
if cpu {
|
||||
Ok("cpu ready")
|
||||
} else {
|
||||
Err(StartupTransportFailure("worker closed its output before ready".into()).into())
|
||||
}
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(result, "cpu ready");
|
||||
assert_eq!(attempts, [false, true]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_failure_never_loops_or_retries_explicit_cpu() {
|
||||
for cpu in [false, true] {
|
||||
let mut attempts = Vec::new();
|
||||
let result: Result<()> = start_with_cpu_fallback(cpu, |selected_cpu| {
|
||||
attempts.push(selected_cpu);
|
||||
Err(StartupTransportFailure("worker exited".into()).into())
|
||||
});
|
||||
let error = format!("{:#}", result.unwrap_err());
|
||||
assert_eq!(attempts, if cpu { vec![true] } else { vec![false, true] });
|
||||
assert!(error.contains("worker exited"));
|
||||
if !cpu {
|
||||
assert!(error.contains("CPU initialization failed"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_startup_protocol_is_not_retried() {
|
||||
let mut attempts = 0;
|
||||
let result: Result<()> = start_with_cpu_fallback(false, |_| {
|
||||
attempts += 1;
|
||||
let mut response = BufReader::new(&b"not JSON\n"[..]);
|
||||
let error = read_response(&mut response).unwrap_err();
|
||||
assert!(error.starts_with("Invalid Gemma worker JSON:"));
|
||||
anyhow::bail!("Invalid Gemma worker startup protocol: {error}");
|
||||
});
|
||||
assert!(result.unwrap_err().to_string().contains("startup protocol"));
|
||||
assert_eq!(attempts, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn incompatible_startup_contract_is_not_retried() {
|
||||
let mut attempts = 0;
|
||||
let result: Result<()> = start_with_cpu_fallback(false, |_| {
|
||||
attempts += 1;
|
||||
let handshake = json!({"ready":true,"model":"other-model","dimensions":768,"device":"cpu","threads":1,"runtime":"other-runtime"});
|
||||
let info: WorkerInfo = serde_json::from_value(handshake)?;
|
||||
ensure!(
|
||||
info.model == "google/embeddinggemma-2",
|
||||
"Gemma worker handshake violates configured model/device contract"
|
||||
);
|
||||
Ok(())
|
||||
});
|
||||
assert!(result
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("model/device contract"));
|
||||
assert_eq!(attempts, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn old_text_worker_handshake_does_not_claim_media_support() {
|
||||
let mut handshake = json!({"ready":true,"model":"google/embeddinggemma-2","dimensions":768,"device":"cpu","threads":1,"runtime":"b11461-no-logits-v1"});
|
||||
|
||||
Reference in new issue
Block a user