From c947a5fca0fa314f76500d9a5c89cfdd3f8292cb Mon Sep 17 00:00:00 2001 From: Jiannan Li Date: Fri, 11 Sep 2026 16:03:24 -0700 Subject: [PATCH] feat(tokenizer): add bounded L0 exact-match encode cache Add CachedTokenizer, a wrapper around Arc that memoizes successful encode() results keyed by the exact input string. Hits return an owned clone of the full Encoding, so token IDs, token strings, offsets and masks are preserved. The cache is bounded by entry count and estimated retained bytes with LRU eviction (lru crate, default features disabled). Results larger than max_entry_bytes are returned but not stored. A single parking_lot::Mutex guards the LRU; tokenization on a miss and the Encoding clone on a hit run outside the lock. encode_batch(), decode and metadata methods delegate to the wrapped tokenizer, and errors are never cached. The wrapped tokenizer must be deterministic with a fixed configuration. Hit, miss, eviction and oversized counters are exported under vllm_tokenizer_cache_*. The entries and bytes gauges hold the total occupancy of every live cache: each instance applies its change to a process-wide aggregate and publishes the totals while holding its own lock, and clear() and Drop subtract its share. Tests use the mock tokenizer and a checked-in byte-level BPE fixture (tests/fixtures/tokenizer/byte_level_bpe.json); gauge semantics are covered with a capturing metrics recorder. A criterion benchmark compares hit, miss, mixed and concurrent workloads against uncached encoding. Construction is library-level only; request-path integration and CLI flags are left for a follow-up. Refs #269, #244 --- Cargo.lock | 7 + Cargo.toml | 6 + benches/tokenizer_cache_benchmark.rs | 304 ++++++ src/metrics.rs | 51 + src/tokenizer/README.md | 76 ++ src/tokenizer/cache.rs | 956 +++++++++++++++++++ src/tokenizer/mod.rs | 2 + tests/fixtures/tokenizer/byte_level_bpe.json | 391 ++++++++ tests/tokenizer_cache_integration.rs | 334 +++++++ tests/tokenizer_cache_metrics.rs | 284 ++++++ 10 files changed, 2411 insertions(+) create mode 100644 benches/tokenizer_cache_benchmark.rs create mode 100644 src/tokenizer/cache.rs create mode 100644 tests/fixtures/tokenizer/byte_level_bpe.json create mode 100644 tests/tokenizer_cache_integration.rs create mode 100644 tests/tokenizer_cache_metrics.rs diff --git a/Cargo.lock b/Cargo.lock index 6b9ecec5..68fe6d2b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2377,6 +2377,12 @@ dependencies = [ "value-bag", ] +[[package]] +name = "lru" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25" + [[package]] name = "lru-slab" version = "0.1.2" @@ -4825,6 +4831,7 @@ dependencies = [ "k8s-openapi", "kube", "lazy_static", + "lru", "metrics", "metrics-exporter-prometheus", "minijinja", diff --git a/Cargo.toml b/Cargo.toml index 9489359c..b4061589 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,6 +79,7 @@ backoff = { version = "0.4", features = ["tokio"] } strum = { version = "0.26", features = ["derive"] } once_cell = "1.21.3" zmq = "0.10.0" +lru = { version = "0.18", default-features = false } rmp-serde = "1.3" [dev-dependencies] @@ -100,6 +101,11 @@ name = "tokenizer_benchmark" harness = false path = "benches/tokenizer_benchmark.rs" +[[bench]] +name = "tokenizer_cache_benchmark" +harness = false +path = "benches/tokenizer_cache_benchmark.rs" + [[bench]] name = "otel_disabled_path" harness = false diff --git a/benches/tokenizer_cache_benchmark.rs b/benches/tokenizer_cache_benchmark.rs new file mode 100644 index 00000000..a77b1051 --- /dev/null +++ b/benches/tokenizer_cache_benchmark.rs @@ -0,0 +1,304 @@ +//! Benchmarks for the exact-match tokenizer cache (`CachedTokenizer`). +//! +//! Compares cached `encode()` against the uncached HuggingFace tokenizer for +//! pure hits, pure misses, mixed hit ratios and concurrent hits, and prints +//! the estimated cache memory per prompt size. +//! +//! Run with: `cargo bench --bench tokenizer_cache_benchmark` +//! +//! The tokenizer is TinyLlama/TinyLlama-1.1B-Chat-v1.0 `tokenizer.json`, +//! fetched into `.tokenizer_cache/` by `common::ensure_tokenizer_cached` on +//! first use. Reported numbers were taken with the file whose SHA-256 is +//! `bcd04f0eadf90287bd26e1a183ac487d8a141b09b06aecb7725bbdd343640f2e`. +//! +//! Timing boundaries: every `encode` group drops the returned `Encoding` +//! outside the timed region (`iter_with_large_drop` / `iter_batched_ref`), +//! and the miss case also builds and drops its fresh cache outside it, so +//! the numbers are per-call costs, not cache lifecycle costs. + +use criterion::{ + black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput, +}; +use std::path::PathBuf; +use std::sync::{Arc, OnceLock}; +use std::thread; +use std::time::Instant; +use vllm_router_rs::tokenizer::{ + cache::{estimate_entry_bytes, CachedTokenizer, TokenizerCacheConfig}, + huggingface::HuggingFaceTokenizer, + traits::*, +}; + +#[path = "../tests/common/mod.rs"] +mod common; +use common::ensure_tokenizer_cached; + +static TOKENIZER_PATH: OnceLock = OnceLock::new(); + +fn tokenizer_path() -> &'static PathBuf { + TOKENIZER_PATH.get_or_init(ensure_tokenizer_cached) +} + +fn load_tokenizer() -> Arc { + Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ) +} + +const SHORT_PROMPT: &str = "What is the capital of France?"; +const MEDIUM_PROMPT: &str = "Write a detailed explanation of quantum computing, including its principles, current applications, and future potential. Be sure to cover both the theoretical foundations and practical implementations."; +const LONG_PROMPT: &str = "You are an expert software engineer. Review the following code and provide detailed feedback on performance optimizations, potential bugs, and architectural improvements. Consider scalability, maintainability, and best practices. The code implements a distributed caching system with the following requirements: 1) High availability across multiple regions, 2) Sub-millisecond latency for cache hits, 3) Automatic failover and recovery, 4) Support for both LRU and LFU eviction policies, 5) Real-time monitoring and alerting. Please analyze each component thoroughly and suggest concrete improvements with code examples where appropriate."; + +fn generate_system_prompt(size: usize) -> String { + let domains = [ + "mathematics", + "physics", + "chemistry", + "biology", + "computer science", + "engineering", + "medicine", + "law", + "economics", + "philosophy", + ]; + let mut prompt = String::from("You are a helpful AI assistant with expertise in "); + let mut i = 0; + while prompt.len() < size { + prompt.push_str(domains[i % domains.len()]); + prompt.push_str(", "); + i += 1; + } + prompt +} + +fn prompts() -> Vec<(&'static str, String)> { + vec![ + ("short_30B", SHORT_PROMPT.to_string()), + ("medium_230B", MEDIUM_PROMPT.to_string()), + ("long_670B", LONG_PROMPT.to_string()), + ("system_4KB", generate_system_prompt(4000)), + ("system_16KB", generate_system_prompt(16000)), + ] +} + +fn roomy_config() -> TokenizerCacheConfig { + TokenizerCacheConfig { + max_entries: 4096, + max_bytes: 1 << 30, + max_entry_bytes: 1 << 30, + } +} + +/// Print the estimated retained bytes for each prompt size once. +fn report_memory(inner: &Arc) { + println!("\nEstimated cache memory per entry (TinyLlama tokenizer):"); + println!( + "{:<14} {:>10} {:>8} {:>14} {:>12}", + "prompt", "input_B", "tokens", "estimate_B", "B/token" + ); + for (name, prompt) in prompts() { + let encoding = inner.encode(&prompt).unwrap(); + let tokens = encoding.token_ids().len(); + let estimate = estimate_entry_bytes(&prompt, &encoding); + println!( + "{:<14} {:>10} {:>8} {:>14} {:>12.1}", + name, + prompt.len(), + tokens, + estimate, + estimate as f64 / tokens.max(1) as f64 + ); + } + println!(); +} + +/// Uncached encode, cache hit, the `Encoding` clone a hit pays for, and a +/// cache miss into an empty cache, per prompt size. Returned encodings and +/// the per-miss caches are dropped outside the timed region. +fn bench_encode_paths(c: &mut Criterion) { + let inner = load_tokenizer(); + report_memory(&inner); + let cached = CachedTokenizer::new(inner.clone(), roomy_config()).unwrap(); + + let mut group = c.benchmark_group("tokenizer_cache/encode"); + for (name, prompt) in prompts() { + let encoding = inner.encode(&prompt).unwrap(); + cached.encode(&prompt).unwrap(); + group.throughput(Throughput::Bytes(prompt.len() as u64)); + + group.bench_with_input(BenchmarkId::new("uncached", name), &prompt, |b, p| { + b.iter_with_large_drop(|| inner.encode(black_box(p)).unwrap()) + }); + group.bench_with_input(BenchmarkId::new("hit", name), &prompt, |b, p| { + b.iter_with_large_drop(|| cached.encode(black_box(p)).unwrap()) + }); + group.bench_with_input(BenchmarkId::new("clone_only", name), &encoding, |b, e| { + b.iter_with_large_drop(|| black_box(e).clone()) + }); + group.bench_with_input(BenchmarkId::new("miss", name), &prompt, |b, p| { + b.iter_batched_ref( + || { + CachedTokenizer::new( + inner.clone(), + TokenizerCacheConfig { + max_entries: 16, + ..roomy_config() + }, + ) + .unwrap() + }, + |cache| cache.encode(black_box(p)).unwrap(), + BatchSize::SmallInput, + ) + }); + } + group.finish(); +} + +/// A hot set of 64 medium prompts served at a fixed hit ratio, with cold +/// prompts unique per call, against the same sequence uncached. +fn bench_mixed_workload(c: &mut Criterion) { + let inner = load_tokenizer(); + let hot: Vec = (0..64) + .map(|i| format!("{MEDIUM_PROMPT} (variant {i})")) + .collect(); + + let mut group = c.benchmark_group("tokenizer_cache/mixed"); + group.throughput(Throughput::Elements(1)); + for hit_pct in [50u64, 90, 99] { + let label = format!("hit{hit_pct}pct"); + + group.bench_with_input( + BenchmarkId::new("cached", &label), + &hit_pct, + |b, &hit_pct| { + b.iter_custom(|iters| { + let cache = CachedTokenizer::new( + inner.clone(), + TokenizerCacheConfig { + max_entries: 1024, + max_bytes: 64 << 20, + max_entry_bytes: 1 << 20, + }, + ) + .unwrap(); + for prompt in &hot { + cache.encode(prompt).unwrap(); + } + let mut cold = 0u64; + let start = Instant::now(); + for i in 0..iters { + if i % 100 < hit_pct { + black_box(cache.encode(&hot[i as usize % hot.len()]).unwrap()); + } else { + cold += 1; + let prompt = format!("{MEDIUM_PROMPT} (cold {cold})"); + black_box(cache.encode(&prompt).unwrap()); + } + } + start.elapsed() + }) + }, + ); + + group.bench_with_input( + BenchmarkId::new("uncached", &label), + &hit_pct, + |b, &hit_pct| { + b.iter_custom(|iters| { + let mut cold = 0u64; + let start = Instant::now(); + for i in 0..iters { + if i % 100 < hit_pct { + black_box(inner.encode(&hot[i as usize % hot.len()]).unwrap()); + } else { + cold += 1; + let prompt = format!("{MEDIUM_PROMPT} (cold {cold})"); + black_box(inner.encode(&prompt).unwrap()); + } + } + start.elapsed() + }) + }, + ); + } + group.finish(); +} + +/// All-hit encodes spread across N threads sharing one cache. Reported time +/// is wall clock per encode across all threads, so flat or rising numbers +/// with more threads indicate lock contention. +fn bench_concurrent_hits(c: &mut Criterion) { + let inner = load_tokenizer(); + let cache = Arc::new(CachedTokenizer::new(inner.clone(), roomy_config()).unwrap()); + let hot: Arc> = Arc::new( + (0..16) + .map(|i| format!("{SHORT_PROMPT} (variant {i})")) + .collect(), + ); + for prompt in hot.iter() { + cache.encode(prompt).unwrap(); + } + + let mut group = c.benchmark_group("tokenizer_cache/concurrent_short"); + group.throughput(Throughput::Elements(1)); + for threads in [1usize, 2, 4, 8] { + group.bench_with_input(BenchmarkId::new("hit", threads), &threads, |b, &threads| { + b.iter_custom(|iters| { + let per_thread = iters.div_ceil(threads as u64); + let start = Instant::now(); + let handles: Vec<_> = (0..threads) + .map(|t| { + let cache = cache.clone(); + let hot = hot.clone(); + thread::spawn(move || { + for i in 0..per_thread as usize { + black_box(cache.encode(&hot[(i + t) % hot.len()]).unwrap()); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + start.elapsed() + }) + }); + group.bench_with_input( + BenchmarkId::new("uncached", threads), + &threads, + |b, &threads| { + b.iter_custom(|iters| { + let per_thread = iters.div_ceil(threads as u64); + let start = Instant::now(); + let handles: Vec<_> = (0..threads) + .map(|t| { + let inner = inner.clone(); + let hot = hot.clone(); + thread::spawn(move || { + for i in 0..per_thread as usize { + black_box(inner.encode(&hot[(i + t) % hot.len()]).unwrap()); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + start.elapsed() + }) + }, + ); + } + group.finish(); +} + +criterion_group!( + benches, + bench_encode_paths, + bench_mixed_workload, + bench_concurrent_hits +); +criterion_main!(benches); diff --git a/src/metrics.rs b/src/metrics.rs index ee13ed13..4a7a1039 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -251,6 +251,32 @@ pub fn init_metrics() { "vllm_tokenizer_factory_load_duration_seconds", "Time to load and initialize tokenizer" ); + + // Tokenizer encode cache metrics + describe_counter!( + "vllm_tokenizer_cache_hits_total", + "Total tokenizer encode calls served from the exact-match cache" + ); + describe_counter!( + "vllm_tokenizer_cache_misses_total", + "Total tokenizer encode calls that ran the underlying tokenizer" + ); + describe_counter!( + "vllm_tokenizer_cache_evictions_total", + "Total tokenizer cache entries evicted to satisfy the entry or byte budget" + ); + describe_counter!( + "vllm_tokenizer_cache_oversized_total", + "Total tokenizer encode results not cached because they exceeded the per-entry byte limit" + ); + describe_gauge!( + "vllm_tokenizer_cache_entries", + "Current number of entries in the tokenizer encode cache" + ); + describe_gauge!( + "vllm_tokenizer_cache_bytes", + "Estimated bytes retained by the tokenizer encode cache" + ); } pub fn start_prometheus(config: PrometheusConfig) { @@ -624,6 +650,31 @@ impl TokenizerMetrics { ) .set(size as f64); } + + // Encode cache metrics + pub fn record_cache_hit() { + counter!("vllm_tokenizer_cache_hits_total").increment(1); + } + + pub fn record_cache_miss() { + counter!("vllm_tokenizer_cache_misses_total").increment(1); + } + + pub fn record_cache_evictions(count: u64) { + counter!("vllm_tokenizer_cache_evictions_total").increment(count); + } + + pub fn record_cache_oversized() { + counter!("vllm_tokenizer_cache_oversized_total").increment(1); + } + + pub fn set_cache_entries(entries: usize) { + gauge!("vllm_tokenizer_cache_entries").set(entries as f64); + } + + pub fn set_cache_bytes(bytes: usize) { + gauge!("vllm_tokenizer_cache_bytes").set(bytes as f64); + } } #[cfg(test)] diff --git a/src/tokenizer/README.md b/src/tokenizer/README.md index d766f8a2..d468372a 100644 --- a/src/tokenizer/README.md +++ b/src/tokenizer/README.md @@ -299,6 +299,7 @@ impl Tokenizer { - Types: `Sequence`, `StopSequenceConfig`, `DecodeStream`, `Encoding`, `TokenizerType` - Chat template: `ChatMessage` - Tokenizer implementations: `HuggingFaceTokenizer`, `TiktokenTokenizer` +- Encode cache: `CachedTokenizer`, `TokenizerCacheConfig`, `TokenizerCacheStats` ### 3.2 traits.rs (Trait Definitions) @@ -775,6 +776,51 @@ impl HuggingFaceTokenizer { - Runtime template modification - Special token handling +### 3.12 cache.rs (Exact-Match Encode Cache) + +**Location**: `src/tokenizer/cache.rs` + +**Purpose:** Bounded L0 cache that memoizes successful `encode()` results by exact input string. Wraps any `Arc` and delegates everything else. + +**Public API:** + +```rust +pub struct TokenizerCacheConfig { + pub max_entries: usize, // default 10_000 + pub max_bytes: usize, // default 64 MiB, estimated retained bytes + pub max_entry_bytes: usize, // default 1 MiB, larger results are not stored +} + +pub struct CachedTokenizer { /* implements Encoder + Decoder + Tokenizer */ } + +impl CachedTokenizer { + pub fn new(inner: Arc, config: TokenizerCacheConfig) -> Result + pub fn stats(&self) -> TokenizerCacheStats // hits, misses, evictions, oversized, entries, bytes + pub fn clear(&self) +} + +pub fn estimate_entry_bytes(input: &str, encoding: &Encoding) -> usize +``` + +**Behavior:** +- Hits return an owned clone of the stored `Encoding` with all backend metadata +- Only successful `encode()` results are stored; errors pass through +- `encode_batch()` neither reads from nor populates the cache +- Each `CachedTokenizer` owns its entries; instances never share +- LRU eviction when either the entry or byte budget is exceeded (`LruCache::sparse`, so the map grows on demand) +- Results larger than `max_entry_bytes` are returned but not stored +- Defaults (10,000 entries, 64 MiB, 1 MiB per entry) are initial values, not tuned against production traffic + +**Preconditions:** the wrapped tokenizer must be deterministic for the same input, and every setting that affects encoding (normalization, truncation, added tokens, BPE dropout or other sampling) must stay fixed for the cache's lifetime. A fixed configuration does not by itself disable stochastic tokenization. + +**Byte estimation:** heap bytes kept alive by an entry, using lengths rather than capacities: the input string, the token ID vector for `Sp`/`Tiktoken`, and for `Hf` the per-token `u32` vectors, word index, offsets, token string headers and token text (overflowing encodings recursively), plus a fixed per-entry overhead (`ENTRY_OVERHEAD_BYTES`). Not counted: allocator padding, hash-map load factor, unused `Vec` capacity, clones held by callers, and evicted encodings kept alive by an in-flight `Arc` from a concurrent hit. The estimate is a lower bound, not a resident-memory bound. + +**Concurrency:** one `parking_lot::Mutex` around the LRU. Lookups clone an `Arc` under the lock and clone the `Encoding` outside it; tokenization on a miss also runs outside the lock. Concurrent misses on the same input both tokenize and the later insert replaces the earlier one. + +**Metrics:** counters sum across instances. The `entries` and `bytes` gauges hold the total occupancy of every live cache in the process: each instance applies its change to a process-wide aggregate and publishes the new totals while still holding its own lock (lock order: instance state, then aggregate), and `clear()` and `Drop` subtract the instance's share. After any `encode`, `clear` or drop the gauges equal the sum of live instances' `stats()`. + +**Not wired yet:** construction is library-level only. Request-path integration and CLI flags are follow-up work. + ## 4. Traits & Contracts ### Core Trait Hierarchy @@ -895,6 +941,7 @@ The `Encoding` enum must: - `tiktoken.rs`: 7 tests - Model detection, encode/decode roundtrip - `chat_template.rs`: 3 tests - Template rendering, loading - `tests.rs`: 9 tests - Cross-module integration +- `cache.rs`: 19 tests - Hit/miss, eviction by count and bytes, oversized bypass, batch bypass, concurrency, byte estimates **Integration Tests (10 tests in tokenizer_integration.rs):** - HuggingFace tokenizer hash verification @@ -907,6 +954,19 @@ The `Encoding` enum must: - Special token handling - Thread safety validation +**Cache Integration Tests (8 tests + 1 ignored in tokenizer_cache_integration.rs):** +- Use `tests/fixtures/tokenizer/byte_level_bpe.json`, a checked-in byte-level BPE tokenizer (256 byte symbols, 34 merges, ``/``/``); no network needed +- Fixture pinned by vocab size, special tokens and token IDs +- Cached vs uncached HuggingFace encodings compared field by field, including overflowing encodings from truncation +- Byte accounting against `estimate_entry_bytes` +- Byte-budget eviction and oversized bypass with real encodings +- Decode and streaming through the `Tokenizer` wrapper +- Concurrent access with eviction +- `#[ignore]` TinyLlama repeat of the comparison plus the hash fixtures (`cargo test --test tokenizer_cache_integration -- --ignored`) + +**Cache Metrics Tests (4 tests in tokenizer_cache_metrics.rs):** +- Capturing `metrics::Recorder` checks the occupancy gauges sum over live instances, follow eviction, `clear()` and `Drop`, and stay consistent under concurrent inserts and clears, including a delayed publication that must not overwrite a later clear + ### Benchmark Suite (tokenizer_benchmark.rs) **Performance Benchmarks (12 benchmark groups):** @@ -928,6 +988,16 @@ The `Encoding` enum must: - Medium: 201 chars (Quantum computing explanation) - Long: 638 chars (Software engineering review) +### Cache Benchmark Suite (tokenizer_cache_benchmark.rs) + +Run with `cargo bench --bench tokenizer_cache_benchmark`. + +1. **encode**: uncached encode, cache hit, `Encoding` clone alone, and miss into an empty cache, per prompt size (30 B to 16 KB) +2. **mixed**: 64 hot prompts at 50/90/99% hit ratio with unique cold prompts, cached vs uncached +3. **concurrent_short**: all-hit encodes across 1/2/4/8 threads sharing one cache, cached vs uncached + +Prints estimated retained bytes per entry for each prompt size before running. + ## 8. Operational Concerns ### Configuration @@ -953,6 +1023,12 @@ The `Encoding` enum must: - `vllm_tokenizer_factory_load_duration_seconds` - `vllm_tokenizer_stop_sequence_detected` - `vllm_tokenizer_stream_incomplete_utf8_total` +- `vllm_tokenizer_cache_hits_total` +- `vllm_tokenizer_cache_misses_total` +- `vllm_tokenizer_cache_evictions_total` +- `vllm_tokenizer_cache_oversized_total` +- `vllm_tokenizer_cache_entries` (gauge, total over live caches) +- `vllm_tokenizer_cache_bytes` (gauge, estimated, total over live caches) **Labels:** - `tokenizer_type`: huggingface, tiktoken, mock diff --git a/src/tokenizer/cache.rs b/src/tokenizer/cache.rs new file mode 100644 index 00000000..e74af3e2 --- /dev/null +++ b/src/tokenizer/cache.rs @@ -0,0 +1,956 @@ +//! Bounded exact-match (L0) cache for tokenizer encode results. +//! +//! [`CachedTokenizer`] wraps any `Arc` and memoizes successful +//! [`Encoder::encode`] results keyed by the exact input string. Every other +//! trait method is delegated to the wrapped tokenizer unchanged. +//! +//! # Semantics +//! +//! - A cache belongs to exactly one tokenizer instance. Two `CachedTokenizer`s +//! never share entries, even when they wrap equivalent tokenizers. +//! - Hits return an owned clone of the stored [`Encoding`], including all +//! metadata (token strings, offsets, masks) carried by the backend. +//! - Only successful `encode()` results are stored. Errors are returned to the +//! caller and leave the cache untouched. +//! - `encode_batch()` neither reads from nor populates the cache. +//! - Tokenization runs outside the cache lock. Two concurrent misses on the +//! same input both tokenize; the later insert replaces the earlier one with +//! an identical value. +//! +//! # Preconditions +//! +//! Caching is only sound when the wrapped tokenizer is deterministic: the +//! same input must always produce the same encoding, and every setting that +//! affects encoding (normalization, truncation, added tokens, BPE dropout or +//! other sampling) must stay fixed for the cache's lifetime. A fixed +//! configuration does not by itself disable stochastic tokenization; callers +//! must make sure it is off before wrapping. +//! +//! # Bounds and eviction +//! +//! The cache is bounded by an entry count ([`TokenizerCacheConfig::max_entries`]) +//! and an estimated retained-byte budget ([`TokenizerCacheConfig::max_bytes`]). +//! When either bound is exceeded, least-recently-used entries are evicted until +//! both hold. Results whose estimate exceeds +//! [`TokenizerCacheConfig::max_entry_bytes`] are returned to the caller but not +//! stored, so a handful of very long inputs cannot flush the working set. +//! +//! # Byte estimation +//! +//! Estimates count heap bytes that an entry keeps alive, using lengths rather +//! than capacities, plus a fixed per-entry bookkeeping overhead +//! ([`ENTRY_OVERHEAD_BYTES`]). Per entry: +//! +//! - the input string: `input.len()` +//! - `Encoding::Sp` and `Encoding::Tiktoken`: the `Vec` header plus +//! `4 bytes × token count` +//! - `Encoding::Hf`: the boxed struct plus, per token, the four `u32` vectors +//! (ids, type ids, special-token mask, attention mask), the `Option` word +//! index, the `(usize, usize)` offset pair, the `String` header, and the token +//! text bytes; overflowing encodings are counted recursively. +//! +//! The estimate is a lower bound on real memory use, not a bound on resident +//! memory: it ignores allocator padding, hash-map load factor, unused `Vec` +//! capacity, clones held by callers, and evicted encodings kept alive by an +//! in-flight `Arc` from a concurrent hit. +//! +//! # Metrics +//! +//! Hits, misses, evictions and oversized skips are counted under +//! `vllm_tokenizer_cache_*_total`. `vllm_tokenizer_cache_entries` and +//! `vllm_tokenizer_cache_bytes` are gauges holding the total occupancy of +//! every live `CachedTokenizer` in the process. Each instance applies its +//! change to a process-wide aggregate and publishes the new totals while +//! still holding its own lock, so after any `encode`, `clear` or drop the +//! gauges equal the sum of the live instances' [`stats`](CachedTokenizer::stats). +//! The series carry no labels; per-cache labels can be added when the cache +//! is wired into the request path. + +use super::traits::{ + Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait, +}; +use crate::metrics::TokenizerMetrics; +use anyhow::{bail, Result}; +use lru::LruCache; +use parking_lot::Mutex; +use std::mem::size_of; +use std::num::NonZeroUsize; +use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +/// Budgets for a [`CachedTokenizer`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct TokenizerCacheConfig { + /// Maximum number of entries retained. Must be at least 1. + pub max_entries: usize, + /// Maximum estimated bytes retained across all entries. Must be at least 1. + pub max_bytes: usize, + /// Encode results whose estimated size exceeds this are not cached. + /// Must be at least 1 and at most `max_bytes`. + pub max_entry_bytes: usize, +} + +impl Default for TokenizerCacheConfig { + fn default() -> Self { + Self { + max_entries: 10_000, + max_bytes: 64 * 1024 * 1024, + max_entry_bytes: 1024 * 1024, + } + } +} + +impl TokenizerCacheConfig { + /// Check that every budget is usable. + pub fn validate(&self) -> Result<()> { + if self.max_entries == 0 { + bail!("tokenizer cache max_entries must be at least 1"); + } + if self.max_bytes == 0 { + bail!("tokenizer cache max_bytes must be at least 1"); + } + if self.max_entry_bytes == 0 { + bail!("tokenizer cache max_entry_bytes must be at least 1"); + } + if self.max_entry_bytes > self.max_bytes { + bail!( + "tokenizer cache max_entry_bytes ({}) must not exceed max_bytes ({})", + self.max_entry_bytes, + self.max_bytes + ); + } + Ok(()) + } +} + +/// Point-in-time counters and occupancy of a [`CachedTokenizer`]. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TokenizerCacheStats { + /// `encode()` calls served from the cache. + pub hits: u64, + /// `encode()` calls that ran the wrapped tokenizer. + pub misses: u64, + /// Entries removed to satisfy the entry or byte budget. + pub evictions: u64, + /// Successful encodes not stored because they exceeded `max_entry_bytes`. + pub oversized: u64, + /// Entries currently retained. + pub entries: usize, + /// Estimated bytes currently retained. + pub bytes: usize, +} + +struct CacheEntry { + encoding: Arc, + bytes: usize, +} + +struct CacheState { + lru: LruCache, + /// Sum of `CacheEntry::bytes` over all entries in `lru`. + bytes: usize, +} + +/// Approximate fixed cost of one cache entry: the LRU node (key, value and two +/// list links) and its hash-map slot (key reference and node pointer). +pub const ENTRY_OVERHEAD_BYTES: usize = + size_of::() + size_of::() + 4 * size_of::(); + +/// Occupancy summed over every live cache in the process. Guarded by a lock +/// so gauge publications are totally ordered and always carry the current +/// total. +struct Occupancy { + entries: usize, + bytes: usize, +} + +static OCCUPANCY: Mutex = Mutex::new(Occupancy { + entries: 0, + bytes: 0, +}); + +/// Apply one instance's state change to the process-wide totals and publish +/// them. Callers hold their own state lock while calling this, so an +/// instance's contribution to the gauges always matches its state. Lock order +/// is instance state, then `OCCUPANCY`. +fn publish_occupancy_delta(entries_delta: isize, bytes_delta: isize) { + if entries_delta == 0 && bytes_delta == 0 { + return; + } + let mut total = OCCUPANCY.lock(); + total.entries = total.entries.saturating_add_signed(entries_delta); + total.bytes = total.bytes.saturating_add_signed(bytes_delta); + TokenizerMetrics::set_cache_entries(total.entries); + TokenizerMetrics::set_cache_bytes(total.bytes); +} + +/// Tokenizer wrapper that memoizes `encode()` results by exact input. +/// +/// See the [module documentation](self) for semantics and bounds. +pub struct CachedTokenizer { + inner: Arc, + config: TokenizerCacheConfig, + state: Mutex, + hits: AtomicU64, + misses: AtomicU64, + evictions: AtomicU64, + oversized: AtomicU64, +} + +impl CachedTokenizer { + /// Wrap `inner` with a cache sized by `config`. + /// + /// Returns an error if `config` fails [`TokenizerCacheConfig::validate`]. + pub fn new(inner: Arc, config: TokenizerCacheConfig) -> Result { + config.validate()?; + let capacity = + NonZeroUsize::new(config.max_entries).expect("validate() guarantees max_entries >= 1"); + Ok(Self { + inner, + config, + state: Mutex::new(CacheState { + // `sparse` grows the map on demand instead of preallocating + // `max_entries` slots, which matters when the byte budget is + // the binding limit. + lru: LruCache::sparse(capacity), + bytes: 0, + }), + hits: AtomicU64::new(0), + misses: AtomicU64::new(0), + evictions: AtomicU64::new(0), + oversized: AtomicU64::new(0), + }) + } + + /// The budgets this cache was built with. + pub fn config(&self) -> &TokenizerCacheConfig { + &self.config + } + + /// The wrapped tokenizer. + pub fn inner(&self) -> &Arc { + &self.inner + } + + /// Number of entries currently retained. + pub fn len(&self) -> usize { + self.state.lock().lru.len() + } + + /// Whether the cache holds no entries. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + /// Snapshot of counters and occupancy. + pub fn stats(&self) -> TokenizerCacheStats { + let (entries, bytes) = { + let state = self.state.lock(); + (state.lru.len(), state.bytes) + }; + TokenizerCacheStats { + hits: self.hits.load(Ordering::Relaxed), + misses: self.misses.load(Ordering::Relaxed), + evictions: self.evictions.load(Ordering::Relaxed), + oversized: self.oversized.load(Ordering::Relaxed), + entries, + bytes, + } + } + + /// Drop every entry. Counters are kept. + pub fn clear(&self) { + let mut guard = self.state.lock(); + let state = &mut *guard; + let entries = state.lru.len(); + let bytes = state.bytes; + state.lru.clear(); + state.bytes = 0; + publish_occupancy_delta(-(entries as isize), -(bytes as isize)); + } + + fn lookup(&self, input: &str) -> Option> { + let mut state = self.state.lock(); + state + .lru + .get(input) + .map(|entry| Arc::clone(&entry.encoding)) + } + + fn insert(&self, input: &str, encoding: Arc, bytes: usize) { + let mut evicted: u64 = 0; + { + let mut guard = self.state.lock(); + let state = &mut *guard; + let entries_before = state.lru.len(); + let bytes_before = state.bytes; + + // Another thread may have inserted this input while we were + // tokenizing. `push` then returns the entry it replaced rather + // than an LRU victim, so account for it as a replacement. + let replaced_bytes = state.lru.peek(input).map(|entry| entry.bytes); + let displaced = state + .lru + .push(input.to_owned(), CacheEntry { encoding, bytes }); + match (replaced_bytes, displaced) { + (Some(previous), Some(_)) => state.bytes -= previous, + (None, Some((_, victim))) => { + state.bytes -= victim.bytes; + evicted += 1; + } + (_, None) => {} + } + state.bytes += bytes; + + // `validate()` bounds every stored entry by `max_bytes`, so this + // loop never reaches the entry just inserted (the MRU). + while state.bytes > self.config.max_bytes { + match state.lru.pop_lru() { + Some((_, victim)) => { + state.bytes -= victim.bytes; + evicted += 1; + } + None => break, + } + } + + publish_occupancy_delta( + state.lru.len() as isize - entries_before as isize, + state.bytes as isize - bytes_before as isize, + ); + } + + if evicted > 0 { + self.evictions.fetch_add(evicted, Ordering::Relaxed); + TokenizerMetrics::record_cache_evictions(evicted); + } + } +} + +impl Drop for CachedTokenizer { + fn drop(&mut self) { + let state = self.state.get_mut(); + publish_occupancy_delta(-(state.lru.len() as isize), -(state.bytes as isize)); + } +} + +impl Encoder for CachedTokenizer { + fn encode(&self, input: &str) -> Result { + if let Some(encoding) = self.lookup(input) { + self.hits.fetch_add(1, Ordering::Relaxed); + TokenizerMetrics::record_cache_hit(); + // Clone outside the lock so a large encoding does not stall + // other readers. + return Ok((*encoding).clone()); + } + self.misses.fetch_add(1, Ordering::Relaxed); + TokenizerMetrics::record_cache_miss(); + + let encoding = self.inner.encode(input)?; + let bytes = estimate_entry_bytes(input, &encoding); + if bytes > self.config.max_entry_bytes { + self.oversized.fetch_add(1, Ordering::Relaxed); + TokenizerMetrics::record_cache_oversized(); + return Ok(encoding); + } + + let shared = Arc::new(encoding); + let result = (*shared).clone(); + self.insert(input, shared, bytes); + Ok(result) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + self.inner.encode_batch(inputs) + } +} + +impl Decoder for CachedTokenizer { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { + self.inner.decode(token_ids, skip_special_tokens) + } +} + +impl TokenizerTrait for CachedTokenizer { + fn vocab_size(&self) -> usize { + self.inner.vocab_size() + } + + fn get_special_tokens(&self) -> &SpecialTokens { + self.inner.get_special_tokens() + } + + fn token_to_id(&self, token: &str) -> Option { + self.inner.token_to_id(token) + } + + fn id_to_token(&self, id: TokenIdType) -> Option { + self.inner.id_to_token(id) + } +} + +/// Estimated bytes a cache entry for `input` and its `encoding` retains. +/// +/// See the [module documentation](self) for what is counted. +pub fn estimate_entry_bytes(input: &str, encoding: &Encoding) -> usize { + ENTRY_OVERHEAD_BYTES + input.len() + estimate_encoding_bytes(encoding) +} + +/// Estimated bytes retained by `encoding` alone. +pub fn estimate_encoding_bytes(encoding: &Encoding) -> usize { + match encoding { + Encoding::Sp(ids) | Encoding::Tiktoken(ids) => { + size_of::>() + ids.len() * size_of::() + } + Encoding::Hf(inner) => { + size_of::() + estimate_hf_bytes(inner) + } + } +} + +fn estimate_hf_bytes(encoding: &tokenizers::tokenizer::Encoding) -> usize { + // ids, type_ids, special_tokens_mask, attention_mask, words, offsets and + // the token `String` headers are all one element per token. + const PER_TOKEN: usize = 4 * size_of::() + + size_of::>() + + size_of::<(usize, usize)>() + + size_of::(); + + let tokens = encoding.get_ids().len() * PER_TOKEN; + let token_text: usize = encoding.get_tokens().iter().map(String::len).sum(); + let overflowing: usize = encoding + .get_overflowing() + .iter() + .map(|overflow| size_of::() + estimate_hf_bytes(overflow)) + .sum(); + // `sequence_ranges` is empty for the single-sequence case and holds one + // `usize -> Range` pair per sequence otherwise. + let sequences = encoding.n_sequences(); + let sequence_ranges = if sequences > 1 { + sequences * (size_of::() + size_of::>()) + } else { + 0 + }; + + tokens + token_text + overflowing + sequence_ranges +} + +#[cfg(test)] +impl CachedTokenizer { + /// Recompute retained bytes from the entries, to check the running total. + fn recomputed_bytes(&self) -> usize { + self.state + .lock() + .lru + .iter() + .map(|(_, entry)| entry.bytes) + .sum() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tokenizer::mock::MockTokenizer; + use crate::tokenizer::Tokenizer; + use std::sync::atomic::AtomicUsize; + use std::thread; + + /// Mock tokenizer that counts calls and fails for inputs containing "fail". + struct CountingTokenizer { + inner: MockTokenizer, + encode_calls: AtomicUsize, + batch_calls: AtomicUsize, + } + + impl CountingTokenizer { + fn new() -> Arc { + Arc::new(Self { + inner: MockTokenizer::new(), + encode_calls: AtomicUsize::new(0), + batch_calls: AtomicUsize::new(0), + }) + } + + fn encode_calls(&self) -> usize { + self.encode_calls.load(Ordering::SeqCst) + } + + fn batch_calls(&self) -> usize { + self.batch_calls.load(Ordering::SeqCst) + } + } + + impl Encoder for CountingTokenizer { + fn encode(&self, input: &str) -> Result { + self.encode_calls.fetch_add(1, Ordering::SeqCst); + if input.contains("fail") { + bail!("simulated encode failure"); + } + self.inner.encode(input) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + self.batch_calls.fetch_add(1, Ordering::SeqCst); + self.inner.encode_batch(inputs) + } + } + + impl Decoder for CountingTokenizer { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { + self.inner.decode(token_ids, skip_special_tokens) + } + } + + impl TokenizerTrait for CountingTokenizer { + fn vocab_size(&self) -> usize { + self.inner.vocab_size() + } + + fn get_special_tokens(&self) -> &SpecialTokens { + self.inner.get_special_tokens() + } + + fn token_to_id(&self, token: &str) -> Option { + self.inner.token_to_id(token) + } + + fn id_to_token(&self, id: TokenIdType) -> Option { + self.inner.id_to_token(id) + } + } + + fn config(max_entries: usize, max_bytes: usize) -> TokenizerCacheConfig { + TokenizerCacheConfig { + max_entries, + max_bytes, + max_entry_bytes: max_bytes, + } + } + + fn cached(inner: &Arc, config: TokenizerCacheConfig) -> CachedTokenizer { + CachedTokenizer::new(inner.clone(), config).unwrap() + } + + /// Estimated entry size for `input` under the mock tokenizer. + fn mock_entry_bytes(input: &str) -> usize { + let encoding = MockTokenizer::new().encode(input).unwrap(); + estimate_entry_bytes(input, &encoding) + } + + #[test] + fn hit_returns_equal_encoding_without_calling_inner() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + let first = cache.encode("Hello world").unwrap(); + let second = cache.encode("Hello world").unwrap(); + + assert_eq!(first.token_ids(), &[1, 2]); + assert_eq!(second.token_ids(), first.token_ids()); + assert_eq!(second.get_hash(), first.get_hash()); + assert_eq!(inner.encode_calls(), 1); + + let stats = cache.stats(); + assert_eq!(stats.hits, 1); + assert_eq!(stats.misses, 1); + assert_eq!(stats.entries, 1); + assert_eq!(stats.bytes, mock_entry_bytes("Hello world")); + } + + #[test] + fn empty_input_is_cached() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + assert!(cache.encode("").unwrap().token_ids().is_empty()); + assert!(cache.encode("").unwrap().token_ids().is_empty()); + + assert_eq!(inner.encode_calls(), 1); + assert_eq!(cache.stats().entries, 1); + } + + #[test] + fn keys_are_exact_bytes() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + // Case, whitespace and Unicode normalization form all change the key. + for input in [ + "Hello", + "hello", + "Hello ", + "caf\u{e9}", + "cafe\u{301}", + "\u{1f600}", + ] { + cache.encode(input).unwrap(); + } + assert_eq!(inner.encode_calls(), 6); + assert_eq!(cache.stats().entries, 6); + + cache.encode("caf\u{e9}").unwrap(); + assert_eq!(inner.encode_calls(), 6); + assert_eq!(cache.stats().hits, 1); + } + + #[test] + fn special_tokens_round_trip_through_cache() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + let first = cache.encode(" Hello ").unwrap(); + let second = cache.encode(" Hello ").unwrap(); + assert_eq!(first.token_ids(), &[1000, 1, 999]); + assert_eq!(second.token_ids(), &[1000, 1, 999]); + assert_eq!(inner.encode_calls(), 1); + } + + #[test] + fn errors_are_not_cached() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + assert!(cache.encode("please fail").is_err()); + assert!(cache.encode("please fail").is_err()); + + assert_eq!(inner.encode_calls(), 2); + let stats = cache.stats(); + assert_eq!(stats.misses, 2); + assert_eq!(stats.hits, 0); + assert_eq!(stats.entries, 0); + assert_eq!(stats.bytes, 0); + } + + #[test] + fn caches_are_isolated_per_instance() { + let inner_a = CountingTokenizer::new(); + let inner_b = CountingTokenizer::new(); + let cache_a = cached(&inner_a, TokenizerCacheConfig::default()); + let cache_b = cached(&inner_b, TokenizerCacheConfig::default()); + + cache_a.encode("Hello").unwrap(); + cache_a.encode("Hello").unwrap(); + assert_eq!(inner_a.encode_calls(), 1); + assert_eq!(inner_b.encode_calls(), 0); + + cache_b.encode("Hello").unwrap(); + assert_eq!(inner_b.encode_calls(), 1); + assert_eq!(cache_a.stats().entries, 1); + assert_eq!(cache_b.stats().entries, 1); + assert_eq!(cache_b.stats().hits, 0); + } + + #[test] + fn evicts_least_recently_used_by_entry_count() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, config(2, usize::MAX)); + + cache.encode("Hello").unwrap(); + cache.encode("world").unwrap(); + cache.encode("Hello").unwrap(); // touch: "world" is now the LRU entry + cache.encode("test").unwrap(); // evicts "world" + assert_eq!(inner.encode_calls(), 3); + + let stats = cache.stats(); + assert_eq!(stats.entries, 2); + assert_eq!(stats.evictions, 1); + assert_eq!( + stats.bytes, + mock_entry_bytes("Hello") + mock_entry_bytes("test") + ); + + cache.encode("Hello").unwrap(); + cache.encode("test").unwrap(); + assert_eq!(inner.encode_calls(), 3, "survivors must still hit"); + + cache.encode("world").unwrap(); + assert_eq!(inner.encode_calls(), 4, "evicted entry must miss"); + } + + #[test] + fn evicts_by_byte_budget() { + let per_entry = mock_entry_bytes("Hello"); + assert_eq!(mock_entry_bytes("world"), per_entry); + assert_eq!(mock_entry_bytes("token"), per_entry); + let inner = CountingTokenizer::new(); + // Room for two entries but not three. + let cache = cached(&inner, config(100, 2 * per_entry + per_entry / 2)); + + cache.encode("Hello").unwrap(); + cache.encode("world").unwrap(); + assert_eq!(cache.stats().evictions, 0); + + cache.encode("token").unwrap(); + let stats = cache.stats(); + assert_eq!(stats.entries, 2); + assert_eq!(stats.evictions, 1); + assert_eq!(stats.bytes, 2 * per_entry); + assert!(stats.bytes <= cache.config().max_bytes); + + cache.encode("Hello").unwrap(); + assert_eq!( + inner.encode_calls(), + 4, + "oldest entry must have been evicted" + ); + assert_eq!(cache.recomputed_bytes(), cache.stats().bytes); + } + + #[test] + fn byte_budget_can_evict_several_small_entries_for_one_larger() { + let small = mock_entry_bytes("Hello"); + let large_input = "Hello world ".repeat(100); + let large = mock_entry_bytes(&large_input); + assert!(large > 3 * small); + let inner = CountingTokenizer::new(); + let cache = cached(&inner, config(100, large + small / 2)); + + cache.encode("Hello").unwrap(); + cache.encode("world").unwrap(); + cache.encode("token").unwrap(); + cache.encode(&large_input).unwrap(); + + let stats = cache.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.evictions, 3); + assert_eq!(stats.bytes, large); + assert_eq!(cache.recomputed_bytes(), large); + } + + #[test] + fn oversized_results_bypass_cache() { + let inner = CountingTokenizer::new(); + let cache = cached( + &inner, + TokenizerCacheConfig { + max_entries: 100, + max_bytes: 1 << 20, + max_entry_bytes: ENTRY_OVERHEAD_BYTES + 1, + }, + ); + + let encoding = cache.encode("Hello world").unwrap(); + assert_eq!(encoding.token_ids(), &[1, 2]); + cache.encode("Hello world").unwrap(); + + assert_eq!(inner.encode_calls(), 2); + let stats = cache.stats(); + assert_eq!(stats.misses, 2); + assert_eq!(stats.oversized, 2); + assert_eq!(stats.entries, 0); + assert_eq!(stats.bytes, 0); + } + + #[test] + fn batch_encode_bypasses_cache() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + let encodings = cache.encode_batch(&["Hello", "world"]).unwrap(); + assert_eq!(encodings.len(), 2); + assert_eq!(inner.batch_calls(), 1); + assert_eq!(cache.stats().entries, 0, "batch must not populate"); + + cache.encode("Hello").unwrap(); + assert_eq!(inner.encode_calls(), 1); + + cache.encode_batch(&["Hello"]).unwrap(); + assert_eq!(inner.batch_calls(), 2, "batch must not read"); + assert_eq!(cache.stats().hits, 0); + } + + #[test] + fn same_key_reinsert_replaces_without_double_counting() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, config(10, usize::MAX)); + let encoding = Arc::new(MockTokenizer::new().encode("Hello").unwrap()); + let bytes = mock_entry_bytes("Hello"); + + // Models two threads that both missed on the same input. + cache.insert("Hello", encoding.clone(), bytes); + cache.insert("Hello", encoding, bytes); + + let stats = cache.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.bytes, bytes); + assert_eq!(stats.evictions, 0); + assert_eq!(cache.recomputed_bytes(), bytes); + } + + #[test] + fn concurrent_access_keeps_accounting_consistent() { + let inner = CountingTokenizer::new(); + let cache = Arc::new(cached(&inner, config(8, usize::MAX))); + let inputs: Vec = (0..16).map(|i| format!("Hello world {i}")).collect(); + let reference = MockTokenizer::new(); + let expected: Vec> = inputs + .iter() + .map(|input| reference.encode(input).unwrap().token_ids().to_vec()) + .collect(); + + let handles: Vec<_> = (0..8) + .map(|t| { + let cache = cache.clone(); + let inputs = inputs.clone(); + let expected = expected.clone(); + thread::spawn(move || { + for i in 0..200 { + let index = (i * 7 + t) % inputs.len(); + let encoding = cache.encode(&inputs[index]).unwrap(); + assert_eq!(encoding.token_ids(), expected[index].as_slice()); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + let stats = cache.stats(); + assert_eq!(stats.hits + stats.misses, 8 * 200); + assert!(stats.entries <= 8); + assert_eq!(stats.entries, cache.len()); + assert_eq!(stats.bytes, cache.recomputed_bytes()); + assert!(stats.evictions > 0); + } + + #[test] + fn clear_drops_entries_and_keeps_counters() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + cache.encode("Hello").unwrap(); + cache.encode("Hello").unwrap(); + assert!(!cache.is_empty()); + + cache.clear(); + assert!(cache.is_empty()); + let stats = cache.stats(); + assert_eq!(stats.entries, 0); + assert_eq!(stats.bytes, 0); + assert_eq!(stats.hits, 1); + assert_eq!(stats.misses, 1); + + cache.encode("Hello").unwrap(); + assert_eq!(inner.encode_calls(), 2); + } + + #[test] + fn delegates_decode_and_metadata() { + let inner = CountingTokenizer::new(); + let cache = cached(&inner, TokenizerCacheConfig::default()); + + assert_eq!(cache.decode(&[1, 2], false).unwrap(), "Hello world"); + assert_eq!(cache.decode(&[1000, 1, 999], true).unwrap(), "Hello"); + assert_eq!(cache.vocab_size(), 8); + assert_eq!( + cache.get_special_tokens().bos_token.as_deref(), + Some("") + ); + assert_eq!(cache.token_to_id("Hello"), Some(1)); + assert_eq!(cache.token_to_id("nope"), None); + assert_eq!(cache.id_to_token(2).as_deref(), Some("world")); + assert_eq!(cache.id_to_token(4242), None); + assert_eq!(cache.stats().entries, 0); + } + + #[test] + fn works_behind_tokenizer_wrapper() { + let inner = CountingTokenizer::new(); + let tokenizer = Tokenizer::from_arc(Arc::new(cached(&inner, Default::default()))); + + let encoding = tokenizer.encode("Hello world").unwrap(); + tokenizer.encode("Hello world").unwrap(); + assert_eq!(inner.encode_calls(), 1); + assert_eq!( + tokenizer.decode(encoding.token_ids(), false).unwrap(), + "Hello world" + ); + } + + #[test] + fn config_validation() { + assert!(TokenizerCacheConfig::default().validate().is_ok()); + assert!(config(0, 1024).validate().is_err()); + assert!(config(1, 0).validate().is_err()); + assert!(TokenizerCacheConfig { + max_entries: 1, + max_bytes: 1024, + max_entry_bytes: 0, + } + .validate() + .is_err()); + assert!(TokenizerCacheConfig { + max_entries: 1, + max_bytes: 1024, + max_entry_bytes: 2048, + } + .validate() + .is_err()); + + let inner = CountingTokenizer::new(); + assert!(CachedTokenizer::new(inner.clone(), config(0, 1)).is_err()); + assert!(CachedTokenizer::new(inner, config(1, 1)).is_ok()); + } + + #[test] + fn estimates_for_id_only_encodings() { + let header = size_of::>(); + assert_eq!( + estimate_encoding_bytes(&Encoding::Sp(vec![1, 2, 3])), + header + 3 * 4 + ); + assert_eq!( + estimate_encoding_bytes(&Encoding::Tiktoken(vec![7; 10])), + header + 10 * 4 + ); + assert_eq!(estimate_encoding_bytes(&Encoding::Sp(vec![])), header); + assert_eq!( + estimate_entry_bytes("abc", &Encoding::Sp(vec![1])), + ENTRY_OVERHEAD_BYTES + 3 + header + 4 + ); + } + + #[test] + fn estimates_for_hf_encodings_count_tokens_text_and_overflow() { + use tokenizers::tokenizer::Encoding as HfEncoding; + + fn hf(tokens: &[&str], overflowing: Vec) -> HfEncoding { + let n = tokens.len(); + HfEncoding::new( + (0..n as u32).collect(), + vec![0; n], + tokens.iter().map(|t| t.to_string()).collect(), + vec![None; n], + vec![(0, 0); n], + vec![0; n], + vec![1; n], + overflowing, + Default::default(), + ) + } + + let per_token = 4 * size_of::() + + size_of::>() + + size_of::<(usize, usize)>() + + size_of::(); + let struct_size = size_of::(); + + let plain = Encoding::Hf(Box::new(hf(&["ab", "c"], vec![]))); + assert_eq!( + estimate_encoding_bytes(&plain), + struct_size + 2 * per_token + 3 + ); + + let nested = Encoding::Hf(Box::new(hf(&["ab", "c"], vec![hf(&["xyz"], vec![])]))); + assert_eq!( + estimate_encoding_bytes(&nested), + struct_size + 2 * per_token + 3 + struct_size + per_token + 3 + ); + + let empty = Encoding::Hf(Box::new(hf(&[], vec![]))); + assert_eq!(estimate_encoding_bytes(&empty), struct_size); + } +} diff --git a/src/tokenizer/mod.rs b/src/tokenizer/mod.rs index 98a23f76..11ce01b2 100644 --- a/src/tokenizer/mod.rs +++ b/src/tokenizer/mod.rs @@ -2,6 +2,7 @@ use anyhow::Result; use std::ops::Deref; use std::sync::Arc; +pub mod cache; pub mod factory; pub mod hub; pub mod mock; @@ -22,6 +23,7 @@ pub mod tiktoken; mod tests; // Re-exports +pub use cache::{CachedTokenizer, TokenizerCacheConfig, TokenizerCacheStats}; pub use factory::{ create_tokenizer, create_tokenizer_async, create_tokenizer_from_file, create_tokenizer_with_chat_template, TokenizerType, diff --git a/tests/fixtures/tokenizer/byte_level_bpe.json b/tests/fixtures/tokenizer/byte_level_bpe.json new file mode 100644 index 00000000..996aa6d9 --- /dev/null +++ b/tests/fixtures/tokenizer/byte_level_bpe.json @@ -0,0 +1,391 @@ +{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [ + { + "id": 290, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 291, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + }, + { + "id": 292, + "content": "", + "single_word": false, + "lstrip": false, + "rstrip": false, + "normalized": false, + "special": true + } + ], + "normalizer": null, + "pre_tokenizer": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "post_processor": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "decoder": { + "type": "ByteLevel", + "add_prefix_space": false, + "trim_offsets": true, + "use_regex": true + }, + "model": { + "type": "BPE", + "dropout": null, + "unk_token": null, + "continuing_subword_prefix": null, + "end_of_word_suffix": null, + "fuse_unk": false, + "byte_fallback": false, + "ignore_merges": false, + "vocab": { + "Ā": 0, + "ā": 1, + "Ă": 2, + "ă": 3, + "Ą": 4, + "ą": 5, + "Ć": 6, + "ć": 7, + "Ĉ": 8, + "ĉ": 9, + "Ċ": 10, + "ċ": 11, + "Č": 12, + "č": 13, + "Ď": 14, + "ď": 15, + "Đ": 16, + "đ": 17, + "Ē": 18, + "ē": 19, + "Ĕ": 20, + "ĕ": 21, + "Ė": 22, + "ė": 23, + "Ę": 24, + "ę": 25, + "Ě": 26, + "ě": 27, + "Ĝ": 28, + "ĝ": 29, + "Ğ": 30, + "ğ": 31, + "Ġ": 32, + "!": 33, + "\"": 34, + "#": 35, + "$": 36, + "%": 37, + "&": 38, + "'": 39, + "(": 40, + ")": 41, + "*": 42, + "+": 43, + ",": 44, + "-": 45, + ".": 46, + "/": 47, + "0": 48, + "1": 49, + "2": 50, + "3": 51, + "4": 52, + "5": 53, + "6": 54, + "7": 55, + "8": 56, + "9": 57, + ":": 58, + ";": 59, + "<": 60, + "=": 61, + ">": 62, + "?": 63, + "@": 64, + "A": 65, + "B": 66, + "C": 67, + "D": 68, + "E": 69, + "F": 70, + "G": 71, + "H": 72, + "I": 73, + "J": 74, + "K": 75, + "L": 76, + "M": 77, + "N": 78, + "O": 79, + "P": 80, + "Q": 81, + "R": 82, + "S": 83, + "T": 84, + "U": 85, + "V": 86, + "W": 87, + "X": 88, + "Y": 89, + "Z": 90, + "[": 91, + "\\": 92, + "]": 93, + "^": 94, + "_": 95, + "`": 96, + "a": 97, + "b": 98, + "c": 99, + "d": 100, + "e": 101, + "f": 102, + "g": 103, + "h": 104, + "i": 105, + "j": 106, + "k": 107, + "l": 108, + "m": 109, + "n": 110, + "o": 111, + "p": 112, + "q": 113, + "r": 114, + "s": 115, + "t": 116, + "u": 117, + "v": 118, + "w": 119, + "x": 120, + "y": 121, + "z": 122, + "{": 123, + "|": 124, + "}": 125, + "~": 126, + "ġ": 127, + "Ģ": 128, + "ģ": 129, + "Ĥ": 130, + "ĥ": 131, + "Ħ": 132, + "ħ": 133, + "Ĩ": 134, + "ĩ": 135, + "Ī": 136, + "ī": 137, + "Ĭ": 138, + "ĭ": 139, + "Į": 140, + "į": 141, + "İ": 142, + "ı": 143, + "IJ": 144, + "ij": 145, + "Ĵ": 146, + "ĵ": 147, + "Ķ": 148, + "ķ": 149, + "ĸ": 150, + "Ĺ": 151, + "ĺ": 152, + "Ļ": 153, + "ļ": 154, + "Ľ": 155, + "ľ": 156, + "Ŀ": 157, + "ŀ": 158, + "Ł": 159, + "ł": 160, + "¡": 161, + "¢": 162, + "£": 163, + "¤": 164, + "¥": 165, + "¦": 166, + "§": 167, + "¨": 168, + "©": 169, + "ª": 170, + "«": 171, + "¬": 172, + "Ń": 173, + "®": 174, + "¯": 175, + "°": 176, + "±": 177, + "²": 178, + "³": 179, + "´": 180, + "µ": 181, + "¶": 182, + "·": 183, + "¸": 184, + "¹": 185, + "º": 186, + "»": 187, + "¼": 188, + "½": 189, + "¾": 190, + "¿": 191, + "À": 192, + "Á": 193, + "Â": 194, + "Ã": 195, + "Ä": 196, + "Å": 197, + "Æ": 198, + "Ç": 199, + "È": 200, + "É": 201, + "Ê": 202, + "Ë": 203, + "Ì": 204, + "Í": 205, + "Î": 206, + "Ï": 207, + "Ð": 208, + "Ñ": 209, + "Ò": 210, + "Ó": 211, + "Ô": 212, + "Õ": 213, + "Ö": 214, + "×": 215, + "Ø": 216, + "Ù": 217, + "Ú": 218, + "Û": 219, + "Ü": 220, + "Ý": 221, + "Þ": 222, + "ß": 223, + "à": 224, + "á": 225, + "â": 226, + "ã": 227, + "ä": 228, + "å": 229, + "æ": 230, + "ç": 231, + "è": 232, + "é": 233, + "ê": 234, + "ë": 235, + "ì": 236, + "í": 237, + "î": 238, + "ï": 239, + "ð": 240, + "ñ": 241, + "ò": 242, + "ó": 243, + "ô": 244, + "õ": 245, + "ö": 246, + "÷": 247, + "ø": 248, + "ù": 249, + "ú": 250, + "û": 251, + "ü": 252, + "ý": 253, + "þ": 254, + "ÿ": 255, + "He": 256, + "lo": 257, + "Hel": 258, + "Hello": 259, + "wo": 260, + "rl": 261, + "Ġwo": 262, + "Ġworl": 263, + "Ġworld": 264, + "le": 265, + "ar": 266, + "in": 267, + "ing": 268, + "Ġle": 269, + "Ġlear": 270, + "Ġlearn": 271, + "Ġlearning": 272, + "is": 273, + "Ġis": 274, + "th": 275, + "the": 276, + "Ġthe": 277, + "an": 278, + "Ġa": 279, + "Ġan": 280, + "Ġand": 281, + "Ġo": 282, + "Ġof": 283, + "Ġt": 284, + "Ġto": 285, + "Ġd": 286, + "Ġde": 287, + "Ġdee": 288, + "Ġdeep": 289 + }, + "merges": [ + "H e", + "l o", + "He l", + "Hel lo", + "w o", + "r l", + "Ġ wo", + "Ġwo rl", + "Ġworl d", + "l e", + "a r", + "i n", + "in g", + "Ġ le", + "Ġle ar", + "Ġlear n", + "Ġlearn ing", + "i s", + "Ġ is", + "t h", + "th e", + "Ġ the", + "a n", + "Ġ a", + "Ġ an", + "Ġan d", + "Ġ o", + "Ġo f", + "Ġ t", + "Ġt o", + "Ġ d", + "Ġd e", + "Ġde e", + "Ġdee p" + ] + } +} diff --git a/tests/tokenizer_cache_integration.rs b/tests/tokenizer_cache_integration.rs new file mode 100644 index 00000000..c90d1021 --- /dev/null +++ b/tests/tokenizer_cache_integration.rs @@ -0,0 +1,334 @@ +//! Integration tests for `CachedTokenizer` with real HuggingFace encodings. +//! +//! The default tests load `tests/fixtures/tokenizer/byte_level_bpe.json`, a +//! small byte-level BPE tokenizer checked into the repository (256 byte +//! symbols, a few merges, ``/``/`` as special tokens). It handles +//! arbitrary UTF-8 and yields `Encoding::Hf` values with tokens, offsets and +//! masks, without network access. One `#[ignore]` test repeats the core +//! comparison with the TinyLlama tokenizer, which is downloaded on first use. + +mod common; + +use std::sync::Arc; +use tokenizers::{Tokenizer as HfTokenizer, TruncationParams}; +use vllm_router_rs::tokenizer::{ + cache::{estimate_entry_bytes, CachedTokenizer, TokenizerCacheConfig}, + huggingface::HuggingFaceTokenizer, + traits::*, + Tokenizer as TokenizerWrapper, +}; + +const FIXTURE: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/tokenizer/byte_level_bpe.json" +); + +const PROMPTS: &[&str] = &[ + "", + " ", + "deep learning is", + "Deep learning is", + "Hello, world!", + " wrapped in special tokens ", + " unknown and repeated markers", + "tabs\tand\nnewlines and double spaces", + "na\u{ef}ve caf\u{e9} \u{2014} accents and dashes", + "\u{1f600}\u{1f603}\u{1f604}\u{1f601}\u{1f606} emoji \u{1f92a}\u{1f47b}", + "\u{4f60}\u{597d}\u{ff0c}\u{4e16}\u{754c} CJK text", + "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor \ + incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud \ + exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.", +]; + +fn load_fixture() -> Arc { + Arc::new(HuggingFaceTokenizer::from_file(FIXTURE).expect("Failed to load fixture tokenizer")) +} + +/// The fixture with truncation enabled, so long inputs produce overflowing +/// encodings. +fn load_fixture_with_overflow() -> Arc { + let mut tokenizer = HfTokenizer::from_file(FIXTURE).expect("Failed to load fixture tokenizer"); + tokenizer + .with_truncation(Some(TruncationParams { + max_length: 8, + stride: 2, + ..TruncationParams::default() + })) + .expect("valid truncation"); + Arc::new(HuggingFaceTokenizer::from_tokenizer(tokenizer)) +} + +fn assert_hf_equal(expected: &Encoding, actual: &Encoding) { + match (expected, actual) { + (Encoding::Hf(expected), Encoding::Hf(actual)) => { + assert_eq!(expected.get_ids(), actual.get_ids()); + assert_eq!(expected.get_tokens(), actual.get_tokens()); + assert_eq!(expected.get_offsets(), actual.get_offsets()); + assert_eq!(expected.get_type_ids(), actual.get_type_ids()); + assert_eq!(expected.get_word_ids(), actual.get_word_ids()); + assert_eq!( + expected.get_special_tokens_mask(), + actual.get_special_tokens_mask() + ); + assert_eq!(expected.get_attention_mask(), actual.get_attention_mask()); + assert_eq!(expected.get_overflowing(), actual.get_overflowing()); + assert_eq!(expected, actual); + } + _ => panic!("expected HuggingFace encodings"), + } +} + +fn assert_cached_matches_uncached(inner: Arc, prompts: &[&str]) { + let cache = CachedTokenizer::new(inner.clone(), TokenizerCacheConfig::default()).unwrap(); + + for prompt in prompts { + let uncached = inner.encode(prompt).unwrap(); + let first = cache.encode(prompt).unwrap(); + let second = cache.encode(prompt).unwrap(); + assert_hf_equal(&uncached, &first); + assert_hf_equal(&uncached, &second); + assert_eq!(uncached.get_hash(), second.get_hash()); + } + + let stats = cache.stats(); + assert_eq!(stats.misses as usize, prompts.len()); + assert_eq!(stats.hits as usize, prompts.len()); + assert_eq!(stats.entries, prompts.len()); + assert_eq!(stats.evictions, 0); + assert_eq!(stats.oversized, 0); +} + +#[test] +fn fixture_is_the_expected_tokenizer() { + let tokenizer = load_fixture(); + assert_eq!(tokenizer.vocab_size(), 290); + assert_eq!( + tokenizer.get_special_tokens().bos_token.as_deref(), + Some("") + ); + assert_eq!( + tokenizer.get_special_tokens().eos_token.as_deref(), + Some("") + ); + assert_eq!( + tokenizer.get_special_tokens().unk_token.as_deref(), + Some("") + ); + + // Pinned output so a silently changed fixture is caught. + let encoding = tokenizer.encode("Hello world").unwrap(); + assert_eq!(encoding.token_ids(), &[259, 264]); + let encoding = tokenizer.encode("deep learning is the").unwrap(); + assert_eq!(encoding.token_ids(), &[100, 101, 101, 112, 272, 274, 277]); + let encoding = tokenizer.encode(" hi ").unwrap(); + assert_eq!(encoding.token_ids(), &[290, 32, 104, 105, 32, 291]); + assert_eq!( + tokenizer.decode(encoding.token_ids(), true).unwrap(), + " hi " + ); +} + +#[test] +fn cached_encodings_match_uncached_including_metadata() { + assert_cached_matches_uncached(load_fixture(), PROMPTS); +} + +#[test] +fn overflowing_encodings_are_preserved() { + let inner = load_fixture_with_overflow(); + let long = PROMPTS[PROMPTS.len() - 1]; + let Encoding::Hf(reference) = inner.encode(long).unwrap() else { + panic!("expected HuggingFace encoding"); + }; + assert!( + !reference.get_overflowing().is_empty(), + "fixture truncation must produce overflowing encodings" + ); + + assert_cached_matches_uncached(inner.clone(), PROMPTS); + + let cache = CachedTokenizer::new(inner.clone(), TokenizerCacheConfig::default()).unwrap(); + let hit = { + cache.encode(long).unwrap(); + cache.encode(long).unwrap() + }; + let Encoding::Hf(hit) = hit else { + panic!("expected HuggingFace encoding"); + }; + assert_eq!(hit.get_overflowing(), reference.get_overflowing()); + + let without_overflow = load_fixture().encode(long).unwrap(); + assert!( + estimate_entry_bytes(long, &Encoding::Hf(reference)) + > estimate_entry_bytes(long, &without_overflow), + "overflowing encodings must count toward the estimate" + ); +} + +#[test] +fn retained_bytes_track_estimates() { + let inner = load_fixture(); + let cache = CachedTokenizer::new(inner.clone(), TokenizerCacheConfig::default()).unwrap(); + + let mut expected_bytes = 0; + for prompt in PROMPTS { + let encoding = inner.encode(prompt).unwrap(); + let estimate = estimate_entry_bytes(prompt, &encoding); + // Every retained byte we know about must be counted. + assert!(estimate >= prompt.len() + encoding.token_ids().len() * 4); + expected_bytes += estimate; + cache.encode(prompt).unwrap(); + assert_eq!(cache.stats().bytes, expected_bytes); + } +} + +#[test] +fn byte_budget_evicts_oldest_real_encodings() { + let inner = load_fixture(); + let short = "deep learning is"; + let long = PROMPTS[PROMPTS.len() - 1]; + let short_bytes = estimate_entry_bytes(short, &inner.encode(short).unwrap()); + let long_bytes = estimate_entry_bytes(long, &inner.encode(long).unwrap()); + assert!(long_bytes > short_bytes); + + let cache = CachedTokenizer::new( + inner, + TokenizerCacheConfig { + max_entries: 100, + max_bytes: long_bytes + short_bytes / 2, + max_entry_bytes: long_bytes, + }, + ) + .unwrap(); + + cache.encode(short).unwrap(); + cache.encode(long).unwrap(); + let stats = cache.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.evictions, 1); + assert_eq!(stats.bytes, long_bytes); + + cache.encode(long).unwrap(); + assert_eq!(cache.stats().hits, 1); + cache.encode(short).unwrap(); + assert_eq!(cache.stats().misses, 3); +} + +#[test] +fn oversized_real_encoding_is_returned_but_not_stored() { + let inner = load_fixture(); + let long = PROMPTS[PROMPTS.len() - 1]; + let long_bytes = estimate_entry_bytes(long, &inner.encode(long).unwrap()); + let cache = CachedTokenizer::new( + inner.clone(), + TokenizerCacheConfig { + max_entries: 100, + max_bytes: 1 << 20, + max_entry_bytes: long_bytes - 1, + }, + ) + .unwrap(); + + assert_hf_equal(&inner.encode(long).unwrap(), &cache.encode(long).unwrap()); + let stats = cache.stats(); + assert_eq!(stats.oversized, 1); + assert_eq!(stats.entries, 0); + assert_eq!(stats.bytes, 0); +} + +#[test] +fn works_behind_wrapper_with_decode_and_streaming() { + let inner = load_fixture(); + let cached = Arc::new(CachedTokenizer::new(inner, TokenizerCacheConfig::default()).unwrap()); + let tokenizer = TokenizerWrapper::from_arc(cached.clone()); + + let prompt = "The quick brown fox jumps over the lazy dog"; + let encoding = tokenizer.encode(prompt).unwrap(); + tokenizer.encode(prompt).unwrap(); + assert_eq!(cached.stats().hits, 1); + + // Byte-level BPE decodes losslessly. + assert_eq!( + tokenizer.decode(encoding.token_ids(), true).unwrap(), + prompt + ); + + let mut stream = tokenizer.decode_stream(&[], true); + let mut text = String::new(); + for &id in encoding.token_ids() { + if let Some(chunk) = stream.step(id).unwrap() { + text.push_str(&chunk); + } + } + if let Some(chunk) = stream.flush().unwrap() { + text.push_str(&chunk); + } + assert_eq!(text, prompt); + + // Batch encoding is delegated and leaves the cache untouched. + let batch = tokenizer.encode_batch(&[prompt, "another prompt"]).unwrap(); + assert_eq!(batch.len(), 2); + assert_eq!(cached.stats().entries, 1); +} + +#[test] +fn concurrent_real_encodes_are_consistent() { + use std::thread; + + let inner = load_fixture(); + let cache = Arc::new( + CachedTokenizer::new( + inner.clone(), + TokenizerCacheConfig { + max_entries: 6, + max_bytes: 1 << 20, + max_entry_bytes: 1 << 20, + }, + ) + .unwrap(), + ); + let expected: Vec = PROMPTS.iter().map(|p| inner.encode(p).unwrap()).collect(); + + let handles: Vec<_> = (0..8) + .map(|t| { + let cache = cache.clone(); + let expected = expected.clone(); + thread::spawn(move || { + for i in 0..100 { + let index = (i * 5 + t) % PROMPTS.len(); + assert_hf_equal(&expected[index], &cache.encode(PROMPTS[index]).unwrap()); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + let stats = cache.stats(); + assert_eq!(stats.hits + stats.misses, 800); + assert!(stats.entries <= 6); + assert!(stats.evictions > 0); +} + +/// Model-level check with the TinyLlama tokenizer. Downloads the file on +/// first use; run with `cargo test --test tokenizer_cache_integration -- --ignored`. +#[test] +#[ignore = "downloads the TinyLlama tokenizer from Hugging Face"] +fn tinyllama_cached_encodings_match_uncached() { + let path = common::ensure_tokenizer_cached(); + let inner: Arc = Arc::new( + HuggingFaceTokenizer::from_file(path.to_str().unwrap()) + .expect("Failed to load TinyLlama tokenizer"), + ); + assert_cached_matches_uncached(inner.clone(), PROMPTS); + + let cache = CachedTokenizer::new(inner, TokenizerCacheConfig::default()).unwrap(); + for _ in 0..2 { + let hashes: Vec = common::TEST_PROMPTS + .iter() + .map(|prompt| cache.encode(prompt).unwrap().get_hash()) + .collect(); + assert_eq!(hashes, common::EXPECTED_HASHES); + } +} diff --git a/tests/tokenizer_cache_metrics.rs b/tests/tokenizer_cache_metrics.rs new file mode 100644 index 00000000..1141f44f --- /dev/null +++ b/tests/tokenizer_cache_metrics.rs @@ -0,0 +1,284 @@ +//! Gauge semantics of `CachedTokenizer`. +//! +//! `vllm_tokenizer_cache_entries` and `vllm_tokenizer_cache_bytes` must equal +//! the total occupancy of every live cache in the process, including after +//! concurrent inserts and clears and after instances are dropped. A capturing +//! `metrics::Recorder` records what the cache publishes. +//! +//! All caches in a process feed one aggregate, so these tests serialize on a +//! lock and drop every cache they create before releasing it. + +use metrics::{ + Counter, Gauge, GaugeFn, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit, +}; +use parking_lot::{Condvar, Mutex}; +use std::collections::HashMap; +use std::sync::Arc; +use std::thread; +use std::time::Duration; +use vllm_router_rs::tokenizer::{ + mock::MockTokenizer, traits::Encoder, CachedTokenizer, TokenizerCacheConfig, +}; + +static SERIAL: Mutex<()> = Mutex::new(()); + +const ENTRIES: &str = "vllm_tokenizer_cache_entries"; +const BYTES: &str = "vllm_tokenizer_cache_bytes"; + +#[derive(Default)] +struct Captured { + values: Mutex>, + /// When set, the next gauge `set` call takes this flag and blocks until + /// `release`. Later calls are unaffected. + block_next_set: Mutex, + released: Mutex, + release_signal: Condvar, +} + +impl Captured { + fn gauges(&self) -> (f64, f64) { + let values = self.values.lock(); + ( + values.get(ENTRIES).copied().unwrap_or(0.0), + values.get(BYTES).copied().unwrap_or(0.0), + ) + } + + /// Make the next gauge `set` block until `release` is called. + fn block_next_set(&self) { + *self.block_next_set.lock() = true; + } + + fn release(&self) { + *self.released.lock() = true; + self.release_signal.notify_all(); + } +} + +struct CapturingGauge { + captured: Arc, + name: String, +} + +impl GaugeFn for CapturingGauge { + fn increment(&self, value: f64) { + *self + .captured + .values + .lock() + .entry(self.name.clone()) + .or_default() += value; + } + + fn decrement(&self, value: f64) { + *self + .captured + .values + .lock() + .entry(self.name.clone()) + .or_default() -= value; + } + + fn set(&self, value: f64) { + let blocked = std::mem::take(&mut *self.captured.block_next_set.lock()); + if blocked { + let mut released = self.captured.released.lock(); + while !*released { + self.captured.release_signal.wait(&mut released); + } + } + self.captured.values.lock().insert(self.name.clone(), value); + } +} + +struct CapturingRecorder(Arc); + +impl Recorder for CapturingRecorder { + fn describe_counter(&self, _: KeyName, _: Option, _: SharedString) {} + fn describe_gauge(&self, _: KeyName, _: Option, _: SharedString) {} + fn describe_histogram(&self, _: KeyName, _: Option, _: SharedString) {} + + fn register_counter(&self, _: &Key, _: &Metadata<'_>) -> Counter { + Counter::noop() + } + + fn register_gauge(&self, key: &Key, _: &Metadata<'_>) -> Gauge { + Gauge::from_arc(Arc::new(CapturingGauge { + captured: self.0.clone(), + name: key.name().to_string(), + })) + } + + fn register_histogram(&self, _: &Key, _: &Metadata<'_>) -> Histogram { + Histogram::noop() + } +} + +fn cache(max_entries: usize) -> CachedTokenizer { + CachedTokenizer::new( + Arc::new(MockTokenizer::new()), + TokenizerCacheConfig { + max_entries, + ..TokenizerCacheConfig::default() + }, + ) + .unwrap() +} + +fn expected(caches: &[&CachedTokenizer]) -> (f64, f64) { + let entries: usize = caches.iter().map(|c| c.stats().entries).sum(); + let bytes: usize = caches.iter().map(|c| c.stats().bytes).sum(); + (entries as f64, bytes as f64) +} + +#[test] +fn gauges_sum_over_live_instances() { + let _serial = SERIAL.lock(); + let captured = Arc::new(Captured::default()); + let recorder = CapturingRecorder(captured.clone()); + + metrics::with_local_recorder(&recorder, || { + let a = cache(10); + let b = cache(10); + a.encode("Hello").unwrap(); + a.encode("world").unwrap(); + b.encode("Hello").unwrap(); + assert_eq!(captured.gauges(), (3.0, expected(&[&a, &b]).1)); + + // Pure hits do not touch the gauges. + a.encode("Hello").unwrap(); + b.encode("Hello").unwrap(); + assert_eq!(captured.gauges(), expected(&[&a, &b])); + + drop(b); + assert_eq!(captured.gauges(), expected(&[&a])); + + a.clear(); + assert_eq!(captured.gauges(), (0.0, 0.0)); + + a.encode("test").unwrap(); + assert_eq!(captured.gauges(), expected(&[&a])); + + drop(a); + assert_eq!(captured.gauges(), (0.0, 0.0)); + }); +} + +#[test] +fn eviction_and_replacement_keep_gauges_exact() { + let _serial = SERIAL.lock(); + let captured = Arc::new(Captured::default()); + let recorder = CapturingRecorder(captured.clone()); + + metrics::with_local_recorder(&recorder, || { + let a = cache(2); + for input in ["Hello", "world", "test", "token"] { + a.encode(input).unwrap(); + assert_eq!(captured.gauges(), expected(&[&a])); + } + assert_eq!(a.stats().evictions, 2); + + // Room for two mock entries but not three, so the byte budget evicts. + let per_entry = a.stats().bytes / 2; + let b = CachedTokenizer::new( + Arc::new(MockTokenizer::new()), + TokenizerCacheConfig { + max_entries: 100, + max_bytes: 2 * per_entry + per_entry / 2, + max_entry_bytes: 2 * per_entry + per_entry / 2, + }, + ) + .unwrap(); + for input in ["Hello", "world", "test"] { + b.encode(input).unwrap(); + assert_eq!(captured.gauges(), expected(&[&a, &b])); + } + assert!(b.stats().evictions > 0); + + drop(a); + drop(b); + assert_eq!(captured.gauges(), (0.0, 0.0)); + }); +} + +/// A publication that stalls inside the recorder must not be overtaken by a +/// later state change: the stalled operation still holds the cache lock, so +/// the other operation and its publication wait behind it. Only the first +/// `set` is blocked, so the second operation publishes freely; that is +/// exactly the interleaving that left a stale value when publication was +/// not ordered with the state change. +#[test] +fn delayed_publish_cannot_overwrite_a_later_clear() { + let _serial = SERIAL.lock(); + let captured = Arc::new(Captured::default()); + let recorder = Arc::new(CapturingRecorder(captured.clone())); + let shared = Arc::new(cache(10)); + + captured.block_next_set(); + let inserter = { + let (recorder, shared) = (recorder.clone(), shared.clone()); + thread::spawn(move || { + metrics::with_local_recorder(&*recorder, || shared.encode("Hello").unwrap()); + }) + }; + // Let the insert reach the recorder and stall there. + thread::sleep(Duration::from_millis(50)); + let clearer = { + let (recorder, shared) = (recorder.clone(), shared.clone()); + thread::spawn(move || metrics::with_local_recorder(&*recorder, || shared.clear())) + }; + thread::sleep(Duration::from_millis(50)); + captured.release(); + inserter.join().unwrap(); + clearer.join().unwrap(); + + let stats = shared.stats(); + assert_eq!( + captured.gauges(), + (stats.entries as f64, stats.bytes as f64), + "gauges must reflect the final state whichever operation ran last" + ); + + metrics::with_local_recorder(&*recorder, || drop(shared)); + assert_eq!(captured.gauges(), (0.0, 0.0)); +} + +#[test] +fn concurrent_inserts_and_clears_end_consistent() { + let _serial = SERIAL.lock(); + let captured = Arc::new(Captured::default()); + let recorder = Arc::new(CapturingRecorder(captured.clone())); + let a = Arc::new(cache(4)); + let b = Arc::new(cache(3)); + let inputs: Vec = (0..16).map(|i| format!("Hello world {i}")).collect(); + + let handles: Vec<_> = (0..8) + .map(|t| { + let (recorder, a, b, inputs) = (recorder.clone(), a.clone(), b.clone(), inputs.clone()); + thread::spawn(move || { + metrics::with_local_recorder(&*recorder, || { + for i in 0..300 { + let target = if (i + t) % 2 == 0 { &a } else { &b }; + if i % 37 == 0 { + target.clear(); + } else { + target.encode(&inputs[(i * 7 + t) % inputs.len()]).unwrap(); + } + } + }) + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + + assert_eq!(captured.gauges(), expected(&[&a, &b])); + assert!(a.stats().evictions > 0 && b.stats().evictions > 0); + + metrics::with_local_recorder(&*recorder, || { + drop(a); + drop(b); + }); + assert_eq!(captured.gauges(), (0.0, 0.0)); +}