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 <noreply@anthropic.com>

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.
This commit is contained in:
Yujong Lee 2026-11-03 15:49:24 +09:00
parent f2e75c8ebf
commit 14d50599fa
28 changed files with 5341 additions and 3903 deletions

View file

@ -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",

View file

@ -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;

View file

@ -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<RateLimits>);
/// `(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<Self> {
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<Self> {
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<Kind> {
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<Self> {
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<Self> {
let (counter, amount, ttl): (PyCounter, PyNumber, Option<u64>) = 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<Self> {
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<Self> {
if let Ok(number) = value.extract::<PyNumber>() {
return Ok(Self(Local::Number(number.0)));
}
if let Ok(text) = value.cast::<PyString>() {
return Ok(Self(Local::Text(text.to_str()?.to_owned())));
}
if let Ok(slots) = value.cast::<PyDict>() {
return slots
.iter()
.map(|(slot, acquired_at)| Ok((slot.extract()?, acquired_at.extract()?)))
.collect::<PyResult<_>>()
.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<Self::Output> {
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<WindowCount>),
Values(Vec<Option<String>>),
Atomic(AtomicReply),
Guarded(GuardedReply),
Acquire(AcquireReply),
Counts(Vec<i64>),
Totals(Vec<f64>),
}
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<Self::Output> {
match self {
Self::WindowCounts(counts) => counts
.into_iter()
.map(
|WindowCount {
window_start,
count,
}| (window_start, count),
)
.collect::<Vec<_>>()
.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::<Vec<_>>(),
)
.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::<i64>::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<Box<dyn Future<Output = Result<Reply, Error>> + 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<Option<Resolve>>,
call_type: &'static str,
}
impl NativeRateLimitPending {
fn new<T: Send + 'static>(
pending: Result<Pending<T>, Error>,
reply: fn(T) -> Reply,
call_type: &'static str,
) -> PyResult<Self> {
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<Bound<'py, PyAny>> {
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<PyScope>) -> Vec<Scope> {
scopes.into_iter().map(|scope| scope.0).collect()
}
fn counters(counters: Vec<PyCounter>) -> Vec<Counter> {
counters.into_iter().map(|counter| counter.0).collect()
}
fn increments(increments: Vec<PyIncrement>) -> Vec<Increment> {
increments
.into_iter()
.map(|increment| increment.0)
.collect()
}
fn request_of(
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> Option<litellm_redis_store::RequestRedis> {
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<PyRef<'_, NativeRedisClient>>,
backfill_ttl_seconds: Option<f64>,
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<PyScope>, force: bool) -> Vec<Vec<usize>> {
self.0.slot_groups(&scopes(scopes_), force)
}
fn local(&self, entries: Vec<PyEntry>) -> PyResult<Vec<Option<PyLocal>>> {
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<f64>) -> 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<f64>,
) -> PyResult<Bound<'py, PyAny>> {
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<PyCounter>,
now: i64,
window_size: i64,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
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<PyCounter>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
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<PyAtomicCounter>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
let request = request_of(request);
let counters_ = counters_
.into_iter()
.map(|counter| counter.0)
.collect::<Vec<_>>();
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<u64>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
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<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
let request = request_of(request);
let gauges = gauges
.into_iter()
.map(|(scope, limit)| Gauge {
scope: scope.0,
limit,
})
.collect::<Vec<_>>();
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<PyScope>,
slot_id: &str,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
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<PyScope>,
slot_ttl_seconds: u64,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
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<PyIncrement>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativeRateLimitPending> {
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<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
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<PyScope>,
slot_id: String,
) -> PyResult<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
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<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
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)
})
}
}

View file

@ -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())
}

View file

@ -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;

View file

@ -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

View file

@ -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;

View file

@ -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<Duration>,
}
#[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<RedisClient>,
memory: InMemoryCache<Local>,
increments: Mutex<()>,
backfill_ttl: Option<Duration>,
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<RedisClient>,
backfill_ttl: Option<Duration>,
redis_ttl: Duration,
) -> Self {
Self::with_memory(
redis,
backfill_ttl,
redis_ttl,
InMemoryCache::new(None, None),
)
}
pub fn with_memory(
redis: Option<RedisClient>,
backfill_ttl: Option<Duration>,
redis_ttl: Duration,
memory: InMemoryCache<Local>,
) -> 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<Vec<usize>> {
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<usize>)>::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<Option<Local>, MemoryError> {
self.memory.get_cache(&entry.key())
}
pub fn set_local(
&self,
entry: &Entry,
value: Local,
ttl: Option<Duration>,
) -> 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<Duration>,
) -> Result<Option<Local>, 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::<Vec<u8>>(&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<Pending<Vec<WindowCount>>, Error> {
Ok(scripts::sliding_window(
&self.pre(request)?,
counters,
now,
window_size,
))
}
pub fn read_counters(
&self,
request: Option<&RequestRedis>,
counters: &[Counter],
) -> Result<Pending<Vec<Option<String>>>, Error> {
Ok(scripts::read_counters(&self.pre(request)?, counters))
}
pub fn check_and_increment(
&self,
request: Option<&RequestRedis>,
counters: &[AtomicCounter],
) -> Result<Pending<AtomicReply>, Error> {
Ok(scripts::check_and_increment(&self.pre(request)?, counters))
}
pub fn guarded_increment(
&self,
request: Option<&RequestRedis>,
increment: &GuardedIncrement,
) -> Result<Pending<GuardedReply>, 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<Pending<AcquireReply>, 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<Pending<Vec<i64>>, 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<Pending<Vec<i64>>, Error> {
Ok(scripts::parallel_count(
&self.pre(request)?,
gauges,
slot_ttl,
))
}
pub fn token_increment(
&self,
request: Option<&RequestRedis>,
increments: &[Increment],
) -> Result<Pending<Vec<f64>>, 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<Self>,
request: &RequestRedis,
increments: Vec<Increment>,
) -> 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<Number, MemoryError> {
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<Option<Vec<Declared<f64>>>, 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<Number, MemoryError> {
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<RedisBatch, Error> {
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<RedisBatch, Error> {
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<Declared<f64>>) -> Option<Error> {
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<Local, crate::DecodeError> {
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",
)),
}
}

View file

@ -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
"#;

View file

@ -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<Script> =
LazyLock::new(|| Script::new(lua::BATCH_RATE_LIMITER));
static BATCH_COUNTER_READ: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::BATCH_COUNTER_READ));
static CHECK_AND_INCREMENT_BY_N: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::CHECK_AND_INCREMENT_BY_N));
static WINDOW_GUARDED_TOKEN_INCREMENT: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::WINDOW_GUARDED_TOKEN_INCREMENT));
static PARALLEL_ACQUIRE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_ACQUIRE));
static PARALLEL_RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_RELEASE));
static PARALLEL_COUNT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_COUNT));
static TOKEN_INCREMENT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::TOKEN_INCREMENT));
/// One script call declared on a batch, decoded into its typed reply once the batch has run.
#[must_use]
pub struct Pending<T> {
reply: Declared<Value>,
decode: fn(Value) -> Result<T, Error>,
}
impl<T> Pending<T> {
fn new(reply: Declared<Value>, decode: fn(Value) -> Result<T, Error>) -> Self {
Self { reply, decode }
}
/// Flushes the batch when this op has not gone out yet.
pub async fn resolve(self) -> Result<T, Error> {
(self.decode)(self.reply.resolve().await?)
}
/// Waits for whoever flushes the batch.
pub async fn settled(self) -> Result<T, Error> {
(self.decode)(self.reply.settled().await?)
}
}
/// `BATCH_RATE_LIMITER_SCRIPT`'s pair for one counter: the window start it saw or opened and the
/// counter after its `+1`.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WindowCount {
pub window_start: String,
pub count: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct AtomicCounter {
pub counter: Counter,
pub limit: i64,
pub increment: i64,
pub ttl_seconds: i64,
pub window_size_seconds: i64,
}
/// `CHECK_AND_INCREMENT_BY_N_SCRIPT`'s reply. `index` is 0-based where the Lua's is 1-based.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AtomicReply {
Applied(Vec<Applied>),
OverLimit {
index: usize,
current: i64,
limit: i64,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Applied {
pub count: i64,
pub window_start: i64,
}
#[derive(Clone, Debug, PartialEq)]
pub struct GuardedIncrement {
pub counter: Counter,
pub expected_window_start: String,
pub amount: Number,
pub ttl: Option<Duration>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct GuardedReply {
pub applied: bool,
pub count: i64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Gauge {
pub scope: Scope,
pub limit: i64,
}
/// `PARALLEL_ACQUIRE_SCRIPT`'s reply. `index` is 0-based where the Lua's is 1-based.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum AcquireReply {
Acquired(Vec<i64>),
OverLimit {
index: usize,
in_flight: i64,
limit: i64,
},
}
pub(super) fn sliding_window(
batch: &RedisBatch,
counters: &[Counter],
now: i64,
window_size: i64,
) -> Pending<Vec<WindowCount>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
Pending::new(
eval(
batch,
&BATCH_RATE_LIMITER,
&keys,
vec![now.to_string(), window_size.to_string()],
),
window_counts,
)
}
/// `BATCH_COUNTER_READ_SCRIPT` over each counter's window and counter keys, in that order.
pub(super) fn read_counters(
batch: &RedisBatch,
counters: &[Counter],
) -> Pending<Vec<Option<String>>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
Pending::new(
eval(batch, &BATCH_COUNTER_READ, &keys, Vec::new()),
optional_texts,
)
}
pub(super) fn check_and_increment(
batch: &RedisBatch,
counters: &[AtomicCounter],
) -> Pending<AtomicReply> {
let keys = counters
.iter()
.flat_map(|counter| [counter.counter.scope.window_key(), counter.counter.key()])
.collect::<Vec<_>>();
let args = counters
.iter()
.flat_map(|counter| {
[
counter.limit,
counter.increment,
counter.ttl_seconds,
counter.window_size_seconds,
]
.map(|arg| arg.to_string())
})
.collect();
Pending::new(
eval(batch, &CHECK_AND_INCREMENT_BY_N, &keys, args),
atomic_reply,
)
}
pub(super) fn guarded_increment(
batch: &RedisBatch,
increment: &GuardedIncrement,
) -> Pending<GuardedReply> {
let keys = [
increment.counter.scope.window_key(),
increment.counter.key(),
];
let args = vec![
increment.expected_window_start.clone(),
python_arg(increment.amount),
ttl_arg(increment.ttl),
];
Pending::new(
eval(batch, &WINDOW_GUARDED_TOKEN_INCREMENT, &keys, args),
guarded_reply,
)
}
pub(super) fn parallel_acquire(
batch: &RedisBatch,
gauges: &[Gauge],
slot_ttl: Duration,
slot_id: &str,
) -> Pending<AcquireReply> {
let keys = gauges
.iter()
.map(|gauge| gauge.scope.gauge_key())
.collect::<Vec<_>>();
let args = gauges
.iter()
.flat_map(|gauge| {
[
gauge.limit.to_string(),
slot_ttl.as_secs().to_string(),
slot_id.to_owned(),
]
})
.collect();
Pending::new(eval(batch, &PARALLEL_ACQUIRE, &keys, args), acquire_reply)
}
pub(super) fn parallel_release(
batch: &RedisBatch,
gauges: &[Scope],
slot_id: &str,
) -> Pending<Vec<i64>> {
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
let args = gauges.iter().map(|_| slot_id.to_owned()).collect();
Pending::new(eval(batch, &PARALLEL_RELEASE, &keys, args), integers)
}
pub(super) fn parallel_count(
batch: &RedisBatch,
gauges: &[Scope],
slot_ttl: Duration,
) -> Pending<Vec<i64>> {
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
let args = gauges
.iter()
.map(|_| slot_ttl.as_secs().to_string())
.collect();
Pending::new(eval(batch, &PARALLEL_COUNT, &keys, args), integers)
}
pub(super) fn token_increment(batch: &RedisBatch, increments: &[Increment]) -> Pending<Vec<f64>> {
let keys = increments
.iter()
.map(|increment| increment.counter.key())
.collect::<Vec<_>>();
let args = increments
.iter()
.flat_map(|increment| [python_arg(increment.amount), ttl_arg(increment.ttl)])
.collect();
Pending::new(eval(batch, &TOKEN_INCREMENT, &keys, args), floats)
}
/// `RedisBatch`'s `_Increment`: `INCRBYFLOAT`, then `EXPIRE` when the increment carries a TTL.
pub(super) fn counter_increment(batch: &RedisBatch, increment: &Increment) -> Declared<f64> {
batch.incr_by_float(
&increment.counter.key(),
increment.amount.into(),
increment.ttl,
)
}
fn eval(
batch: &RedisBatch,
script: &Script,
keys: &[String],
args: Vec<String>,
) -> Declared<Value> {
let keys = keys.iter().map(String::as_str).collect::<Vec<_>>();
batch.eval(script, &keys, args)
}
/// The bytes redis-py sends for a Python `int` or `float` argument.
fn python_arg(amount: Number) -> String {
match amount {
Number::Int(value) => value.to_string(),
Number::Float(value) => float_repr(value),
}
}
/// `ttl or 0`: the scripts read a zero TTL as "leave the expiry alone".
fn ttl_arg(ttl: Option<Duration>) -> String {
ttl.map_or(0, |ttl| ttl.as_secs()).to_string()
}
fn items(reply: Value) -> Result<Vec<Value>, Error> {
match reply {
Value::Array(items) => Ok(items),
_ => Err(Error::Reply("is not a list")),
}
}
fn integer(value: &Value) -> Result<i64, Error> {
match value {
Value::Int(value) => Ok(*value),
_ => Err(Error::Reply("holds a non-integer where a count belongs")),
}
}
fn text(value: &Value) -> Result<Option<String>, Error> {
match value {
Value::Nil => Ok(None),
Value::BulkString(bytes) => String::from_utf8(bytes.clone())
.map(Some)
.map_err(|_| Error::Reply("holds a value that is not UTF-8")),
Value::SimpleString(text) => Ok(Some(text.clone())),
Value::Int(value) => Ok(Some(value.to_string())),
_ => Err(Error::Reply("holds a nested value")),
}
}
fn float(reply: Value) -> Result<f64, Error> {
text(&reply)?
.and_then(|text| text.parse().ok())
.ok_or(Error::Reply("is not a float"))
}
fn index(raw: &Value) -> Result<usize, Error> {
usize::try_from(integer(raw)? - 1).map_err(|_| Error::Reply("names an index below 1"))
}
fn window_counts(reply: Value) -> Result<Vec<WindowCount>, Error> {
items(reply)?
.chunks(2)
.map(|pair| match pair {
[window_start, count] => Ok(WindowCount {
window_start: text(window_start)?
.ok_or(Error::Reply("holds a nil window start"))?,
count: integer(count)?,
}),
_ => Err(Error::Reply("has an odd length")),
})
.collect()
}
fn optional_texts(reply: Value) -> Result<Vec<Option<String>>, Error> {
items(reply)?.iter().map(text).collect()
}
fn integers(reply: Value) -> Result<Vec<i64>, Error> {
items(reply)?.iter().map(integer).collect()
}
fn floats(reply: Value) -> Result<Vec<f64>, Error> {
items(reply)?.into_iter().map(float).collect()
}
fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
let items = items(reply)?;
match items.split_first() {
Some((Value::Int(0), rest)) => rest
.chunks(2)
.map(|pair| match pair {
[count, window_start] => Ok(Applied {
count: integer(count)?,
window_start: integer(window_start)?,
}),
_ => Err(Error::Reply("has an odd length")),
})
.collect::<Result<_, _>>()
.map(AtomicReply::Applied),
Some((Value::Int(1), [index_value, current, limit])) => Ok(AtomicReply::OverLimit {
index: index(index_value)?,
current: integer(current)?,
limit: integer(limit)?,
}),
_ => Err(Error::Reply("has no status")),
}
}
fn guarded_reply(reply: Value) -> Result<GuardedReply, Error> {
match items(reply)?.as_slice() {
[applied, count] => Ok(GuardedReply {
applied: integer(applied)? == 1,
count: integer(count)?,
}),
_ => Err(Error::Reply("is not an (applied, count) pair")),
}
}
fn acquire_reply(reply: Value) -> Result<AcquireReply, Error> {
let items = items(reply)?;
match items.split_first() {
Some((Value::Int(0), rest)) => rest
.iter()
.map(integer)
.collect::<Result<_, _>>()
.map(AcquireReply::Acquired),
Some((Value::Int(1), [index_value, in_flight, limit])) => Ok(AcquireReply::OverLimit {
index: index(index_value)?,
in_flight: integer(in_flight)?,
limit: integer(limit)?,
}),
_ => Err(Error::Reply("has no status")),
}
}

View file

@ -0,0 +1,902 @@
mod support;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Operation, RedisClient, Script, Ttl};
use litellm_redis_store::{
Number, RequestRedis,
rate_limits::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement,
GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount,
},
router_cooldowns::{POLICY as COOLDOWN_POLICY, RouterCooldowns},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, connect_observed, within};
use tokio::runtime::Handle;
const SLOT_TTL: Duration = Duration::from_secs(3600);
struct WireScope {
descriptor: &'static str,
value: &'static str,
window: &'static str,
requests: &'static str,
tokens: &'static str,
gauge: &'static str,
}
const SCOPES: [WireScope; 4] = [
WireScope {
descriptor: "api_key",
value: "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
window: "{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:window",
requests: "{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:requests",
tokens: "{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:tokens",
gauge: "{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:max_parallel_requests",
},
WireScope {
descriptor: "team_member",
value: "team-1:user@example.com",
window: "{team_member:team-1:user@example.com}:window",
requests: "{team_member:team-1:user@example.com}:requests",
tokens: "{team_member:team-1:user@example.com}:tokens",
gauge: "{team_member:team-1:user@example.com}:max_parallel_requests",
},
WireScope {
descriptor: "model_per_key",
value: "sk-hash:anthropic/claude-sonnet-5-5",
window: "{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:window",
requests: "{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:requests",
tokens: "{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:tokens",
gauge: "{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:max_parallel_requests",
},
WireScope {
descriptor: "tag",
value: "tier}é",
window: "{tag:tier}é}:window",
requests: "{tag:tier}é}:requests",
tokens: "{tag:tier}é}:tokens",
gauge: "{tag:tier}é}:max_parallel_requests",
},
];
const SCRIPTS: [(&str, &str); 8] = [
(
"batch_rate_limiter",
"1e1d6e80e9488215da3f8d4aefae5ee3fbe71663",
),
(
"batch_counter_read",
"b272298fae5e83182e7719c6be8d1ef76a62465f",
),
(
"check_and_increment_by_n",
"cbf6fb1b9184f1834b44b05bdeda3806fbc2da41",
),
(
"window_guarded_token_increment",
"fec058661fefb81df4e5da0cc045df5f4a3f0a00",
),
(
"parallel_acquire",
"86c5a792cfc23df2842c413908b3eb3f2b811ac3",
),
(
"parallel_release",
"cb232fec078c0b66691670a16b84079139035b4e",
),
("parallel_count", "9ae6ae0ca54e4c1930356965275e1d442b2f25ba"),
(
"token_increment",
"bacdccafbf0cb85105a45154294f340c6275fe3f",
),
];
fn scope(descriptor: &str, value: &str) -> Scope {
Scope {
descriptor: descriptor.to_owned(),
value: value.to_owned(),
}
}
fn counter(scope: &Scope, kind: Kind) -> Counter {
Counter {
scope: scope.clone(),
kind,
}
}
fn atomic(scope: &Scope, kind: Kind, limit: i64, increment: i64) -> AtomicCounter {
AtomicCounter {
counter: counter(scope, kind),
limit,
increment,
ttl_seconds: 60,
window_size_seconds: 60,
}
}
/// `BaseCache.default_ttl`, what `RedisCache` gives an increment that names no TTL.
const REDIS_TTL: Duration = Duration::from_secs(60);
fn store(deployment: &RedisDeployment) -> Arc<RateLimits> {
Arc::new(RateLimits::new(Some(connect(deployment)), None, REDIS_TTL))
}
async fn raw_text(raw: &RedisClient, key: &str) -> Option<String> {
raw.get(key).await.unwrap()
}
async fn zcard(raw: &RedisClient, key: &str) -> i64 {
raw.eval(
&Script::new("return redis.call('ZCARD', KEYS[1])"),
&[key],
Vec::<String>::new(),
)
.await
.unwrap()
}
async fn redis_now(raw: &RedisClient) -> i64 {
let time: Vec<String> = raw
.eval(
&Script::new("return redis.call('TIME')"),
&["any"],
Vec::<String>::new(),
)
.await
.unwrap();
time[0].parse().unwrap()
}
async fn cli(deployment: &RedisDeployment, args: &[&str]) -> Vec<String> {
match deployment {
RedisDeployment::Standalone(server) => vec![server.cli(args).await.unwrap()],
RedisDeployment::Cluster(cluster) => cluster.cli_outputs(args).await.unwrap(),
}
}
#[rstest]
#[tokio::test]
async fn every_script_runs_mains_bytes_under_mains_keys(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let now = redis_now(&raw).await;
for wire in &SCOPES {
let key = scope(wire.descriptor, wire.value);
let requests = counter(&key, Kind::Requests);
let tokens = counter(&key, Kind::Tokens);
limits
.sliding_window(None, std::slice::from_ref(&requests), now, 60)
.unwrap()
.resolve()
.await
.unwrap();
limits
.check_and_increment(None, &[atomic(&key, Kind::Tokens, 1000, 40)])
.unwrap()
.resolve()
.await
.unwrap();
limits
.token_increment(
None,
&[Increment {
counter: tokens.clone(),
amount: Number::Float(2.5),
ttl: Some(Duration::from_secs(60)),
}],
)
.unwrap()
.resolve()
.await
.unwrap();
limits
.parallel_acquire(
None,
&[Gauge {
scope: key.clone(),
limit: 5,
}],
SLOT_TTL,
"slot-1",
)
.unwrap()
.resolve()
.await
.unwrap();
assert_eq!(raw_text(&raw, wire.window).await, Some(now.to_string()));
assert_eq!(raw_text(&raw, wire.requests).await, Some("1".to_owned()));
assert_eq!(raw_text(&raw, wire.tokens).await, Some("42.5".to_owned()));
assert_eq!(zcard(&raw, wire.gauge).await, 1);
assert_eq!(
limits
.read_counters(None, &[requests, tokens])
.unwrap()
.resolve()
.await
.unwrap(),
vec![
Some(now.to_string()),
Some("1".to_owned()),
Some(now.to_string()),
Some("42.5".to_owned()),
]
);
}
let guarded = scope("api_key", "sha-check");
limits
.guarded_increment(
None,
&GuardedIncrement {
counter: counter(&guarded, Kind::Tokens),
expected_window_start: "0".to_owned(),
amount: Number::Int(1),
ttl: None,
},
)
.unwrap()
.resolve()
.await
.unwrap();
for pending in [
limits.parallel_release(None, std::slice::from_ref(&guarded), "slot-1"),
limits.parallel_count(None, std::slice::from_ref(&guarded), SLOT_TTL),
] {
pending.unwrap().resolve().await.unwrap();
}
for (name, sha) in SCRIPTS {
let cached = cli(&deployment, &["SCRIPT", "EXISTS", sha]).await;
assert!(
cached.iter().any(|reply| reply.trim() == "1"),
"{name} ran under a different SHA than main's {sha}"
);
}
}
#[rstest]
#[case::epoch("1730000000", Number::Int(1730000000))]
#[case::int("7", Number::Int(7))]
#[case::float("7.5", Number::Float(7.5))]
#[case::negative("-3", Number::Int(-3))]
#[tokio::test]
async fn reads_values_python_wrote_and_backfills_memory(
#[case] raw_value: &str,
#[case] expected: Number,
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let wire = &SCOPES[0];
raw.set(wire.tokens, raw_value, Duration::from_secs(60))
.await
.unwrap();
let entry = Entry::Counter(counter(&scope(wire.descriptor, wire.value), Kind::Tokens));
assert_eq!(
limits
.get(&entry, Some(Duration::from_secs(30)))
.await
.unwrap(),
Some(Local::Number(expected))
);
assert_eq!(limits.local(&entry).unwrap(), Some(Local::Number(expected)));
}
#[rstest]
#[tokio::test]
async fn sliding_window_opens_counts_and_resets_with_its_sibling(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let key = scope("api_key", "k1");
let requests = counter(&key, Kind::Requests);
let run = |now: i64| {
let pending = limits
.sliding_window(None, std::slice::from_ref(&requests), now, 60)
.unwrap();
async move { pending.resolve().await.unwrap() }
};
assert_eq!(
run(1_000).await,
vec![WindowCount {
window_start: "1000".to_owned(),
count: 1
}]
);
assert_eq!(run(1_059).await[0].count, 2);
raw.set("{api_key:k1}:tokens", "500", Duration::from_secs(60))
.await
.unwrap();
let Ttl::Expires(ttl) = raw.ttl("{api_key:k1}:requests").await.unwrap() else {
panic!("the counter has no TTL");
};
assert!(within(ttl, 60), "{ttl:?}");
assert_eq!(
run(1_060).await,
vec![WindowCount {
window_start: "1060".to_owned(),
count: 1
}]
);
assert_eq!(raw_text(&raw, "{api_key:k1}:tokens").await, None);
}
#[rstest]
#[tokio::test]
async fn check_and_increment_is_all_or_nothing(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let key = scope("model_per_key", "k1:gpt");
let now = redis_now(&raw).await;
let both = |requests: i64, tokens: i64| {
[
atomic(&key, Kind::Requests, 2, requests),
atomic(&key, Kind::Tokens, 100, tokens),
]
};
let first = limits
.check_and_increment(None, &both(1, 60))
.unwrap()
.resolve()
.await
.unwrap();
let AtomicReply::Applied(applied) = first else {
panic!("{first:?}");
};
assert_eq!(applied.iter().map(|a| a.count).collect::<Vec<_>>(), [1, 60]);
assert!(applied.iter().all(|a| (a.window_start - now).abs() <= 2));
assert_eq!(
limits
.check_and_increment(None, &both(1, 41))
.unwrap()
.resolve()
.await
.unwrap(),
AtomicReply::OverLimit {
index: 1,
current: 60,
limit: 100
}
);
assert_eq!(
raw_text(&raw, "{model_per_key:k1:gpt}:requests").await,
Some("1".to_owned())
);
assert_eq!(
raw_text(&raw, "{model_per_key:k1:gpt}:tokens").await,
Some("60".to_owned())
);
let at_limit = limits
.check_and_increment(None, &both(1, 40))
.unwrap()
.resolve()
.await
.unwrap();
assert!(
matches!(at_limit, AtomicReply::Applied(ref applied) if applied[0].count == 2 && applied[1].count == 100)
);
assert_eq!(
limits
.check_and_increment(None, &both(0, 0))
.unwrap()
.resolve()
.await
.unwrap(),
AtomicReply::OverLimit {
index: 0,
current: 2,
limit: 2
}
);
}
#[rstest]
#[tokio::test]
async fn guarded_increment_only_settles_the_window_it_was_reserved_in(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let key = scope("team", "t1");
let tokens = counter(&key, Kind::Tokens);
let reply = limits
.check_and_increment(None, &[atomic(&key, Kind::Tokens, 1000, 300)])
.unwrap()
.resolve()
.await
.unwrap();
let AtomicReply::Applied(applied) = reply else {
panic!("{reply:?}");
};
let Applied { window_start, .. } = applied[0];
let guarded = |expected_window_start: String, amount: i64| {
let pending = limits
.guarded_increment(
None,
&GuardedIncrement {
counter: tokens.clone(),
expected_window_start,
amount: Number::Int(amount),
ttl: Some(Duration::from_secs(60)),
},
)
.unwrap();
async move { pending.resolve().await.unwrap() }
};
assert_eq!(
guarded(window_start.to_string(), -120).await,
GuardedReply {
applied: true,
count: 180
}
);
assert_eq!(
guarded((window_start - 60).to_string(), -180).await,
GuardedReply {
applied: false,
count: 180
}
);
assert_eq!(
raw_text(&raw, "{team:t1}:tokens").await,
Some("180".to_owned())
);
}
#[rstest]
#[tokio::test]
async fn gauge_slots_are_acquired_released_by_id_and_expire(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let limits = store(&deployment);
let key = scope("api_key", "k1");
let gauges = [Gauge {
scope: key.clone(),
limit: 2,
}];
let acquire = |slot_id: &'static str, slot_ttl: Duration| {
let pending = limits
.parallel_acquire(None, &gauges, slot_ttl, slot_id)
.unwrap();
async move { pending.resolve().await.unwrap() }
};
let release = |slot_id: &'static str| {
let pending = limits
.parallel_release(None, std::slice::from_ref(&key), slot_id)
.unwrap();
async move { pending.resolve().await.unwrap() }
};
let count = |slot_ttl: Duration| {
let pending = limits
.parallel_count(None, std::slice::from_ref(&key), slot_ttl)
.unwrap();
async move { pending.resolve().await.unwrap() }
};
assert_eq!(
acquire("a", SLOT_TTL).await,
AcquireReply::Acquired(vec![1])
);
assert_eq!(
acquire("b", SLOT_TTL).await,
AcquireReply::Acquired(vec![2])
);
assert_eq!(
acquire("c", SLOT_TTL).await,
AcquireReply::OverLimit {
index: 0,
in_flight: 2,
limit: 2
}
);
assert_eq!(release("never-acquired").await, vec![2]);
assert_eq!(release("a").await, vec![1]);
assert_eq!(release("a").await, vec![1]);
assert_eq!(count(SLOT_TTL).await, vec![1]);
tokio::time::sleep(Duration::from_millis(2_100)).await;
assert_eq!(count(Duration::from_secs(1)).await, vec![0]);
assert_eq!(
acquire("d", SLOT_TTL).await,
AcquireReply::Acquired(vec![1])
);
}
#[rstest]
#[tokio::test]
async fn token_increment_sets_a_ttl_only_on_a_key_without_one(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let key = scope("user", "u1");
let increment = |amount: Number| Increment {
counter: counter(&key, Kind::Tokens),
amount,
ttl: Some(Duration::from_secs(60)),
};
raw.set("{user:u1}:tokens", "10", Duration::from_secs(600))
.await
.unwrap();
let totals = limits
.token_increment(
None,
&[increment(Number::Float(0.5)), increment(Number::Int(-3))],
)
.unwrap()
.resolve()
.await
.unwrap();
assert_eq!(totals, vec![10.5, 7.5]);
let Ttl::Expires(ttl) = raw.ttl("{user:u1}:tokens").await.unwrap() else {
panic!("the counter lost its TTL");
};
assert!(within(ttl, 600), "{ttl:?}");
}
#[rstest]
#[tokio::test]
async fn slot_groups_follow_the_cluster_slot_of_each_hash_tag(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let limits = store(&deployment);
let scopes = [
scope("api_key", "a"),
scope("team", "t"),
scope("api_key", "a"),
scope("user", "u"),
];
let slot = |key: String| {
let deployment = &deployment;
async move {
match deployment {
RedisDeployment::Cluster(cluster) => cluster
.cli_outputs(&["CLUSTER", "KEYSLOT", &key])
.await
.unwrap()[0]
.trim()
.to_owned(),
RedisDeployment::Standalone(_) => key,
}
}
};
let mut expected: Vec<(String, Vec<usize>)> = Vec::new();
for (index, scope) in scopes.iter().enumerate() {
let slot = slot(format!("{{{}:{}}}:window", scope.descriptor, scope.value)).await;
match expected.iter_mut().find(|(seen, _)| *seen == slot) {
Some((_, members)) => members.push(index),
None => expected.push((slot, vec![index])),
}
}
let by_slot = expected
.into_iter()
.map(|(_, members)| members)
.collect::<Vec<_>>();
let grouped = limits.slot_groups(&scopes, false);
match deployment {
RedisDeployment::Cluster(_) => assert_eq!(grouped, by_slot),
RedisDeployment::Standalone(_) => assert_eq!(grouped, vec![vec![0, 1, 2, 3]]),
}
assert_eq!(
limits.slot_groups(&scopes, true),
vec![vec![0, 2], vec![1], vec![3]]
);
}
#[rstest]
#[tokio::test]
async fn gauges_of_different_slots_in_one_call_fail_only_on_a_cluster(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let limits = store(&deployment);
let gauges = ["a", "b", "c", "d"].map(|value| Gauge {
scope: scope("api_key", value),
limit: 5,
});
let acquired = limits
.parallel_acquire(None, &gauges, SLOT_TTL, "slot")
.unwrap()
.resolve()
.await;
match deployment {
RedisDeployment::Standalone(_) => {
assert_eq!(acquired.unwrap(), AcquireReply::Acquired(vec![1, 1, 1, 1]))
}
RedisDeployment::Cluster(_) => assert!(acquired.is_err(), "{acquired:?}"),
}
}
#[tokio::test]
async fn a_cross_slot_token_script_on_the_post_call_batch_falls_back_to_plain_increments() {
let deployment = RedisDeployment::start(Topology::Cluster).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let increments = ["a", "b", "c", "d"].map(|value| Increment {
counter: counter(&scope("api_key", value), Kind::Tokens),
amount: Number::Int(25),
ttl: Some(Duration::from_secs(60)),
});
limits
.token_increment_post_call(&request, increments.to_vec())
.unwrap();
assert_eq!(raw_text(&raw, "{api_key:a}:tokens").await, None);
request.flush_post().await;
for value in ["a", "b", "c", "d"] {
let key = format!("{{api_key:{value}}}:tokens");
let written = litellm_redis_testing::eventually(async || raw_text(&raw, &key).await).await;
assert_eq!(written, "25");
}
let local = limits
.local(&Entry::Counter(counter(
&scope("api_key", "a"),
Kind::Tokens,
)))
.unwrap();
assert_eq!(local, Some(Local::Number(Number::Int(25))));
}
#[rstest]
#[tokio::test]
async fn a_post_call_increment_from_a_thread_without_a_runtime_reaches_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let increment = Increment {
counter: counter(&scope("end_user", "e1"), Kind::Tokens),
amount: Number::Int(40),
ttl: Some(Duration::from_secs(60)),
};
std::thread::spawn({
let (limits, request, increment) =
(Arc::clone(&limits), request.clone(), increment.clone());
move || {
limits
.increment_many_post_call(&request, std::slice::from_ref(&increment))
.unwrap()
}
})
.join()
.expect("a Python thread has no Tokio context, and declaring there must not panic");
request.flush_post().await;
let written =
litellm_redis_testing::eventually(async || raw_text(&raw, "{end_user:e1}:tokens").await)
.await;
assert_eq!(written, "40");
}
#[rstest]
#[tokio::test]
async fn post_call_increments_count_in_memory_now_and_reach_redis_with_the_batch(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let increment = Increment {
counter: counter(&scope("end_user", "e1"), Kind::Tokens),
amount: Number::Int(-40),
ttl: Some(Duration::from_secs(60)),
};
raw.set("{end_user:e1}:tokens", "100", Duration::from_secs(600))
.await
.unwrap();
limits
.increment_many_post_call(&request, std::slice::from_ref(&increment))
.unwrap();
assert_eq!(
limits
.local(&Entry::Counter(increment.counter.clone()))
.unwrap(),
Some(Local::Number(Number::Int(-40)))
);
assert_eq!(
raw_text(&raw, "{end_user:e1}:tokens").await,
Some("100".to_owned())
);
request.flush_post().await;
let written = litellm_redis_testing::eventually(async || {
raw_text(&raw, "{end_user:e1}:tokens")
.await
.filter(|value| value != "100")
})
.await;
assert_eq!(written, "60");
let Ttl::Expires(ttl) = raw.ttl("{end_user:e1}:tokens").await.unwrap() else {
panic!("the counter lost its TTL");
};
// RedisBatch refreshes the TTL in a pipeline, and on a cluster runs async_increment, which
// keeps the TTL the counter already had
match topology {
Topology::Standalone => assert!(within(ttl, 60), "{ttl:?}"),
Topology::Cluster => assert!(ttl > Duration::from_secs(60), "{ttl:?}"),
}
}
#[rstest]
#[tokio::test]
async fn increment_keeps_an_existing_ttl_and_counts_in_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
let increment = Increment {
counter: counter(&scope("organization", "o1"), Kind::Tokens),
amount: Number::Int(5),
ttl: Some(Duration::from_secs(60)),
};
raw.set("{organization:o1}:tokens", "10", Duration::from_secs(600))
.await
.unwrap();
assert_eq!(
limits.increment(&increment).await.unwrap(),
Number::Float(15.0)
);
let Ttl::Expires(ttl) = raw.ttl("{organization:o1}:tokens").await.unwrap() else {
panic!("the counter lost its TTL");
};
assert!(within(ttl, 600), "{ttl:?}");
assert_eq!(
limits
.local(&Entry::Counter(increment.counter.clone()))
.unwrap(),
Some(Local::Number(Number::Int(5)))
);
}
#[tokio::test]
async fn increment_answers_with_memory_while_redis_is_down() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let limits = store(&deployment);
let increment = Increment {
counter: counter(&scope("user", "u1"), Kind::Tokens),
amount: Number::Int(7),
ttl: Some(Duration::from_secs(60)),
};
deployment.stop().await.unwrap();
assert_eq!(limits.increment(&increment).await.unwrap(), Number::Int(7));
assert_eq!(limits.increment(&increment).await.unwrap(), Number::Int(14));
assert!(
limits
.sliding_window(None, std::slice::from_ref(&increment.counter), 1_000, 60)
.unwrap()
.resolve()
.await
.is_err()
);
}
#[tokio::test]
async fn scripts_need_redis() {
let limits = RateLimits::new(None, None, REDIS_TTL);
assert!(
limits
.sliding_window(
None,
&[counter(&scope("api_key", "k"), Kind::Requests)],
0,
60
)
.is_err()
);
assert_eq!(
limits.slot_groups(&[scope("a", "1"), scope("b", "2")], false),
vec![vec![0, 1]]
);
}
#[rstest]
#[tokio::test]
async fn the_limiter_script_and_the_armed_cooldown_read_share_one_pipeline(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = connect_observed(&deployment);
let cooldowns = RouterCooldowns::new(Some(client.clone()), COOLDOWN_POLICY);
let limits = RateLimits::new(Some(client), None, REDIS_TTL);
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let armed = cooldowns
.declare_active(&request, &["dep-a".to_owned()])
.unwrap();
let admission = limits
.sliding_window(
Some(&request),
&[counter(&scope("api_key", "k1"), Kind::Requests)],
1_000,
60,
)
.unwrap();
let counts = admission.resolve().await.unwrap();
let active = cooldowns.resolve_active(armed).await.unwrap();
assert_eq!(counts[0].count, 1);
assert!(active.is_empty());
let mut seen = trips.seen();
seen.sort_by_key(|operation| format!("{operation:?}"));
let expected = match topology {
Topology::Standalone => vec![Operation::Pipeline { ops: 2 }],
Topology::Cluster => vec![Operation::Command("MGET"), Operation::Script],
};
assert_eq!(seen, expected);
}
// DualCache.async_increment_cache without a TTL still passes RedisCache.default_ttl, so
// async_increment sends INCRBYFLOAT and TTL, then EXPIRE on a key that has none, never a script
#[rstest]
#[tokio::test]
async fn an_increment_without_a_ttl_takes_the_redis_default_like_async_increment(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let limits = store(&deployment);
deployment
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
.await
.unwrap();
limits
.increment(&Increment {
counter: counter(&scope("api_key", "k"), Kind::Tokens),
amount: Number::Int(3),
ttl: None,
})
.await
.unwrap();
let Ttl::Expires(ttl) = raw.ttl("{api_key:k}:tokens").await.unwrap() else {
panic!("the counter has no TTL");
};
assert!(within(ttl, 60), "{ttl:?}");
assert_eq!(support::calls(&deployment, "incrbyfloat").await, 1);
assert_eq!(support::calls(&deployment, "expire").await, 1);
assert_eq!(
support::calls(&deployment, "eval").await + support::calls(&deployment, "evalsha").await,
0
);
}

View file

@ -59,10 +59,12 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROJECT_ITPM_DESCRIPTOR_KEY,
PROJECT_OTPM_DESCRIPTOR_KEY,
ReservationAwareIncrementOperation,
ReservationWindowIdentity,
WindowGuardedIncrement,
get_or_create_request_stash,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitScope
from litellm.types.llms.base import LiteLLMBaseModel
if TYPE_CHECKING:
@ -835,9 +837,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
return None
try:
window_start: Final = _WINDOW_START_ADAPTER.validate_python(
await self.parallel_request_limiter.internal_usage_cache.async_get_cache(
key=f"{{{descriptor_key}:{status.get('descriptor_value') or ''}}}:window",
litellm_parent_otel_span=parent_otel_span,
await self.parallel_request_limiter.rate_limits.get(
RateLimitScope(descriptor_key, status.get("descriptor_value") or "").window
),
strict=True,
)
@ -849,8 +850,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
self,
descriptors: Sequence["RateLimitDescriptor"],
tokens: int,
reservation_windows: frozenset[tuple[str, str, Literal["redis", "local"]]],
) -> tuple[ReservationAwareIncrementOperation, ...]:
reservation_windows: frozenset[ReservationWindowIdentity],
) -> tuple[WindowGuardedIncrement, ...]:
"""Refund operations for the daily token counters this batch charged.
The v3 limiter's failure hook applies them when the submission fails
@ -860,26 +861,19 @@ class _PROXY_BatchRateLimiter(CustomLogger):
"""
if tokens <= 0 or not reservation_windows:
return ()
tpd_descriptors_by_counter: Final[Mapping[str, RateLimitDescriptor]] = MappingProxyType(
{
self.parallel_request_limiter.create_rate_limit_keys(
descriptor["key"], descriptor["value"], "tokens"
): descriptor
for descriptor in descriptors
if descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX)
}
tpd_counters: Final = frozenset(
RateLimitScope(descriptor["key"], descriptor["value"]).counter("tokens")
for descriptor in descriptors
if descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX)
)
return tuple(
ReservationAwareIncrementOperation(
key=counter_key,
increment_value=-tokens,
ttl=BATCH_TPD_WINDOW_SECONDS,
window_key=f"{{{descriptor['key']}:{descriptor['value']}}}:window",
WindowGuardedIncrement(
increment=CounterIncrement(counter=counter, amount=-tokens, ttl=BATCH_TPD_WINDOW_SECONDS),
expected_window_start=window_start,
reservation_backend=backend,
)
for counter_key, window_start, backend in sorted(reservation_windows)
if (descriptor := tpd_descriptors_by_counter.get(counter_key)) is not None
for counter, window_start, backend in sorted(reservation_windows)
if counter in tpd_counters
)
@with_service_target("rate_limits")

View file

@ -38,6 +38,7 @@ from litellm.proxy.hooks.rate_limiter_utils import (
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitCounter, RateLimitScope
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral, LLMResponseTypes
@ -102,8 +103,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
async def _get_saturation_value_from_cache(
self,
counter_key: str,
) -> str | None:
counter: RateLimitCounter,
) -> int | float | str | None:
"""
Get saturation value with configurable local cache TTL.
@ -111,19 +112,15 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
TTL is configurable via litellm.priority_reservation_settings.saturation_check_cache_ttl
Args:
counter_key: The cache key for the saturation counter
counter: The saturation counter
Returns:
Counter value as string, or None if not found
Counter value, or None if not found
"""
local_cache_ttl: Final = self._get_saturation_check_cache_ttl()
return await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False,
ttl=local_cache_ttl,
value: Final = await self.v3_limiter.rate_limits.get(
counter, backfill_ttl=self._get_saturation_check_cache_ttl()
)
return None if isinstance(value, dict) else value
def _get_priority_weight(self, priority: str | None, model_info: ModelGroupInfo | None = None) -> float:
"""Get the weight for a given priority from litellm.priority_reservation"""
@ -250,14 +247,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Query RPM saturation - always read from Redis for multi-node consistency
if model_group_info.rpm is not None and model_group_info.rpm > 0:
# Use v3 limiter's key format: {key:value}:rate_limit_type
counter_key = self.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model,
rate_limit_type="requests",
)
# Query Redis directly for current counter value (skip local cache for consistency)
counter_value = await self._get_saturation_value_from_cache(counter_key=counter_key)
counter_value = await self._get_saturation_value_from_cache(
RateLimitScope("model_saturation_check", model).counter("requests")
)
if counter_value is not None:
current_requests: Final = int(counter_value)
@ -270,14 +263,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Query TPM saturation
if model_group_info.tpm is not None and model_group_info.tpm > 0:
counter_key = self.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model,
rate_limit_type="tokens",
counter_value = await self._get_saturation_value_from_cache(
RateLimitScope("model_saturation_check", model).counter("tokens")
)
counter_value = await self._get_saturation_value_from_cache(counter_key=counter_key)
if counter_value is not None:
current_tokens: Final = float(counter_value)
tpm_saturation: Final = current_tokens / model_group_info.tpm
@ -417,7 +406,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
Raises:
HTTPException: If any limit is exceeded
"""
import json
saturation_threshold: Final = _get_priority_settings().saturation_threshold
should_enforce_priority: Final = saturation >= saturation_threshold
@ -464,9 +452,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
parent_otel_span=user_api_key_dict.parent_otel_span,
)
verbose_proxy_logger.debug(
"Atomic check+increment response: %s", json.dumps(atomic_response, indent=2, default=list)
)
verbose_proxy_logger.debug("Atomic check+increment response: %s", atomic_response)
if atomic_response["overall_code"] == "OVER_LIMIT":
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model)
@ -706,7 +692,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import Usage
try:
@ -745,21 +730,11 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
if total_tokens == 0:
return
# Create pipeline operations for token increments
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = []
# Model-wide token tracking (model_saturation_check)
model_token_key: Final = self.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model_group,
rate_limit_type="tokens",
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=model_token_key,
increment_value=total_tokens,
ttl=self.v3_limiter.window_size,
)
model_tokens: Final = CounterIncrement(
counter=RateLimitScope("model_saturation_check", model_group).counter("tokens"),
amount=total_tokens,
ttl=self.v3_limiter.window_size,
)
# Priority-specific token tracking (priority_model)
@ -775,17 +750,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
else:
priority_key = f"{model_group}:default_pool"
priority_token_key: Final = self.v3_limiter.create_rate_limit_keys(
key="priority_model",
value=priority_key,
rate_limit_type="tokens",
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=priority_token_key,
increment_value=total_tokens,
pipeline_operations: Final = (
model_tokens,
CounterIncrement(
counter=RateLimitScope("priority_model", priority_key).counter("tokens"),
amount=total_tokens,
ttl=self.v3_limiter.window_size,
)
),
)
# Execute token increments with TTL preservation

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,404 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import Final, Literal, Protocol, TypeAlias, TypeVar
from pydantic import TypeAdapter
from litellm.caching.redis_batch import active_request_redis_batches
from litellm.rust_bridge._native import NativeRateLimits, NativeRedisClient, NativeRequestRedis
_R = TypeVar("_R")
CounterKind: TypeAlias = Literal["requests", "tokens"]
LocalValue: TypeAlias = int | float | str | dict[str, float]
@dataclass(frozen=True, slots=True, order=True)
class RateLimitScope:
"""A rate-limit descriptor; every key of its window shares the ``{key:value}`` hash tag."""
descriptor_key: str
descriptor_value: str
def counter(self, kind: CounterKind) -> RateLimitCounter:
return RateLimitCounter(self, kind)
@property
def window(self) -> RateLimitWindow:
return RateLimitWindow(self)
@property
def gauge(self) -> RateLimitGauge:
return RateLimitGauge(self)
@dataclass(frozen=True, slots=True)
class RateLimitWindow:
scope: RateLimitScope
@dataclass(frozen=True, slots=True, order=True)
class RateLimitCounter:
scope: RateLimitScope
kind: CounterKind
@property
def window(self) -> RateLimitWindow:
return RateLimitWindow(self.scope)
@dataclass(frozen=True, slots=True)
class RateLimitGauge:
scope: RateLimitScope
RateLimitEntry: TypeAlias = RateLimitWindow | RateLimitCounter | RateLimitGauge
@dataclass(frozen=True, slots=True)
class CounterIncrement:
counter: RateLimitCounter
amount: int | float
ttl: int | None
@dataclass(frozen=True, slots=True)
class AtomicCounter:
counter: RateLimitCounter
limit: int
increment: int
ttl: int
window_size: int
@dataclass(frozen=True, slots=True)
class WindowCount:
window_start: str
count: int
@dataclass(frozen=True, slots=True)
class AtomicApplied:
"""Per counter, the count after the increment and the start of the window it landed in."""
counters: tuple[tuple[int, int], ...]
@dataclass(frozen=True, slots=True)
class AtomicOverLimit:
index: int
current: int
limit: int
AtomicReply: TypeAlias = AtomicApplied | AtomicOverLimit
@dataclass(frozen=True, slots=True)
class GaugesAcquired:
in_flight: tuple[int, ...]
@dataclass(frozen=True, slots=True)
class GaugeOverLimit:
index: int
in_flight: int
limit: int
AcquireReply: TypeAlias = GaugesAcquired | GaugeOverLimit
@dataclass(frozen=True, slots=True)
class GuardedReply:
applied: bool
count: int
_WINDOW_COUNTS: Final = TypeAdapter(list[tuple[str, int]])
_VALUES: Final = TypeAdapter(list[str | None])
_COUNTS: Final = TypeAdapter(list[int])
_TOTALS: Final = TypeAdapter(list[float])
_GUARDED: Final = TypeAdapter(tuple[bool, int])
_ATOMIC: Final = TypeAdapter(tuple[tuple[int, int, int] | None, list[tuple[int, int]]])
_ACQUIRE: Final = TypeAdapter(tuple[tuple[int, int, int] | None, list[int]])
NativeScope: TypeAlias = tuple[str, str]
NativeCounter: TypeAlias = tuple[str, str, CounterKind]
NativeEntry: TypeAlias = tuple[Literal["window", "requests", "tokens", "max_parallel_requests"], str, str]
NativeIncrement: TypeAlias = tuple[NativeCounter, int | float, int | None]
NativeAtomicCounter: TypeAlias = tuple[NativeCounter, int, int, int, int]
class PendingScript(Protocol):
def resolve(self) -> Awaitable[object]: ...
class RateLimitsHandle(Protocol):
def has_redis(self) -> bool: ...
def slot_groups(self, scopes_: Sequence[NativeScope], force: bool) -> list[list[int]]: ...
def local(self, entries: Sequence[NativeEntry]) -> list[LocalValue | None]: ...
def set_local(self, entry: NativeEntry, value: LocalValue, ttl_seconds: float | None) -> None: ...
def get(self, entry: NativeEntry, *, backfill_ttl_seconds: float | None = None) -> Awaitable[LocalValue | None]: ...
def sliding_window(
self,
counters_: Sequence[NativeCounter],
now: int,
window_size: int,
*,
request: NativeRequestRedis | None = None,
) -> PendingScript: ...
def read_counters(
self, counters_: Sequence[NativeCounter], *, request: NativeRequestRedis | None = None
) -> PendingScript: ...
def check_and_increment(
self, counters_: Sequence[NativeAtomicCounter], *, request: NativeRequestRedis | None = None
) -> PendingScript: ...
def guarded_increment(
self,
counter: NativeCounter,
expected_window_start: str,
amount: float,
ttl_seconds: int | None,
*,
request: NativeRequestRedis | None = None,
) -> PendingScript: ...
def parallel_acquire(
self,
gauges: Sequence[tuple[NativeScope, int]],
slot_ttl_seconds: int,
slot_id: str,
*,
request: NativeRequestRedis | None = None,
) -> PendingScript: ...
def parallel_release(
self, gauges: Sequence[NativeScope], slot_id: str, *, request: NativeRequestRedis | None = None
) -> PendingScript: ...
def parallel_count(
self, gauges: Sequence[NativeScope], slot_ttl_seconds: int, *, request: NativeRequestRedis | None = None
) -> PendingScript: ...
def token_increment(
self, increments_: Sequence[NativeIncrement], *, request: NativeRequestRedis | None = None
) -> PendingScript: ...
def token_increment_post_call(
self, request: NativeRequestRedis, increments_: Sequence[NativeIncrement]
) -> Awaitable[None]: ...
def parallel_release_post_call(
self, request: NativeRequestRedis, gauges: Sequence[NativeScope], slot_id: str
) -> Awaitable[None]: ...
def increment(self, increment: NativeIncrement) -> Awaitable[int | float]: ...
def increment_many(self, increments_: Sequence[NativeIncrement]) -> Awaitable[None]: ...
def increment_many_post_call(
self, request: NativeRequestRedis, increments_: Sequence[NativeIncrement]
) -> Awaitable[None]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, float | None, int], RateLimitsHandle]
def _scope(scope: RateLimitScope) -> NativeScope:
return (scope.descriptor_key, scope.descriptor_value)
def _counter(counter: RateLimitCounter) -> NativeCounter:
return (counter.scope.descriptor_key, counter.scope.descriptor_value, counter.kind)
def _entry(entry: RateLimitEntry) -> NativeEntry:
match entry:
case RateLimitWindow(scope):
return ("window", scope.descriptor_key, scope.descriptor_value)
case RateLimitGauge(scope):
return ("max_parallel_requests", scope.descriptor_key, scope.descriptor_value)
case RateLimitCounter(scope, kind):
return (kind, scope.descriptor_key, scope.descriptor_value)
def _increment(increment: CounterIncrement) -> NativeIncrement:
return (_counter(increment.counter), increment.amount, increment.ttl)
def _atomic(counter: AtomicCounter) -> NativeAtomicCounter:
return (_counter(counter.counter), counter.limit, counter.increment, counter.ttl, counter.window_size)
async def _window_counts(pending: Awaitable[object]) -> tuple[WindowCount, ...]:
return tuple(WindowCount(start, count) for start, count in _WINDOW_COUNTS.validate_python(await pending))
async def _values(pending: Awaitable[object]) -> tuple[str | None, ...]:
return tuple(_VALUES.validate_python(await pending))
async def _counts(pending: Awaitable[object]) -> tuple[int, ...]:
return tuple(_COUNTS.validate_python(await pending))
async def _totals(pending: Awaitable[object]) -> tuple[float, ...]:
return tuple(_TOTALS.validate_python(await pending))
async def _guarded(pending: Awaitable[object]) -> GuardedReply:
applied, count = _GUARDED.validate_python(await pending)
return GuardedReply(applied=applied, count=count)
async def _atomic_reply(pending: Awaitable[object]) -> AtomicReply:
over_limit, applied = _ATOMIC.validate_python(await pending)
if over_limit is None:
return AtomicApplied(tuple(applied))
index, current, limit = over_limit
return AtomicOverLimit(index=index, current=current, limit=limit)
async def _acquire_reply(pending: Awaitable[object]) -> AcquireReply:
over_limit, in_flight = _ACQUIRE.validate_python(await pending)
if over_limit is None:
return GaugesAcquired(tuple(in_flight))
index, count, limit = over_limit
return GaugeOverLimit(index=index, in_flight=count, limit=limit)
class RateLimitStore:
"""The v3 limiter's keyspace on the Rust store. Inside a request scope every script rides the request's
native pre-call batch, and ``declare_*`` returns it without awaiting so several share one trip (None
outside a scope, where the plain method runs on a batch of its own)."""
def __init__(
self,
client: NativeRedisClient | None,
backfill_ttl: float | None = None,
redis_ttl: int = 60,
*,
handle: HandleFactory = NativeRateLimits,
) -> None:
self._handle: Final = handle(client, backfill_ttl, redis_ttl)
@property
def has_redis(self) -> bool:
return self._handle.has_redis()
@property
def has_post_call_batch(self) -> bool:
"""Whether a ``*_post_call`` write would ride the request's post-call batch rather than be refused."""
return self.has_redis and active_request_redis_batches() is not None
def slot_groups(self, scopes: Sequence[RateLimitScope], force: bool) -> tuple[tuple[int, ...], ...]:
return tuple(tuple(group) for group in self._handle.slot_groups(tuple(map(_scope, scopes)), force))
def local(self, entries: Sequence[RateLimitEntry]) -> tuple[LocalValue | None, ...]:
return tuple(self._handle.local(tuple(map(_entry, entries))))
def set_local(self, entry: RateLimitEntry, value: LocalValue, ttl: float | None) -> None:
self._handle.set_local(_entry(entry), value, ttl)
async def get(self, entry: RateLimitEntry, backfill_ttl: float | None = None) -> LocalValue | None:
"""``DualCache.async_get_cache``: memory, else Redis written back to memory. A memory hit, or a miss
without Redis, answers without yielding to the event loop, as ``DualCache`` did."""
native: Final = _entry(entry)
local: Final = self._handle.local((native,))[0]
if local is not None or not self.has_redis:
return local
return await self._handle.get(native, backfill_ttl_seconds=backfill_ttl)
def declare_sliding_window(
self, counters: Sequence[RateLimitCounter], now: int, window_size: int
) -> Awaitable[tuple[WindowCount, ...]] | None:
return self._declare(
lambda request: self._handle.sliding_window(
tuple(map(_counter, counters)), now, window_size, request=request
),
_window_counts,
)
async def sliding_window(
self, counters: Sequence[RateLimitCounter], now: int, window_size: int
) -> tuple[WindowCount, ...]:
return await _window_counts(
self._handle.sliding_window(tuple(map(_counter, counters)), now, window_size).resolve()
)
async def read_counters(self, counters: Sequence[RateLimitCounter]) -> tuple[str | None, ...]:
"""Each counter's window start then its count, as ``BATCH_COUNTER_READ_SCRIPT`` reads them."""
native: Final = tuple(map(_counter, counters))
return await _values(self._handle.read_counters(native).resolve())
def declare_check_and_increment(self, counters: Sequence[AtomicCounter]) -> Awaitable[AtomicReply] | None:
native: Final = tuple(map(_atomic, counters))
return self._declare(lambda request: self._handle.check_and_increment(native, request=request), _atomic_reply)
async def check_and_increment(self, counters: Sequence[AtomicCounter]) -> AtomicReply:
return await _atomic_reply(self._handle.check_and_increment(tuple(map(_atomic, counters))).resolve())
async def guarded_increment(
self, counter: RateLimitCounter, expected_window_start: str, amount: float, ttl: int | None
) -> GuardedReply:
native: Final = _counter(counter)
return await _guarded(self._handle.guarded_increment(native, expected_window_start, amount, ttl).resolve())
async def parallel_acquire(
self, gauges: Sequence[tuple[RateLimitScope, int]], slot_ttl: int, slot_id: str
) -> AcquireReply:
native: Final = tuple((_scope(scope), limit) for scope, limit in gauges)
return await _acquire_reply(self._handle.parallel_acquire(native, slot_ttl, slot_id).resolve())
async def parallel_release(self, gauges: Sequence[RateLimitScope], slot_id: str) -> tuple[int, ...]:
native: Final = tuple(map(_scope, gauges))
return await _counts(self._handle.parallel_release(native, slot_id).resolve())
async def parallel_count(self, gauges: Sequence[RateLimitScope], slot_ttl: int) -> tuple[int, ...]:
native: Final = tuple(map(_scope, gauges))
return await _counts(self._handle.parallel_count(native, slot_ttl).resolve())
async def token_increment(self, increments: Sequence[CounterIncrement]) -> tuple[float, ...]:
return await _totals(self._handle.token_increment(tuple(map(_increment, increments))).resolve())
async def token_increment_post_call(self, increments: Sequence[CounterIncrement]) -> bool:
"""``TOKEN_INCREMENT_SCRIPT`` on the request's post-call batch, falling back to ``increment_many`` when
the script fails; False outside a request scope or without Redis."""
request: Final = self._post_call_request()
if request is None:
return False
await self._handle.token_increment_post_call(request, tuple(map(_increment, increments)))
return True
async def parallel_release_post_call(self, gauges: Sequence[RateLimitScope], slot_id: str) -> bool:
request: Final = self._post_call_request()
if request is None:
return False
await self._handle.parallel_release_post_call(request, tuple(map(_scope, gauges)), slot_id)
return True
async def increment(self, increment: CounterIncrement) -> int | float:
"""``DualCache.async_increment_cache``: memory, then Redis keeping an existing TTL."""
return await self._handle.increment(_increment(increment))
async def increment_many(self, increments: Sequence[CounterIncrement]) -> None:
"""``DualCache.async_increment_cache_pipeline``: memory, then one Redis pipeline refreshing each TTL."""
await self._handle.increment_many(tuple(map(_increment, increments)))
async def increment_many_post_call(self, increments: Sequence[CounterIncrement]) -> None:
"""Memory now; the Redis increments ride the request's post-call batch, or go out at once without one."""
request: Final = self._post_call_request()
native: Final = tuple(map(_increment, increments))
if request is None:
await self._handle.increment_many(native)
return
await self._handle.increment_many_post_call(request, native)
def _post_call_request(self) -> NativeRequestRedis | None:
batches: Final = active_request_redis_batches()
if batches is None or not self.has_redis:
return None
return batches.native_post_call()
def _declare(
self,
declare: Callable[[NativeRequestRedis], PendingScript],
decode: Callable[[Awaitable[object]], Awaitable[_R]],
) -> Awaitable[_R] | None:
batches: Final = active_request_redis_batches()
if batches is None or not self.has_redis:
return None
return decode(batches.native_result(declare(batches.native())))

View file

@ -19,7 +19,6 @@ class NativeRedisStoreError(Exception): ...
def trace_encode_error(message: str) -> bytes: ...
def bundled_model_cost_map() -> bytes: ...
def native_redis_breaker_states() -> list[str]: ...
@final
class NativeClickHouseSpendConfig:
def __new__(cls, database: str, url: str, retention_days: int) -> NativeClickHouseSpendConfig: ...
@ -85,6 +84,8 @@ class NativeRouterUsage:
minute: str,
tokens: float,
requests: int,
*,
request: NativeRequestRedis | None = None,
) -> Future[None]: ...
def get_many(
self, counters: Sequence[tuple[str, str | None, str, Literal["tpm", "rpm"]]]
@ -92,6 +93,85 @@ class NativeRouterUsage:
def flush_local_counters(self) -> None: ...
def snapshot(self) -> dict[str, float]: ...
_RateLimitScope = tuple[str, str]
_RateLimitCounter = tuple[str, str, Literal["requests", "tokens"]]
_RateLimitEntry = tuple[Literal["window", "requests", "tokens", "max_parallel_requests"], str, str]
_RateLimitLocal = int | float | str | dict[str, float]
_RateLimitIncrement = tuple[_RateLimitCounter, int | float, int | None]
_RateLimitAtomicCounter = tuple[_RateLimitCounter, int, int, int, int]
@final
class NativeRateLimitPending:
def resolve(self) -> Future[object]: ...
@final
class NativeRateLimits:
def __new__(
cls,
client: NativeRedisClient | None,
backfill_ttl_seconds: float | None = None,
redis_ttl_seconds: int = 60,
/,
) -> NativeRateLimits: ...
def has_redis(self) -> bool: ...
def slot_groups(self, scopes_: Sequence[_RateLimitScope], force: bool) -> list[list[int]]: ...
def local(self, entries: Sequence[_RateLimitEntry]) -> list[_RateLimitLocal | None]: ...
def set_local(self, entry: _RateLimitEntry, value: _RateLimitLocal, ttl_seconds: float | None) -> None: ...
def get(
self, entry: _RateLimitEntry, *, backfill_ttl_seconds: float | None = None
) -> Future[_RateLimitLocal | None]: ...
def sliding_window(
self,
counters_: Sequence[_RateLimitCounter],
now: int,
window_size: int,
*,
request: NativeRequestRedis | None = None,
) -> NativeRateLimitPending: ...
def read_counters(
self, counters_: Sequence[_RateLimitCounter], *, request: NativeRequestRedis | None = None
) -> NativeRateLimitPending: ...
def check_and_increment(
self, counters_: Sequence[_RateLimitAtomicCounter], *, request: NativeRequestRedis | None = None
) -> NativeRateLimitPending: ...
def guarded_increment(
self,
counter: _RateLimitCounter,
expected_window_start: str,
amount: float,
ttl_seconds: int | None,
*,
request: NativeRequestRedis | None = None,
) -> NativeRateLimitPending: ...
def parallel_acquire(
self,
gauges: Sequence[tuple[_RateLimitScope, int]],
slot_ttl_seconds: int,
slot_id: str,
*,
request: NativeRequestRedis | None = None,
) -> NativeRateLimitPending: ...
def parallel_release(
self, gauges: Sequence[_RateLimitScope], slot_id: str, *, request: NativeRequestRedis | None = None
) -> NativeRateLimitPending: ...
def parallel_count(
self, gauges: Sequence[_RateLimitScope], slot_ttl_seconds: int, *, request: NativeRequestRedis | None = None
) -> NativeRateLimitPending: ...
def token_increment(
self, increments_: Sequence[_RateLimitIncrement], *, request: NativeRequestRedis | None = None
) -> NativeRateLimitPending: ...
def token_increment_post_call(
self, request: NativeRequestRedis, increments_: Sequence[_RateLimitIncrement]
) -> Future[None]: ...
def parallel_release_post_call(
self, request: NativeRequestRedis, gauges: Sequence[_RateLimitScope], slot_id: str
) -> Future[None]: ...
def increment(self, increment: _RateLimitIncrement) -> Future[int | float]: ...
def increment_many(self, increments_: Sequence[_RateLimitIncrement]) -> Future[None]: ...
def increment_many_post_call(
self, request: NativeRequestRedis, increments_: Sequence[_RateLimitIncrement]
) -> Future[None]: ...
@final
class NativeAuthCache:
def __new__(
@ -215,7 +295,6 @@ def messages(
def amessages(
call: NativeCall,
) -> Execution: ...
@final
class ResponsesWebSocketConnection:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...

View file

@ -10,6 +10,7 @@ import datetime
import hashlib
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final
import pytest
@ -28,17 +29,17 @@ from litellm.caching.redis_batch import (
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PARALLEL_RELEASE_SCRIPT,
TOKEN_INCREMENT_SCRIPT,
ParallelSlotAcquisition,
PROXY_MaxParallelRequestsHandler_v3,
RequestRateLimiterStash,
)
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement
from litellm.proxy.utils import InternalUsageCache
from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitScope
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import ModelResponse
from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits
from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache
@ -97,20 +98,31 @@ def _names(client: FakeClient, index: int = 0) -> list[str]:
return [command[0] for command in client.pipelines[index].commands]
def _limiter(redis_cache: FakeRedisCache) -> PROXY_MaxParallelRequestsHandler_v3:
dual_cache = DualCache()
dual_cache.attach_redis_cache(redis_cache)
return PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(dual_cache=dual_cache))
def _slot_stash(slot_id: str, *counter_keys: str) -> RequestRateLimiterStash:
return RequestRateLimiterStash(
parallel_slot=ParallelSlotAcquisition(slot_id=slot_id, counter_keys=list(counter_keys))
def _limiter(fakes: FakeRateLimits) -> PROXY_MaxParallelRequestsHandler_v3:
return PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(dual_cache=DualCache()), rate_limit_store=fakes.store()
)
def _token_ops(*keys: str) -> list[RedisPipelineIncrementOperation]:
return [RedisPipelineIncrementOperation(key=key, increment_value=10, ttl=60) for key in keys]
_K1: Final = RateLimitScope("api_key", "k1")
def _slot_stash(slot_id: str, *gauges: RateLimitScope) -> RequestRateLimiterStash:
return RequestRateLimiterStash(parallel_slot=ParallelSlotAcquisition(slot_id=slot_id, gauges=gauges))
def _token_ops(*scopes: RateLimitScope) -> list[CounterIncrement]:
return [CounterIncrement(counter=scope.counter("tokens"), amount=10, ttl=60) for scope in scopes]
def _scripts(calls: list[tuple[str, list[str], list[object]]], reply: object) -> Callable[..., Awaitable[object]]:
async def script(keys: list[str], args: list[object]) -> object:
calls.append(("script", keys, args))
if isinstance(reply, Exception):
raise reply
return reply
return script
def _response_cache(redis_cache: FakeRedisCache) -> Cache:
@ -139,25 +151,28 @@ def _tpm_kwargs() -> Mapping[str, object]:
@pytest.mark.asyncio
async def test_every_post_call_owner_rides_one_pipeline_that_goes_out_when_the_callbacks_are_done():
async def test_post_call_owners_wait_for_the_callbacks_and_the_limiter_declares_on_the_native_post_call_batch():
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
limiter = _limiter(redis_cache)
calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake scripts
fakes = FakeRateLimits(token_increment_script=_scripts(calls, [10.0]), parallel_release_script=_scripts(calls, [0]))
limiter = _limiter(fakes)
tpm, router_cache = _tpm_router(redis_cache)
with request_redis_batch_scope():
await tpm.async_log_success_event(_tpm_kwargs(), None, None, None)
await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{api_key:k1}:tokens"))
await limiter._release_stashed_parallel_slot(
_slot_stash("slot-1", "{api_key:k1}:parallel"), None, in_logging_callback=True
)
await limiter.async_increment_tokens_with_ttl_preservation(_token_ops(_K1))
await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True)
assert client.pipelines == [] # nothing goes out while the callbacks are still declaring
await flush_post_call_redis_batches()
assert fakes.post_call == ["token_increment", "parallel_release"]
assert calls == [
("script", ["{api_key:k1}:tokens"], [10, 60]),
("script", ["{api_key:k1}:max_parallel_requests"], ["slot-1"]),
]
assert len(client.pipelines) == 1
assert _names(client) == ["INCRBYFLOAT", "EXPIRE", "EVALSHA", "EVALSHA"]
evalshas = [c for c in client.pipelines[0].commands if c[0] == "EVALSHA"]
assert [c[1] for c in evalshas] == [sha_of(TOKEN_INCREMENT_SCRIPT), sha_of(PARALLEL_RELEASE_SCRIPT)]
assert _names(client) == ["INCRBYFLOAT", "EXPIRE"]
assert redis_cache.alone == []
assert (
await router_cache.in_memory_cache.async_get_cache(
@ -206,99 +221,52 @@ async def test_a_chat_response_written_through_the_handler_dual_cache_is_in_memo
assert (direct_set[0], direct_set[1], direct_set[2]["ttl"]) == ("SET", cache_key, 120)
@pytest.mark.asyncio
async def test_a_failed_operation_fails_only_its_owner_and_the_owner_applies_its_own_fallback():
def replies(command: tuple[object, ...]) -> object:
if command[0] == "EVALSHA" and command[3] == "{api_key:k1}:tokens":
return Exception("ERR Lua")
return _ok_replies(command)
client = FakeClient(replies)
redis_cache = PostCallFakeRedisCache(client)
limiter = _limiter(redis_cache)
with request_redis_batch_scope():
await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{api_key:k1}:tokens"))
await limiter.async_increment_tokens_with_ttl_preservation(_token_ops("{team:t1}:tokens"))
await flush_post_call_redis_batches()
assert len(client.pipelines) == 1
# the failed group falls back to the plain increment (memory + Redis), the healthy group does not
assert redis_cache.alone == [("INCRBYFLOAT", "{api_key:k1}:tokens", 10)]
assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{api_key:k1}:tokens") == 10
assert await limiter.internal_usage_cache.dual_cache.in_memory_cache.async_get_cache("{team:t1}:tokens") is None
@pytest.mark.asyncio
async def test_a_failed_slot_release_script_releases_the_slot_in_memory():
def replies(command: tuple[object, ...]) -> object:
if command[0] == "EVALSHA":
return Exception("ERR Lua")
return _ok_replies(command)
redis_cache = PostCallFakeRedisCache(FakeClient(replies))
limiter = _limiter(redis_cache)
memory = limiter.internal_usage_cache.dual_cache.in_memory_cache
await memory.async_set_cache("{api_key:k1}:parallel", {"slot-1": 1.0, "slot-2": 1.0})
calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake script
fakes = FakeRateLimits(parallel_release_script=_scripts(calls, Exception("ERR Lua")))
limiter = _limiter(fakes)
limiter.rate_limits.set_local(_K1.gauge, {"slot-1": 1.0, "slot-2": 1.0}, None)
with request_redis_batch_scope():
await limiter._release_stashed_parallel_slot(
_slot_stash("slot-1", "{api_key:k1}:parallel"), None, in_logging_callback=True
)
await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True)
await flush_post_call_redis_batches()
assert await memory.async_get_cache("{api_key:k1}:parallel") == {"slot-2": 1.0}
class DirectScriptFakeRedisCache(PostCallFakeRedisCache):
"""Records the release script a pre-response caller runs outside the pipeline."""
def async_register_script(self, script: str) -> Callable[..., Awaitable[object]]:
async def run(keys: Sequence[str], args: Sequence[object]) -> object:
self.alone.append(("EVALSHA", tuple(keys), tuple(args)))
return [0 for _ in keys]
return run
assert fakes.post_call == ["parallel_release"]
assert limiter.rate_limits.local((_K1.gauge,)) == ({"slot-2": 1.0},)
@pytest.mark.asyncio
async def test_a_slot_released_before_the_response_reaches_redis_at_once_not_on_the_pipeline():
client = FakeClient(_ok_replies)
redis_cache = DirectScriptFakeRedisCache(client)
limiter = _limiter(redis_cache)
memory = limiter.internal_usage_cache.dual_cache.in_memory_cache
await memory.async_set_cache("{api_key:k1}:parallel", {"slot-1": 1.0})
async def test_a_slot_released_before_the_response_reaches_redis_at_once_not_on_the_post_call_batch():
calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake script
fakes = FakeRateLimits(parallel_release_script=_scripts(calls, [0]))
limiter = _limiter(fakes)
limiter.rate_limits.set_local(_K1.gauge, {"slot-1": 1.0}, None)
with request_redis_batch_scope():
await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", "{api_key:k1}:parallel"), None)
assert redis_cache.alone == [("EVALSHA", ("{api_key:k1}:parallel",), ("slot-1",))]
assert await memory.async_get_cache("{api_key:k1}:parallel") == 0
await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None)
assert calls == [("script", ["{api_key:k1}:max_parallel_requests"], ["slot-1"])]
assert limiter.rate_limits.local((_K1.gauge,)) == (0,)
await flush_post_call_redis_batches()
assert client.pipelines == []
assert fakes.post_call == []
@pytest.mark.asyncio
async def test_a_released_slot_is_free_locally_at_once_and_the_older_redis_count_does_not_overwrite_the_gauge():
def replies(command: tuple[object, ...]) -> object:
if command[0] == "EVALSHA":
return [2]
return _ok_replies(command)
redis_cache = PostCallFakeRedisCache(FakeClient(replies))
limiter = _limiter(redis_cache)
memory = limiter.internal_usage_cache.dual_cache.in_memory_cache
await memory.async_set_cache("{api_key:k1}:parallel", {"slot-1": 1.0, "slot-2": 1.0, "slot-3": 1.0})
calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake script
fakes = FakeRateLimits(parallel_release_script=_scripts(calls, [2]))
limiter = _limiter(fakes)
limiter.rate_limits.set_local(_K1.gauge, {"slot-1": 1.0, "slot-2": 1.0, "slot-3": 1.0}, None)
with request_redis_batch_scope():
await limiter._release_stashed_parallel_slot(
_slot_stash("slot-1", "{api_key:k1}:parallel"), None, in_logging_callback=True
)
assert await memory.async_get_cache("{api_key:k1}:parallel") == {"slot-2": 1.0, "slot-3": 1.0}
await memory.async_set_cache("{api_key:k1}:parallel", {"slot-2": 1.0, "slot-3": 1.0, "slot-4": 1.0})
await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True)
assert limiter.rate_limits.local((_K1.gauge,)) == ({"slot-2": 1.0, "slot-3": 1.0},)
limiter.rate_limits.set_local(_K1.gauge, {"slot-2": 1.0, "slot-3": 1.0, "slot-4": 1.0}, None)
await flush_post_call_redis_batches()
assert await memory.async_get_cache("{api_key:k1}:parallel") == {"slot-2": 1.0, "slot-3": 1.0, "slot-4": 1.0}
assert calls == [("script", ["{api_key:k1}:max_parallel_requests"], ["slot-1"])]
assert limiter.rate_limits.local((_K1.gauge,)) == ({"slot-2": 1.0, "slot-3": 1.0, "slot-4": 1.0},)
@pytest.mark.asyncio

View file

@ -4,7 +4,6 @@ router's usage counter read, auth identity and spend counters all join the reque
from __future__ import annotations
import asyncio
import hashlib
import json
from collections.abc import Callable, Sequence
from dataclasses import dataclass
@ -23,12 +22,12 @@ from litellm.proxy._types import LiteLLM_TeamTableCachedObj, LiteLLM_UserTable
from litellm.proxy.auth.auth_checks import cache_team_object
from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back, prefetch_identity_keys
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
CHECK_AND_INCREMENT_BY_N_SCRIPT,
PROXY_MaxParallelRequestsHandler_v3,
RateLimitDescriptor,
RateLimitUnverifiableError,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.redis.stores.rate_limits import RateLimitScope
from litellm.redis.stores.router_cooldowns import ActiveCooldowns, RouterCooldownStore
from litellm.router_api import RoutingContext
from litellm.router_utils.routing_prefetch import ArmedRead, TierRead, arm_cooldowns, read_cooldowns
@ -37,6 +36,7 @@ from litellm.router_utils.routing_read_batch import (
RoutingPrefetch,
)
from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache
from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits
from .test_redis_batch import (
FakeClient,
@ -55,30 +55,47 @@ _MODEL_GROUP = "claude"
_FAR_FUTURE = 4_102_444_800.0 # 2100-01-01, a cooldown stamped then is still active
def sha_of(script: str) -> str:
return hashlib.sha1(script.encode()).hexdigest() # noqa: S324
_OK: Final = [0, 1, 1700000000]
def _limiter(redis_cache: FakeRedisCache, fail_closed: bool = False) -> PROXY_MaxParallelRequestsHandler_v3:
dual_cache = DualCache()
limiter = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(dual_cache=dual_cache),
def _over_limit(current: int, limit: int) -> list[int]:
return [1, 1, current, limit]
def _check_replies(replies: dict[str, object]) -> Callable[..., Any]:
"""``CHECK_AND_INCREMENT_BY_N_SCRIPT`` per descriptor window: OK unless ``replies`` holds another reply or an
exception for it."""
async def check(keys: list[str], args: list[object]) -> object:
reply: Final = replies.get(keys[0], _OK)
if isinstance(reply, Exception):
raise reply
return reply
return check
def _limiter(
replies: dict[str, object] | None = None, fail_closed: bool = False
) -> tuple[PROXY_MaxParallelRequestsHandler_v3, FakeRateLimits]:
fakes: Final = FakeRateLimits(
check_and_increment_by_n_script=_check_replies(replies or {}),
window_guarded_token_increment_script=AsyncMock(return_value=[1, 0]),
)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(dual_cache=DualCache()),
fail_closed_resolver=lambda: fail_closed,
rate_limit_store=fakes.store(),
)
dual_cache.attach_redis_cache(redis_cache) # after init: the fake has no server to register scripts on
limiter.check_and_increment_by_n_script = AsyncMock(
side_effect=AssertionError("descriptor groups must ride the request pipeline")
)
limiter.window_guarded_token_increment_script = AsyncMock(return_value=[1, 0])
return limiter
return limiter, fakes
def _descriptor(key: str, value: str, rpm: int) -> RateLimitDescriptor:
return {"key": key, "value": value, "rate_limit": {"requests_per_unit": rpm}}
def _refunds(limiter: PROXY_MaxParallelRequestsHandler_v3) -> list[tuple[str, float]]:
refund_script = limiter.window_guarded_token_increment_script
def _refunds(fakes: FakeRateLimits) -> list[tuple[str, float]]:
refund_script: Final = fakes.window_guarded_token_increment_script
assert isinstance(refund_script, AsyncMock)
return [(call.kwargs["keys"][1], call.kwargs["args"][1]) for call in refund_script.await_args_list]
@ -89,8 +106,6 @@ def _declared_before_the_first_flush(natives: FakeNatives) -> list[str]:
def _lua_ok_replies(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA":
return [0, 1, 1700000000] # OK: one counter, new_counter=1, window_start
if command[0] == "MGET":
return [None for _ in command[1:]]
if command[0] == "SET":
@ -99,16 +114,16 @@ def _lua_ok_replies(command: tuple[Any, ...]) -> Any:
@pytest.mark.asyncio
async def test_descriptor_lua_calls_share_one_pipeline_and_each_keeps_its_result():
client = FakeClient(_lua_ok_replies)
limiter = _limiter(FakeRedisCache(client))
async def test_descriptor_lua_calls_share_one_native_pipeline_and_each_keeps_its_result():
limiter, _fakes = _limiter()
natives: Final = FakeNatives()
descriptors = [
_descriptor("api_key", "k1", 10),
_descriptor("model_per_key", "k1:gpt", 5),
_descriptor("team", "t1", 20),
]
with request_redis_batch_scope():
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
response = await limiter.atomic_check_and_increment_by_n(
descriptors=descriptors,
increments=[{"requests": 1}, {"requests": 1}, {"requests": 1}],
@ -116,23 +131,12 @@ async def test_descriptor_lua_calls_share_one_pipeline_and_each_keeps_its_result
assert response["overall_code"] == "OK"
assert [s["descriptor_key"] for s in response["statuses"]] == ["api_key", "model_per_key", "team"]
assert len(client.pipelines) == 1
evalshas = [c for c in client.pipelines[0].commands if c[0] == "EVALSHA"]
assert len(evalshas) == 3
assert {c[1] for c in evalshas} == {sha_of(CHECK_AND_INCREMENT_BY_N_SCRIPT)}
assert [c[3] for c in evalshas] == ["{api_key:k1}:window", "{model_per_key:k1:gpt}:window", "{team:t1}:window"]
assert _declared_before_the_first_flush(natives) == ["declare check_and_increment"] * 3
@pytest.mark.asyncio
async def test_an_over_limit_descriptor_in_the_pipeline_refunds_the_groups_that_were_applied():
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA" and command[3] == "{team:t1}:window":
return [1, 1, 21, 20] # OVER_LIMIT on its first counter
return _lua_ok_replies(command)
client = FakeClient(replies)
redis_cache = FakeRedisCache(client)
limiter = _limiter(redis_cache)
limiter, fakes = _limiter({"{team:t1}:window": _over_limit(21, 20)})
with request_redis_batch_scope():
response = await limiter.atomic_check_and_increment_by_n(
@ -142,20 +146,12 @@ async def test_an_over_limit_descriptor_in_the_pipeline_refunds_the_groups_that_
assert response["overall_code"] == "OVER_LIMIT"
assert response["statuses"][0]["descriptor_key"] == "team"
assert _refunds(limiter) == [("{api_key:k1}:requests", -1.0)]
assert len(client.pipelines) == 1
assert _refunds(fakes) == [("{api_key:k1}:requests", -1)]
@pytest.mark.asyncio
async def test_an_over_limit_descriptor_also_refunds_the_groups_the_pipeline_incremented_after_it():
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA" and command[3] == "{api_key:k1}:window":
return [1, 1, 11, 10] # OVER_LIMIT on the first group; the later groups already incremented
return _lua_ok_replies(command)
client = FakeClient(replies)
redis_cache = FakeRedisCache(client)
limiter = _limiter(redis_cache)
limiter, fakes = _limiter({"{api_key:k1}:window": _over_limit(11, 10)})
with request_redis_batch_scope():
response = await limiter.atomic_check_and_increment_by_n(
@ -169,22 +165,14 @@ async def test_an_over_limit_descriptor_also_refunds_the_groups_the_pipeline_inc
assert response["overall_code"] == "OVER_LIMIT"
assert response["statuses"][0]["descriptor_key"] == "api_key"
assert _refunds(limiter) == [("{team:t1}:requests", -1.0), ("{model_per_key:k1:gpt}:requests", -1.0)]
assert len(client.pipelines) == 1
assert _refunds(fakes) == [("{team:t1}:requests", -1), ("{model_per_key:k1:gpt}:requests", -1)]
@pytest.mark.asyncio
async def test_a_redis_denial_stands_when_another_pipelined_group_fails():
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA" and command[3] == "{api_key:k1}:window":
return [1, 1, 11, 10] # OVER_LIMIT
if command[0] == "EVALSHA" and command[3] == "{team:t1}:window":
return ValueError("script blew up")
return _lua_ok_replies(command)
client = FakeClient(replies)
redis_cache = FakeRedisCache(client)
limiter = _limiter(redis_cache)
limiter, fakes = _limiter(
{"{api_key:k1}:window": _over_limit(11, 10), "{team:t1}:window": ValueError("script blew up")}
)
with request_redis_batch_scope():
response = await limiter.atomic_check_and_increment_by_n(
@ -198,20 +186,12 @@ async def test_a_redis_denial_stands_when_another_pipelined_group_fails():
assert response["overall_code"] == "OVER_LIMIT" # not the in-memory fallback's verdict
assert response["statuses"][0]["descriptor_key"] == "api_key"
assert _refunds(limiter) == [("{model_per_key:k1:gpt}:requests", -1.0)]
assert len(client.pipelines) == 1
assert _refunds(fakes) == [("{model_per_key:k1:gpt}:requests", -1)]
@pytest.mark.asyncio
async def test_one_failed_lua_group_refunds_the_other_pipelined_groups_and_falls_back_to_in_memory():
def replies(command: tuple[Any, ...]) -> Any:
if command[0] == "EVALSHA" and command[3] == "{api_key:k1}:window":
return ValueError("script blew up")
return _lua_ok_replies(command)
client = FakeClient(replies)
redis_cache = FakeRedisCache(client)
limiter = _limiter(redis_cache)
limiter, fakes = _limiter({"{api_key:k1}:window": ValueError("script blew up")})
with request_redis_batch_scope():
response = await limiter.atomic_check_and_increment_by_n(
@ -221,32 +201,26 @@ async def test_one_failed_lua_group_refunds_the_other_pipelined_groups_and_falls
assert response["overall_code"] == "OK"
assert len(response["statuses"]) == 2 # in-memory enforcement covered both descriptors
assert _refunds(limiter) == [("{team:t1}:requests", -1.0)]
assert len(client.pipelines) == 1
assert _refunds(fakes) == [("{team:t1}:requests", -1)]
assert limiter.rate_limits.local((RateLimitScope("api_key", "k1").counter("requests"),)) == (1,)
@pytest.mark.parametrize(
"client, refunded",
"replies, refunded",
[
({"{api_key:k1}:window": ValueError("script blew up")}, [("{team:t1}:requests", -1)]),
(
FakeClient(
lambda command: (
ValueError("script blew up")
if command[0] == "EVALSHA" and command[3] == "{api_key:k1}:window"
else _lua_ok_replies(command)
)
),
[("{team:t1}:requests", -1.0)],
{"{api_key:k1}:window": ConnectionError("redis down"), "{team:t1}:window": ConnectionError("redis down")},
[],
),
(FakeClient(_lua_ok_replies, fail=ConnectionError("redis down")), []),
],
ids=["one_group_failed", "pipeline_failed"],
)
@pytest.mark.asyncio
async def test_fail_closed_rejects_when_a_pipelined_lua_group_cannot_be_verified(
client: FakeClient, refunded: list[tuple[str, float]]
replies: dict[str, object], refunded: list[tuple[str, float]]
):
limiter = _limiter(FakeRedisCache(client), fail_closed=True)
limiter, fakes = _limiter(replies, fail_closed=True)
with request_redis_batch_scope(), pytest.raises(RateLimitUnverifiableError) as exc:
await limiter.atomic_check_and_increment_by_n(
@ -255,14 +229,14 @@ async def test_fail_closed_rejects_when_a_pipelined_lua_group_cannot_be_verified
)
assert exc.value.status_code == 503
assert _refunds(limiter) == refunded
assert len(client.pipelines) == 1
assert _refunds(fakes) == refunded
@pytest.mark.asyncio
async def test_a_pipeline_failure_refunds_nothing_and_falls_back_to_in_memory_enforcement():
client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down"))
limiter = _limiter(FakeRedisCache(client))
limiter, fakes = _limiter(
{"{api_key:k1}:window": ConnectionError("redis down"), "{team:t1}:window": ConnectionError("redis down")}
)
with request_redis_batch_scope():
response = await limiter.atomic_check_and_increment_by_n(
@ -272,23 +246,69 @@ async def test_a_pipeline_failure_refunds_nothing_and_falls_back_to_in_memory_en
assert response["overall_code"] == "OK"
assert len(response["statuses"]) == 2
assert len(client.pipelines) == 1
assert _refunds(fakes) == []
@pytest.mark.asyncio
async def test_without_a_request_scope_descriptor_groups_run_the_script_directly_as_before():
client = FakeClient(_lua_ok_replies)
limiter = _limiter(FakeRedisCache(client))
limiter.check_and_increment_by_n_script = AsyncMock(return_value=[0, 1, 1700000000])
async def test_without_a_request_scope_descriptor_groups_run_one_at_a_time_and_stop_at_the_first_denial():
calls: Final[list[str]] = [] # mutable-ok: one window key per script call
async def check(keys: list[str], args: list[object]) -> object:
calls.append(keys[0])
return _over_limit(11, 10) if keys[0] == "{team:t1}:window" else _OK
limiter, fakes = _limiter()
fakes.check_and_increment_by_n_script = check
response = await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10), _descriptor("team", "t1", 20)],
increments=[{"requests": 1}, {"requests": 1}],
descriptors=[
_descriptor("api_key", "k1", 10),
_descriptor("team", "t1", 20),
_descriptor("model_per_key", "k1:gpt", 5),
],
increments=[{"requests": 1}, {"requests": 1}, {"requests": 1}],
)
assert response["overall_code"] == "OVER_LIMIT"
assert calls == ["{api_key:k1}:window", "{team:t1}:window"]
assert _refunds(fakes) == [("{api_key:k1}:requests", -1)]
@dataclass
class _FakeCooldownsHandle:
"""A connected cooldown store whose ``declare_active`` lands on the request's fake native batch."""
def declare_active(
self, request: FakeNativeRequest, model_ids: Sequence[str]
) -> FakeNativePending[ActiveCooldowns]:
request.events.append("declare cooldowns")
return FakeNativePending(request, lambda: ())
@pytest.mark.asyncio
async def test_the_admission_script_and_the_armed_cooldown_read_go_out_in_one_native_pipeline():
async def window(keys: list[str], args: list[object]) -> object:
return ["1700000000", 1]
fakes: Final = FakeRateLimits(batch_rate_limiter_script=window)
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(dual_cache=DualCache()), rate_limit_store=fakes.store()
)
cooldowns: Final = RouterCooldownStore(
FakeRedisCache(FakeClient(_lua_ok_replies)), handle=lambda _c: _FakeCooldownsHandle()
) # pyright: ignore[reportArgumentType] # test double
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
arm_cooldowns(cooldowns, ["dep-a"])
response = await limiter.should_rate_limit([_descriptor("api_key", "k1", 10)])
active = await read_cooldowns(cooldowns, ["dep-a"], None)
assert response["overall_code"] == "OK"
assert limiter.check_and_increment_by_n_script.await_count == 2
assert client.pipelines == []
assert active == ()
assert len(natives.built) == 1
assert _declared_before_the_first_flush(natives) == ["declare cooldowns", "declare sliding_window"]
assert request.prefetched == {}
def _deployment(deployment_id: str) -> dict:
@ -374,24 +394,19 @@ async def _route(router: Router) -> str:
@pytest.mark.asyncio
@pytest.mark.usefixtures("pinned_clock")
async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own():
async def test_armed_routing_read_is_one_pipeline_and_routing_issues_no_read_of_its_own():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
limiter = _limiter(redis_cache)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10), _descriptor("team", "t1", 20)],
increments=[{"requests": 1}, {"requests": 1}],
)
deployment = await _route(router)
assert deployment in {"dep-a", "dep-b"}
assert len(client.pipelines) == 1
commands = client.pipelines[0].commands
assert [c[0] for c in commands] == ["MGET", "EVALSHA", "EVALSHA"]
assert [c[0] for c in commands] == ["MGET"]
assert frozenset(commands[0][1:]) == _counter_keys(router)
assert redis_cache.alone == []
@ -504,14 +519,10 @@ async def test_an_abandoned_prefetch_still_backfills_the_counters_it_read():
redis_cache: Final = FakeRedisCache(client)
router: Final = _router(redis_cache)
client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100})
limiter: Final = _limiter(redis_cache)
with request_redis_batch_scope():
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
await request.flush_all()
after_abandoned: Final = len(_counter_mgets(router, client, redis_cache))
with request_redis_batch_scope():
@ -587,30 +598,23 @@ async def test_simple_shuffle_arms_no_python_prefetch_and_reads_no_cooldown_key(
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache, routing_strategy="simple-shuffle")
limiter = _limiter(redis_cache)
with request_redis_batch_scope() as request:
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
assert set(request.prefetched) == {"routing_cooldowns"}
await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
deployment = await _route(router)
assert deployment in {"dep-a", "dep-b"}
assert [[c[0] for c in pipeline.commands] for pipeline in client.pipelines] == [["EVALSHA"]]
assert client.pipelines == []
assert redis_cache.alone == []
@pytest.mark.asyncio
@pytest.mark.usefixtures("pinned_clock")
@pytest.mark.parametrize("with_limiter", [True, False])
async def test_requests_within_the_read_interval_read_the_counters_from_redis_once(with_limiter: bool):
async def test_requests_within_the_read_interval_read_the_counters_from_redis_once():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
limiter = _limiter(redis_cache)
request_round_trips: list[tuple[int, int]] = [] # mutable-ok: one entry per simulated request
for _ in range(3):
@ -618,17 +622,11 @@ async def test_requests_within_the_read_interval_read_the_counters_from_redis_on
alone_count = len(redis_cache.alone)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
if with_limiter:
await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
await _route(router)
request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count))
assert len(_counter_mgets(router, client, redis_cache)) == 1
if not with_limiter:
assert request_round_trips[1:] == [(0, 0), (0, 0)]
assert request_round_trips[1:] == [(0, 0), (0, 0)]
@pytest.mark.asyncio
@ -712,29 +710,6 @@ async def test_two_backends_flush_concurrently_one_pipeline_each():
assert len(a_client.pipelines) == 1 and len(b_client.pipelines) == 1
@pytest.mark.asyncio
@pytest.mark.usefixtures("pinned_clock")
async def test_a_single_lua_group_rides_the_pipeline_with_the_armed_routing_read():
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
router = _router(redis_cache)
limiter = _limiter(redis_cache)
with request_redis_batch_scope():
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
await limiter.atomic_check_and_increment_by_n(
descriptors=[_descriptor("api_key", "k1", 10)],
increments=[{"requests": 1}],
)
await router.async_get_available_deployment(
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
)
assert len(client.pipelines) == 1
assert [c[0] for c in client.pipelines[0].commands] == ["MGET", "EVALSHA"]
assert redis_cache.alone == []
class _SameServerCache(FakeRedisCache):
def __init__(self, client: FakeClient, namespace: str | None = None, **redis_kwargs: object) -> None:
super().__init__(client, namespace)
@ -874,17 +849,6 @@ async def test_identity_prefetch_is_one_read_after_which_hits_and_misses_alike_c
assert cache.get_local("end_user_id:eu-miss") is None
@dataclass
class _FakeCooldownsHandle:
"""A connected cooldown store whose ``declare_active`` lands on the request's fake native batch."""
def declare_active(
self, request: FakeNativeRequest, model_ids: Sequence[str]
) -> FakeNativePending[ActiveCooldowns]:
request.events.append("declare cooldowns")
return FakeNativePending(request, lambda: ())
@pytest.mark.asyncio
async def test_the_auth_prefetch_and_the_armed_cooldown_read_go_out_in_one_native_pipeline():
cache, _fake = shared_cache(FakeAuthCache(redis={"key-hit": {"k": "key-hit"}}))

View file

@ -28,6 +28,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.redis.stores.rate_limits import RateLimitScope
from litellm.types.llms.openai import ChatCompletionUserMessage, LiteLLMBatchCreateRequest
@ -48,8 +49,7 @@ def _make_limiters(clock: _Clock | None = None):
async def _counter(internal_usage_cache, rate_limiter, descriptor_key, value, rate_limit_type):
cache_key = rate_limiter.create_rate_limit_keys(descriptor_key, value, rate_limit_type)
raw = await internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None, local_only=True)
(raw,) = rate_limiter.rate_limits.local((RateLimitScope(descriptor_key, value).counter(rate_limit_type),))
return int(raw or 0)
@ -72,7 +72,7 @@ async def test_batch_over_rpm_and_tpm_but_under_tpd_is_accepted():
@pytest.mark.asyncio
async def test_cumulative_batch_tokens_over_tpd_returns_429_with_remaining_daily_window():
window_start = datetime(2026, 9, 13, 8, 0, 0)
window_start = datetime(2026, 9, 13, 8, 0, 0, tzinfo=timezone.utc)
clock = _Clock(window_start)
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock)
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("tpd-key-2"), rpm_limit=1, tpd_limit=1000)
@ -82,7 +82,7 @@ async def test_cumulative_batch_tokens_over_tpd_returns_429_with_remaining_daily
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
clock.now = datetime(2026, 9, 13, 11, 0, 0)
clock.now = datetime(2026, 9, 13, 11, 0, 0, tzinfo=timezone.utc)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,

View file

@ -19,6 +19,11 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler,
)
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.redis.stores.rate_limits import RateLimitScope
from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits
class TimeController:
@ -92,9 +97,9 @@ async def test_priority_weight_allocation(monkeypatch):
expected_high_tpm = int(total_tpm * 0.9) # 900
actual_high_tpm = high_descriptor["rate_limit"]["tokens_per_unit"]
assert (
actual_high_tpm == expected_high_tpm
), f"High priority should get {expected_high_tpm} TPM (90%), got {actual_high_tpm}"
assert actual_high_tpm == expected_high_tpm, (
f"High priority should get {expected_high_tpm} TPM (90%), got {actual_high_tpm}"
)
assert high_descriptor["value"] == f"{model}:high"
# Test low priority allocation
@ -112,17 +117,15 @@ async def test_priority_weight_allocation(monkeypatch):
expected_low_tpm = int(total_tpm * 0.1) # 100
actual_low_tpm = low_descriptor["rate_limit"]["tokens_per_unit"]
assert (
actual_low_tpm == expected_low_tpm
), f"Low priority should get {expected_low_tpm} TPM (10%), got {actual_low_tpm}"
assert actual_low_tpm == expected_low_tpm, (
f"Low priority should get {expected_low_tpm} TPM (10%), got {actual_low_tpm}"
)
assert low_descriptor["value"] == f"{model}:low"
# Verify the ratio is 9:1, not 1:1 (equal splitting)
ratio = actual_high_tpm / actual_low_tpm
expected_ratio = 9.0
assert (
abs(ratio - expected_ratio) < 0.1
), f"High:Low ratio should be {expected_ratio}:1, got {ratio}:1"
assert abs(ratio - expected_ratio) < 0.1, f"High:Low ratio should be {expected_ratio}:1, got {ratio}:1"
@pytest.mark.asyncio
@ -188,8 +191,7 @@ async def test_concurrent_priority_requests(monkeypatch):
descriptor = descriptors[0]
# Each high priority user should get 900 TPM, not divided by 3
assert descriptor["rate_limit"]["tokens_per_unit"] == 900, (
f"High priority user {user.user_id} should get 900 TPM, "
f"got {descriptor['rate_limit']['tokens_per_unit']}"
f"High priority user {user.user_id} should get 900 TPM, got {descriptor['rate_limit']['tokens_per_unit']}"
)
assert descriptor["value"] == f"{model}:high"
@ -205,8 +207,7 @@ async def test_concurrent_priority_requests(monkeypatch):
descriptor = descriptors[0]
# Each low priority user should get 100 TPM, not divided by 2
assert descriptor["rate_limit"]["tokens_per_unit"] == 100, (
f"Low priority user {user.user_id} should get 100 TPM, "
f"got {descriptor['rate_limit']['tokens_per_unit']}"
f"Low priority user {user.user_id} should get 100 TPM, got {descriptor['rate_limit']['tokens_per_unit']}"
)
assert descriptor["value"] == f"{model}:low"
@ -228,9 +229,7 @@ async def test_100_concurrent_priority_requests(time_controller, monkeypatch):
litellm.priority_reservation = {"high": 0.9, "low": 0.1}
dual_cache = DualCache()
handler = DynamicRateLimitHandler(
internal_usage_cache=dual_cache, time_provider=time_controller.now
)
handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache, time_provider=time_controller.now)
model = "stress-test-model"
total_tpm = 1000
@ -278,22 +277,20 @@ async def test_100_concurrent_priority_requests(time_controller, monkeypatch):
priority=priority,
)
assert (
len(descriptors) == 1
), f"User {user.user_id} should have exactly 1 descriptor"
assert len(descriptors) == 1, f"User {user.user_id} should have exactly 1 descriptor"
descriptor = descriptors[0]
# Validate TPM allocation
actual_tpm = descriptor["rate_limit"]["tokens_per_unit"]
assert (
actual_tpm == expected_tpm
), f"User {user.user_id} ({priority}) should get {expected_tpm} TPM, got {actual_tpm}"
assert actual_tpm == expected_tpm, (
f"User {user.user_id} ({priority}) should get {expected_tpm} TPM, got {actual_tpm}"
)
# Validate RPM allocation
actual_rpm = descriptor["rate_limit"]["requests_per_unit"]
assert (
actual_rpm == expected_rpm
), f"User {user.user_id} ({priority}) should get {expected_rpm} RPM, got {actual_rpm}"
assert actual_rpm == expected_rpm, (
f"User {user.user_id} ({priority}) should get {expected_rpm} RPM, got {actual_rpm}"
)
# Validate descriptor key
assert descriptor["value"] == f"{model}:{priority}"
@ -312,9 +309,7 @@ async def test_100_concurrent_priority_requests(time_controller, monkeypatch):
# Split into batches to simulate requests over 10 seconds
batch_size = 10 # 10 requests per batch
batches = [
all_users[i : i + batch_size] for i in range(0, len(all_users), batch_size)
]
batches = [all_users[i : i + batch_size] for i in range(0, len(all_users), batch_size)]
all_results = []
@ -333,47 +328,29 @@ async def test_100_concurrent_priority_requests(time_controller, monkeypatch):
total_duration = end_time - start_time
# Validate that the test ran over approximately 10 seconds
assert (
total_duration >= 9.0
), f"Test should take ~10 seconds, took {total_duration:.2f}s"
assert total_duration >= 9.0, f"Test should take ~10 seconds, took {total_duration:.2f}s"
assert total_duration <= 15.0, f"Test took too long: {total_duration:.2f}s"
# Validate all requests were successful
successful_results = [
r for r in all_results if isinstance(r, dict) and r.get("success")
]
assert (
len(successful_results) == 100
), f"Expected 100 successful results, got {len(successful_results)}"
successful_results = [r for r in all_results if isinstance(r, dict) and r.get("success")]
assert len(successful_results) == 100, f"Expected 100 successful results, got {len(successful_results)}"
# Validate priority distribution
high_priority_results = [r for r in successful_results if r["priority"] == "high"]
low_priority_results = [r for r in successful_results if r["priority"] == "low"]
assert (
len(high_priority_results) == 70
), f"Expected 70 high priority results, got {len(high_priority_results)}"
assert (
len(low_priority_results) == 30
), f"Expected 30 low priority results, got {len(low_priority_results)}"
assert len(high_priority_results) == 70, f"Expected 70 high priority results, got {len(high_priority_results)}"
assert len(low_priority_results) == 30, f"Expected 30 low priority results, got {len(low_priority_results)}"
# Validate all high priority users got correct allocation
for result in high_priority_results:
assert (
result["tpm"] == 900
), f"High priority user {result['user_id']} got {result['tpm']} TPM, expected 900"
assert (
result["rpm"] == 450
), f"High priority user {result['user_id']} got {result['rpm']} RPM, expected 450"
assert result["tpm"] == 900, f"High priority user {result['user_id']} got {result['tpm']} TPM, expected 900"
assert result["rpm"] == 450, f"High priority user {result['user_id']} got {result['rpm']} RPM, expected 450"
# Validate all low priority users got correct allocation
for result in low_priority_results:
assert (
result["tpm"] == 100
), f"Low priority user {result['user_id']} got {result['tpm']} TPM, expected 100"
assert (
result["rpm"] == 50
), f"Low priority user {result['user_id']} got {result['rpm']} RPM, expected 50"
assert result["tpm"] == 100, f"Low priority user {result['user_id']} got {result['tpm']} TPM, expected 100"
assert result["rpm"] == 50, f"Low priority user {result['user_id']} got {result['rpm']} RPM, expected 50"
print(f"✅ Successfully processed 100 concurrent requests in {total_duration:.2f}s")
print(f" - 70 high priority users: 900 TPM, 450 RPM each")
@ -429,9 +406,7 @@ async def test_concurrent_pre_call_hooks_stress(monkeypatch):
return 1800 # 1800/2000 = 90% saturation
return None
async def mock_should_rate_limit(
descriptors, parent_otel_span=None, read_only=False
):
async def mock_should_rate_limit(descriptors, parent_otel_span=None, read_only=False):
"""Mock rate limiter that handles saturation-aware descriptors."""
descriptor = descriptors[0]
descriptor_key = descriptor["key"]
@ -540,9 +515,7 @@ async def test_concurrent_pre_call_hooks_stress(monkeypatch):
)
# If no exception, request was allowed
successful_requests.append(
{"user_id": user.user_id, "priority": priority, "result": "allowed"}
)
successful_requests.append({"user_id": user.user_id, "priority": priority, "result": "allowed"})
return {
"status": "success",
"user_id": user.user_id,
@ -551,9 +524,7 @@ async def test_concurrent_pre_call_hooks_stress(monkeypatch):
except Exception as e:
# Request was rate limited
rate_limited_requests.append(
{"user_id": user.user_id, "priority": priority, "error": str(e)}
)
rate_limited_requests.append({"user_id": user.user_id, "priority": priority, "error": str(e)})
return {
"status": "rate_limited",
"user_id": user.user_id,
@ -563,68 +534,44 @@ async def test_concurrent_pre_call_hooks_stress(monkeypatch):
# Run all 50 requests concurrently with patches applied to the entire batch
start_time = time.time()
with (
patch.object(
handler.v3_limiter, "should_rate_limit", side_effect=mock_should_rate_limit
),
patch.object(
handler.internal_usage_cache, "async_get_cache", side_effect=mock_get_cache
),
patch.object(handler.v3_limiter, "should_rate_limit", side_effect=mock_should_rate_limit),
patch.object(handler.internal_usage_cache, "async_get_cache", side_effect=mock_get_cache),
):
tasks = [make_request(user_data) for user_data in users]
results = await asyncio.gather(*tasks, return_exceptions=True)
end_time = time.time()
# Analyze results
successful_count = len(
[r for r in results if isinstance(r, dict) and r["status"] == "success"]
)
rate_limited_count = len(
[r for r in results if isinstance(r, dict) and r["status"] == "rate_limited"]
)
successful_count = len([r for r in results if isinstance(r, dict) and r["status"] == "success"])
rate_limited_count = len([r for r in results if isinstance(r, dict) and r["status"] == "rate_limited"])
# Validate that premium users were mostly successful (priority worked)
premium_results = [
r for r in results if isinstance(r, dict) and r["priority"] == "premium"
]
premium_results = [r for r in results if isinstance(r, dict) and r["priority"] == "premium"]
premium_success = len([r for r in premium_results if r["status"] == "success"])
standard_results = [
r for r in results if isinstance(r, dict) and r["priority"] == "standard"
]
standard_results = [r for r in results if isinstance(r, dict) and r["priority"] == "standard"]
standard_success = len([r for r in standard_results if r["status"] == "success"])
# Premium users should have higher success rate due to priority
premium_success_rate = (
premium_success / len(premium_results) if premium_results else 0
)
standard_success_rate = (
standard_success / len(standard_results) if standard_results else 0
)
premium_success_rate = premium_success / len(premium_results) if premium_results else 0
standard_success_rate = standard_success / len(standard_results) if standard_results else 0
assert (
premium_success_rate >= 0.9
), f"Premium success rate should be >= 90%, got {premium_success_rate:.2%}"
assert (
standard_success_rate >= 0.5
), f"Standard success rate should be >= 50% (with 30% random limiting, allows for variance), got {standard_success_rate:.2%}"
assert premium_success_rate >= 0.9, f"Premium success rate should be >= 90%, got {premium_success_rate:.2%}"
assert standard_success_rate >= 0.5, (
f"Standard success rate should be >= 50% (with 30% random limiting, allows for variance), got {standard_success_rate:.2%}"
)
# Allow for the case where both are 100% due to timing/mocking issues
# The test is inherently flaky due to random behavior
if premium_success_rate < 1.0 or standard_success_rate < 1.0:
assert (
premium_success_rate >= standard_success_rate
), "Premium should have >= success rate than standard"
assert premium_success_rate >= standard_success_rate, "Premium should have >= success rate than standard"
total_duration = end_time - start_time
print(f"✅ Processed 50 concurrent pre-call hooks in {total_duration:.2f}s")
print(
f" - Premium users: {premium_success}/{len(premium_results)} success ({premium_success_rate:.1%})"
)
print(
f" - Standard users: {standard_success}/{len(standard_results)} success ({standard_success_rate:.1%})"
)
print(f" - Total successful: {successful_count}/50 ({successful_count/50:.1%})")
print(f" - Premium users: {premium_success}/{len(premium_results)} success ({premium_success_rate:.1%})")
print(f" - Standard users: {standard_success}/{len(standard_results)} success ({standard_success_rate:.1%})")
print(f" - Total successful: {successful_count}/50 ({successful_count / 50:.1%})")
print(f" - Priority system working: Premium > Standard success rates")
@ -732,28 +679,24 @@ async def test_fake_calls_case_1_no_rate_limiting_at_capacity(monkeypatch):
print(f"Test Case 1 - Saturation-Aware Rate Limiting:")
print(f" - Duration: {end_time - start_time:.2f}s")
print(f" - Key A: {successful_requests['key_a']}/1 successful (reserved 75 RPM)")
print(
f" - Key B: {successful_requests['key_b']}/100 successful (reserved 25 RPM)"
)
print(f" - Key B: {successful_requests['key_b']}/100 successful (reserved 25 RPM)")
print(f" - Total successful: {total_successful}/101")
print(f" - Total rate limited: {total_rate_limited}/101")
# Key A should get its 1 request
assert (
successful_requests["key_a"] == 1
), f"Key A should get 1 request, got {successful_requests['key_a']}"
assert successful_requests["key_a"] == 1, f"Key A should get 1 request, got {successful_requests['key_a']}"
# Key B can send until saturation hits 50% (which is ~50 total requests)
# After that, strict mode enforces its 25 RPM reservation
# Due to race conditions in concurrent execution, allow 45-52 successful requests
assert (
45 <= successful_requests["key_b"] <= 52
), f"Key B should get ~49 requests (45-52), got {successful_requests['key_b']}"
assert 45 <= successful_requests["key_b"] <= 52, (
f"Key B should get ~49 requests (45-52), got {successful_requests['key_b']}"
)
# Verify approximately half of key_b requests were rate limited
assert (
rate_limited_requests["key_b"] >= 45
), f"Key B should have ≥45 rate limited requests, got {rate_limited_requests['key_b']}"
assert rate_limited_requests["key_b"] >= 45, (
f"Key B should have ≥45 rate limited requests, got {rate_limited_requests['key_b']}"
)
@pytest.mark.asyncio
@ -855,32 +798,24 @@ async def test_fake_calls_case_2_priority_queue_during_saturation(monkeypatch):
print(f"Test Case 2 - Priority Queue Behavior During Saturation:")
print(f" - Duration: {end_time - start_time:.2f}s")
print(
f" - Key A: {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})"
)
print(
f" - Key B: {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})"
)
print(f" - Key A: {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})")
print(f" - Key B: {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})")
print(f" - Total successful: {total_successful}/400")
# Key A should get significantly more requests than Key B (75:25 ratio)
assert (
key_a_success_rate > key_b_success_rate
), f"Key A should have higher success rate: {key_a_success_rate:.1%} vs {key_b_success_rate:.1%}"
assert key_a_success_rate > key_b_success_rate, (
f"Key A should have higher success rate: {key_a_success_rate:.1%} vs {key_b_success_rate:.1%}"
)
# Check ratio is approximately 3:1 (75:25)
if total_successful > 0:
key_a_share = successful_requests["key_a"] / total_successful
expected_key_a_share = 0.75
print(
f" - Key A got {key_a_share:.1%} of successful requests (expected ~75%)"
)
print(f" - Key A got {key_a_share:.1%} of successful requests (expected ~75%)")
# Allow tolerance for timing effects
assert (
abs(key_a_share - expected_key_a_share) < 0.2
), f"Key A share should be ~75%, got {key_a_share:.1%}"
assert abs(key_a_share - expected_key_a_share) < 0.2, f"Key A share should be ~75%, got {key_a_share:.1%}"
@pytest.mark.asyncio
@ -1003,22 +938,12 @@ async def test_fake_calls_case_3_spillover_capacity_default_keys(monkeypatch):
print(f" - Total successful: {total_successful}/600")
# Key A should get the most requests (75% of capacity)
assert (
successful_requests["key_a"] > successful_requests["key_b"]
), "Key A should get more than Key B"
assert (
successful_requests["key_a"] > successful_requests["key_c"]
), "Key A should get more than Key C"
assert (
successful_requests["key_a"] > successful_requests["key_d"]
), "Key A should get more than Key D"
assert successful_requests["key_a"] > successful_requests["key_b"], "Key A should get more than Key B"
assert successful_requests["key_a"] > successful_requests["key_c"], "Key A should get more than Key C"
assert successful_requests["key_a"] > successful_requests["key_d"], "Key A should get more than Key D"
# Default keys should get similar amounts (spillover capacity)
avg_default = (
successful_requests["key_b"]
+ successful_requests["key_c"]
+ successful_requests["key_d"]
) / 3
avg_default = (successful_requests["key_b"] + successful_requests["key_c"] + successful_requests["key_d"]) / 3
print(f" - Average default key success: {avg_default:.1f}")
@ -1124,32 +1049,27 @@ async def test_fake_calls_case_4_over_allocated_with_normalization(monkeypatch):
print(f"Test Case 4 - Over-Allocated Priority Reservations with Normalization:")
print(f" - Duration: {end_time - start_time:.2f}s")
print(
f" - Key A (0.60): {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})"
)
print(
f" - Key B (0.80): {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})"
)
print(f" - Key A (0.60): {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})")
print(f" - Key B (0.80): {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})")
print(f" - Total successful: {total_successful}/400")
# With saturation-aware behavior:
# 1. Verify total capacity is reasonably bounded (not all 400 requests succeed)
assert (
total_successful < 300
), f"Total requests should be bounded by saturation detection, got {total_successful}/400"
assert total_successful < 300, (
f"Total requests should be bounded by saturation detection, got {total_successful}/400"
)
# 2. Verify significant rate limiting occurred (at least 50% blocked)
assert (
total_successful < 200
), f"At least 50% of requests should be rate limited, got {total_successful}/400 successful"
assert total_successful < 200, (
f"At least 50% of requests should be rate limited, got {total_successful}/400 successful"
)
# 3. Verify both keys got some requests through (normalization is working)
assert successful_requests["key_a"] > 0, "Key A should get some requests"
assert successful_requests["key_b"] > 0, "Key B should get some requests"
print(
f" - Normalization test PASSED: Both priorities got requests, "
f"total bounded to {total_successful} (under 200)"
f" - Normalization test PASSED: Both priorities got requests, total bounded to {total_successful} (under 200)"
)
@ -1273,24 +1193,16 @@ async def test_fake_calls_case_5_default_value_priority_reservation(monkeypatch)
print(f" - Total successful: {total_successful}/600")
# Verify priority ordering: A > B > C ≈ D
assert (
successful_requests["key_a"] > successful_requests["key_b"]
), "Key A should get more than Key B"
assert (
successful_requests["key_b"] > successful_requests["key_c"]
), "Key B should get more than Key C"
assert successful_requests["key_a"] > successful_requests["key_b"], "Key A should get more than Key B"
assert successful_requests["key_b"] > successful_requests["key_c"], "Key B should get more than Key C"
# Key C and Key D should get similar amounts (both have 0.05 priority)
key_c_vs_d_ratio = successful_requests["key_c"] / max(
successful_requests["key_d"], 1
)
key_c_vs_d_ratio = successful_requests["key_c"] / max(successful_requests["key_d"], 1)
print(f" - Key C vs Key D ratio: {key_c_vs_d_ratio:.2f} (expected ~1.0)")
if total_successful > 0:
key_a_share = successful_requests["key_a"] / total_successful
print(
f" - Key A got {key_a_share:.1%} of successful requests (expected ~55-62%)"
)
print(f" - Key A got {key_a_share:.1%} of successful requests (expected ~55-62%)")
@pytest.mark.asyncio
@ -1342,15 +1254,9 @@ async def test_default_priority_shared_pool(monkeypatch):
user_c.user_id = "user_c"
# Get descriptors for each
desc_a = handler._create_priority_based_descriptors(
model=model, user_api_key_dict=user_a, priority=None
)
desc_b = handler._create_priority_based_descriptors(
model=model, user_api_key_dict=user_b, priority=None
)
desc_c = handler._create_priority_based_descriptors(
model=model, user_api_key_dict=user_c, priority=None
)
desc_a = handler._create_priority_based_descriptors(model=model, user_api_key_dict=user_a, priority=None)
desc_b = handler._create_priority_based_descriptors(model=model, user_api_key_dict=user_b, priority=None)
desc_c = handler._create_priority_based_descriptors(model=model, user_api_key_dict=user_c, priority=None)
# All should use the SAME shared pool key
assert desc_a[0]["value"] == f"{model}:default_pool"
@ -1365,9 +1271,7 @@ async def test_default_priority_shared_pool(monkeypatch):
# Verify explicit priority uses different pool
user_prod = UserAPIKeyAuth()
user_prod.metadata = {"priority": "prod"}
desc_prod = handler._create_priority_based_descriptors(
model=model, user_api_key_dict=user_prod, priority="prod"
)
desc_prod = handler._create_priority_based_descriptors(model=model, user_api_key_dict=user_prod, priority="prod")
assert desc_prod[0]["value"] == f"{model}:prod"
assert desc_prod[0]["rate_limit"]["requests_per_unit"] == 75
@ -1421,8 +1325,8 @@ async def test_async_log_success_event_increments_by_actual_tokens(monkeypatch):
for op in pipeline_operations:
increment_calls.append(
{
"key": op["key"],
"increment_value": op["increment_value"],
"counter": op.counter,
"increment_value": op.amount,
}
)
@ -1460,24 +1364,22 @@ async def test_async_log_success_event_increments_by_actual_tokens(monkeypatch):
)
# Verify increments happened with actual token count (60 total tokens)
assert (
len(increment_calls) == 2
), f"Expected 2 increment calls, got {len(increment_calls)}"
assert len(increment_calls) == 2, f"Expected 2 increment calls, got {len(increment_calls)}"
# Both should increment by 50 (total_tokens, since rate_limit_type defaults to 'total')
for call in increment_calls:
assert (
call["increment_value"] == 60
), f"Expected increment of 60 tokens, got {call['increment_value']} for key {call['key']}"
assert call["increment_value"] == 60, (
f"Expected increment of 60 tokens, got {call['increment_value']} for counter {call['counter']}"
)
# Verify correct keys were used
keys = [call["key"] for call in increment_calls]
assert any(
"model_saturation_check" in k for k in keys
), "Should increment model_saturation_check"
assert any(
"priority_model" in k and "dev" in k for k in keys
), "Should increment priority_model with 'dev' priority"
# Verify correct counters were used
counters = [call["counter"] for call in increment_calls]
assert RateLimitScope("model_saturation_check", model).counter("tokens") in counters, (
"Should increment model_saturation_check"
)
assert RateLimitScope("priority_model", f"{model}:dev").counter("tokens") in counters, (
"Should increment priority_model with 'dev' priority"
)
@pytest.mark.asyncio
@ -1518,66 +1420,44 @@ async def test_saturation_check_cache_ttl_configuration(monkeypatch):
handler.update_variables(llm_router=llm_router)
# Verify the TTL getter returns configured value
assert (
handler._get_saturation_check_cache_ttl() == 5
), "TTL should be configurable via priority_reservation_settings"
# Track async_get_cache calls to verify TTL is passed
get_cache_calls = []
original_get_cache = handler.internal_usage_cache.async_get_cache
async def mock_get_cache(
key, litellm_parent_otel_span=None, local_only=False, **kwargs
):
get_cache_calls.append(
{
"key": key,
"ttl": kwargs.get("ttl"),
"local_only": local_only,
}
)
return None # Simulate cache miss
handler.internal_usage_cache.async_get_cache = mock_get_cache
# Call _get_saturation_value_from_cache
counter_key = handler.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model,
rate_limit_type="requests",
assert handler._get_saturation_check_cache_ttl() == 5, (
"TTL should be configurable via priority_reservation_settings"
)
await handler._get_saturation_value_from_cache(counter_key=counter_key)
# Track read-through (memory then Redis) gets to verify TTL is passed
get_cache_calls = []
# Verify async_get_cache was called with the configured TTL
class _RecordingRateLimits(FakeRateLimits):
async def get(self, entry, *, backfill_ttl_seconds=None):
get_cache_calls.append({"entry": entry, "ttl": backfill_ttl_seconds})
return None # Simulate cache miss
handler.v3_limiter = PROXY_MaxParallelRequestsHandler_v3(
handler.internal_usage_cache,
rate_limit_store=_RecordingRateLimits().store(),
)
# Call _get_saturation_value_from_cache
counter = RateLimitScope("model_saturation_check", model).counter("requests")
await handler._get_saturation_value_from_cache(counter)
# Verify the read-through get was called with the configured TTL
assert len(get_cache_calls) == 1, "Expected 1 cache call"
assert (
get_cache_calls[0]["ttl"] == 5
), f"Expected TTL of 5 seconds, got {get_cache_calls[0]['ttl']}"
assert (
get_cache_calls[0]["local_only"] is False
), "Should check Redis (local_only=False) for multi-node consistency"
assert get_cache_calls[0]["entry"] == ("requests", "model_saturation_check", model)
assert get_cache_calls[0]["ttl"] == 5, f"Expected TTL of 5 seconds, got {get_cache_calls[0]['ttl']}"
# Test with different TTL value
get_cache_calls.clear()
litellm.priority_reservation_settings.saturation_check_cache_ttl = 30
await handler._get_saturation_value_from_cache(counter_key=counter_key)
await handler._get_saturation_value_from_cache(counter)
assert (
get_cache_calls[0]["ttl"] == 30
), f"TTL should update to 30 seconds, got {get_cache_calls[0]['ttl']}"
assert get_cache_calls[0]["ttl"] == 30, f"TTL should update to 30 seconds, got {get_cache_calls[0]['ttl']}"
print("Saturation check cache TTL test passed:")
print(
" - TTL is configurable via priority_reservation_settings.saturation_check_cache_ttl"
)
print(
" - TTL is passed to async_get_cache for local cache expiration control"
)
print(
" - local_only=False ensures Redis is checked for multi-node consistency"
)
print(" - TTL is configurable via priority_reservation_settings.saturation_check_cache_ttl")
print(" - TTL is passed to the read-through get for local cache expiration control")
finally:
# Restore original TTL
@ -1623,7 +1503,7 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(mon
async def mock_increment(pipeline_operations, parent_otel_span=None):
for op in pipeline_operations:
incremented_keys.append(op["key"])
incremented_keys.append(op.counter.scope)
handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment
@ -1662,19 +1542,13 @@ async def test_async_log_success_event_uses_team_priority_from_auth_metadata(mon
)
# Verify the priority_model key uses 'team_priority' (not 'default_pool')
priority_keys = [k for k in incremented_keys if "priority_model" in k]
assert (
len(priority_keys) == 1
), f"Expected 1 priority_model key, got {len(priority_keys)}"
priority_keys = [scope.descriptor_value for scope in incremented_keys if scope.descriptor_key == "priority_model"]
assert len(priority_keys) == 1, f"Expected 1 priority_model key, got {len(priority_keys)}"
# The key should contain 'team_priority', not 'default_pool'
assert "team_priority" in priority_keys[0], (
f"Expected priority key to use 'team_priority' from user_api_key_auth_metadata, "
f"got key: {priority_keys[0]}"
# The key should use 'team_priority', not 'default_pool'
assert priority_keys[0] == f"{model}:team_priority", (
f"Expected priority key to use 'team_priority' from user_api_key_auth_metadata, got key: {priority_keys[0]}"
)
assert (
"default_pool" not in priority_keys[0]
), f"Priority key should NOT use 'default_pool', should use team's priority. Got: {priority_keys[0]}"
@pytest.mark.asyncio

File diff suppressed because it is too large Load diff

View file

@ -34,6 +34,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.redis.stores.rate_limits import RateLimitScope
def _make_phase1_barrier(num_concurrent: int, timeout: float = 0.1):
@ -89,9 +90,7 @@ async def test_batch_limiter_concurrent_bypasses_tpm_via_toctou():
dual_cache = DualCache()
internal_usage_cache = InternalUsageCache(dual_cache=dual_cache)
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=internal_usage_cache
)
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
@ -142,9 +141,7 @@ async def test_batch_limiter_uses_atomic_check_and_increment():
"""
dual_cache = DualCache()
internal_usage_cache = InternalUsageCache(dual_cache=dual_cache)
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=internal_usage_cache
)
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
@ -176,13 +173,11 @@ async def test_batch_limiter_uses_atomic_check_and_increment():
)
assert "atomic_check_and_increment_by_n" in call_log, (
f"Batch limiter must route through atomic_check_and_increment_by_n. "
f"Calls observed: {call_log}"
f"Batch limiter must route through atomic_check_and_increment_by_n. Calls observed: {call_log}"
)
legacy_calls = [c for c in call_log if c.startswith("should_rate_limit(")]
assert not legacy_calls, (
f"Batch limiter must not call should_rate_limit directly (legacy "
f"two-phase pattern). Observed: {legacy_calls}"
f"Batch limiter must not call should_rate_limit directly (legacy two-phase pattern). Observed: {legacy_calls}"
)
@ -337,8 +332,7 @@ async def test_dynamic_rate_limiter_v3_uses_atomic_check_and_increment(monkeypat
"separate-increment pattern)."
)
assert "model_saturation_check" in atomic_descriptors_observed[0], (
f"Expected model_saturation_check in atomic descriptor set. "
f"Got: {atomic_descriptors_observed}"
f"Expected model_saturation_check in atomic descriptor set. Got: {atomic_descriptors_observed}"
)
@ -356,9 +350,7 @@ async def test_batch_zero_token_consumes_rpm_only():
"""
dual_cache = DualCache()
internal_usage_cache = InternalUsageCache(dual_cache=dual_cache)
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=internal_usage_cache
)
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
@ -379,18 +371,8 @@ async def test_batch_zero_token_consumes_rpm_only():
)
# Inspect counters: RPM key incremented to 3, TPM key absent (or 0).
rpm_key = rate_limiter.create_rate_limit_keys(
"api_key", user_api_key_dict.api_key or "", "requests"
)
tpm_key = rate_limiter.create_rate_limit_keys(
"api_key", user_api_key_dict.api_key or "", "tokens"
)
rpm_val = await internal_usage_cache.async_get_cache(
key=rpm_key, litellm_parent_otel_span=None, local_only=True
)
tpm_val = await internal_usage_cache.async_get_cache(
key=tpm_key, litellm_parent_otel_span=None, local_only=True
)
scope = RateLimitScope("api_key", user_api_key_dict.api_key or "")
rpm_val, tpm_val = rate_limiter.rate_limits.local((scope.counter("requests"), scope.counter("tokens")))
assert int(rpm_val or 0) == 3, f"RPM counter must reach 3, got {rpm_val}"
assert tpm_val in (
None,
@ -479,6 +461,4 @@ async def test_dynamic_rate_limiter_v3_fails_closed_on_unknown_descriptor(monkey
priority="high",
saturation=0.0,
)
assert (
exc.value.status_code == 429
), f"Expected 429 fail-closed on unknown descriptor; got {exc.value.status_code}"
assert exc.value.status_code == 429, f"Expected 429 fail-closed on unknown descriptor; got {exc.value.status_code}"

File diff suppressed because it is too large Load diff

View file

@ -33,6 +33,7 @@ from litellm.proxy.mcp_server.auth.user_api_key_auth_mcp import (
_agent_capped_servers,
is_mcp_admitted_user_subject,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.types.agents import AgentCaller
@ -9089,7 +9090,7 @@ class TestUserSubjectTeamUnion:
assert billed is not None and billed.team_id == "t1", "throttling and billing pick the same source"
descriptors: list = []
limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock())
limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache()))
limiter._add_mcp_per_team_rate_limit_descriptor(auth, "srv1", descriptors)
charged = {d["value"]: d["rate_limit"]["requests_per_unit"] for d in descriptors}
assert charged == {"t1:srv1": 5}, "only the attributing team's bucket is charged"
@ -12452,7 +12453,6 @@ async def test_personal_key_display_name_billing_uses_original_tool_grant(
async def test_personal_key_display_name_billing_uses_allowed_servers(
monkeypatch: pytest.MonkeyPatch, personal_key_owner_teams: UserAPIKeyAuth
) -> None:
from litellm.proxy import proxy_server
from litellm.proxy.auth import auth_checks
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.types.mcp import MCPTransport
@ -12845,7 +12845,7 @@ async def test_personal_key_team_mcp_rpm_limit_charges_granting_team(
stamped: Final = personal_key_owner_teams.model_copy(update={"mcp_source_team_rpm_limits": limits})
descriptors: list[RateLimitDescriptor] = []
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock())
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache()))
limiter._add_mcp_per_team_rate_limit_descriptor(stamped, server_b.name, descriptors)
assert len(descriptors) == 1
assert descriptors[0]["key"] == "mcp_per_team"
@ -13224,8 +13224,8 @@ async def test_admin_owned_team_key_reaches_only_its_team_mcp_servers(
expected_servers: set[str],
credential_shape: dict[str, object],
) -> None:
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
from litellm.proxy.auth import auth_checks
from litellm.proxy.mcp_server.mcp_server_manager import global_mcp_server_manager
async def _get_team_object(team_id: str, **_: object) -> LiteLLM_TeamTable:
return _personal_key_team(team_id) if team_id in _PERSONAL_KEY_TEAMS else LiteLLM_TeamTable(team_id=team_id)

View file

@ -12722,6 +12722,7 @@ async def test_completed_synthetic_stream_reuses_capacity_without_freeing_anothe
from litellm.proxy.common_utils.sse_keepalive import SSE_COMMENT_PING_BYTES
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.redis.stores.rate_limits import RateLimitScope
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
def clock() -> datetime.datetime:
@ -12734,7 +12735,7 @@ async def test_completed_synthetic_stream_reuses_capacity_without_freeing_anothe
logger: Final = ProxyLogging(user_api_key_cache=cache)
logger.proxy_hook_mapping["parallel_request_limiter"] = limiter
auth: Final = ProxyUserAPIKeyAuth(api_key=hash_token("sk-synthetic-stream"), max_parallel_requests=2)
counter: Final = f"{{api_key:{auth.api_key}}}:max_parallel_requests"
gauge: Final = RateLimitScope("api_key", auth.api_key).gauge
monkeypatch.setattr(litellm, "callbacks", [])
monkeypatch.setattr(proxy_server, "proxy_logging_obj", logger)
@ -12790,7 +12791,7 @@ async def test_completed_synthetic_stream_reuses_capacity_without_freeing_anothe
await contextvars.Context().run(asyncio.create_task, admit("still-streaming"))
for index in range(4):
await contextvars.Context().run(asyncio.create_task, request(f"completed-{index}"))
occupied: Final = await cache.async_get_cache(counter)
occupied: Final = limiter.rate_limits.local((gauge,))[0]
assert isinstance(occupied, dict) and len(occupied) == 1, occupied

View file

@ -0,0 +1,261 @@
"""A ``RateLimitsHandle`` whose memory tier is the real native store and whose Redis scripts are Python fakes
called the way ``main`` called its registered Lua: ``script(keys=[...], args=[...])`` with ``main``'s key strings,
answering ``main``'s raw replies. Tests keep asserting the wire shape while the limiter talks to the store."""
from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field
from typing import Final
from pydantic import TypeAdapter
from litellm.redis.stores.rate_limits import (
LocalValue,
NativeAtomicCounter,
NativeCounter,
NativeEntry,
NativeIncrement,
NativeScope,
RateLimitStore,
)
from litellm.rust_bridge._native import NativeRateLimits
RawScript = Callable[..., Awaitable[object]]
RawValue = int | float | str | bytes | None
_RAW: Final = TypeAdapter(list[RawValue])
def scope_key(scope: NativeScope, suffix: str) -> str:
"""``main``'s ``f"{{{key}:{value}}}:{suffix}"``."""
return f"{{{scope[0]}:{scope[1]}}}:{suffix}"
def counter_key(counter: NativeCounter) -> str:
return scope_key((counter[0], counter[1]), counter[2])
def window_key(counter: NativeCounter) -> str:
return scope_key((counter[0], counter[1]), "window")
class Resolved:
"""A script call; declared on a request that logs ``events`` (``FakeNativeRequest``), it notes the declaration
there and refuses to answer before that request's pre-call batch was flushed."""
def __init__(self, reply: Callable[[], Awaitable[object]], request: object, name: str) -> None:
self._reply: Final = reply
self._events: Final = _events(request)
if self._events is not None:
self._events.append(f"declare {name}")
def resolve(self) -> Awaitable[object]:
if self._events is not None and "flush_pre" not in self._events:
raise ConnectionError("resolved before its batch was flushed")
return asyncio.ensure_future(self._reply())
def _events(request: object) -> list[str] | None:
events: Final = getattr(request, "events", None)
return events if isinstance(events, list) else None # pyright: ignore[reportUnknownVariableType] # duck-typed test request
def _int(value: RawValue) -> int:
assert value is not None
return int(value)
def _text(value: RawValue) -> str | None:
if value is None:
return None
if isinstance(value, bytes):
return value.decode()
return str(value)
def _statused(raw: list[RawValue]) -> tuple[tuple[int, int, int] | None, list[RawValue]]:
if _int(raw[0]) == 1:
return (_int(raw[1]) - 1, _int(raw[2]), _int(raw[3])), []
return None, raw[1:]
@dataclass
class FakeRateLimits:
"""Assign a fake to a ``*_script`` attribute to give the limiter a Redis that runs it; a script the test
gave no fake for fails loudly."""
connected: bool = True
batch_rate_limiter_script: RawScript | None = None
batch_counter_read_script: RawScript | None = None
check_and_increment_by_n_script: RawScript | None = None
window_guarded_token_increment_script: RawScript | None = None
parallel_acquire_script: RawScript | None = None
parallel_release_script: RawScript | None = None
parallel_count_script: RawScript | None = None
token_increment_script: RawScript | None = None
redis_increments: list[tuple[str, float, int | None]] = field(default_factory=list)
post_call: list[str] = field(default_factory=list)
memory: NativeRateLimits = field(default_factory=lambda: NativeRateLimits(None))
def store(self) -> RateLimitStore:
return RateLimitStore(None, handle=lambda _client, _ttl, _redis_ttl: self)
async def _run(
self, script: RawScript | None, name: str, keys: Sequence[str], args: Sequence[object]
) -> list[RawValue]:
if script is None:
raise AssertionError(f"the test gave the limiter no fake for {name}")
return _RAW.validate_python(await script(keys=list(keys), args=list(args)))
def has_redis(self) -> bool:
return self.connected
def slot_groups(self, scopes_: Sequence[NativeScope], force: bool) -> list[list[int]]:
return self.memory.slot_groups(scopes_, force)
def local(self, entries: Sequence[NativeEntry]) -> list[LocalValue | None]:
return self.memory.local(entries)
def set_local(self, entry: NativeEntry, value: LocalValue, ttl_seconds: float | None) -> None:
self.memory.set_local(entry, value, ttl_seconds)
async def get(self, entry: NativeEntry, *, backfill_ttl_seconds: float | None = None) -> LocalValue | None:
return self.memory.local((entry,))[0]
def sliding_window(
self, counters_: Sequence[NativeCounter], now: int, window_size: int, *, request: object = None
) -> Resolved:
keys: Final = [key for counter in counters_ for key in (window_key(counter), counter_key(counter))]
async def reply() -> object:
raw: Final = await self._run(
self.batch_rate_limiter_script, "batch_rate_limiter_script", keys, [now, window_size]
)
return [(_text(raw[index]), _int(raw[index + 1])) for index in range(0, len(raw), 2)]
return Resolved(reply, request, "sliding_window")
def read_counters(self, counters_: Sequence[NativeCounter], *, request: object = None) -> Resolved:
keys: Final = [key for counter in counters_ for key in (window_key(counter), counter_key(counter))]
async def reply() -> object:
raw: Final = await self._run(self.batch_counter_read_script, "batch_counter_read_script", keys, [])
return [_text(value) for value in raw]
return Resolved(reply, request, "read_counters")
def check_and_increment(self, counters_: Sequence[NativeAtomicCounter], *, request: object = None) -> Resolved:
keys: Final = [key for counter, *_ in counters_ for key in (window_key(counter), counter_key(counter))]
args: Final = [arg for _counter, *numbers in counters_ for arg in numbers]
async def reply() -> object:
over_limit, rest = _statused(
await self._run(self.check_and_increment_by_n_script, "check_and_increment_by_n_script", keys, args)
)
return over_limit, [(_int(rest[index]), _int(rest[index + 1])) for index in range(0, len(rest), 2)]
return Resolved(reply, request, "check_and_increment")
def guarded_increment(
self,
counter: NativeCounter,
expected_window_start: str,
amount: float,
ttl_seconds: int | None,
*,
request: object = None,
) -> Resolved:
async def reply() -> object:
raw: Final = await self._run(
self.window_guarded_token_increment_script,
"window_guarded_token_increment_script",
[window_key(counter), counter_key(counter)],
[expected_window_start, amount, ttl_seconds or 0],
)
return _int(raw[0]) == 1, _int(raw[1])
return Resolved(reply, request, "guarded_increment")
def parallel_acquire(
self,
gauges: Sequence[tuple[NativeScope, int]],
slot_ttl_seconds: int,
slot_id: str,
*,
request: object = None,
) -> Resolved:
keys: Final = [scope_key(scope, "max_parallel_requests") for scope, _limit in gauges]
args: Final = [arg for _scope, limit in gauges for arg in (limit, slot_ttl_seconds, slot_id)]
async def reply() -> object:
over_limit, rest = _statused(
await self._run(self.parallel_acquire_script, "parallel_acquire_script", keys, args)
)
return over_limit, [_int(value) for value in rest]
return Resolved(reply, request, "parallel_acquire")
def parallel_release(self, gauges: Sequence[NativeScope], slot_id: str, *, request: object = None) -> Resolved:
keys: Final = [scope_key(scope, "max_parallel_requests") for scope in gauges]
async def reply() -> object:
raw: Final = await self._run(
self.parallel_release_script, "parallel_release_script", keys, [slot_id for _ in gauges]
)
return [_int(value) for value in raw]
return Resolved(reply, request, "parallel_release")
def parallel_count(
self, gauges: Sequence[NativeScope], slot_ttl_seconds: int, *, request: object = None
) -> Resolved:
keys: Final = [scope_key(scope, "max_parallel_requests") for scope in gauges]
async def reply() -> object:
raw: Final = await self._run(
self.parallel_count_script, "parallel_count_script", keys, [slot_ttl_seconds for _ in gauges]
)
return [_int(value) for value in raw]
return Resolved(reply, request, "parallel_count")
def token_increment(self, increments_: Sequence[NativeIncrement], *, request: object = None) -> Resolved:
keys: Final = [counter_key(counter) for counter, _amount, _ttl in increments_]
args: Final = [arg for _counter, amount, ttl in increments_ for arg in (amount, ttl or 0)]
async def reply() -> object:
raw: Final = await self._run(self.token_increment_script, "token_increment_script", keys, args)
return [float(_text(value) or 0) for value in raw]
return Resolved(reply, request, "token_increment")
async def token_increment_post_call(self, request: object, increments_: Sequence[NativeIncrement]) -> None:
self.post_call.append("token_increment")
try:
await self.token_increment(increments_).resolve()
except Exception: # noqa: BLE001 # the store's post-call path falls back exactly like this
await self.increment_many(increments_)
async def parallel_release_post_call(self, request: object, gauges: Sequence[NativeScope], slot_id: str) -> None:
self.post_call.append("parallel_release")
try:
await self.parallel_release(gauges, slot_id).resolve()
except Exception: # noqa: BLE001 # the store's post-call path only logs this
return
async def increment(self, increment: NativeIncrement) -> float:
total: Final = await self.memory.increment(increment)
if self.connected:
counter, amount, ttl = increment
self.redis_increments.append((counter_key(counter), amount, ttl))
return total
async def increment_many(self, increments_: Sequence[NativeIncrement]) -> None:
await self.memory.increment_many(increments_)
if self.connected:
self.redis_increments.extend((counter_key(counter), amount, ttl) for counter, amount, ttl in increments_)
async def increment_many_post_call(self, request: object, increments_: Sequence[NativeIncrement]) -> None:
self.post_call.append("increment_many")
await self.increment_many(increments_)

View file

@ -0,0 +1,158 @@
import binascii
from dataclasses import dataclass
from typing import Final
import pytest
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.rate_limits import (
AtomicCounter,
CounterIncrement,
RateLimitScope,
RateLimitStore,
)
from tests.unit.caching.test_redis_batch import FakeNatives
@dataclass(frozen=True, slots=True)
class _WireScope:
descriptor_key: str
descriptor_value: str
window: str
requests: str
tokens: str
max_parallel_requests: str
_WIRE_SCOPES: Final = (
_WireScope(
descriptor_key="api_key",
descriptor_value="88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
window="{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:window",
requests="{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:requests",
tokens="{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:tokens",
max_parallel_requests=(
"{api_key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b}:max_parallel_requests"
),
),
_WireScope(
descriptor_key="team_member",
descriptor_value="team-1:user@example.com",
window="{team_member:team-1:user@example.com}:window",
requests="{team_member:team-1:user@example.com}:requests",
tokens="{team_member:team-1:user@example.com}:tokens",
max_parallel_requests="{team_member:team-1:user@example.com}:max_parallel_requests",
),
_WireScope(
descriptor_key="model_per_key",
descriptor_value="sk-hash:anthropic/claude-sonnet-5-5",
window="{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:window",
requests="{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:requests",
tokens="{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:tokens",
max_parallel_requests="{model_per_key:sk-hash:anthropic/claude-sonnet-5-5}:max_parallel_requests",
),
_WireScope(
descriptor_key="tag",
descriptor_value="tier}é",
window="{tag:tier}é}:window",
requests="{tag:tier}é}:requests",
tokens="{tag:tier}é}:tokens",
max_parallel_requests="{tag:tier}é}:max_parallel_requests",
),
)
def _mains_key(key: str, value: str, rate_limit_type: str) -> str:
"""``create_rate_limit_keys`` and the window key as ``main``'s limiter formatted them."""
return f"{{{key}:{value}}}:{rate_limit_type}"
def _mains_keyslot(key: str) -> int:
"""``keyslot_for_redis_cluster`` from ``main``'s limiter."""
start: Final = key.find("{")
if start != -1:
end: Final = key.find("}", start + 1)
if end != -1 and end != start + 1:
key = key[start + 1 : end]
return binascii.crc_hqx(key.encode("utf-8"), 0) % 16384
@pytest.mark.parametrize("scope", _WIRE_SCOPES, ids=("api_key", "team_member", "model_per_key", "tag"))
def test_the_wire_keys_are_the_keys_mains_limiter_formats(scope: _WireScope) -> None:
assert (scope.window, scope.requests, scope.tokens, scope.max_parallel_requests) == tuple(
_mains_key(scope.descriptor_key, scope.descriptor_value, suffix)
for suffix in ("window", "requests", "tokens", "max_parallel_requests")
)
@pytest.mark.parametrize(("raw", "decoded"), (("1730000000", 1730000000), ("7", 7), ("7.5", 7.5), ("-3", -3)))
def test_the_wire_values_decode_as_mains_redis_cache_reads_them(raw: str, decoded: float) -> None:
redis_cache: Final = RedisCache.__new__(RedisCache)
read: Final = redis_cache._get_cache_logic(raw.encode()) # pyright: ignore[reportPrivateUsage] # main's read codec is the contract
assert read == decoded
assert type(read) is type(decoded)
@pytest.mark.parametrize("force", [True, False])
def test_slot_groups_follow_mains_hash_tag_slots_only_when_grouping_is_forced(force: bool):
store: Final = RateLimitStore(None)
scopes: Final = (
*(RateLimitScope(s.descriptor_key, s.descriptor_value) for s in _WIRE_SCOPES),
RateLimitScope("api_key", "a"),
RateLimitScope("team", "t"),
RateLimitScope("api_key", "a"),
RateLimitScope("tag", ""),
)
slots: Final = [_mains_keyslot(_mains_key(s.descriptor_key, s.descriptor_value, "window")) for s in scopes]
expected: Final = tuple(
tuple(index for index, slot in enumerate(slots) if slot == group_slot) for group_slot in dict.fromkeys(slots)
)
grouped: Final = store.slot_groups(scopes, force)
assert grouped == (expected if force else (tuple(range(len(scopes))),))
def test_the_memory_tier_keeps_each_value_type_it_was_given():
store: Final = RateLimitStore(None)
scope: Final = RateLimitScope("api_key", "k1")
store.set_local(scope.window, "1730000000", 60)
store.set_local(scope.counter("requests"), 3, 60)
store.set_local(scope.counter("tokens"), 2.5, 60)
store.set_local(scope.gauge, {"slot-1": 1730000000.5}, 60)
values: Final = store.local((scope.window, scope.counter("requests"), scope.counter("tokens"), scope.gauge))
assert values == ("1730000000", 3, 2.5, {"slot-1": 1730000000.5})
assert [type(value) for value in values] == [str, int, float, dict]
assert store.local((RateLimitScope("api_key", "k2").counter("requests"),)) == (None,)
@pytest.mark.asyncio
async def test_without_redis_increments_add_in_memory_as_python_numbers_do():
store: Final = RateLimitStore(None)
tokens: Final = RateLimitScope("team", "t1").counter("tokens")
assert await store.increment(CounterIncrement(tokens, 5, 60)) == 5
await store.increment_many((CounterIncrement(tokens, 2.5, 60), CounterIncrement(tokens, -1, 60)))
await store.increment_many_post_call((CounterIncrement(tokens, 1, 60),))
assert store.local((tokens,)) == (7.5,)
assert not store.has_redis
assert not store.has_post_call_batch
@pytest.mark.asyncio
async def test_without_redis_nothing_is_declared_on_the_request_and_post_call_writes_are_refused():
store: Final = RateLimitStore(None)
counter: Final = RateLimitScope("api_key", "k1").counter("requests")
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
assert store.declare_sliding_window((counter,), 0, 60) is None
assert store.declare_check_and_increment((AtomicCounter(counter, 10, 1, 60, 60),)) is None
assert not await store.token_increment_post_call((CounterIncrement(counter, 1, 60),))
assert not await store.parallel_release_post_call((counter.scope,), "slot")
assert natives.built == []

View file

@ -22,11 +22,13 @@ import pytest
from fastapi import HTTPException
import litellm
from litellm.caching.caching import DualCache
from litellm.exceptions import RateLimitError, RateLimitErrorCategory, RateLimitType
from litellm.proxy.common_utils.proxy_rate_limit_error import (
ProxyRateLimitError,
map_v3_rate_limit_type,
)
from litellm.proxy.utils import InternalUsageCache
class TestRateLimitErrorCategory:
@ -421,13 +423,12 @@ class TestProxyHooksActuallyRaiseProxyRateLimitError:
def test_parallel_request_limiter_v3_handle_rate_limit_error_raises(self):
"""v3 parallel_request_limiter's ``_handle_rate_limit_error`` must
translate an OVER_LIMIT response into a ProxyRateLimitError."""
from unittest.mock import MagicMock
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
handler = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=MagicMock())
handler = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(DualCache()))
# Minimal fabricated OVER_LIMIT response. The helper only reads a
# handful of fields off `status` and ignores everything else.
response = {
@ -1171,14 +1172,13 @@ class TestProxyHooksWireTypeCorrectly:
`map_v3_rate_limit_type(status["rate_limit_type"])` call site so
coverage tools see the new wiring as exercised.
"""
from unittest.mock import MagicMock
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
handler = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=MagicMock(),
internal_usage_cache=InternalUsageCache(DualCache()),
)
# Minimal RateLimitResponse + descriptors shape that the handler
# reads. We only need one OVER_LIMIT status to drive the raise.
@ -1222,14 +1222,13 @@ class TestProxyHooksWireTypeCorrectly:
@pytest.mark.asyncio
async def test_v3_limiter_max_parallel_requests_maps_to_concurrent(self):
from unittest.mock import MagicMock
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
handler = PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=MagicMock(),
internal_usage_cache=InternalUsageCache(DualCache()),
)
response = {
"overall_code": "OVER_LIMIT",