feat(embeddings): add parallel session pooling

- Add rayon crate for parallel computation
- Implement session pooling for concurrent execution
- Enable parallel tokenization and inference
- Optimize default batch size and thread configuration
- Add environment variables for session and thread tuning

Author: donhardman

embeddings/Cargo.lock

@@ -13,6 +13,7 @@ serde_json = "1.0.114"
 serde = "1.0.197"
 rand = "0.8.5"
 reqwest = { version = "0.12.8", default-features = false, features = ["blocking", "json", "rustls-tls"] }
+rayon = "1.10.0"
 
 candle-core = "0.9.2"
 candle-nn = "0.9.2"

embeddings/Cargo.toml

@@ -30,10 +30,7 @@ use std::sync::{Arc, Mutex, OnceLock};
 use tokenizers::Tokenizer;
 
 /// Default batch size per ONNX forward pass.
-const DEFAULT_BATCH_SIZE: usize = 32;
-
-/// Default intra-op threads: 0 = ORT default (all cores).
-const DEFAULT_INTRA_THREADS: usize = 0;
+const DEFAULT_BATCH_SIZE: usize = 8;
 
 fn batch_size() -> usize {
     std::env::var("EMBEDDINGS_BATCH_SIZE")
@@ -42,43 +39,28 @@ fn batch_size() -> usize {
         .unwrap_or(DEFAULT_BATCH_SIZE)
 }
 
-fn intra_threads() -> usize {
-    std::env::var("EMBEDDINGS_INTRA_THREADS")
-        .ok()
-        .and_then(|v| v.parse().ok())
-        .unwrap_or(DEFAULT_INTRA_THREADS)
+fn available_cpus() -> usize {
+    std::thread::available_parallelism()
+        .map(|n| n.get())
+        .unwrap_or(2)
 }
 
-/// Thread-safe wrapper around `ort::session::Session`.
-/// ORT's C API `OrtSession::Run` is documented as thread-safe for concurrent calls.
-/// The Rust `ort` crate requires `&mut self` for `run()` which is overly conservative.
-/// This wrapper provides `run(&self)` using `UnsafeCell` to allow parallel inference
-/// on a single loaded model without duplicating weights in memory.
-struct SyncSession {
-    inner: std::cell::UnsafeCell<ort::session::Session>,
+/// Number of parallel ORT sessions. Default: num_cpus / 2 (min 1).
+/// Moderate oversubscription (total_threads ≈ 2×cores) fills idle gaps between ORT ops.
+fn num_sessions() -> usize {
+    std::env::var("EMBEDDINGS_NUM_SESSIONS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or_else(|| (available_cpus() / 2).max(1))
 }
 
-// SAFETY: ORT's OrtSession::Run is thread-safe for concurrent calls.
-// The session state is internally synchronized by ORT.
-unsafe impl Sync for SyncSession {}
-unsafe impl Send for SyncSession {}
-
-impl SyncSession {
-    fn new(session: ort::session::Session) -> Self {
-        Self {
-            inner: std::cell::UnsafeCell::new(session),
-        }
-    }
-
-    /// Run inference on this session. Safe to call from multiple threads concurrently.
-    /// SAFETY: relies on ORT's documented thread-safety of OrtSession::Run.
-    fn run<'s, 'i, 'v: 'i, const N: usize>(
-        &'s self,
-        input_values: impl Into<ort::session::SessionInputs<'i, 'v, N>>,
-    ) -> ort::Result<ort::session::SessionOutputs<'s>> {
-        // SAFETY: ORT guarantees thread-safe concurrent Run() on the same session.
-        unsafe { &mut *self.inner.get() }.run(input_values)
-    }
+/// Intra-op threads per session. Default: 4.
+/// With num_sessions = cpus/2, total threads = cpus/2 * 4 = 2×cpus — optimal oversubscription.
+fn intra_threads() -> usize {
+    std::env::var("EMBEDDINGS_INTRA_THREADS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(4)
 }
 
 /// Model architecture type - determines pooling strategy
@@ -777,15 +759,35 @@ impl QuantizedEmbeddingModel {
 }
 
 /// ONNX embedding model using ORT (onnxruntime) for optimized inference.
-/// Uses a single session with concurrent inference via SyncSession wrapper.
+/// Supports single session (default) or session pool for parallel inference.
 pub struct OnnxEmbeddingModel {
-    session: SyncSession,
+    sessions: Vec<Mutex<ort::session::Session>>,
     tokenizer: Tokenizer,
     max_input_len: usize,
     hidden_size: usize,
 }
 
 impl OnnxEmbeddingModel {
+    fn build_session(
+        onnx_path: &PathBuf,
+        it: usize,
+    ) -> Result<ort::session::Session, Box<dyn Error>> {
+        ort::session::Session::builder()
+            .map_err(|_| LibError::OnnxModelEvalFailed)?
+            .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)
+            .map_err(|_| LibError::OnnxModelEvalFailed)?
+            .with_intra_threads(it)
+            .map_err(|_| LibError::OnnxModelEvalFailed)?
+            .with_intra_op_spinning(false)
+            .map_err(|_| LibError::OnnxModelEvalFailed)?
+            .with_flush_to_zero()
+            .map_err(|_| LibError::OnnxModelEvalFailed)?
+            .with_approximate_gelu()
+            .map_err(|_| LibError::OnnxModelEvalFailed)?
+            .commit_from_file(onnx_path)
+            .map_err(|_| -> Box<dyn Error> { Box::new(LibError::ModelWeightsLoadFailed) })
+    }
+
     pub fn new(
         onnx_path: PathBuf,
         model_info: LocalModelInfo,
@@ -802,23 +804,15 @@ impl OnnxEmbeddingModel {
         tokenizer.with_padding(None);
         let _ = tokenizer.with_truncation(None);
 
-        let session = ort::session::Session::builder()
-            .map_err(|_| LibError::OnnxModelEvalFailed)?
-            .with_optimization_level(ort::session::builder::GraphOptimizationLevel::Level3)
-            .map_err(|_| LibError::OnnxModelEvalFailed)?
-            .with_intra_threads(intra_threads())
-            .map_err(|_| LibError::OnnxModelEvalFailed)?
-            .with_intra_op_spinning(false)
-            .map_err(|_| LibError::OnnxModelEvalFailed)?
-            .with_flush_to_zero()
-            .map_err(|_| LibError::OnnxModelEvalFailed)?
-            .with_approximate_gelu()
-            .map_err(|_| LibError::OnnxModelEvalFailed)?
-            .commit_from_file(&onnx_path)
-            .map_err(|_| LibError::ModelWeightsLoadFailed)?;
+        let it = intra_threads();
+        let ns = num_sessions();
+        let mut sessions = Vec::with_capacity(ns);
+        for _ in 0..ns {
+            sessions.push(Mutex::new(Self::build_session(&onnx_path, it)?));
+        }
 
         Ok(Self {
-            session: SyncSession::new(session),
+            sessions,
             tokenizer,
             max_input_len,
             hidden_size,
@@ -827,7 +821,7 @@ impl OnnxEmbeddingModel {
 
     /// Run a single ONNX forward pass on one batch.
     fn run_batch(
-        session: &SyncSession,
+        session: &mut ort::session::Session,
         batch: &[Vec<u32>],
     ) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
         let batch_size = batch.len();
@@ -909,39 +903,77 @@ impl OnnxEmbeddingModel {
         Ok(embeddings)
     }
 
-    /// Pipelined predict: tokenizer thread prepares batch N+1 while ORT infers batch N.
-    fn predict_pipelined(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
+    fn tokenize_batch(
+        tokenizer: &Tokenizer,
+        batch: &[&str],
+        max_input: usize,
+    ) -> Result<Vec<Vec<u32>>, LibError> {
+        let truncated: Vec<&str> = batch
+            .iter()
+            .map(|t| pre_truncate_text(t, max_input))
+            .collect();
+        let encoded = tokenizer
+            .encode_batch(truncated, true)
+            .map_err(|_| LibError::ModelTokenizerEncodeFailed)?;
+        Ok(encoded
+            .iter()
+            .map(|enc| {
+                let ids = enc.get_ids();
+                ids[..ids.len().min(max_input)].to_vec()
+            })
+            .collect())
+    }
+
+    /// Predict using session pool with rayon (parallel tokenize+infer per batch).
+    fn predict_pool(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
+        use rayon::prelude::*;
+
         let bs = batch_size();
         let max_input = self.max_input_len;
-        let tokenizer = &self.tokenizer;
-        let session = &self.session;
+        let num_sessions = self.sessions.len();
+
+        let text_batches: Vec<&[&str]> = texts.chunks(bs).collect();
 
+        let results: Vec<Result<Vec<Vec<f32>>, LibError>> = text_batches
+            .par_iter()
+            .enumerate()
+            .map(|(i, batch)| {
+                let chunks = Self::tokenize_batch(&self.tokenizer, batch, max_input)?;
+                let mut session = self.sessions[i % num_sessions].lock().unwrap();
+                Self::run_batch(&mut session, &chunks).map_err(|_| LibError::OnnxModelEvalFailed)
+            })
+            .collect();
+
+        let mut all_embeddings = Vec::with_capacity(texts.len());
+        for result in results {
+            all_embeddings.extend(result.map_err(|e| -> Box<dyn Error> { Box::new(e) })?);
+        }
+        Ok(all_embeddings)
+    }
+
+    /// Pipelined predict: tokenizer thread prepares batch N+1 while ORT infers batch N.
+    fn predict_pipelined(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
         if texts.is_empty() {
             return Ok(Vec::new());
         }
 
-        let text_batches: Vec<&[&str]> = texts.chunks(bs).collect();
-        if text_batches.is_empty() {
-            return Ok(Vec::new());
+        // Session pool mode: use rayon parallel dispatch
+        if self.sessions.len() > 1 {
+            return self.predict_pool(texts);
         }
 
+        // Single session mode: pipeline tokenization with inference
+        let bs = batch_size();
+        let max_input = self.max_input_len;
+        let tokenizer = &self.tokenizer;
+
+        let text_batches: Vec<&[&str]> = texts.chunks(bs).collect();
+
         // Single batch — no pipeline overhead needed
         if text_batches.len() == 1 {
-            let truncated: Vec<&str> = text_batches[0]
-                .iter()
-                .map(|t| pre_truncate_text(t, max_input))
-                .collect();
-            let encoded = tokenizer
-                .encode_batch(truncated, true)
-                .map_err(|_| LibError::ModelTokenizerEncodeFailed)?;
-            let chunks: Vec<Vec<u32>> = encoded
-                .iter()
-                .map(|enc| {
-                    let ids = enc.get_ids();
-                    ids[..ids.len().min(max_input)].to_vec()
-                })
-                .collect();
-            return Self::run_batch(session, &chunks);
+            let chunks = Self::tokenize_batch(tokenizer, text_batches[0], max_input)
+                .map_err(|e| -> Box<dyn Error> { Box::new(e) })?;
+            return Self::run_batch(&mut self.sessions[0].lock().unwrap(), &chunks);
         }
 
         // Pipeline: tokenizer thread feeds batches, main thread runs inference
@@ -952,30 +984,20 @@ impl OnnxEmbeddingModel {
         std::thread::scope(|s| {
             s.spawn(move || {
                 for batch in &text_batches {
-                    let truncated: Vec<&str> = batch
-                        .iter()
-                        .map(|t| pre_truncate_text(t, max_input))
-                        .collect();
-                    let encoded = match tokenizer.encode_batch(truncated, true) {
-                        Ok(enc) => enc,
+                    match Self::tokenize_batch(tokenizer, batch, max_input) {
+                        Ok(chunks) => {
+                            if tx.send(chunks).is_err() {
+                                return;
+                            }
+                        }
                         Err(_) => return,
-                    };
-                    let chunks: Vec<Vec<u32>> = encoded
-                        .iter()
-                        .map(|enc| {
-                            let ids = enc.get_ids();
-                            ids[..ids.len().min(max_input)].to_vec()
-                        })
-                        .collect();
-                    if tx.send(chunks).is_err() {
-                        return;
                     }
                 }
-                // tx drops here — closes channel, rx loop ends
             });
 
+            let mut session = self.sessions[0].lock().unwrap();
             for chunks in rx {
-                match Self::run_batch(session, &chunks) {
+                match Self::run_batch(&mut session, &chunks) {
                     Ok(embs) => all_embeddings.extend(embs),
                     Err(_) => {
                         infer_error = Some(LibError::OnnxModelEvalFailed);