feat: switch inference engine from tract-onnx to ort-tract (#3)

Cargo.lock

@@ -1886,6 +1886,7 @@ dependencies = [
  "macros",
  "maxminddb",
  "mime_guess",
+ "ndarray 0.17.2",
  "network-types",
  "ort",
  "ort-tract",

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"
 

net-guardia/src/core/infrastructure/mod.rs

@@ -36,7 +36,7 @@ impl AppServices {
     pub fn new(app_config: Arc<AppConfig>, inference_config: Arc<InferenceConfig>) -> Result<Self, Error> {
         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 {

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<f32>) -> TractResult<f32> {
-        let result = self
-            .models
-            .deep_autoencoder
-            .run(tvec![input.clone().into_tensor().into()])?;
+    fn run_autoencoder(&self, input: &Array3<f32>) -> Result<f32, String> {
+        let mut session = self.models.deep_autoencoder.lock().map_err(|_| {
+            log!(MLError::SessionLockPoisoned);
+            MLError::SessionLockPoisoned.to_string()
+        })?;
 
-        let output = result[0]
-            .to_array_view::<f32>()?
-            .into_dimensionality::<tract_ndarray::Ix3>()?;
+        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::<f32>().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;
 

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<RunnableModel>,
 }
 
 impl MLModels {
-    pub fn load_models(
-        app_config: &Arc<AppConfig>,
-        inference_config: &Arc<InferenceConfig>,
-    ) -> Result<Self, MLError> {
+    pub fn load_models(app_config: &Arc<AppConfig>) -> Result<Self, MLError> {
+        // 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<RunnableModel, MLError> {
+    fn load_lstm_ae(model_name: &str) -> Result<Session, MLError> {
         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(),

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,

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<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>;
+pub type RunnableModel = Session;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ClipParams {
@@ -158,4 +158,4 @@ impl AlertMessage {
             ae_score: result.ae_score,
         }
     }
-}
+}