From b8900518d159dd0ebf6bcada3a9bb711bd87953c Mon Sep 17 00:00:00 2001 From: DaLaw2 Date: Sat, 18 Apr 2026 17:14:52 +0800 Subject: [PATCH] fix(ml): enforce 5s wall-clock budget on ONNX load MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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(path, budget, f)` — spawns `f` on a fresh thread, waits `budget`, returns the loader's `Result` 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) --- net-guardia/src/core/ml/model_loader.rs | 86 ++++++++++++++++++++++++- net-guardia/src/model/error/ml.rs | 4 ++ 2 files changed, 87 insertions(+), 3 deletions(-) diff --git a/net-guardia/src/core/ml/model_loader.rs b/net-guardia/src/core/ml/model_loader.rs index 5c4df04..2ca6dfa 100644 --- a/net-guardia/src/core/ml/model_loader.rs +++ b/net-guardia/src/core/ml/model_loader.rs @@ -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(path: PathBuf, budget: Duration, f: F) -> Result +where + F: FnOnce() -> Result + 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 = + 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 = 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 = 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] diff --git a/net-guardia/src/model/error/ml.rs b/net-guardia/src/model/error/ml.rs index 004a7c9..447cda8 100644 --- a/net-guardia/src/model/error/ml.rs +++ b/net-guardia/src/model/error/ml.rs @@ -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,