From 14d50599fa262f69cd531d17fe8a8029bcfb5ffb Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 3 Nov 2026 15:49:24 +0900 Subject: [PATCH] feat(redis): move the v3 rate limiter keyspace onto the Rust store The limiter's keys ({descriptor:value}:requests, :tokens, :window and the :max_parallel_requests gauge), its eight Lua scripts and its memory tier now live in litellm-redis-store (rate_limits.rs). The scripts are byte for byte main's, so both stacks share one script SHA during a rolling deploy. The Python limiter keeps its admission logic and talks to a thin store in litellm/redis/stores/rate_limits.py with domain ids (RateLimitScope, RateLimitCounter) instead of key strings. Inside a request scope its scripts ride the request's native pre-call batch, so the first admission script and the armed cooldown read go out in one pipeline. Post-call token scripts, window-guarded settlements, slot releases and refund increments ride the native post-call batch, and their fallbacks run once it settles. The batch and dynamic limiters read and write the keyspace through the same store. Rust tests run every script against real Redis, standalone and cluster, including over-limit, refunds, gauge slot expiry and cross-slot behavior. Matching inline wire cases are checked from both sides. Co-Authored-By: Claude Opus 5.5 The limiter sends the commands main does. Counter reads, slot acquire, count and the pre-response release are direct EVALSHA calls with the counter-read groups one after another, window-guarded token adjusts run one by one on their own trips, plain increments are INCRBYFLOAT with EXPIRE, and an increment without a TTL takes RedisCache.default_ttl like async_increment. --- litellm-rust/crates/python-bridge/src/lib.rs | 4 + .../python-bridge/src/redis/store/mod.rs | 1 + .../src/redis/store/rate_limits.rs | 604 +++++ .../crates/redis-client/src/client.rs | 8 +- litellm-rust/crates/redis-client/src/lib.rs | 1 + litellm-rust/crates/redis-store/AGENTS.md | 5 +- litellm-rust/crates/redis-store/src/lib.rs | 1 + .../crates/redis-store/src/rate_limits.rs | 489 ++++ .../crates/redis-store/src/rate_limits/lua.rs | 260 ++ .../redis-store/src/rate_limits/scripts.rs | 391 +++ .../crates/redis-store/tests/rate_limits.rs | 902 +++++++ litellm/proxy/hooks/batch_rate_limiter.py | 36 +- .../proxy/hooks/dynamic_rate_limiter_v3.py | 77 +- .../hooks/parallel_request_limiter_v3.py | 1700 ++++--------- litellm/redis/stores/rate_limits.py | 404 ++++ litellm/rust_bridge/_native.pyi | 83 +- .../test_request_redis_batch_post_call.py | 156 +- .../test_request_redis_batch_pre_call.py | 278 +-- .../proxy/hooks/test_batch_rate_limiter.py | 8 +- .../hooks/test_dynamic_rate_limiter_v3.py | 392 +-- .../hooks/test_parallel_request_limiter_v3.py | 2144 ++++++----------- .../proxy/hooks/test_rate_limiter_toctou.py | 40 +- tests/unit/proxy/hooks/test_tpm_concurrent.py | 817 ++----- .../auth/test_user_api_key_auth_mcp.py | 8 +- .../proxy/test_common_request_processing.py | 5 +- tests/unit/redis/stores/fake_rate_limits.py | 261 ++ tests/unit/redis/stores/test_rate_limits.py | 158 ++ .../unit/test_rate_limit_error_unification.py | 11 +- 28 files changed, 5341 insertions(+), 3903 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs create mode 100644 litellm-rust/crates/redis-store/src/rate_limits.rs create mode 100644 litellm-rust/crates/redis-store/src/rate_limits/lua.rs create mode 100644 litellm-rust/crates/redis-store/src/rate_limits/scripts.rs create mode 100644 litellm-rust/crates/redis-store/tests/rate_limits.rs create mode 100644 litellm/redis/stores/rate_limits.py create mode 100644 tests/unit/redis/stores/fake_rate_limits.py create mode 100644 tests/unit/redis/stores/test_rate_limits.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 87d70d0799c..e41f25113db 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -43,6 +43,8 @@ mod _native { #[pymodule_export] use crate::redis::store::auth_cache::NativeAuthCache; #[pymodule_export] + use crate::redis::store::rate_limits::{NativeRateLimitPending, NativeRateLimits}; + #[pymodule_export] use crate::redis::store::request::NativeRequestRedis; #[pymodule_export] use crate::redis::store::router_usage::NativeRouterUsage; @@ -128,6 +130,8 @@ mod tests { "NativeRedisClient", "NativeRequestRedis", "NativeAuthCache", + "NativeRateLimits", + "NativeRateLimitPending", "NativePending", "NativeRouterCooldowns", "NativeRouterUsage", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index 42dcc91b5f9..0dcd65dee17 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod auth_cache; mod policy; +pub(crate) mod rate_limits; pub(crate) mod request; pub(crate) mod router_cooldowns; pub(crate) mod router_usage; diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs new file mode 100644 index 00000000000..540d984c6e6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -0,0 +1,604 @@ +use std::{ + pin::Pin, + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; + +use litellm_redis_store::rate_limits::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Error, Gauge, + GuardedIncrement, GuardedReply, Increment, Kind, Local, Pending, RateLimits, Scope, + WindowCount, +}; +use pyo3::{ + exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + prelude::*, + types::{PyDict, PyString}, +}; + +use super::{PyNumber, call, request::NativeRequestRedis, store_error}; +use crate::execution::run_async_value; +use crate::redis::client::NativeRedisClient; + +/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`; +/// every Redis call is a future. +#[pyclass(frozen, name = "NativeRateLimits")] +pub(crate) struct NativeRateLimits(Arc); + +/// `(descriptor_key, descriptor_value)`. +struct PyScope(Scope); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyScope { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (descriptor, value): (String, String) = value.extract()?; + Ok(Self(Scope { descriptor, value })) + } +} + +/// `(descriptor_key, descriptor_value, "requests" | "tokens")`. +struct PyCounter(Counter); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (descriptor, value, kind): (String, String, String) = value.extract()?; + Ok(Self(Counter { + scope: Scope { descriptor, value }, + kind: kind_of(&kind)?, + })) + } +} + +fn kind_of(kind: &str) -> PyResult { + match kind { + "requests" => Ok(Kind::Requests), + "tokens" => Ok(Kind::Tokens), + _ => Err(PyValueError::new_err(format!( + "unknown rate-limit counter {kind:?}" + ))), + } +} + +/// `("window" | "requests" | "tokens" | "max_parallel_requests", descriptor_key, descriptor_value)`. +struct PyEntry(Entry); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyEntry { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (slot, descriptor, value): (String, String, String) = value.extract()?; + let scope = Scope { descriptor, value }; + Ok(Self(match slot.as_str() { + "window" => Entry::Window(scope), + "max_parallel_requests" => Entry::Gauge(scope), + kind => Entry::Counter(Counter { + scope, + kind: kind_of(kind)?, + }), + })) + } +} + +/// `(counter, amount, ttl_seconds | None)`. +struct PyIncrement(Increment); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyIncrement { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (counter, amount, ttl): (PyCounter, PyNumber, Option) = value.extract()?; + Ok(Self(Increment { + counter: counter.0, + amount: amount.0, + ttl: ttl.map(Duration::from_secs), + })) + } +} + +/// `(counter, limit, increment, ttl_seconds, window_size_seconds)`. +struct PyAtomicCounter(AtomicCounter); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyAtomicCounter { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let (counter, limit, increment, ttl_seconds, window_size_seconds): ( + PyCounter, + i64, + i64, + i64, + i64, + ) = value.extract()?; + Ok(Self(AtomicCounter { + counter: counter.0, + limit, + increment, + ttl_seconds, + window_size_seconds, + })) + } +} + +/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry. +struct PyLocal(Local); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyLocal { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + if let Ok(number) = value.extract::() { + return Ok(Self(Local::Number(number.0))); + } + if let Ok(text) = value.cast::() { + return Ok(Self(Local::Text(text.to_str()?.to_owned()))); + } + if let Ok(slots) = value.cast::() { + return slots + .iter() + .map(|(slot, acquired_at)| Ok((slot.extract()?, acquired_at.extract()?))) + .collect::>() + .map(|slots| Self(Local::Slots(slots))); + } + Err(PyTypeError::new_err( + "a rate-limit memory value is an int, a float, a str or a dict of slot timestamps", + )) + } +} + +impl<'py> IntoPyObject<'py> for PyLocal { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + Ok(match self.0 { + Local::Number(number) => PyNumber(number).into_pyobject(py)?, + Local::Text(text) => text.into_pyobject(py)?.into_any(), + Local::Slots(slots) => { + let registry = PyDict::new(py); + for (slot, acquired_at) in slots { + registry.set_item(slot, acquired_at)?; + } + registry.into_any() + } + }) + } +} + +enum Reply { + WindowCounts(Vec), + Values(Vec>), + Atomic(AtomicReply), + Guarded(GuardedReply), + Acquire(AcquireReply), + Counts(Vec), + Totals(Vec), +} + +impl<'py> IntoPyObject<'py> for Reply { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult { + match self { + Self::WindowCounts(counts) => counts + .into_iter() + .map( + |WindowCount { + window_start, + count, + }| (window_start, count), + ) + .collect::>() + .into_pyobject(py), + Self::Values(values) => values.into_pyobject(py), + Self::Atomic(AtomicReply::Applied(applied)) => ( + None::<(usize, i64, i64)>, + applied + .into_iter() + .map( + |Applied { + count, + window_start, + }| (count, window_start), + ) + .collect::>(), + ) + .into_pyobject(py) + .map(Bound::into_any), + Self::Atomic(AtomicReply::OverLimit { + index, + current, + limit, + }) => (Some((index, current, limit)), Vec::<(i64, i64)>::new()) + .into_pyobject(py) + .map(Bound::into_any), + Self::Guarded(GuardedReply { applied, count }) => { + (applied, count).into_pyobject(py).map(Bound::into_any) + } + Self::Acquire(AcquireReply::Acquired(counts)) => (None::<(usize, i64, i64)>, counts) + .into_pyobject(py) + .map(Bound::into_any), + Self::Acquire(AcquireReply::OverLimit { + index, + in_flight, + limit, + }) => (Some((index, in_flight, limit)), Vec::::new()) + .into_pyobject(py) + .map(Bound::into_any), + Self::Counts(counts) => counts.into_pyobject(py), + Self::Totals(totals) => totals.into_pyobject(py), + } + } +} + +type Resolve = Pin> + Send>>; + +/// One script call declared on a batch: the request's pre-call batch, or a batch of its own when +/// no request was passed. `resolve` flushes that batch if it has not gone out, once. +#[pyclass(frozen, name = "NativeRateLimitPending")] +pub(crate) struct NativeRateLimitPending { + resolve: Mutex>, + call_type: &'static str, +} + +impl NativeRateLimitPending { + fn new( + pending: Result, Error>, + reply: fn(T) -> Reply, + call_type: &'static str, + ) -> PyResult { + let pending = pending.map_err(rate_limit_error)?; + Ok(Self { + resolve: Mutex::new(Some(Box::pin( + async move { pending.resolve().await.map(reply) }, + ))), + call_type, + }) + } +} + +#[pymethods] +impl NativeRateLimitPending { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + let resolve = self + .resolve + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take() + .ok_or_else(|| { + PyRuntimeError::new_err("this rate-limit script was already resolved") + })?; + let scope = call(py, self.call_type)?; + run_async_value(py, async move { + scope.scope(resolve).await.map_err(rate_limit_error) + }) + } +} + +fn rate_limit_error(error: Error) -> PyErr { + PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string()) +} + +fn scopes(scopes: Vec) -> Vec { + scopes.into_iter().map(|scope| scope.0).collect() +} + +fn counters(counters: Vec) -> Vec { + counters.into_iter().map(|counter| counter.0).collect() +} + +fn increments(increments: Vec) -> Vec { + increments + .into_iter() + .map(|increment| increment.0) + .collect() +} + +fn request_of( + request: Option>, +) -> Option { + request.map(|request| request.0.clone()) +} + +#[pymethods] +impl NativeRateLimits { + #[new] + #[pyo3(signature = (client, backfill_ttl_seconds = None, redis_ttl_seconds = 60))] + fn new( + client: Option>, + backfill_ttl_seconds: Option, + redis_ttl_seconds: u64, + ) -> Self { + let backfill_ttl = backfill_ttl_seconds.map(|ttl| Duration::from_secs_f64(ttl.max(0.0))); + Self(Arc::new(RateLimits::new( + client.map(|client| client.0.clone()), + backfill_ttl, + Duration::from_secs(redis_ttl_seconds), + ))) + } + + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn slot_groups(&self, scopes_: Vec, force: bool) -> Vec> { + self.0.slot_groups(&scopes(scopes_), force) + } + + fn local(&self, entries: Vec) -> PyResult>> { + entries + .iter() + .map(|entry| { + self.0 + .local(&entry.0) + .map(|local| local.map(PyLocal)) + .map_err(store_error) + }) + .collect() + } + + #[pyo3(signature = (entry, value, ttl_seconds))] + fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option) -> PyResult<()> { + self.0 + .set_local( + &entry.0, + value.0, + ttl_seconds.map(|ttl| Duration::from_secs_f64(ttl.max(0.0))), + ) + .map_err(store_error) + } + + #[pyo3(signature = (entry, *, backfill_ttl_seconds = None))] + fn get<'py>( + &self, + py: Python<'py>, + entry: PyEntry, + backfill_ttl_seconds: Option, + ) -> PyResult> { + let store = self.0.clone(); + let backfill_ttl = backfill_ttl_seconds.map(|ttl| Duration::from_secs_f64(ttl.max(0.0))); + let scope = call(py, "async_get_cache")?; + run_async_value(py, async move { + scope + .scope(store.get(&entry.0, backfill_ttl)) + .await + .map(|local| local.map(PyLocal)) + .map_err(store_error) + }) + } + + #[pyo3(signature = (counters_, now, window_size, *, request = None))] + fn sliding_window( + &self, + counters_: Vec, + now: i64, + window_size: i64, + request: Option>, + ) -> PyResult { + let request = request_of(request); + NativeRateLimitPending::new( + self.0 + .sliding_window(request.as_ref(), &counters(counters_), now, window_size), + Reply::WindowCounts, + "batch_rate_limiter_script", + ) + } + + #[pyo3(signature = (counters_, *, request = None))] + fn read_counters( + &self, + counters_: Vec, + request: Option>, + ) -> PyResult { + let request = request_of(request); + NativeRateLimitPending::new( + self.0.read_counters(request.as_ref(), &counters(counters_)), + Reply::Values, + "batch_counter_read_script", + ) + } + + #[pyo3(signature = (counters_, *, request = None))] + fn check_and_increment( + &self, + counters_: Vec, + request: Option>, + ) -> PyResult { + let request = request_of(request); + let counters_ = counters_ + .into_iter() + .map(|counter| counter.0) + .collect::>(); + NativeRateLimitPending::new( + self.0.check_and_increment(request.as_ref(), &counters_), + Reply::Atomic, + "check_and_increment_by_n_script", + ) + } + + #[pyo3(signature = (counter, expected_window_start, amount, ttl_seconds, *, request = None))] + fn guarded_increment( + &self, + counter: PyCounter, + expected_window_start: String, + amount: PyNumber, + ttl_seconds: Option, + request: Option>, + ) -> PyResult { + let request = request_of(request); + let increment = GuardedIncrement { + counter: counter.0, + expected_window_start, + amount: amount.0, + ttl: ttl_seconds.map(Duration::from_secs), + }; + NativeRateLimitPending::new( + self.0.guarded_increment(request.as_ref(), &increment), + Reply::Guarded, + "window_guarded_token_increment_script", + ) + } + + #[pyo3(signature = (gauges, slot_ttl_seconds, slot_id, *, request = None))] + fn parallel_acquire( + &self, + gauges: Vec<(PyScope, i64)>, + slot_ttl_seconds: u64, + slot_id: &str, + request: Option>, + ) -> PyResult { + let request = request_of(request); + let gauges = gauges + .into_iter() + .map(|(scope, limit)| Gauge { + scope: scope.0, + limit, + }) + .collect::>(); + NativeRateLimitPending::new( + self.0.parallel_acquire( + request.as_ref(), + &gauges, + Duration::from_secs(slot_ttl_seconds), + slot_id, + ), + Reply::Acquire, + "parallel_acquire_script", + ) + } + + #[pyo3(signature = (gauges, slot_id, *, request = None))] + fn parallel_release( + &self, + gauges: Vec, + slot_id: &str, + request: Option>, + ) -> PyResult { + let request = request_of(request); + NativeRateLimitPending::new( + self.0 + .parallel_release(request.as_ref(), &scopes(gauges), slot_id), + Reply::Counts, + "parallel_release_script", + ) + } + + #[pyo3(signature = (gauges, slot_ttl_seconds, *, request = None))] + fn parallel_count( + &self, + gauges: Vec, + slot_ttl_seconds: u64, + request: Option>, + ) -> PyResult { + let request = request_of(request); + NativeRateLimitPending::new( + self.0.parallel_count( + request.as_ref(), + &scopes(gauges), + Duration::from_secs(slot_ttl_seconds), + ), + Reply::Counts, + "parallel_count_script", + ) + } + + #[pyo3(signature = (increments_, *, request = None))] + fn token_increment( + &self, + increments_: Vec, + request: Option>, + ) -> PyResult { + let request = request_of(request); + NativeRateLimitPending::new( + self.0 + .token_increment(request.as_ref(), &increments(increments_)), + Reply::Totals, + "token_increment_script", + ) + } + + fn token_increment_post_call<'py>( + &self, + py: Python<'py>, + request: PyRef<'_, NativeRequestRedis>, + increments_: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.0.clone(); + let increments_ = increments(increments_); + run_async_value(py, async move { + store + .token_increment_post_call(&request, increments_) + .map_err(rate_limit_error) + }) + } + + fn parallel_release_post_call<'py>( + &self, + py: Python<'py>, + request: PyRef<'_, NativeRequestRedis>, + gauges: Vec, + slot_id: String, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.0.clone(); + let gauges = scopes(gauges); + run_async_value(py, async move { + store + .parallel_release_post_call(&request, &gauges, &slot_id) + .map_err(rate_limit_error) + }) + } + + fn increment<'py>( + &self, + py: Python<'py>, + increment: PyIncrement, + ) -> PyResult> { + let store = self.0.clone(); + let scope = call(py, "async_increment")?; + run_async_value(py, async move { + scope + .scope(store.increment(&increment.0)) + .await + .map(PyNumber) + .map_err(store_error) + }) + } + + fn increment_many<'py>( + &self, + py: Python<'py>, + increments_: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments_ = increments(increments_); + let scope = call(py, "async_increment_pipeline")?; + run_async_value(py, async move { + scope + .scope(store.increment_many(&increments_)) + .await + .map_err(store_error) + }) + } + + fn increment_many_post_call<'py>( + &self, + py: Python<'py>, + request: PyRef<'_, NativeRequestRedis>, + increments_: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.0.clone(); + let increments_ = increments(increments_); + run_async_value(py, async move { + store + .increment_many_post_call(&request, &increments_) + .map_err(store_error) + }) + } +} diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 69c0121d5e3..a9b4d9deb1e 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -258,14 +258,14 @@ impl RedisClient { Pipeline::new(self.0.namespace.clone()) } - pub fn same_client(&self, other: &Self) -> bool { - Arc::ptr_eq(&self.0, &other.0) - } - pub fn is_cluster(&self) -> bool { matches!(self.0.connection, Connection::Cluster(_)) } + pub fn same_client(&self, other: &Self) -> bool { + Arc::ptr_eq(&self.0, &other.0) + } + pub fn batch(&self) -> RedisBatch { RedisBatch::new(self.clone()) } diff --git a/litellm-rust/crates/redis-client/src/lib.rs b/litellm-rust/crates/redis-client/src/lib.rs index 6064a81892d..7eaa3e21b57 100644 --- a/litellm-rust/crates/redis-client/src/lib.rs +++ b/litellm-rust/crates/redis-client/src/lib.rs @@ -23,5 +23,6 @@ pub use namespace::Namespace; pub use pipeline::Pipeline; pub use redis::{ BasicAuth, ErrorKind as RedisErrorKind, RedisError, StreamingCredentialsProvider, Value, + cluster_routing::Slot, }; pub use script::Script; diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index 31d3b445155..c65a3f7e0f1 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -1,8 +1,9 @@ -- One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys +- One module per Redis family (`router_cooldowns.rs`, `router_usage.rs`, `rate_limits.rs`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys - Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec - A read-through family builds its own `MemoryTier` (`tier.rs`, `main`'s `DualCache` over a `RedisCache`) from a `TierPolicy` the host passes in. Its `pub const POLICY` is `main`'s defaults, and the Python host overrides what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient` - Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not -- `tests/fixtures/router_wire.json` is the wire contract, checked here against real Redis and in `tests/unit/redis/stores/` against Python's own codec. Change both sides together +- Wire contracts are hand-written cases with expected values, inlined as the same literals in this crate's test and in `tests/unit/redis/stores/`, checked here against real Redis and there against Python's own codec. Change both sides together +- `rate_limits.rs` keeps the v3 limiter's Lua byte for byte (`rate_limits/lua.rs`) so pods on either stack share one script SHA; the admission decisions stay in the Python limiter - `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. A family exposes a request-aware variant that declares on it (`declare_active`) next to its direct method, which runs the same declaration on a batch of its own. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here - Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect - Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index b53df6d46fd..6d58aeb043f 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -2,6 +2,7 @@ pub mod auth_cache; mod codec; mod error; mod number; +pub mod rate_limits; mod request; pub mod router_cooldowns; pub mod router_usage; diff --git a/litellm-rust/crates/redis-store/src/rate_limits.rs b/litellm-rust/crates/redis-store/src/rate_limits.rs new file mode 100644 index 00000000000..95db7b5fe42 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits.rs @@ -0,0 +1,489 @@ +mod lua; +mod scripts; + +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; + +use futures_util::future::join_all; +use litellm_cache::Error as MemoryError; +use litellm_cache_memory::InMemoryCache; +use litellm_python_compat::Value; +use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient, Slot}; + +pub use self::scripts::{ + AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, + Pending, WindowCount, +}; +use crate::{Number, RequestRedis, codec::python_value}; + +/// A rate-limit descriptor, `{descriptor:value}`, the hash tag every key of its window shares. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Scope { + pub descriptor: String, + pub value: String, +} + +impl Scope { + fn key(&self, suffix: &str) -> String { + format!("{{{}:{}}}:{suffix}", self.descriptor, self.value) + } + + fn window_key(&self) -> String { + self.key("window") + } + + fn gauge_key(&self) -> String { + self.key("max_parallel_requests") + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum Kind { + Requests, + Tokens, +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct Counter { + pub scope: Scope, + pub kind: Kind, +} + +impl Counter { + fn key(&self) -> String { + self.scope.key(match self.kind { + Kind::Requests => "requests", + Kind::Tokens => "tokens", + }) + } +} + +/// One value of the limiter's keyspace, as both tiers address it. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Entry { + Window(Scope), + Counter(Counter), + Gauge(Scope), +} + +impl Entry { + fn key(&self) -> String { + match self { + Self::Window(scope) => scope.window_key(), + Self::Counter(counter) => counter.key(), + Self::Gauge(scope) => scope.gauge_key(), + } + } +} + +/// What the memory tier holds: counters and mirrored gauge counts, window starts, and the per-slot +/// registry a gauge keeps while Redis is unreachable. +#[derive(Clone, Debug, PartialEq)] +pub enum Local { + Number(Number), + Text(String), + Slots(Vec<(String, f64)>), +} + +/// A `RedisPipelineIncrementOperation`. +#[derive(Clone, Debug, PartialEq)] +pub struct Increment { + pub counter: Counter, + pub amount: Number, + pub ttl: Option, +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error(transparent)] + Redis(#[from] litellm_redis_client::Error), + #[error("rate limiter script reply {0}")] + Reply(&'static str), + #[error("the rate limiter has no Redis to run its scripts on")] + NoRedis, +} + +/// The v3 rate limiter's keyspace: its memory tier, its Lua scripts and the plain counter writes +/// `DualCache` made for it on `main`. The admission logic stays with the caller. +pub struct RateLimits { + redis: Option, + memory: InMemoryCache, + increments: Mutex<()>, + backfill_ttl: Option, + redis_ttl: Duration, +} + +impl RateLimits { + /// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces, and + /// `redis_ttl` its `RedisCache.default_ttl`, which an increment without a TTL takes. + pub fn new( + redis: Option, + backfill_ttl: Option, + redis_ttl: Duration, + ) -> Self { + Self::with_memory( + redis, + backfill_ttl, + redis_ttl, + InMemoryCache::new(None, None), + ) + } + + pub fn with_memory( + redis: Option, + backfill_ttl: Option, + redis_ttl: Duration, + memory: InMemoryCache, + ) -> Self { + Self { + redis, + memory, + increments: Mutex::new(()), + backfill_ttl, + redis_ttl, + } + } + + pub fn has_redis(&self) -> bool { + self.redis.is_some() + } + + /// `_group_keys_by_hash_tag`: indices of `scopes` sharing a cluster slot, in first-seen order, + /// or one group when the client is not a cluster and grouping is not forced. + pub fn slot_groups(&self, scopes: &[Scope], force: bool) -> Vec> { + let by_slot = force || self.redis.as_ref().is_some_and(RedisClient::is_cluster); + if !by_slot { + return vec![(0..scopes.len()).collect()]; + } + scopes + .iter() + .enumerate() + .fold( + Vec::<(Slot, Vec)>::new(), + |mut groups, (index, scope)| { + let slot = Slot::for_key(scope.window_key()); + match groups.iter_mut().find(|(seen, _)| *seen == slot) { + Some((_, members)) => members.push(index), + None => groups.push((slot, vec![index])), + } + groups + }, + ) + .into_iter() + .map(|(_, members)| members) + .collect() + } + + pub fn local(&self, entry: &Entry) -> Result, MemoryError> { + self.memory.get_cache(&entry.key()) + } + + pub fn set_local( + &self, + entry: &Entry, + value: Local, + ttl: Option, + ) -> Result<(), MemoryError> { + self.memory.set_cache(entry.key(), value, ttl).map(drop) + } + + /// `DualCache.async_get_cache`: memory, else Redis written back to memory for `backfill_ttl`, + /// or the store's default. A Redis failure reads as a miss, as `DualCache` logs and swallows it. + pub async fn get( + &self, + entry: &Entry, + backfill_ttl: Option, + ) -> Result, MemoryError> { + let key = entry.key(); + if let Some(local) = self.memory.get_cache(&key)? { + return Ok(Some(local)); + } + let Some(redis) = &self.redis else { + return Ok(None); + }; + let raw = match redis.get::>(&key).await { + Ok(raw) => raw, + Err(error) => { + tracing::error!(%error, "LiteLLM Cache: exception in async_get_cache"); + return Ok(None); + } + }; + let Some(value) = raw + .map(|raw| decode(&raw)) + .transpose() + .unwrap_or_else(|error| { + tracing::error!(%error, "LiteLLM Cache: exception in async_get_cache"); + None + }) + else { + return Ok(None); + }; + self.memory + .set_cache(key, value.clone(), backfill_ttl.or(self.backfill_ttl))?; + Ok(Some(value)) + } + + pub fn sliding_window( + &self, + request: Option<&RequestRedis>, + counters: &[Counter], + now: i64, + window_size: i64, + ) -> Result>, Error> { + Ok(scripts::sliding_window( + &self.pre(request)?, + counters, + now, + window_size, + )) + } + + pub fn read_counters( + &self, + request: Option<&RequestRedis>, + counters: &[Counter], + ) -> Result>>, Error> { + Ok(scripts::read_counters(&self.pre(request)?, counters)) + } + + pub fn check_and_increment( + &self, + request: Option<&RequestRedis>, + counters: &[AtomicCounter], + ) -> Result, Error> { + Ok(scripts::check_and_increment(&self.pre(request)?, counters)) + } + + pub fn guarded_increment( + &self, + request: Option<&RequestRedis>, + increment: &GuardedIncrement, + ) -> Result, Error> { + Ok(scripts::guarded_increment(&self.pre(request)?, increment)) + } + + pub fn parallel_acquire( + &self, + request: Option<&RequestRedis>, + gauges: &[Gauge], + slot_ttl: Duration, + slot_id: &str, + ) -> Result, Error> { + Ok(scripts::parallel_acquire( + &self.pre(request)?, + gauges, + slot_ttl, + slot_id, + )) + } + + pub fn parallel_release( + &self, + request: Option<&RequestRedis>, + gauges: &[Scope], + slot_id: &str, + ) -> Result>, Error> { + Ok(scripts::parallel_release( + &self.pre(request)?, + gauges, + slot_id, + )) + } + + pub fn parallel_count( + &self, + request: Option<&RequestRedis>, + gauges: &[Scope], + slot_ttl: Duration, + ) -> Result>, Error> { + Ok(scripts::parallel_count( + &self.pre(request)?, + gauges, + slot_ttl, + )) + } + + pub fn token_increment( + &self, + request: Option<&RequestRedis>, + increments: &[Increment], + ) -> Result>, Error> { + Ok(scripts::token_increment(&self.pre(request)?, increments)) + } + + /// `TOKEN_INCREMENT_SCRIPT` on the request's post-call batch. A failed script falls back to + /// `increment_many` for these increments once the batch has gone out. + pub fn token_increment_post_call( + self: &Arc, + request: &RequestRedis, + increments: Vec, + ) -> Result<(), Error> { + let pending = scripts::token_increment(&self.post(request)?, &increments); + let store = Arc::clone(self); + request.spawn(async move { + let Err(error) = pending.settled().await else { + return; + }; + tracing::warn!(%error, "TTL preservation failed, falling back to regular pipeline"); + if let Err(error) = store.increment_many(&increments).await { + tracing::warn!(%error, "in-memory increment failed"); + } + }); + Ok(()) + } + + /// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed + /// the slot in memory. + pub fn parallel_release_post_call( + &self, + request: &RequestRedis, + gauges: &[Scope], + slot_id: &str, + ) -> Result<(), Error> { + let pending = scripts::parallel_release(&self.post(request)?, gauges, slot_id); + request.spawn(async move { + if let Err(error) = pending.settled().await { + tracing::warn!( + %error, + "parallel_release_script failed, the slot stays released in memory only" + ); + } + }); + Ok(()) + } + + /// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL. A + /// Redis failure answers with the memory count. + pub async fn increment(&self, increment: &Increment) -> Result { + let local = self.add_local(increment)?; + let Some(redis) = &self.redis else { + return Ok(local); + }; + let key = increment.counter.key(); + let ttl = increment.ttl.unwrap_or(self.redis_ttl); + let total = redis + .incr_by_float(&key, increment.amount.into(), Expiry::KeepExisting(ttl)) + .await + .map_err(Error::from); + Ok(total.map(Number::Float).unwrap_or_else(|error| { + tracing::warn!( + %error, + "Redis async_increment_cache failed, falling back to in-memory result" + ); + local + })) + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline of `INCRBYFLOAT` + /// with `EXPIRE` per increment. + pub async fn increment_many(&self, increments: &[Increment]) -> Result<(), MemoryError> { + increments + .iter() + .try_for_each(|increment| self.add_local(increment).map(drop))?; + let Some(redis) = &self.redis else { + return Ok(()); + }; + let pipeline = increments + .iter() + .fold(redis.pipeline(), |pipeline, increment| { + pipeline.incr_by_float( + &increment.counter.key(), + increment.amount.into(), + increment.ttl, + ) + }); + if let Some(error) = redis.run(&pipeline).await.into_iter().find_map(Result::err) { + tracing::warn!( + %error, + "Redis async_increment_cache_pipeline failed, falling back to in-memory result" + ); + } + Ok(()) + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, the Redis increments on + /// the request's post-call batch. + pub fn increment_many_post_call( + &self, + request: &RequestRedis, + increments: &[Increment], + ) -> Result<(), MemoryError> { + let Some(declared) = self.declare_increments(increments, |client| request.post(client))? + else { + return Ok(()); + }; + request.spawn(async move { + if let Some(error) = settle(declared).await { + tracing::warn!(%error, "deferred Redis increment failed"); + } + }); + Ok(()) + } + + fn declare_increments( + &self, + increments: &[Increment], + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + ) -> Result>>, MemoryError> { + increments + .iter() + .try_for_each(|increment| self.add_local(increment).map(drop))?; + let Some(redis) = &self.redis else { + return Ok(None); + }; + let batch = batch_for(redis); + Ok(Some( + increments + .iter() + .map(|increment| scripts::counter_increment(&batch, increment)) + .collect(), + )) + } + + /// `InMemoryCache.increment_cache`: `(current or 0) + amount` under one lock. + fn add_local(&self, increment: &Increment) -> Result { + let _guard = self + .increments + .lock() + .unwrap_or_else(PoisonError::into_inner); + let key = increment.counter.key(); + let current = match self.memory.get_cache(&key)? { + None => Number::Int(0), + Some(Local::Number(current)) => current, + Some(_) => return Err(MemoryError::InvalidEntry(None)), + }; + let total = current + increment.amount; + self.memory + .set_cache(key, Local::Number(total), increment.ttl)?; + Ok(total) + } + + fn pre(&self, request: Option<&RequestRedis>) -> Result { + let redis = self.redis.as_ref().ok_or(Error::NoRedis)?; + Ok(request.map_or_else(|| redis.batch(), |request| request.pre(redis))) + } + + fn post(&self, request: &RequestRedis) -> Result { + let redis = self.redis.as_ref().ok_or(Error::NoRedis)?; + Ok(request.post(redis)) + } +} + +/// Every increment's outcome; the first failure when there is one. `flush` sends the batch first. +async fn settle(declared: Vec>) -> Option { + let outcomes = join_all(declared.into_iter().map(Declared::settled)).await; + outcomes.into_iter().find_map(Result::err).map(Error::from) +} + +fn decode(raw: &[u8]) -> Result { + match python_value(raw)? { + Value::Str(text) => Ok(Local::Text(text)), + value => Number::from_value(&value) + .map(Local::Number) + .ok_or(crate::DecodeError::Shape( + "a rate-limit value is a number or a string", + )), + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/lua.rs b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs new file mode 100644 index 00000000000..96451212d35 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs @@ -0,0 +1,260 @@ +//! The limiter's Lua, byte for byte as `main`'s Python registers it, so both share one script SHA. + +pub(super) const BATCH_RATE_LIMITER: &str = r#" +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) + +-- Process each window/counter pair +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local increment_value = 1 + + -- Check if window exists and is valid + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + -- Reset window and counter + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment_value) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) -- window_start + table.insert(results, increment_value) -- counter + else + local counter = redis.call('INCR', counter_key) + -- This happens when window_key exists but counter_key doesn't (e.g., tokens key + -- created after requests key when both share the same window_key) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) -- window_start + table.insert(results, counter) -- counter + end +end + +return results +"#; + +pub(super) const BATCH_COUNTER_READ: &str = r#"return redis.call('MGET', unpack(KEYS))"#; + +pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#" +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) +-- +-- Return on success: +-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local descriptor_count = #KEYS / 2 +local reset_windows = {} + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + local blocked + if increment > 0 then + blocked = current_counter + increment > limit + else + blocked = current_counter >= limit + end + if blocked then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter, window_start } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_expired = descriptor_state[i][1] + local active_window_start + + if window_expired then + active_window_start = now + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + active_window_start = tonumber(descriptor_state[i][3]) + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end + table.insert(results, active_window_start) +end + +return results +"#; + +pub(super) const WINDOW_GUARDED_TOKEN_INCREMENT: &str = r#" +local results = {} +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local arg_base = ((i - 1) / 2) * 3 + 1 + local expected_window_start = ARGV[arg_base] + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local active_window_start = redis.call('GET', window_key) + + if active_window_start and active_window_start == expected_window_start then + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, 1) + table.insert(results, new_counter) + else + table.insert(results, 0) + table.insert(results, tonumber(redis.call('GET', counter_key) or 0)) + end +end +return results +"#; + +pub(super) const PARALLEL_ACQUIRE: &str = r#" +-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. +-- Each gauge key is a sorted set of per-request slot ids scored by acquire +-- time (Redis server clock). In-flight requests are counted by ZCARD after +-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM +-- counters the gauge is never reset while requests are in flight, a +-- rejected request never occupies a slot, and a slot leaked by a crashed +-- worker self-heals after the slot TTL even under continuous traffic. +-- +-- KEYS: one gauge zset key per descriptor. +-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id). +-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +for i = 1, #KEYS do + local limit = tonumber(ARGV[(i - 1) * 3 + 1]) + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl) + local in_flight = redis.call('ZCARD', KEYS[i]) + if in_flight + 1 > limit then + return { 1, i, in_flight, limit } + end +end +local results = { 0 } +for i = 1, #KEYS do + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + local slot_id = ARGV[(i - 1) * 3 + 3] + redis.call('ZADD', KEYS[i], now, slot_id) + redis.call('EXPIRE', KEYS[i], slot_ttl) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const PARALLEL_RELEASE: &str = r#" +-- Release one slot per gauge key by removing this request's slot id. +-- ZREM of an absent member (or key) is a no-op, so a release without a +-- matching acquire (proxy-side rejection, double-fired callback, slot +-- already expired) can never free a slot owned by another request. +-- KEYS: gauge zset keys. ARGV: per-key slot_id. +-- Returns the remaining in-flight count per key. +local results = {} +for i = 1, #KEYS do + redis.call('ZREM', KEYS[i], ARGV[i]) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const PARALLEL_COUNT: &str = r#" +-- Read the current in-flight count per gauge key (prunes expired slots +-- first so leaked slots do not inflate the reading). +-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local results = {} +for i = 1, #KEYS do + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i])) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +"#; + +pub(super) const TOKEN_INCREMENT: &str = r#" +local results = {} + +-- Process each key/increment_value/ttl triplet +for i = 1, #KEYS do + local key = KEYS[i] + local increment_value = tonumber(ARGV[i * 2 - 1]) + local ttl_seconds = tonumber(ARGV[i * 2]) + + -- Increment the value + local new_value = redis.call('INCRBYFLOAT', key, increment_value) + + -- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL + -- ttl_seconds can be 0 (no TTL) or positive (set TTL) + if ttl_seconds and ttl_seconds > 0 then + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl_seconds) + end + end + + table.insert(results, new_value) +end + +return results +"#; diff --git a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs new file mode 100644 index 00000000000..0346dccbca1 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs @@ -0,0 +1,391 @@ +use std::{sync::LazyLock, time::Duration}; + +use litellm_python_compat::repr::float_repr; +use litellm_redis_client::{Declared, RedisBatch, Script, Value}; + +use super::{Counter, Error, Increment, Scope, lua}; +use crate::Number; + +static BATCH_RATE_LIMITER: LazyLock