Share ORT sessions across pipeline handles: N concurrent diarizations at one engine's VRAM

src/inference.rs

@@ -17,8 +17,29 @@ pub use segmentation::{SegmentationError, SegmentationModel};
 pub(crate) mod coreml;
 
 use ort::ep;
+use ort::session::Session;
 use ort::session::builder::SessionBuilder;
 
+/// One ORT session shared across pipeline handles (diar-native patch, T9a).
+///
+/// `Session::run` takes `&mut self` in ort 2.0.0-rc.12 even though the ORT C API's `Run`
+/// is thread-safe, so cross-handle sharing goes through a mutex held for exactly one
+/// inference call. Weights and the session's arena are loaded once; every handle cloned
+/// via `clone_shared` re-uses them and pays only for its own scratch buffers.
+pub(crate) type SharedSession = std::sync::Arc<std::sync::Mutex<Session>>;
+
+pub(crate) fn share_session(session: Session) -> SharedSession {
+    std::sync::Arc::new(std::sync::Mutex::new(session))
+}
+
+pub(crate) fn lock_session(session: &SharedSession) -> std::sync::MutexGuard<'_, Session> {
+    // A poisoned lock means another handle panicked mid-run; the session itself has no
+    // torn state to protect (ORT `Run` is atomic at the C API level), so keep serving.
+    session
+        .lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+}
+
 #[cfg(all(feature = "load-dynamic", not(target_arch = "wasm32")))]
 static ORT_RUNTIME_INIT: OnceLock<Result<(), OrtRuntimeError>> = OnceLock::new();
 

src/inference/embedding.rs

@@ -5,11 +5,11 @@ use std::sync::Arc;
 
 #[cfg(feature = "coreml")]
 use crate::inference::coreml::{CachedInputShape, CoreMlModel, SharedCoreMlModel};
-use crate::inference::{ExecutionMode, ModelLoadError};
+use crate::inference::{ExecutionMode, ModelLoadError, SharedSession};
 use ndarray::{Array2, Array3, s};
 #[cfg(feature = "coreml")]
 use objc2_core_ml::MLComputeUnits;
-use ort::session::{HasSelectedOutputs, RunOptions, Session};
+use ort::session::{HasSelectedOutputs, RunOptions};
 
 mod batch;
 #[cfg(feature = "coreml")]
@@ -48,9 +48,13 @@ const FBANK_FRAMES: usize = 998;
 const FBANK_FEATURES: usize = 80;
 const MASK_FRAMES: usize = 589;
 
+/// One masked audio window for [`EmbeddingModel::embed_batch`]
 pub struct MaskedEmbeddingInput<'a> {
+    /// 16 kHz mono samples (padded/truncated to the model window internally)
     pub audio: &'a [f32],
+    /// Frame-level speaker weights over the segmentation mask grid
     pub mask: &'a [f32],
+    /// Optional overlap-cleaned mask, preferred when it keeps enough weight
     pub clean_mask: Option<&'a [f32]>,
 }
 
@@ -71,18 +75,38 @@ struct EmbeddingMeta {
 }
 
 struct OrtEmbeddingState {
-    session: Session,
-    primary_batched_session: Option<Session>,
-    split_fbank_session: Option<Session>,
-    split_fbank_batched_session: Option<Session>,
-    split_tail_session: Option<Session>,
-    split_tail_batched_session: Option<Session>,
-    split_primary_tail_batched_session: Option<Session>,
-    multi_mask_session: Option<Session>,
-    multi_mask_batched_session: Option<Session>,
+    session: SharedSession,
+    primary_batched_session: Option<SharedSession>,
+    split_fbank_session: Option<SharedSession>,
+    split_fbank_batched_session: Option<SharedSession>,
+    split_tail_session: Option<SharedSession>,
+    split_tail_batched_session: Option<SharedSession>,
+    split_primary_tail_batched_session: Option<SharedSession>,
+    multi_mask_session: Option<SharedSession>,
+    multi_mask_batched_session: Option<SharedSession>,
+    // Per-handle: carries a preallocated output tensor, so it must never be shared
+    // across concurrent runs.
     primary_batch_run_options: Option<RunOptions<HasSelectedOutputs>>,
 }
 
+impl OrtEmbeddingState {
+    fn fresh_primary_run_options(
+        has_primary_batched: bool,
+    ) -> Result<Option<RunOptions<HasSelectedOutputs>>, ort::Error> {
+        has_primary_batched
+            .then(|| {
+                let mut opts = preallocated_run_options(
+                    PRIMARY_BATCH_SIZE,
+                    256,
+                    "primary batched embedding output",
+                )?;
+                let _ = opts.disable_device_sync();
+                Ok::<RunOptions<HasSelectedOutputs>, ort::Error>(opts)
+            })
+            .transpose()
+    }
+}
+
 #[cfg(feature = "coreml")]
 struct CoreMlEmbeddingState {
     #[cfg(feature = "coreml")]
@@ -136,6 +160,40 @@ struct EmbeddingBuffers {
     split_primary_weights_batch_buffer: Array2<f32>,
 }
 
+impl EmbeddingBuffers {
+    fn fresh() -> Self {
+        Self {
+            multi_mask_fbank_buffer: Array3::zeros((
+                MULTI_MASK_BATCH_SIZE,
+                FBANK_FRAMES,
+                FBANK_FEATURES,
+            )),
+            multi_mask_masks_buffer: Array2::zeros((
+                MULTI_MASK_BATCH_SIZE * NUM_SPEAKERS,
+                MASK_FRAMES,
+            )),
+            waveform_buffer: Array3::zeros((1, 1, 160_000)),
+            weights_buffer: Array2::zeros((1, 589)),
+            primary_batch_waveform_buffer: Array3::zeros((PRIMARY_BATCH_SIZE, 1, 160_000)),
+            primary_batch_weights_buffer: Array2::zeros((PRIMARY_BATCH_SIZE, 589)),
+            split_waveform_buffer: Array3::zeros((1, 1, 160_000)),
+            split_fbank_batch_buffer: Array3::zeros((FBANK_BATCH_SIZE, 1, 160_000)),
+            split_feature_batch_buffer: Array3::zeros((
+                CHUNK_SPEAKER_BATCH_SIZE,
+                FBANK_FRAMES,
+                FBANK_FEATURES,
+            )),
+            split_weights_batch_buffer: Array2::zeros((CHUNK_SPEAKER_BATCH_SIZE, 589)),
+            split_primary_feature_batch_buffer: Array3::zeros((
+                PRIMARY_BATCH_SIZE,
+                FBANK_FRAMES,
+                FBANK_FEATURES,
+            )),
+            split_primary_weights_batch_buffer: Array2::zeros((PRIMARY_BATCH_SIZE, 589)),
+        }
+    }
+}
+
 /// WeSpeaker speaker embedding model with split-backend and chunk embedding support
 pub struct EmbeddingModel {
     meta: EmbeddingMeta,
@@ -159,6 +217,45 @@ impl EmbeddingModel {
         Self::with_mode_and_config(model_path, mode, &crate::pipeline::RuntimeConfig::default())
     }
 
+    /// Cheap handle over the same ORT sessions with fresh scratch buffers.
+    ///
+    /// All 9 ORT sessions (weights + arenas — the VRAM) are shared through
+    /// [`SharedSession`]; the staging buffers and the preallocated-output
+    /// `RunOptions` are re-created per handle, so N handles cost N × scratch
+    /// (~130 MB host RAM), not N × engine (weights + arenas). Handles may run concurrently — each
+    /// inference call serializes on its session's mutex for exactly one `Run`.
+    #[cfg(not(feature = "coreml"))]
+    pub fn clone_shared(&self) -> Result<Self, ort::Error> {
+        Ok(Self {
+            meta: EmbeddingMeta {
+                model_path: self.meta.model_path.clone(),
+                mode: self.meta.mode,
+                sample_rate: self.meta.sample_rate,
+                window_samples: self.meta.window_samples,
+                mask_frames: self.meta.mask_frames,
+                min_num_samples: self.meta.min_num_samples,
+            },
+            ort: OrtEmbeddingState {
+                session: std::sync::Arc::clone(&self.ort.session),
+                primary_batched_session: self.ort.primary_batched_session.clone(),
+                split_fbank_session: self.ort.split_fbank_session.clone(),
+                split_fbank_batched_session: self.ort.split_fbank_batched_session.clone(),
+                split_tail_session: self.ort.split_tail_session.clone(),
+                split_tail_batched_session: self.ort.split_tail_batched_session.clone(),
+                split_primary_tail_batched_session: self
+                    .ort
+                    .split_primary_tail_batched_session
+                    .clone(),
+                multi_mask_session: self.ort.multi_mask_session.clone(),
+                multi_mask_batched_session: self.ort.multi_mask_batched_session.clone(),
+                primary_batch_run_options: OrtEmbeddingState::fresh_primary_run_options(
+                    self.ort.primary_batched_session.is_some(),
+                )?,
+            },
+            buffers: EmbeddingBuffers::fresh(),
+        })
+    }
+
     /// Audio sample rate in Hz (16000)
     pub fn sample_rate(&self) -> usize {
         self.meta.sample_rate

src/inference/embedding/batch.rs

@@ -8,6 +8,7 @@ use super::{
     NUM_SPEAKERS, PRIMARY_BATCH_SIZE, SplitTailInput, array2_from_shape_vec, array2_slice_mut,
     array3_slice_mut, first_output, select_mask,
 };
+use crate::inference::lock_session;
 
 impl EmbeddingModel {
     /// Extract speaker embeddings for a batch of masked audio windows
@@ -18,7 +19,7 @@ impl EmbeddingModel {
         if let Some(sess) = self
             .ort
             .primary_batched_session
-            .as_mut()
+            .as_ref()
             .filter(|_| inputs.len() == PRIMARY_BATCH_SIZE)
         {
             for (batch_idx, input) in inputs.iter().enumerate() {
@@ -48,6 +49,7 @@ impl EmbeddingModel {
                 TensorRef::from_array_view(self.buffers.primary_batch_weights_buffer.view())?;
             let ort_inputs =
                 ort::inputs!["waveform" => waveform_tensor, "weights" => weights_tensor];
+            let mut sess = lock_session(sess);
             let outputs = if let Some(opts) = &self.ort.primary_batch_run_options {
                 sess.run_with_options(ort_inputs, opts)?
             } else {
@@ -146,12 +148,14 @@ impl EmbeddingModel {
                 TensorRef::from_array_view(self.buffers.multi_mask_fbank_buffer.view())?;
             let masks_tensor =
                 TensorRef::from_array_view(self.buffers.multi_mask_masks_buffer.view())?;
-            let outputs = self
-                .ort
-                .multi_mask_batched_session
-                .as_mut()
-                .ok_or_else(|| ort::Error::new("missing multi-mask batched session"))?
-                .run(ort::inputs!["fbank" => fbank_tensor, "masks" => masks_tensor])?;
+            let mut session = lock_session(
+                self.ort
+                    .multi_mask_batched_session
+                    .as_ref()
+                    .ok_or_else(|| ort::Error::new("missing multi-mask batched session"))?,
+            );
+            let outputs =
+                session.run(ort::inputs!["fbank" => fbank_tensor, "masks" => masks_tensor])?;
             let output = first_output(outputs.values(), "multi-mask batched output")?;
             let (_shape, data) = output.try_extract_tensor::<f32>()?;
             let batch = array2_from_shape_vec(
@@ -177,12 +181,14 @@ impl EmbeddingModel {
                     .slice(s![mask_start..mask_end, ..]);
                 let fbank_tensor = TensorRef::from_array_view(fbank_slice.view())?;
                 let masks_tensor = TensorRef::from_array_view(masks_slice.view())?;
-                let outputs = self
-                    .ort
-                    .multi_mask_session
-                    .as_mut()
-                    .ok_or_else(|| ort::Error::new("missing multi-mask session"))?
-                    .run(ort::inputs!["fbank" => fbank_tensor, "masks" => masks_tensor])?;
+                let mut session = lock_session(
+                    self.ort
+                        .multi_mask_session
+                        .as_ref()
+                        .ok_or_else(|| ort::Error::new("missing multi-mask session"))?,
+                );
+                let outputs =
+                    session.run(ort::inputs!["fbank" => fbank_tensor, "masks" => masks_tensor])?;
                 let output = first_output(outputs.values(), "multi-mask output")?;
                 let (_shape, data) = output.try_extract_tensor::<f32>()?;
                 for (local_idx, row_idx) in (mask_start..mask_end).enumerate() {
@@ -266,12 +272,14 @@ impl EmbeddingModel {
             TensorRef::from_array_view(self.buffers.split_primary_feature_batch_buffer.view())?;
         let weights_tensor =
             TensorRef::from_array_view(self.buffers.split_primary_weights_batch_buffer.view())?;
-        let outputs = self
-            .ort
-            .split_primary_tail_batched_session
-            .as_mut()
-            .ok_or_else(|| ort::Error::new("missing primary tail batched session"))?
-            .run(ort::inputs!["fbank" => fbank_tensor, "weights" => weights_tensor])?;
+        let mut session = lock_session(
+            self.ort
+                .split_primary_tail_batched_session
+                .as_ref()
+                .ok_or_else(|| ort::Error::new("missing primary tail batched session"))?,
+        );
+        let outputs =
+            session.run(ort::inputs!["fbank" => fbank_tensor, "weights" => weights_tensor])?;
         let output = first_output(outputs.values(), "primary tail batched output")?;
         let (_shape, data) = output.try_extract_tensor::<f32>()?;
         let batch = array2_from_shape_vec(

src/inference/embedding/fbank.rs

@@ -4,6 +4,7 @@ use ort::value::TensorRef;
 #[cfg(feature = "coreml")]
 use super::tensor::array3_slice;
 use super::{EmbeddingModel, FBANK_BATCH_SIZE, array2_from_shape_vec, first_output};
+use crate::inference::lock_session;
 
 impl EmbeddingModel {
     /// Compute fbank features for a single audio chunk via the split fbank model
@@ -40,12 +41,13 @@ impl EmbeddingModel {
 
         let waveform_tensor =
             TensorRef::from_array_view(self.buffers.split_waveform_buffer.view())?;
-        let outputs = self
-            .ort
-            .split_fbank_session
-            .as_mut()
-            .ok_or_else(|| ort::Error::new("missing split fbank session"))?
-            .run(ort::inputs!["waveform" => waveform_tensor])?;
+        let mut session = lock_session(
+            self.ort
+                .split_fbank_session
+                .as_ref()
+                .ok_or_else(|| ort::Error::new("missing split fbank session"))?,
+        );
+        let outputs = session.run(ort::inputs!["waveform" => waveform_tensor])?;
         let output = first_output(outputs.values(), "chunk fbank output")?;
         let (shape, data) = output.try_extract_tensor::<f32>()?;
         let frames = shape[1] as usize;
@@ -97,12 +99,13 @@ impl EmbeddingModel {
 
             let waveform_tensor =
                 TensorRef::from_array_view(self.buffers.split_fbank_batch_buffer.view())?;
-            let outputs = self
-                .ort
-                .split_fbank_batched_session
-                .as_mut()
-                .ok_or_else(|| ort::Error::new("missing split fbank batched session"))?
-                .run(ort::inputs!["waveform" => waveform_tensor])?;
+            let mut session = lock_session(
+                self.ort
+                    .split_fbank_batched_session
+                    .as_ref()
+                    .ok_or_else(|| ort::Error::new("missing split fbank batched session"))?,
+            );
+            let outputs = session.run(ort::inputs!["waveform" => waveform_tensor])?;
             let output = first_output(outputs.values(), "batched chunk fbank output")?;
             let (shape, data) = output.try_extract_tensor::<f32>()?;
             Self::push_fbank_batch_results(