diff --git a/Cargo.lock b/Cargo.lock index 3438e2e..06ba71a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1886,6 +1886,7 @@ dependencies = [ "macros", "maxminddb", "mime_guess", + "ndarray 0.17.2", "network-types", "ort", "ort-tract", diff --git a/net-guardia/Cargo.toml b/net-guardia/Cargo.toml index a66c575..3ed697f 100644 --- a/net-guardia/Cargo.toml +++ b/net-guardia/Cargo.toml @@ -43,6 +43,7 @@ rusqlite = { version = "0.32", features = ["bundled"] } ort-tract = "0.3.0+0.22" ort = { version = "=2.0.0-rc.12", default-features = false, features = ["std", "ndarray", "alternative-backend"] } +ndarray = "0.17" #csv = "1.4.0" #anyhow = "1.0.100" diff --git a/net-guardia/src/core/infrastructure/mod.rs b/net-guardia/src/core/infrastructure/mod.rs index 9818dbe..4d853a4 100644 --- a/net-guardia/src/core/infrastructure/mod.rs +++ b/net-guardia/src/core/infrastructure/mod.rs @@ -36,7 +36,7 @@ impl AppServices { pub fn new(app_config: Arc, inference_config: Arc) -> Result { let health = SystemHealth::new(app_config.clone())?; - let ml_models = Arc::new(MLModels::load_models(&app_config, &inference_config)?); + let ml_models = Arc::new(MLModels::load_models(&app_config)?); let ml_alert = Arc::new(MLAlert::new()); let traffic_logger = if app_config.traffic_logging_mode { diff --git a/net-guardia/src/detection/ml/inference.rs b/net-guardia/src/detection/ml/inference.rs index 900d386..dc0ebf9 100644 --- a/net-guardia/src/detection/ml/inference.rs +++ b/net-guardia/src/detection/ml/inference.rs @@ -1,9 +1,10 @@ use std::collections::VecDeque; use std::collections::HashMap; -use std::sync::Mutex; +use std::sync::{Arc, Mutex}; use macros::log; -use tract_onnx::prelude::*; +use ndarray::Array3; +use ort::{inputs, value::TensorRef}; use super::config_loader::InferenceConfig; use super::feature_extractor::FlowFeatures; @@ -59,7 +60,7 @@ impl Inference { if buf.len() > window_size { buf.pop_front(); } - + println!( "Buffer [{}->{}]: {}/{} | contents: {:?}", flow.flow_key.src_ip, @@ -87,7 +88,7 @@ impl Inference { // Build 3D tensor (1, window_size, features) let feat_len = self.config.num_ae_features(); - let ae_input = tract_ndarray::Array3::from_shape_fn( + let ae_input = Array3::from_shape_fn( (1, window_size, feat_len), |(_, t, f)| sequence[t][f], ); @@ -132,19 +133,21 @@ impl Inference { features.features.iter().map(|&x| x as f32).collect() } - fn run_autoencoder(&self, input: &tract_ndarray::Array3) -> TractResult { - let result = self - .models - .deep_autoencoder - .run(tvec![input.clone().into_tensor().into()])?; + fn run_autoencoder(&self, input: &Array3) -> Result { + let mut session = self.models.deep_autoencoder.lock().map_err(|_| { + log!(MLError::SessionLockPoisoned); + MLError::SessionLockPoisoned.to_string() + })?; - let output = result[0] - .to_array_view::()? - .into_dimensionality::()?; + let outputs = session + .run(inputs![TensorRef::from_array_view(input).map_err(|e| e.to_string())?]) + .map_err(|e| e.to_string())?; + + let output = outputs[0].try_extract_array::().map_err(|e| e.to_string())?; // MSE between input and reconstructed output - let input_view = input.view(); - let diff = &input_view - &output; + let input_dyn = input.view().into_dyn(); + let diff = &input_dyn - &output; let n = (self.config.window_size * self.config.num_ae_features()) as f32; let mse = (&diff * &diff).sum() / n; diff --git a/net-guardia/src/detection/ml/model_loader.rs b/net-guardia/src/detection/ml/model_loader.rs index 4625a1f..61d85e7 100644 --- a/net-guardia/src/detection/ml/model_loader.rs +++ b/net-guardia/src/detection/ml/model_loader.rs @@ -1,55 +1,46 @@ -use tract_onnx::prelude::*; use std::path::PathBuf; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; + +use ort::session::{Session, builder::GraphOptimizationLevel}; use crate::core::infrastructure::app_config::AppConfig; use crate::model::error::ml::MLError; use crate::model::ml_detection::RunnableModel; -use super::config_loader::InferenceConfig; - pub struct MLModels { - pub deep_autoencoder: RunnableModel, + pub deep_autoencoder: Mutex, } impl MLModels { - pub fn load_models( - app_config: &Arc, - inference_config: &Arc, - ) -> Result { + pub fn load_models(app_config: &Arc) -> Result { + // Register ort-tract as the ORT execution backend before creating any sessions. + // Returns false only if already initialized, which is harmless. + ort::set_api(ort_tract::api()); + Ok(Self { - deep_autoencoder: Self::load_lstm_ae( - &app_config.deep_autoencoder_name, - inference_config.window_size, - inference_config.num_ae_features(), - )?, + deep_autoencoder: Mutex::new(Self::load_lstm_ae(&app_config.deep_autoencoder_name)?), }) } - fn load_lstm_ae(model_name: &str, window_size: usize, features: usize) -> Result { + fn load_lstm_ae(model_name: &str) -> Result { let model_path = PathBuf::from(env!("ARTIFACTCS_PATH")).join(model_name); - let mut model = onnx() - .model_for_path(&model_path) - .map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })?; - - // LSTM AE: input shape = (batch=1, seq_len=window_size, features) - model - .set_input_fact(0, f32::fact(&[1, window_size, features]).into()) - .map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })?; - - model - .into_optimized() + Session::builder() .map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })? - .into_runnable() + .with_optimization_level(GraphOptimizationLevel::Disable) + .map_err(|_| MLError::ModelLoadFailed { path: model_path.clone() })? + .commit_from_file(&model_path) .map_err(|_| MLError::ModelLoadFailed { path: model_path }) } pub fn get_model_info(&self, name: &str) -> String { match name { "deep_autoencoder" => { - let inputs = self.deep_autoencoder.model().inputs.len(); - let outputs = self.deep_autoencoder.model().outputs.len(); + let Ok(session) = self.deep_autoencoder.lock() else { + return format!("{}: lock poisoned", name); + }; + let inputs = session.inputs().len(); + let outputs = session.outputs().len(); format!("{}: inputs: {}, outputs: {}", name, inputs, outputs) } _ => "unknown model".to_string(), diff --git a/net-guardia/src/model/error/ml.rs b/net-guardia/src/model/error/ml.rs index 3b547d2..648e122 100644 --- a/net-guardia/src/model/error/ml.rs +++ b/net-guardia/src/model/error/ml.rs @@ -28,6 +28,10 @@ traceable! { #[error("Inference flow buffers mutex poisoned; skipping operation")] InferenceLockPoisoned => tracing::Level::ERROR, + #[no_source] + #[error("Inference session mutex poisoned; skipping operation")] + SessionLockPoisoned => tracing::Level::ERROR, + #[no_source] #[error("Failed to archive CSV '{path}': {reason}")] TrafficLogArchiveFailed { path: String, reason: String } => tracing::Level::WARN, diff --git a/net-guardia/src/model/ml_detection.rs b/net-guardia/src/model/ml_detection.rs index bd5d9a1..51bf809 100644 --- a/net-guardia/src/model/ml_detection.rs +++ b/net-guardia/src/model/ml_detection.rs @@ -1,11 +1,11 @@ use common::model::event::{Event, TcpFlags}; +use ort::session::Session; use serde::{Deserialize, Serialize}; -use tract_onnx::prelude::{Graph, SimplePlan, TypedFact, TypedOp}; use crate::model::direction::Direction; use crate::utils::packet_parser::{format_ipv4, format_ipv6}; -pub type RunnableModel = SimplePlan, Graph>>; +pub type RunnableModel = Session; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ClipParams { @@ -158,4 +158,4 @@ impl AlertMessage { ae_score: result.ae_score, } } -} +} \ No newline at end of file