gte 0.0.14 → 0.0.15

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 37ad2ef3f640b8bbaefa14cae29541f329c3b99950a2867b9054b8e8854ca242
-  data.tar.gz: 0b54c757ca510ccc8644e0d3c1519aa7b6576fa1da4ba9012535e0b0b7d598dc
+  metadata.gz: 9edaf702ddb74aa67328a218eef5ce5c211868ae9e59aac312dd4bd471c5ba51
+  data.tar.gz: fb7ab71c6ee482afc2db321b2a84222956716da906d8df5eee1e962eb26dbd7f
 SHA512:
-  metadata.gz: 5d67fc8c73aa2b162bf804082accc849ce972b4df51e2ec86bb7d977cd760e2bf5b760d7de141b5fbd5b6f68dd0e51d2bd15c84d1cb8dc9a8f4bf424b62490ba
-  data.tar.gz: c6cefec4d42f7ca72980e4d3454447aa681f05f0b72f96ceec226b5efd5ae32bade13d85ae55f032098d7cb7f72b530d8332f1eaf25fa364d0c33703da13e043
+  metadata.gz: 8ddb7b0a4f2b84c5d48c4007b51143b8275980332fe48c47ae444e47dc71691cb24f9dcbb386f17eb80e58ae30d981544f377be238d3872f11ee0c64be404ed7
+  data.tar.gz: 95173dde395043b1a792c8954548a7334d6bc1e528c39f4d54d9bed6796e949806e1c359e42b17d403b316f46005a31f0033785055fa6dc725a50a8e2599117c

data/VERSION

@@ -1 +1 @@
-0.0.14
+0.0.15

data/ext/gte/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gte"
-version = "0.0.14"
+version = "0.0.15"
 edition = "2021"
 authors = ["elcuervo <elcuervo@elcuervo.net>"]
 license = "MIT"
@@ -22,6 +22,7 @@ ruby-ffi = ["dep:magnus", "dep:rb-sys"]
 rb-sys = { version = "0.9", features = ["stable-api-compiled-fallback"], optional = true }
 magnus = { version = "0.8", optional = true }
 ort = { version = "=2.0.0-rc.12", features = ["ndarray", "xnnpack"] }
+parking_lot = "0.12"
 tokenizers = "0.21.0"
 ndarray = "0.17"
 serde_json = "1"

data/ext/gte/src/session.rs

@@ -6,44 +6,57 @@ use crate::tokenizer::Tokenized;
 use ndarray::{Array2, ArrayView2, ArrayViewD, Ix2};
 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 parking_lot::Mutex;
 use std::path::{Path, PathBuf};
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
 // ---------------------------------------------------------------------------
-// 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.
+// Lazy session pool — starts with 1 session, grows on contention, capped.
+//
+// Pool max is resolved in order:
+//   1. GTE_SESSION_POOL_SIZE env var (explicit override)
+//   2. Auto: 2 (conservative: 2× pure Ruby memory at peak, no OOM risk)
+//
+// At idle the pool holds 1 session (same memory as pure Ruby's single
+// OnnxRuntime::Model).  When all existing sessions are busy and the cap
+// hasn't been reached, a new session is created on-demand.
 // ---------------------------------------------------------------------------
 
-static NEXT_POOL_ID: AtomicUsize = AtomicUsize::new(1);
-
-struct SessionRecipe {
-    model_path: PathBuf,
-    build_config: ModelConfig,
+fn resolve_pool_cap() -> usize {
+    if let Some(n) =
+        std::env::var("GTE_SESSION_POOL_SIZE").ok().and_then(|v| v.trim().parse::<usize>().ok()).filter(|&n| n > 0)
+    {
+        return n;
+    }
+    2
 }
 
-thread_local! {
-    static SESSIONS: RefCell<HashMap<usize, Session>> = RefCell::new(HashMap::new());
+pub struct SessionPool {
+    inner: Mutex<PoolInner>,
+    next_idx: AtomicUsize,
+    cap: usize,
 }
 
-pub struct SessionPool {
-    pool_id: usize,
-    recipe: SessionRecipe,
+struct PoolInner {
+    sessions: Vec<Arc<Mutex<Session>>>,
+    model_path: PathBuf,
+    build_config: ModelConfig,
 }
 
 impl SessionPool {
     pub fn new(initial: Session, model_path: &Path, build_config: &ModelConfig) -> Result<Self> {
-        let pool_id = NEXT_POOL_ID.fetch_add(1, Ordering::Relaxed);
-
-        SESSIONS.with(|map| {
-            _ = map.borrow_mut().insert(pool_id, initial);
-        });
+        let cap = resolve_pool_cap();
+        let sessions = vec![Arc::new(Mutex::new(initial))];
 
         Ok(Self {
-            pool_id,
-            recipe: SessionRecipe { model_path: model_path.to_path_buf(), build_config: build_config.clone() },
+            inner: Mutex::new(PoolInner {
+                sessions,
+                model_path: model_path.to_path_buf(),
+                build_config: build_config.clone(),
+            }),
+            next_idx: AtomicUsize::new(0),
+            cap,
         })
     }
 
@@ -55,17 +68,59 @@ impl SessionPool {
     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)
+        const SPIN_LIMIT: u32 = 64;
+
+        loop {
+            // Snapshot the pool under the outer lock so the scan below
+            // doesn't contend on that lock at all.
+            let arcs: Vec<Arc<Mutex<Session>>> = {
+                let inner = self.inner.lock();
+                inner.sessions.clone()
+            };
+            let len = arcs.len();
+            let start = self.next_idx.fetch_add(1, Ordering::Relaxed) % len;
+
+            for offset in 0..len {
+                let idx = (start + offset) % len;
+                if let Some(mut guard) = arcs[idx].try_lock() {
+                    return f(&mut guard);
+                }
+            }
+
+            // All sessions busy — try to grow the pool
+            let grew = {
+                let mut inner = self.inner.lock();
+                if inner.sessions.len() < self.cap {
+                    match build_session(&inner.model_path, &inner.build_config) {
+                        Ok(session) => {
+                            inner.sessions.push(Arc::new(Mutex::new(session)));
+                            true
+                        }
+                        Err(e) => return Err(e),
+                    }
+                } else {
+                    false
                 }
             };
-            f(session)
-        })
+
+            if grew {
+                continue;
+            }
+
+            // At cap — spin briefly, then block on a session
+            let idx = self.next_idx.fetch_add(1, Ordering::Relaxed) % len;
+            let arc = Arc::clone(&arcs[idx]);
+
+            for _ in 0..SPIN_LIMIT {
+                if let Some(mut guard) = arc.try_lock() {
+                    return f(&mut guard);
+                }
+                std::hint::spin_loop();
+            }
+
+            let mut guard = arc.lock();
+            return f(&mut guard);
+        }
     }
 }
 

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gte
 version: !ruby/object:Gem::Version
-  version: 0.0.14
+  version: 0.0.15
 platform: ruby
 authors:
 - elcuervo