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