docs,changelog, cleanup

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

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

@@ -1,9 +1,9 @@
 use anyhow::{bail, Result};
 use cervo::asset::AssetData;
 use cervo::core::epsilon::EpsilonInjectorWrapper;
-use cervo::core::model::{BaseCase, Model, ModelWrapper};
 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;
@@ -169,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,
@@ -225,11 +225,11 @@ pub fn build_inputs_from_desc(
 }
 
 fn do_run(
-    wrapper: impl ModelWrapper,
+    wrapper: impl InfererWrapper + 'static,
     inferer: impl Inferer + 'static,
     config: &Args,
 ) -> Result<Vec<Record>> {
-    let mut model = Model::new(wrapper, Box::new(inferer) as Box<dyn Inferer>);
+    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() {
@@ -246,7 +246,7 @@ fn do_run(
         };
 
         model = model
-            .with_new_policy(Box::new(inferer) as Box<dyn Inferer>)
+            .with_new_inferer(Box::new(inferer) as Box<dyn Inferer>)
             .map_err(|(_, e)| e)?;
 
         let shapes = model.input_shapes().to_vec();
@@ -274,7 +274,7 @@ fn do_run(
 }
 
 fn run_apply_epsilon_config(
-    wrapper: impl ModelWrapper,
+    wrapper: impl InfererWrapper + 'static,
     inferer: impl Inferer + 'static,
     config: &Args,
 ) -> Result<Vec<Record>> {
@@ -287,7 +287,7 @@ fn run_apply_epsilon_config(
 }
 
 fn run_apply_recurrent(
-    wrapper: impl ModelWrapper,
+    wrapper: impl InfererWrapper + 'static,
     inferer: impl Inferer + 'static,
     config: &Args,
 ) -> Result<Vec<Record>> {
@@ -328,7 +328,7 @@ pub(super) fn run(config: Args) -> Result<()> {
         }
     };
 
-    let records = run_apply_recurrent(BaseCase, inferer, &config)?;
+    let records = run_apply_recurrent(BaseWrapper, inferer, &config)?;
 
     // Print JSON
     if matches!(config.output, OutputFormat::Json) {

crates/cervo-core/src/epsilon.rs

@@ -8,7 +8,7 @@ Utilities for filling noise inputs for an inference model.
 
 use std::cell::RefCell;
 
-use crate::{batcher::ScratchPadView, inferer::Inferer, prelude::ModelWrapper};
+use crate::{batcher::ScratchPadView, inferer::Inferer, prelude::InfererWrapper};
 use anyhow::{bail, Result};
 use perchance::PerchanceContext;
 use rand::thread_rng;
@@ -212,21 +212,21 @@ where
         self.inner.raw_output_shapes()
     }
 
-    fn begin_agent(&mut self, id: u64) {
+    fn begin_agent(&self, id: u64) {
         self.inner.begin_agent(id);
     }
 
-    fn end_agent(&mut self, id: u64) {
+    fn end_agent(&self, id: u64) {
         self.inner.end_agent(id);
     }
 }
 
-pub struct EpsilonInjectorWrapper<Inner: ModelWrapper, NG: NoiseGenerator> {
+pub struct EpsilonInjectorWrapper<Inner: InfererWrapper, NG: NoiseGenerator> {
     inner: Inner,
     state: EpsilonInjectorState<NG>,
 }
 
-impl<Inner: ModelWrapper> EpsilonInjectorWrapper<Inner, HighQualityNoiseGenerator> {
+impl<Inner: InfererWrapper> EpsilonInjectorWrapper<Inner, HighQualityNoiseGenerator> {
     /// Wraps the provided `inferer` to automatically generate noise for the input named by `key`.
     ///
     /// This function will use [`HighQualityNoiseGenerator`] as the noise source.
@@ -245,7 +245,7 @@ impl<Inner: ModelWrapper> EpsilonInjectorWrapper<Inner, HighQualityNoiseGenerato
 
 impl<Inner, NG> EpsilonInjectorWrapper<Inner, NG>
 where
-    Inner: ModelWrapper,
+    Inner: InfererWrapper,
     NG: NoiseGenerator,
 {
     /// Create a new injector for the provided `key`, using the custom `generator` as the noise source.
@@ -284,12 +284,12 @@ where
     }
 }
 
-impl<Inner, NG> ModelWrapper for EpsilonInjectorWrapper<Inner, NG>
+impl<Inner, NG> InfererWrapper for EpsilonInjectorWrapper<Inner, NG>
 where
-    Inner: ModelWrapper,
+    Inner: InfererWrapper,
     NG: NoiseGenerator,
 {
-    fn invoke(&self, inferer: &impl Inferer, batch: &mut ScratchPadView<'_>) -> anyhow::Result<()> {
+    fn invoke(&self, inferer: &dyn Inferer, batch: &mut ScratchPadView<'_>) -> anyhow::Result<()> {
         self.inner.invoke(inferer, batch)?;
         let total_count = self.state.count * batch.len();
         let output = batch.input_slot_mut(self.state.index);
@@ -306,11 +306,11 @@ where
         self.inner.output_shapes(inferer)
     }
 
-    fn begin_agent(&self, id: u64) {
-        self.inner.begin_agent(id);
+    fn begin_agent(&self, inferer: &dyn Inferer, id: u64) {
+        self.inner.begin_agent(inferer, id);
     }
 
-    fn end_agent(&self, id: u64) {
-        self.inner.end_agent(id);
+    fn end_agent(&self, inferer: &dyn Inferer, id: u64) {
+        self.inner.end_agent(inferer, id);
     }
 }

crates/cervo-core/src/inferer.rs

@@ -112,8 +112,8 @@ pub trait Inferer {
     /// Retrieve the name and shapes of the model outputs.
     fn raw_output_shapes(&self) -> &[(String, Vec<usize>)];
 
-    fn begin_agent(&mut self, id: u64);
-    fn end_agent(&mut self, id: u64);
+    fn begin_agent(&self, id: u64);
+    fn end_agent(&self, id: u64);
 }
 
 /// Helper trait to provide helper functions for loadable models.
@@ -242,12 +242,12 @@ impl Inferer for Box<dyn Inferer + Send> {
         self.as_ref().raw_output_shapes()
     }
 
-    fn begin_agent(&mut self, id: u64) {
-        self.as_mut().begin_agent(id);
+    fn begin_agent(&self, id: u64) {
+        self.as_ref().begin_agent(id);
     }
 
-    fn end_agent(&mut self, id: u64) {
-        self.as_mut().end_agent(id);
+    fn end_agent(&self, id: u64) {
+        self.as_ref().end_agent(id);
     }
 }
 
@@ -268,11 +268,11 @@ impl Inferer for Box<dyn Inferer> {
         self.as_ref().raw_output_shapes()
     }
 
-    fn begin_agent(&mut self, id: u64) {
-        self.as_mut().begin_agent(id);
+    fn begin_agent(&self, id: u64) {
+        self.as_ref().begin_agent(id);
     }
 
-    fn end_agent(&mut self, id: u64) {
-        self.as_mut().end_agent(id);
+    fn end_agent(&self, id: u64) {
+        self.as_ref().end_agent(id);
     }
 }

crates/cervo-core/src/inferer/basic.rs

@@ -94,6 +94,6 @@ impl Inferer for BasicInferer {
         &self.model_api.outputs
     }
 
-    fn begin_agent(&mut self, _id: u64) {}
-    fn end_agent(&mut self, _id: u64) {}
+    fn begin_agent(&self, _id: u64) {}
+    fn end_agent(&self, _id: u64) {}
 }

crates/cervo-core/src/inferer/dynamic.rs

@@ -110,6 +110,6 @@ impl Inferer for DynamicInferer {
         &self.model_api.outputs
     }
 
-    fn begin_agent(&mut self, _id: u64) {}
-    fn end_agent(&mut self, _id: u64) {}
+    fn begin_agent(&self, _id: u64) {}
+    fn end_agent(&self, _id: u64) {}
 }

crates/cervo-core/src/inferer/fixed.rs

@@ -111,8 +111,8 @@ impl Inferer for FixedBatchInferer {
         &self.model_api.outputs
     }
 
-    fn begin_agent(&mut self, _id: u64) {}
-    fn end_agent(&mut self, _id: u64) {}
+    fn begin_agent(&self, _id: u64) {}
+    fn end_agent(&self, _id: u64) {}
 }
 
 struct BatchedModel {

crates/cervo-core/src/inferer/memoizing.rs

@@ -175,6 +175,6 @@ impl Inferer for MemoizingDynamicInferer {
         &self.model_api.outputs
     }
 
-    fn begin_agent(&mut self, _id: u64) {}
-    fn end_agent(&mut self, _id: u64) {}
+    fn begin_agent(&self, _id: u64) {}
+    fn end_agent(&self, _id: u64) {}
 }

crates/cervo-core/src/lib.rs

@@ -15,9 +15,9 @@ pub use tract_hir;
 pub mod batcher;
 pub mod epsilon;
 pub mod inferer;
-pub mod model;
 mod model_api;
 pub mod recurrent;
+pub mod wrapper;
 
 /// Most core utilities are re-exported here.
 pub mod prelude {
@@ -30,7 +30,7 @@ pub mod prelude {
         InfererProvider, MemoizingDynamicInferer, Response, State,
     };
 
-    pub use super::model::ModelWrapper;
     pub use super::model_api::ModelApi;
     pub use super::recurrent::{RecurrentInfo, RecurrentTracker};
+    pub use super::wrapper::{InfererWrapper, InfererWrapperExt, IntoStateful, StatefulInferer};
 }