Port to ort 2.0.0-rc.12 (ndarray 0.17)

Cargo.toml

@@ -12,8 +12,8 @@ repository = "https://github.com/fbilhaut/orp"
 documentation = "https://docs.rs/fbilhaut/orp"
 
 [dependencies]
-ort = { version = "=2.0.0-rc.9" }
-ndarray = { version = "0.16.0" }
+ort = { version = "=2.0.0-rc.12" }
+ndarray = { version = "0.17" }
 composable = { version = "0.9.0" }
 
 [features]

src/model.rs

@@ -1,42 +1,56 @@
 use std::collections::HashSet;
 use std::path::Path;
+use std::sync::Mutex;
 use composable::Composable;
-use ort::session::{Session, SessionInputs, SessionOutputs, builder::GraphOptimizationLevel};
+use ort::session::{Session, SessionOutputs, builder::GraphOptimizationLevel};
 use crate::error::UnexpectedModelSchemaError;
 
 use super::Result;
 use super::params::RuntimeParameters;
 use super::pipeline::Pipeline;
 
 
+/// Strips the recovery payload from an `ort` builder error so it satisfies
+/// `Box<dyn Error + Send + Sync>` (since ort 2.0.0-rc.10, builder methods
+/// return `Error<SessionBuilder>` whose payload is neither `Send` nor `Sync`).
+fn ort_err<R>(e: ort::Error<R>) -> Box<dyn std::error::Error + Send + Sync> {
+    Box::new(ort::Error::new(e.to_string()))
+}
+
+
 /// A `Model` can load an ONNX model, and run it using the provided pipeline.
-pub struct Model {    
-    session: Session,
+///
+/// Since ort 2.0.0-rc.10, `Session::run` requires `&mut Session`; the session
+/// is wrapped in a `Mutex` to preserve this crate's `&self` inference API.
+/// The lock is held per `run` call, so concurrent callers serialize on
+/// inference (ONNX Runtime parallelizes internally via intra-op threads).
+pub struct Model {
+    session: Mutex<Session>,
 }
 
 
-impl Model {    
+impl Model {
     pub fn new<P: AsRef<Path>>(model_path: P, params: RuntimeParameters) -> Result<Self> {
         let session = Session::builder()?
-            .with_intra_threads(params.threads())?
-            .with_execution_providers(params.into_execution_providers())?
-            .with_optimization_level(GraphOptimizationLevel::Level3)?            
+            .with_intra_threads(params.threads()).map_err(ort_err)?
+            .with_execution_providers(params.into_execution_providers()).map_err(ort_err)?
+            .with_optimization_level(GraphOptimizationLevel::Level3).map_err(ort_err)?
             .commit_from_file(model_path)?;
 
         Ok(Self {
-            session,
+            session: Mutex::new(session),
         })
     }
 
     pub fn new_from_bytes(model_bytes: &[u8], params: RuntimeParameters) -> Result<Self> {
         let session = Session::builder()?
-            .with_intra_threads(params.threads())?
-            .with_execution_providers(params.into_execution_providers())?
-            .with_optimization_level(GraphOptimizationLevel::Level3)?
+            .with_intra_threads(params.threads()).map_err(ort_err)?
+            .with_execution_providers(params.into_execution_providers()).map_err(ort_err)?
+            .with_optimization_level(GraphOptimizationLevel::Level3).map_err(ort_err)?
             .commit_from_memory(model_bytes)?;
 
         Ok(Self {
-            session
+            session: Mutex::new(session),
         })
     }
 
@@ -47,9 +61,20 @@ impl Model {
         // pre-process
         let (input, context) = pipeline.pre_processor(params).apply(input)?;
         // inference
-        let output = self.run(input)?;                
+        let output = {
+            let mut session = self.session.lock().map_err(|e| e.to_string())?;
+            let output = session.run(input)?;
+            // SAFETY: `SessionOutputs<'r>` only borrows the output *names*
+            // (`&'r str`) from the session's outlet metadata; the values are
+            // owned `DynValue`s. Those name strings live on the heap inside
+            // the `Session`, which is owned by `self` (borrowed for `'a`)
+            // and is neither dropped nor structurally mutated while `'a` is
+            // live. The transmute only widens the guard-scoped borrow of
+            // those names to `'a`.
+            unsafe { std::mem::transmute::<SessionOutputs<'_>, SessionOutputs<'a>>(output) }
+        };
         // post-process
-        let output = pipeline.post_processor(params).apply((output, context))?;        
+        let output = pipeline.post_processor(params).apply((output, context))?;
         // ok
         Ok(output)
     }
@@ -60,46 +85,44 @@ impl Model {
 
     /// Writes various model properties from metadata and input/output tensors
     pub fn inspect<W: std::io::Write>(&self, mut writer: W) -> Result<()> {
-        let metadata = self.session.metadata()?;
-        writeln!(writer, "NAME: {}", metadata.name()?)?;
-        writeln!(writer, "PRODUCER: {}", metadata.producer()?)?;
-        writeln!(writer, "VERSION: {}", metadata.version()?)?;
+        let session = self.session.lock().map_err(|e| e.to_string())?;
+        let metadata = session.metadata()?;
+        writeln!(writer, "NAME: {}", metadata.name().unwrap_or_default())?;
+        writeln!(writer, "PRODUCER: {}", metadata.producer().unwrap_or_default())?;
+        writeln!(writer, "VERSION: {}", metadata.version().unwrap_or_default())?;
         writeln!(writer, "INPUTS:")?;
-        for input in &self.session.inputs {
-            writeln!(writer, "\t{}: {:?}", input.name, input.input_type)?;
+        for input in session.inputs() {
+            writeln!(writer, "\t{}: {:?}", input.name(), input.dtype())?;
         }
         writeln!(writer, "OUTPUTS:")?;
-        for input in &self.session.outputs {
-            writeln!(writer, "\t{}: {:?}", input.name, input.output_type)?;
+        for input in session.outputs() {
+            writeln!(writer, "\t{}: {:?}", input.name(), input.dtype())?;
         }
         Ok(())
     }
 
     /// Check model schema wrt. pipeline expectations
     fn check_schema<'a, P: Pipeline<'a>>(&'a self, pipeline: &P, params: &P::Parameters) -> Result<()> {
+        let session = self.session.lock().map_err(|e| e.to_string())?;
         if let Some(expected_inputs) = pipeline.expected_inputs(params) {
             // inputs should be exactly the same sets
             let expected_inputs = &expected_inputs.collect();
-            let actual_inputs: HashSet<_> = self.session.inputs.iter().map(|i| i.name.as_str()).collect();
+            let actual_inputs: HashSet<_> = session.inputs().iter().map(|i| i.name()).collect();
             if !actual_inputs.eq(expected_inputs) {
                 return UnexpectedModelSchemaError::new_for_input(expected_inputs, &actual_inputs).into_err();
             }
         }
-        if let Some(expected_outputs) = pipeline.expected_outputs(params) {            
+        if let Some(expected_outputs) = pipeline.expected_outputs(params) {
             // for outputs, we just check that the expected ones are present (but having others is ok)
             let expected_outputs = &expected_outputs.collect();
-            let actual_outputs: HashSet<_> = self.session.outputs.iter().map(|i| i.name.as_str()).collect();
+            let actual_outputs: HashSet<_> = session.outputs().iter().map(|i| i.name()).collect();
             if !actual_outputs.is_superset(&expected_outputs) {
                 return UnexpectedModelSchemaError::new_for_output(expected_outputs, &actual_outputs).into_err();
             }
         }
         Ok(())
     }
 
-    fn run(&self, input: SessionInputs<'_, '_>) -> Result<SessionOutputs<'_, '_>> {
-        Ok(self.session.run(input)?)
-    }
-
 }
 
 
@@ -122,4 +145,4 @@ impl<'a, P: Pipeline<'a>> Composable<P::Input, P::Output> for ComposableModel<'a
     fn apply(&self, input: P::Input) -> Result<P::Output> {
         self.model.inference(input, self.pipeline, self.params)
     }
-}
+}

src/pipeline.rs

@@ -39,8 +39,8 @@ impl<'a, I, C, T: Composable<I, (SessionInputs<'a, 'a>, C)>> PreProcessor<'a, I,
 
 
 /// Defines a generic post-processor
-pub trait PostProcessor<'a, O, C>: Composable<(SessionOutputs<'a, 'a>, C), O> {}
-impl<'a, O, C, T: Composable<(SessionOutputs<'a, 'a>, C), O>> PostProcessor<'a, O, C> for T {}
+pub trait PostProcessor<'a, O, C>: Composable<(SessionOutputs<'a>, C), O> {}
+impl<'a, O, C, T: Composable<(SessionOutputs<'a>, C), O>> PostProcessor<'a, O, C> for T {}
 
 
 /// Owns a pipeline, and references a model and some parameters to implement `Composable`