mirror of
https://github.com/DaLaw2/NetGuardia.git
synced 2026-08-24 14:10:28 +09:00
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:
parent
8aeadc7c64
commit
b8900518d1
@ -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]
|
||||
|
||||
@ -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,
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user