gigatoken 0.2.0 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 071b3abb52387367f392e88e50d92c7413f5f6c70dae24accf6d378ad6e78fd8
-  data.tar.gz: 26da76745e877b6ee30821ab5fcae50aefdb6ecddcd660e9b8c7b48b4d94bd06
+  metadata.gz: eee37e720c1b0fba0305398b731c3bd909af1a6ebc1ae7393e924c66fd629377
+  data.tar.gz: 966c52f87dcff1b0cda6218c45b0cbc5362a02655cb0c570034b4d91249a2518
 SHA512:
-  metadata.gz: 402d1e74262ffd0d4ff8d6f6162c944747b75c567249f788fdf46d2c5e09cd001f1cd66d037156848023fb5aad62087c73def5c4d3134d2aba7a898575ce2631
-  data.tar.gz: c43c48249a2c7bd0d8146e20d830975dae6e3acec93b1bc42e125fdd72050da38d02d0a10068b6ad208c85bd1790a54629b3b2394c1b12762feddce834909995
+  metadata.gz: 75eda453e48567dda4b39765031ed1dab4e5d12715507f0515e6742de2dd0ce67a9c4b5e4c1581664bce956b01cf8978cc9198ee3878253bd26287ba0d558aa2
+  data.tar.gz: 3a4e18c3e8603ae524cfa48752177277297c94343623571ad32bb87434eacfa3456264769f9b3daf5421ad5d615de47502607db17e6ded024e8c9b35a542b9c1

data/Cargo.lock

@@ -771,7 +771,7 @@ dependencies = [
 
 [[package]]
 name = "gigatoken-rb"
-version = "0.2.0"
+version = "0.2.1"
 dependencies = [
  "gigatoken",
  "magnus",

data/ext/gigatoken/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gigatoken-rb"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2021"
 
 [lib]

data/ext/gigatoken/src/sentencepiece.rs

@@ -13,7 +13,7 @@
 //! `Gigatoken::Error` instead of ever calling `str::from_utf8_unchecked` on
 //! Ruby-supplied bytes.
 
-use std::cell::RefCell;
+use std::sync::Mutex;
 
 use gigatoken_rs::input::file_source::DocFormat;
 use gigatoken_rs::{EncodeState, SentencePieceBPE, sp_encode_docs_ragged, sp_encode_files_docs, sp_encode_files_docs_serial};
@@ -35,10 +35,13 @@ fn require_utf8<'a>(ruby: &Ruby, bytes: &'a [u8]) -> Result<&'a str, Error> {
 #[magnus::wrap(class = "Gigatoken::Native::SentencePieceTokenizer", free_immediately, size)]
 pub struct SentencePieceTokenizer {
     // `SentencePieceBPE`'s encode methods take `&self` (only `EncodeState`
-    // is mutated), so this `RefCell` is never `borrow_mut`'d — see the
-    // builder report's DISAGREEMENTS for why it's here anyway.
-    tokenizer: RefCell<SentencePieceBPE>,
-    state: RefCell<EncodeState>,
+    // is mutated), so the model needs no interior mutability at all — every
+    // path here reads it, including the ones that release the GVL.
+    tokenizer: SentencePieceBPE,
+    // The one mutable piece. A `Mutex` rather than a `RefCell` so the wrapped
+    // object is `Sync`: Ruby hands the same instance to every thread, and a
+    // `RefCell` shared that way is unsound (see `BPETokenizer`'s lock).
+    state: Mutex<EncodeState>,
 }
 
 impl SentencePieceTokenizer {
@@ -48,8 +51,8 @@ impl SentencePieceTokenizer {
         let tokenizer = crate::cache::apply_max_cache_bytes_sp(tokenizer);
         let state = EncodeState::with_budget(tokenizer.max_cache_bytes());
         Self {
-            tokenizer: RefCell::new(tokenizer),
-            state: RefCell::new(state),
+            tokenizer,
+            state: Mutex::new(state),
         }
     }
 
@@ -58,8 +61,8 @@ impl SentencePieceTokenizer {
         let bytes = unsafe { input.as_slice() };
         let text = require_utf8(ruby, bytes)?;
         let mut ids: Vec<u32> = Vec::new();
-        let mut state = rb_self.state.borrow_mut();
-        rb_self.tokenizer.borrow().encode_raw_cb(&mut state, text, &mut |tokens| {
+        let mut state = rb_self.state.lock().unwrap_or_else(|e| e.into_inner());
+        rb_self.tokenizer.encode_raw_cb(&mut state, text, &mut |tokens| {
             ids.extend(tokens.iter().map(|&t| u32::from(t)))
         });
         Ok(ids)
@@ -84,8 +87,7 @@ impl SentencePieceTokenizer {
             })
             .collect::<Result<_, _>>()?;
         let doc_refs: Vec<&str> = docs.iter().map(String::as_str).collect();
-        let tokenizer = rb_self.tokenizer.borrow();
-        let tokenizer: &SentencePieceBPE = &tokenizer;
+        let tokenizer: &SentencePieceBPE = &rb_self.tokenizer;
         Ok(without_gvl(|| sp_encode_docs_ragged(tokenizer, &doc_refs)))
     }
 
@@ -130,8 +132,7 @@ impl SentencePieceTokenizer {
             }
         }
 
-        let tokenizer = rb_self.tokenizer.borrow();
-        let tokenizer: &SentencePieceBPE = &tokenizer;
+        let tokenizer: &SentencePieceBPE = &rb_self.tokenizer;
         let encoded: std::io::Result<(Vec<u32>, Vec<i64>)> = without_gvl(|| {
             sources::encode_files_ragged(&source, parallel, |files, format| {
                 for &region in files {
@@ -166,16 +167,16 @@ impl SentencePieceTokenizer {
     fn decode(ruby: &Ruby, rb_self: &Self, tokens: RArray) -> Result<RString, Error> {
         let ids: Vec<u32> = tokens.to_vec()?;
         let ids: Vec<_> = ids.into_iter().map(Into::into).collect();
-        let bytes = rb_self.tokenizer.borrow().decode(&ids);
+        let bytes = rb_self.tokenizer.decode(&ids);
         Ok(binary_string(ruby, &bytes))
     }
 
     fn vocab_size(&self) -> usize {
-        self.tokenizer.borrow().vocab_size()
+        self.tokenizer.vocab_size()
     }
 
     fn vocab(ruby: &Ruby, rb_self: &Self) -> Result<RHash, Error> {
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = &rb_self.tokenizer;
         let hash = ruby.hash_new();
         for (id, bytes) in tokenizer.vocab_entries() {
             hash.aset(id, binary_string(ruby, bytes))?;
@@ -184,7 +185,7 @@ impl SentencePieceTokenizer {
     }
 
     fn merges(ruby: &Ruby, rb_self: &Self) -> Result<RArray, Error> {
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = &rb_self.tokenizer;
         let entries = tokenizer.merge_entries();
         let result = ruby.ary_new_capa(entries.len());
         for (a, b) in entries {
@@ -196,7 +197,7 @@ impl SentencePieceTokenizer {
     /// Cached unit entries on the single-document `encode` path's state
     /// (batch encoders are per-call); see `BPETokenizer::cache_entries`.
     fn cache_entries(&self) -> usize {
-        self.state.borrow().cache_size()
+        self.state.lock().unwrap_or_else(|e| e.into_inner()).cache_size()
     }
 }
 

data/ext/gigatoken/src/tokenizer.rs

@@ -3,7 +3,7 @@
 //! the core crate's `src/lib.rs` (the `python` feature), minus the
 //! numpy/awkward-array machinery that has no Ruby analog.
 
-use std::cell::RefCell;
+use std::sync::RwLock;
 use std::collections::{HashMap, HashSet};
 use std::os::raw::c_long;
 
@@ -353,14 +353,14 @@ fn marshal_inputs(inputs: RArray) -> Result<InputDocs, Error> {
 
 #[magnus::wrap(class = "Gigatoken::Native::BPETokenizer", free_immediately, size)]
 pub struct BPETokenizer {
-    tokenizer: RefCell<Tokenizer>,
+    tokenizer: RwLock<Tokenizer>,
     workers: WorkerPool,
 }
 
 impl BPETokenizer {
     pub(crate) fn from_tokenizer(tokenizer: Tokenizer) -> Self {
         Self {
-            tokenizer: RefCell::new(crate::cache::apply_max_cache_bytes(tokenizer)),
+            tokenizer: RwLock::new(crate::cache::apply_max_cache_bytes(tokenizer)),
             workers: WorkerPool::new(),
         }
     }
@@ -407,14 +407,81 @@ impl BPETokenizer {
         }
     }
 
+    /// Shared access to the tokenizer, for everything that only reads it —
+    /// the batch paths, `decode`, `vocab`, `merges`, the size accessors.
+    ///
+    /// Readers never exclude each other, which is what makes the long holds
+    /// safe: `encode_batch`/`encode_files` keep this across a GVL release
+    /// (they run the core pool over `&Tokenizer`), and any other Ruby thread
+    /// reading meanwhile just proceeds. The only exclusive holder is
+    /// [`Self::encode`], which is short.
+    fn read_tokenizer(&self) -> std::sync::RwLockReadGuard<'_, Tokenizer> {
+        self.tokenizer.read().unwrap_or_else(|e| e.into_inner())
+    }
+
+    /// Encode one string, mutating the tokenizer's pretoken cache — the only
+    /// exclusive use of the lock.
+    ///
+    /// Uncontended (every single-threaded caller, and the common case under
+    /// threads) this takes the fast path: grab the write guard, encode against
+    /// the Ruby string's own bytes, never release the GVL. Identical cost to
+    /// the pre-lock version plus one uncontended atomic.
+    ///
+    /// Contended, the writer is waiting on a reader that will hold the lock
+    /// for as long as a batch encode takes. Blocking there while holding the
+    /// GVL would stall every other Ruby thread in the VM, so instead the input
+    /// is copied and the whole wait-and-encode moves inside `without_gvl`.
+    /// The copy is what makes that sound: no Ruby `VALUE` and no `RString`
+    /// buffer may outlive the release (see `marshal_inputs`), and the guard is
+    /// taken and dropped inside the closure, so it never crosses OS threads
+    /// even when the scheduler offloads it (see `gvl`).
     fn encode(&self, input: RString) -> Vec<u32> {
-        // SAFETY: read-only, for the duration of this synchronous call.
-        let bytes = unsafe { input.as_slice() };
-        let mut out = Vec::new();
-        self.tokenizer
-            .borrow_mut()
-            .encode_with_added_tokens_flat(bytes, &mut out);
-        out
+        if let Ok(mut tokenizer) = self.tokenizer.try_write() {
+            // SAFETY: read-only, for the duration of this synchronous call,
+            // with no GVL release in between.
+            let bytes = unsafe { input.as_slice() };
+            let mut out = Vec::new();
+            tokenizer.encode_with_added_tokens_flat(bytes, &mut out);
+            return out;
+        }
+
+        self.encode_contended(input)
+    }
+
+    /// The contended half of [`Self::encode`], outlined and `#[cold]`.
+    ///
+    /// Keeping this out of `encode`'s body is a measured requirement, not
+    /// tidiness: the workspace builds with `lto = "fat"`, so the core encode
+    /// routine inlines into `encode`, and inlining is sensitive to the caller's
+    /// size. Written inline, this second path measured slower on single
+    /// encodes — no lock overhead, just a flipped inlining decision. Outlined,
+    /// `encode`'s hot body is the original three lines behind a `try_write`.
+    ///
+    /// Before you re-inline this "to simplify": rerun the evidence rather than
+    /// trusting a number. `ruby -Ilib bench/encode_ab.rb` with the attributes
+    /// stripped and again with them restored, and read
+    /// `docs/rb/benchmarks.md` first — no size resolves this on the hardware
+    /// measured so far. The instrument is honest (an interleaved same-build
+    /// run never calls a size faster or slower, at any size) and has power to
+    /// catch a couple-percent effect reliably, but the attributes' real
+    /// effect is small enough that even the tightest floor (medium,
+    /// well under 3%) swallows it more often than not. Removing the
+    /// attributes measures slower at medium and large, in the direction the
+    /// outlining was added to prevent, but neither delta clears the noise
+    /// floor. Keep it outlined on that direction and on the original
+    /// inline-regression measurement, not on a pinned-down magnitude.
+    #[cold]
+    #[inline(never)]
+    fn encode_contended(&self, input: RString) -> Vec<u32> {
+        // SAFETY: copied before any GVL release, so nothing Ruby-owned is
+        // captured by the closure below.
+        let owned = unsafe { input.as_slice() }.to_vec();
+        without_gvl(move || {
+            let mut tokenizer = self.tokenizer.write().unwrap_or_else(|e| e.into_inner());
+            let mut out = Vec::new();
+            tokenizer.encode_with_added_tokens_flat(&owned, &mut out);
+            out
+        })
     }
 
     /// Encode a batch on the core worker pool, with the GVL released for the
@@ -425,7 +492,7 @@ impl BPETokenizer {
     fn encode_batch_ragged(rb_self: &Self, inputs: RArray) -> Result<(Vec<u32>, Vec<i64>), Error> {
         let marshaled = marshal_inputs(inputs)?;
         let doc_slices = marshaled.as_slices();
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = rb_self.read_tokenizer();
         let tokenizer: &Tokenizer = &tokenizer;
         let workers = &rb_self.workers;
         Ok(without_gvl(|| encode_docs_ragged(workers, tokenizer, &doc_slices)))
@@ -474,7 +541,7 @@ impl BPETokenizer {
         // the duration of the gather below (see the allocation above).
         let dest = unsafe { GatherBuf::new(ptr, total_bytes) };
 
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = rb_self.read_tokenizer();
         let tokenizer: &Tokenizer = &tokenizer;
         let workers = &rb_self.workers;
         match without_gvl(|| encode_docs_into(workers, tokenizer, &doc_slices, dest)) {
@@ -515,7 +582,7 @@ impl BPETokenizer {
         };
 
         let source = sources::resolve(ruby, source)?;
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = rb_self.read_tokenizer();
         let tokenizer: &Tokenizer = &tokenizer;
         let workers = &rb_self.workers;
         let encoded: std::io::Result<(Vec<u32>, Vec<i64>)> = without_gvl(|| {
@@ -547,16 +614,16 @@ impl BPETokenizer {
     fn decode(ruby: &Ruby, rb_self: &Self, tokens: RArray) -> Result<RString, Error> {
         let ids: Vec<u32> = tokens.to_vec()?;
         let ids: Vec<_> = ids.into_iter().map(Into::into).collect();
-        let bytes: Vec<u8> = rb_self.tokenizer.borrow().decode(&ids).collect();
+        let bytes: Vec<u8> = rb_self.read_tokenizer().decode(&ids).collect();
         Ok(binary_string(ruby, &bytes))
     }
 
     fn vocab_size(&self) -> usize {
-        self.tokenizer.borrow().vocab_size()
+        self.read_tokenizer().vocab_size()
     }
 
     fn vocab(ruby: &Ruby, rb_self: &Self) -> Result<RHash, Error> {
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = rb_self.read_tokenizer();
         let hash = ruby.hash_new();
         for (id, bytes) in tokenizer.vocab_entries() {
             hash.aset(id, binary_string(ruby, bytes))?;
@@ -565,7 +632,7 @@ impl BPETokenizer {
     }
 
     fn merges(ruby: &Ruby, rb_self: &Self) -> Result<RArray, Error> {
-        let tokenizer = rb_self.tokenizer.borrow();
+        let tokenizer = rb_self.read_tokenizer();
         let entries = tokenizer.merge_entries();
         let result = ruby.ary_new_capa(entries.len());
         for (a, b) in entries {
@@ -578,7 +645,7 @@ impl BPETokenizer {
     /// drops back toward vocab-seed level when a budgeted cache wipes (see
     /// `Gigatoken.max_cache_bytes`).
     fn cache_entries(&self) -> usize {
-        self.tokenizer.borrow().cache_entries()
+        self.read_tokenizer().cache_entries()
     }
 }
 

data/lib/gigatoken/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Gigatoken
-  VERSION = "0.2.0"
+  VERSION = "0.2.1"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gigatoken
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.2.1
 platform: ruby
 authors:
 - Eric Jacobs