Harden Candle::Reranker for Benchmate: revision pinning, bounded batching, concurrency safety

.github/workflows/build.yml

@@ -112,6 +112,10 @@ jobs:
             spec/reranker_spec.rb
             spec/reranker_truncation_spec.rb
             spec/reranker_max_length_spec.rb
+            spec/reranker_revision_spec.rb
+            spec/reranker_batching_spec.rb
+            spec/reranker_concurrency_spec.rb
+            spec/reranker_parity_spec.rb
             spec/ner_spec.rb
             spec/ner_core_functionality_spec.rb
             spec/ner_comprehensive_spec.rb

README.md

@@ -634,9 +634,92 @@ Choose based on your needs:
 
 5. **Memory Usage**:
    - Model size: ~125MB
-   - Each batch processes all documents simultaneously
+   - Documents are processed in bounded micro-batches (`batch_size:`, default 16) rather
+     than one tensor for every candidate — see [Production Usage](#production-usage-immutable-revisions-batching--concurrency) below
    - Consider batching if you have many documents
 
+## Production Usage: Immutable Revisions, Batching & Concurrency
+
+For running a reranker in-process as a long-lived, shared singleton (e.g. one instance
+memoized behind a Puma/Rails app, reused across requests and threads), a few options
+matter beyond the basics above.
+
+### Pinning an Immutable Revision
+
+By default `from_pretrained` tracks the repository's default branch ("main"), matching
+prior versions of Red Candle. Pass `revision:` to pin the model, tokenizer, and config to
+an exact HuggingFace revision (a commit SHA, tag, or branch) — production code should
+almost always do this, so a model repo update can never silently change what's loaded:
+
+```ruby
+reranker = Candle::Reranker.from_pretrained(
+  "cross-encoder/ms-marco-MiniLM-L6-v2",
+  device: Candle::Device.cpu,
+  revision: "233902d25c440f23af6f7d6e94d2946bac0bee0a"
+)
+
+reranker.options["revision"] # => "233902d25c440f23af6f7d6e94d2946bac0bee0a"
+```
+
+When `revision:` is given, files are fetched from exactly that revision — there is no
+fallback to the default branch. Both `options` and `inspect` surface the active revision,
+and load/inference errors include the model id and revision to make failures easy to
+diagnose. Downloaded files are cached locally by `hf-hub` the same way regardless of
+revision (see [`docs/HUGGINGFACE.md`](docs/HUGGINGFACE.md)), so a pinned revision does not
+mean re-downloading on every process start.
+
+### Bounded Micro-Batching
+
+`batch_size:` (default `16`) controls how many query-document pairs are processed per
+forward pass. Passing 100 documents with `batch_size: 16` runs 7 micro-batches
+internally instead of building one tensor sized for all 100 candidates at once — this
+keeps peak memory bounded on CPU regardless of how many candidates a caller sends in one
+`#rerank` call. Results are identical to a single large batch: each document's global
+`doc_id` is preserved, and the final ranking is sorted once, globally, by descending
+score (ties broken by ascending original index). `batch_size` must be a positive integer;
+anything else raises `ArgumentError`.
+
+```ruby
+reranker = Candle::Reranker.from_pretrained(
+  "cross-encoder/ms-marco-MiniLM-L6-v2", revision: "233902d25c440f23af6f7d6e94d2946bac0bee0a",
+  batch_size: 16
+)
+reranker.rerank(query, one_hundred_documents) # processed as 7 micro-batches of ≤16
+```
+
+See [`bench/reranker_bench.rb`](bench/reranker_bench.rb) and
+[`docs/RERANKER_BENCHMARKS.md`](docs/RERANKER_BENCHMARKS.md) for latency/throughput/memory
+measurements across candidate counts, `max_length`, and `batch_size`.
+
+### Thread Safety & Singleton Reuse
+
+A single `Candle::Reranker` instance is safe to call concurrently from multiple Ruby
+threads for **BERT, XLM-RoBERTa, DeBERTa, and ModernBERT** models (which includes
+`cross-encoder/ms-marco-MiniLM-L6-v2`) — this is exactly the memoized-singleton-shared-
+across-Puma-threads pattern. `#rerank` releases the GVL for tokenization, inference,
+scoring, and sorting, so concurrent calls run in true parallel — there is no internal
+lock serializing them, because these architectures have no mutable state to protect
+(weights are loaded once and never mutated; candle's tensor storage is designed for
+concurrent reads).
+
+**Qwen3** reranker models are the one exception: they hold a KV cache that is genuinely
+mutated during each forward pass, so concurrent calls on one Qwen3 reranker instance are
+internally serialized with a narrowly-scoped lock (acquired without holding the GVL, so
+waiting threads don't block unrelated Ruby work). If you need concurrent Qwen3 reranking
+throughput, use one reranker instance per thread instead of sharing one.
+
+```ruby
+# Safe: one reranker, shared across request threads
+RERANKER = Candle::Reranker.from_pretrained(
+  "cross-encoder/ms-marco-MiniLM-L6-v2",
+  revision: "233902d25c440f23af6f7d6e94d2946bac0bee0a",
+  device: Candle::Device.cpu
+).freeze
+```
+
+See `spec/reranker_concurrency_spec.rb` for a regression test that exercises this with 8
+threads making concurrent calls through one instance.
+
 ## Tokenizer
 
 Red-Candle provides direct access to tokenizers for text preprocessing and analysis. This is useful for understanding how models process text, debugging issues, and building custom NLP pipelines.

bench/reranker_bench.rb

@@ -0,0 +1,177 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# Reproducible reranker benchmark for Candle::Reranker, sized for Benchmate's target
+# model/workload: cross-encoder/ms-marco-MiniLM-L6-v2, up to 100 candidates per call,
+# CPU-only inference.
+#
+# Run on Linux amd64 CPU (the production target) with:
+#
+#   bundle exec ruby bench/reranker_bench.rb
+#
+# See docs/RERANKER_BENCHMARKS.md for measured results and how to interpret them —
+# including which platform they were measured on. Writes bench/reranker_bench_results.csv.
+#
+# Flags:
+#   --quick   Small matrix + few iterations, for a fast sanity check of the script itself.
+
+require "bundler/setup"
+require "candle"
+require "benchmark"
+
+QUICK = ARGV.include?("--quick")
+
+MODEL_ID = "cross-encoder/ms-marco-MiniLM-L6-v2"
+REVISION = "233902d25c440f23af6f7d6e94d2946bac0bee0a"
+
+CANDIDATE_COUNTS = QUICK ? [10, 100] : [10, 25, 50, 100]
+MAX_LENGTHS      = QUICK ? [128, 512] : [128, 256, 512]
+BATCH_SIZES      = QUICK ? [16] : [8, 16, 32]
+WARM_ITERATIONS  = QUICK ? 5 : 15
+CONCURRENCY_LEVELS = [1, 4]
+CONCURRENT_CALLS_PER_THREAD = QUICK ? 3 : 6
+MEMORY_GROWTH_ITERATIONS = QUICK ? 10 : 50
+
+QUERY = "What are the benefits of regular exercise for cardiovascular health?"
+
+# Synthetic documents of varying, realistic length so max_length truncation and batching
+# both get meaningfully exercised. Deterministic (no RNG) so runs are comparable.
+TOPICS = [
+  "Regular exercise improves cardiovascular health and reduces the risk of heart disease over time.",
+  "The stock market fluctuates based on interest rates, earnings reports, and investor sentiment.",
+  "Photosynthesis converts sunlight, water, and carbon dioxide into glucose and oxygen in plants.",
+  "The Roman Empire spanned three continents at its height and lasted for over a thousand years.",
+  "Deep learning models require large datasets and significant compute for effective training.",
+  "Coral reefs support roughly a quarter of all marine species despite covering a tiny fraction of the ocean floor."
+].freeze
+
+def build_documents(count)
+  (0...count).map do |i|
+    base = TOPICS[i % TOPICS.length]
+    # Vary length: every third document gets padded out with extra sentences, to mix
+    # short and long sequences within a micro-batch (the case the padding-direction fix
+    # in ext/candle/src/ruby/reranker.rs targets).
+    extra = (i % 3 == 0) ? " #{base} #{base}" : ""
+    "#{base}#{extra} (candidate #{i})"
+  end
+end
+
+def rss_kb(pid = Process.pid)
+  `ps -o rss= -p #{pid}`.to_i
+end
+
+def percentile(sorted_values, pct)
+  return sorted_values.first if sorted_values.length == 1
+  idx = ((pct / 100.0) * (sorted_values.length - 1)).round
+  sorted_values[idx]
+end
+
+results = []
+
+puts "=" * 100
+puts "Candle::Reranker benchmark — #{MODEL_ID} @ #{REVISION}"
+puts "Platform: #{RUBY_PLATFORM}, Ruby #{RUBY_VERSION}"
+puts "=" * 100
+
+MAX_LENGTHS.each do |max_length|
+  BATCH_SIZES.each do |batch_size|
+    print "\nLoading reranker (max_length=#{max_length}, batch_size=#{batch_size})... "
+    cold_load_s = Benchmark.realtime do
+      $reranker = Candle::Reranker.from_pretrained(
+        MODEL_ID, device: Candle::Device.cpu, revision: REVISION,
+        max_length: max_length, batch_size: batch_size
+      )
+    end
+    reranker = $reranker
+    puts "done in #{(cold_load_s * 1000).round(1)}ms"
+
+    CANDIDATE_COUNTS.each do |count|
+      documents = build_documents(count)
+
+      # Warm up (exclude from timing — page cache, tokenizer JIT-ish paths, etc.)
+      2.times { reranker.rerank(QUERY, documents) }
+
+      latencies_ms = []
+      rss_before = rss_kb
+      WARM_ITERATIONS.times do
+        t = Benchmark.realtime { reranker.rerank(QUERY, documents) }
+        latencies_ms << (t * 1000)
+      end
+      rss_after = rss_kb
+
+      sorted = latencies_ms.sort
+      p50 = percentile(sorted, 50)
+      p95 = percentile(sorted, 95)
+      throughput = count / (p50 / 1000.0)
+
+      row = {
+        max_length: max_length,
+        batch_size: batch_size,
+        candidate_count: count,
+        cold_load_ms: (cold_load_s * 1000).round(1),
+        warm_p50_ms: p50.round(2),
+        warm_p95_ms: p95.round(2),
+        throughput_docs_per_sec: throughput.round(1),
+        rss_before_kb: rss_before,
+        rss_after_kb: rss_after,
+        rss_growth_kb: rss_after - rss_before,
+        concurrency_1_p50_ms: nil,
+        concurrency_4_p50_ms: nil
+      }
+
+      # Single- vs multi-threaded caller latency, only for the "realistic" combo (largest
+      # candidate count) to keep total runtime bounded — this is the shape that matters for
+      # Benchmate's concurrent-Puma-threads use case.
+      if count == CANDIDATE_COUNTS.max
+        CONCURRENCY_LEVELS.each do |n_threads|
+          per_thread_latencies = Array.new(n_threads) { [] }
+          Benchmark.realtime do
+            threads = Array.new(n_threads) do |ti|
+              Thread.new do
+                CONCURRENT_CALLS_PER_THREAD.times do
+                  t = Benchmark.realtime { reranker.rerank(QUERY, documents) }
+                  per_thread_latencies[ti] << (t * 1000)
+                end
+              end
+            end
+            threads.each(&:join)
+          end
+          all_latencies = per_thread_latencies.flatten.sort
+          row[:"concurrency_#{n_threads}_p50_ms"] = percentile(all_latencies, 50).round(2)
+        end
+      end
+
+      results << row
+      printf(
+        "  count=%3d  p50=%7.2fms  p95=%7.2fms  throughput=%7.1f docs/s  rss_growth=%+dkB\n",
+        count, row[:warm_p50_ms], row[:warm_p95_ms], row[:throughput_docs_per_sec], row[:rss_growth_kb]
+      )
+    end
+  end
+end
+
+# Repeated-call memory growth: hammer one representative config many times, check RSS
+# before vs after to catch leaks (mmap'd weights should be shared/reused, not re-allocated).
+puts "\n" + "=" * 100
+puts "Repeated-call memory growth (#{MEMORY_GROWTH_ITERATIONS} calls, max_length=512, batch_size=16, 100 docs)"
+reranker = Candle::Reranker.from_pretrained(
+  MODEL_ID, device: Candle::Device.cpu, revision: REVISION, max_length: 512, batch_size: 16
+)
+documents = build_documents(100)
+reranker.rerank(QUERY, documents) # warm up
+GC.start
+rss_start = rss_kb
+MEMORY_GROWTH_ITERATIONS.times { reranker.rerank(QUERY, documents) }
+GC.start
+rss_end = rss_kb
+puts "RSS before: #{rss_start}kB, RSS after: #{rss_end}kB, growth: #{rss_end - rss_start}kB " \
+     "(#{((rss_end - rss_start).to_f / MEMORY_GROWTH_ITERATIONS).round(2)}kB/call)"
+
+csv_path = File.join(__dir__, "reranker_bench_results.csv")
+File.open(csv_path, "w") do |f|
+  f.puts results.first.keys.join(",")
+  results.each { |r| f.puts r.values.map { |v| v.nil? ? "" : v }.join(",") }
+end
+puts "\nWrote #{csv_path}"
+puts "\nPeak RSS observed: #{results.map { |r| r[:rss_after_kb] }.max}kB " \
+     "(~#{(results.map { |r| r[:rss_after_kb] }.max / 1024.0 / 1024.0).round(3)} GiB)"

bench/reranker_bench_results.csv

@@ -0,0 +1,37 @@
+max_length,batch_size,candidate_count,cold_load_ms,warm_p50_ms,warm_p95_ms,throughput_docs_per_sec,rss_before_kb,rss_after_kb,rss_growth_kb,concurrency_1_p50_ms,concurrency_4_p50_ms
+128,8,10,64.6,151.45,199.83,66.0,232176,239904,7728,,
+128,8,25,64.6,312.8,413.49,79.9,240032,245200,5168,,
+128,8,50,64.6,622.05,780.42,80.4,245424,246784,1360,,
+128,8,100,64.6,1149.64,1478.54,87.0,201344,198416,-2928,1202.29,2009.65
+128,16,10,35.9,101.07,119.29,98.9,410176,410272,96,,
+128,16,25,35.9,264.99,279.29,94.3,434576,435360,784,,
+128,16,50,35.9,582.22,640.32,85.9,435552,436208,656,,
+128,16,100,35.9,1067.88,1223.12,93.6,436400,437344,944,1190.51,1842.18
+128,32,10,36.3,111.64,137.0,89.6,603104,603232,128,,
+128,32,25,36.3,248.47,288.09,100.6,647584,647904,320,,
+128,32,50,36.3,510.26,596.26,98.0,722656,722912,256,,
+128,32,100,36.3,935.82,1081.2,106.9,722928,723216,288,866.71,1526.36
+256,8,10,31.9,107.48,118.29,93.0,742400,742400,0,,
+256,8,25,31.9,237.38,264.1,105.3,742400,742448,48,,
+256,8,50,31.9,512.66,540.85,97.5,742464,742576,112,,
+256,8,100,31.9,927.52,984.48,107.8,742592,743072,480,926.8,1989.26
+256,16,10,36.7,89.15,95.86,112.2,732416,732512,96,,
+256,16,25,36.7,219.65,227.16,113.8,732528,732576,48,,
+256,16,50,36.7,436.81,450.92,114.5,732624,635056,-97568,,
+256,16,100,36.7,870.73,891.18,114.8,635152,516640,-118512,1080.17,1542.28
+256,32,10,34.3,100.46,120.23,99.5,648304,648304,0,,
+256,32,25,34.3,219.98,250.08,113.6,692336,692320,-16,,
+256,32,50,34.3,436.74,480.98,114.5,770080,715504,-54576,,
+256,32,100,34.3,824.73,848.9,121.3,715536,655856,-59680,827.95,1189.96
+512,8,10,28.5,94.87,98.69,105.4,801536,801536,0,,
+512,8,25,28.5,228.59,230.75,109.4,801536,801488,-48,,
+512,8,50,28.5,461.29,518.59,108.4,801488,657344,-144144,,
+512,8,100,28.5,904.93,944.05,110.5,657360,552880,-104480,915.26,1390.22
+512,16,10,33.8,87.74,88.97,114.0,561584,561600,16,,
+512,16,25,33.8,216.48,218.95,115.5,574864,574816,-48,,
+512,16,50,33.8,495.13,556.36,101.0,574816,489072,-85744,,
+512,16,100,33.8,1163.27,1290.06,86.0,489072,489072,0,1127.09,1919.75
+512,32,10,34.9,132.07,137.91,75.7,666368,666368,0,,
+512,32,25,34.9,246.78,264.73,101.3,710400,710352,-48,,
+512,32,50,34.9,545.18,570.46,91.7,784816,784784,-32,,
+512,32,100,34.9,1061.55,1157.3,94.2,784912,706224,-78688,1173.48,1767.18

bench/reranker_memory_growth.rb

@@ -0,0 +1,44 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+# Isolated repeated-call memory growth check for Candle::Reranker — run standalone (own
+# process, nothing else loaded first) so RSS deltas reflect only this reranker's behavior,
+# not GC cleanup of unrelated objects from a prior sweep. See docs/RERANKER_BENCHMARKS.md.
+#
+#   bundle exec ruby bench/reranker_memory_growth.rb
+
+require "bundler/setup"
+require "candle"
+
+MODEL_ID = "cross-encoder/ms-marco-MiniLM-L6-v2"
+REVISION = "233902d25c440f23af6f7d6e94d2946bac0bee0a"
+ITERATIONS = 50
+
+def rss_kb
+  `ps -o rss= -p #{Process.pid}`.to_i
+end
+
+QUERY = "What are the benefits of regular exercise for cardiovascular health?"
+documents = (0...100).map { |i| "Regular exercise improves cardiovascular health and reduces disease risk. (candidate #{i})" }
+
+reranker = Candle::Reranker.from_pretrained(
+  MODEL_ID, device: Candle::Device.cpu, revision: REVISION, max_length: 512, batch_size: 16
+)
+reranker.rerank(QUERY, documents) # warm up
+GC.start(full_mark: true, immediate_sweep: true)
+rss_after_load = rss_kb
+
+samples = [rss_after_load]
+ITERATIONS.times do |i|
+  reranker.rerank(QUERY, documents)
+  samples << rss_kb if (i + 1) % 10 == 0
+end
+GC.start(full_mark: true, immediate_sweep: true)
+rss_final = rss_kb
+
+puts "RSS after load + warmup: #{rss_after_load}kB"
+puts "RSS every 10 calls: #{samples.join(' -> ')}kB"
+puts "RSS after #{ITERATIONS} calls (pre-GC): #{samples.last}kB"
+puts "RSS after #{ITERATIONS} calls (post-GC.start): #{rss_final}kB"
+puts "Net growth (post-GC): #{rss_final - rss_after_load}kB " \
+     "(#{((rss_final - rss_after_load).to_f / ITERATIONS).round(2)}kB/call)"