Add Model + ModelWrapper abstraction to split wrapper state from policy (#64)

WIP

Author: tgolsson

CHANGELOG.md

@@ -9,6 +9,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 <!-- next-header -->
 ## [Unreleased] - ReleaseDate
+
+- Breaking: `Inferer.begin_agent` and `Inferer.end_agent` now take
+  `&self`, changed from a mutable reference.
+
+### Wrapper rework
+
+To support a wider variety of uses, we have implemented a new category
+of wrappers that do not require ownership of the inferer. This allows
+for more flexible usage patterns, where the inferer policy can be
+replaced in a live application without losing any state kept in
+wrappers.
+
+This change is currently non-breaking and is implemented separately
+from the old wrapper system.
+
 ## [0.8.0] - 2025-05-28
 - Added a new `RecurrentTracker` wrapper to handle recurrent
   inputs/outputs if the recurrent data is only needed durign network

Cargo.lock

@@ -1,7 +1,9 @@
 use anyhow::{bail, Result};
 use cervo::asset::AssetData;
-use cervo::core::prelude::{Batcher, Inferer, InfererExt, State};
-use cervo::core::recurrent::{RecurrentInfo, RecurrentTracker};
+use cervo::core::epsilon::EpsilonInjectorWrapper;
+use cervo::core::prelude::{Batcher, Inferer, State};
+use cervo::core::recurrent::{RecurrentInfo, RecurrentTrackerWrapper};
+use cervo::core::wrapper::{BaseWrapper, InfererWrapper, InfererWrapperExt};
 use clap::Parser;
 use clap::ValueEnum;
 use serde::Serialize;
@@ -167,7 +169,7 @@ struct Record {
     total: f64,
 }
 
-fn execute_load_metrics<I: Inferer>(
+fn execute_load_metrics<I: Inferer + 'static>(
     batch_size: usize,
     data: HashMap<u64, State<'_>>,
     count: usize,
@@ -222,85 +224,112 @@ pub fn build_inputs_from_desc(
         .collect()
 }
 
-fn do_run(mut inferer: impl Inferer, batch_size: usize, config: &Args) -> Result<Record> {
-    let shapes = inferer.input_shapes().to_vec();
-    let observations = build_inputs_from_desc(batch_size as u64, &shapes);
-    for id in 0..batch_size {
-        inferer.begin_agent(id as u64);
-    }
-    let res = execute_load_metrics(batch_size, observations, config.count, &mut inferer)?;
-    for id in 0..batch_size {
-        inferer.end_agent(id as u64);
+fn do_run(
+    wrapper: impl InfererWrapper + 'static,
+    inferer: impl Inferer + 'static,
+    config: &Args,
+) -> Result<Vec<Record>> {
+    let mut model = wrapper.wrap(Box::new(inferer) as Box<dyn Inferer>);
+
+    let mut records = Vec::with_capacity(config.batch_sizes.len());
+    for batch_size in config.batch_sizes.clone() {
+        let mut reader = File::open(&config.file)?;
+        let inferer = if cervo::nnef::is_nnef_tar(&config.file) {
+            cervo::nnef::builder(&mut reader).build_fixed(&[batch_size])?
+        } else {
+            match config.file.extension().and_then(|ext| ext.to_str()) {
+                Some("onnx") => cervo::onnx::builder(&mut reader).build_fixed(&[batch_size])?,
+                Some("crvo") => AssetData::deserialize(&mut reader)?.load_fixed(&[batch_size])?,
+                Some(other) => bail!("unknown file type {:?}", other),
+                None => bail!("missing file extension {:?}", config.file),
+            }
+        };
+
+        model = model
+            .with_new_inferer(Box::new(inferer) as Box<dyn Inferer>)
+            .map_err(|(_, e)| e)?;
+
+        let shapes = model.input_shapes().to_vec();
+        let observations = build_inputs_from_desc(batch_size as u64, &shapes);
+        for id in 0..batch_size {
+            model.begin_agent(id as u64);
+        }
+        let res = execute_load_metrics(batch_size, observations, config.count, &mut model)?;
+
+        // Print Text
+        if matches!(config.output, OutputFormat::Text) {
+            println!(
+                "Batch Size {}: {:.2} ms ± {:.2} per element, {:.2} ms total",
+                res.batch_size, res.mean, res.stddev, res.total,
+            );
+        }
+
+        records.push(res);
+        for id in 0..batch_size {
+            model.end_agent(id as u64);
+        }
     }
 
-    Ok(res)
+    Ok(records)
 }
 
 fn run_apply_epsilon_config(
-    inferer: impl Inferer,
-    batch_size: usize,
+    wrapper: impl InfererWrapper + 'static,
+    inferer: impl Inferer + 'static,
     config: &Args,
-) -> Result<Record> {
+) -> Result<Vec<Record>> {
     if let Some(epsilon) = config.with_epsilon.as_ref() {
-        let inferer = inferer.with_default_epsilon(epsilon)?;
-        do_run(inferer, batch_size, config)
+        let wrapper = EpsilonInjectorWrapper::wrap(wrapper, &inferer, epsilon)?;
+        do_run(wrapper, inferer, config)
     } else {
-        do_run(inferer, batch_size, config)
+        do_run(wrapper, inferer, config)
     }
 }
 
-fn run_apply_recurrent(inferer: impl Inferer, batch_size: usize, config: &Args) -> Result<Record> {
+fn run_apply_recurrent(
+    wrapper: impl InfererWrapper + 'static,
+    inferer: impl Inferer + 'static,
+    config: &Args,
+) -> Result<Vec<Record>> {
     if let Some(recurrent) = config.recurrent.as_ref() {
         if matches!(recurrent, RecurrentConfig::None) {
-            run_apply_epsilon_config(inferer, batch_size, config)
+            run_apply_epsilon_config(wrapper, inferer, config)
         } else {
-            let inferer = match recurrent {
+            let wrapper = match recurrent {
                 RecurrentConfig::None => unreachable!(),
-                RecurrentConfig::Auto => RecurrentTracker::wrap(inferer),
+                RecurrentConfig::Auto => RecurrentTrackerWrapper::wrap(wrapper, &inferer),
                 RecurrentConfig::Mapped(map) => {
                     let infos = map
                         .iter()
                         .cloned()
                         .map(|(inkey, outkey)| RecurrentInfo { inkey, outkey })
                         .collect::<Vec<_>>();
-                    RecurrentTracker::new(inferer, infos)
+                    RecurrentTrackerWrapper::new(wrapper, &inferer, infos)
                 }
             }?;
 
-            run_apply_epsilon_config(inferer, batch_size, config)
+            run_apply_epsilon_config(wrapper, inferer, config)
         }
     } else {
-        run_apply_epsilon_config(inferer, batch_size, config)
+        run_apply_epsilon_config(wrapper, inferer, config)
     }
 }
 
 pub(super) fn run(config: Args) -> Result<()> {
-    let mut records: Vec<Record> = Vec::new();
-    for batch_size in config.batch_sizes.clone() {
-        let mut reader = File::open(&config.file)?;
-        let inferer = if cervo::nnef::is_nnef_tar(&config.file) {
-            cervo::nnef::builder(&mut reader).build_fixed(&[batch_size])?
-        } else {
-            match config.file.extension().and_then(|ext| ext.to_str()) {
-                Some("onnx") => cervo::onnx::builder(&mut reader).build_fixed(&[batch_size])?,
-                Some("crvo") => AssetData::deserialize(&mut reader)?.load_fixed(&[batch_size])?,
-                Some(other) => bail!("unknown file type {:?}", other),
-                None => bail!("missing file extension {:?}", config.file),
-            }
-        };
-
-        let record = run_apply_recurrent(inferer, batch_size, &config)?;
-
-        // Print Text
-        if matches!(config.output, OutputFormat::Text) {
-            println!(
-                "Batch Size {}: {:.2} ms ± {:.2} per element, {:.2} ms total",
-                record.batch_size, record.mean, record.stddev, record.total,
-            );
+    let mut reader = File::open(&config.file)?;
+    let inferer = if cervo::nnef::is_nnef_tar(&config.file) {
+        cervo::nnef::builder(&mut reader).build_basic()?
+    } else {
+        match config.file.extension().and_then(|ext| ext.to_str()) {
+            Some("onnx") => cervo::onnx::builder(&mut reader).build_basic()?,
+            Some("crvo") => AssetData::deserialize(&mut reader)?.load_basic()?,
+            Some(other) => bail!("unknown file type {:?}", other),
+            None => bail!("missing file extension {:?}", config.file),
         }
+    };
+
+    let records = run_apply_recurrent(BaseWrapper, inferer, &config)?;
 
-        records.push(record);
-    }
     // Print JSON
     if matches!(config.output, OutputFormat::Json) {
         let json = serde_json::to_string_pretty(&records)?;

crates/cervo-cli/src/commands/benchmark.rs

@@ -84,15 +84,7 @@ pub(super) fn run(config: Args) -> Result<()> {
 
     let elapsed = if let Some(epsilon) = config.with_epsilon.as_ref() {
         let inferer = inferer.with_default_epsilon(epsilon)?;
-        // TODO[TSolberg]: Issue #31.
-        let shapes = inferer
-            .raw_input_shapes()
-            .iter()
-            .filter(|(k, _)| k.as_str() != epsilon)
-            .cloned()
-            .collect::<Vec<_>>();
-
-        let observations = build_inputs_from_desc(config.batch_size as u64, &shapes);
+        let observations = build_inputs_from_desc(config.batch_size as u64, inferer.input_shapes());
 
         if config.print_input {
             print_input(&observations);