Merge pull request #46 from nbigaouette/42-tie-session-lifetime-to-environment

Store a reference to environment in session to tie its lifetime

Author: nbigaouette

onnxruntime/src/environment.rs

@@ -135,7 +135,7 @@ impl Environment {
     /// Create a new [`SessionBuilder`](../session/struct.SessionBuilder.html)
     /// used to create a new ONNX session.
     pub fn new_session_builder(&self) -> Result<SessionBuilder> {
-        SessionBuilder::new(self.clone())
+        SessionBuilder::new(self)
     }
 }
 

onnxruntime/src/session.rs

@@ -62,15 +62,15 @@ use crate::{download::AvailableOnnxModel, error::OrtDownloadError};
 /// # }
 /// ```
 #[derive(Debug)]
-pub struct SessionBuilder {
-    env: Environment,
+pub struct SessionBuilder<'a> {
+    env: &'a Environment,
     session_options_ptr: *mut sys::OrtSessionOptions,
 
     allocator: AllocatorType,
     memory_type: MemType,
 }
 
-impl Drop for SessionBuilder {
+impl<'a> Drop for SessionBuilder<'a> {
     #[tracing::instrument]
     fn drop(&mut self) {
         debug!("Dropping the session options.");
@@ -79,8 +79,8 @@ impl Drop for SessionBuilder {
     }
 }
 
-impl SessionBuilder {
-    pub(crate) fn new(env: Environment) -> Result<SessionBuilder> {
+impl<'a> SessionBuilder<'a> {
+    pub(crate) fn new(env: &'a Environment) -> Result<SessionBuilder<'a>> {
         let mut session_options_ptr: *mut sys::OrtSessionOptions = std::ptr::null_mut();
         let status = unsafe { g_ort().CreateSessionOptions.unwrap()(&mut session_options_ptr) };
 
@@ -97,7 +97,7 @@ impl SessionBuilder {
     }
 
     /// Configure the session to use a number of threads
-    pub fn with_number_threads(self, num_threads: i16) -> Result<SessionBuilder> {
+    pub fn with_number_threads(self, num_threads: i16) -> Result<SessionBuilder<'a>> {
         // FIXME: Pre-built binaries use OpenMP, set env variable instead
 
         // We use a u16 in the builder to cover the 16-bits positive values of a i32.
@@ -113,7 +113,7 @@ impl SessionBuilder {
     pub fn with_optimization_level(
         self,
         opt_level: GraphOptimizationLevel,
-    ) -> Result<SessionBuilder> {
+    ) -> Result<SessionBuilder<'a>> {
         // Sets graph optimization level
         unsafe {
             g_ort().SetSessionGraphOptimizationLevel.unwrap()(
@@ -127,47 +127,44 @@ impl SessionBuilder {
     /// Set the session's allocator
     ///
     /// Defaults to [`AllocatorType::Arena`](../enum.AllocatorType.html#variant.Arena)
-    pub fn with_allocator(mut self, allocator: AllocatorType) -> Result<SessionBuilder> {
+    pub fn with_allocator(mut self, allocator: AllocatorType) -> Result<SessionBuilder<'a>> {
         self.allocator = allocator;
         Ok(self)
     }
 
     /// Set the session's memory type
     ///
     /// Defaults to [`MemType::Default`](../enum.MemType.html#variant.Default)
-    pub fn with_memory_type(mut self, memory_type: MemType) -> Result<SessionBuilder> {
+    pub fn with_memory_type(mut self, memory_type: MemType) -> Result<SessionBuilder<'a>> {
         self.memory_type = memory_type;
         Ok(self)
     }
 
     /// Download an ONNX pre-trained model from the [ONNX Model Zoo](https://github.com/onnx/models) and commit the session
     #[cfg(feature = "model-fetching")]
-    pub fn with_model_downloaded<M>(self, model: M) -> Result<Session>
+    pub fn with_model_downloaded<M>(self, model: M) -> Result<Session<'a>>
     where
         M: Into<AvailableOnnxModel>,
     {
         self.with_model_downloaded_monomorphized(model.into())
     }
 
     #[cfg(feature = "model-fetching")]
-    fn with_model_downloaded_monomorphized(self, model: AvailableOnnxModel) -> Result<Session> {
+    fn with_model_downloaded_monomorphized(self, model: AvailableOnnxModel) -> Result<Session<'a>> {
         let download_dir = env::current_dir().map_err(OrtDownloadError::IoError)?;
         let downloaded_path = model.download_to(download_dir)?;
-        self.with_model_from_file_monomorphized(downloaded_path.as_ref())
+        self.with_model_from_file(downloaded_path)
     }
 
     // TODO: Add all functions changing the options.
     //       See all OrtApi methods taking a `options: *mut OrtSessionOptions`.
 
     /// Load an ONNX graph from a file and commit the session
-    pub fn with_model_from_file<P>(self, model_filepath: P) -> Result<Session>
+    pub fn with_model_from_file<P>(self, model_filepath_ref: P) -> Result<Session<'a>>
     where
-        P: AsRef<Path>,
+        P: AsRef<Path> + 'a,
     {
-        self.with_model_from_file_monomorphized(model_filepath.as_ref())
-    }
-
-    fn with_model_from_file_monomorphized(self, model_filepath: &Path) -> Result<Session> {
+        let model_filepath = model_filepath_ref.as_ref();
         let mut session_ptr: *mut sys::OrtSession = std::ptr::null_mut();
 
         if !model_filepath.exists() {
@@ -224,6 +221,7 @@ impl SessionBuilder {
             .collect::<Result<Vec<Output>>>()?;
 
         Ok(Session {
+            env: self.env,
             session_ptr,
             allocator_ptr,
             memory_info,
@@ -233,14 +231,14 @@ impl SessionBuilder {
     }
 
     /// Load an ONNX graph from memory and commit the session
-    pub fn with_model_from_memory<B>(self, model_bytes: B) -> Result<Session>
+    pub fn with_model_from_memory<B>(self, model_bytes: B) -> Result<Session<'a>>
     where
         B: AsRef<[u8]>,
     {
         self.with_model_from_memory_monomorphized(model_bytes.as_ref())
     }
 
-    fn with_model_from_memory_monomorphized(self, model_bytes: &[u8]) -> Result<Session> {
+    fn with_model_from_memory_monomorphized(self, model_bytes: &[u8]) -> Result<Session<'a>> {
         let mut session_ptr: *mut sys::OrtSession = std::ptr::null_mut();
 
         let env_ptr: *const sys::OrtEnv = self.env.env_ptr();
@@ -279,6 +277,7 @@ impl SessionBuilder {
             .collect::<Result<Vec<Output>>>()?;
 
         Ok(Session {
+            env: self.env,
             session_ptr,
             allocator_ptr,
             memory_info,
@@ -290,7 +289,8 @@ impl SessionBuilder {
 
 /// Type storing the session information, built from an [`Environment`](environment/struct.Environment.html)
 #[derive(Debug)]
-pub struct Session {
+pub struct Session<'a> {
+    env: &'a Environment,
     session_ptr: *mut sys::OrtSession,
     allocator_ptr: *mut sys::OrtAllocator,
     memory_info: MemoryInfo,
@@ -348,7 +348,7 @@ impl Output {
     }
 }
 
-impl Drop for Session {
+impl<'a> Drop for Session<'a> {
     #[tracing::instrument]
     fn drop(&mut self) {
         debug!("Dropping the session.");
@@ -360,7 +360,7 @@ impl Drop for Session {
     }
 }
 
-impl Session {
+impl<'a> Session<'a> {
     /// Run the input data through the ONNX graph, performing inference.
     ///
     /// Note that ONNX models can have multiple inputs; a `Vec<_>` is thus
@@ -562,7 +562,7 @@ impl Session {
 
 /// This module contains dangerous functions working on raw pointers.
 /// Those functions are only to be used from inside the
-/// `SessionBuilder::with_model_from_file_monomorphized()` method.
+/// `SessionBuilder::with_model_from_file()` method.
 mod dangerous {
     use super::*;