fix(ml): enforce 5s wall-clock budget on ONNX load

A malformed or maliciously-crafted ONNX can wedge tract's graph solver
during `model_for_path` / `into_optimized`, blocking the model-watcher
reload path (and the boot path via `AppServices::new`) indefinitely.
Introduce `ONNX_LOAD_TIMEOUT = 5s` and wrap `loader_inner` in a helper
that runs it on a dedicated OS thread with `mpsc::recv_timeout`:

- `load_with_timeout<F, R>(path, budget, f)` — spawns `f` on a fresh
  thread, waits `budget`, returns the loader's `Result<R, MLError>` or
  `MLError::ModelLoadTimeout { path, seconds }` when the worker misses
  the deadline. The worker thread is detached on timeout and allowed
  to finish on its own; the leaked thread's allocations drop when it
  completes. That cost is fine for a low-frequency action gated behind
  admin upload + manifest validation.
- A `Disconnected` channel error (sender dropped before send) maps to
  `ModelLoadFailed` with a clear detail string, so panics inside the
  loader surface as bounded failures instead of hangs.
- New `MLError::ModelLoadTimeout { path, seconds }` variant with a
  dedicated tracing level so reload dashboards can colour it distinctly.

`build_adapter` stays synchronous — the timeout lives in the sync path,
so neither `AppServices::new` (sync bootstrap) nor the model watcher
(sync reload callback) needs an async refactor to benefit.

Tests: 3 new (fast loader passes through, timeout surfaces
`ModelLoadTimeout` with preserved path, real loader errors propagate
unchanged without being swallowed by the timeout wrapper) — 210 pass
total. clippy --package net-guardia -- -D warnings clean.

Closes I-region I-4.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
DaLaw2 2026-04-18 17:14:52 +08:00
parent 8aeadc7c64
commit b8900518d1
2 changed files with 87 additions and 3 deletions

View File

@ -7,8 +7,9 @@
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Instant;
use std::sync::{Arc, mpsc};
use std::thread;
use std::time::{Duration, Instant};
use macros::log;
use tract_onnx::prelude::*;
@ -23,6 +24,13 @@ use crate::model::error::ml::MLError;
use crate::model::log::ml::MLLog;
use crate::model::system::config::MLInferenceConfig;
/// Wall-clock budget for a single ONNX parse + optimize + runnable chain.
/// A malformed or maliciously-crafted model can wedge tract's graph solver;
/// the timeout keeps an admin-triggered upload from blocking the watcher
/// indefinitely. Five seconds is generous for the models that currently
/// ship (<10MB) while still bounding pathological inputs.
const ONNX_LOAD_TIMEOUT: Duration = Duration::from_secs(5);
/// Build an `MLModelAdapter` by loading the ONNX file(s) the manifest names,
/// validating shape against the inference config's feature counts, and
/// wrapping the underlying `RunnableModel`s in `Arc` for zero-copy swap.
@ -124,13 +132,42 @@ fn loader(model_path: &Path, model_name: &str, features: usize, batch_size: usiz
log!(MLLog::ModelLoading(model_name.to_string(), features, batch_size));
let start = Instant::now();
let result = loader_inner(model_path, model_name, features, batch_size);
let path_for_thread = model_path.to_path_buf();
let name_for_thread = model_name.to_string();
let result = load_with_timeout(model_path.to_path_buf(), ONNX_LOAD_TIMEOUT, move || {
loader_inner(&path_for_thread, &name_for_thread, features, batch_size)
});
let elapsed_ms = start.elapsed().as_millis() as u64;
log!(MLLog::ModelLoadComplete(model_name.to_string(), elapsed_ms));
result
}
/// Run `f` on a dedicated OS thread with a wall-clock cap. Prevents a
/// pathological ONNX from wedging `tract`'s graph solver and holding up
/// the watcher / bootstrap indefinitely. On timeout the worker thread
/// is detached — it will finish on its own and drop its state; the cost
/// of one leaked thread is acceptable for a low-frequency operation
/// gated behind admin upload + manifest validation.
fn load_with_timeout<F, R>(path: PathBuf, budget: Duration, f: F) -> Result<R, MLError>
where
F: FnOnce() -> Result<R, MLError> + Send + 'static,
R: Send + 'static,
{
let (tx, rx) = mpsc::channel();
thread::spawn(move || {
let _ = tx.send(f());
});
match rx.recv_timeout(budget) {
Ok(result) => result,
Err(mpsc::RecvTimeoutError::Timeout) => Err(MLError::ModelLoadTimeout(path, budget.as_secs())),
Err(mpsc::RecvTimeoutError::Disconnected) => Err(MLError::ModelLoadFailed(
path,
"loader thread disconnected before completing".to_string(),
)),
}
}
fn loader_inner(
model_path: &Path,
model_name: &str,
@ -206,6 +243,49 @@ mod tests {
assert_eq!(dim, 32, "v10 classifier ONNX input dim changed unexpectedly");
}
#[test]
fn load_with_timeout_passes_fast_loader() {
let result: Result<u32, MLError> =
load_with_timeout(PathBuf::from("/tmp/fast.onnx"), Duration::from_millis(500), || Ok(42));
assert_eq!(result.ok(), Some(42));
}
#[test]
fn load_with_timeout_returns_timeout_error_when_budget_exceeded() {
let path = PathBuf::from("/tmp/slow.onnx");
let result: Result<u32, MLError> = load_with_timeout(path.clone(), Duration::from_millis(50), || {
thread::sleep(Duration::from_millis(500));
Ok(42)
});
match result {
Err(MLError::ModelLoadTimeout {
path: got_path,
seconds,
}) => {
assert_eq!(got_path, path);
assert_eq!(seconds, 0, "budget < 1s rounds to 0 on `as_secs`");
}
other => panic!("expected ModelLoadTimeout, got {other:?}"),
}
}
#[test]
fn load_with_timeout_propagates_loader_error() {
// Failures surface unchanged; timeout wrapping must not swallow them.
let path = PathBuf::from("/tmp/bad.onnx");
let err_path = path.clone();
let result: Result<u32, MLError> = load_with_timeout(path, Duration::from_millis(500), move || {
Err(MLError::ModelLoadFailed(
err_path,
"synthetic parse failure".to_string(),
))
});
match result {
Err(MLError::ModelLoadFailed { err, .. }) => assert!(err.contains("synthetic")),
other => panic!("expected ModelLoadFailed, got {other:?}"),
}
}
/// Smoke: the shipped manifest + sidecar must load and produce a
/// `MultiTask` adapter with both underlying models.
#[test]

View File

@ -27,6 +27,10 @@ traceable! {
#[error("Feature count mismatch for {model:?}: manifest declares {declared}, ONNX input expects {onnx_dim}")]
FeatureMismatch { model: PathBuf, declared: usize, onnx_dim: usize } => tracing::Level::ERROR,
#[no_source]
#[error("Model load timed out after {seconds}s: {path:?}")]
ModelLoadTimeout { path: PathBuf, seconds: u64 } => tracing::Level::ERROR,
#[no_source]
#[error("Unknown feature '{name}' — not registered in FEATURE_REGISTRY")]
UnknownFeature { name: String } => tracing::Level::ERROR,