perf(embeddings): rotate sessions via round-robin

Author: donhardman

embeddings/src/model/local.rs

@@ -780,6 +780,7 @@ pub struct OnnxEmbeddingModel {
     tokenizer: Tokenizer,
     max_input_len: usize,
     hidden_size: usize,
+    next_session: std::sync::atomic::AtomicUsize,
 }
 
 impl OnnxEmbeddingModel {
@@ -831,9 +832,17 @@ impl OnnxEmbeddingModel {
             tokenizer,
             max_input_len,
             hidden_size,
+            next_session: std::sync::atomic::AtomicUsize::new(0),
         })
     }
 
+    /// Pick next session round-robin for concurrent callers.
+    fn next_session_idx(&self) -> usize {
+        self.next_session
+            .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
+            % self.sessions.len()
+    }
+
     /// Run a single ONNX forward pass on one batch.
     fn run_batch(
         session: &mut ort::session::Session,
@@ -984,11 +993,12 @@ impl OnnxEmbeddingModel {
 
         let text_batches: Vec<&[&str]> = texts.chunks(bs).collect();
 
-        // Single batch — no pipeline overhead needed
+        // Single batch — pick session round-robin for concurrent callers
         if text_batches.len() == 1 {
+            let idx = self.next_session_idx();
             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);
+            return Self::run_batch(&mut self.sessions[idx].lock().unwrap(), &chunks);
         }
 
         // Pipeline: tokenizer thread feeds batches, main thread runs inference
@@ -1010,7 +1020,8 @@ impl OnnxEmbeddingModel {
                 }
             });
 
-            let mut session = self.sessions[0].lock().unwrap();
+            let idx = self.next_session_idx();
+            let mut session = self.sessions[idx].lock().unwrap();
             for chunks in rx {
                 match Self::run_batch(&mut session, &chunks) {
                     Ok(embs) => all_embeddings.extend(embs),