gte 0.0.12 → 0.0.14

data/ext/gte/src/session.rs

@@ -4,221 +4,151 @@ use crate::pipeline::{extract_output_tensor, InputTensors};
 use crate::postprocess::mean_pool;
 use crate::tokenizer::Tokenized;
 use ndarray::{Array2, ArrayView2, ArrayViewD, Ix2};
-use ort::execution_providers::{
-    CoreMLExecutionProvider, ExecutionProviderDispatch, XNNPACKExecutionProvider,
-};
-use ort::session::Session;
+use ort::execution_providers::{CoreMLExecutionProvider, ExecutionProviderDispatch, XNNPACKExecutionProvider};
+use ort::session::{OutputSelector, RunOptions, Session};
+use std::cell::RefCell;
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Condvar, Mutex};
-
-pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Result<Session> {
-    let opt_level = match config.optimization_level {
-        0 => ort::session::builder::GraphOptimizationLevel::Disable,
-        1 => ort::session::builder::GraphOptimizationLevel::Level1,
-        2 => ort::session::builder::GraphOptimizationLevel::Level2,
-        _ => ort::session::builder::GraphOptimizationLevel::Level3,
-    };
-
-    fn ort_err(e: impl std::fmt::Display) -> GteError {
-        GteError::Ort(e.to_string())
-    }
-
-    let mut builder = Session::builder()
-        .map_err(ort_err)?
-        .with_optimization_level(opt_level)
-        .map_err(ort_err)?;
-
-    let providers = preferred_execution_providers(config.execution_providers.as_deref());
-    if !providers.is_empty() {
-        builder = builder
-            .with_execution_providers(providers)
-            .map_err(ort_err)?;
-    }
-
-    builder.commit_from_file(model_path).map_err(ort_err)
-}
 
 // ---------------------------------------------------------------------------
-// Session pool
+// Thread-local session storage — each OS thread lazily creates its own ONNX
+// session the first time it calls into a given pool.  No Mutex, no contention.
 // ---------------------------------------------------------------------------
 
-fn pool_capacity() -> usize {
-    let available = std::thread::available_parallelism()
-        .map(|n| n.get())
-        .unwrap_or(1);
-    parse_pool_capacity_override().map_or(available, |cap| cap.min(available).max(1))
+static NEXT_POOL_ID: AtomicUsize = AtomicUsize::new(1);
+
+struct SessionRecipe {
+    model_path: PathBuf,
+    build_config: ModelConfig,
 }
 
-fn parse_pool_capacity_override() -> Option<usize> {
-    let raw = std::env::var("GTE_SESSION_POOL_CAP").ok()?;
-    let parsed = raw.trim().parse::<usize>().ok()?;
-    (parsed > 0).then_some(parsed)
+thread_local! {
+    static SESSIONS: RefCell<HashMap<usize, Session>> = RefCell::new(HashMap::new());
 }
 
 pub struct SessionPool {
-    sessions: Mutex<Vec<Session>>,
-    available: Condvar,
-    created: AtomicUsize,
-    capacity: usize,
-    model_path: PathBuf,
-    build_config: ModelConfig,
+    pool_id: usize,
+    recipe: SessionRecipe,
 }
 
 impl SessionPool {
-    pub fn new(initial: Session, model_path: PathBuf, build_config: ModelConfig) -> Self {
-        let capacity = pool_capacity();
-        Self {
-            sessions: Mutex::new(vec![initial]),
-            available: Condvar::new(),
-            created: AtomicUsize::new(1),
-            capacity,
-            model_path,
-            build_config,
-        }
-    }
-
-    pub fn acquire(&self) -> Result<PooledSession<'_>> {
-        if let Some(session) = self.take_available_session() {
-            return Ok(PooledSession {
-                pool: self,
-                session: Some(session),
-            });
-        }
+    pub fn new(initial: Session, model_path: &Path, build_config: &ModelConfig) -> Result<Self> {
+        let pool_id = NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed);
 
-        if let Some(session) = self.try_grow()? {
-            return Ok(PooledSession {
-                pool: self,
-                session: Some(session),
-            });
-        }
+        SESSIONS.with(|map| {
+            _ = map.borrow_mut().insert(pool_id, initial);
+        });
 
-        let session = self.wait_for_session();
-        Ok(PooledSession {
-            pool: self,
-            session: Some(session),
+        Ok(Self {
+            pool_id,
+            recipe: SessionRecipe { model_path: model_path.to_path_buf(), build_config: build_config.clone() },
         })
     }
 
-    fn release(&self, session: Session) {
-        self.sessions.lock().unwrap().push(session);
-        self.available.notify_one();
+    pub fn run(&self, tokenized: &Tokenized, config: &ModelConfig) -> Result<Array2<f32>> {
+        self.with_session(|session| run_session(session, tokenized, config))
     }
 
-    fn take_available_session(&self) -> Option<Session> {
-        self.sessions.lock().unwrap().pop()
+    pub fn with_session<F, R>(&self, f: F) -> Result<R>
+    where
+        F: FnOnce(&mut Session) -> Result<R>,
+    {
+        SESSIONS.with(|map| {
+            let mut map = map.borrow_mut();
+            let session = match map.entry(self.pool_id) {
+                Entry::Occupied(e) => e.into_mut(),
+                Entry::Vacant(e) => {
+                    let session = build_session(&self.recipe.model_path, &self.recipe.build_config)?;
+                    e.insert(session)
+                }
+            };
+            f(session)
+        })
     }
+}
 
-    fn try_grow(&self) -> Result<Option<Session>> {
-        let grew = self
-            .created
-            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
-                (count < self.capacity).then_some(count + 1)
-            });
-        if grew.is_err() {
-            return Ok(None);
-        }
+// ---------------------------------------------------------------------------
+// Session construction
+// ---------------------------------------------------------------------------
 
-        match build_session(&self.model_path, &self.build_config) {
-            Ok(session) => Ok(Some(session)),
-            Err(error) => {
-                self.created.fetch_sub(1, Ordering::AcqRel);
-                Err(error)
-            }
-        }
+pub fn build_session<P: AsRef<Path>>(model_path: P, config: &ModelConfig) -> Result<Session> {
+    fn ort_err(e: impl std::fmt::Display) -> GteError {
+        GteError::Ort(e.to_string())
     }
 
-    fn wait_for_session(&self) -> Session {
-        let mut lock = self.sessions.lock().unwrap();
-        loop {
-            if let Some(session) = lock.pop() {
-                return session;
-            }
-            lock = self.available.wait(lock).unwrap();
-        }
-    }
-}
+    let opt_level = match config.optimization_level {
+        0 => ort::session::builder::GraphOptimizationLevel::Disable,
+        1 => ort::session::builder::GraphOptimizationLevel::Level1,
+        2 => ort::session::builder::GraphOptimizationLevel::Level2,
+        _ => ort::session::builder::GraphOptimizationLevel::Level3,
+    };
 
-pub struct PooledSession<'a> {
-    pool: &'a SessionPool,
-    session: Option<Session>,
-}
+    let mut builder = Session::builder().map_err(ort_err)?.with_optimization_level(opt_level).map_err(ort_err)?;
 
-impl std::ops::Deref for PooledSession<'_> {
-    type Target = Session;
-    fn deref(&self) -> &Session {
-        self.session.as_ref().unwrap()
-    }
-}
+    let intra_threads = std::env::var("GTE_INTRA_OP_NUM_THREADS")
+        .ok()
+        .and_then(|v| v.trim().parse::<usize>().ok())
+        .unwrap_or_else(|| std::thread::available_parallelism().map(|n| n.get().min(4)).unwrap_or(1));
+    builder = builder.with_intra_threads(intra_threads).map_err(ort_err)?;
 
-impl std::ops::DerefMut for PooledSession<'_> {
-    fn deref_mut(&mut self) -> &mut Session {
-        self.session.as_mut().unwrap()
-    }
-}
+    let inter_threads =
+        std::env::var("GTE_INTER_OP_NUM_THREADS").ok().and_then(|v| v.trim().parse::<usize>().ok()).unwrap_or(1);
+    builder = builder.with_inter_threads(inter_threads).map_err(ort_err)?;
 
-impl Drop for PooledSession<'_> {
-    fn drop(&mut self) {
-        if let Some(s) = self.session.take() {
-            self.pool.release(s);
-        }
+    let providers = match config.execution_providers.as_deref() {
+        Some(override_val) => preferred_execution_providers(Some(override_val)),
+        None => auto_detect_providers(),
+    };
+    if !providers.is_empty() {
+        builder = builder.with_execution_providers(providers).map_err(ort_err)?;
     }
-}
-
-// ---------------------------------------------------------------------------
 
-fn preferred_execution_providers(order_override: Option<&str>) -> Vec<ExecutionProviderDispatch> {
-    let order = resolve_provider_order(order_override);
+    builder.commit_from_file(model_path).map_err(ort_err)
+}
 
+fn auto_detect_providers() -> Vec<ExecutionProviderDispatch> {
     let mut providers = Vec::new();
-    for provider in parse_provider_registrations(order.as_str()) {
-        match provider {
-            "xnnpack" => {
-                providers.push(XNNPACKExecutionProvider::default().build().fail_silently())
-            }
-            "coreml" => providers.push(CoreMLExecutionProvider::default().build().fail_silently()),
-            _ => {}
-        }
-    }
+    #[cfg(target_arch = "aarch64")]
+    providers.push(XNNPACKExecutionProvider::default().build().fail_silently());
     providers
 }
 
-fn resolve_provider_order(order_override: Option<&str>) -> String {
-    let env_order = std::env::var("GTE_EXECUTION_PROVIDERS").ok();
-    resolve_provider_order_with_env(order_override, env_order.as_deref())
-}
-
-fn resolve_provider_order_with_env(
-    order_override: Option<&str>,
-    env_order: Option<&str>,
-) -> String {
-    order_override
-        .or(env_order)
-        .unwrap_or("cpu")
-        .to_ascii_lowercase()
-}
+fn preferred_execution_providers(order_override: Option<&str>) -> Vec<ExecutionProviderDispatch> {
+    let order = match order_override {
+        Some(s) => s.to_ascii_lowercase(),
+        None => return auto_detect_providers(),
+    };
 
-fn parse_provider_registrations(order: &str) -> Vec<&str> {
-    let mut providers = Vec::new();
-    for provider in order.split(',').map(str::trim).filter(|p| !p.is_empty()) {
-        match provider {
-            "xnnpack" | "coreml" => providers.push(provider),
-            "none" | "cpu" => {}
-            _ => {}
-        }
+    if order.is_empty() || order == "cpu" || order == "none" {
+        return Vec::new();
     }
+
+    let providers: Vec<_> = order
+        .split(',')
+        .map(str::trim)
+        .filter(|p| !p.is_empty())
+        .filter_map(|provider| match provider {
+            "xnnpack" => Some(XNNPACKExecutionProvider::default().build().fail_silently()),
+            "coreml" => Some(CoreMLExecutionProvider::default().build().fail_silently()),
+            _ => None,
+        })
+        .collect();
     providers
 }
 
-pub fn run_session(
-    session: &mut Session,
-    tokenized: &Tokenized,
-    config: &ModelConfig,
-) -> Result<Array2<f32>> {
+// ---------------------------------------------------------------------------
+// Run a single inference
+// ---------------------------------------------------------------------------
+
+pub fn run_session(session: &mut Session, tokenized: &Tokenized, config: &ModelConfig) -> Result<Array2<f32>> {
     let input_tensors = InputTensors::from_tokenized(tokenized, config.with_attention_mask)?;
-    let outputs = session
-        .run(input_tensors.inputs)
-        .map_err(|e| GteError::Ort(e.to_string()))?;
+    let run_opts = RunOptions::new()
+        .map_err(|e| GteError::Ort(e.to_string()))?
+        .with_outputs(OutputSelector::no_default().with(config.output_tensor.as_str()));
+    let outputs =
+        session.run_with_options(input_tensors.inputs, &run_opts).map_err(|e| GteError::Ort(e.to_string()))?;
     let array = extract_output_tensor(&outputs, config.output_tensor.as_str())?;
 
     extract_embeddings(array, input_tensors.attention_mask, config)
@@ -234,26 +164,21 @@ fn extract_embeddings(
             let shape = array.shape();
             if shape.len() != 3 || idx >= shape[1] {
                 return Err(GteError::Inference(format!(
-                    "token extraction index {} out of bounds for output shape {:?}",
-                    idx, shape
+                    "token extraction index {idx} out of bounds for output shape {shape:?}"
                 )));
             }
             Ok(array.slice(ndarray::s![.., idx, ..]).into_owned())
         }
         ExtractorMode::MeanPool => {
             let ndim = array.ndim();
-            let hidden_states = array.into_dimensionality::<ndarray::Ix3>().map_err(|_| {
-                GteError::Inference(format!(
-                    "mean pooling requires rank-3 output, got rank {}",
-                    ndim
-                ))
-            })?;
+            let hidden_states = array
+                .into_dimensionality::<ndarray::Ix3>()
+                .map_err(|_| GteError::Inference(format!("mean pooling requires rank-3 output, got rank {ndim}")))?;
             mean_pool(hidden_states, attention_mask)
         }
-        ExtractorMode::Raw => array
-            .into_dimensionality::<Ix2>()
-            .map(|view| view.to_owned())
-            .map_err(|e| GteError::Shape(e.to_string())),
+        ExtractorMode::Raw => {
+            array.into_dimensionality::<Ix2>().map(|view| view.to_owned()).map_err(|e| GteError::Shape(e.to_string()))
+        }
     }
 }
 
@@ -262,10 +187,22 @@ mod tests {
     use crate::model_config::{ExtractorMode, ModelConfig, PaddingMode};
     use ndarray::{array, ArrayView2};
 
-    use super::{
-        extract_embeddings, parse_pool_capacity_override, parse_provider_registrations,
-        resolve_provider_order_with_env,
-    };
+    use super::extract_embeddings;
+
+    fn resolve_provider_order_with_env(order_override: Option<&str>, env_order: Option<&str>) -> String {
+        order_override.or(env_order).unwrap_or("cpu").to_ascii_lowercase()
+    }
+
+    fn parse_provider_registrations(order: &str) -> Vec<&str> {
+        let mut providers = Vec::new();
+        for provider in order.split(',').map(str::trim).filter(|p| !p.is_empty()) {
+            match provider {
+                "xnnpack" | "coreml" => providers.push(provider),
+                _ => {}
+            }
+        }
+        providers
+    }
 
     fn test_config(mode: ExtractorMode) -> ModelConfig {
         ModelConfig {
@@ -277,6 +214,8 @@ mod tests {
             with_attention_mask: true,
             optimization_level: 3,
             execution_providers: None,
+            lowercase_input: false,
+            max_input_chars: None,
         }
     }
 
@@ -306,93 +245,45 @@ mod tests {
 
     #[test]
     fn resolve_provider_order_prefers_override() {
-        assert_eq!(
-            resolve_provider_order_with_env(Some("xnnpack"), Some("coreml")),
-            "xnnpack"
-        );
+        assert_eq!(resolve_provider_order_with_env(Some("xnnpack"), Some("coreml")), "xnnpack");
         assert_eq!(resolve_provider_order_with_env(Some("CPU"), None), "cpu");
     }
 
     #[test]
     fn resolve_provider_order_falls_back_to_env_then_cpu_default() {
-        assert_eq!(
-            resolve_provider_order_with_env(None, Some("coreml")),
-            "coreml"
-        );
+        assert_eq!(resolve_provider_order_with_env(None, Some("coreml")), "coreml");
         assert_eq!(resolve_provider_order_with_env(None, None), "cpu");
     }
 
-    #[test]
-    fn parse_pool_capacity_override_uses_positive_integer_only() {
-        unsafe {
-            std::env::remove_var("GTE_SESSION_POOL_CAP");
-        }
-        assert_eq!(parse_pool_capacity_override(), None);
-
-        unsafe {
-            std::env::set_var("GTE_SESSION_POOL_CAP", "0");
-        }
-        assert_eq!(parse_pool_capacity_override(), None);
-
-        unsafe {
-            std::env::set_var("GTE_SESSION_POOL_CAP", "4");
-        }
-        assert_eq!(parse_pool_capacity_override(), Some(4));
-
-        unsafe {
-            std::env::set_var("GTE_SESSION_POOL_CAP", "abc");
-        }
-        assert_eq!(parse_pool_capacity_override(), None);
-
-        unsafe {
-            std::env::remove_var("GTE_SESSION_POOL_CAP");
-        }
-    }
-
     #[test]
     fn extract_embeddings_raw_copies_only_final_matrix() {
         let output = array![[1.0f32, 2.0], [3.0, 4.0]];
-        let extracted = extract_embeddings(
-            output.view().into_dyn(),
-            empty_attention_mask(),
-            &test_config(ExtractorMode::Raw),
-        )
-        .unwrap();
+        let extracted =
+            extract_embeddings(output.view().into_dyn(), empty_attention_mask(), &test_config(ExtractorMode::Raw))
+                .unwrap();
 
         assert_eq!(extracted, output);
     }
 
     #[test]
     fn extract_embeddings_token_selects_without_copying_full_sequence() {
-        let output = array![
-            [[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]],
-            [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]
-        ];
+        let output = array![[[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]], [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]]];
         let expected = array![[3.0f32, 4.0], [9.0, 10.0]];
-        let extracted = extract_embeddings(
-            output.view().into_dyn(),
-            empty_attention_mask(),
-            &test_config(ExtractorMode::Token(1)),
-        )
-        .unwrap();
+        let extracted =
+            extract_embeddings(output.view().into_dyn(), empty_attention_mask(), &test_config(ExtractorMode::Token(1)))
+                .unwrap();
 
         assert_eq!(extracted, expected);
     }
 
     #[test]
     fn extract_embeddings_mean_pool_uses_output_view_and_attention_mask() {
-        let output = array![
-            [[1.0f32, 3.0], [5.0, 7.0], [100.0, 100.0]],
-            [[2.0, 4.0], [6.0, 8.0], [10.0, 12.0]]
-        ];
+        let output = array![[[1.0f32, 3.0], [5.0, 7.0], [100.0, 100.0]], [[2.0, 4.0], [6.0, 8.0], [10.0, 12.0]]];
         let attention_mask = array![[1_i64, 1, 0], [0, 1, 1]];
         let expected = array![[3.0f32, 5.0], [8.0, 10.0]];
-        let extracted = extract_embeddings(
-            output.view().into_dyn(),
-            attention_mask.view(),
-            &test_config(ExtractorMode::MeanPool),
-        )
-        .unwrap();
+        let extracted =
+            extract_embeddings(output.view().into_dyn(), attention_mask.view(), &test_config(ExtractorMode::MeanPool))
+                .unwrap();
 
         assert_eq!(extracted, expected);
     }