revert: using transmute to ignore mut refs of Sessions instead of Mutex lock

Author: kallebysantos

ext/ai/lib.rs

@@ -40,6 +40,8 @@ use tracing::error;
 use tracing::trace_span;
 use tracing::Instrument;
 
+use crate::onnxruntime::session::as_mut_session;
+
 deno_core::extension!(
   ai,
   ops = [
@@ -213,14 +215,10 @@ async fn init_gte(state: Rc<RefCell<OpState>>) -> Result<(), Error> {
         &*token_type_ids,
       ))?;
 
-      let Ok(mut guard) = session.lock() else {
-        let err = anyhow!("failed to lock session");
-        error!(reason = ?err);
-        return Err(err);
-      };
+      let session = unsafe { as_mut_session(&session) };
 
       let outputs = trace_span!("infer_gte").in_scope(|| {
-        guard.run(inputs! {
+        session.run(inputs! {
           "input_ids" => input_ids_array,
           "token_type_ids" => token_type_ids_array,
           "attention_mask" => attention_mask_array,

ext/ai/onnxruntime/mod.rs

@@ -10,7 +10,6 @@ use std::cell::RefCell;
 use std::collections::HashMap;
 use std::rc::Rc;
 use std::sync::Arc;
-use std::sync::Mutex;
 
 use anyhow::anyhow;
 use anyhow::Context;
@@ -33,6 +32,8 @@ use tokio::sync::oneshot;
 use tracing::debug;
 use tracing::trace;
 
+use crate::onnxruntime::session::as_mut_session;
+
 #[op2(async)]
 #[to_v8]
 pub async fn op_ai_ort_init_session(
@@ -56,11 +57,8 @@ pub async fn op_ai_ort_init_session(
   };
 
   let mut state = state.borrow_mut();
-  let mut sessions = {
-    state
-      .try_take::<Vec<Arc<Mutex<Session>>>>()
-      .unwrap_or_default()
-  };
+  let mut sessions =
+    { state.try_take::<Vec<Arc<Session>>>().unwrap_or_default() };
 
   sessions.push(model.get_session());
   state.put(sessions);
@@ -107,12 +105,9 @@ pub async fn op_ai_ort_run_session(
     JsRuntime::op_state_from(state)
       .borrow_mut()
       .spawn_cpu_accumul_blocking_scope(move || {
-        let Ok(mut session_guard) = model_session.lock() else {
-          let _ = tx.send(Err(anyhow!("failed to lock model session")));
-          return;
-        };
+        let session = unsafe { as_mut_session(&model_session) };
 
-        let outputs = match session_guard.run(input_values) {
+        let outputs = match session.run(input_values) {
           Ok(v) => v,
           Err(err) => {
             let _ = tx.send(Err(anyhow::Error::from(err)));

ext/ai/onnxruntime/model.rs

@@ -1,12 +1,9 @@
-use std::sync::Arc;
-use std::sync::Mutex;
-
-use anyhow::anyhow;
 use anyhow::Result;
 use deno_core::serde_v8::to_v8;
 use deno_core::ToV8;
 use ort::session::Session;
 use reqwest::Url;
+use std::sync::Arc;
 
 use super::session::get_session;
 use super::session::load_session_from_bytes;
@@ -29,26 +26,21 @@ impl std::fmt::Display for ModelInfo {
 #[derive(Debug)]
 pub struct Model {
   info: ModelInfo,
-  session: Arc<Mutex<Session>>,
+  session: Arc<Session>,
 }
 
 impl Model {
   fn new(session_with_id: SessionWithId) -> Result<Self> {
     let (input_names, output_names) = {
-      let Ok(session_guard) = session_with_id.session.lock() else {
-        return Err(anyhow!(
-          "Could not lock model session {}",
-          session_with_id.id
-        ));
-      };
-
-      let input_names = session_guard
+      let session = { session_with_id.session.clone() };
+
+      let input_names = session
         .inputs
         .iter()
         .map(|input| input.name.clone())
         .collect::<Vec<_>>();
 
-      let output_names = session_guard
+      let output_names = session
         .outputs
         .iter()
         .map(|output| output.name.clone())
@@ -71,7 +63,7 @@ impl Model {
     self.info.clone()
   }
 
-  pub fn get_session(&self) -> Arc<Mutex<Session>> {
+  pub fn get_session(&self) -> Arc<Session> {
     self.session.clone()
   }
 

ext/ai/onnxruntime/session.rs

@@ -5,7 +5,6 @@ use once_cell::sync::Lazy;
 use reqwest::Url;
 use std::hash::Hasher;
 use std::sync::Arc;
-use std::sync::Mutex;
 use tokio_util::compat::FuturesAsyncWriteCompatExt;
 use tracing::debug;
 use tracing::instrument;
@@ -25,17 +24,16 @@ use ort::session::Session;
 
 use crate::onnx::ensure_onnx_env_init;
 
-static SESSIONS: Lazy<DashMap<String, Arc<Mutex<Session>>>> =
-  Lazy::new(DashMap::new);
+static SESSIONS: Lazy<DashMap<String, Arc<Session>>> = Lazy::new(DashMap::new);
 
 #[derive(Debug)]
 pub struct SessionWithId {
   pub(crate) id: String,
-  pub(crate) session: Arc<Mutex<Session>>,
+  pub(crate) session: Arc<Session>,
 }
 
-impl From<(String, Arc<Mutex<Session>>)> for SessionWithId {
-  fn from(value: (String, Arc<Mutex<Session>>)) -> Self {
+impl From<(String, Arc<Session>)> for SessionWithId {
+  fn from(value: (String, Arc<Session>)) -> Self {
     Self {
       id: value.0,
       session: value.1,
@@ -50,7 +48,7 @@ impl std::fmt::Display for SessionWithId {
 }
 
 impl SessionWithId {
-  pub fn into_split(self) -> (String, Arc<Mutex<Session>>) {
+  pub fn into_split(self) -> (String, Arc<Session>) {
     (self.id, self.session)
   }
 }
@@ -106,7 +104,7 @@ fn get_execution_providers() -> Vec<ExecutionProviderDispatch> {
   [cpu].to_vec()
 }
 
-fn create_session(model_bytes: &[u8]) -> Result<Arc<Mutex<Session>>, Error> {
+fn create_session(model_bytes: &[u8]) -> Result<Arc<Session>, Error> {
   let session = {
     if let Some(err) = ensure_onnx_env_init() {
       return Err(anyhow!("failed to create onnx environment: {err}"));
@@ -117,7 +115,14 @@ fn create_session(model_bytes: &[u8]) -> Result<Arc<Mutex<Session>>, Error> {
       .commit_from_memory(model_bytes)?
   };
 
-  Ok(Arc::new(Mutex::new(session)))
+  Ok(Arc::new(session))
+}
+
+#[allow(mutable_transmutes)]
+#[allow(clippy::mut_from_ref)]
+pub(crate) unsafe fn as_mut_session(session: &Arc<Session>) -> &mut Session {
+  // SAFETY: CPU EP https://github.com/pykeio/ort/issues/402#issuecomment-2949993914
+  unsafe { std::mem::transmute::<&Session, &mut Session>(&session.clone()) }
 }
 
 #[instrument(level = "debug", skip_all, fields(model_bytes = model_bytes.len()), err)]
@@ -174,7 +179,7 @@ pub(crate) async fn load_session_from_url(
   Ok((session_id, session).into())
 }
 
-pub(crate) async fn get_session(id: &str) -> Option<Arc<Mutex<Session>>> {
+pub(crate) async fn get_session(id: &str) -> Option<Arc<Session>> {
   SESSIONS.get(id).map(|value| value.pair().1.clone())
 }