scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

refactor(redis): send every store op through one Send and commands as values

Browse source 
Each redis-client command is one constructor returning Command<T>, sent by RedisClient, Pipeline or
RedisBatch. Each redis-store family becomes a folder with its keys and Python-facing types, every op
takes a Send (Now, PreCall, PostCall), and every family works without Redis until one is attached.
The Python store shims are replaced by the native types.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-05 13:04:09 +09:00
parent 7a78b136a8
commit 011d46cf09
104 changed files with 5416 additions and 5106 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
litellm-rust/crates/gateway/src/redis/mod.rs
View file

@ -261,7 +261,7 @@ mod tests {
use std::{collections::HashMap, time::Duration};
use litellm_config::Config;
use litellm_redis_client::{Standalone, StaticAuth};
use litellm_redis_client::{Standalone, StaticAuth, commands};
use litellm_redis_testing::RedisServer;
use litellm_redis_types::{Auth, HealthCheck, Passthrough};
use rstest::rstest;
@ -400,16 +400,16 @@ mod tests {
.unwrap();
assert_eq!(
cache.get::<String>("rate-limit").await.unwrap().as_deref(),
cache.send(commands::get::<String>("rate-limit")).await.unwrap().as_deref(),
Some("1")
);
assert_eq!(
cache.get::<String>("spend").await.unwrap().as_deref(),
cache.send(commands::get::<String>("spend")).await.unwrap().as_deref(),
Some("2")
);
assert_eq!(
unnamespaced
.get::<String>("responses:rate-limit")
.send(commands::get::<String>("responses:rate-limit"))
.await
.unwrap()
.as_deref(),
@ -417,7 +417,7 @@ mod tests {
);
assert_eq!(
unnamespaced
.get::<String>("responses:spend")
.send(commands::get::<String>("responses:spend"))
.await
.unwrap()
.as_deref(),
@ -425,12 +425,12 @@ mod tests {
);
assert!(
unnamespaced
.get::<String>("rate-limit")
.send(commands::get::<String>("rate-limit"))
.await
.unwrap()
.is_none()
);
assert!(unnamespaced.get::<String>("spend").await.unwrap().is_none());
assert!(unnamespaced.send(commands::get::<String>("spend")).await.unwrap().is_none());
}
#[rstest]

3
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -45,7 +45,7 @@ mod _native {
#[pymodule_export]
use crate::redis::store::rate_limits::NativeRateLimits;
#[pymodule_export]
use crate::redis::store::request::NativeRequestRedis;
use crate::redis::store::request::{NativeRequestRedis, NativeSend};
#[pymodule_export]
use crate::redis::store::router_usage::NativeRouterUsage;
#[pymodule_export]
@ -129,6 +129,7 @@ mod tests {
"NativeDiagnosticProcessor",
"NativeRedisClient",
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
"NativeRateLimits",
"NativePending",

120
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs
View file

@ -1,26 +1,15 @@
use std::{sync::Arc, time::Duration};
use std::sync::Arc;
use litellm_python_compat::Value;
use litellm_redis_store::auth_cache::{AuthCache, DeclaredWrite, Error, Settings};
use litellm_redis_store::auth_cache::{AuthCache, Settings};
use pyo3::prelude::*;
use super::{
NativePending, PyNumber, kind_error, redis_async, redis_sync, request::NativeRequestRedis,
store_error, ttl, value::PyValue,
NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote, request::NativeSend,
ttl, value::PyValue,
};
use crate::redis::client::NativeRedisClient;
fn auth_cache_error(error: Error) -> PyErr {
kind_error(error.kind(), &error)
}
/// A `set_pre_call` or `delete_pre_call`, settled once and raising when Redis refused it.
fn declared_write(declared: DeclaredWrite) -> NativePending {
NativePending::new("async_set_cache", auth_cache_error, async move {
declared.resolve().await
})
}
fn values(values: Vec<Option<Value>>) -> Vec<Option<PyValue>> {
values.into_iter().map(|value| value.map(PyValue)).collect()
}
@ -41,11 +30,10 @@ impl NativeAuthCache {
redis_ttl_seconds: u64,
) -> Self {
Self(Arc::new(AuthCache::new(
client.map(|client| client.0.clone()),
client.map(|client| remote(&client, redis_ttl_seconds)),
Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds),
redis_ttl: Duration::from_secs(redis_ttl_seconds),
},
)))
}
@ -57,14 +45,14 @@ impl NativeAuthCache {
fn get<'py>(&self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_get_cache", store_error, async move {
redis_async(py, "async_get_cache", kind_error, async move {
store.get(&key).await.map(|value| value.map(PyValue))
})
}
fn get_sync(&self, py: Python<'_>, key: String) -> PyResult<Option<PyValue>> {
let store = self.0.clone();
redis_sync(py, "get_cache", store_error, async move {
redis_sync(py, "get_cache", kind_error, async move {
store.get(&key).await.map(|value| value.map(PyValue))
})
}
@ -78,7 +66,7 @@ impl NativeAuthCache {
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_set_cache", store_error, async move {
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&key, value.0, ttl(ttl_seconds)).await
})
}
@ -92,37 +80,30 @@ impl NativeAuthCache {
ttl_seconds: Option<f64>,
) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "set_cache", store_error, async move {
redis_sync(py, "set_cache", kind_error, async move {
store.set_str(&key, value.0, ttl(ttl_seconds)).await
})
}
/// Memory now and the `SET` on the request's pre-call batch; `None` without Redis.
#[pyo3(signature = (request, key, value, ttl_seconds = None))]
fn set_pre_call(
/// Memory now and the `SET` on the batch `send` picks.
#[pyo3(signature = (key, value, ttl_seconds, send))]
fn declare_set(
&self,
request: &NativeRequestRedis,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<Option<NativePending>> {
Ok(self
send: &NativeSend,
) -> PyResult<NativePending> {
let written = self
.0
.set_pre_call(&request.0, &key, value.0, ttl(ttl_seconds))
.map_err(store_error)?
.map(declared_write))
.declare_set(send.send(), &key, value.0, ttl(ttl_seconds))
.map_err(kind_error)?;
Ok(pending("async_set_cache", kind_error, written, std::convert::identity))
}
fn delete_pre_call(
&self,
request: &NativeRequestRedis,
key: String,
) -> PyResult<Option<NativePending>> {
Ok(self
.0
.delete_pre_call(&request.0, &key)
.map_err(store_error)?
.map(declared_write))
fn declare_delete(&self, key: String, send: &NativeSend) -> PyResult<NativePending> {
let deleted = self.0.declare_delete(send.send(), &key).map_err(kind_error)?;
Ok(pending("async_set_cache", kind_error, deleted, std::convert::identity))
}
#[pyo3(signature = (entries, ttl_seconds = None))]
@ -137,70 +118,45 @@ impl NativeAuthCache {
.into_iter()
.map(|(key, value)| (key, value.0))
.collect();
redis_async(py, "async_set_cache_pipeline", store_error, async move {
redis_async(py, "async_set_cache_pipeline", kind_error, async move {
store.set_many(entries, ttl(ttl_seconds)).await
})
}
/// On the request's pre-call batch when `request` is given, returning before anything is sent.
/// On a request batch it returns before anything is sent, and `Now` is a pipeline of its own.
fn write_back<'py>(
&self,
py: Python<'py>,
request: Option<&NativeRequestRedis>,
entries: Vec<(String, PyValue, Option<f64>)>,
send: &NativeSend,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let request = request.map(|request| request.0.clone());
let entries = entries
.into_iter()
.map(|(key, value, seconds)| (key, value.0, ttl(seconds)))
.collect();
redis_async(
py,
"async_set_cache_pipeline_with_ttls",
auth_cache_error,
async move { store.write_back(request.as_ref(), entries).await },
)
let written = self.0.write_back(send.send(), entries).map_err(kind_error)?;
redis_async(py, "async_set_cache_pipeline_with_ttls", kind_error, written)
}
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_delete_cache", auth_cache_error, async move {
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&keys).await
})
}
fn delete_sync(&self, py: Python<'_>, keys: Vec<String>) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "delete_cache", auth_cache_error, async move {
redis_sync(py, "delete_cache", kind_error, async move {
store.delete(&keys).await
})
}
/// Reads of keys the caller found missing in memory, on the request's pre-call batch; `None`
/// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None`
/// without Redis.
fn declare_get_many(
&self,
request: &NativeRequestRedis,
keys: Vec<String>,
) -> Option<NativePending> {
self.0.declare_get_many(&request.0, &keys).map(|declared| {
NativePending::new("async_batch_get_cache", auth_cache_error, async move {
declared.resolve().await.map(values)
})
})
}
/// One round trip for keys missing in memory; resolves to `None` without Redis.
fn get_many<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_batch_get_cache", auth_cache_error, async move {
store
.get_many(&keys)
.await
.transpose()
.map(|read| read.map(values))
})
fn get_many(&self, keys: Vec<String>, send: &NativeSend) -> Option<NativePending> {
let read = self.0.get_many(send.send(), &keys)?;
Some(pending("async_batch_get_cache", kind_error, read, values))
}
fn increment_if_present<'py>(
@ -213,32 +169,32 @@ impl NativeAuthCache {
redis_async(
py,
"async_increment_if_present",
auth_cache_error,
kind_error,
async move { store.increment_if_present(&key, amount.0).await },
)
}
fn get_local(&self, key: &str) -> PyResult<Option<PyValue>> {
Ok(self.0.get_local(key).map_err(store_error)?.map(PyValue))
Ok(self.0.get_local(key).map_err(kind_error)?.map(PyValue))
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set_local(&self, key: &str, value: PyValue, ttl_seconds: Option<f64>) -> PyResult<()> {
self.0
.set_local(key, value.0, ttl(ttl_seconds))
.map_err(store_error)
.map_err(kind_error)
}
fn delete_local(&self, key: &str) -> PyResult<()> {
self.0.delete_local(key).map_err(store_error)
self.0.delete_local(key).map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(store_error)
self.0.flush_local().map_err(kind_error)
}
/// `(management objects, user-key objects)` held in memory.
fn len_local(&self) -> PyResult<(usize, usize)> {
self.0.len_local().map_err(store_error)
self.0.len_local().map_err(kind_error)
}
}

71
litellm-rust/crates/python-bridge/src/redis/store/mod.rs
View file

@ -7,6 +7,8 @@ pub(crate) mod router_usage;
pub(crate) mod spend_counters;
mod value;
use value::PyValue;
use std::{
pin::Pin,
sync::{Mutex, PoisonError},
@ -14,16 +16,16 @@ use std::{
};
use litellm_host_python::PythonContext;
use litellm_redis_store::{ErrorKind, Number};
use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote};
use pyo3::{
IntoPyObjectExt,
exceptions::{PyRuntimeError, PyTypeError},
exceptions::{PyRuntimeError, PyTypeError, PyValueError},
prelude::*,
types::{PyBool, PyFloat, PyInt},
};
use crate::execution::{run_async_value, run_sync_value};
use crate::redis::events::RedisCall;
use crate::redis::{client::NativeRedisClient, events::RedisCall};
/// The labels `RedisCache` reports for this operation, in the caller's context so its
/// `service_target` reaches the service event.
@ -73,6 +75,14 @@ fn ttl(seconds: Option<f64>) -> Option<Duration> {
seconds.map(|seconds| Duration::from_secs_f64(seconds.max(0.0)))
}
/// A store's Redis with the `RedisCache.default_ttl` Python attached it with.
fn remote(client: &NativeRedisClient, redis_ttl_seconds: u64) -> Remote {
Remote {
client: client.0.clone(),
default_ttl: Duration::from_secs(redis_ttl_seconds),
}
}
type Resolve = Pin<Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>>;
/// A Redis call declared on a request's batch. `resolve` flushes that batch if it has not gone
@ -84,7 +94,7 @@ pub(crate) struct NativePending {
}
impl NativePending {
fn new<T, E: 'static>(
pub(crate) fn new<T, E: 'static>(
call_type: &'static str,
error: fn(E) -> PyErr,
reply: impl Future<Output = Result<T, E>> + Send + 'static,
@ -123,15 +133,31 @@ pyo3::create_exception!(
"A native Redis store operation failed. Args are (kind, message), kind spelled as litellm-redis-store's ErrorKind."
);
fn kind_error(kind: ErrorKind, error: &(dyn std::error::Error + 'static)) -> PyErr {
let kind: &'static str = kind.into();
NativeRedisStoreError::new_err((kind, litellm_error::Report::new(error).to_string()))
/// What the spend counter and auth cache stores raise.
fn kind_error(error: Error) -> PyErr {
let kind: &'static str = error.kind().into();
NativeRedisStoreError::new_err((kind, litellm_error::Report::new(&error).to_string()))
}
fn store_error(error: litellm_cache::Error) -> PyErr {
/// What the router stores and the rate limiter raise.
fn runtime_error(error: Error) -> PyErr {
PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string())
}
/// A store op declared on the batch the caller's `NativeSend` picked, answered as `reply` shapes it.
fn pending<T, R>(
call_type: &'static str,
error: fn(Error) -> PyErr,
declared: Pending<T>,
reply: fn(T) -> R,
) -> NativePending
where
T: Send + 'static,
R: for<'a> IntoPyObject<'a> + Send + 'static,
{
NativePending::new(call_type, error, async move { declared.await.map(reply) })
}
/// A Python `int` or `float` across the bridge, keeping which one it is.
pub(crate) struct PyNumber(Number);
@ -161,3 +187,32 @@ impl<'py> IntoPyObject<'py> for PyNumber {
})
}
}
/// A store type whose Python side is plain data, converted through its `FromValue` and `IntoValue`
/// so the bridge never names a field. A shape the store refuses is `ValueError`.
pub(crate) struct PyData<T>(pub(crate) T);
impl<'a, 'py, T: FromValue> FromPyObject<'a, 'py> for PyData<T> {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let PyValue(value) = value.extract()?;
T::from_value(&value)
.map(Self)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
impl<'py, T: IntoValue> IntoPyObject<'py> for PyData<T> {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
PyValue(self.0.into_value()).into_pyobject(py)
}
}
fn unwrap_all<T>(values: Vec<PyData<T>>) -> Vec<T> {
values.into_iter().map(|value| value.0).collect()
}

366
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs
View file

@ -1,17 +1,23 @@
use std::{sync::Arc, time::Duration};
use litellm_redis_store::rate_limits::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Error, Gauge,
GuardedIncrement, GuardedReply, Increment, Local, Pending, RateLimits, Scope, WindowCount,
use litellm_redis_store::{
Pending,
rate_limits::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge,
GuardedIncrement, GuardedReply, Increment, Local, RateLimits, Scope, WindowCount,
},
};
use pyo3::{
exceptions::{PyRuntimeError, PyTypeError, PyValueError},
exceptions::PyTypeError,
prelude::*,
types::{PyDict, PyString},
};
use super::{NativePending, PyNumber, redis_async, request::NativeRequestRedis, store_error, ttl};
use crate::execution::run_async_value;
use super::{
NativePending, PyData, PyNumber, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
runtime_error, ttl, unwrap_all,
};
use crate::redis::client::NativeRedisClient;
/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`;
@ -19,91 +25,6 @@ use crate::redis::client::NativeRedisClient;
#[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, kind)`, the kind spelled as its key suffix.
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.parse().map_err(|_| unknown_suffix(&kind))?,
}))
}
}
fn unknown_suffix(suffix: &str) -> PyErr {
PyValueError::new_err(format!("unknown rate-limit key suffix {suffix:?}"))
}
/// `(key_suffix, descriptor_key, descriptor_value)`, read by `Entry::new`.
struct PyEntry(Entry);
impl<'a, 'py> FromPyObject<'a, 'py> for PyEntry {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let (suffix, descriptor, value): (String, String, String) = value.extract()?;
Entry::new(&suffix, Scope { descriptor, value })
.map(Self)
.map_err(|_| unknown_suffix(&suffix))
}
}
/// `(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);
@ -201,48 +122,13 @@ fn acquire(reply: AcquireReply) -> (OverLimit, Vec<i64>) {
}
}
/// One script call declared on a batch: the request's pre-call batch, or a batch of its own when
/// no request was passed. Resolving it flushes that batch if it has not gone out.
fn pending<T, R>(
call_type: &'static str,
pending: Result<Pending<T>, Error>,
reply: fn(T) -> R,
) -> PyResult<NativePending>
/// One script call declared on the batch the caller's `NativeSend` picked; `None` without Redis.
fn script<T, R>(call_type: &'static str, declared: Option<Pending<T>>, reply: fn(T) -> R) -> Option<NativePending>
where
T: Send + 'static,
R: for<'a> IntoPyObject<'a> + Send + 'static,
{
let pending = pending.map_err(rate_limit_error)?;
Ok(NativePending::new(
call_type,
rate_limit_error,
async move { pending.resolve().await.map(reply) },
))
}
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())
Some(super::pending(call_type, runtime_error, declared?, reply))
}
#[pymethods]
@ -254,51 +140,47 @@ impl NativeRateLimits {
backfill_ttl_seconds: Option<f64>,
redis_ttl_seconds: u64,
) -> Self {
let backfill_ttl = ttl(backfill_ttl_seconds);
Self(Arc::new(RateLimits::new(
client.map(|client| client.0.clone()),
backfill_ttl,
Duration::from_secs(redis_ttl_seconds),
client.map(|client| remote(&client, redis_ttl_seconds)),
ttl(backfill_ttl_seconds),
)))
}
#[getter]
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 slot_groups(&self, scopes: Vec<PyData<Scope>>, force: bool) -> Vec<Vec<usize>> {
self.0.slot_groups(&unwrap_all(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()
fn get_local(&self, entry: PyData<Entry>) -> PyResult<Option<PyLocal>> {
Ok(self.0.get_local(&entry.0).map_err(runtime_error)?.map(PyLocal))
}
#[pyo3(signature = (entry, value, ttl_seconds))]
fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option<f64>) -> PyResult<()> {
fn set_local(
&self,
entry: PyData<Entry>,
value: PyLocal,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
self.0
.set_local(&entry.0, value.0, ttl(ttl_seconds))
.map_err(store_error)
.map_err(runtime_error)
}
#[pyo3(signature = (entry, *, backfill_ttl_seconds = None))]
fn get<'py>(
&self,
py: Python<'py>,
entry: PyEntry,
entry: PyData<Entry>,
backfill_ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let backfill_ttl = ttl(backfill_ttl_seconds);
redis_async(py, "async_get_cache", store_error, async move {
redis_async(py, "async_get_cache", runtime_error, async move {
store
.get(&entry.0, backfill_ttl)
.await
@ -306,99 +188,79 @@ impl NativeRateLimits {
})
}
#[pyo3(signature = (counters_, now, window_size, *, request = None))]
fn sliding_window(
&self,
counters_: Vec<PyCounter>,
counters: Vec<PyData<Counter>>,
now: i64,
window_size: i64,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
pending(
send: &NativeSend,
) -> Option<NativePending> {
script(
"batch_rate_limiter_script",
self.0
.sliding_window(request.as_ref(), &counters(counters_), now, window_size),
.sliding_window(send.send(), &unwrap_all(counters), now, window_size),
window_counts,
)
}
#[pyo3(signature = (counters_, *, request = None))]
fn read_counters(
&self,
counters_: Vec<PyCounter>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
pending(
counters: Vec<PyData<Counter>>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"batch_counter_read_script",
self.0.read_counters(request.as_ref(), &counters(counters_)),
self.0.read_counters(send.send(), &unwrap_all(counters)),
std::convert::identity,
)
}
#[pyo3(signature = (counters_, *, request = None))]
fn check_and_increment(
&self,
counters_: Vec<PyAtomicCounter>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
let counters_ = counters_
.into_iter()
.map(|counter| counter.0)
.collect::<Vec<_>>();
pending(
counters: Vec<PyData<AtomicCounter>>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"check_and_increment_by_n_script",
self.0.check_and_increment(request.as_ref(), &counters_),
self.0
.check_and_increment(send.send(), &unwrap_all(counters)),
atomic,
)
}
#[pyo3(signature = (counter, expected_window_start, amount, ttl_seconds, *, request = None))]
fn guarded_increment(
&self,
counter: PyCounter,
counter: PyData<Counter>,
expected_window_start: String,
amount: PyNumber,
ttl_seconds: Option<u64>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
send: &NativeSend,
) -> Option<NativePending> {
let increment = GuardedIncrement {
counter: counter.0,
expected_window_start,
amount: amount.0,
ttl: ttl_seconds.map(Duration::from_secs),
};
pending(
script(
"window_guarded_token_increment_script",
self.0.guarded_increment(request.as_ref(), &increment),
self.0.guarded_increment(send.send(), &increment),
guarded,
)
}
#[pyo3(signature = (gauges, slot_ttl_seconds, slot_id, *, request = None))]
fn parallel_acquire(
&self,
gauges: Vec<(PyScope, i64)>,
gauges: Vec<PyData<Gauge>>,
slot_ttl_seconds: u64,
slot_id: &str,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
let gauges = gauges
.into_iter()
.map(|(scope, limit)| Gauge {
scope: scope.0,
limit,
})
.collect::<Vec<_>>();
pending(
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_acquire_script",
self.0.parallel_acquire(
request.as_ref(),
&gauges,
send.send(),
&unwrap_all(gauges),
Duration::from_secs(slot_ttl_seconds),
slot_id,
),
@ -406,96 +268,76 @@ impl NativeRateLimits {
)
}
#[pyo3(signature = (gauges, slot_id, *, request = None))]
fn parallel_release(
&self,
gauges: Vec<PyScope>,
gauges: Vec<PyData<Scope>>,
slot_id: &str,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
pending(
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_release_script",
self.0
.parallel_release(request.as_ref(), &scopes(gauges), slot_id),
.parallel_release(send.send(), &unwrap_all(gauges), slot_id),
std::convert::identity,
)
}
#[pyo3(signature = (gauges, slot_ttl_seconds, *, request = None))]
fn parallel_count(
&self,
gauges: Vec<PyScope>,
gauges: Vec<PyData<Scope>>,
slot_ttl_seconds: u64,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
pending(
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_count_script",
self.0.parallel_count(
request.as_ref(),
&scopes(gauges),
send.send(),
&unwrap_all(gauges),
Duration::from_secs(slot_ttl_seconds),
),
std::convert::identity,
)
}
#[pyo3(signature = (increments_, *, request = None))]
fn token_increment(
&self,
increments_: Vec<PyIncrement>,
request: Option<PyRef<'_, NativeRequestRedis>>,
) -> PyResult<NativePending> {
let request = request_of(request);
pending(
increments: Vec<PyData<Increment>>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"token_increment_script",
self.0
.token_increment(request.as_ref(), &increments(increments_)),
self.0.token_increment(send.send(), &unwrap_all(increments)),
std::convert::identity,
)
}
fn token_increment_post_call<'py>(
/// `False` without Redis.
fn token_increment_post_call(
&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)
})
request: &NativeRequestRedis,
increments: Vec<PyData<Increment>>,
) -> bool {
self.0.token_increment_post_call(&request.0, unwrap_all(increments))
}
fn parallel_release_post_call<'py>(
/// `False` without Redis.
fn parallel_release_post_call(
&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)
})
request: &NativeRequestRedis,
gauges: Vec<PyData<Scope>>,
slot_id: &str,
) -> bool {
self.0
.parallel_release_post_call(&request.0, &unwrap_all(gauges), slot_id)
}
fn increment<'py>(
&self,
py: Python<'py>,
increment: PyIncrement,
increment: PyData<Increment>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment", store_error, async move {
redis_async(py, "async_increment", runtime_error, async move {
store.increment(&increment.0).await.map(PyNumber)
})
}
@ -503,28 +345,22 @@ impl NativeRateLimits {
fn increment_many<'py>(
&self,
py: Python<'py>,
increments_: Vec<PyIncrement>,
increments: Vec<PyData<Increment>>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments_ = increments(increments_);
redis_async(py, "async_increment_pipeline", store_error, async move {
store.increment_many(&increments_).await
let increments = unwrap_all(increments);
redis_async(py, "async_increment_pipeline", runtime_error, async move {
store.increment_many(&increments).await
})
}
fn increment_many_post_call<'py>(
fn increment_many_post_call(
&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)
})
request: &NativeRequestRedis,
increments: Vec<PyData<Increment>>,
) -> PyResult<()> {
self.0
.increment_many_post_call(&request.0, &unwrap_all(increments))
.map_err(runtime_error)
}
}

21
litellm-rust/crates/python-bridge/src/redis/store/request.rs
View file

@ -1,6 +1,6 @@
use std::time::Duration;
use litellm_redis_store::RequestRedis;
use litellm_redis_store::{RequestRedis, Send};
use pyo3::prelude::*;
use super::call;
@ -39,3 +39,22 @@ impl NativeRequestRedis {
})
}
}
/// The Python face of `Send`: which batch a store op rides. Each variant carries exactly what it
/// needs, so there is no phase string and no phase without a request.
#[pyclass(frozen, name = "NativeSend")]
pub(crate) enum NativeSend {
Now(),
PreCall(Py<NativeRequestRedis>),
PostCall(Py<NativeRequestRedis>),
}
impl NativeSend {
pub(crate) fn send(&self) -> Send<'_> {
match self {
Self::Now() => Send::Now,
Self::PreCall(request) => Send::PreCall(&request.get().0),
Self::PostCall(request) => Send::PostCall(&request.get().0),
}
}
}

106
litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs
View file

@ -1,15 +1,17 @@
use std::sync::Arc;
use litellm_redis_store::router_cooldowns::{Cooldown, RouterCooldowns};
use litellm_redis_store::{
Send,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use pyo3::{
exceptions::PyValueError,
prelude::*,
types::{PyDict, PyTuple},
};
use super::{
NativePending, PyNumber, policy::router_tier_policies, redis_async, redis_sync,
request::NativeRequestRedis, store_error, ttl, value::PyValue,
NativePending, PyData, PyNumber, pending, policy::router_tier_policies, redis_sync, remote,
request::NativeSend, runtime_error, ttl,
};
use crate::redis::client::NativeRedisClient;
@ -18,30 +20,6 @@ use crate::redis::client::NativeRedisClient;
#[pyclass(frozen, name = "NativeRouterCooldowns")]
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
/// A `CooldownCacheValue` dict, whose shape `Cooldown` owns.
pub(crate) struct PyCooldown(Cooldown);
impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let PyValue(value) = value.extract()?;
Cooldown::from_value(&value)
.map(Self)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
impl<'py> IntoPyObject<'py> for PyCooldown {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
PyValue(self.0.to_value()).into_pyobject(py)
}
}
/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read.
struct PyActive(Vec<(String, Cooldown)>);
@ -68,7 +46,7 @@ fn active_to_py(
py,
[
model_id.into_pyobject(py)?.into_any(),
PyCooldown(cooldown).into_pyobject(py)?.into_any(),
PyData(cooldown).into_pyobject(py)?,
],
)
})
@ -78,22 +56,29 @@ fn active_to_py(
#[pymethods]
impl NativeRouterCooldowns {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let client = client.map(|client| client.0.clone());
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>, redis_ttl_seconds: u64) -> PyResult<Self> {
let remote = client.map(|client| remote(&client, redis_ttl_seconds));
let policy = router_tier_policies(py)?.cooldowns;
Ok(Self(Arc::new(RouterCooldowns::new(client, policy))))
Ok(Self(Arc::new(RouterCooldowns::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn attach(&self, client: PyRef<'_, NativeRedisClient>, redis_ttl_seconds: u64) -> bool {
self.0.attach(remote(&client, redis_ttl_seconds))
}
fn put(
&self,
py: Python<'_>,
model_id: String,
cooldown: PyCooldown,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = ttl(Some(ttl_seconds)).unwrap_or_default();
redis_sync(py, "set_cache", store_error, async move {
redis_sync(py, "set_cache", runtime_error, async move {
store.put(&model_id, &cooldown.0, ttl).await
})
}
@ -107,52 +92,33 @@ impl NativeRouterCooldowns {
) -> PyResult<i64> {
let store = self.0.clone();
let ttl = ttl(Some(ttl_seconds)).unwrap_or_default();
redis_sync(py, "increment_cache", store_error, async move {
redis_sync(py, "increment_cache", runtime_error, async move {
store
.increment_allowed_fails(&model_id, suffix.as_deref(), ttl)
.await
})
}
fn active<'py>(
/// The deployments still cooling down, read on the batch `send` picks.
fn active(&self, model_ids: Vec<String>, send: &NativeSend) -> PyResult<NativePending> {
let active = self.0.active(send.send(), &model_ids).map_err(runtime_error)?;
Ok(pending("async_batch_get_cache", runtime_error, active, PyActive))
}
/// `active` on its own with the GIL released, for the router's sync completion path.
fn active_sync<'py>(
&self,
py: Python<'py>,
model_ids: Vec<String>,
) -> PyResult<Vec<Bound<'py, PyTuple>>> {
let store = self.0.clone();
let active = redis_sync(py, "batch_get_cache", store_error, async move {
store.active(&model_ids).await
})?;
let active = self.0.active(Send::Now, &model_ids).map_err(runtime_error)?;
let active = redis_sync(py, "batch_get_cache", runtime_error, active)?;
active_to_py(py, active)
}
fn aactive<'py>(&self, py: Python<'py>, model_ids: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_batch_get_cache", store_error, async move {
store.active(&model_ids).await.map(PyActive)
})
}
fn declare_active(
&self,
request: &NativeRequestRedis,
model_ids: Vec<String>,
) -> PyResult<NativePending> {
let declared = self
.0
.declare_active(&request.0, &model_ids)
.map_err(store_error)?;
let store = self.0.clone();
Ok(NativePending::new(
"async_batch_get_cache",
store_error,
async move { store.resolve_active(declared).await.map(PyActive) },
))
}
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
redis_sync(py, "batch_get_cache", store_error, async move {
redis_sync(py, "batch_get_cache", runtime_error, async move {
store
.minimum(&model_ids)
.await
@ -160,18 +126,18 @@ impl NativeRouterCooldowns {
})
}
fn evict_local(&self, model_id: &str) -> PyResult<()> {
self.0.evict_local(model_id).map_err(store_error)
fn delete_local(&self, model_id: &str) -> PyResult<()> {
self.0.delete_local(model_id).map_err(runtime_error)
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(store_error)
self.0.flush_local_counters().map_err(runtime_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, cooldown) in self.0.snapshot().map_err(store_error)? {
snapshot.set_item(key, PyCooldown(cooldown))?;
for (key, cooldown) in self.0.snapshot().map_err(runtime_error)? {
snapshot.set_item(key, PyData(cooldown))?;
}
Ok(snapshot)
}

57
litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs
View file

@ -1,54 +1,40 @@
use std::{sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, RouterUsage};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use pyo3::{prelude::*, types::PyDict};
use super::{PyNumber, policy::router_tier_policies, redis_async, store_error};
use super::{PyData, PyNumber, policy::router_tier_policies, redis_async, remote, runtime_error, unwrap_all};
use crate::redis::client::NativeRedisClient;
/// One router's global TPM/RPM counters. Every caller on `main` is async.
#[pyclass(frozen, name = "NativeRouterUsage")]
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
/// `(deployment_id, model, minute, kind)`, the kind spelled as its key segment.
pub(crate) struct PyCounter(Counter);
impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let (deployment_id, model, minute, kind): (String, Option<String>, String, String) =
value.extract()?;
let kind = kind
.parse()
.map_err(|_| PyValueError::new_err(format!("unknown usage counter kind {kind:?}")))?;
Ok(Self(Counter {
deployment_id,
model,
minute,
kind,
}))
}
}
#[pymethods]
impl NativeRouterUsage {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let client = client.map(|client| client.0.clone());
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>, redis_ttl_seconds: u64) -> PyResult<Self> {
let remote = client.map(|client| remote(&client, redis_ttl_seconds));
let policy = router_tier_policies(py)?.usage;
Ok(Self(Arc::new(RouterUsage::new(client, policy))))
Ok(Self(Arc::new(RouterUsage::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn attach(&self, client: PyRef<'_, NativeRedisClient>, redis_ttl_seconds: u64) -> bool {
self.0.attach(remote(&client, redis_ttl_seconds))
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyCounter,
counter: PyData<Counter>,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment", store_error, async move {
redis_async(py, "async_increment", runtime_error, async move {
store
.increment(&counter.0, value.0, Duration::from_secs(ttl_seconds))
.await
@ -67,7 +53,7 @@ impl NativeRouterUsage {
requests: i64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_pipeline", store_error, async move {
redis_async(py, "async_increment_pipeline", runtime_error, async move {
store
.increment_deployment(
&deployment_id,
@ -83,14 +69,11 @@ impl NativeRouterUsage {
fn get_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PyCounter>,
counters: Vec<PyData<Counter>>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = counters
.into_iter()
.map(|counter| counter.0)
.collect::<Vec<_>>();
redis_async(py, "async_batch_get_cache", store_error, async move {
let counters = unwrap_all(counters);
redis_async(py, "async_batch_get_cache", runtime_error, async move {
let read = store.get_many(&counters).await?;
Ok(read
.into_iter()
@ -100,12 +83,12 @@ impl NativeRouterUsage {
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(store_error)
self.0.flush_local_counters().map_err(runtime_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, count) in self.0.snapshot().map_err(store_error)? {
for (key, count) in self.0.snapshot().map_err(runtime_error)? {
snapshot.set_item(key, PyNumber(count))?;
}
Ok(snapshot)

144
litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs
View file

@ -1,53 +1,27 @@
use std::{convert::Infallible, sync::Arc, time::Duration};
use std::sync::Arc;
use litellm_redis_store::spend_counters::{Error, LostSeed, SpendCounter, SpendCounters};
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, SpendCounter, SpendCounters},
};
use pyo3::{prelude::*, types::PyDict};
use super::{
NativePending, PyNumber, kind_error, redis_async, request::NativeRequestRedis, store_error, ttl,
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl, unwrap_all,
};
use crate::redis::client::NativeRedisClient;
fn spend_error(error: Error) -> PyErr {
kind_error(error.kind(), &error)
}
fn infallible(error: Infallible) -> PyErr {
match error {}
}
/// A `spend:*` key, refused when it names anything outside the spend counter keyspace.
pub(crate) struct PySpendCounter(SpendCounter);
impl<'a, 'py> FromPyObject<'a, 'py> for PySpendCounter {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let key: String = value.extract()?;
SpendCounter::parse(&key)
.map(Self)
.map_err(|error| PyValueError::new_err(error.to_string()))
}
}
fn counters(counters: Vec<PySpendCounter>) -> Vec<SpendCounter> {
counters.into_iter().map(|counter| counter.0).collect()
}
fn increments(
increments: Vec<(PySpendCounter, PyNumber)>,
) -> Vec<(SpendCounter, litellm_redis_store::Number)> {
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
increments
.into_iter()
.map(|(counter, amount)| (counter.0, amount.0))
.collect()
}
fn numbers(values: Vec<Option<litellm_redis_store::Number>>) -> Vec<Option<PyNumber>> {
values
.into_iter()
.map(|value| value.map(PyNumber))
.collect()
fn numbers(values: Option<Vec<Option<Number>>>) -> Option<Vec<Option<PyNumber>>> {
values.map(|values| values.into_iter().map(|value| value.map(PyNumber)).collect())
}
/// The proxy's spend counters, Redis first with this worker's memory copy behind it.
@ -60,55 +34,41 @@ impl NativeSpendCounters {
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn new(client: Option<PyRef<'_, NativeRedisClient>>, redis_ttl_seconds: u64) -> Self {
Self(Arc::new(SpendCounters::new(
client.map(|client| client.0.clone()),
Duration::from_secs(redis_ttl_seconds),
client.map(|client| remote(&client, redis_ttl_seconds)),
)))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
#[pyo3(signature = (client, redis_ttl_seconds = 60))]
fn attach(&self, client: PyRef<'_, NativeRedisClient>, redis_ttl_seconds: u64) -> bool {
self.0.attach(remote(&client, redis_ttl_seconds))
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
fn read<'py>(&self, py: Python<'py>, counter: PySpendCounter) -> PyResult<Bound<'py, PyAny>> {
fn read<'py>(&self, py: Python<'py>, counter: PyData<SpendCounter>) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_get_cache", store_error, async move {
redis_async(py, "async_get_cache", kind_error, async move {
let reading = store.read(&counter.0).await?;
let source: &'static str = reading.into();
Ok((source, reading.value().map(PyNumber)))
})
}
fn read_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PySpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_async(py, "async_batch_get_cache", infallible, async move {
Ok(store.read_many(&counters).await.map(numbers))
})
}
fn declare_read_many(
&self,
request: &NativeRequestRedis,
counters: Vec<PySpendCounter>,
) -> Option<NativePending> {
self.0
.declare_read_many_pre_call(&request.0, &self::counters(counters))
.map(|declared| {
NativePending::new("async_batch_get_cache", infallible, async move {
Ok(declared.resolve().await.map(numbers))
})
})
/// One `MGET` on the batch `send` picks; `None` without Redis. It resolves to `None` when the
/// read failed, so each caller reads on its own.
fn read_many(&self, counters: Vec<PyData<SpendCounter>>, send: &NativeSend) -> Option<NativePending> {
let read = self.0.read_many(send.send(), &unwrap_all(counters))?;
Some(pending("async_batch_get_cache", kind_error, read, numbers))
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
counter: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment", spend_error, async move {
redis_async(py, "async_increment", kind_error, async move {
store.increment(&counter.0, amount.0).await.map(PyNumber)
})
}
@ -116,11 +76,11 @@ impl NativeSpendCounters {
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<(PySpendCounter, PyNumber)>,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments = self::increments(increments);
redis_async(py, "async_increment_pipeline", spend_error, async move {
redis_async(py, "async_increment_pipeline", kind_error, async move {
let totals = store.increment_many(&increments).await?;
Ok(totals.into_iter().map(PyNumber).collect::<Vec<_>>())
})
@ -130,7 +90,7 @@ impl NativeSpendCounters {
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PySpendCounter, PyNumber)>,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> bool {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
@ -139,11 +99,11 @@ impl NativeSpendCounters {
fn set_max<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
counter: PyData<SpendCounter>,
spend: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_set_max", spend_error, async move {
redis_async(py, "async_set_max", kind_error, async move {
store
.set_max(&counter.0, spend.0)
.await
@ -155,7 +115,7 @@ impl NativeSpendCounters {
fn seed<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
counter: PyData<SpendCounter>,
spend: PyNumber,
increment_if_lost: bool,
) -> PyResult<Bound<'py, PyAny>> {
@ -165,7 +125,7 @@ impl NativeSpendCounters {
} else {
LostSeed::Answer
};
redis_async(py, "async_set_cache", spend_error, async move {
redis_async(py, "async_set_cache", kind_error, async move {
store.seed(&counter.0, spend.0, lost).await.map(PyNumber)
})
}
@ -174,12 +134,12 @@ impl NativeSpendCounters {
fn set<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_set_cache", spend_error, async move {
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&counter.0, spend.0, ttl(ttl_seconds)).await
})
}
@ -187,11 +147,11 @@ impl NativeSpendCounters {
fn delete<'py>(
&self,
py: Python<'py>,
counters: Vec<PySpendCounter>,
counters: Vec<PyData<SpendCounter>>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_async(py, "async_delete_cache", spend_error, async move {
let counters = unwrap_all(counters);
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&counters).await
})
}
@ -199,52 +159,52 @@ impl NativeSpendCounters {
fn refresh_ttl<'py>(
&self,
py: Python<'py>,
counter: PySpendCounter,
counter: PyData<SpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_refresh_ttl", spend_error, async move {
redis_async(py, "async_refresh_ttl", kind_error, async move {
store.refresh_ttl(&counter.0).await
})
}
fn get_local(&self, counter: PySpendCounter) -> PyResult<Option<PyNumber>> {
fn get_local(&self, counter: PyData<SpendCounter>) -> PyResult<Option<PyNumber>> {
Ok(self
.0
.get_local(&counter.0)
.map_err(store_error)?
.map_err(kind_error)?
.map(PyNumber))
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set_local(
&self,
counter: PySpendCounter,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
self.0
.set_local(&counter.0, spend.0, ttl(ttl_seconds))
.map_err(store_error)
.map_err(kind_error)
}
fn delete_local(&self, counter: PySpendCounter) -> PyResult<()> {
self.0.delete_local(&counter.0).map_err(store_error)
fn delete_local(&self, counter: PyData<SpendCounter>) -> PyResult<()> {
self.0.delete_local(&counter.0).map_err(kind_error)
}
fn increment_local(&self, counter: PySpendCounter, amount: PyNumber) -> PyResult<PyNumber> {
fn increment_local(&self, counter: PyData<SpendCounter>, amount: PyNumber) -> PyResult<PyNumber> {
self.0
.increment_local(&counter.0, amount.0)
.map(PyNumber)
.map_err(store_error)
.map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(store_error)
self.0.flush_local().map_err(kind_error)
}
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
let snapshot = PyDict::new(py);
for (key, spend) in self.0.snapshot().map_err(store_error)? {
for (key, spend) in self.0.snapshot().map_err(kind_error)? {
snapshot.set_item(key, PyNumber(spend))?;
}
Ok(snapshot)

11
litellm-rust/crates/redis-client/AGENTS.md
View file

@ -23,21 +23,23 @@
- The rule never changes, because Python and Rust pods share keys in the same Redis
## Commands
- One typed method per command, built with `redis::cmd`. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
- A command is a value, one constructor in `commands.rs` returning `Command<T>`: its alone and pipelined forms and its reply decode. `RedisClient::send`, `Pipeline::add` and `RedisBatch::add` are the only entry points, and each applies its namespace to the keys once. A new command is one constructor, never a method on an executor. `ping` and `subscribe` stay methods. Empty batches (MGET, DEL) return without a round trip, since Redis rejects an empty MGET
- A store attaches its value codec with `Command::map`, never a keyspace codec inside this crate
- On cluster, redis-rs splits MGET and DEL by slot and returns MGET values in input order. Other multi-key commands are not split, so a new one needs a cross-slot test
- Every command is the one `RedisCache` sends, argument for argument, so `MONITOR` on a mixed fleet shows one shape. A difference from main's Python is a defect even when the result is the same
- TTLs are `Duration` and are sent in whole seconds (`SET ... EX`, `EXPIRE`, `TTL`), like redis-py's `int(ttl)`. `set_nx` sends `EX` before `NX`, the order redis-py builds
- Numbers go out as `Amount`, written the way redis-py encodes them: an `int` with `str()` and a `float` with `repr()` from `python-compat`
- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. In a pipeline `incr_by_float` is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when a TTL is given
- Increments are plain commands, never a script. `incr_by` is `increment_cache`: `INCRBY`, `TTL`, then `EXPIRE` only on a key with no TTL, each its own round trip. `incr_by_float` alone is `async_increment`: `INCRBYFLOAT` with `EXPIRE` (`Expiry::Refresh`) or `TTL` (`Expiry::KeepExisting`) in one round trip, and for `KeepExisting` a third, `EXPIRE`, only on a key with no TTL. Pipelined it is `RedisBatch`'s `_Increment`, `INCRBYFLOAT` then `EXPIRE` when an expiry is given, `KeepExisting` lowered to `Refresh`. `Plan` keeps only the cases whose alone and pipelined forms differ
## Scripts and pipelines
- A `Script` declares every key it touches in KEYS. On cluster an undeclared key works or fails depending on which node owns its slot, so it can pass tests and break in production
- `eval` goes through redis-rs `invoke_async`, which loads the script on NOSCRIPT. Inside a pipeline a script is sent as `EVALSHA`, like `RedisBatch`. A `NOSCRIPT` reply means that op ran nothing, so it alone runs again after the pipeline through `invoke_async` (`EVALSHA`, `SCRIPT LOAD`, `EVALSHA`), like `RedisBatch`'s `_settle_alone`. It then lands after the ops around it, as it does in Python
- `Entry` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice
- `Plan::Script` keeps a script call apart from plain commands, and the one resend takes only a `ScriptCall` whose reply was `NOSCRIPT`. Nothing can resend a pipeline or a command that already ran, so a retry can never apply an increment twice
- `RedisBatch::pipelined` is one `Pipeline` whose ops answer as typed `Declared<T>`s. A store that sends its own pipeline uses it rather than decoding `run`'s raw `Value`s
- A pipeline is never atomic. Every op gets its own `Result`, from `.ignore_errors()`, so one rejected op leaves the others alone. A transport failure fails every op of that round trip
- redis-rs rejects an async cluster pipeline whose ops span slots. `Pipeline` groups ops by the slot of their first key and runs the slots concurrently, returning results in input order. Within a slot, runs of idempotent ops (GET, SET, DEL, EXPIRE) share a request and every script op gets its own, sent in order. redis-rs retries a cluster request whole, so a script op batched with others would be applied twice when only a neighbour failed
- Pipeline ops take one key each. A multi-key op would span slots inside one group. An op can send more than one command for that key (`INCRBYFLOAT` then `EXPIRE`), and settles from the first reply, failing if any of its replies failed
- Command builders live in `commands.rs` and are shared by `RedisClient` and `Pipeline`, so a command is written once
- Command constructors live in `commands.rs` and every executor sends them, so a command is written once
## Pub/sub
- Classic `PUBLISH` / `SUBSCRIBE`, never sharded, since redis-py publishes and subscribes through one plain client and cluster broadcasts classic messages to every node. Payloads are raw bytes, and the codec belongs to the store that owns the channel
@ -50,6 +52,7 @@
## Batches
- `RedisBatch` is litellm's `RedisBatch` (`litellm/caching/redis_batch.py`). It keeps each declared op as data and, on a standalone or Sentinel Redis, sends them later as one `Pipeline`, so it inherits every pipeline rule above. On a cluster it never pipelines and runs each op on its own, concurrently, the way `_settle_alone` calls the `RedisCache` method behind it. An increment there keeps a TTL the key already has, as `async_increment` does, where the pipelined one refreshes it like `_Increment`
- `mget` is one `MGET` of the keys in order with repeats read once, like `_MGet`. Use it instead of one `get` per key whenever Python reads several keys in one call It knows nothing about requests. The host that owns a request (Axum handler, pyo3 adapter) builds one and passes it down as a value, never through a task-local
- `RedisBatch::alone` runs every op on its own like the cluster path, so a fresh alone batch is "send it now" and `Declared<T>` is the one reply handle on every route
- Each op returns a `Declared<T>` backed by a oneshot. `flush` sends everything declared so far, one flush at a time, and an empty flush makes no round trip. `resolve` flushes only when its op has not gone out yet, so an op declared after another caller resolved costs a second round trip. Stages that know their reads up front declare all of them and call `flush` once
- An op whose flush was cancelled resolves as `Transport`, because it may already have reached Redis
- `settled` waits for whichever flush carries the op and never sends one itself, for follow-ups that must not force a flush. `RedisClient::same_client` tells clones of one client apart from a second connection, so a host can keep one batch per client

252
litellm-rust/crates/redis-client/src/batch.rs
View file

@ -1,160 +1,78 @@
use std::{
io,
sync::{Arc, Mutex, MutexGuard},
time::Duration,
};
use futures_util::future::join_all;
use redis::{FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value};
use redis::{RedisError, Value};
use tokio::sync::oneshot::{self, error::TryRecvError};
use crate::{Amount, Error, Expiry, Pipeline, RedisClient, Script};
use crate::{Command, Error, Pipeline, RedisClient, commands::Op};
/// Ops declared by many callers and sent together, like litellm's `RedisBatch`: one pipeline on a
/// standalone or Sentinel Redis, and each op on its own on a cluster.
/// standalone or Sentinel Redis, and each op on its own on a cluster or in an alone batch.
#[derive(Clone)]
pub struct RedisBatch(Arc<Shared>);
#[derive(Clone, Copy)]
enum Mode {
/// litellm's `RedisBatch`: one pipeline, or each op on its own on a cluster.
Batch,
/// Each op on its own, the way the `RedisCache` method behind it sends it.
Alone,
/// One redis-py pipeline on every topology, split by slot on a cluster.
Pipeline,
}
struct Shared {
client: RedisClient,
mode: Mode,
pending: Mutex<Pending>,
flushing: tokio::sync::Mutex<()>,
}
#[derive(Default)]
struct Pending {
ops: Vec<BatchOp>,
ops: Vec<Op>,
replies: Vec<oneshot::Sender<Result<Value, Error>>>,
}
/// One declared op, kept as data so a flush can pipeline it or run it alone.
enum BatchOp {
Get(String),
MGet(Vec<String>),
Set {
key: String,
value: Vec<u8>,
ttl: Duration,
},
Del(String),
Expire {
key: String,
ttl: Duration,
},
IncrByFloat {
key: String,
amount: Amount,
ttl: Option<Duration>,
},
Eval {
script: Script,
keys: Vec<String>,
args: Vec<Vec<u8>>,
},
}
type Decode<T> = Box<dyn FnOnce(Value) -> Result<T, Error> + Send>;
#[must_use]
pub struct Declared<T> {
batch: RedisBatch,
reply: oneshot::Receiver<Result<Value, Error>>,
decode: fn(Value) -> Result<T, Error>,
decode: Decode<T>,
}
impl RedisBatch {
pub(crate) fn new(client: RedisClient) -> Self {
Self::build(client, Mode::Batch)
}
/// A batch whose every op goes out on its own, the way the `RedisCache` method behind it sends
/// it, so a fresh one is "send it now": the first `resolve` sends it.
pub fn alone(client: RedisClient) -> Self {
Self::build(client, Mode::Alone)
}
/// A batch sent as one redis-py pipeline, like `async_increment_pipeline`, on every topology.
pub fn pipelined(client: RedisClient) -> Self {
Self::build(client, Mode::Pipeline)
}
fn build(client: RedisClient, mode: Mode) -> Self {
Self(Arc::new(Shared {
client,
mode,
pending: Mutex::new(Pending::default()),
flushing: tokio::sync::Mutex::new(()),
}))
}
pub fn get<T: FromRedisValue>(&self, key: &str) -> Declared<Option<T>> {
self.declare(BatchOp::Get(key.to_owned()), decode)
}
/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`.
pub fn mget<T: FromRedisValue>(&self, keys: &[&str]) -> Declared<Vec<Option<T>>> {
let mut unique: Vec<String> = Vec::with_capacity(keys.len());
for key in keys {
if !unique.iter().any(|seen| seen == key) {
unique.push((*key).to_owned());
}
}
self.declare(BatchOp::MGet(unique), decode)
}
pub fn set(&self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Declared<()> {
let value = value.to_redis_args().concat();
self.declare(
BatchOp::Set {
key: key.to_owned(),
value,
ttl,
},
decode,
)
}
pub fn del(&self, key: &str) -> Declared<usize> {
self.declare(BatchOp::Del(key.to_owned()), decode)
}
pub fn expire(&self, key: &str, ttl: Duration) -> Declared<bool> {
self.declare(
BatchOp::Expire {
key: key.to_owned(),
ttl,
},
decode,
)
}
/// `RedisBatch.increment`: `INCRBYFLOAT`, then `EXPIRE` when a TTL is given.
pub fn incr_by_float(&self, key: &str, amount: Amount, ttl: Option<Duration>) -> Declared<f64> {
self.declare(
BatchOp::IncrByFloat {
key: key.to_owned(),
amount,
ttl,
},
decode,
)
}
pub fn eval<T: FromRedisValue>(
&self,
script: &Script,
keys: &[&str],
args: impl ToRedisArgs,
) -> Declared<T> {
self.declare(
BatchOp::Eval {
script: script.clone(),
keys: keys.iter().map(|key| (*key).to_owned()).collect(),
args: args.to_redis_args(),
},
decode,
)
}
pub async fn flush(&self) {
let _flushing = self.0.flushing.lock().await;
let Pending { ops, replies } = std::mem::take(&mut *self.pending());
if replies.is_empty() {
return;
}
let client = &self.0.client;
let results = match client.is_cluster() {
true => join_all(ops.iter().map(|op| op.run_alone(client))).await,
false => client.run(&pipeline(client, &ops)).await,
};
for (reply, result) in replies.into_iter().zip(results) {
let _ = reply.send(result);
}
}
fn declare<T>(&self, op: BatchOp, decode: fn(Value) -> Result<T, Error>) -> Declared<T> {
pub fn add<T>(&self, command: Command<T>) -> Declared<T> {
let (op, decode) = command.prepare(self.0.client.namespace());
let (sender, reply) = oneshot::channel();
let mut pending = self.pending();
pending.ops.push(op);
@ -166,6 +84,31 @@ impl RedisBatch {
}
}
pub async fn flush(&self) {
let _flushing = self.0.flushing.lock().await;
let Pending { ops, replies } = std::mem::take(&mut *self.pending());
if replies.is_empty() {
return;
}
let client = &self.0.client;
let alone = match self.0.mode {
Mode::Alone => true,
Mode::Batch => client.is_cluster(),
Mode::Pipeline => false,
};
let results = match alone {
true => join_all(ops.into_iter().map(|op| client.run_alone(op))).await,
false => {
client
.run(&Pipeline::of(client.namespace().clone(), ops))
.await
}
};
for (reply, result) in replies.into_iter().zip(results) {
let _ = reply.send(result);
}
}
fn pending(&self) -> MutexGuard<'_, Pending> {
self.0
.pending
@ -174,74 +117,6 @@ impl RedisBatch {
}
}
fn pipeline(client: &RedisClient, ops: &[BatchOp]) -> Pipeline {
ops.iter().fold(client.pipeline(), |pipeline, op| match op {
BatchOp::Get(key) => pipeline.get(key),
BatchOp::MGet(keys) => {
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
pipeline.mget(&keys)
}
BatchOp::Set { key, value, ttl } => pipeline.set(key, value.as_slice(), *ttl),
BatchOp::Del(key) => pipeline.del(key),
BatchOp::Expire { key, ttl } => pipeline.expire(key, *ttl),
BatchOp::IncrByFloat { key, amount, ttl } => pipeline.incr_by_float(key, *amount, *ttl),
BatchOp::Eval { script, keys, args } => {
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
pipeline.eval(script, &keys, args.as_slice())
}
})
}
impl BatchOp {
/// What `RedisBatch` runs for an op on a cluster, where it never pipelines: the `RedisCache`
/// method behind it, so an increment there keeps a TTL it already has like `async_increment`.
async fn run_alone(&self, client: &RedisClient) -> Result<Value, Error> {
match self {
Self::Get(key) => client
.get::<Value>(key)
.await
.map(|value| value.unwrap_or(Value::Nil)),
Self::MGet(keys) => {
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
let values = client.mget::<Value>(&keys).await?;
Ok(Value::Array(
values
.into_iter()
.map(|value| value.unwrap_or(Value::Nil))
.collect(),
))
}
Self::Set { key, value, ttl } => client
.set(key, value.as_slice(), *ttl)
.await
.map(|()| Value::Okay),
Self::Del(key) => client
.del(&[key])
.await
.map(|deleted| Value::Int(i64::try_from(deleted).unwrap_or(i64::MAX))),
Self::Expire { key, ttl } => client
.expire(key, *ttl)
.await
.map(|set| Value::Int(i64::from(set))),
Self::IncrByFloat { key, amount, ttl } => {
let value = match ttl {
Some(ttl) => {
client
.incr_by_float(key, *amount, Expiry::KeepExisting(*ttl))
.await?
}
None => client.incr_by_float_without_ttl(key, *amount).await?,
};
Ok(Value::Double(value))
}
Self::Eval { script, keys, args } => {
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
client.eval::<Value>(script, &keys, args.as_slice()).await
}
}
}
}
impl<T> Declared<T> {
pub async fn resolve(mut self) -> Result<T, Error> {
match self.reply.try_recv() {
@ -257,10 +132,11 @@ impl<T> Declared<T> {
let reply = self.reply.await.unwrap_or_else(|_| Err(abandoned()));
reply.and_then(self.decode)
}
}
fn decode<T: FromRedisValue>(value: Value) -> Result<T, Error> {
T::from_redis_value(value).map_err(|error| Error::Decode(error.into()))
/// Every reply in order, flushing their batch once if nobody has yet.
pub async fn all(declared: impl IntoIterator<Item = Self>) -> Vec<Result<T, Error>> {
join_all(declared.into_iter().map(Self::resolve)).await
}
}
fn abandoned() -> Error {

241
litellm-rust/crates/redis-client/src/client.rs
View file

@ -6,11 +6,13 @@ use std::{
use futures_util::stream::BoxStream;
use litellm_circuit_breaker::{Admission, Breaker, Failure, State};
use redis::{Cmd, FromRedisValue, RedisError, ToRedisArgs, ToSingleRedisArg, Value};
use redis::{Cmd, RedisError, Value};
use crate::{
Amount, CallEvent, ConnectError, Error, Expiry, FailureKind, Namespace, Operation, Outcome,
Pipeline, RedisBatch, RedisConfig, RedisEvents, Script, commands, connection::Connection,
CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline,
RedisBatch, RedisConfig, RedisEvents,
commands::{self, Op, Plan, Trips},
connection::Connection,
pubsub::Subscriber,
};
@ -74,126 +76,70 @@ impl RedisClient {
}
pub async fn ping(&self) -> Result<(), Error> {
self.query("PING", &commands::ping()).await
}
pub async fn get<T: FromRedisValue>(&self, key: &str) -> Result<Option<T>, Error> {
self.query("GET", &commands::get(self.key(key))).await
}
pub async fn mget<T: FromRedisValue>(&self, keys: &[&str]) -> Result<Vec<Option<T>>, Error> {
if keys.is_empty() {
return Ok(Vec::new());
}
self.query("MGET", &commands::mget(self.keys(keys))).await
}
pub async fn set(
&self,
key: &str,
value: impl ToSingleRedisArg,
ttl: Duration,
) -> Result<(), Error> {
self.query("SET", &commands::set(self.key(key), value, ttl))
.await
}
pub async fn set_nx(
&self,
key: &str,
value: impl ToSingleRedisArg,
ttl: Duration,
) -> Result<bool, Error> {
let reply: Option<String> = self
.query("SET", &commands::set_nx(self.key(key), value, ttl))
.await?;
Ok(reply.is_some())
}
pub async fn del(&self, keys: &[&str]) -> Result<usize, Error> {
if keys.is_empty() {
return Ok(0);
}
self.query("DEL", &commands::del(self.keys(keys))).await
}
pub async fn expire(&self, key: &str, ttl: Duration) -> Result<bool, Error> {
self.query("EXPIRE", &commands::expire(self.key(key), ttl))
.await
}
pub async fn ttl(&self, key: &str) -> Result<Ttl, Error> {
let reply: i64 = self.query("TTL", &commands::ttl(self.key(key))).await?;
Ok(match reply {
-2 => Ttl::Missing,
-1 => Ttl::Persistent,
seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())),
})
}
/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no
/// TTL, each its own round trip.
pub async fn incr_by(&self, key: &str, delta: i64, ttl: Duration) -> Result<i64, Error> {
let key = self.key(key);
self.guarded(Operation::Command("INCRBY"), async {
let mut connection = self.0.connection.clone();
let value: i64 = commands::incr_by(key.clone(), delta)
.query_async(&mut connection)
.await?;
let current: i64 = commands::ttl(key.clone())
.query_async(&mut connection)
.await?;
if current == -1 {
commands::expire(key, ttl)
.exec_async(&mut connection)
.await?;
}
Ok(value)
self.guarded(Operation::Command("PING"), async {
commands::ping()
.query_async(&mut self.0.connection.clone())
.await
})
.await
}
/// `INCRBYFLOAT` alone, `_incrbyfloat_with_ttl` without a TTL.
pub(crate) async fn incr_by_float_without_ttl(
&self,
key: &str,
amount: Amount,
) -> Result<f64, Error> {
self.query(
"INCRBYFLOAT",
&commands::incr_by_float(self.key(key), amount),
)
.await
/// Sends one command on its own, the way the `RedisCache` method behind it sends it.
pub async fn send<T>(&self, command: Command<T>) -> Result<T, Error> {
let (op, decode) = command.prepare(&self.0.namespace);
self.run_alone(op).await.and_then(decode)
}
/// `RedisCache.async_increment`: `INCRBYFLOAT` and its TTL command in one round trip, and for
/// [`Expiry::KeepExisting`] a third, `EXPIRE`, only when the key has no TTL.
pub async fn incr_by_float(
&self,
key: &str,
amount: Amount,
expiry: Expiry,
) -> Result<f64, Error> {
let key = self.key(key);
self.guarded(Operation::Command("INCRBYFLOAT"), async {
pub(crate) async fn run_alone(&self, op: Op) -> Result<Value, Error> {
let Op {
operation, plan, ..
} = op;
self.guarded(operation, async {
let mut connection = self.0.connection.clone();
let mut pipeline = redis::pipe();
pipeline.add_command(commands::incr_by_float(key.clone(), amount));
match expiry {
Expiry::Refresh(ttl) => {
let (value, _): (f64, bool) = pipeline
.add_command(commands::expire(key, ttl))
.query_async(&mut connection)
.await?;
Ok(value)
match plan {
Plan::Reply(value) => Ok(value),
Plan::Commands(commands) => match <[Cmd; 1]>::try_from(commands) {
Ok([command]) => command.query_async(&mut connection).await,
Err(commands) => {
let mut pipeline = redis::pipe();
for command in commands {
pipeline.add_command(command);
}
let replies: Vec<Value> = pipeline.query_async(&mut connection).await?;
Ok(replies.into_iter().next().unwrap_or(Value::Nil))
}
},
Plan::Script(call) => {
call.script
.invocation(&call.keys, &call.args)
.invoke_async(&mut connection)
.await
}
Expiry::KeepExisting(ttl) => {
let (value, current): (f64, i64) = pipeline
.add_command(commands::ttl(key.clone()))
.query_async(&mut connection)
.await?;
Plan::KeepingTtl {
increment,
key,
ttl,
trips,
} => {
let (value, current): (Value, i64) = match trips {
Trips::Shared => {
redis::pipe()
.add_command(increment)
.add_command(commands::ttl_cmd(&key))
.query_async(&mut connection)
.await?
}
Trips::Separate => (
increment.query_async(&mut connection).await?,
commands::ttl_cmd(&key)
.query_async(&mut connection)
.await?,
),
};
if current == -1 {
commands::expire(key, ttl)
redis::cmd("EXPIRE")
.arg(&key)
.arg(ttl.as_secs())
.exec_async(&mut connection)
.await?;
}
@ -204,60 +150,20 @@ impl RedisClient {
.await
}
pub async fn eval<T: FromRedisValue>(
&self,
script: &Script,
keys: &[&str],
args: impl ToRedisArgs,
) -> Result<T, Error> {
let keys = self.keys(keys);
let args = args.to_redis_args();
self.guarded(Operation::Script, async {
script
.invocation(&keys, &args)
.invoke_async(&mut self.0.connection.clone())
.await
})
.await
}
/// The channel is sent as given, without the namespace, like redis-py's `publish`.
pub async fn publish(&self, channel: &str, payload: &[u8]) -> Result<(), Error> {
self.query("PUBLISH", &commands::publish(channel, payload))
.await
}
/// Messages published on `channel`, on a connection of its own. After the connection is lost
/// it reconnects and resubscribes, and messages published meanwhile are not delivered.
pub fn subscribe(&self, channel: &str) -> BoxStream<'static, Vec<u8>> {
self.0.subscriber.clone().messages(channel.to_owned())
}
/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the
/// scripts `RedisCache` runs that way rather than through `register_script`.
pub async fn eval_source<T: FromRedisValue>(
&self,
source: &str,
keys: &[&str],
args: impl ToRedisArgs,
) -> Result<T, Error> {
let keys = self.keys(keys);
let command = redis::cmd("EVAL")
.arg(source)
.arg(keys.len())
.arg(keys)
.arg(args)
.to_owned();
self.guarded(Operation::Script, async {
command.query_async(&mut self.0.connection.clone()).await
})
.await
}
pub fn pipeline(&self) -> Pipeline {
Pipeline::new(self.0.namespace.clone())
}
pub(crate) fn namespace(&self) -> &Namespace {
&self.0.namespace
}
pub fn is_cluster(&self) -> bool {
matches!(self.0.connection, Connection::Cluster(_))
}
@ -305,25 +211,6 @@ impl RedisClient {
results
}
fn key(&self, key: &str) -> String {
self.0.namespace.apply(key)
}
fn keys(&self, keys: &[&str]) -> Vec<String> {
keys.iter().map(|key| self.key(key)).collect()
}
async fn query<T: FromRedisValue>(
&self,
name: &'static str,
command: &Cmd,
) -> Result<T, Error> {
self.guarded(Operation::Command(name), async {
command.query_async(&mut self.0.connection.clone()).await
})
.await
}
async fn guarded<T>(
&self,
operation: Operation,

370
litellm-rust/crates/redis-client/src/commands.rs
View file

@ -1,7 +1,11 @@
use std::time::Duration;
use litellm_python_compat::repr::float_repr;
use redis::{Cmd, RedisWrite, ToRedisArgs, ToSingleRedisArg};
use redis::{
Cmd, ErrorKind, FromRedisValue, RedisError, RedisWrite, ToRedisArgs, ToSingleRedisArg, Value,
};
use crate::{Error, Namespace, Operation, Script, Ttl, script::ScriptCall};
/// How an increment treats the key's TTL, the two ways `RedisCache` increments with one.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -9,6 +13,7 @@ pub enum Expiry {
/// `async_increment(refresh_ttl=True)` and every pipelined increment: `EXPIRE` after each write.
Refresh(Duration),
/// `async_increment(refresh_ttl=False)`: read the `TTL` and `EXPIRE` only a key that has none.
/// Pipelined it is lowered to [`Expiry::Refresh`], like `RedisBatch.increment`.
KeepExisting(Duration),
}
@ -30,58 +35,357 @@ impl ToRedisArgs for Amount {
impl ToSingleRedisArg for Amount {}
pub(crate) fn ping() -> Cmd {
redis::cmd("PING")
type Build = Box<dyn FnOnce(Vec<String>) -> Plan + Send>;
type Decode<T> = Box<dyn FnOnce(Value) -> Result<T, Error> + Send>;
/// One Redis command as a value: what it sends alone and pipelined, and how its reply decodes. Every
/// executor (`RedisClient::send`, `Pipeline::add`, `RedisBatch::add`) applies its namespace to the
/// keys once.
#[must_use]
pub struct Command<T> {
operation: Operation,
keys: Vec<String>,
idempotent: bool,
build: Build,
decode: Decode<T>,
}
pub(crate) fn publish(channel: &str, payload: &[u8]) -> Cmd {
redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned()
/// What a command sends once its keys carry the namespace. Only the cases whose alone and pipelined
/// forms differ get a variant of their own.
pub(crate) enum Plan {
/// Answered without a round trip, like an empty `MGET` or `DEL`.
Reply(Value),
/// Sent together; the reply is the first command's, failing if any failed.
Commands(Vec<Cmd>),
/// Alone through `invoke_async`, pipelined as `EVALSHA` resent once on `NOSCRIPT`.
Script(ScriptCall),
/// Alone the increment reads the `TTL` and sends `EXPIRE` only to a key without one, pipelined it
/// is the increment then `EXPIRE`.
KeepingTtl {
increment: Cmd,
key: String,
ttl: Duration,
trips: Trips,
},
}
pub(crate) fn get(key: String) -> Cmd {
redis::cmd("GET").arg(key).to_owned()
#[derive(Clone, Copy)]
pub(crate) enum Trips {
/// `async_increment`: the increment and its `TTL` in one round trip.
Shared,
/// `increment_cache`: the increment and its `TTL` each on its own round trip.
Separate,
}
pub(crate) fn mget(keys: Vec<String>) -> Cmd {
redis::cmd("MGET").arg(keys).to_owned()
/// A command with the namespace applied, ready for an executor to send.
pub(crate) struct Op {
pub(crate) operation: Operation,
pub(crate) key: String,
pub(crate) idempotent: bool,
pub(crate) plan: Plan,
}
// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())`
pub(crate) fn set(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("EX")
.arg(ttl.as_secs())
.to_owned()
impl Plan {
/// What a pipeline sends for this op.
pub(crate) fn pipelined(&self) -> Vec<Cmd> {
match self {
Self::Reply(_) => Vec::new(),
Self::Commands(commands) => commands.clone(),
Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)],
Self::KeepingTtl {
increment, key, ttl, ..
} => vec![increment.clone(), expire_cmd(key, *ttl)],
}
}
}
pub(crate) fn set_nx(key: String, value: impl ToSingleRedisArg, ttl: Duration) -> Cmd {
redis::cmd("SET")
.arg(key)
.arg(value)
.arg("EX")
.arg(ttl.as_secs())
.arg("NX")
.to_owned()
impl<T> Command<T> {
fn new(
operation: Operation,
keys: Vec<String>,
idempotent: bool,
build: impl FnOnce(Vec<String>) -> Plan + Send + 'static,
decode: impl FnOnce(Value) -> Result<T, Error> + Send + 'static,
) -> Self {
Self {
operation,
keys,
idempotent,
build: Box::new(build),
decode: Box::new(decode),
}
}
/// Attaches a codec over the decoded reply, such as a store's value format.
pub fn map<U: 'static>(self, decode: fn(T) -> Result<U, Error>) -> Command<U>
where
T: 'static,
{
let inner = self.decode;
Command {
operation: self.operation,
keys: self.keys,
idempotent: self.idempotent,
build: self.build,
decode: Box::new(move |value| inner(value).and_then(decode)),
}
}
pub(crate) fn prepare(self, namespace: &Namespace) -> (Op, Decode<T>) {
let keys: Vec<String> = self.keys.iter().map(|key| namespace.apply(key)).collect();
let key = keys.first().cloned().unwrap_or_default();
let op = Op {
operation: self.operation,
key,
idempotent: self.idempotent,
plan: (self.build)(keys),
};
(op, self.decode)
}
}
pub(crate) fn del(keys: Vec<String>) -> Cmd {
redis::cmd("DEL").arg(keys).to_owned()
fn decode<T: FromRedisValue>(value: Value) -> Result<T, Error> {
T::from_redis_value(value).map_err(|error| Error::Decode(error.into()))
}
pub(crate) fn expire(key: String, ttl: Duration) -> Cmd {
fn owned(keys: &[&str]) -> Vec<String> {
keys.iter().map(|key| (*key).to_owned()).collect()
}
fn single(keys: Vec<String>) -> String {
keys.into_iter().next().unwrap_or_default()
}
fn expire_cmd(key: &str, ttl: Duration) -> Cmd {
redis::cmd("EXPIRE").arg(key).arg(ttl.as_secs()).to_owned()
}
pub(crate) fn ttl(key: String) -> Cmd {
pub(crate) fn ttl_cmd(key: &str) -> Cmd {
redis::cmd("TTL").arg(key).to_owned()
}
pub(crate) fn incr_by(key: String, delta: i64) -> Cmd {
redis::cmd("INCRBY").arg(key).arg(delta).to_owned()
/// A reply whose shape a codec over it refuses, as `what` describes.
pub fn unexpected_reply(what: &'static str) -> Error {
Error::Decode(RedisError::from((ErrorKind::UnexpectedReturnType, what)))
}
pub(crate) fn incr_by_float(key: String, amount: Amount) -> Cmd {
redis::cmd("INCRBYFLOAT").arg(key).arg(amount).to_owned()
pub fn get<T: FromRedisValue + 'static>(key: &str) -> Command<Option<T>> {
Command::new(
Operation::Command("GET"),
vec![key.to_owned()],
true,
|keys| Plan::Commands(vec![redis::cmd("GET").arg(single(keys)).to_owned()]),
decode,
)
}
/// One `MGET` of `keys` in order, repeated keys read once, like `RedisBatch.mget`. No keys is no
/// round trip, since Redis rejects an empty `MGET`. Pipelined, every key must share a hash slot on a
/// cluster, and a batch never pipelines there.
pub fn mget<T: FromRedisValue + 'static>(keys: &[&str]) -> Command<Vec<Option<T>>> {
let unique = keys.iter().fold(Vec::<String>::new(), |mut unique, key| {
if !unique.iter().any(|seen| seen == key) {
unique.push((*key).to_owned());
}
unique
});
Command::new(
Operation::Command("MGET"),
unique,
true,
|keys| match keys.is_empty() {
true => Plan::Reply(Value::Array(Vec::new())),
false => Plan::Commands(vec![redis::cmd("MGET").arg(keys).to_owned()]),
},
decode,
)
}
// redis-py's `set(ex=...)` takes whole seconds, as `int(ttl)` or a `timedelta`'s `int(total_seconds())`
pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()> {
let value = value.to_redis_args().concat();
Command::new(
Operation::Command("SET"),
vec![key.to_owned()],
true,
move |keys| {
Plan::Commands(vec![
redis::cmd("SET")
.arg(single(keys))
.arg(value)
.arg("EX")
.arg(ttl.as_secs())
.to_owned(),
])
},
decode,
)
}
/// `SET ... EX ... NX`, `EX` before `NX` as redis-py builds it. `true` when this call wrote the key.
pub fn set_nx(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<bool> {
let value = value.to_redis_args().concat();
Command::new(
Operation::Command("SET"),
vec![key.to_owned()],
false,
move |keys| {
Plan::Commands(vec![
redis::cmd("SET")
.arg(single(keys))
.arg(value)
.arg("EX")
.arg(ttl.as_secs())
.arg("NX")
.to_owned(),
])
},
|value| decode::<Option<String>>(value).map(|reply| reply.is_some()),
)
}
/// One multi-key `DEL`, no round trip for no keys. Pipelined, every key must share a hash slot.
pub fn del(keys: &[&str]) -> Command<usize> {
Command::new(
Operation::Command("DEL"),
owned(keys),
true,
|keys| match keys.is_empty() {
true => Plan::Reply(Value::Int(0)),
false => Plan::Commands(vec![redis::cmd("DEL").arg(keys).to_owned()]),
},
decode,
)
}
pub fn expire(key: &str, ttl: Duration) -> Command<bool> {
Command::new(
Operation::Command("EXPIRE"),
vec![key.to_owned()],
true,
move |keys| Plan::Commands(vec![expire_cmd(&single(keys), ttl)]),
decode,
)
}
pub fn ttl(key: &str) -> Command<Ttl> {
Command::new(
Operation::Command("TTL"),
vec![key.to_owned()],
true,
|keys| Plan::Commands(vec![ttl_cmd(&single(keys))]),
|value| {
Ok(match decode::<i64>(value)? {
-2 => Ttl::Missing,
-1 => Ttl::Persistent,
seconds => Ttl::Expires(Duration::from_secs(seconds.unsigned_abs())),
})
},
)
}
/// `RedisCache.increment_cache`: `INCRBY`, then `TTL`, then `EXPIRE` only when the key has no TTL,
/// each its own round trip.
pub fn incr_by(key: &str, delta: i64, ttl: Duration) -> Command<i64> {
Command::new(
Operation::Command("INCRBY"),
vec![key.to_owned()],
false,
move |keys| {
let key = single(keys);
Plan::KeepingTtl {
increment: redis::cmd("INCRBY").arg(&key).arg(delta).to_owned(),
key,
ttl,
trips: Trips::Separate,
}
},
decode,
)
}
/// `INCRBYFLOAT`, with the TTL handled as `expiry` says, or left alone without one. Alone it is
/// `RedisCache.async_increment`, pipelined `RedisBatch.increment`.
pub fn incr_by_float(key: &str, amount: Amount, expiry: Option<Expiry>) -> Command<f64> {
Command::new(
Operation::Command("INCRBYFLOAT"),
vec![key.to_owned()],
false,
move |keys| {
let key = single(keys);
let increment = redis::cmd("INCRBYFLOAT").arg(&key).arg(amount).to_owned();
match expiry {
None => Plan::Commands(vec![increment]),
Some(Expiry::Refresh(ttl)) => {
Plan::Commands(vec![increment, expire_cmd(&key, ttl)])
}
Some(Expiry::KeepExisting(ttl)) => Plan::KeepingTtl {
increment,
key,
ttl,
trips: Trips::Shared,
},
}
},
decode,
)
}
/// A registered script, like redis-py's `Script`. Every key it touches must be in `keys`.
pub fn eval<T: FromRedisValue + 'static>(
script: &Script,
keys: &[&str],
args: impl ToRedisArgs,
) -> Command<T> {
let script = script.clone();
let args = args.to_redis_args();
Command::new(
Operation::Script,
owned(keys),
false,
move |keys| Plan::Script(ScriptCall { script, keys, args }),
decode,
)
}
/// A script sent with its source as `EVAL` every time, like redis-py's `Redis.eval`, for the scripts
/// `RedisCache` runs that way rather than through `register_script`.
pub fn eval_source<T: FromRedisValue + 'static>(
source: &'static str,
keys: &[&str],
args: impl ToRedisArgs,
) -> Command<T> {
let args = args.to_redis_args();
Command::new(
Operation::Script,
owned(keys),
false,
move |keys| {
Plan::Commands(vec![
redis::cmd("EVAL")
.arg(source)
.arg(keys.len())
.arg(keys)
.arg(args)
.to_owned(),
])
},
decode,
)
}
/// The channel is sent as given, without the namespace, like redis-py's `publish`.
pub fn publish(channel: &str, payload: &[u8]) -> Command<()> {
let command = redis::cmd("PUBLISH").arg(channel).arg(payload).to_owned();
Command::new(
Operation::Command("PUBLISH"),
Vec::new(),
false,
move |_| Plan::Commands(vec![command]),
|_| Ok(()),
)
}
pub(crate) fn ping() -> Cmd {
redis::cmd("PING")
}

4
litellm-rust/crates/redis-client/src/lib.rs
View file

@ -1,6 +1,6 @@
mod batch;
mod client;
mod commands;
pub mod commands;
mod config;
mod connection;
mod error;
@ -12,7 +12,7 @@ mod script;
pub use batch::{Declared, RedisBatch};
pub use client::{RedisClient, Ttl, WeakRedisClient};
pub use commands::{Amount, Expiry};
pub use commands::{Amount, Command, Expiry};
pub use config::{Credentials, CredentialsProvider, DEFAULT_TIMEOUT, RedisConfig, Timeouts};
pub use error::{ConnectError, Error, KeyError};
pub use events::{CallEvent, FailureKind, Operation, Outcome, RedisEvents};

117
litellm-rust/crates/redis-client/src/pipeline.rs
View file

@ -1,13 +1,13 @@
use std::{collections::HashMap, time::Duration};
use std::collections::HashMap;
use futures_util::future::join_all;
use redis::{
Cmd, ErrorKind, RedisError, RedisResult, ServerErrorKind, ToRedisArgs, ToSingleRedisArg, Value,
cluster_routing::Slot,
};
use redis::{ErrorKind, RedisError, RedisResult, ServerErrorKind, Value, cluster_routing::Slot};
use crate::{
Amount, Error, Namespace, Script, commands, connection::Connection, script::ScriptCall,
Command, Error, Namespace,
commands::{Op, Plan},
connection::Connection,
script::ScriptCall,
};
pub struct Pipeline {
@ -15,29 +15,6 @@ pub struct Pipeline {
ops: Vec<Op>,
}
struct Op {
entry: Entry,
key: String,
idempotent: bool,
}
/// What a pipeline sends for one op. Only a [`ScriptCall`] can be sent again, so a retry can never
/// repeat a command that already ran.
enum Entry {
/// The op's reply is the first command's, like `RedisBatch`'s `INCRBYFLOAT` plus `EXPIRE`.
Commands(Vec<Cmd>),
Script(ScriptCall),
}
impl Entry {
fn commands(&self) -> Vec<Cmd> {
match self {
Self::Commands(commands) => commands.clone(),
Self::Script(call) => vec![call.script.evalsha(&call.keys, &call.args)],
}
}
}
type Settled = (usize, Result<Value, Error>);
impl Pipeline {
@ -48,67 +25,14 @@ impl Pipeline {
}
}
pub fn get(self, key: &str) -> Self {
let key = self.namespace.apply(key);
self.commands(vec![commands::get(key.clone())], key, true)
pub(crate) fn of(namespace: Namespace, ops: Vec<Op>) -> Self {
Self { namespace, ops }
}
/// One `MGET`, so every key must share a hash slot on a cluster. Only `RedisBatch` sends it,
/// and a batch never pipelines on a cluster.
pub(crate) fn mget(self, keys: &[&str]) -> Self {
let keys: Vec<String> = keys.iter().map(|key| self.namespace.apply(key)).collect();
let key = keys.first().cloned().unwrap_or_default();
self.commands(vec![commands::mget(keys)], key, true)
}
pub fn set(self, key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Self {
let key = self.namespace.apply(key);
self.commands(vec![commands::set(key.clone(), value, ttl)], key, true)
}
pub fn del(self, key: &str) -> Self {
let key = self.namespace.apply(key);
self.commands(vec![commands::del(vec![key.clone()])], key, true)
}
pub fn expire(self, key: &str, ttl: Duration) -> Self {
let key = self.namespace.apply(key);
self.commands(vec![commands::expire(key.clone(), ttl)], key, true)
}
/// `INCRBYFLOAT`, then `EXPIRE` when a TTL is given, like `RedisBatch` and
/// `async_increment_pipeline` send it.
pub fn incr_by_float(self, key: &str, amount: Amount, ttl: Option<Duration>) -> Self {
let key = self.namespace.apply(key);
let increment = commands::incr_by_float(key.clone(), amount);
let commands = match ttl {
Some(ttl) => vec![increment, commands::expire(key.clone(), ttl)],
None => vec![increment],
};
self.commands(commands, key, false)
}
pub fn eval(self, script: &Script, keys: &[&str], args: impl ToRedisArgs) -> Self {
let keys: Vec<String> = keys.iter().map(|key| self.namespace.apply(key)).collect();
let key = keys.first().cloned().unwrap_or_default();
let call = ScriptCall {
script: script.clone(),
keys,
args: args.to_redis_args(),
};
self.push(Entry::Script(call), key, false)
}
fn commands(self, commands: Vec<Cmd>, key: String, idempotent: bool) -> Self {
self.push(Entry::Commands(commands), key, idempotent)
}
fn push(mut self, entry: Entry, key: String, idempotent: bool) -> Self {
self.ops.push(Op {
entry,
key,
idempotent,
});
/// Appends one command in its pipelined form; `RedisClient::run` answers it in order, undecoded.
pub fn add<T>(mut self, command: Command<T>) -> Self {
let (op, _) = command.prepare(&self.namespace);
self.ops.push(op);
self
}
@ -175,8 +99,8 @@ impl Pipeline {
index: usize,
reply: RedisResult<Value>,
) -> (usize, RedisResult<Value>) {
match (&self.ops[index].entry, reply) {
(Entry::Script(call), Err(error)) if is_missing_script(&error) => {
match (&self.ops[index].plan, reply) {
(Plan::Script(call), Err(error)) if is_missing_script(&error) => {
(index, rerun_missing_script(connection, call).await)
}
(_, reply) => (index, reply),
@ -190,7 +114,7 @@ impl Pipeline {
let widths: Vec<usize> = indices
.iter()
.map(|&index| {
let commands = self.ops[index].entry.commands();
let commands = self.ops[index].plan.pipelined();
let width = commands.len();
for command in commands {
pipeline.add_command(command);
@ -198,8 +122,10 @@ impl Pipeline {
width
})
.collect();
let replies: RedisResult<Vec<RedisResult<Value>>> =
pipeline.query_async(&mut connection).await;
let replies: RedisResult<Vec<RedisResult<Value>>> = match pipeline.is_empty() {
true => Ok(Vec::new()),
false => pipeline.query_async(&mut connection).await,
};
let replies = match replies {
Ok(replies) => replies,
Err(error) => {
@ -213,7 +139,10 @@ impl Pipeline {
let per_op: Vec<(usize, RedisResult<Value>)> = indices
.iter()
.zip(widths)
.map(|(&index, width)| (index, first_reply(replies.by_ref().take(width))))
.map(|(&index, width)| match &self.ops[index].plan {
Plan::Reply(value) => (index, Ok(value.clone())),
_ => (index, first_reply(replies.by_ref().take(width))),
})
.collect();
join_all(
per_op

74
litellm-rust/crates/redis-client/tests/batch.rs
View file

@ -7,14 +7,15 @@ use std::{
use litellm_circuit_breaker::{State, Transition};
use litellm_redis_client::{
Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Namespace, Operation, RedisClient,
RedisConfig, RedisEvents, Script, Timeouts, Ttl,
Amount, CallEvent, Credentials, Declared, DEFAULT_TIMEOUT, Error, Expiry, Namespace, Operation, RedisClient,
RedisBatch, RedisConfig, RedisEvents, Script, Timeouts, Ttl,
};
use litellm_redis_client::commands;
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, endpoint};
const MINUTE: Option<Duration> = Some(Duration::from_secs(60));
const MINUTE: Option<Expiry> = Some(Expiry::KeepExisting(Duration::from_secs(60)));
// RedisBatch pipelines a standalone flush and runs each op on its own on a cluster
fn expected_trips(topology: Topology, alone: &[Operation]) -> Vec<Operation> {
@ -64,13 +65,13 @@ async fn declared_ops_wait_for_the_flush_and_share_one_round_trip(
let server = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = observed(&server);
let batch = client.batch();
let counter = batch.incr_by_float("counter", Amount::Int(3), MINUTE);
let total = batch.incr_by_float("total", Amount::Float(1.5), MINUTE);
let written = batch.set("other", "v", Duration::from_secs(60));
let read = batch.get::<String>("value");
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(3), MINUTE));
let total = batch.add(commands::incr_by_float("total", Amount::Float(1.5), MINUTE));
let written = batch.add(commands::set("other", "v", Duration::from_secs(60)));
let read = batch.add(commands::get::<String>("value"));
let outside = connect(&server, "ns", DEFAULT_TIMEOUT);
assert_eq!(outside.get::<i64>("counter").await.unwrap(), None);
assert_eq!(outside.send(commands::get::<i64>("counter")).await.unwrap(), None);
assert!(trips.seen().is_empty());
batch.flush().await;
@ -104,12 +105,12 @@ async fn resolving_before_a_flush_sends_everything_declared_so_far(
let server = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = observed(&server);
let batch = client.batch();
let first = batch.incr_by_float("a", Amount::Int(1), MINUTE);
let second = batch.incr_by_float("b", Amount::Int(2), MINUTE);
let first = batch.add(commands::incr_by_float("a", Amount::Int(1), MINUTE));
let second = batch.add(commands::incr_by_float("b", Amount::Int(2), MINUTE));
let increment = Operation::Command("INCRBYFLOAT");
assert_eq!(second.resolve().await.unwrap(), 2.0);
let late = batch.incr_by_float("a", Amount::Int(10), MINUTE);
let late = batch.add(commands::incr_by_float("a", Amount::Int(10), MINUTE));
assert_eq!(first.resolve().await.unwrap(), 1.0);
assert_eq!(
trips.seen(),
@ -126,6 +127,22 @@ async fn resolving_before_a_flush_sends_everything_declared_so_far(
assert_eq!(trips.seen().len(), both.len());
}
#[rstest]
#[tokio::test]
async fn a_pipelined_batch_is_one_pipeline_on_every_topology(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let server = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = observed(&server);
let batch = RedisBatch::pipelined(client);
let declared = ["a", "b", "c"].map(|key| batch.add(commands::incr_by_float(key, Amount::Int(2), MINUTE)));
let totals = Declared::all(declared).await;
assert_eq!(totals.into_iter().map(Result::unwrap).collect::<Vec<_>>(), [2.0; 3]);
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 3 }]);
}
#[rstest]
#[tokio::test]
async fn each_declared_op_gets_its_own_typed_result(
@ -134,16 +151,19 @@ async fn each_declared_op_gets_its_own_typed_result(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.set("text", "not a number", Duration::from_secs(60))
.send(commands::set("text", "not a number", Duration::from_secs(60)))
.await
.unwrap();
let batch = client.batch();
let rejected = batch.incr_by_float("text", Amount::Int(1), MINUTE);
let undecodable = batch.get::<i64>("text");
let missing = batch.get::<String>("missing");
let script =
batch.eval::<Vec<String>>(&Script::new("return {KEYS[1], ARGV[1]}"), &["k"], "arg");
let removed = batch.del("text");
let rejected = batch.add(commands::incr_by_float("text", Amount::Int(1), MINUTE));
let undecodable = batch.add(commands::get::<i64>("text"));
let missing = batch.add(commands::get::<String>("missing"));
let script = batch.add(commands::eval::<Vec<String>>(
&Script::new("return {KEYS[1], ARGV[1]}"),
&["k"],
"arg",
));
let removed = batch.add(commands::del(&["text"]));
batch.flush().await;
@ -162,7 +182,7 @@ async fn an_op_whose_flush_was_cancelled_resolves_as_transport(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", Duration::from_secs(5));
let batch = client.batch();
let counter = batch.incr_by_float("counter", Amount::Int(1), MINUTE);
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE));
server.freeze().await.unwrap();
let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await;
@ -182,7 +202,7 @@ async fn a_settled_op_waits_for_someone_else_to_flush_and_sends_nothing_itself(
let batch = client.batch();
let counter = tokio::spawn(
batch
.incr_by_float("counter", Amount::Int(4), MINUTE)
.add(commands::incr_by_float("counter", Amount::Int(4), MINUTE))
.settled(),
);
@ -205,7 +225,7 @@ async fn a_settled_op_whose_flush_was_cancelled_resolves_as_transport() {
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
let client = connect(&server, "", Duration::from_secs(5));
let batch = client.batch();
let counter = batch.incr_by_float("counter", Amount::Int(1), MINUTE);
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), MINUTE));
server.freeze().await.unwrap();
let cancelled = tokio::time::timeout(Duration::from_millis(200), batch.flush()).await;
@ -233,15 +253,15 @@ async fn mget_reads_each_key_once_in_order_in_one_command(
) {
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client.set("a", "1", Duration::from_secs(60)).await.unwrap();
client.set("c", "3", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("a", "1", Duration::from_secs(60))).await.unwrap();
client.send(commands::set("c", "3", Duration::from_secs(60))).await.unwrap();
server
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
.await
.unwrap();
let batch = client.batch();
let values = batch.mget::<String>(&["c", "missing", "a", "c"]);
let values = batch.add(commands::mget::<String>(&["c", "missing", "a", "c"]));
batch.flush().await;
assert_eq!(
@ -266,16 +286,16 @@ async fn a_batched_increment_treats_the_ttl_like_redis_batch(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.set("counter", "1", Duration::from_secs(30))
.send(commands::set("counter", "1", Duration::from_secs(30)))
.await
.unwrap();
let batch = client.batch();
let counter = batch.incr_by_float("counter", Amount::Int(1), Some(Duration::from_secs(600)));
let counter = batch.add(commands::incr_by_float("counter", Amount::Int(1), Some(Expiry::KeepExisting(Duration::from_secs(600)))));
batch.flush().await;
assert_eq!(counter.resolve().await.unwrap(), 2.0);
let Ttl::Expires(left) = client.ttl("counter").await.unwrap() else {
let Ttl::Expires(left) = client.send(commands::ttl("counter")).await.unwrap() else {
panic!("expected an expiring counter")
};
assert_eq!(left <= Duration::from_secs(30), keeps_the_ttl, "{left:?}");

167
litellm-rust/crates/redis-client/tests/client.rs
View file

@ -13,15 +13,16 @@ use litellm_circuit_breaker::{Breaker, Config, State};
use litellm_redis_client::{
Amount, CallEvent, Credentials, DEFAULT_TIMEOUT, Error, Expiry, FailureKind, Namespace,
Operation, Outcome, RedisClient, RedisConfig, RedisEvents, RedisUrl, Script, Standalone,
Target, Timeouts, Ttl,
Target, Timeouts, Ttl, Value,
};
use litellm_redis_client::commands;
use litellm_secret_value::SecretValue;
use support::{KEEP, connect, connect_guarded, connections, endpoint};
async fn persist(client: &RedisClient, key: &str) {
let set = Script::new("return redis.call('SET', KEYS[1], '0')");
client
.eval::<()>(&set, &[key], Vec::<String>::new())
.send(commands::eval::<()>(&set, &[key], Vec::<String>::new()))
.await
.unwrap();
}
@ -37,12 +38,12 @@ const MINUTE: Duration = Duration::from_secs(60);
async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry) {
match (counter, expiry) {
(Counter::IncrBy, Expiry::KeepExisting(ttl)) => {
client.incr_by(key, 1, ttl).await.unwrap();
client.send(commands::incr_by(key, 1, ttl)).await.unwrap();
}
(Counter::IncrBy, Expiry::Refresh(_)) => unreachable!("increment_cache never refreshes"),
(Counter::IncrByFloat, expiry) => {
client
.incr_by_float(key, Amount::Float(0.5), expiry)
.send(commands::incr_by_float(key, Amount::Float(0.5), Some(expiry)))
.await
.unwrap();
}
@ -52,7 +53,7 @@ async fn write(client: &RedisClient, counter: Counter, key: &str, expiry: Expiry
async fn persist_as(client: &RedisClient, key: &str) {
let source = "return redis.call('SET', KEYS[1], '1')";
client
.eval::<redis::Value>(&Script::new(source), &[key], Vec::<String>::new())
.send(commands::eval::<redis::Value>(&Script::new(source), &[key], Vec::<String>::new()))
.await
.unwrap();
}
@ -79,7 +80,7 @@ async fn a_counter_write_expires_a_new_key(
write(&client, counter, "fresh", expiry).await;
let left = remaining(client.ttl("fresh").await.unwrap());
let left = remaining(client.send(commands::ttl("fresh")).await.unwrap());
assert!(left > Duration::from_secs(50) && left <= MINUTE, "{left:?}");
}
@ -124,9 +125,9 @@ async fn keep_existing_leaves_a_running_ttl_and_refresh_resets_it(
)
.await;
assert!(remaining(client.ttl("kept").await.unwrap()) <= Duration::from_secs(30));
assert!(remaining(client.ttl("kept_int").await.unwrap()) <= Duration::from_secs(30));
assert!(remaining(client.ttl("refreshed").await.unwrap()) > Duration::from_secs(500));
assert!(remaining(client.send(commands::ttl("kept")).await.unwrap()) <= Duration::from_secs(30));
assert!(remaining(client.send(commands::ttl("kept_int")).await.unwrap()) <= Duration::from_secs(30));
assert!(remaining(client.send(commands::ttl("refreshed")).await.unwrap()) > Duration::from_secs(500));
}
#[rstest]
@ -141,7 +142,7 @@ async fn keep_existing_expires_a_key_that_had_no_ttl(
write(&client, counter, "persistent", KEEP).await;
let left = remaining(client.ttl("persistent").await.unwrap());
let left = remaining(client.send(commands::ttl("persistent")).await.unwrap());
assert!(left > Duration::from_secs(50), "{left:?}");
}
@ -157,21 +158,21 @@ async fn counters_send_the_commands_redis_cache_sends(
.await
.unwrap();
client.incr_by("int", 1, MINUTE).await.unwrap();
client.send(commands::incr_by("int", 1, MINUTE)).await.unwrap();
client
.incr_by_float("kept", Amount::Int(1), KEEP)
.send(commands::incr_by_float("kept", Amount::Int(1), Some(KEEP)))
.await
.unwrap();
client
.incr_by_float("refreshed", Amount::Float(0.5), Expiry::Refresh(MINUTE))
.send(commands::incr_by_float("refreshed", Amount::Float(0.5), Some(Expiry::Refresh(MINUTE))))
.await
.unwrap();
client
.run(
&client
.pipeline()
.incr_by_float("piped", Amount::Float(0.25), Some(MINUTE))
.incr_by_float("bare", Amount::Float(0.25), None),
.add(commands::incr_by_float("piped", Amount::Float(0.25), Some(Expiry::Refresh(MINUTE))))
.add(commands::incr_by_float("bare", Amount::Float(0.25), None)),
)
.await;
@ -204,13 +205,13 @@ async fn integer_increments_preserve_every_bit(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.set("integer", initial, Duration::from_secs(60))
.send(commands::set("integer", initial, Duration::from_secs(60)))
.await
.unwrap();
let expected = initial + delta;
let result = client.incr_by("integer", delta, MINUTE).await.unwrap();
let result = client.send(commands::incr_by("integer", delta, MINUTE)).await.unwrap();
assert_eq!(result, expected);
assert_eq!(client.get::<i64>("integer").await.unwrap(), Some(expected));
assert_eq!(client.send(commands::get::<i64>("integer")).await.unwrap(), Some(expected));
}
#[rstest]
@ -240,12 +241,12 @@ async fn outage_events_report_the_failure_kind(
}
let result = if pipeline {
client
.run(&client.pipeline().get("missing"))
.run(&client.pipeline().add(commands::get::<Value>("missing")))
.await
.remove(0)
.map(drop)
} else {
client.get::<String>("missing").await.map(drop)
client.send(commands::get::<String>("missing")).await.map(drop)
};
if timeout {
server.thaw().await.unwrap();
@ -290,19 +291,19 @@ async fn every_call_reports_its_operation_and_outcome(
})
.unwrap();
client
.set("list", "item", Duration::from_secs(60))
.send(commands::set("list", "item", Duration::from_secs(60)))
.await
.unwrap();
let _ = client.incr_by("list", 1, Duration::from_secs(60)).await;
let _ = client.get::<String>("missing").await;
let _ = client.run(&client.pipeline().get("a").get("b")).await;
let _ = client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await;
let _ = client.send(commands::get::<String>("missing")).await;
let _ = client.run(&client.pipeline().add(commands::get::<Value>("a")).add(commands::get::<Value>("b"))).await;
let failures = client
.run(&client.pipeline().incr_by_float(
.run(&client.pipeline().add(commands::incr_by_float(
"list",
Amount::Int(1),
Some(Duration::from_secs(60)),
))
Some(Expiry::Refresh(Duration::from_secs(60))),
)))
.await;
assert!(matches!(failures[0], Err(Error::Server(_))));
@ -341,15 +342,15 @@ async fn namespaces_isolate_keys_on_the_server(
let b = connect(&server, "b", DEFAULT_TIMEOUT);
let bare = connect(&server, "", DEFAULT_TIMEOUT);
a.set("x", "1", Duration::from_secs(60)).await.unwrap();
a.send(commands::set("x", "1", Duration::from_secs(60))).await.unwrap();
assert_eq!(a.get::<String>("x").await.unwrap().as_deref(), Some("1"));
assert_eq!(a.get::<String>("a:x").await.unwrap().as_deref(), Some("1"));
assert_eq!(a.send(commands::get::<String>("x")).await.unwrap().as_deref(), Some("1"));
assert_eq!(a.send(commands::get::<String>("a:x")).await.unwrap().as_deref(), Some("1"));
assert_eq!(
bare.get::<String>("a:x").await.unwrap().as_deref(),
bare.send(commands::get::<String>("a:x")).await.unwrap().as_deref(),
Some("1")
);
assert_eq!(b.get::<String>("x").await.unwrap(), None);
assert_eq!(b.send(commands::get::<String>("x")).await.unwrap(), None);
}
#[rstest]
@ -360,16 +361,16 @@ async fn a_rejected_command_keeps_the_connection(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.set("list", "item", Duration::from_secs(60))
.send(commands::set("list", "item", Duration::from_secs(60)))
.await
.unwrap();
let before = connections(&server).await;
let rejected = client.incr_by("list", 1, Duration::from_secs(60)).await;
let rejected = client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await;
assert!(matches!(rejected, Err(Error::Server(_))), "{rejected:?}");
assert_eq!(
client.get::<String>("list").await.unwrap().as_deref(),
client.send(commands::get::<String>("list")).await.unwrap().as_deref(),
Some("item")
);
assert_eq!(connections(&server).await, before);
@ -383,14 +384,14 @@ async fn a_reply_of_another_type_is_a_decode_error_and_a_missing_key_is_none(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.set("float", "15.0", Duration::from_secs(60))
.send(commands::set("float", "15.0", Duration::from_secs(60)))
.await
.unwrap();
let decoded = client.get::<i64>("float").await;
let decoded = client.send(commands::get::<i64>("float")).await;
assert!(matches!(decoded, Err(Error::Decode(_))), "{decoded:?}");
assert_eq!(client.get::<i64>("missing").await.unwrap(), None);
assert_eq!(client.send(commands::get::<i64>("missing")).await.unwrap(), None);
}
#[rstest]
@ -401,16 +402,16 @@ async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", Duration::from_millis(500));
client
.set("first", "1", Duration::from_secs(60))
.send(commands::set("first", "1", Duration::from_secs(60)))
.await
.unwrap();
client
.set("second", "2", Duration::from_secs(60))
.send(commands::set("second", "2", Duration::from_secs(60)))
.await
.unwrap();
server.freeze().await.unwrap();
let timed_out = client.get::<String>("first").await;
let timed_out = client.send(commands::get::<String>("first")).await;
server.thaw().await.unwrap();
assert!(
@ -418,7 +419,7 @@ async fn a_timed_out_command_does_not_leak_its_reply_into_the_next_one(
"{timed_out:?}"
);
assert_eq!(
client.get::<String>("second").await.unwrap().as_deref(),
client.send(commands::get::<String>("second")).await.unwrap().as_deref(),
Some("2")
);
}
@ -435,7 +436,7 @@ async fn the_same_client_recovers_after_redis_restarts(
let recovered = eventually(async || {
client
.incr_by("after_restart", 1, Duration::from_secs(60))
.send(commands::incr_by("after_restart", 1, Duration::from_secs(60)))
.await
.ok()
})
@ -476,10 +477,10 @@ async fn empty_batches_succeed_without_asking_redis(
let client = connect(&server, "", DEFAULT_TIMEOUT);
assert_eq!(
client.mget::<String>(&[]).await.unwrap(),
client.send(commands::mget::<String>(&[])).await.unwrap(),
Vec::<Option<String>>::new()
);
assert_eq!(client.del(&[]).await.unwrap(), 0);
assert_eq!(client.send(commands::del(&[])).await.unwrap(), 0);
}
#[rstest]
@ -489,11 +490,11 @@ async fn mget_and_del_span_slots_in_input_order(
) {
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client.set("a", "1", Duration::from_secs(60)).await.unwrap();
client.set("c", "3", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("a", "1", Duration::from_secs(60))).await.unwrap();
client.send(commands::set("c", "3", Duration::from_secs(60))).await.unwrap();
let values = client.mget::<String>(&["c", "b", "a"]).await.unwrap();
let deleted = client.del(&["a", "b", "c"]).await.unwrap();
let values = client.send(commands::mget::<String>(&["c", "b", "a"])).await.unwrap();
let deleted = client.send(commands::del(&["a", "b", "c"])).await.unwrap();
assert_eq!(
values,
@ -501,7 +502,7 @@ async fn mget_and_del_span_slots_in_input_order(
);
assert_eq!(deleted, 2);
assert_eq!(
client.mget::<String>(&["a", "c"]).await.unwrap(),
client.send(commands::mget::<String>(&["a", "c"])).await.unwrap(),
vec![None, None]
);
}
@ -515,18 +516,18 @@ async fn set_nx_only_writes_a_missing_key(
let client = connect(&server, "", DEFAULT_TIMEOUT);
let first = client
.set_nx("lock", "owner-1", Duration::from_secs(60))
.send(commands::set_nx("lock", "owner-1", Duration::from_secs(60)))
.await
.unwrap();
let second = client
.set_nx("lock", "owner-2", Duration::from_secs(60))
.send(commands::set_nx("lock", "owner-2", Duration::from_secs(60)))
.await
.unwrap();
assert!(first);
assert!(!second);
assert_eq!(
client.get::<String>("lock").await.unwrap().as_deref(),
client.send(commands::get::<String>("lock")).await.unwrap().as_deref(),
Some("owner-1")
);
}
@ -540,33 +541,33 @@ async fn ttl_reports_missing_persistent_and_remaining_time(
let client = connect(&server, "", DEFAULT_TIMEOUT);
persist(&client, "persistent").await;
client
.set("expiring", "1", Duration::from_secs(60))
.send(commands::set("expiring", "1", Duration::from_secs(60)))
.await
.unwrap();
client
.incr_by_float("expired_later", Amount::Float(0.5), KEEP)
.send(commands::incr_by_float("expired_later", Amount::Float(0.5), Some(KEEP)))
.await
.unwrap();
assert!(
client
.expire("expired_later", Duration::from_secs(30))
.send(commands::expire("expired_later", Duration::from_secs(30)))
.await
.unwrap()
);
assert_eq!(client.ttl("missing").await.unwrap(), Ttl::Missing);
assert_eq!(client.ttl("persistent").await.unwrap(), Ttl::Persistent);
assert_eq!(client.send(commands::ttl("missing")).await.unwrap(), Ttl::Missing);
assert_eq!(client.send(commands::ttl("persistent")).await.unwrap(), Ttl::Persistent);
assert!(matches!(
client.ttl("expiring").await.unwrap(),
client.send(commands::ttl("expiring")).await.unwrap(),
Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60)
));
assert!(matches!(
client.ttl("expired_later").await.unwrap(),
client.send(commands::ttl("expired_later")).await.unwrap(),
Ttl::Expires(left) if left > Duration::from_secs(20) && left <= Duration::from_secs(30)
));
assert!(
!client
.expire("missing", Duration::from_secs(30))
.send(commands::expire("missing", Duration::from_secs(30)))
.await
.unwrap()
);
@ -581,13 +582,13 @@ async fn incr_by_float_accumulates_fractions(
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.incr_by_float("spend", Amount::Float(0.25), KEEP)
.send(commands::incr_by_float("spend", Amount::Float(0.25), Some(KEEP)))
.await
.unwrap();
assert_eq!(
client
.incr_by_float("spend", Amount::Float(0.5), KEEP)
.send(commands::incr_by_float("spend", Amount::Float(0.5), Some(KEEP)))
.await
.unwrap(),
0.75
@ -608,7 +609,7 @@ async fn concurrent_commands_on_one_cloned_client_all_apply(
tokio::spawn(async move {
for _ in 0..10 {
client
.incr_by("shared", 1, Duration::from_secs(60))
.send(commands::incr_by("shared", 1, Duration::from_secs(60)))
.await
.unwrap();
}
@ -619,7 +620,7 @@ async fn concurrent_commands_on_one_cloned_client_all_apply(
task.await.unwrap();
}
assert_eq!(client.get::<i64>("shared").await.unwrap(), Some(1000));
assert_eq!(client.send(commands::get::<i64>("shared")).await.unwrap(), Some(1000));
}
fn breaker(failure_threshold: u32) -> Breaker {
@ -648,10 +649,10 @@ async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it(
);
server.freeze().await.unwrap();
let first = client.get::<String>("k").await;
let second = client.get::<String>("k").await;
let first = client.send(commands::get::<String>("k")).await;
let second = client.send(commands::get::<String>("k")).await;
let started = Instant::now();
let third = client.get::<String>("k").await;
let third = client.send(commands::get::<String>("k")).await;
let rejected_in = started.elapsed();
server.thaw().await.unwrap();
@ -660,7 +661,7 @@ async fn an_open_breaker_rejects_without_waiting_on_redis_and_a_probe_closes_it(
assert!(matches!(third, Err(Error::BreakerOpen)), "{third:?}");
assert!(rejected_in < Duration::from_millis(50), "{rejected_in:?}");
let recovered = eventually(async || client.get::<String>("k").await.ok()).await;
let recovered = eventually(async || client.send(commands::get::<String>("k")).await.ok()).await;
assert_eq!(recovered, None);
assert!(
matches!(breaker.state(), State::Closed { .. }),
@ -678,14 +679,14 @@ async fn rejected_commands_never_open_the_breaker(
let breaker = breaker(2);
let client = connect_guarded(&server, "", DEFAULT_TIMEOUT, Some(breaker.clone()));
client
.set("list", "item", Duration::from_secs(60))
.send(commands::set("list", "item", Duration::from_secs(60)))
.await
.unwrap();
let rejected = [
client.incr_by("list", 1, Duration::from_secs(60)).await,
client.incr_by("list", 1, Duration::from_secs(60)).await,
client.incr_by("list", 1, Duration::from_secs(60)).await,
client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await,
client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await,
client.send(commands::incr_by("list", 1, Duration::from_secs(60))).await,
];
assert!(
@ -709,8 +710,8 @@ async fn an_open_breaker_fails_every_pipeline_op_without_sending(
let client = connect_guarded(&server, "", Duration::from_millis(300), Some(breaker(1)));
let pipeline = client
.pipeline()
.incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60)))
.incr_by_float("c", Amount::Int(1), Some(Duration::from_secs(60)));
.add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::incr_by_float("c", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))));
server.freeze().await.unwrap();
let failed = client.run(&pipeline).await;
@ -770,13 +771,13 @@ async fn a_host_endpoint_connects_to_its_database() {
let via_url = connect(url(format!("{}/3", server.url()))).unwrap();
in_db_3
.set("k", "3", Duration::from_secs(60))
.send(commands::set("k", "3", Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(in_db_0.get::<String>("k").await.unwrap(), None);
assert_eq!(in_db_0.send(commands::get::<String>("k")).await.unwrap(), None);
assert_eq!(
via_url.get::<String>("k").await.unwrap().as_deref(),
via_url.send(commands::get::<String>("k")).await.unwrap().as_deref(),
Some("3")
);
}
@ -826,13 +827,13 @@ async fn an_idle_client_pings_every_node_before_the_next_command_like_redis_py()
unreachable!("started a cluster")
};
let client = health_checked(&server, Duration::from_secs(3), DEFAULT_TIMEOUT, None, None);
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap();
let before = ping_calls(cluster).await;
client.set("k", "2", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("k", "2", Duration::from_secs(60))).await.unwrap();
let busy = ping_calls(cluster).await;
tokio::time::sleep(Duration::from_millis(3500)).await;
client.set("k", "3", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("k", "3", Duration::from_secs(60))).await.unwrap();
let idle = ping_calls(cluster).await;
assert_eq!(busy, before);
@ -856,11 +857,11 @@ async fn a_failed_health_check_is_neither_a_call_nor_a_breaker_failure() {
Some(breaker.clone()),
Some(recorded.clone()),
);
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap();
server.freeze().await.unwrap();
tokio::time::sleep(Duration::from_millis(300)).await;
let failed = client.get::<String>("k").await;
let failed = client.send(commands::get::<String>("k")).await;
server.thaw().await.unwrap();
assert!(matches!(failed, Err(Error::Transport(_))), "{failed:?}");

11
litellm-rust/crates/redis-client/tests/connection.rs
View file

@ -12,6 +12,7 @@ use litellm_redis_client::{
Credentials, CredentialsProvider, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel,
Standalone, Target, Timeouts,
};
use litellm_redis_client::commands;
use litellm_secret_value::SecretValue;
use support::{Rotating, connect_as, connections};
@ -38,7 +39,7 @@ async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect
&server,
Credentials::Provider(CredentialsProvider(provider.clone())),
);
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap();
let before = connections(&server).await;
provider.rotate("bob", "second");
@ -55,7 +56,7 @@ async fn a_pushed_credential_reauthenticates_the_live_connection_and_a_reconnect
.cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"])
.await
.unwrap();
let after_reconnect = eventually(async || client.get::<String>("k").await.ok().flatten()).await;
let after_reconnect = eventually(async || client.send(commands::get::<String>("k")).await.ok().flatten()).await;
let reconnected = connections(&server).await;
@ -157,7 +158,7 @@ async fn sentinel_writes_follow_the_master_through_a_failover() {
let sentinel = RedisSentinel::start().await.unwrap();
let client = sentinel_client(&sentinel);
client
.set("before", "1", Duration::from_secs(600))
.send(commands::set("before", "1", Duration::from_secs(600)))
.await
.unwrap();
eventually(async || {
@ -176,7 +177,7 @@ async fn sentinel_writes_follow_the_master_through_a_failover() {
let promoted = sentinel.failover().await.unwrap();
eventually(async || {
client
.set("after", "2", Duration::from_secs(600))
.send(commands::set("after", "2", Duration::from_secs(600)))
.await
.ok()
})
@ -184,7 +185,7 @@ async fn sentinel_writes_follow_the_master_through_a_failover() {
assert_eq!(promoted, sentinel.replica_port());
assert_eq!(
client.get::<String>("before").await.unwrap().as_deref(),
client.send(commands::get::<String>("before")).await.unwrap().as_deref(),
Some("1")
);
let promoted = redis::Client::open(format!("redis://127.0.0.1:{promoted}")).unwrap();

55
litellm-rust/crates/redis-client/tests/pipeline.rs
View file

@ -3,10 +3,11 @@ mod support;
use std::time::{Duration, Instant};
use litellm_redis_client::Value;
use litellm_redis_client::commands;
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Script};
use litellm_redis_client::{Amount, DEFAULT_TIMEOUT, Error, Expiry, Script};
use support::connect;
async fn begin_migration(cluster: &litellm_redis_testing::RedisCluster, key: &str) {
@ -90,7 +91,7 @@ async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() {
let client = connect(&server, "", DEFAULT_TIMEOUT);
let counter = "{migration}:counter";
client
.set(counter, 0, Duration::from_secs(60))
.send(commands::set(counter, 0, Duration::from_secs(60)))
.await
.unwrap();
let RedisDeployment::Cluster(cluster) = &server else {
@ -103,16 +104,16 @@ async fn a_retrying_neighbor_cannot_reapply_a_cluster_counter() {
.run(
&client
.pipeline()
.incr_by_float(counter, Amount::Int(1), Some(Duration::from_secs(60)))
.eval(
.add(commands::incr_by_float(counter, Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::eval::<Value>(
&both_keys,
&[counter, "{migration}:missing"],
Vec::<String>::new(),
),
)),
)
.await;
assert_eq!(client.get::<i64>(counter).await.unwrap(), Some(1));
assert_eq!(client.send(commands::get::<i64>(counter)).await.unwrap(), Some(1));
assert_eq!(
results[0].as_ref().unwrap(),
&Value::BulkString(b"1".to_vec())
@ -132,18 +133,18 @@ async fn each_op_settles_on_its_own_in_input_order(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
client
.set("text", "not a number", Duration::from_secs(60))
.send(commands::set("text", "not a number", Duration::from_secs(60)))
.await
.unwrap();
let pipeline = client
.pipeline()
.incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60)))
.incr_by_float("text", Amount::Int(1), Some(Duration::from_secs(60)))
.incr_by_float("c", Amount::Int(5), Some(Duration::from_secs(60)))
.set("b", "2", Duration::from_secs(60))
.get("b")
.incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60)))
.get("missing");
.add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::incr_by_float("text", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::incr_by_float("c", Amount::Int(5), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::set("b", "2", Duration::from_secs(60)))
.add(commands::get::<Value>("b"))
.add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::get::<Value>("missing"));
let results = client.run(&pipeline).await;
@ -184,7 +185,7 @@ async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py(
let server = RedisDeployment::start(Topology::Standalone).await.unwrap();
let client = connect(&server, "ns", DEFAULT_TIMEOUT);
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.eval::<i64>(&incr, &["warmup"], 1).await.unwrap();
client.send(commands::eval::<i64>(&incr, &["warmup"], 1)).await.unwrap();
if !cached {
server
.cli_on_every_node(&["SCRIPT", "FLUSH"])
@ -193,14 +194,14 @@ async fn a_lost_script_reruns_after_the_rest_of_its_pipeline_like_redis_py(
}
let pipeline = client
.pipeline()
.eval(&incr, &["counter"], 2)
.get("counter");
.add(commands::eval::<Value>(&incr, &["counter"], 2))
.add(commands::get::<Value>("counter"));
let results = client.run(&pipeline).await;
assert_eq!(results[0].as_ref().unwrap(), &Value::Int(2));
assert_eq!(results[1].as_ref().unwrap(), &read_in_the_same_pipeline);
assert_eq!(client.get::<i64>("counter").await.unwrap(), Some(2));
assert_eq!(client.send(commands::get::<i64>("counter")).await.unwrap(), Some(2));
}
#[rstest]
@ -212,8 +213,8 @@ async fn an_unreachable_server_fails_every_op_as_transport(
let client = connect(&server, "", Duration::from_millis(500));
let pipeline = client
.pipeline()
.incr_by_float("a", Amount::Int(1), Some(Duration::from_secs(60)))
.incr_by_float("c", Amount::Int(1), Some(Duration::from_secs(60)));
.add(commands::incr_by_float("a", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))))
.add(commands::incr_by_float("c", Amount::Int(1), Some(Expiry::Refresh(Duration::from_secs(60)))));
server.freeze().await.unwrap();
let results = client.run(&pipeline).await;
@ -271,20 +272,20 @@ async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_los
.run(
&client
.pipeline()
.incr_by_float(
.add(commands::incr_by_float(
"{a}spend",
Amount::Float(0.5),
Some(Duration::from_secs(60)),
)
.eval(&count, &["{a}"], Vec::<String>::new())
.eval(&count, &["b"], Vec::<String>::new()),
Some(Expiry::Refresh(Duration::from_secs(60))),
))
.add(commands::eval::<Value>(&count, &["{a}"], Vec::<String>::new()))
.add(commands::eval::<Value>(&count, &["b"], Vec::<String>::new())),
)
.await;
let second = client
.run(
&client
.pipeline()
.eval(&count, &["{a}"], Vec::<String>::new()),
.add(commands::eval::<Value>(&count, &["{a}"], Vec::<String>::new())),
)
.await;
@ -302,7 +303,7 @@ async fn a_pipeline_sends_scripts_by_sha_and_reruns_only_the_ones_the_server_los
Value::Int(2)
]
);
assert_eq!(client.get::<f64>("{a}spend").await.unwrap(), Some(0.5));
assert_eq!(client.send(commands::get::<f64>("{a}spend")).await.unwrap(), Some(0.5));
let stats = support::command_stats(&deployment).await;
assert!(support::calls(&stats, "evalsha") > 0, "{stats}");
assert_eq!(support::calls(&stats, "eval"), 0, "{stats}");

5
litellm-rust/crates/redis-client/tests/pubsub.rs
View file

@ -8,6 +8,7 @@ use litellm_redis_client::{
Credentials, DEFAULT_TIMEOUT, Error, Namespace, Node, RedisClient, RedisConfig, Sentinel,
Target, Timeouts,
};
use litellm_redis_client::commands;
use litellm_redis_testing::{RedisDeployment, RedisSentinel, SERVICE_NAME, Topology, eventually};
use rstest::rstest;
use support::{connect, connect_guarded};
@ -35,7 +36,7 @@ async fn delivered(
payload: &[u8],
) {
eventually(async || {
publisher.publish(channel, payload).await.ok()?;
publisher.send(commands::publish(channel, payload)).await.ok()?;
let deadline = tokio::time::Instant::now() + Duration::from_millis(200);
while let Ok(Some(message)) = tokio::time::timeout_at(deadline, received.recv()).await {
if message == payload {
@ -147,7 +148,7 @@ async fn an_open_breaker_rejects_a_publish_without_sending_it(
server.freeze().await.unwrap();
let opened = [client.ping().await, client.ping().await];
let started = Instant::now();
let rejected = client.publish(CHANNEL, b"payload").await;
let rejected = client.send(commands::publish(CHANNEL, b"payload")).await;
let rejected_in = started.elapsed();
server.thaw().await.unwrap();

23
litellm-rust/crates/redis-client/tests/script.rs
View file

@ -6,6 +6,7 @@ use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use litellm_redis_client::Script;
use litellm_redis_client::commands;
use litellm_redis_client::{DEFAULT_TIMEOUT, Error, Ttl};
use support::connect;
@ -18,10 +19,10 @@ async fn script_keys_are_namespaced(
let client = connect(&server, "a", DEFAULT_TIMEOUT);
let bare = connect(&server, "", DEFAULT_TIMEOUT);
let read = Script::new("return redis.call('GET', KEYS[1])");
bare.set("a:x", "1", Duration::from_secs(60)).await.unwrap();
bare.send(commands::set("a:x", "1", Duration::from_secs(60))).await.unwrap();
let value: Option<String> = client
.eval(&read, &["x"], Vec::<String>::new())
.send(commands::eval(&read, &["x"], Vec::<String>::new()))
.await
.unwrap();
@ -37,14 +38,14 @@ async fn a_failing_script_is_a_server_error_and_keeps_the_connection(
let client = connect(&server, "", DEFAULT_TIMEOUT);
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client
.set("list", "item", Duration::from_secs(60))
.send(commands::set("list", "item", Duration::from_secs(60)))
.await
.unwrap();
let failed = client.eval::<i64>(&incr, &["list"], 1).await;
let failed = client.send(commands::eval::<i64>(&incr, &["list"], 1)).await;
assert!(matches!(failed, Err(Error::Server(_))), "{failed:?}");
assert_eq!(client.eval::<i64>(&incr, &["counter"], 3).await.unwrap(), 3);
assert_eq!(client.send(commands::eval::<i64>(&incr, &["counter"], 3)).await.unwrap(), 3);
}
#[rstest]
@ -55,14 +56,14 @@ async fn a_flushed_script_is_loaded_again(
let server = RedisDeployment::start(topology).await.unwrap();
let client = connect(&server, "", DEFAULT_TIMEOUT);
let incr = Script::new("return redis.call('INCRBY', KEYS[1], ARGV[1])");
client.eval::<i64>(&incr, &["counter"], 1).await.unwrap();
client.send(commands::eval::<i64>(&incr, &["counter"], 1)).await.unwrap();
server
.cli_on_every_node(&["SCRIPT", "FLUSH"])
.await
.unwrap();
assert_eq!(client.eval::<i64>(&incr, &["counter"], 1).await.unwrap(), 2);
assert_eq!(client.send(commands::eval::<i64>(&incr, &["counter"], 1)).await.unwrap(), 2);
}
const INCREMENT_WITH_FLOOR: &str = "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) \
@ -86,21 +87,21 @@ async fn pythons_increment_with_floor_script_floors_at_zero_and_expires_in_one_c
let count = match pipelined {
false => client
.eval::<i64>(&floor, &["slots"], (-5, 60))
.send(commands::eval::<i64>(&floor, &["slots"], (-5, 60)))
.await
.unwrap(),
true => {
let results = client
.run(&client.pipeline().eval(&floor, &["slots"], (-5, 60)))
.run(&client.pipeline().add(commands::eval::<redis::Value>(&floor, &["slots"], (-5, 60))))
.await;
redis::from_redis_value(results[0].as_ref().unwrap().clone()).unwrap()
}
};
assert_eq!(count, 0);
assert_eq!(client.get::<i64>("slots").await.unwrap(), Some(0));
assert_eq!(client.send(commands::get::<i64>("slots")).await.unwrap(), Some(0));
assert!(matches!(
client.ttl("slots").await.unwrap(),
client.send(commands::ttl("slots")).await.unwrap(),
Ttl::Expires(left) if left > Duration::from_secs(50) && left <= Duration::from_secs(60)
));
}

13
litellm-rust/crates/redis-connect/src/refresh.rs
View file

@ -236,6 +236,7 @@ fn refresh_after(expires_at: Option<SystemTime>, now: SystemTime) -> Duration {
#[cfg(test)]
mod tests {
use litellm_redis_client::commands;
use std::{
collections::BTreeSet,
sync::{
@ -636,7 +637,7 @@ mod tests {
health_check: None,
})
.unwrap();
client.set("k", "1", Duration::from_secs(60)).await.unwrap();
client.send(commands::set("k", "1", Duration::from_secs(60))).await.unwrap();
let before = server.clients().await.unwrap();
let after = eventually(async || {
@ -654,14 +655,14 @@ mod tests {
assert_eq!(connection_users(&before), BTreeSet::from(["alice"]));
assert_eq!(connection_ids(&after), connection_ids(&before));
assert_eq!(
client.get::<String>("k").await.unwrap().as_deref(),
client.send(commands::get::<String>("k")).await.unwrap().as_deref(),
Some("1")
);
server
.cli_on_every_node(&["CLIENT", "KILL", "USER", "bob"])
.await
.unwrap();
let value = eventually(async || client.get::<String>("k").await.ok().flatten()).await;
let value = eventually(async || client.send(commands::get::<String>("k")).await.ok().flatten()).await;
let reconnected = server.clients().await.unwrap();
assert_eq!(value, "1");
@ -715,7 +716,7 @@ mod tests {
})
.unwrap();
client
.set("before", "1", Duration::from_secs(60))
.send(commands::set("before", "1", Duration::from_secs(60)))
.await
.unwrap();
let master = sentinel.master_port().await.unwrap();
@ -737,7 +738,7 @@ mod tests {
.unwrap();
let promoted = sentinel.failover().await.unwrap();
eventually(async || client.set("after", "2", Duration::from_secs(60)).await.ok()).await;
eventually(async || client.send(commands::set("after", "2", Duration::from_secs(60))).await.ok()).await;
let on_promoted = sentinel.clients_on(promoted).await.unwrap();
assert_eq!(promoted, sentinel.replica_port());
@ -747,7 +748,7 @@ mod tests {
"{on_promoted:?}"
);
assert_eq!(
client.get::<String>("after").await.unwrap().as_deref(),
client.send(commands::get::<String>("after")).await.unwrap().as_deref(),
Some("2")
);
}

7
litellm-rust/crates/redis-connect/tests/connect.rs
View file

@ -1,3 +1,4 @@
use litellm_redis_client::commands;
use std::{collections::BTreeMap, sync::Arc, time::Duration};
use litellm_redis_connect::{
@ -122,14 +123,14 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() {
});
let client = connect(resolved(kwargs), services(), None).unwrap();
client
.set("before", "1", Duration::from_secs(600))
.send(commands::set("before", "1", Duration::from_secs(600)))
.await
.unwrap();
let promoted = sentinel.failover().await.unwrap();
eventually(async || {
client
.set("after", "2", Duration::from_secs(600))
.send(commands::set("after", "2", Duration::from_secs(600)))
.await
.ok()
})
@ -138,7 +139,7 @@ async fn python_sentinel_kwargs_follow_the_master_through_a_failover() {
assert_eq!(promoted, sentinel.replica_port());
assert_eq!(sentinel.master_port().await.unwrap(), promoted);
assert_eq!(
client.get::<String>("after").await.unwrap().as_deref(),
client.send(commands::get::<String>("after")).await.unwrap().as_deref(),
Some("2")
);
}

10
litellm-rust/crates/redis-store/AGENTS.md
View file

@ -1,9 +1,13 @@
- 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
- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (Python-facing values with their `FromValue` / `IntoValue`). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` with its attach-once slot, the `DualCache` read-through and chunked deletes (`remote.rs`) and the Python codecs (`codec.rs`, `value.rs`)
- An op declared on a `Send` returns `Pending<T>`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on
- A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s
- 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`
- 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 `Remote` (or none), never a cache, and a family owns one `RemoteSlot` that its memory partitions and `MemoryTier` reads share. Redis-only families are plain functions over `&RedisClient`
- Every family works with no Redis, answering from memory. `attach` adds the Redis tier at most once and keeps memory, as `DualCache.attach_redis_cache` does, so a host builds the store at startup and never rebuilds it to connect
- 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
- Key suffixes and wire spellings derive from strum on the family's enums, and a type whose Python side is plain data implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`, so the bridge never names a field or a spelling
- 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. A direct script call goes out alone instead, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. 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
- `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. Each op exists once and takes a `Send` (`Now`, `PreCall`, `PostCall`) picked by the caller, and the store calls `send.batch(&client)`. A store never encodes a phase in a method name. A post-call method exists only where a follow-up runs after the write settles (`token_increment_post_call`), as a thin wrapper over the shared op. `Send::Now` is an alone batch, so a script call outside a request goes out alone, like a redis-py `Script` outside a pipeline, so `NOSCRIPT` costs `SCRIPT LOAD` and one more `EVALSHA` the way it does on main. 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

474
litellm-rust/crates/redis-store/src/auth_cache.rs
View file

@ -1,474 +0,0 @@
use std::time::Duration;
use litellm_cache::Error as MemoryError;
use litellm_cache_memory::InMemoryCache;
use litellm_python_compat::{Value, json, repr::to_str};
use litellm_redis_client::{Declared, RedisBatch, RedisClient};
use crate::tier::{Clock, system_clock};
use crate::{DecodeError, ErrorKind, Number, RequestRedis, codec::python_value};
/// `InMemoryCache`'s TTL for a write that names none and has no `default_in_memory_ttl`.
pub const MEMORY_TTL: Duration = Duration::from_secs(600);
/// `BaseCache.default_ttl`, which `RedisCache` keeps when `litellm.default_redis_ttl` is unset.
pub const REDIS_TTL: Duration = Duration::from_secs(60);
pub const MEMORY_ENTRIES: usize = 200;
/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it.
const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])";
/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `DualCache.async_delete_cache_keys` deletes at once.
const DELETE_CHUNK: usize = 1000;
/// `UserApiKeyCache`'s configuration: each memory partition's size, `default_in_memory_ttl`, and the
/// `RedisCache.default_ttl` of the Redis it was attached to.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Settings {
pub memory_entries: usize,
pub memory_ttl: Option<Duration>,
pub redis_ttl: Duration,
}
impl Default for Settings {
fn default() -> Self {
Self {
memory_entries: MEMORY_ENTRIES,
memory_ttl: None,
redis_ttl: REDIS_TTL,
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error(transparent)]
Redis(#[from] litellm_redis_client::Error),
#[error(transparent)]
Memory(#[from] MemoryError),
#[error(transparent)]
Decode(#[from] DecodeError),
#[error("the value is not JSON serializable: {0}")]
Encode(#[from] litellm_python_compat::Error),
}
impl Error {
pub fn kind(&self) -> ErrorKind {
match self {
Self::Redis(error) => error.into(),
Self::Decode(_) => ErrorKind::Decode,
Self::Encode(_) => ErrorKind::Encode,
Self::Memory(_) => ErrorKind::Memory,
}
}
}
/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every
/// other object under a prefixed key.
pub fn is_key_object(key: &str) -> bool {
key.len() == 64
&& key
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}
/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds.
fn whole_seconds(ttl: Duration) -> Duration {
Duration::from_secs(ttl.as_secs())
}
fn decode(raw: Option<Vec<u8>>) -> Result<Option<Value>, DecodeError> {
raw.as_deref().map(python_value).transpose()
}
/// The proxy's auth object cache with `enable_redis_auth_cache` on: two memory partitions, so churn
/// in management objects cannot evict user-key objects, in front of one Redis shared by every pod.
/// Keys are the logical cache keys `UserApiKeyCache` callers already use; values are what
/// `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`.
pub struct AuthCache {
key_objects: InMemoryCache<Value>,
objects: InMemoryCache<Value>,
remote: Option<RedisClient>,
settings: Settings,
}
impl AuthCache {
pub fn new(client: Option<RedisClient>, settings: Settings) -> Self {
Self::with_clock(client, settings, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, settings: Settings, now: Clock) -> Self {
let memory = || {
let now = now.clone();
InMemoryCache::with_clock(Some(settings.memory_entries), Some(MEMORY_TTL), move || {
now()
})
};
Self {
key_objects: memory(),
objects: memory(),
remote: client,
settings,
}
}
pub fn has_redis(&self) -> bool {
self.remote.is_some()
}
pub fn settings(&self) -> Settings {
self.settings
}
fn memory(&self, key: &str) -> &InMemoryCache<Value> {
if is_key_object(key) {
&self.key_objects
} else {
&self.objects
}
}
/// `DualCache`'s rule: a write without a TTL takes `default_in_memory_ttl` in memory and Redis.
fn memory_ttl(&self, ttl: Option<Duration>) -> Option<Duration> {
ttl.or(self.settings.memory_ttl)
}
fn redis_ttl(&self, ttl: Option<Duration>) -> Duration {
whole_seconds(self.memory_ttl(ttl).unwrap_or(self.settings.redis_ttl))
}
pub fn get_local(&self, key: &str) -> Result<Option<Value>, MemoryError> {
self.memory(key).get_cache(key)
}
/// A memory write with `ttl`, or `default_in_memory_ttl` when it names none.
pub fn set_local(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
self.memory(key)
.set_cache(key, value, self.memory_ttl(ttl))
.map(drop)
}
pub fn delete_local(&self, key: &str) -> Result<(), MemoryError> {
self.memory(key).delete_cache(key)
}
pub fn flush_local(&self) -> Result<(), MemoryError> {
self.key_objects.flush_cache()?;
self.objects.flush_cache()
}
/// Entries held by the management-object and the user-key partitions.
pub fn len_local(&self) -> Result<(usize, usize), MemoryError> {
Ok((self.objects.len()?, self.key_objects.len()?))
}
/// `DualCache.async_get_cache`: memory, then Redis, backfilling memory. A failed Redis read is a
/// miss, as `DualCache` logs and swallows it.
pub async fn get(&self, key: &str) -> Result<Option<Value>, MemoryError> {
if let Some(value) = self.get_local(key)? {
return Ok(Some(value));
}
let Some(remote) = &self.remote else {
return Ok(None);
};
let read = remote
.get::<Vec<u8>>(key)
.await
.map_err(Error::from)
.and_then(|raw| Ok(decode(raw)?));
match read {
Ok(Some(value)) => {
self.set_local(key, value.clone(), None)?;
Ok(Some(value))
}
Ok(None) => Ok(None),
Err(error) => {
tracing::warn!(%error, "auth cache Redis read failed, treating it as a miss");
Ok(None)
}
}
}
/// `DualCache.async_set_cache`: memory, then `SET` of `json.dumps(value)`. Redis failures and
/// values `json.dumps` refuses only log, as there.
pub async fn set(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
let encoded = json::dumps(&value);
self.set_local(key, value, ttl)?;
let Some(remote) = &self.remote else {
return Ok(());
};
let written = match encoded {
Ok(bytes) => remote
.set(key, bytes, self.redis_ttl(ttl))
.await
.map_err(Error::from),
Err(error) => Err(error.into()),
};
if let Err(error) = written {
tracing::warn!(%error, "writing an auth cache entry to Redis failed");
}
Ok(())
}
/// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`.
pub async fn set_str(
&self,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
let encoded = to_str(&value);
self.set_local(key, value, ttl)?;
let Some(remote) = &self.remote else {
return Ok(());
};
if let Err(error) = remote.set(key, encoded, self.redis_ttl(ttl)).await {
tracing::warn!(%error, "writing an auth cache entry to Redis failed");
}
Ok(())
}
/// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the request's pre-call batch,
/// sent with whatever read goes out next. `None` without Redis.
pub fn set_pre_call(
&self,
request: &RequestRedis,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<Option<DeclaredWrite>, MemoryError> {
let encoded = json::dumps(&value);
let Some(remote) = &self.remote else {
return Ok(None);
};
self.set_local(key, value, ttl)?;
Ok(Some(match encoded {
Ok(bytes) => {
DeclaredWrite::Set(request.pre(remote).set(key, bytes, self.redis_ttl(ttl)))
}
Err(error) => DeclaredWrite::Unencodable(error),
}))
}
/// `DualCache.async_delete_cache_pre_call`. `None` without Redis.
pub fn delete_pre_call(
&self,
request: &RequestRedis,
key: &str,
) -> Result<Option<DeclaredWrite>, MemoryError> {
let Some(remote) = &self.remote else {
return Ok(None);
};
self.delete_local(key)?;
Ok(Some(DeclaredWrite::Delete(request.pre(remote).del(key))))
}
/// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then
/// the rest in another, each with one TTL.
pub async fn set_many(
&self,
entries: Vec<(String, Value)>,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
let (key_objects, others): (Vec<_>, Vec<_>) =
entries.into_iter().partition(|(key, _)| is_key_object(key));
for partition in [key_objects, others] {
if !partition.is_empty() {
self.set_partition(partition, ttl).await?;
}
}
Ok(())
}
async fn set_partition(
&self,
entries: Vec<(String, Value)>,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
let mut pipeline = self.remote.as_ref().map(RedisClient::pipeline);
for (key, value) in entries {
if let Some(writes) = pipeline.take() {
pipeline = Some(match json::dumps(&value) {
Ok(bytes) => writes.set(&key, bytes, self.redis_ttl(ttl)),
Err(error) => {
tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only");
writes
}
});
}
self.set_local(&key, value, ttl)?;
}
if let (Some(remote), Some(pipeline)) = (&self.remote, pipeline) {
for reply in remote.run(&pipeline).await {
if let Err(error) = reply {
tracing::warn!(%error, "writing an auth cache entry to Redis failed");
}
}
}
Ok(())
}
/// The auth prefetch write-back, `async_set_cache_pipeline_with_ttls`: an entry without a TTL
/// takes `default_in_memory_ttl` in memory but the connection's default in Redis. On the
/// request's pre-call batch when there is one, else in a round trip of its own.
pub async fn write_back(
&self,
request: Option<&RequestRedis>,
entries: Vec<(String, Value, Option<Duration>)>,
) -> Result<(), Error> {
let Some(remote) = &self.remote else {
for (key, value, ttl) in entries {
self.set_local(&key, value, ttl)?;
}
return Ok(());
};
if let Some(request) = request {
let batch = request.pre(remote);
for (key, value, ttl) in entries {
let redis_ttl = whole_seconds(ttl.unwrap_or(self.settings.redis_ttl));
drop(batch.set(&key, json::dumps(&value)?, redis_ttl));
self.set_local(&key, value, ttl)?;
}
return Ok(());
}
let mut pipeline = remote.pipeline();
for (key, value, ttl) in entries {
let redis_ttl = whole_seconds(ttl.unwrap_or(self.settings.redis_ttl));
pipeline = pipeline.set(&key, json::dumps(&value)?, redis_ttl);
self.set_local(&key, value, ttl)?;
}
for reply in remote.run(&pipeline).await {
reply?;
}
Ok(())
}
/// `UserApiKeyCache.async_delete_cache_keys`: memory, then the key-object partition and the
/// rest as two concurrent `DualCache.async_delete_cache_keys`, each a multi-key `DEL` per
/// chunk of 1000. Both run even when one fails.
pub async fn delete(&self, keys: &[String]) -> Result<(), Error> {
for key in keys {
self.delete_local(key)?;
}
let Some(remote) = &self.remote else {
return Ok(());
};
let (key_objects, others): (Vec<&str>, Vec<&str>) = keys
.iter()
.map(String::as_str)
.partition(|key| is_key_object(key));
let (key_objects, others) = tokio::join!(
delete_chunks(remote, &key_objects),
delete_chunks(remote, &others)
);
key_objects.and(others)
}
/// Reads for keys the caller found missing in memory, on the request's pre-call batch. `None`
/// without Redis.
pub fn declare_get_many(
&self,
request: &RequestRedis,
keys: &[String],
) -> Option<DeclaredGets> {
Some(self.declare_gets(&request.pre(self.remote.as_ref()?), keys))
}
/// `RedisCache.async_batch_get_cache` for keys missing in memory, in one round trip. `None`
/// without Redis.
pub async fn get_many(&self, keys: &[String]) -> Option<Result<Vec<Option<Value>>, Error>> {
Some(
self.declare_gets(&self.remote.as_ref()?.batch(), keys)
.resolve()
.await,
)
}
fn declare_gets(&self, batch: &RedisBatch, keys: &[String]) -> DeclaredGets {
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
DeclaredGets(batch.mget::<Vec<u8>>(&keys))
}
/// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds.
/// `None` without Redis or for an absent entry.
pub async fn increment_if_present(
&self,
key: &str,
amount: Number,
) -> Result<Option<f64>, Error> {
let Some(remote) = &self.remote else {
return Ok(None);
};
let total: Option<Vec<u8>> = remote
.eval_source(
INCREMENT_IF_PRESENT,
&[key],
litellm_redis_client::Amount::from(amount),
)
.await?;
total
.as_deref()
.map(|raw| {
std::str::from_utf8(raw)
.ok()
.and_then(|text| text.trim().parse::<f64>().ok())
.ok_or(Error::Decode(DecodeError::Shape(
"INCRBYFLOAT returned a non-number",
)))
})
.transpose()
}
}
/// A write declared on a request's pre-call batch, settled once the batch goes out.
pub enum DeclaredWrite {
Set(Declared<()>),
Delete(Declared<usize>),
/// Kept in memory only, as `json.dumps` refused the value.
Unencodable(litellm_python_compat::Error),
}
impl DeclaredWrite {
/// Flushes the batch if nobody has yet.
pub async fn resolve(self) -> Result<(), Error> {
match self {
Self::Set(written) => written.resolve().await?,
Self::Delete(deleted) => drop(deleted.resolve().await?),
Self::Unencodable(error) => {
tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only");
}
}
Ok(())
}
}
/// Reads declared on a batch, answered once it goes out.
pub struct DeclaredGets(Declared<Vec<Option<Vec<u8>>>>);
impl DeclaredGets {
/// Flushes the batch if nobody has yet. Any failed read or undecodable value fails the whole
/// read, as a failed `MGET` does.
pub async fn resolve(self) -> Result<Vec<Option<Value>>, Error> {
self.0
.resolve()
.await?
.into_iter()
.map(|reply| Ok(decode(reply)?))
.collect()
}
}
async fn delete_chunks(remote: &RedisClient, keys: &[&str]) -> Result<(), Error> {
for chunk in keys.chunks(DELETE_CHUNK) {
remote.del(chunk).await?;
}
Ok(())
}

8
litellm-rust/crates/redis-store/src/auth_cache/keys.rs Normal file
View file

@ -0,0 +1,8 @@
/// `is_user_key_cache_key`: user-key objects are cached under a bare `hash_token` digest, every
/// other object under a prefixed key.
pub fn is_key_object(key: &str) -> bool {
key.len() == 64
&& key
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}

292
litellm-rust/crates/redis-store/src/auth_cache/mod.rs Normal file
View file

@ -0,0 +1,292 @@
mod keys;
mod types;
use std::time::Duration;
use litellm_python_compat::{Value, json, repr::to_str};
use litellm_redis_client::{Declared, RedisBatch, commands};
pub use self::keys::is_key_object;
pub use self::types::Settings;
use crate::{
Error, Number, Pending, Send,
codec::python_value,
memory::{Clock, Memory, system_clock},
remote::{self, Remote, RemoteSlot},
request::{pending, ready},
};
/// `RedisCache.async_increment_if_present`, byte for byte, sent as `EVAL` like it.
const INCREMENT_IF_PRESENT: &str = "if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])";
/// The proxy's auth object cache with `enable_redis_auth_cache` on: two memory partitions, so churn
/// in management objects cannot evict user-key objects, in front of one Redis shared by every pod.
/// Keys are the logical cache keys `UserApiKeyCache` callers already use; values are what
/// `json.loads` returns, written as `json.dumps`, like `RedisCache.async_set_cache`.
pub struct AuthCache {
key_objects: Memory<Value>,
objects: Memory<Value>,
remote: RemoteSlot,
settings: Settings,
}
impl AuthCache {
pub fn new(remote: Option<Remote>, settings: Settings) -> Self {
Self::with_clock(remote, settings, system_clock())
}
pub fn with_clock(remote: Option<Remote>, settings: Settings, now: Clock) -> Self {
Self {
key_objects: Memory::new(Some(settings.memory_entries), &now),
objects: Memory::new(Some(settings.memory_entries), &now),
remote: RemoteSlot::new(remote),
settings,
}
}
pub fn has_redis(&self) -> bool {
self.remote.get().is_some()
}
pub fn settings(&self) -> Settings {
self.settings
}
pub fn get_local(&self, key: &str) -> Result<Option<Value>, Error> {
Ok(self.memory(key).get(key)?)
}
/// A memory write with `ttl`, or `default_in_memory_ttl` when it names none.
pub fn set_local(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
Ok(self.memory(key).set(key, value, self.memory_ttl(ttl))?)
}
pub fn delete_local(&self, key: &str) -> Result<(), Error> {
Ok(self.memory(key).delete(key)?)
}
pub fn flush_local(&self) -> Result<(), Error> {
self.key_objects.flush()?;
Ok(self.objects.flush()?)
}
/// Entries held by the management-object and the user-key partitions.
pub fn len_local(&self) -> Result<(usize, usize), Error> {
Ok((self.objects.len()?, self.key_objects.len()?))
}
/// `DualCache.async_get_cache`, backfilling memory for `default_in_memory_ttl`.
pub async fn get(&self, key: &str) -> Result<Option<Value>, Error> {
self.remote
.read_through(self.memory(key), key, python_value, self.settings.memory_ttl)
.await
}
/// `DualCache.async_set_cache`: memory, then `SET` of `json.dumps(value)`. Redis failures and
/// values `json.dumps` refuses only log, as there.
pub async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
let encoded = json::dumps(&value);
self.set_local(key, value, ttl)?;
let Some(remote) = self.remote.get() else {
return Ok(());
};
let written = match encoded {
Ok(bytes) => remote
.client
.send(commands::set(key, bytes, self.redis_ttl(remote, ttl)))
.await
.map_err(Error::from),
Err(error) => Err(error.into()),
};
if let Err(error) = written {
tracing::warn!(%error, "writing an auth cache entry to Redis failed");
}
Ok(())
}
/// `DualCache.set_cache`, whose `RedisCache.set_cache` writes `str(value)`.
pub async fn set_str(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
let encoded = to_str(&value);
self.set_local(key, value, ttl)?;
let Some(remote) = self.remote.get() else {
return Ok(());
};
let write = commands::set(key, encoded, self.redis_ttl(remote, ttl));
if let Err(error) = remote.client.send(write).await {
tracing::warn!(%error, "writing an auth cache entry to Redis failed");
}
Ok(())
}
/// `DualCache.async_set_cache_pre_call`: memory now, the `SET` on the batch `send` picks.
pub fn declare_set(
&self,
send: Send<'_>,
key: &str,
value: Value,
ttl: Option<Duration>,
) -> Result<Pending<()>, Error> {
let encoded = json::dumps(&value);
self.set_local(key, value, ttl)?;
let Some(remote) = self.remote.get() else {
return Ok(ready(()));
};
match encoded {
Ok(bytes) => Ok(pending(
send.batch(&remote.client)
.add(commands::set(key, bytes, self.redis_ttl(remote, ttl))),
)),
Err(error) => {
tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only");
Ok(ready(()))
}
}
}
/// `DualCache.async_delete_cache_pre_call`, memory now and the `DEL` on the batch `send` picks.
pub fn declare_delete(&self, send: Send<'_>, key: &str) -> Result<Pending<()>, Error> {
self.delete_local(key)?;
let Some(batch) = self.remote.batch(send) else {
return Ok(ready(()));
};
let deleted = batch.add(commands::del(&[key]));
Ok(Box::pin(async move { Ok(deleted.resolve().await.map(drop)?) }))
}
/// `UserApiKeyCache.async_set_cache_pipeline`: the key-object entries in one pipeline, then
/// the rest in another, each with one TTL.
pub async fn set_many(&self, entries: Vec<(String, Value)>, ttl: Option<Duration>) -> Result<(), Error> {
let (key_objects, others): (Vec<_>, Vec<_>) =
entries.into_iter().partition(|(key, _)| is_key_object(key));
for partition in [key_objects, others] {
self.set_partition(partition, ttl).await?;
}
Ok(())
}
/// The auth prefetch write-back, `async_set_cache_pipeline_with_ttls`: memory now, and an entry
/// without a TTL takes `default_in_memory_ttl` in memory but the connection's default in Redis.
/// On a request batch the `SET`s are declared and nothing is left to wait for. [`Send::Now`] is
/// one pipeline of its own, as `async_set_cache_pipeline_with_ttls` sends it, rather than one
/// round trip per entry.
pub fn write_back(
&self,
send: Send<'_>,
entries: Vec<(String, Value, Option<Duration>)>,
) -> Result<Pending<()>, Error> {
let Some(remote) = self.remote.get() else {
for (key, value, ttl) in entries {
self.set_local(&key, value, ttl)?;
}
return Ok(ready(()));
};
let sets = entries
.into_iter()
.map(|(key, value, ttl)| {
let redis_ttl = ttl.unwrap_or(remote.default_ttl);
let set = commands::set(&key, json::dumps(&value)?, redis_ttl);
self.set_local(&key, value, ttl)?;
Ok(set)
})
.collect::<Result<Vec<_>, Error>>()?;
let Send::Now = send else {
let batch = send.batch(&remote.client);
sets.into_iter().for_each(|set| drop(batch.add(set)));
return Ok(ready(()));
};
let pipeline = RedisBatch::pipelined(remote.client.clone());
let written: Vec<_> = sets.into_iter().map(|set| pipeline.add(set)).collect();
Ok(Box::pin(async move {
Declared::all(written).await.into_iter().try_for_each(|reply| Ok(reply?))
}))
}
/// `UserApiKeyCache.async_delete_cache_keys`: memory, then the key-object partition and the
/// rest as two concurrent `DualCache.async_delete_cache_keys`, each a multi-key `DEL` per
/// chunk of 1000. Both run even when one fails.
pub async fn delete(&self, keys: &[String]) -> Result<(), Error> {
keys.iter().try_for_each(|key| self.delete_local(key))?;
let Some(client) = self.remote.client() else {
return Ok(());
};
let (key_objects, others): (Vec<String>, Vec<String>) =
keys.iter().cloned().partition(|key| is_key_object(key));
let (key_objects, others) = tokio::join!(
remote::delete(client, &key_objects),
remote::delete(client, &others)
);
key_objects.and(others)
}
/// `RedisCache.async_batch_get_cache` for keys the caller found missing in memory, one `MGET` on
/// the batch `send` picks. `None` without Redis. Any failed read or undecodable value fails the
/// whole read, as a failed `MGET` does.
pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Option<Pending<Vec<Option<Value>>>> {
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
let read = self.remote.batch(send)?.add(commands::mget::<Vec<u8>>(&keys));
Some(Box::pin(async move {
read.resolve()
.await?
.into_iter()
.map(|raw| Ok(types::decode(raw)?))
.collect()
}))
}
/// `RedisCache.async_increment_if_present`: `INCRBYFLOAT` only an entry Redis still holds.
/// `None` without Redis or for an absent entry.
pub async fn increment_if_present(&self, key: &str, amount: Number) -> Result<Option<f64>, Error> {
let Some(client) = self.remote.client() else {
return Ok(None);
};
let amount = litellm_redis_client::Amount::from(amount);
Ok(client.send(commands::eval_source(INCREMENT_IF_PRESENT, &[key], amount)).await?)
}
fn memory(&self, key: &str) -> &Memory<Value> {
if is_key_object(key) {
&self.key_objects
} else {
&self.objects
}
}
/// `DualCache`'s rule: a write without a TTL takes `default_in_memory_ttl` in memory and Redis.
fn memory_ttl(&self, ttl: Option<Duration>) -> Option<Duration> {
ttl.or(self.settings.memory_ttl)
}
fn redis_ttl(&self, remote: &Remote, ttl: Option<Duration>) -> Duration {
self.memory_ttl(ttl).unwrap_or(remote.default_ttl)
}
async fn set_partition(&self, entries: Vec<(String, Value)>, ttl: Option<Duration>) -> Result<(), Error> {
let target = self
.remote
.get()
.map(|remote| (RedisBatch::pipelined(remote.client.clone()), self.redis_ttl(remote, ttl)));
let written = entries
.into_iter()
.map(|(key, value)| {
let set = match (&target, json::dumps(&value)) {
(Some((pipeline, redis_ttl)), Ok(bytes)) => {
Some(pipeline.add(commands::set(&key, bytes, *redis_ttl)))
}
(Some(_), Err(error)) => {
tracing::warn!(%error, "an auth cache value is not JSON serializable, kept in memory only");
None
}
(None, _) => None,
};
self.set_local(&key, value, ttl)?;
Ok(set)
})
.collect::<Result<Vec<_>, Error>>()?;
for reply in Declared::all(written.into_iter().flatten()).await {
if let Err(error) = reply {
tracing::warn!(%error, "writing an auth cache entry to Redis failed");
}
}
Ok(())
}
}

26
litellm-rust/crates/redis-store/src/auth_cache/types.rs Normal file
View file

@ -0,0 +1,26 @@
use std::time::Duration;
use litellm_python_compat::Value;
use crate::{DecodeError, codec::python_value};
/// `UserApiKeyCache`'s memory configuration: each partition's size and `default_in_memory_ttl`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Settings {
pub memory_entries: usize,
pub memory_ttl: Option<Duration>,
}
impl Default for Settings {
/// `InMemoryCache`'s 200 entries.
fn default() -> Self {
Self {
memory_entries: 200,
memory_ttl: None,
}
}
}
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Value>, DecodeError> {
raw.as_deref().map(python_value).transpose()
}

7
litellm-rust/crates/redis-store/src/codec.rs
View file

@ -1,5 +1,7 @@
use litellm_python_compat::{Value, json, literal};
use crate::Number;
#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
#[error("Redis value is not UTF-8")]
@ -18,3 +20,8 @@ pub(crate) fn python_value(raw: &[u8]) -> Result<Value, DecodeError> {
Err(_) => literal::literal_eval(text).map_err(DecodeError::Unparsed),
}
}
/// A number written by Python, `shape` naming what it was meant to be when it is not one.
pub(crate) fn number(raw: &[u8], shape: &'static str) -> Result<Number, DecodeError> {
Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape(shape))
}

25
litellm-rust/crates/redis-store/src/error.rs
View file

@ -1,3 +1,28 @@
use crate::DecodeError;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error(transparent)]
Redis(#[from] litellm_redis_client::Error),
#[error(transparent)]
Memory(#[from] litellm_cache::Error),
#[error(transparent)]
Decode(#[from] DecodeError),
#[error("the value is not JSON serializable: {0}")]
Encode(#[from] litellm_python_compat::Error),
}
impl Error {
pub fn kind(&self) -> ErrorKind {
match self {
Self::Redis(error) => error.into(),
Self::Memory(_) => ErrorKind::Memory,
Self::Decode(_) => ErrorKind::Decode,
Self::Encode(_) => ErrorKind::Encode,
}
}
}
/// The failure class a host reports with a store error, `timeout` and `breaker_open` being the ones a
/// caller may drop.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::IntoStaticStr)]

12
litellm-rust/crates/redis-store/src/lib.rs
View file

@ -1,16 +1,22 @@
pub mod auth_cache;
mod codec;
mod error;
mod memory;
mod number;
pub mod rate_limits;
mod remote;
mod request;
pub mod router_cooldowns;
pub mod router_usage;
pub mod spend_counters;
mod tier;
mod value;
pub use codec::DecodeError;
pub use error::ErrorKind;
pub use error::{Error, ErrorKind};
pub use memory::{Clock, system_clock};
pub use number::Number;
pub use request::RequestRedis;
pub use tier::{Clock, TierPolicy};
pub use remote::Remote;
pub use request::{Pending, RequestRedis, Send};
pub use tier::TierPolicy;
pub use value::{FromValue, IntoValue};

84
litellm-rust/crates/redis-store/src/memory.rs Normal file
View file

@ -0,0 +1,84 @@
use std::{
sync::{Arc, Mutex, MutexGuard, PoisonError},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::Error;
use litellm_cache_memory::InMemoryCache;
pub type Clock = Arc<dyn Fn() -> Duration + Send + Sync>;
pub fn system_clock() -> Clock {
Arc::new(|| {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
/// A store's memory tier: Python's `InMemoryCache`, with its 200 entries and 600 s TTL unless the
/// store names others, plus one lock for every read-modify-write on it.
pub struct Memory<V: Clone> {
cache: InMemoryCache<V>,
updates: Mutex<()>,
}
impl<V: Clone> Memory<V> {
pub fn new(entries: Option<usize>, now: &Clock) -> Self {
let now = now.clone();
Self {
cache: InMemoryCache::with_clock(entries, None, move || now()),
updates: Mutex::new(()),
}
}
pub fn get(&self, key: &str) -> Result<Option<V>, Error> {
self.cache.get_cache(key)
}
pub fn set(&self, key: &str, value: V, ttl: Option<Duration>) -> Result<(), Error> {
self.cache.set_cache(key, value, ttl).map(drop)
}
pub fn delete(&self, key: &str) -> Result<(), Error> {
self.cache.delete_cache(key)
}
pub fn flush(&self) -> Result<(), Error> {
self.cache.flush_cache()
}
pub fn len(&self) -> Result<usize, Error> {
self.cache.len()
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
self.cache.expires_at(key)
}
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
self.cache.entries()
}
/// Read-modify-write of one entry, atomic against every other update of this tier.
pub fn update(
&self,
key: &str,
ttl: Option<Duration>,
update: impl FnOnce(Option<V>) -> Result<V, Error>,
) -> Result<V, Error> {
let _guard = self.guard();
let value = update(self.get(key)?)?;
self.set(key, value.clone(), ttl)?;
Ok(value)
}
/// Held across a read and a conditional write that `update` cannot express.
pub fn guard(&self) -> MutexGuard<'_, ()> {
lock(&self.updates)
}
}
pub(crate) fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}

8
litellm-rust/crates/redis-store/src/number.rs
View file

@ -1,6 +1,6 @@
use std::{cmp::Ordering, ops::Add};
use litellm_python_compat::{BigInt, Value};
use litellm_python_compat::{BigInt, Value, json, repr};
/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently.
#[derive(Clone, Copy, Debug, PartialEq)]
@ -41,6 +41,12 @@ impl Number {
}
}
/// `json.dumps`, which `RedisCache.async_set_cache` writes.
pub fn to_json(self) -> String {
let value = self.to_value();
json::dumps(&value).unwrap_or_else(|_| repr::to_str(&value))
}
pub fn from_value(value: &Value) -> Option<Self> {
match value {
Value::Int(value) => i64::try_from(value).ok().map(Self::Int),

508
litellm-rust/crates/redis-store/src/rate_limits.rs
View file

@ -1,508 +0,0 @@
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};
use self::scripts::Route;
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(GAUGE)
}
}
const WINDOW: &str = "window";
const GAUGE: &str = "max_parallel_requests";
/// A counter's key suffix, which is also how the Python limiter names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "requests")]
Requests,
#[strum(serialize = "tokens")]
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(self.kind.into())
}
}
/// 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 {
/// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a
/// counter's `Kind`.
pub fn new(suffix: &str, scope: Scope) -> Result<Self, strum::ParseError> {
match suffix {
WINDOW => Ok(Self::Window(scope)),
GAUGE => Ok(Self::Gauge(scope)),
kind => kind
.parse()
.map(|kind| Self::Counter(Counter { scope, kind })),
}
}
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<Route, Error> {
let redis = self.redis.as_ref().ok_or(Error::NoRedis)?;
Ok(request.map_or_else(
|| Route::Direct(redis.clone()),
|request| Route::Batch(request.pre(redis)),
))
}
fn post(&self, request: &RequestRedis) -> Result<Route, Error> {
let redis = self.redis.as_ref().ok_or(Error::NoRedis)?;
Ok(Route::Batch(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",
)),
}
}

74
litellm-rust/crates/redis-store/src/rate_limits/keys.rs Normal file
View file

@ -0,0 +1,74 @@
/// 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 {
pub(super) fn key(&self, suffix: &str) -> String {
format!("{{{}:{}}}:{suffix}", self.descriptor, self.value)
}
pub(super) fn window_key(&self) -> String {
self.key(WINDOW)
}
pub(super) fn gauge_key(&self) -> String {
self.key(GAUGE)
}
}
const WINDOW: &str = "window";
const GAUGE: &str = "max_parallel_requests";
/// A counter's key suffix, which is also how the Python limiter names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "requests")]
Requests,
#[strum(serialize = "tokens")]
Tokens,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct Counter {
pub scope: Scope,
pub kind: Kind,
}
impl Counter {
pub(super) fn key(&self) -> String {
self.scope.key(self.kind.into())
}
}
/// 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 {
/// The entry the Python limiter names by its key suffix: `window`, `max_parallel_requests`, or a
/// counter's `Kind`.
pub fn new(suffix: &str, scope: Scope) -> Result<Self, strum::ParseError> {
match suffix {
WINDOW => Ok(Self::Window(scope)),
GAUGE => Ok(Self::Gauge(scope)),
kind => kind
.parse()
.map(|kind| Self::Counter(Counter { scope, kind })),
}
}
pub(super) fn key(&self) -> String {
match self {
Self::Window(scope) => scope.window_key(),
Self::Counter(counter) => counter.key(),
Self::Gauge(scope) => scope.gauge_key(),
}
}
}

259
litellm-rust/crates/redis-store/src/rate_limits/mod.rs Normal file
View file

@ -0,0 +1,259 @@
mod keys;
mod lua;
mod scripts;
mod types;
use std::{sync::Arc, time::Duration};
use futures_util::future::join_all;
use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient, Slot, commands};
pub use self::keys::{Counter, Entry, Kind, Scope};
pub use self::types::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply,
Increment, Local, WindowCount,
};
use crate::{
Error, Number, Pending, RequestRedis, Send,
memory::{Clock, Memory, system_clock},
remote::{Remote, RemoteSlot},
request::pending,
};
/// 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. A script needs
/// Redis, so without one it declares nothing and answers `None`.
pub struct RateLimits {
memory: Memory<Local>,
remote: RemoteSlot,
backfill_ttl: Option<Duration>,
}
impl RateLimits {
/// `backfill_ttl` is the `default_in_memory_ttl` of the `DualCache` this tier replaces. An
/// increment without a TTL takes the remote's default.
pub fn new(remote: Option<Remote>, backfill_ttl: Option<Duration>) -> Self {
Self::with_clock(remote, backfill_ttl, system_clock())
}
pub fn with_clock(remote: Option<Remote>, backfill_ttl: Option<Duration>, now: Clock) -> Self {
Self {
memory: Memory::new(None, &now),
remote: RemoteSlot::new(remote),
backfill_ttl,
}
}
pub fn has_redis(&self) -> bool {
self.remote.get().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.remote.client().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 get_local(&self, entry: &Entry) -> Result<Option<Local>, Error> {
Ok(self.memory.get(&entry.key())?)
}
pub fn set_local(&self, entry: &Entry, value: Local, ttl: Option<Duration>) -> Result<(), Error> {
Ok(self.memory.set(&entry.key(), value, ttl)?)
}
/// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default.
pub async fn get(&self, entry: &Entry, backfill_ttl: Option<Duration>) -> Result<Option<Local>, Error> {
let backfill_ttl = backfill_ttl.or(self.backfill_ttl);
self.remote
.read_through(&self.memory, &entry.key(), types::decode, backfill_ttl)
.await
}
pub fn sliding_window(
&self,
send: Send<'_>,
counters: &[Counter],
now: i64,
window_size: i64,
) -> Option<Pending<Vec<WindowCount>>> {
Some(pending(scripts::sliding_window(&self.batch(send)?, counters, now, window_size)))
}
pub fn read_counters(&self, send: Send<'_>, counters: &[Counter]) -> Option<Pending<Vec<Option<String>>>> {
Some(pending(scripts::read_counters(&self.batch(send)?, counters)))
}
pub fn check_and_increment(&self, send: Send<'_>, counters: &[AtomicCounter]) -> Option<Pending<AtomicReply>> {
Some(pending(scripts::check_and_increment(&self.batch(send)?, counters)))
}
pub fn guarded_increment(&self, send: Send<'_>, increment: &GuardedIncrement) -> Option<Pending<GuardedReply>> {
Some(pending(scripts::guarded_increment(&self.batch(send)?, increment)))
}
pub fn parallel_acquire(
&self,
send: Send<'_>,
gauges: &[Gauge],
slot_ttl: Duration,
slot_id: &str,
) -> Option<Pending<AcquireReply>> {
Some(pending(scripts::parallel_acquire(&self.batch(send)?, gauges, slot_ttl, slot_id)))
}
pub fn parallel_release(&self, send: Send<'_>, gauges: &[Scope], slot_id: &str) -> Option<Pending<Vec<i64>>> {
Some(pending(self.declare_release(send, gauges, slot_id)?))
}
pub fn parallel_count(&self, send: Send<'_>, gauges: &[Scope], slot_ttl: Duration) -> Option<Pending<Vec<i64>>> {
Some(pending(scripts::parallel_count(&self.batch(send)?, gauges, slot_ttl)))
}
pub fn token_increment(&self, send: Send<'_>, increments: &[Increment]) -> Option<Pending<Vec<f64>>> {
Some(pending(self.declare_token_increment(send, 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. `false` without Redis.
pub fn token_increment_post_call(self: &Arc<Self>, request: &RequestRedis, increments: Vec<Increment>) -> bool {
let Some(declared) = self.declare_token_increment(Send::PostCall(request), &increments) else {
return false;
};
let store = Arc::clone(self);
request.spawn(async move {
let Err(error) = declared.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");
}
});
true
}
/// `PARALLEL_RELEASE_SCRIPT` on the request's post-call batch; the caller has already freed
/// the slot in memory. `false` without Redis.
pub fn parallel_release_post_call(&self, request: &RequestRedis, gauges: &[Scope], slot_id: &str) -> bool {
let Some(declared) = self.declare_release(Send::PostCall(request), gauges, slot_id) else {
return false;
};
request.spawn(async move {
if let Err(error) = declared.settled().await {
tracing::warn!(%error, "parallel_release_script failed, the slot stays released in memory only");
}
});
true
}
/// `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, Error> {
let local = self.add_local(increment)?;
let Some(remote) = self.remote.get() else {
return Ok(local);
};
let ttl = increment.ttl.unwrap_or(remote.default_ttl);
let total = remote
.client
.send(commands::incr_by_float(
&increment.counter.key(),
increment.amount.into(),
Some(Expiry::KeepExisting(ttl)),
))
.await;
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<(), Error> {
increments
.iter()
.try_for_each(|increment| self.add_local(increment).map(drop))?;
let Some(client) = self.remote.client() else {
return Ok(());
};
let pipeline = RedisBatch::pipelined(client.clone());
let totals = increments.iter().map(|increment| {
pipeline.add(commands::incr_by_float(
&increment.counter.key(),
increment.amount.into(),
increment.ttl.map(Expiry::Refresh),
))
});
let totals = Declared::all(totals.collect::<Vec<_>>()).await;
if let Some(error) = totals.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<(), Error> {
increments
.iter()
.try_for_each(|increment| self.add_local(increment).map(drop))?;
let Some(batch) = self.remote.batch(Send::PostCall(request)) else {
return Ok(());
};
let declared = increments
.iter()
.map(|increment| scripts::counter_increment(&batch, increment))
.collect::<Vec<_>>();
request.spawn(async move {
let outcomes = join_all(declared.into_iter().map(Declared::settled)).await;
if let Some(error) = outcomes.into_iter().find_map(Result::err) {
tracing::warn!(%error, "deferred Redis increment failed");
}
});
Ok(())
}
fn declare_release(&self, send: Send<'_>, gauges: &[Scope], slot_id: &str) -> Option<Declared<Vec<i64>>> {
Some(scripts::parallel_release(&self.batch(send)?, gauges, slot_id))
}
fn declare_token_increment(&self, send: Send<'_>, increments: &[Increment]) -> Option<Declared<Vec<f64>>> {
Some(scripts::token_increment(&self.batch(send)?, increments))
}
/// `InMemoryCache.increment_cache`: `(current or 0) + amount` under one lock.
fn add_local(&self, increment: &Increment) -> Result<Number, Error> {
let key = increment.counter.key();
let _guard = self.memory.guard();
let total = match self.memory.get(&key)? {
None => increment.amount,
Some(Local::Number(current)) => current + increment.amount,
Some(_) => return Err(litellm_cache::Error::InvalidEntry(None).into()),
};
self.memory.set(&key, Local::Number(total), increment.ttl)?;
Ok(total)
}
fn batch(&self, send: Send<'_>) -> Option<RedisBatch> {
self.remote.batch(send)
}
}

231
litellm-rust/crates/redis-store/src/rate_limits/scripts.rs
View file

@ -1,10 +1,16 @@
use std::{sync::LazyLock, time::Duration};
use futures_util::future::BoxFuture;
use litellm_python_compat::repr::float_repr;
use litellm_redis_client::{Declared, RedisBatch, RedisClient, Script, Value};
use litellm_redis_client::{
Command, Declared, Error, Expiry, RedisBatch, Script, Value,
commands::{self, unexpected_reply},
};
use super::{Counter, Error, Increment, Scope, lua};
use super::{
keys::{Counter, Scope},
lua,
types::{AcquireReply, Applied, AtomicCounter, AtomicReply, Gauge, GuardedIncrement, GuardedReply, Increment, WindowCount},
};
use crate::Number;
static BATCH_RATE_LIMITER: LazyLock<Script> =
@ -20,151 +26,32 @@ static PARALLEL_RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PA
static PARALLEL_COUNT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_COUNT));
static TOKEN_INCREMENT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::TOKEN_INCREMENT));
/// Where a script call goes: a batch's pipeline, or alone like a redis-py `Script` called outside
/// one, which answers `NOSCRIPT` with `SCRIPT LOAD` and `EVALSHA`.
pub(super) enum Route {
Batch(RedisBatch),
Direct(RedisClient),
}
enum Reply {
Declared(Declared<Value>),
Direct(BoxFuture<'static, Result<Value, litellm_redis_client::Error>>),
}
/// One script call, decoded into its typed reply once it has run.
#[must_use]
pub struct Pending<T> {
reply: Reply,
decode: fn(Value) -> Result<T, Error>,
}
impl<T> Pending<T> {
fn new(reply: Reply, 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> {
let reply = match self.reply {
Reply::Declared(declared) => declared.resolve().await?,
Reply::Direct(call) => call.await?,
};
(self.decode)(reply)
}
/// Waits for whoever flushes the batch.
pub async fn settled(self) -> Result<T, Error> {
let reply = match self.reply {
Reply::Declared(declared) => declared.settled().await?,
Reply::Direct(call) => call.await?,
};
(self.decode)(reply)
}
}
/// `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(
route: &Route,
batch: &RedisBatch,
counters: &[Counter],
now: i64,
window_size: i64,
) -> Pending<Vec<WindowCount>> {
) -> Declared<Vec<WindowCount>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
Pending::new(
eval(
route,
&BATCH_RATE_LIMITER,
&keys,
vec![now.to_string(), window_size.to_string()],
),
window_counts,
)
batch.add(eval(&BATCH_RATE_LIMITER, &keys, vec![now.to_string(), window_size.to_string()],).map(window_counts))
}
/// `BATCH_COUNTER_READ_SCRIPT` over each counter's window and counter keys, in that order.
pub(super) fn read_counters(route: &Route, counters: &[Counter]) -> Pending<Vec<Option<String>>> {
pub(super) fn read_counters(batch: &RedisBatch, counters: &[Counter]) -> Declared<Vec<Option<String>>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
Pending::new(
eval(route, &BATCH_COUNTER_READ, &keys, Vec::new()),
optional_texts,
)
batch.add(eval(&BATCH_COUNTER_READ, &keys, Vec::new()).map(optional_texts))
}
pub(super) fn check_and_increment(
route: &Route,
batch: &RedisBatch,
counters: &[AtomicCounter],
) -> Pending<AtomicReply> {
) -> Declared<AtomicReply> {
let keys = counters
.iter()
.flat_map(|counter| [counter.counter.scope.window_key(), counter.counter.key()])
@ -181,16 +68,13 @@ pub(super) fn check_and_increment(
.map(|arg| arg.to_string())
})
.collect();
Pending::new(
eval(route, &CHECK_AND_INCREMENT_BY_N, &keys, args),
atomic_reply,
)
batch.add(eval(&CHECK_AND_INCREMENT_BY_N, &keys, args).map(atomic_reply))
}
pub(super) fn guarded_increment(
route: &Route,
batch: &RedisBatch,
increment: &GuardedIncrement,
) -> Pending<GuardedReply> {
) -> Declared<GuardedReply> {
let keys = [
increment.counter.scope.window_key(),
increment.counter.key(),
@ -200,18 +84,15 @@ pub(super) fn guarded_increment(
python_arg(increment.amount),
ttl_arg(increment.ttl),
];
Pending::new(
eval(route, &WINDOW_GUARDED_TOKEN_INCREMENT, &keys, args),
guarded_reply,
)
batch.add(eval(&WINDOW_GUARDED_TOKEN_INCREMENT, &keys, args).map(guarded_reply))
}
pub(super) fn parallel_acquire(
route: &Route,
batch: &RedisBatch,
gauges: &[Gauge],
slot_ttl: Duration,
slot_id: &str,
) -> Pending<AcquireReply> {
) -> Declared<AcquireReply> {
let keys = gauges
.iter()
.map(|gauge| gauge.scope.gauge_key())
@ -226,33 +107,33 @@ pub(super) fn parallel_acquire(
]
})
.collect();
Pending::new(eval(route, &PARALLEL_ACQUIRE, &keys, args), acquire_reply)
batch.add(eval(&PARALLEL_ACQUIRE, &keys, args).map(acquire_reply))
}
pub(super) fn parallel_release(
route: &Route,
batch: &RedisBatch,
gauges: &[Scope],
slot_id: &str,
) -> Pending<Vec<i64>> {
) -> Declared<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(route, &PARALLEL_RELEASE, &keys, args), integers)
batch.add(eval(&PARALLEL_RELEASE, &keys, args).map(integers))
}
pub(super) fn parallel_count(
route: &Route,
batch: &RedisBatch,
gauges: &[Scope],
slot_ttl: Duration,
) -> Pending<Vec<i64>> {
) -> Declared<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(route, &PARALLEL_COUNT, &keys, args), integers)
batch.add(eval(&PARALLEL_COUNT, &keys, args).map(integers))
}
pub(super) fn token_increment(route: &Route, increments: &[Increment]) -> Pending<Vec<f64>> {
pub(super) fn token_increment(batch: &RedisBatch, increments: &[Increment]) -> Declared<Vec<f64>> {
let keys = increments
.iter()
.map(|increment| increment.counter.key())
@ -261,32 +142,22 @@ pub(super) fn token_increment(route: &Route, increments: &[Increment]) -> Pendin
.iter()
.flat_map(|increment| [python_arg(increment.amount), ttl_arg(increment.ttl)])
.collect();
Pending::new(eval(route, &TOKEN_INCREMENT, &keys, args), floats)
batch.add(eval(&TOKEN_INCREMENT, &keys, args).map(floats))
}
/// `RedisBatch`'s `_Increment`: `INCRBYFLOAT`, then `EXPIRE` when the increment carries a TTL.
/// `RedisBatch`'s `_Increment`: `INCRBYFLOAT`, then `EXPIRE` when the increment carries a TTL. On
/// its own it keeps a TTL the key already has, like `async_increment`.
pub(super) fn counter_increment(batch: &RedisBatch, increment: &Increment) -> Declared<f64> {
batch.incr_by_float(
batch.add(commands::incr_by_float(
&increment.counter.key(),
increment.amount.into(),
increment.ttl,
)
increment.ttl.map(Expiry::KeepExisting),
))
}
fn eval(route: &Route, script: &Script, keys: &[String], args: Vec<String>) -> Reply {
match route {
Route::Batch(batch) => {
let keys = keys.iter().map(String::as_str).collect::<Vec<_>>();
Reply::Declared(batch.eval(script, &keys, args))
}
Route::Direct(client) => {
let (client, script, keys) = (client.clone(), script.clone(), keys.to_vec());
Reply::Direct(Box::pin(async move {
let keys = keys.iter().map(String::as_str).collect::<Vec<_>>();
client.eval(&script, &keys, args).await
}))
}
}
fn eval(script: &Script, keys: &[String], args: Vec<String>) -> Command<Value> {
let keys = keys.iter().map(String::as_str).collect::<Vec<_>>();
commands::eval(script, &keys, args)
}
/// The bytes redis-py sends for a Python `int` or `float` argument.
@ -305,14 +176,14 @@ fn ttl_arg(ttl: Option<Duration>) -> String {
fn items(reply: Value) -> Result<Vec<Value>, Error> {
match reply {
Value::Array(items) => Ok(items),
_ => Err(Error::Reply("is not a list")),
_ => Err(unexpected_reply("rate limiter script 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")),
_ => Err(unexpected_reply("rate limiter script reply holds a non-integer where a count belongs")),
}
}
@ -321,21 +192,21 @@ fn text(value: &Value) -> Result<Option<String>, Error> {
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")),
.map_err(|_| unexpected_reply("rate limiter script 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")),
_ => Err(unexpected_reply("rate limiter script 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"))
.ok_or_else(|| unexpected_reply("rate limiter script 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"))
usize::try_from(integer(raw)? - 1).map_err(|_| unexpected_reply("rate limiter script reply names an index below 1"))
}
fn window_counts(reply: Value) -> Result<Vec<WindowCount>, Error> {
@ -344,10 +215,10 @@ fn window_counts(reply: Value) -> Result<Vec<WindowCount>, Error> {
.map(|pair| match pair {
[window_start, count] => Ok(WindowCount {
window_start: text(window_start)?
.ok_or(Error::Reply("holds a nil window start"))?,
.ok_or_else(|| unexpected_reply("rate limiter script reply holds a nil window start"))?,
count: integer(count)?,
}),
_ => Err(Error::Reply("has an odd length")),
_ => Err(unexpected_reply("rate limiter script reply has an odd length")),
})
.collect()
}
@ -374,7 +245,7 @@ fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
count: integer(count)?,
window_start: integer(window_start)?,
}),
_ => Err(Error::Reply("has an odd length")),
_ => Err(unexpected_reply("rate limiter script reply has an odd length")),
})
.collect::<Result<_, _>>()
.map(AtomicReply::Applied),
@ -383,7 +254,7 @@ fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
current: integer(current)?,
limit: integer(limit)?,
}),
_ => Err(Error::Reply("has no status")),
_ => Err(unexpected_reply("rate limiter script reply has no status")),
}
}
@ -393,7 +264,7 @@ fn guarded_reply(reply: Value) -> Result<GuardedReply, Error> {
applied: integer(applied)? == 1,
count: integer(count)?,
}),
_ => Err(Error::Reply("is not an (applied, count) pair")),
_ => Err(unexpected_reply("rate limiter script reply is not an (applied, count) pair")),
}
}
@ -410,6 +281,6 @@ fn acquire_reply(reply: Value) -> Result<AcquireReply, Error> {
in_flight: integer(in_flight)?,
limit: integer(limit)?,
}),
_ => Err(Error::Reply("has no status")),
_ => Err(unexpected_reply("rate limiter script reply has no status")),
}
}

191
litellm-rust/crates/redis-store/src/rate_limits/types.rs Normal file
View file

@ -0,0 +1,191 @@
use std::time::Duration;
use litellm_python_compat::Value;
use super::keys::{Counter, Entry, Scope};
use crate::{DecodeError, FromValue, Number, codec::python_value, value};
/// `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,
},
}
const SCOPE: &str = "a rate-limit scope is (descriptor_key, descriptor_value)";
const COUNTER: &str =
"a rate-limit counter is (descriptor_key, descriptor_value, \"requests\" | \"tokens\")";
const ENTRY: &str = "a rate-limit entry is (key_suffix, descriptor_key, descriptor_value)";
const INCREMENT: &str = "a rate-limit increment is (counter, amount, ttl_seconds | None)";
const ATOMIC: &str =
"an atomic counter is (counter, limit, increment, ttl_seconds, window_size_seconds)";
const GAUGE_LIMIT: &str = "a gauge is (scope, limit)";
/// `(descriptor_key, descriptor_value)`.
impl FromValue for Scope {
fn from_value(raw: &Value) -> Result<Self, DecodeError> {
let [descriptor, value] = value::fields(raw, SCOPE)?;
Ok(Self {
descriptor: value::text(descriptor, SCOPE)?,
value: value::text(value, SCOPE)?,
})
}
}
/// `(descriptor_key, descriptor_value, kind)`, the kind spelled as its key suffix.
impl FromValue for Counter {
fn from_value(raw: &Value) -> Result<Self, DecodeError> {
let [descriptor, value, kind] = value::fields(raw, COUNTER)?;
Ok(Self {
scope: Scope {
descriptor: value::text(descriptor, COUNTER)?,
value: value::text(value, COUNTER)?,
},
kind: value::text(kind, COUNTER)?
.parse()
.map_err(|_| DecodeError::Shape(COUNTER))?,
})
}
}
/// `(key_suffix, descriptor_key, descriptor_value)`, read by `Entry::new`.
impl FromValue for Entry {
fn from_value(raw: &Value) -> Result<Self, DecodeError> {
let [suffix, descriptor, value] = value::fields(raw, ENTRY)?;
let scope = Scope {
descriptor: value::text(descriptor, ENTRY)?,
value: value::text(value, ENTRY)?,
};
Self::new(&value::text(suffix, ENTRY)?, scope).map_err(|_| DecodeError::Shape(ENTRY))
}
}
/// `(counter, amount, ttl_seconds | None)`.
impl FromValue for Increment {
fn from_value(raw: &Value) -> Result<Self, DecodeError> {
let [counter, amount, ttl] = value::fields(raw, INCREMENT)?;
Ok(Self {
counter: Counter::from_value(counter)?,
amount: value::number(amount, INCREMENT)?,
ttl: match ttl {
Value::None => None,
ttl => Some(Duration::from_secs(
u64::try_from(value::integer(ttl, INCREMENT)?)
.map_err(|_| DecodeError::Shape(INCREMENT))?,
)),
},
})
}
}
/// `(counter, limit, increment, ttl_seconds, window_size_seconds)`.
impl FromValue for AtomicCounter {
fn from_value(raw: &Value) -> Result<Self, DecodeError> {
let [counter, limit, increment, ttl_seconds, window_size_seconds] =
value::fields(raw, ATOMIC)?;
Ok(Self {
counter: Counter::from_value(counter)?,
limit: value::integer(limit, ATOMIC)?,
increment: value::integer(increment, ATOMIC)?,
ttl_seconds: value::integer(ttl_seconds, ATOMIC)?,
window_size_seconds: value::integer(window_size_seconds, ATOMIC)?,
})
}
}
/// `(scope, limit)`.
impl FromValue for Gauge {
fn from_value(raw: &Value) -> Result<Self, DecodeError> {
let [scope, limit] = value::fields(raw, GAUGE_LIMIT)?;
Ok(Self {
scope: Scope::from_value(scope)?,
limit: value::integer(limit, GAUGE_LIMIT)?,
})
}
}
/// 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>,
}
pub(super) fn decode(raw: &[u8]) -> Result<Local, DecodeError> {
match python_value(raw)? {
Value::Str(text) => Ok(Local::Text(text)),
value => Number::from_value(&value)
.map(Local::Number)
.ok_or(DecodeError::Shape(
"a rate-limit value is a number or a string",
)),
}
}

82
litellm-rust/crates/redis-store/src/remote.rs Normal file
View file

@ -0,0 +1,82 @@
use std::{sync::OnceLock, time::Duration};
use litellm_redis_client::{RedisBatch, RedisClient, commands};
use crate::{Error, Send, memory::Memory};
/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `DualCache.async_delete_cache_keys` deletes at once.
const DELETE_CHUNK: usize = 1000;
/// The Redis a store shares with every pod, with the `RedisCache.default_ttl` a write takes when it
/// names none.
#[derive(Clone)]
pub struct Remote {
pub client: RedisClient,
pub default_ttl: Duration,
}
/// A store's Redis, absent until the host attaches one, as `DualCache.attach_redis_cache` leaves
/// memory in place and only adds the second tier. It is attached at most once.
pub(crate) struct RemoteSlot(OnceLock<Remote>);
impl RemoteSlot {
pub(crate) fn new(remote: Option<Remote>) -> Self {
Self(remote.map(OnceLock::from).unwrap_or_default())
}
/// `false` when a Redis was already attached, which keeps it.
pub(crate) fn attach(&self, remote: Remote) -> bool {
self.0.set(remote).is_ok()
}
pub(crate) fn get(&self) -> Option<&Remote> {
self.0.get()
}
pub(crate) fn client(&self) -> Option<&RedisClient> {
self.get().map(|remote| &remote.client)
}
pub(crate) fn batch(&self, send: Send<'_>) -> Option<RedisBatch> {
Some(send.batch(self.client()?))
}
/// `DualCache.async_get_cache`: memory, else Redis written back to memory for `backfill_ttl`. A
/// failed or undecodable Redis read is a miss, as `DualCache` logs and swallows it.
pub(crate) async fn read_through<V: Clone, E: Into<Error>>(
&self,
memory: &Memory<V>,
key: &str,
decode: impl FnOnce(&[u8]) -> Result<V, E>,
backfill_ttl: Option<Duration>,
) -> Result<Option<V>, Error> {
if let Some(value) = memory.get(key)? {
return Ok(Some(value));
}
let Some(client) = self.client() else {
return Ok(None);
};
let read = client.send(commands::get::<Vec<u8>>(key)).await;
let value = read
.map_err(Error::from)
.and_then(|raw| raw.as_deref().map(decode).transpose().map_err(Into::into))
.unwrap_or_else(|error| {
tracing::error!(%error, "LiteLLM Cache: exception in async_get_cache");
None
});
let Some(value) = value else {
return Ok(None);
};
memory.set(key, value.clone(), backfill_ttl)?;
Ok(Some(value))
}
}
/// `DualCache.async_delete_cache_keys`: one multi-key `DEL` per chunk, in order.
pub(crate) async fn delete(client: &RedisClient, keys: &[String]) -> Result<(), Error> {
for chunk in keys.chunks(DELETE_CHUNK) {
let chunk: Vec<&str> = chunk.iter().map(String::as_str).collect();
client.send(commands::del(&chunk)).await?;
}
Ok(())
}

45
litellm-rust/crates/redis-store/src/request.rs
View file

@ -1,13 +1,15 @@
use std::{
future::Future,
sync::{Arc, Mutex, MutexGuard, PoisonError},
sync::{Arc, Mutex},
time::Duration,
};
use futures_util::future::join_all;
use litellm_redis_client::{RedisBatch, RedisClient};
use futures_util::future::{BoxFuture, join_all};
use litellm_redis_client::{Declared, RedisBatch, RedisClient};
use tokio::{runtime::Handle, sync::oneshot, task::JoinHandle};
use crate::memory::lock;
/// One request's Redis batches: a pre-call batch for the reads it needs before answering and a
/// post-call batch for the writes nobody waits on, one of each per Redis client. The host that owns
/// the request builds one and passes it down; it knows nothing else about the request.
@ -56,7 +58,7 @@ impl RequestRedis {
/// Runs work that follows this request's post-call batch, such as mirroring settled totals
/// into memory, on the request's runtime, from any thread. `flush_post` waits for it, so the
/// shutdown drain does not drop a fallback write still in flight.
pub fn spawn(&self, task: impl Future<Output = ()> + Send + 'static) {
pub fn spawn(&self, task: impl Future<Output = ()> + std::marker::Send + 'static) {
let follow_up = self.0.runtime.spawn(task);
lock(&self.0.follow_ups).push(follow_up);
}
@ -116,6 +118,37 @@ impl RequestRedis {
}
}
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
/// Which batch an op rides, picked by the caller because only it knows when it needs the answer.
/// The store owns its client, so it calls `batch`, the one place a request and a client meet.
#[derive(Clone, Copy)]
pub enum Send<'a> {
/// On its own, sent at the first `resolve`.
Now,
/// The request's pre-call batch, sent at the first `resolve` or `flush_pre`.
PreCall(&'a RequestRedis),
/// The request's post-call batch, sent at `flush_post` or the post-call deadline.
PostCall(&'a RequestRedis),
}
impl Send<'_> {
pub fn batch(&self, client: &RedisClient) -> RedisBatch {
match self {
Self::Now => RedisBatch::alone(client.clone()),
Self::PreCall(request) => request.pre(client),
Self::PostCall(request) => request.post(client),
}
}
}
/// A store op declared on the batch its `Send` picked: building it sends nothing, and awaiting it
/// flushes that batch if nobody has yet.
pub type Pending<T> = BoxFuture<'static, Result<T, crate::Error>>;
pub(crate) fn pending<T: std::marker::Send + 'static>(declared: Declared<T>) -> Pending<T> {
Box::pin(async move { Ok(declared.resolve().await?) })
}
/// An op answered without Redis, from memory.
pub(crate) fn ready<T: std::marker::Send + 'static>(value: T) -> Pending<T> {
Box::pin(std::future::ready(Ok(value)))
}

290
litellm-rust/crates/redis-store/src/router_cooldowns.rs
View file

@ -1,290 +0,0 @@
use std::time::Duration;
use litellm_cache::Error as MemoryError;
use litellm_cache_memory::InMemoryCache;
use litellm_python_compat::{Value, repr::repr};
use litellm_redis_client::{RedisBatch, RedisClient};
use crate::{
DecodeError, Number, RequestRedis,
codec::python_value,
tier::{Clock, DeclaredRead, MemoryTier, OnBreakerOpen, TierPolicy, system_clock},
};
/// `main`'s cooldown tier: `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` (1 s), `InMemoryCache`'s
/// 600 s default for backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(1),
backfill_ttl: Duration::from_secs(600),
on_breaker_open: OnBreakerOpen::ServeMemory,
tracked_reads: 1000,
};
const MEMORY_ENTRIES: usize = 200;
const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60);
const CORRECTION_SLACK: Duration = Duration::from_secs(5);
#[derive(Clone, Debug, PartialEq)]
pub struct Cooldown {
pub exception_received: String,
pub status_code: String,
pub timestamp: Number,
pub cooldown_time: Number,
}
impl Cooldown {
fn remaining(&self, now: Duration) -> f64 {
self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64()
}
/// The `CooldownCacheValue` dict the router reads and writes.
pub fn to_value(&self) -> Value {
let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value);
Value::Dict(vec![
field(
"exception_received",
Value::Str(self.exception_received.clone()),
),
field("status_code", Value::Str(self.status_code.clone())),
field("timestamp", self.timestamp.to_value()),
field("cooldown_time", self.cooldown_time.to_value()),
])
}
pub fn from_value(value: &Value) -> Result<Self, DecodeError> {
let Value::Dict(fields) = value else {
return Err(DecodeError::Shape("a cooldown is a dict"));
};
let field = |name: &str| {
fields
.iter()
.find(|(key, _)| matches!(key, Value::Str(key) if key == name))
.map(|(_, value)| value)
};
let text = |name: &'static str| match field(name) {
Some(Value::Str(text)) => Ok(text.clone()),
_ => Err(DecodeError::Shape(name)),
};
let number = |name: &'static str| {
field(name)
.and_then(Number::from_value)
.ok_or(DecodeError::Shape(name))
};
Ok(Self {
exception_received: text("exception_received")?,
status_code: text("status_code")?,
timestamp: number("timestamp")?,
cooldown_time: number("cooldown_time")?,
})
}
/// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`.
pub fn encode(&self) -> String {
repr(&self.to_value())
}
/// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does.
pub fn decode(raw: &[u8]) -> Result<Self, DecodeError> {
Self::from_value(&python_value(raw)?)
}
}
fn key(model_id: &str) -> String {
format!("deployment:{model_id}:cooldown")
}
fn keys(model_ids: &[String]) -> Vec<String> {
model_ids.iter().map(|model_id| key(model_id)).collect()
}
fn allowed_fails_key(model_id: &str, suffix: Option<&str>) -> String {
match suffix.filter(|suffix| !suffix.is_empty()) {
Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"),
None => format!("deployment:{model_id}:allowed_fails"),
}
}
fn memory<V: Clone>(policy: TierPolicy, now: &Clock) -> InMemoryCache<V> {
let now = now.clone();
InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || {
now()
})
}
/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds.
fn whole_seconds(ttl: Duration) -> Duration {
Duration::from_secs(ttl.as_secs())
}
/// An `active` read declared on a request batch, answered by `resolve_active`.
pub struct DeclaredCooldowns {
model_ids: Vec<String>,
read: DeclaredRead,
}
/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod.
pub struct RouterCooldowns {
cooldowns: MemoryTier<InMemoryCache<Cooldown>, RedisClient>,
fails: MemoryTier<InMemoryCache<i64>, RedisClient>,
}
impl RouterCooldowns {
pub fn new(client: Option<RedisClient>, policy: TierPolicy) -> Self {
Self::with_clock(client, policy, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, policy: TierPolicy, now: Clock) -> Self {
Self {
cooldowns: MemoryTier::read_through(
memory(policy, &now),
client.clone(),
policy,
now.clone(),
),
fails: MemoryTier::read_through(memory(policy, &now), client, policy, now),
}
}
/// Memory first, then a `SET ... EX int(ttl)` whose failure only logs, like `DualCache.set_cache`.
pub async fn put(
&self,
model_id: &str,
cooldown: &Cooldown,
ttl: Duration,
) -> Result<(), MemoryError> {
let key = key(model_id);
self.cooldowns.set_memory(&key, cooldown.clone(), ttl)?;
if let Some(remote) = self.cooldowns.remote()
&& let Err(error) = remote
.set(&key, cooldown.encode(), whole_seconds(ttl))
.await
{
tracing::warn!(%error, "writing a cooldown to Redis failed");
}
Ok(())
}
/// The fleet-wide fail count. A Redis failure falls back to this router's own count.
pub async fn increment_allowed_fails(
&self,
model_id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<i64, MemoryError> {
let key = allowed_fails_key(model_id, suffix);
let local = self
.fails
.update_memory(&key, ttl, |count| count.unwrap_or(0) + 1)?;
let Some(remote) = self.fails.remote() else {
return Ok(local);
};
match remote.incr_by(&key, 1, whole_seconds(ttl)).await {
Ok(count) => Ok(count),
Err(error) => {
tracing::warn!(%error, "allowed_fails counter fell back to this worker's in-memory count");
Ok(local)
}
}
}
/// The deployments still cooling down. An entry whose cooldown already ended is dropped
/// from memory, and one memory would keep past its end is shortened, at most to 60 s.
pub async fn active(
&self,
model_ids: &[String],
) -> Result<Vec<(String, Cooldown)>, MemoryError> {
let declared = self.declare(RedisClient::batch, model_ids)?;
self.resolve_active(declared).await
}
/// `active`, with its Redis read declared on the request's pre-call batch.
pub fn declare_active(
&self,
request: &RequestRedis,
model_ids: &[String],
) -> Result<DeclaredCooldowns, MemoryError> {
self.declare(|client| request.pre(client), model_ids)
}
pub async fn resolve_active(
&self,
declared: DeclaredCooldowns,
) -> Result<Vec<(String, Cooldown)>, MemoryError> {
let DeclaredCooldowns { model_ids, read } = declared;
let read = self.cooldowns.resolve_many(read, Cooldown::decode).await?;
let now = self.cooldowns.now();
let corrected = model_ids
.into_iter()
.zip(read)
.filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown)))
.map(|(model_id, cooldown)| {
Ok(self
.correct(&model_id, &cooldown, now)?
.then_some((model_id, cooldown)))
})
.collect::<Result<Vec<_>, MemoryError>>()?;
Ok(corrected.into_iter().flatten().collect())
}
/// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it.
pub async fn minimum(&self, model_ids: &[String]) -> Result<Option<Number>, MemoryError> {
let read = self.read(model_ids).await?;
Ok(Number::min(
read.into_iter()
.flatten()
.map(|cooldown| cooldown.cooldown_time),
))
}
pub fn evict_local(&self, model_id: &str) -> Result<(), MemoryError> {
self.cooldowns.delete_memory(&key(model_id))
}
pub fn flush_local_counters(&self) -> Result<(), MemoryError> {
self.fails.flush_memory()
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Cooldown)>, MemoryError> {
self.cooldowns.memory().entries()
}
async fn read(&self, model_ids: &[String]) -> Result<Vec<Option<Cooldown>>, MemoryError> {
self.cooldowns
.get_many(&keys(model_ids), Cooldown::decode)
.await
}
fn declare(
&self,
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
model_ids: &[String],
) -> Result<DeclaredCooldowns, MemoryError> {
Ok(DeclaredCooldowns {
model_ids: model_ids.to_vec(),
read: self.cooldowns.declare_many(batch_for, &keys(model_ids))?,
})
}
fn correct(
&self,
model_id: &str,
cooldown: &Cooldown,
now: Duration,
) -> Result<bool, MemoryError> {
let key = key(model_id);
let remaining = cooldown.remaining(now);
if remaining <= 0.0 {
self.cooldowns.delete_memory(&key)?;
return Ok(false);
}
let remaining = Duration::from_secs_f64(remaining);
let memory = self.cooldowns.memory();
if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) {
self.cooldowns.delete_memory(&key)?;
self.cooldowns
.set_memory(&key, cooldown.clone(), remaining.min(MAX_CORRECTED_TTL))?;
}
Ok(true)
}
}

14
litellm-rust/crates/redis-store/src/router_cooldowns/keys.rs Normal file
View file

@ -0,0 +1,14 @@
pub(super) fn cooldown(model_id: &str) -> String {
format!("deployment:{model_id}:cooldown")
}
pub(super) fn cooldowns(model_ids: &[String]) -> Vec<String> {
model_ids.iter().map(|model_id| cooldown(model_id)).collect()
}
pub(super) fn allowed_fails(model_id: &str, suffix: Option<&str>) -> String {
match suffix.filter(|suffix| !suffix.is_empty()) {
Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"),
None => format!("deployment:{model_id}:allowed_fails"),
}
}

158
litellm-rust/crates/redis-store/src/router_cooldowns/mod.rs Normal file
View file

@ -0,0 +1,158 @@
mod keys;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::commands;
pub use self::types::Cooldown;
use crate::{
Error, Number, Pending, Send,
memory::{Clock, Memory, system_clock},
remote::{Remote, RemoteSlot},
tier::{MemoryTier, OnBreakerOpen, TierPolicy},
};
/// `main`'s cooldown tier: `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` (1 s), `InMemoryCache`'s
/// 600 s default for backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(1),
backfill_ttl: Duration::from_secs(600),
on_breaker_open: OnBreakerOpen::ServeMemory,
tracked_reads: 1000,
};
const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60);
const CORRECTION_SLACK: Duration = Duration::from_secs(5);
/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod.
pub struct RouterCooldowns {
cooldowns: MemoryTier<Cooldown>,
fails: Memory<i64>,
remote: RemoteSlot,
}
impl RouterCooldowns {
pub fn new(remote: Option<Remote>, policy: TierPolicy) -> Self {
Self::with_clock(remote, policy, system_clock())
}
pub fn with_clock(remote: Option<Remote>, policy: TierPolicy, now: Clock) -> Self {
Self {
fails: Memory::new(None, &now),
cooldowns: MemoryTier::read_through(policy, now),
remote: RemoteSlot::new(remote),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.remote.attach(remote)
}
/// Memory first, then a `SET ... EX int(ttl)` whose failure only logs, like `DualCache.set_cache`.
pub async fn put(&self, model_id: &str, cooldown: &Cooldown, ttl: Duration) -> Result<(), Error> {
let key = keys::cooldown(model_id);
self.cooldowns.memory().set(&key, cooldown.clone(), Some(ttl))?;
if let Some(client) = self.remote.client()
&& let Err(error) = client.send(commands::set(&key, cooldown.encode(), ttl)).await
{
tracing::warn!(%error, "writing a cooldown to Redis failed");
}
Ok(())
}
/// The fleet-wide fail count. A Redis failure falls back to this router's own count.
pub async fn increment_allowed_fails(
&self,
model_id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<i64, Error> {
let key = keys::allowed_fails(model_id, suffix);
let local = self
.fails
.update(&key, Some(ttl), |count| Ok(count.unwrap_or(0) + 1))?;
let Some(client) = self.remote.client() else {
return Ok(local);
};
match client.send(commands::incr_by(&key, 1, ttl)).await {
Ok(count) => Ok(count),
Err(error) => {
tracing::warn!(%error, "allowed_fails counter fell back to this worker's in-memory count");
Ok(local)
}
}
}
/// The deployments still cooling down, read on the batch `send` picks. An entry whose cooldown
/// already ended is dropped from memory, and one memory would keep past its end is shortened, at
/// most to 60 s.
pub fn active(
self: &Arc<Self>,
send: Send<'_>,
model_ids: &[String],
) -> Result<Pending<Vec<(String, Cooldown)>>, Error> {
let read = self
.cooldowns
.declare_many(&self.remote, send, &keys::cooldowns(model_ids))?;
let model_ids = model_ids.to_vec();
let store = Arc::clone(self);
Ok(Box::pin(async move {
let read = store.cooldowns.resolve_many(read, Cooldown::decode).await?;
let now = store.cooldowns.now();
let active = model_ids
.into_iter()
.zip(read)
.filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown)))
.map(|(model_id, cooldown)| {
Ok(store
.correct(&model_id, &cooldown, now)?
.then_some((model_id, cooldown)))
})
.collect::<Result<Vec<_>, Error>>()?;
Ok(active.into_iter().flatten().collect())
}))
}
/// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it.
pub async fn minimum(&self, model_ids: &[String]) -> Result<Option<Number>, Error> {
let read = self
.cooldowns
.get_many(&self.remote, &keys::cooldowns(model_ids), Cooldown::decode)
.await?;
Ok(Number::min(
read.into_iter().flatten().map(|cooldown| cooldown.cooldown_time),
))
}
pub fn delete_local(&self, model_id: &str) -> Result<(), Error> {
Ok(self.cooldowns.memory().delete(&keys::cooldown(model_id))?)
}
/// Drops this router's allowed-fails counts, leaving its cooldowns in place.
pub fn flush_local_counters(&self) -> Result<(), Error> {
Ok(self.fails.flush()?)
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Cooldown)>, Error> {
Ok(self.cooldowns.memory().entries()?)
}
fn correct(&self, model_id: &str, cooldown: &Cooldown, now: Duration) -> Result<bool, Error> {
let key = keys::cooldown(model_id);
let memory = self.cooldowns.memory();
let remaining = cooldown.remaining(now);
if remaining <= 0.0 {
memory.delete(&key)?;
return Ok(false);
}
let remaining = Duration::from_secs_f64(remaining);
if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) {
memory.delete(&key)?;
memory.set(&key, cooldown.clone(), Some(remaining.min(MAX_CORRECTED_TTL)))?;
}
Ok(true)
}
}

71
litellm-rust/crates/redis-store/src/router_cooldowns/types.rs Normal file
View file

@ -0,0 +1,71 @@
use std::time::Duration;
use litellm_python_compat::{Value, repr::repr};
use crate::{DecodeError, FromValue, IntoValue, Number, codec::python_value};
#[derive(Clone, Debug, PartialEq)]
pub struct Cooldown {
pub exception_received: String,
pub status_code: String,
pub timestamp: Number,
pub cooldown_time: Number,
}
impl Cooldown {
pub(super) fn remaining(&self, now: Duration) -> f64 {
self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64()
}
/// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`.
pub fn encode(&self) -> String {
repr(&self.clone().into_value())
}
/// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does.
pub fn decode(raw: &[u8]) -> Result<Self, DecodeError> {
Self::from_value(&python_value(raw)?)
}
}
/// The `CooldownCacheValue` dict the router reads and writes.
impl IntoValue for Cooldown {
fn into_value(self) -> Value {
let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value);
Value::Dict(vec![
field("exception_received", Value::Str(self.exception_received)),
field("status_code", Value::Str(self.status_code)),
field("timestamp", self.timestamp.to_value()),
field("cooldown_time", self.cooldown_time.to_value()),
])
}
}
impl FromValue for Cooldown {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
let Value::Dict(fields) = value else {
return Err(DecodeError::Shape("a cooldown is a dict"));
};
let field = |name: &str| {
fields
.iter()
.find(|(key, _)| matches!(key, Value::Str(key) if key == name))
.map(|(_, value)| value)
};
let text = |name: &'static str| match field(name) {
Some(Value::Str(text)) => Ok(text.clone()),
_ => Err(DecodeError::Shape(name)),
};
let number = |name: &'static str| {
field(name)
.and_then(Number::from_value)
.ok_or(DecodeError::Shape(name))
};
Ok(Self {
exception_received: text("exception_received")?,
status_code: text("status_code")?,
timestamp: number("timestamp")?,
cooldown_time: number("cooldown_time")?,
})
}
}

222
litellm-rust/crates/redis-store/src/router_usage.rs
View file

@ -1,222 +0,0 @@
use std::time::Duration;
use litellm_cache::Error as MemoryError;
use litellm_cache_memory::InMemoryCache;
use litellm_redis_client::{Error, Expiry, RedisClient, RedisError, RedisErrorKind, Value};
use crate::{
DecodeError, Number,
codec::python_value,
tier::{Clock, MemoryTier, OnBreakerOpen, TierPolicy, system_clock},
};
/// `main`'s usage tier: `DualCache`'s 10 s default refresh, `InMemoryCache`'s 600 s default for
/// backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(10),
backfill_ttl: Duration::from_secs(600),
on_breaker_open: OnBreakerOpen::ServeMemory,
tracked_reads: 1000,
};
const MEMORY_ENTRIES: usize = 200;
/// `RoutingArgs.ttl`, the window of a per-minute counter.
pub const COUNTER_TTL: Duration = Duration::from_secs(60);
/// A counter's key segment, which is also how the Python router names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "tpm")]
Tpm,
#[strum(serialize = "rpm")]
Rpm,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Counter {
pub deployment_id: String,
pub model: Option<String>,
pub minute: String,
pub kind: Kind,
}
impl Counter {
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
fn key(&self) -> String {
let kind: &str = self.kind.into();
let model = self.model.as_deref().unwrap_or("None");
format!(
"global_router:{}:{model}:{kind}:{}",
self.deployment_id, self.minute
)
}
}
fn decode(raw: &[u8]) -> Result<Number, DecodeError> {
Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a usage counter is a number"))
}
fn whole_seconds(ttl: Duration) -> Duration {
Duration::from_secs(ttl.as_secs())
}
/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod.
pub struct RouterUsage(MemoryTier<InMemoryCache<Number>, RedisClient>);
impl RouterUsage {
pub fn new(client: Option<RedisClient>, policy: TierPolicy) -> Self {
Self::with_clock(client, policy, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, policy: TierPolicy, now: Clock) -> Self {
let clock = now.clone();
let memory =
InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || {
clock()
});
Self(MemoryTier::read_through(memory, client, policy, now))
}
/// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL.
/// A Redis failure answers with the memory count.
pub async fn increment(
&self,
counter: &Counter,
value: Number,
ttl: Duration,
) -> Result<Number, MemoryError> {
let key = counter.key();
let local = self.add_local(&key, value, ttl)?;
let Some(remote) = self.0.remote() else {
return Ok(local);
};
match remote
.incr_by_float(&key, value.into(), Expiry::KeepExisting(whole_seconds(ttl)))
.await
{
Ok(total) => Ok(Number::Float(total)),
Err(error) => {
tracing::warn!(%error, "Redis usage increment failed, falling back to the in-memory count");
Ok(local)
}
}
}
/// One deployment's tokens and requests, `async_increment_pipeline` on its own round trip:
/// memory, then `INCRBYFLOAT` and `EXPIRE 60` for each counter, mirrored back into memory as
/// `int` totals.
pub async fn increment_deployment(
&self,
deployment_id: &str,
model: Option<&str>,
minute: &str,
tokens: Number,
requests: i64,
) -> Result<(), MemoryError> {
let increment = Increment::new(deployment_id, model, minute, tokens, requests);
let local = increment
.counts
.iter()
.map(|(key, value)| self.add_local(key, *value, COUNTER_TTL))
.collect::<Result<Vec<_>, _>>()?;
let Some(remote) = self.0.remote() else {
return Ok(());
};
let pipeline = increment
.counts
.iter()
.fold(remote.pipeline(), |pipeline, (key, value)| {
pipeline.incr_by_float(key, (*value).into(), Some(COUNTER_TTL))
});
let totals = remote
.run(&pipeline)
.await
.into_iter()
.map(|total| total.and_then(float_reply))
.collect();
let keys: Vec<String> = increment.counts.into_iter().map(|(key, _)| key).collect();
self.mirror(&keys, local, totals)
}
pub async fn get_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, MemoryError> {
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
self.0.get_many(&keys, decode).await
}
pub fn flush_local_counters(&self) -> Result<(), MemoryError> {
self.0.flush_memory()
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, MemoryError> {
self.0.memory().entries()
}
fn mirror(
&self,
keys: &[String],
local: Vec<Number>,
totals: Result<Vec<f64>, Error>,
) -> Result<(), MemoryError> {
let totals = match totals {
Ok(totals) => totals.into_iter().map(Number::Float).collect(),
Err(error) => {
tracing::warn!(%error, "Redis usage pipeline failed, keeping the in-memory counts");
local
}
};
keys.iter().zip(totals).try_for_each(|(key, total)| {
self.0
.set_memory(key, Number::Int(total.truncate()), COUNTER_TTL)
})
}
fn add_local(&self, key: &str, value: Number, ttl: Duration) -> Result<Number, MemoryError> {
self.0.update_memory(key, ttl, |current| {
current.unwrap_or(Number::Int(0)) + value
})
}
}
struct Increment {
counts: [(String, Number); 2],
}
impl Increment {
fn new(
deployment_id: &str,
model: Option<&str>,
minute: &str,
tokens: Number,
requests: i64,
) -> Self {
let counter = |kind| Counter {
deployment_id: deployment_id.to_owned(),
model: model.map(str::to_owned),
minute: minute.to_owned(),
kind,
};
Self {
counts: [
(counter(Kind::Tpm).key(), tokens),
(counter(Kind::Rpm).key(), Number::Int(requests)),
],
}
}
}
fn float_reply(value: Value) -> Result<f64, Error> {
let total = match &value {
Value::BulkString(bytes) => std::str::from_utf8(bytes)
.ok()
.and_then(|text| text.parse().ok()),
Value::Double(total) => Some(*total),
_ => None,
};
total.ok_or_else(|| {
Error::Decode(RedisError::from((
RedisErrorKind::UnexpectedReturnType,
"INCRBYFLOAT did not answer a number",
)))
})
}

28
litellm-rust/crates/redis-store/src/router_usage/keys.rs Normal file
View file

@ -0,0 +1,28 @@
/// A counter's key segment, which is also how the Python router names it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "tpm")]
Tpm,
#[strum(serialize = "rpm")]
Rpm,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Counter {
pub deployment_id: String,
pub model: Option<String>,
pub minute: String,
pub kind: Kind,
}
impl Counter {
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
pub(super) fn key(&self) -> String {
let kind: &str = self.kind.into();
let model = self.model.as_deref().unwrap_or("None");
format!(
"global_router:{}:{model}:{kind}:{}",
self.deployment_id, self.minute
)
}
}

139
litellm-rust/crates/redis-store/src/router_usage/mod.rs Normal file
View file

@ -0,0 +1,139 @@
mod keys;
mod types;
use std::time::Duration;
use litellm_redis_client::{Declared, Expiry, RedisBatch, commands};
pub use self::keys::{Counter, Kind};
use crate::{
Error, Number,
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
tier::{MemoryTier, OnBreakerOpen, TierPolicy},
};
/// `main`'s usage tier: `DualCache`'s 10 s default refresh, `InMemoryCache`'s 600 s default for
/// backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(10),
backfill_ttl: Duration::from_secs(600),
on_breaker_open: OnBreakerOpen::ServeMemory,
tracked_reads: 1000,
};
/// `RoutingArgs.ttl`, the window of a per-minute counter.
pub const COUNTER_TTL: Duration = Duration::from_secs(60);
/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod.
pub struct RouterUsage {
usage: MemoryTier<Number>,
remote: RemoteSlot,
}
impl RouterUsage {
pub fn new(remote: Option<Remote>, policy: TierPolicy) -> Self {
Self::with_clock(remote, policy, system_clock())
}
pub fn with_clock(remote: Option<Remote>, policy: TierPolicy, now: Clock) -> Self {
Self {
usage: MemoryTier::read_through(policy, now),
remote: RemoteSlot::new(remote),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.remote.attach(remote)
}
/// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL.
/// A Redis failure answers with the memory count.
pub async fn increment(&self, counter: &Counter, value: Number, ttl: Duration) -> Result<Number, Error> {
let key = counter.key();
let local = self.add_local(&key, value, ttl)?;
let Some(client) = self.remote.client() else {
return Ok(local);
};
let increment = commands::incr_by_float(&key, value.into(), Some(Expiry::KeepExisting(ttl)));
match client.send(increment).await {
Ok(total) => Ok(Number::Float(total)),
Err(error) => {
tracing::warn!(%error, "Redis usage increment failed, falling back to the in-memory count");
Ok(local)
}
}
}
/// One deployment's tokens and requests, `async_increment_pipeline` on its own round trip:
/// memory, then `INCRBYFLOAT` and `EXPIRE 60` for each counter, mirrored back into memory as
/// `int` totals.
pub async fn increment_deployment(
&self,
deployment_id: &str,
model: Option<&str>,
minute: &str,
tokens: Number,
requests: i64,
) -> Result<(), Error> {
let counter = |kind| Counter {
deployment_id: deployment_id.to_owned(),
model: model.map(str::to_owned),
minute: minute.to_owned(),
kind,
};
let counts = [
(counter(Kind::Tpm).key(), tokens),
(counter(Kind::Rpm).key(), Number::Int(requests)),
];
let local = counts
.iter()
.map(|(key, value)| self.add_local(key, *value, COUNTER_TTL))
.collect::<Result<Vec<_>, _>>()?;
let Some(client) = self.remote.client() else {
return Ok(());
};
let pipeline = RedisBatch::pipelined(client.clone());
let totals = counts.iter().map(|(key, value)| {
pipeline.add(commands::incr_by_float(key, (*value).into(), Some(Expiry::Refresh(COUNTER_TTL))))
});
let totals = Declared::all(totals.collect::<Vec<_>>())
.await
.into_iter()
.collect::<Result<Vec<_>, _>>();
let totals = match totals {
Ok(totals) => totals.into_iter().map(Number::Float).collect(),
Err(error) => {
tracing::warn!(%error, "Redis usage pipeline failed, keeping the in-memory counts");
local
}
};
counts.iter().zip(totals).try_for_each(|((key, _), total)| {
self.usage
.memory()
.set(key, Number::Int(total.truncate()), Some(COUNTER_TTL))
})?;
Ok(())
}
pub async fn get_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, Error> {
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
Ok(self.usage.get_many(&self.remote, &keys, types::decode).await?)
}
pub fn flush_local_counters(&self) -> Result<(), Error> {
Ok(self.usage.memory().flush()?)
}
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
/// in-memory router state.
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
Ok(self.usage.memory().entries()?)
}
fn add_local(&self, key: &str, value: Number, ttl: Duration) -> Result<Number, Error> {
Ok(self.usage.memory().update(key, Some(ttl), |current| {
Ok(current.unwrap_or(Number::Int(0)) + value)
})?)
}
}

24
litellm-rust/crates/redis-store/src/router_usage/types.rs Normal file
View file

@ -0,0 +1,24 @@
use litellm_python_compat::Value;
use super::keys::Counter;
use crate::{DecodeError, FromValue, Number, codec, value};
/// `(deployment_id, model, minute, kind)`, the kind spelled as its key segment.
impl FromValue for Counter {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
const SHAPE: &str = "a usage counter is (deployment_id, model, minute, \"tpm\" | \"rpm\")";
let [deployment_id, model, minute, kind] = value::fields(value, SHAPE)?;
Ok(Self {
deployment_id: value::text(deployment_id, SHAPE)?,
model: value::optional_text(model, SHAPE)?,
minute: value::text(minute, SHAPE)?,
kind: value::text(kind, SHAPE)?
.parse()
.map_err(|_| DecodeError::Shape(SHAPE))?,
})
}
}
pub(super) fn decode(raw: &[u8]) -> Result<Number, DecodeError> {
codec::number(raw, "a usage counter is a number")
}

575
litellm-rust/crates/redis-store/src/spend_counters.rs
View file

@ -1,575 +0,0 @@
use std::{
sync::{Arc, Mutex, MutexGuard, PoisonError},
time::Duration,
};
use litellm_cache::Error as MemoryError;
use litellm_cache_memory::InMemoryCache;
use litellm_python_compat::{json, repr::to_str};
use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient};
use crate::tier::{Clock, system_clock};
use crate::{DecodeError, ErrorKind, Number, RequestRedis, codec::python_value};
/// `InMemoryCache`'s defaults, which `spend_counter_cache` runs its memory tier with.
const MEMORY_ENTRIES: usize = 200;
pub const MEMORY_TTL: Duration = Duration::from_secs(600);
/// `spend_counter_cache`'s `default_in_memory_ttl`, which only its memory-only increment applies.
pub const LOCAL_INCREMENT_TTL: Duration = Duration::from_secs(60);
/// `RedisCache`'s TTL when `litellm.default_redis_ttl` is unset.
pub const REDIS_TTL: Duration = Duration::from_secs(60);
/// `RedisCache.async_set_max`, byte for byte, sent as `EVAL` like it: a counter only ever rises to
/// the authoritative spend.
const SET_MAX: &str = concat!(
"local cur = redis.call('GET', KEYS[1]) ",
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]) ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end ",
"return ARGV[1] end ",
"return cur"
);
/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, the chunk `DualCache.async_delete_cache_keys` deletes at once.
const DELETE_CHUNK: usize = 1000;
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SpendCounter {
Key {
token: String,
window: Option<String>,
},
Team {
team_id: String,
window: Option<String>,
},
TeamMember {
user_id: String,
team_id: String,
},
User(String),
EndUser(String),
Org(String),
Tag(String),
Project(String),
ModelAccessGroup(String),
ShadowEval(String),
}
#[derive(Debug, thiserror::Error)]
#[error("{0:?} is not a spend counter")]
pub struct NotASpendCounter(pub String);
impl SpendCounter {
pub fn key(&self) -> String {
let windowed = |scope: &str, id: &str, window: &Option<String>| match window {
Some(window) => format!("spend:{scope}:{id}:window:{window}"),
None => format!("spend:{scope}:{id}"),
};
match self {
Self::Key { token, window } => windowed("key", token, window),
Self::Team { team_id, window } => windowed("team", team_id, window),
Self::TeamMember { user_id, team_id } => {
format!("spend:team_member:{user_id}:{team_id}")
}
Self::User(id) => format!("spend:user:{id}"),
Self::EndUser(id) => format!("spend:end_user:{id}"),
Self::Org(id) => format!("spend:org:{id}"),
Self::Tag(name) => format!("spend:tag:{name}"),
Self::Project(id) => format!("spend:project:{id}"),
Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"),
Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"),
}
}
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
let invalid = || NotASpendCounter(key.to_owned());
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
let windowed = || match id.rsplit_once(":window:") {
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),
_ => (id.to_owned(), None),
};
Ok(match scope {
"key" => {
let (token, window) = windowed();
Self::Key { token, window }
}
"team" => {
let (team_id, window) = windowed();
Self::Team { team_id, window }
}
"team_member" => {
let (user_id, team_id) = id.rsplit_once(':').ok_or_else(invalid)?;
Self::TeamMember {
user_id: user_id.to_owned(),
team_id: team_id.to_owned(),
}
}
"user" => Self::User(id.to_owned()),
"end_user" => Self::EndUser(id.to_owned()),
"org" => Self::Org(id.to_owned()),
"tag" => Self::Tag(id.to_owned()),
"project" => Self::Project(id.to_owned()),
"model_access_group" => Self::ModelAccessGroup(id.to_owned()),
"shadow_eval" => Self::ShadowEval(id.to_owned()),
_ => return Err(invalid()),
})
}
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error(transparent)]
Redis(#[from] litellm_redis_client::Error),
#[error(transparent)]
Memory(#[from] MemoryError),
#[error(transparent)]
Decode(#[from] DecodeError),
}
impl Error {
pub fn kind(&self) -> ErrorKind {
match self {
Self::Redis(error) => error.into(),
Self::Decode(_) => ErrorKind::Decode,
Self::Memory(_) => ErrorKind::Memory,
}
}
}
/// Where a read was answered, named the way the Python store reports it. Redis is cross-pod truth,
/// including a clean miss; memory only holds this pod's writes and answers only when there is no
/// Redis or it failed.
#[derive(Clone, Copy, Debug, PartialEq, strum::IntoStaticStr)]
pub enum Reading {
#[strum(serialize = "redis")]
Redis(Option<Number>),
#[strum(serialize = "memory")]
Memory(Option<Number>),
}
impl Reading {
pub fn value(self) -> Option<Number> {
match self {
Self::Redis(value) | Self::Memory(value) => value,
}
}
}
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LostSeed {
/// `SpendCounterReseed.coalesced`: answer with the seed without writing it.
Answer,
/// `SpendCounterReseed.coalesced_window`: `INCRBYFLOAT` the seed, keeping any TTL.
Increment,
}
fn decode(raw: &[u8]) -> Result<Number, DecodeError> {
Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a spend counter is a number"))
}
fn decode_reply(raw: Option<Vec<u8>>) -> Result<Option<Number>, DecodeError> {
raw.as_deref().map(decode).transpose()
}
/// `json.dumps` of the value, as `RedisCache.async_set_cache` writes it.
fn json_bytes(value: Number) -> String {
json::dumps(&value.to_value()).unwrap_or_else(|_| to_str(&value.to_value()))
}
fn whole_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
}
/// The proxy's spend counters: Redis shared by every pod, and this pod's memory copy that answers
/// only when Redis cannot.
pub struct SpendCounters {
memory: InMemoryCache<Number>,
remote: Option<RedisClient>,
redis_ttl: Duration,
updates: Mutex<()>,
}
impl SpendCounters {
pub fn new(client: Option<RedisClient>, redis_ttl: Duration) -> Self {
Self::with_clock(client, redis_ttl, system_clock())
}
pub fn with_clock(client: Option<RedisClient>, redis_ttl: Duration, now: Clock) -> Self {
Self {
memory: InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(MEMORY_TTL), move || {
now()
}),
remote: client,
redis_ttl,
updates: Mutex::new(()),
}
}
pub fn has_redis(&self) -> bool {
self.remote.is_some()
}
pub fn redis_ttl(&self) -> Duration {
self.redis_ttl
}
/// `read_spend_counter_cache_value`: Redis first, memory when Redis is absent or failed.
pub async fn read(&self, counter: &SpendCounter) -> Result<Reading, MemoryError> {
if let Some(remote) = &self.remote {
match remote
.get::<Vec<u8>>(&counter.key())
.await
.map_err(Error::from)
.and_then(|raw| decode_reply(raw).map_err(Error::from))
{
Ok(value) => return Ok(Reading::Redis(value)),
Err(error) => {
tracing::debug!(%error, "spend counter Redis read failed, falling back to memory");
}
}
}
Ok(Reading::Memory(self.get_local(counter)?))
}
/// One round trip for many counters; `None` when there is no Redis or the read failed, so each
/// caller reads on its own.
pub async fn read_many(&self, counters: &[SpendCounter]) -> Option<Vec<Option<Number>>> {
let declared = self.declare_read_many(RedisClient::batch, counters)?;
declared.resolve().await
}
/// `read_many` declared on the request's pre-call batch, so it goes out with the request's other
/// pre-call reads.
pub fn declare_read_many_pre_call(
&self,
request: &RequestRedis,
counters: &[SpendCounter],
) -> Option<DeclaredReads> {
self.declare_read_many(|client| request.pre(client), counters)
}
fn declare_read_many(
&self,
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
counters: &[SpendCounter],
) -> Option<DeclaredReads> {
let batch = batch_for(self.remote.as_ref()?);
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
Some(DeclaredReads(batch.mget::<Vec<u8>>(&keys)))
}
/// `increment_spend_counter_cache`: `INCRBYFLOAT` refreshing the TTL, mirrored into memory. A
/// failed increment may or may not have landed, so the counter is dropped everywhere before the
/// error returns. Without Redis it is a memory increment.
pub async fn increment(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
let Some(remote) = &self.remote else {
return Ok(self.increment_local(counter, amount)?);
};
let key = counter.key();
match remote
.incr_by_float(&key, amount.into(), Expiry::Refresh(self.redis_ttl))
.await
{
Ok(total) => {
let total = Number::Float(total);
self.memory.set_cache(key, total, None)?;
Ok(total)
}
Err(error) => {
self.invalidate(std::slice::from_ref(counter)).await;
Err(error.into())
}
}
}
/// `run_spend_counter_pipeline`: every increment in one round trip, totals in order. The caller
/// decides what a failure means for counters whose increment may or may not have landed.
pub async fn increment_many(
&self,
increments: &[(SpendCounter, Number)],
) -> Result<Vec<Number>, Error> {
let Some(remote) = &self.remote else {
return increments
.iter()
.map(|(counter, amount)| Ok(self.increment_local(counter, *amount)?))
.collect();
};
let pipeline = increments
.iter()
.fold(remote.pipeline(), |pipeline, (counter, amount)| {
pipeline.incr_by_float(&counter.key(), (*amount).into(), Some(self.redis_ttl))
});
let replies = remote.run(&pipeline).await;
let mut totals = Vec::with_capacity(replies.len());
for ((counter, _), reply) in increments.iter().zip(replies) {
let total = Number::Float(float_reply(reply?)?);
self.memory.set_cache(counter.key(), total, None)?;
totals.push(total);
}
Ok(totals)
}
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
pub fn increment_many_post_call(
self: &Arc<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
) -> bool {
let Some(remote) = &self.remote else {
return false;
};
let declared = self.declare_increments(&request.post(remote), increments);
let counters = increments
.iter()
.map(|(counter, _)| counter.clone())
.collect::<Vec<_>>();
let store = Arc::clone(self);
request.spawn(async move {
for ((key, total), counter) in declared.into_iter().zip(counters) {
match total.settled().await {
Ok(total) => {
if let Err(error) = store.memory.set_cache(key, Number::Float(total), None)
{
tracing::warn!(%error, "mirroring a spend counter total into memory failed");
}
}
Err(error) => {
tracing::warn!(%error, %key, "spend counter increment did not land in the post-call pipeline; invalidating it");
store.invalidate(std::slice::from_ref(&counter)).await;
}
}
}
});
true
}
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
/// `None` without Redis, else what Redis holds afterwards.
pub async fn set_max(
&self,
counter: &SpendCounter,
spend: Number,
) -> Result<Option<Number>, Error> {
let key = counter.key();
{
let _guard = lock(&self.updates);
let raise = self
.memory
.get_cache(&key)?
.is_none_or(|cached| cached.as_f64() < spend.as_f64());
if raise {
self.memory.set_cache(key.clone(), spend, None)?;
}
}
let Some(remote) = &self.remote else {
return Ok(None);
};
let stored: Option<Vec<u8>> = remote
.eval_source(
SET_MAX,
&[&key],
(
to_str(&spend.to_value()),
whole_seconds(self.redis_ttl).to_string(),
),
)
.await?;
Ok(stored
.as_deref()
.map(|raw| {
std::str::from_utf8(raw)
.ok()
.and_then(|text| text.trim().parse::<f64>().ok())
.map(Number::Float)
.ok_or(DecodeError::Shape("set_max returned a non-number"))
})
.transpose()?)
}
/// `SpendCounterReseed.coalesced(_window)`: seed a cold counter with `SET NX` so one pod wins,
/// answer with whatever the counter holds, and mirror it into memory. Without Redis, memory
/// keeps the larger of the seed and what it held.
pub async fn seed(
&self,
counter: &SpendCounter,
spend: Number,
lost: LostSeed,
) -> Result<Number, Error> {
let key = counter.key();
let Some(remote) = &self.remote else {
let _guard = lock(&self.updates);
let seeded = match self.memory.get_cache(&key)? {
Some(cached) if cached.as_f64() > spend.as_f64() => Number::Float(cached.as_f64()),
_ => spend,
};
self.memory.set_cache(key, seeded, None)?;
return Ok(seeded);
};
let current = if remote
.set_nx(&key, json_bytes(spend), self.redis_ttl)
.await?
{
spend
} else {
match decode_reply(remote.get::<Vec<u8>>(&key).await?)? {
Some(winner) => Number::Float(winner.as_f64()),
None => match lost {
LostSeed::Answer => spend,
LostSeed::Increment => Number::Float(
remote
.incr_by_float(&key, spend.into(), Expiry::KeepExisting(self.redis_ttl))
.await?,
),
},
}
};
self.memory.set_cache(key, current, None)?;
Ok(current)
}
/// A spend reset: memory, then Redis, each for `ttl` or its own default.
pub async fn set(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), Error> {
let key = counter.key();
self.memory.set_cache(key.clone(), spend, ttl)?;
let Some(remote) = &self.remote else {
return Ok(());
};
remote
.set(&key, json_bytes(spend), ttl.unwrap_or(self.redis_ttl))
.await?;
Ok(())
}
/// Drop counters from memory, then from Redis in one round trip.
pub async fn delete(&self, counters: &[SpendCounter]) -> Result<(), Error> {
for counter in counters {
self.memory.delete_cache(&counter.key())?;
}
let Some(remote) = &self.remote else {
return Ok(());
};
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
for chunk in keys.chunks(DELETE_CHUNK) {
let chunk: Vec<&str> = chunk.iter().map(String::as_str).collect();
remote.del(&chunk).await?;
}
Ok(())
}
/// `refresh_spend_counter_ttl`: `false` without Redis or for a counter Redis no longer holds.
pub async fn refresh_ttl(&self, counter: &SpendCounter) -> Result<bool, Error> {
let Some(remote) = &self.remote else {
return Ok(false);
};
Ok(remote.expire(&counter.key(), self.redis_ttl).await?)
}
pub fn get_local(&self, counter: &SpendCounter) -> Result<Option<Number>, MemoryError> {
self.memory.get_cache(&counter.key())
}
pub fn set_local(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), MemoryError> {
self.memory.set_cache(counter.key(), spend, ttl).map(drop)
}
pub fn delete_local(&self, counter: &SpendCounter) -> Result<(), MemoryError> {
self.memory.delete_cache(&counter.key())
}
/// `DualCache.async_increment_cache(local_only=True)`.
pub fn increment_local(
&self,
counter: &SpendCounter,
amount: Number,
) -> Result<Number, MemoryError> {
let key = counter.key();
let _guard = lock(&self.updates);
let total = self.memory.get_cache(&key)?.unwrap_or(Number::Int(0)) + amount;
self.memory
.set_cache(key, total, Some(LOCAL_INCREMENT_TTL))?;
Ok(total)
}
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, MemoryError> {
self.memory.entries()
}
pub fn flush_local(&self) -> Result<(), MemoryError> {
self.memory.flush_cache()
}
/// `invalidate_spend_counter`: memory always, Redis best effort.
pub async fn invalidate(&self, counters: &[SpendCounter]) {
if let Err(error) = self.delete(counters).await {
tracing::debug!(%error, "dropping spend counters from Redis failed");
}
}
fn declare_increments(
&self,
batch: &RedisBatch,
increments: &[(SpendCounter, Number)],
) -> Vec<(String, Declared<f64>)> {
increments
.iter()
.map(|(counter, amount)| {
let key = counter.key();
let total = batch.incr_by_float(&key, (*amount).into(), Some(self.redis_ttl));
(key, total)
})
.collect()
}
}
/// One `MGET` declared on a batch, answered once it goes out.
pub struct DeclaredReads(Declared<Vec<Option<Vec<u8>>>>);
impl DeclaredReads {
/// Flushes the batch if nobody has yet. `None` when the read or a decode failed.
pub async fn resolve(self) -> Option<Vec<Option<Number>>> {
self.0
.resolve()
.await
.map_err(Error::from)
.and_then(|replies| {
replies
.into_iter()
.map(|raw| decode_reply(raw).map_err(Error::from))
.collect::<Result<Vec<_>, _>>()
})
.inspect_err(|error| {
tracing::debug!(%error, "spend counter batch read failed, falling back to per-key reads");
})
.ok()
}
}
fn lock(mutex: &Mutex<()>) -> MutexGuard<'_, ()> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}
fn float_reply(value: litellm_redis_client::Value) -> Result<f64, Error> {
let total = match &value {
litellm_redis_client::Value::BulkString(bytes) => std::str::from_utf8(bytes)
.ok()
.and_then(|text| text.parse().ok()),
litellm_redis_client::Value::Double(total) => Some(*total),
_ => None,
};
Ok(total.ok_or(DecodeError::Shape("INCRBYFLOAT did not answer a number"))?)
}

85
litellm-rust/crates/redis-store/src/spend_counters/keys.rs Normal file
View file

@ -0,0 +1,85 @@
/// One budget's spend counter. `parse` and `key` round-trip `main`'s `spend:*` keys exactly.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum SpendCounter {
Key {
token: String,
window: Option<String>,
},
Team {
team_id: String,
window: Option<String>,
},
TeamMember {
user_id: String,
team_id: String,
},
User(String),
EndUser(String),
Org(String),
Tag(String),
Project(String),
ModelAccessGroup(String),
ShadowEval(String),
}
#[derive(Debug, thiserror::Error)]
#[error("{0:?} is not a spend counter")]
pub struct NotASpendCounter(pub String);
impl SpendCounter {
pub fn key(&self) -> String {
let windowed = |scope: &str, id: &str, window: &Option<String>| match window {
Some(window) => format!("spend:{scope}:{id}:window:{window}"),
None => format!("spend:{scope}:{id}"),
};
match self {
Self::Key { token, window } => windowed("key", token, window),
Self::Team { team_id, window } => windowed("team", team_id, window),
Self::TeamMember { user_id, team_id } => {
format!("spend:team_member:{user_id}:{team_id}")
}
Self::User(id) => format!("spend:user:{id}"),
Self::EndUser(id) => format!("spend:end_user:{id}"),
Self::Org(id) => format!("spend:org:{id}"),
Self::Tag(name) => format!("spend:tag:{name}"),
Self::Project(id) => format!("spend:project:{id}"),
Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"),
Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"),
}
}
pub fn parse(key: &str) -> Result<Self, NotASpendCounter> {
let invalid = || NotASpendCounter(key.to_owned());
let rest = key.strip_prefix("spend:").ok_or_else(invalid)?;
let (scope, id) = rest.split_once(':').ok_or_else(invalid)?;
let windowed = || match id.rsplit_once(":window:") {
Some((id, window)) if !window.is_empty() => (id.to_owned(), Some(window.to_owned())),
_ => (id.to_owned(), None),
};
Ok(match scope {
"key" => {
let (token, window) = windowed();
Self::Key { token, window }
}
"team" => {
let (team_id, window) = windowed();
Self::Team { team_id, window }
}
"team_member" => {
let (user_id, team_id) = id.rsplit_once(':').ok_or_else(invalid)?;
Self::TeamMember {
user_id: user_id.to_owned(),
team_id: team_id.to_owned(),
}
}
"user" => Self::User(id.to_owned()),
"end_user" => Self::EndUser(id.to_owned()),
"org" => Self::Org(id.to_owned()),
"tag" => Self::Tag(id.to_owned()),
"project" => Self::Project(id.to_owned()),
"model_access_group" => Self::ModelAccessGroup(id.to_owned()),
"shadow_eval" => Self::ShadowEval(id.to_owned()),
_ => return Err(invalid()),
})
}
}

308
litellm-rust/crates/redis-store/src/spend_counters/mod.rs Normal file
View file

@ -0,0 +1,308 @@
mod keys;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_python_compat::repr::to_str;
use litellm_redis_client::{Declared, Expiry, RedisBatch, commands};
pub use self::keys::{NotASpendCounter, SpendCounter};
pub use self::types::{LostSeed, Reading};
use crate::{
Error, Number, Pending, RequestRedis, Send,
memory::{Clock, Memory, system_clock},
remote::{self, Remote, RemoteSlot},
};
/// `spend_counter_cache`'s `default_in_memory_ttl`, which only its memory-only increment applies.
pub const LOCAL_INCREMENT_TTL: Duration = Duration::from_secs(60);
/// `RedisCache.async_set_max`, byte for byte, sent as `EVAL` like it: a counter only ever rises to
/// the authoritative spend.
const SET_MAX: &str = concat!(
"local cur = redis.call('GET', KEYS[1]) ",
"if cur == false or tonumber(cur) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]) ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end ",
"return ARGV[1] end ",
"return cur"
);
/// The proxy's spend counters: Redis shared by every pod, and this pod's memory copy that answers
/// only when Redis cannot.
pub struct SpendCounters {
memory: Memory<Number>,
remote: RemoteSlot,
}
impl SpendCounters {
pub fn new(remote: Option<Remote>) -> Self {
Self::with_clock(remote, system_clock())
}
pub fn with_clock(remote: Option<Remote>, now: Clock) -> Self {
Self {
memory: Memory::new(None, &now),
remote: RemoteSlot::new(remote),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.remote.attach(remote)
}
pub fn has_redis(&self) -> bool {
self.remote.get().is_some()
}
/// `read_spend_counter_cache_value`: Redis first, memory when Redis is absent or failed.
pub async fn read(&self, counter: &SpendCounter) -> Result<Reading, Error> {
if let Some(client) = self.remote.client() {
let read = client.send(commands::get::<Vec<u8>>(&counter.key())).await;
match read.map_err(Error::from).and_then(|raw| Ok(types::decode(raw)?)) {
Ok(value) => return Ok(Reading::Redis(value)),
Err(error) => {
tracing::debug!(%error, "spend counter Redis read failed, falling back to memory");
}
}
}
Ok(Reading::Memory(self.get_local(counter)?))
}
/// One `MGET` for many counters on the batch `send` picks; `None` without Redis. The read
/// resolves to `None` when it failed, so each caller reads on its own.
pub fn read_many(
&self,
send: Send<'_>,
counters: &[SpendCounter],
) -> Option<Pending<Option<Vec<Option<Number>>>>> {
let batch = self.remote.batch(send)?;
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
let keys: Vec<&str> = keys.iter().map(String::as_str).collect();
let read = batch.add(commands::mget::<Vec<u8>>(&keys));
Some(Box::pin(async move {
let decoded = read.resolve().await.map_err(Error::from).and_then(|replies| {
replies
.into_iter()
.map(|raw| Ok(types::decode(raw)?))
.collect::<Result<Vec<_>, Error>>()
});
Ok(decoded
.inspect_err(|error| {
tracing::debug!(%error, "spend counter batch read failed, falling back to per-key reads");
})
.ok())
}))
}
/// `increment_spend_counter_cache`: `INCRBYFLOAT` refreshing the TTL, mirrored into memory. A
/// failed increment may or may not have landed, so the counter is dropped everywhere before the
/// error returns. Without Redis it is a memory increment.
pub async fn increment(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
let Some(remote) = self.remote.get() else {
return self.increment_local(counter, amount);
};
let key = counter.key();
let increment = commands::incr_by_float(&key, amount.into(), Some(Expiry::Refresh(remote.default_ttl)));
match remote.client.send(increment).await {
Ok(total) => {
let total = Number::Float(total);
self.memory.set(&key, total, None)?;
Ok(total)
}
Err(error) => {
self.invalidate(std::slice::from_ref(counter)).await;
Err(error.into())
}
}
}
/// `run_spend_counter_pipeline`: every increment in one round trip, totals in order. The caller
/// decides what a failure means for counters whose increment may or may not have landed.
pub async fn increment_many(&self, increments: &[(SpendCounter, Number)]) -> Result<Vec<Number>, Error> {
let Some(remote) = self.remote.get() else {
return increments
.iter()
.map(|(counter, amount)| self.increment_local(counter, *amount))
.collect();
};
let pipeline = RedisBatch::pipelined(remote.client.clone());
let totals = increments.iter().map(|(counter, amount)| {
pipeline.add(commands::incr_by_float(
&counter.key(),
(*amount).into(),
Some(Expiry::Refresh(remote.default_ttl)),
))
});
let totals = Declared::all(totals.collect::<Vec<_>>()).await;
increments
.iter()
.zip(totals)
.map(|((counter, _), total)| {
let total = Number::Float(total?);
self.memory.set(&counter.key(), total, None)?;
Ok(total)
})
.collect()
}
/// `increment_many` on the request's post-call batch. Each total reaches memory once the batch
/// goes out, and a counter whose increment failed is dropped everywhere. `false` without Redis.
pub fn increment_many_post_call(
self: &Arc<Self>,
request: &RequestRedis,
increments: &[(SpendCounter, Number)],
) -> bool {
let Some(remote) = self.remote.get() else {
return false;
};
let batch = Send::PostCall(request).batch(&remote.client);
let declared = increments
.iter()
.map(|(counter, amount)| {
let total = batch.add(commands::incr_by_float(
&counter.key(),
(*amount).into(),
Some(Expiry::KeepExisting(remote.default_ttl)),
));
(counter.clone(), total)
})
.collect::<Vec<_>>();
let store = Arc::clone(self);
request.spawn(async move {
for (counter, total) in declared {
match total.settled().await {
Ok(total) => {
if let Err(error) = store.memory.set(&counter.key(), Number::Float(total), None) {
tracing::warn!(%error, "mirroring a spend counter total into memory failed");
}
}
Err(error) => {
let key = counter.key();
tracing::warn!(%error, %key, "spend counter increment did not land in the post-call pipeline; invalidating it");
store.invalidate(std::slice::from_ref(&counter)).await;
}
}
}
});
true
}
/// `_repair_stale_spend_counter`: raise memory and Redis to `spend`, never lowering either.
/// `None` without Redis, else what Redis holds afterwards.
pub async fn set_max(&self, counter: &SpendCounter, spend: Number) -> Result<Option<Number>, Error> {
let key = counter.key();
{
let _guard = self.memory.guard();
if self.memory.get(&key)?.is_none_or(|cached| cached.as_f64() < spend.as_f64()) {
self.memory.set(&key, spend, None)?;
}
}
let Some(remote) = self.remote.get() else {
return Ok(None);
};
let args = (to_str(&spend.to_value()), remote.default_ttl.as_secs().to_string());
let stored: Option<f64> = remote.client.send(commands::eval_source(SET_MAX, &[&key], args)).await?;
Ok(stored.map(Number::Float))
}
/// `SpendCounterReseed.coalesced(_window)`: seed a cold counter with `SET NX` so one pod wins,
/// answer with whatever the counter holds, and mirror it into memory. Without Redis, memory
/// keeps the larger of the seed and what it held.
pub async fn seed(&self, counter: &SpendCounter, spend: Number, lost: LostSeed) -> Result<Number, Error> {
let key = counter.key();
let Some(Remote { client, default_ttl }) = self.remote.get() else {
return Ok(self.memory.update(&key, None, |cached| {
Ok(match cached {
Some(cached) if cached.as_f64() > spend.as_f64() => Number::Float(cached.as_f64()),
_ => spend,
})
})?);
};
let current = if client.send(commands::set_nx(&key, spend.to_json(), *default_ttl)).await? {
spend
} else {
match types::decode(client.send(commands::get::<Vec<u8>>(&key)).await?)? {
Some(winner) => Number::Float(winner.as_f64()),
None => match lost {
LostSeed::Answer => spend,
LostSeed::Increment => {
let increment = commands::incr_by_float(
&key,
spend.into(),
Some(Expiry::KeepExisting(*default_ttl)),
);
Number::Float(client.send(increment).await?)
}
},
}
};
self.memory.set(&key, current, None)?;
Ok(current)
}
/// A spend reset: memory, then Redis, each for `ttl` or its own default.
pub async fn set(&self, counter: &SpendCounter, spend: Number, ttl: Option<Duration>) -> Result<(), Error> {
let key = counter.key();
self.memory.set(&key, spend, ttl)?;
let Some(remote) = self.remote.get() else {
return Ok(());
};
let ttl = ttl.unwrap_or(remote.default_ttl);
remote.client.send(commands::set(&key, spend.to_json(), ttl)).await?;
Ok(())
}
/// Drop counters from memory, then from Redis in one round trip.
pub async fn delete(&self, counters: &[SpendCounter]) -> Result<(), Error> {
let keys: Vec<String> = counters.iter().map(SpendCounter::key).collect();
keys.iter().try_for_each(|key| self.memory.delete(key))?;
match self.remote.client() {
Some(client) => remote::delete(client, &keys).await,
None => Ok(()),
}
}
/// `refresh_spend_counter_ttl`: `false` without Redis or for a counter Redis no longer holds.
pub async fn refresh_ttl(&self, counter: &SpendCounter) -> Result<bool, Error> {
let Some(remote) = self.remote.get() else {
return Ok(false);
};
Ok(remote.client.send(commands::expire(&counter.key(), remote.default_ttl)).await?)
}
pub fn get_local(&self, counter: &SpendCounter) -> Result<Option<Number>, Error> {
Ok(self.memory.get(&counter.key())?)
}
pub fn set_local(&self, counter: &SpendCounter, spend: Number, ttl: Option<Duration>) -> Result<(), Error> {
Ok(self.memory.set(&counter.key(), spend, ttl)?)
}
pub fn delete_local(&self, counter: &SpendCounter) -> Result<(), Error> {
Ok(self.memory.delete(&counter.key())?)
}
/// `DualCache.async_increment_cache(local_only=True)`.
pub fn increment_local(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
Ok(self.memory.update(&counter.key(), Some(LOCAL_INCREMENT_TTL), |total| {
Ok(total.unwrap_or(Number::Int(0)) + amount)
})?)
}
pub fn flush_local(&self) -> Result<(), Error> {
Ok(self.memory.flush()?)
}
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
Ok(self.memory.entries()?)
}
/// `invalidate_spend_counter`: memory always, Redis best effort.
pub async fn invalidate(&self, counters: &[SpendCounter]) {
if let Err(error) = self.delete(counters).await {
tracing::debug!(%error, "dropping spend counters from Redis failed");
}
}
}

48
litellm-rust/crates/redis-store/src/spend_counters/types.rs Normal file
View file

@ -0,0 +1,48 @@
use litellm_python_compat::Value;
use super::keys::SpendCounter;
use crate::{DecodeError, FromValue, Number, codec};
/// The `spend:*` key Python names a counter by.
impl FromValue for SpendCounter {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
let Value::Str(key) = value else {
return Err(DecodeError::Shape("a spend counter is a spend:* key"));
};
Self::parse(key).map_err(|_| DecodeError::Shape("a spend counter is a spend:* key"))
}
}
/// Where a read was answered, named the way the Python store reports it. Redis is cross-pod truth,
/// including a clean miss; memory only holds this pod's writes and answers only when there is no
/// Redis or it failed.
#[derive(Clone, Copy, Debug, PartialEq, strum::IntoStaticStr)]
pub enum Reading {
#[strum(serialize = "redis")]
Redis(Option<Number>),
#[strum(serialize = "memory")]
Memory(Option<Number>),
}
impl Reading {
pub fn value(self) -> Option<Number> {
match self {
Self::Redis(value) | Self::Memory(value) => value,
}
}
}
/// Who wins when a seed's `SET NX` lost and the winner's value is gone by the time it is read.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LostSeed {
/// `SpendCounterReseed.coalesced`: answer with the seed without writing it.
Answer,
/// `SpendCounterReseed.coalesced_window`: `INCRBYFLOAT` the seed, keeping any TTL.
Increment,
}
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Number>, DecodeError> {
raw.as_deref()
.map(|raw| codec::number(raw, "a spend counter is a number"))
.transpose()
}

122
litellm-rust/crates/redis-store/src/tier.rs
View file

@ -1,14 +1,16 @@
use std::{
collections::{HashMap, VecDeque},
sync::{Arc, Mutex, MutexGuard},
time::{Duration, SystemTime, UNIX_EPOCH},
sync::Mutex,
time::Duration,
};
use litellm_redis_client::{Declared, Error as RedisError, RedisBatch, RedisClient};
use litellm_cache::Error;
use litellm_redis_client::{Declared, Error as RedisError, commands};
use litellm_cache::{BaseCache, DeleteCache, Error, ExactCacheContext, FlushCache};
pub type Clock = Arc<dyn Fn() -> Duration + Send + Sync>;
use crate::{
memory::{Clock, Memory, lock},
remote::RemoteSlot,
};
/// How a read-through family uses its memory tier, declared once by the family.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@ -29,102 +31,60 @@ pub enum OnBreakerOpen {
}
/// The memory tier in front of Redis that a read-through family builds for itself, as Python's
/// `DualCache` with a `RedisCache` does.
pub struct MemoryTier<L1, L2> {
memory: L1,
remote: Option<L2>,
/// `DualCache` with a `RedisCache` does. The family owns the Redis and passes it to each read.
pub struct MemoryTier<V: Clone> {
memory: Memory<V>,
policy: TierPolicy,
reads: Mutex<RecentReads>,
updates: Mutex<()>,
now: Clock,
}
pub fn system_clock() -> Clock {
Arc::new(|| {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
})
}
impl<V, L1> MemoryTier<L1, RedisClient>
where
V: Clone + Send + Sync + 'static,
L1: BaseCache<Value = V, Context = ExactCacheContext> + DeleteCache + FlushCache,
{
pub fn read_through(
memory: L1,
remote: Option<RedisClient>,
policy: TierPolicy,
now: Clock,
) -> Self {
impl<V: Clone + Send + Sync + 'static> MemoryTier<V> {
pub fn read_through(policy: TierPolicy, now: Clock) -> Self {
Self {
memory,
remote,
memory: Memory::new(None, &now),
reads: Mutex::new(RecentReads::new(policy.tracked_reads)),
policy,
updates: Mutex::new(()),
now,
}
}
pub fn memory(&self) -> &L1 {
pub fn memory(&self) -> &Memory<V> {
&self.memory
}
pub fn remote(&self) -> Option<&RedisClient> {
self.remote.as_ref()
}
pub fn now(&self) -> Duration {
(self.now)()
}
pub fn set_memory(&self, key: &str, value: V, ttl: Duration) -> Result<(), Error> {
self.memory
.set_cache(key, value, &ExactCacheContext { ttl: Some(ttl) })
}
/// Read-modify-write of one memory entry, atomic against every other update of this tier.
pub fn update_memory(
&self,
key: &str,
ttl: Duration,
update: impl FnOnce(Option<V>) -> V,
) -> Result<V, Error> {
let _guard = lock(&self.updates)?;
let value = update(self.memory.get_cache(key, &ExactCacheContext::default())?);
self.set_memory(key, value.clone(), ttl)?;
Ok(value)
}
/// Python's `DualCache.async_batch_get_cache`: memory first, then one round trip for the misses
/// whose refresh is due, written back to memory. A failed Redis read serves memory, as
/// `RedisCache` swallows it, and only a breaker refusal releases the refresh it reserved.
pub async fn get_many<E: std::fmt::Display>(
&self,
remote: &RemoteSlot,
keys: &[String],
decode: impl Fn(&[u8]) -> Result<V, E>,
) -> Result<Vec<Option<V>>, Error> {
let declared = self.declare_many(RedisClient::batch, keys)?;
let declared = self.declare_many(remote, crate::Send::Now, keys)?;
self.resolve_many(declared, decode).await
}
/// The first half of `get_many`: reserves the due refreshes and declares their GETs on the
/// batch `batch_for` picks for this tier's client, without sending anything.
/// The first half of `get_many`: reserves the due refreshes and declares their `MGET` on the
/// batch `send` picks for the family's client, without sending anything.
pub fn declare_many(
&self,
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
remote: &RemoteSlot,
send: crate::Send<'_>,
keys: &[String],
) -> Result<DeclaredRead, Error> {
let fetch = match &self.remote {
let fetch = match remote.client() {
None => None,
Some(remote) => {
Some(client) => {
let reservation = self.reserve(keys, &self.read_memory(keys)?)?;
(!reservation.keys.is_empty()).then(|| {
let batch = batch_for(remote);
let keys: Vec<&str> = reservation.keys.iter().map(String::as_str).collect();
let replies = batch.mget::<Vec<u8>>(&keys);
let replies = send.batch(client).add(commands::mget::<Vec<u8>>(&keys));
Fetch {
reservation,
replies,
@ -153,7 +113,7 @@ where
else {
return self.read_memory(&keys);
};
let fetched = settle(replies).await;
let fetched = replies.resolve().await;
let local = self.read_memory(&keys)?;
let fetched = match fetched {
Ok(fetched) => fetched,
@ -191,31 +151,21 @@ where
.collect()
}
pub fn delete_memory(&self, key: &str) -> Result<(), Error> {
self.memory.delete_cache(key)
}
pub fn flush_memory(&self) -> Result<(), Error> {
self.memory.flush_cache()
}
fn read_memory(&self, keys: &[String]) -> Result<Vec<Option<V>>, Error> {
let context = ExactCacheContext::default();
keys.iter()
.map(|key| self.memory.get_cache(key, &context))
.collect()
keys.iter().map(|key| self.memory.get(key)).collect()
}
fn backfill(&self, key: &str, value: Option<V>) -> Result<Option<V>, Error> {
if let Some(value) = &value {
self.set_memory(key, value.clone(), self.policy.backfill_ttl)?;
self.memory
.set(key, value.clone(), Some(self.policy.backfill_ttl))?;
}
Ok(value)
}
fn reserve(&self, keys: &[String], local: &[Option<V>]) -> Result<Reservation, Error> {
let now = self.now();
let mut reads = lock(&self.reads)?;
let mut reads = lock(&self.reads);
let due = keys
.iter()
.zip(local)
@ -235,7 +185,7 @@ where
}
fn release(&self, reservation: Reservation) -> Result<(), Error> {
let mut reads = lock(&self.reads)?;
let mut reads = lock(&self.reads);
for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) {
match previous {
Some(read) => reads.insert(key, read),
@ -246,18 +196,8 @@ where
}
}
fn lock<T>(mutex: &Mutex<T>) -> Result<MutexGuard<'_, T>, Error> {
Ok(mutex.lock()?)
}
async fn settle(
replies: Declared<Vec<Option<Vec<u8>>>>,
) -> Result<Vec<Option<Vec<u8>>>, RedisError> {
replies.resolve().await
}
/// The GETs `declare_many` put on a batch, with the refreshes they reserved.
pub struct DeclaredRead {
pub(crate) struct DeclaredRead {
keys: Vec<String>,
fetch: Option<Fetch>,
}

54
litellm-rust/crates/redis-store/src/value.rs Normal file
View file

@ -0,0 +1,54 @@
use litellm_python_compat::Value;
use crate::{DecodeError, Number};
/// A store type read from the plain Python data its host sends, so the host never names a field.
pub trait FromValue: Sized {
fn from_value(value: &Value) -> Result<Self, DecodeError>;
}
/// A store type written as the plain Python data its host hands back.
pub trait IntoValue {
fn into_value(self) -> Value;
}
/// The fields of a Python tuple of exactly `N` items.
pub(crate) fn fields<'a, const N: usize>(
value: &'a Value,
shape: &'static str,
) -> Result<&'a [Value; N], DecodeError> {
match value {
Value::Tuple(items) | Value::List(items) => {
items.as_slice().try_into().map_err(|_| DecodeError::Shape(shape))
}
_ => Err(DecodeError::Shape(shape)),
}
}
pub(crate) fn text(value: &Value, shape: &'static str) -> Result<String, DecodeError> {
match value {
Value::Str(text) => Ok(text.clone()),
_ => Err(DecodeError::Shape(shape)),
}
}
pub(crate) fn optional_text(
value: &Value,
shape: &'static str,
) -> Result<Option<String>, DecodeError> {
match value {
Value::None => Ok(None),
value => text(value, shape).map(Some),
}
}
pub(crate) fn integer(value: &Value, shape: &'static str) -> Result<i64, DecodeError> {
match value {
Value::Int(value) => i64::try_from(value).map_err(|_| DecodeError::Shape(shape)),
_ => Err(DecodeError::Shape(shape)),
}
}
pub(crate) fn number(value: &Value, shape: &'static str) -> Result<Number, DecodeError> {
Number::from_value(value).ok_or(DecodeError::Shape(shape))
}

91
litellm-rust/crates/redis-store/tests/auth_cache.rs
View file

@ -4,14 +4,16 @@ use std::time::Duration;
use litellm_python_compat::Value;
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_client::commands;
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
Number, Remote, RequestRedis,
auth_cache::{self, AuthCache, Settings},
router_cooldowns::{self, RouterCooldowns},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, connect_observed, within};
use support::{connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const LONG_DEADLINE: Duration = Duration::from_secs(60);
@ -123,7 +125,7 @@ fn settings(memory_ttl: Option<u64>) -> Settings {
}
fn store(client: Option<RedisClient>) -> AuthCache {
AuthCache::new(client, settings(Some(60)))
AuthCache::new(client.map(remote), settings(Some(60)))
}
fn team(id: &str) -> Value {
@ -131,7 +133,7 @@ fn team(id: &str) -> Value {
}
async fn redis_ttl(raw: &RedisClient, key: &str) -> Duration {
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
remaining
@ -173,12 +175,12 @@ async fn writes_mains_json_and_str_bytes(
let key = format!("team_id:t{index}");
store.set(&key, value.clone(), None).await.unwrap();
assert_eq!(
raw.get::<String>(&key).await.unwrap().as_deref(),
raw.send(commands::get::<String>(&key)).await.unwrap().as_deref(),
Some(json)
);
store.set_str(&key, value, None).await.unwrap();
assert_eq!(
raw.get::<String>(&key).await.unwrap().as_deref(),
raw.send(commands::get::<String>(&key)).await.unwrap().as_deref(),
Some(printed)
);
}
@ -195,10 +197,10 @@ async fn reads_every_encoding_main_writes_and_backfills_memory(
for (index, (bytes, expected)) in read_back().into_iter().enumerate() {
let key = format!("user_id:u{index}");
raw.set(&key, bytes, Duration::from_secs(60)).await.unwrap();
raw.send(commands::set(&key, bytes, Duration::from_secs(60))).await.unwrap();
assert_eq!(store.get(&key).await.unwrap(), Some(expected.clone()));
raw.del(&[&key]).await.unwrap();
raw.send(commands::del(&[&key])).await.unwrap();
assert_eq!(store.get(&key).await.unwrap(), Some(expected), "{key}");
}
assert_eq!(store.get("user_id:absent").await.unwrap(), None);
@ -212,10 +214,12 @@ async fn ttls_follow_dual_cache_and_the_write_back_rule(
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = AuthCache::new(
Some(connect(&deployment)),
Some(Remote {
client: connect(&deployment),
default_ttl: Duration::from_secs(45),
}),
Settings {
memory_ttl: Some(Duration::from_secs(20)),
redis_ttl: Duration::from_secs(45),
..Settings::default()
},
);
@ -231,12 +235,13 @@ async fn ttls_follow_dual_cache_and_the_write_back_rule(
store.set("team_id:default", team("b"), None).await.unwrap();
store
.write_back(
None,
RedisSend::Now,
vec![
("user-1".to_owned(), team("c"), None),
("user-2".to_owned(), team("d"), Some(Duration::from_secs(5))),
],
)
.unwrap()
.await
.unwrap();
@ -285,7 +290,7 @@ async fn a_value_json_refuses_stays_in_memory_and_never_reaches_redis(
.unwrap();
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(raw.get::<String>("team_id:set").await.unwrap(), None);
assert_eq!(raw.send(commands::get::<String>("team_id:set")).await.unwrap(), None);
}
#[rstest]
@ -296,32 +301,29 @@ async fn pre_call_writes_and_reads_ride_one_pipeline_with_the_cooldown_read(
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = RouterCooldowns::new(Some(client.clone()), router_cooldowns::POLICY);
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(Some(remote(client.clone())), router_cooldowns::POLICY));
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
raw.set(KEY_OBJECT, "{\"token\": \"t\"}", Duration::from_secs(60))
raw.send(commands::set(KEY_OBJECT, "{\"token\": \"t\"}", Duration::from_secs(60)))
.await
.unwrap();
raw.set("team_alias:alpha", "\"stale\"", Duration::from_secs(60))
raw.send(commands::set("team_alias:alpha", "\"stale\"", Duration::from_secs(60)))
.await
.unwrap();
let reads = store
.declare_get_many(
&request,
.get_many(RedisSend::PreCall(&request),
&[KEY_OBJECT.to_owned(), "end_user_id:absent".to_owned()],
)
.unwrap();
let active = cooldowns
.declare_active(&request, &["dep-a".to_owned()])
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let written = store
.set_pre_call(&request, "team_id:t1", team("t1"), None)
.unwrap()
.declare_set(RedisSend::PreCall(&request), "team_id:t1", team("t1"), None)
.unwrap();
let deleted = store
.delete_pre_call(&request, "team_alias:alpha")
.unwrap()
.declare_delete(RedisSend::PreCall(&request), "team_alias:alpha")
.unwrap();
assert!(trips.seen().is_empty());
@ -341,21 +343,21 @@ async fn pre_call_writes_and_reads_ride_one_pipeline_with_the_cooldown_read(
};
assert_eq!(seen, expected);
assert_eq!(
reads.resolve().await.unwrap(),
reads.await.unwrap(),
vec![
Some(Value::Dict(vec![(Value::from("token"), Value::from("t"))])),
None
]
);
assert_eq!(cooldowns.resolve_active(active).await.unwrap(), vec![]);
written.resolve().await.unwrap();
deleted.resolve().await.unwrap();
assert_eq!(active.await.unwrap(), vec![]);
written.await.unwrap();
deleted.await.unwrap();
assert_eq!(
raw.get::<String>("team_id:t1").await.unwrap(),
raw.send(commands::get::<String>("team_id:t1")).await.unwrap(),
Some("{\"team_id\": \"t1\"}".to_owned())
);
assert!(within(redis_ttl(&raw, "team_id:t1").await, 60));
assert_eq!(raw.get::<String>("team_alias:alpha").await.unwrap(), None);
assert_eq!(raw.send(commands::get::<String>("team_alias:alpha")).await.unwrap(), None);
assert_eq!(trips.seen().len(), expected.len());
}
@ -372,13 +374,14 @@ async fn a_write_back_on_the_request_waits_for_the_next_flush(
store
.write_back(
Some(&request),
RedisSend::PreCall(&request),
vec![(
"user-1".to_owned(),
team("u"),
Some(Duration::from_secs(42)),
)],
)
.unwrap()
.await
.unwrap();
@ -412,7 +415,7 @@ async fn an_eviction_broadcast_drops_only_this_pods_copy_and_the_next_read_refil
this_pod.delete_local("team_id:t1").unwrap();
assert_eq!(this_pod.get("team_id:t1").await.unwrap(), Some(team("new")));
assert!(raw.get::<String>("team_id:t1").await.unwrap().is_some());
assert!(raw.send(commands::get::<String>("team_id:t1")).await.unwrap().is_some());
}
#[rstest]
@ -431,13 +434,13 @@ async fn delete_drops_memory_and_redis(
)
.await
.unwrap();
assert!(raw.get::<String>("team_id:t1").await.unwrap().is_some());
assert!(raw.send(commands::get::<String>("team_id:t1")).await.unwrap().is_some());
store.delete(&keys).await.unwrap();
for key in &keys {
assert_eq!(store.get_local(key).unwrap(), None);
assert_eq!(raw.get::<String>(key).await.unwrap(), None);
assert_eq!(raw.send(commands::get::<String>(key)).await.unwrap(), None);
}
}
@ -458,7 +461,7 @@ async fn increment_if_present_only_raises_an_entry_redis_holds(
None
);
assert_eq!(
raw.get::<String>("litellm-proxy-budget:spend")
raw.send(commands::get::<String>("litellm-proxy-budget:spend"))
.await
.unwrap(),
None
@ -484,18 +487,24 @@ async fn without_redis_every_operation_is_memory_only() {
store.set("team_id:t1", team("t"), None).await.unwrap();
store
.declare_set(RedisSend::PreCall(&request), "team_id:t2", team("t"), None)
.unwrap()
.await
.unwrap();
assert_eq!(store.get_local("team_id:t2").unwrap(), Some(team("t")));
store
.declare_delete(RedisSend::PreCall(&request), "team_id:t2")
.unwrap()
.await
.unwrap();
assert_eq!(store.get_local("team_id:t2").unwrap(), None);
assert!(
store
.set_pre_call(&request, "team_id:t2", team("t"), None)
.unwrap()
.get_many(RedisSend::PreCall(&request), &["team_id:t1".to_owned()])
.is_none()
);
assert!(
store
.declare_get_many(&request, &["team_id:t1".to_owned()])
.is_none()
);
assert!(store.get_many(&["team_id:t1".to_owned()]).await.is_none());
assert!(store.get_many(RedisSend::Now, &["team_id:t1".to_owned()]).is_none());
assert_eq!(
store
.increment_if_present("k", Number::Float(1.0))

158
litellm-rust/crates/redis-store/tests/rate_limits.rs
View file

@ -3,6 +3,8 @@ mod support;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Operation, RedisClient, Script, Ttl};
use litellm_redis_client::commands;
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
rate_limits::{
@ -13,7 +15,7 @@ use litellm_redis_store::{
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, connect_observed, within};
use support::{connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const SLOT_TTL: Duration = Duration::from_secs(3600);
@ -118,34 +120,31 @@ fn atomic(scope: &Scope, kind: Kind, limit: i64, increment: i64) -> AtomicCounte
}
}
/// `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))
Arc::new(RateLimits::new(Some(remote(connect(deployment))), None))
}
async fn raw_text(raw: &RedisClient, key: &str) -> Option<String> {
raw.get(key).await.unwrap()
raw.send(commands::get(key)).await.unwrap()
}
async fn zcard(raw: &RedisClient, key: &str) -> i64 {
raw.eval(
raw.send(commands::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(
.send(commands::eval(
&Script::new("return redis.call('TIME')"),
&["any"],
Vec::<String>::new(),
)
))
.await
.unwrap();
time[0].parse().unwrap()
@ -173,20 +172,17 @@ async fn every_script_runs_mains_bytes_under_mains_keys(
let requests = counter(&key, Kind::Requests);
let tokens = counter(&key, Kind::Tokens);
limits
.sliding_window(None, std::slice::from_ref(&requests), now, 60)
.sliding_window(RedisSend::Now, std::slice::from_ref(&requests), now, 60)
.unwrap()
.resolve()
.await
.unwrap();
limits
.check_and_increment(None, &[atomic(&key, Kind::Tokens, 1000, 40)])
.check_and_increment(RedisSend::Now, &[atomic(&key, Kind::Tokens, 1000, 40)])
.unwrap()
.resolve()
.await
.unwrap();
limits
.token_increment(
None,
.token_increment(RedisSend::Now,
&[Increment {
counter: tokens.clone(),
amount: Number::Float(2.5),
@ -194,12 +190,10 @@ async fn every_script_runs_mains_bytes_under_mains_keys(
}],
)
.unwrap()
.resolve()
.await
.unwrap();
limits
.parallel_acquire(
None,
.parallel_acquire(RedisSend::Now,
&[Gauge {
scope: key.clone(),
limit: 5,
@ -208,7 +202,6 @@ async fn every_script_runs_mains_bytes_under_mains_keys(
"slot-1",
)
.unwrap()
.resolve()
.await
.unwrap();
@ -218,9 +211,8 @@ async fn every_script_runs_mains_bytes_under_mains_keys(
assert_eq!(zcard(&raw, wire.gauge).await, 1);
assert_eq!(
limits
.read_counters(None, &[requests, tokens])
.read_counters(RedisSend::Now, &[requests, tokens])
.unwrap()
.resolve()
.await
.unwrap(),
vec![
@ -234,8 +226,7 @@ async fn every_script_runs_mains_bytes_under_mains_keys(
let guarded = scope("api_key", "sha-check");
limits
.guarded_increment(
None,
.guarded_increment(RedisSend::Now,
&GuardedIncrement {
counter: counter(&guarded, Kind::Tokens),
expected_window_start: "0".to_owned(),
@ -244,14 +235,13 @@ async fn every_script_runs_mains_bytes_under_mains_keys(
},
)
.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),
limits.parallel_release(RedisSend::Now, std::slice::from_ref(&guarded), "slot-1"),
limits.parallel_count(RedisSend::Now, std::slice::from_ref(&guarded), SLOT_TTL),
] {
pending.unwrap().resolve().await.unwrap();
pending.unwrap().await.unwrap();
}
for (name, sha) in SCRIPTS {
let cached = cli(&deployment, &["SCRIPT", "EXISTS", sha]).await;
@ -296,7 +286,7 @@ async fn reads_values_python_wrote_and_backfills_memory(
let raw = connect(&deployment);
let limits = store(&deployment);
let wire = &SCOPES[0];
raw.set(wire.tokens, raw_value, Duration::from_secs(60))
raw.send(commands::set(wire.tokens, raw_value, Duration::from_secs(60)))
.await
.unwrap();
let entry = Entry::Counter(counter(&scope(wire.descriptor, wire.value), Kind::Tokens));
@ -308,7 +298,7 @@ async fn reads_values_python_wrote_and_backfills_memory(
.unwrap(),
Some(Local::Number(expected))
);
assert_eq!(limits.local(&entry).unwrap(), Some(Local::Number(expected)));
assert_eq!(limits.get_local(&entry).unwrap(), Some(Local::Number(expected)));
}
#[rstest]
@ -323,9 +313,9 @@ async fn sliding_window_opens_counts_and_resets_with_its_sibling(
let requests = counter(&key, Kind::Requests);
let run = |now: i64| {
let pending = limits
.sliding_window(None, std::slice::from_ref(&requests), now, 60)
.sliding_window(RedisSend::Now, std::slice::from_ref(&requests), now, 60)
.unwrap();
async move { pending.resolve().await.unwrap() }
async move { pending.await.unwrap() }
};
assert_eq!(
@ -336,10 +326,10 @@ async fn sliding_window_opens_counts_and_resets_with_its_sibling(
}]
);
assert_eq!(run(1_059).await[0].count, 2);
raw.set("{api_key:k1}:tokens", "500", Duration::from_secs(60))
raw.send(commands::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 {
let Ttl::Expires(ttl) = raw.send(commands::ttl("{api_key:k1}:requests")).await.unwrap() else {
panic!("the counter has no TTL");
};
assert!(within(ttl, 60), "{ttl:?}");
@ -372,9 +362,8 @@ async fn check_and_increment_is_all_or_nothing(
};
let first = limits
.check_and_increment(None, &both(1, 60))
.check_and_increment(RedisSend::Now, &both(1, 60))
.unwrap()
.resolve()
.await
.unwrap();
let AtomicReply::Applied(applied) = first else {
@ -385,9 +374,8 @@ async fn check_and_increment_is_all_or_nothing(
assert_eq!(
limits
.check_and_increment(None, &both(1, 41))
.check_and_increment(RedisSend::Now, &both(1, 41))
.unwrap()
.resolve()
.await
.unwrap(),
AtomicReply::OverLimit {
@ -406,9 +394,8 @@ async fn check_and_increment_is_all_or_nothing(
);
let at_limit = limits
.check_and_increment(None, &both(1, 40))
.check_and_increment(RedisSend::Now, &both(1, 40))
.unwrap()
.resolve()
.await
.unwrap();
assert!(
@ -416,9 +403,8 @@ async fn check_and_increment_is_all_or_nothing(
);
assert_eq!(
limits
.check_and_increment(None, &both(0, 0))
.check_and_increment(RedisSend::Now, &both(0, 0))
.unwrap()
.resolve()
.await
.unwrap(),
AtomicReply::OverLimit {
@ -440,9 +426,8 @@ async fn guarded_increment_only_settles_the_window_it_was_reserved_in(
let key = scope("team", "t1");
let tokens = counter(&key, Kind::Tokens);
let reply = limits
.check_and_increment(None, &[atomic(&key, Kind::Tokens, 1000, 300)])
.check_and_increment(RedisSend::Now, &[atomic(&key, Kind::Tokens, 1000, 300)])
.unwrap()
.resolve()
.await
.unwrap();
let AtomicReply::Applied(applied) = reply else {
@ -451,8 +436,7 @@ async fn guarded_increment_only_settles_the_window_it_was_reserved_in(
let Applied { window_start, .. } = applied[0];
let guarded = |expected_window_start: String, amount: i64| {
let pending = limits
.guarded_increment(
None,
.guarded_increment(RedisSend::Now,
&GuardedIncrement {
counter: tokens.clone(),
expected_window_start,
@ -461,7 +445,7 @@ async fn guarded_increment_only_settles_the_window_it_was_reserved_in(
},
)
.unwrap();
async move { pending.resolve().await.unwrap() }
async move { pending.await.unwrap() }
};
assert_eq!(
@ -498,21 +482,21 @@ async fn gauge_slots_are_acquired_released_by_id_and_expire(
}];
let acquire = |slot_id: &'static str, slot_ttl: Duration| {
let pending = limits
.parallel_acquire(None, &gauges, slot_ttl, slot_id)
.parallel_acquire(RedisSend::Now, &gauges, slot_ttl, slot_id)
.unwrap();
async move { pending.resolve().await.unwrap() }
async move { pending.await.unwrap() }
};
let release = |slot_id: &'static str| {
let pending = limits
.parallel_release(None, std::slice::from_ref(&key), slot_id)
.parallel_release(RedisSend::Now, std::slice::from_ref(&key), slot_id)
.unwrap();
async move { pending.resolve().await.unwrap() }
async move { pending.await.unwrap() }
};
let count = |slot_ttl: Duration| {
let pending = limits
.parallel_count(None, std::slice::from_ref(&key), slot_ttl)
.parallel_count(RedisSend::Now, std::slice::from_ref(&key), slot_ttl)
.unwrap();
async move { pending.resolve().await.unwrap() }
async move { pending.await.unwrap() }
};
assert_eq!(
@ -558,22 +542,20 @@ async fn token_increment_sets_a_ttl_only_on_a_key_without_one(
amount,
ttl: Some(Duration::from_secs(60)),
};
raw.set("{user:u1}:tokens", "10", Duration::from_secs(600))
raw.send(commands::set("{user:u1}:tokens", "10", Duration::from_secs(600)))
.await
.unwrap();
let totals = limits
.token_increment(
None,
.token_increment(RedisSend::Now,
&[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 {
let Ttl::Expires(ttl) = raw.send(commands::ttl("{user:u1}:tokens")).await.unwrap() else {
panic!("the counter lost its TTL");
};
assert!(within(ttl, 600), "{ttl:?}");
@ -643,9 +625,8 @@ async fn gauges_of_different_slots_in_one_call_fail_only_on_a_cluster(
});
let acquired = limits
.parallel_acquire(None, &gauges, SLOT_TTL, "slot")
.parallel_acquire(RedisSend::Now, &gauges, SLOT_TTL, "slot")
.unwrap()
.resolve()
.await;
match deployment {
@ -668,9 +649,7 @@ async fn a_cross_slot_token_script_on_the_post_call_batch_falls_back_to_plain_in
ttl: Some(Duration::from_secs(60)),
});
limits
.token_increment_post_call(&request, increments.to_vec())
.unwrap();
assert!(limits.token_increment_post_call(&request, increments.to_vec()));
assert_eq!(raw_text(&raw, "{api_key:a}:tokens").await, None);
request.flush_post().await;
@ -680,7 +659,7 @@ async fn a_cross_slot_token_script_on_the_post_call_batch_falls_back_to_plain_in
assert_eq!(written, "25");
}
let local = limits
.local(&Entry::Counter(counter(
.get_local(&Entry::Counter(counter(
&scope("api_key", "a"),
Kind::Tokens,
)))
@ -736,7 +715,7 @@ async fn post_call_increments_count_in_memory_now_and_reach_redis_with_the_batch
amount: Number::Int(-40),
ttl: Some(Duration::from_secs(60)),
};
raw.set("{end_user:e1}:tokens", "100", Duration::from_secs(600))
raw.send(commands::set("{end_user:e1}:tokens", "100", Duration::from_secs(600)))
.await
.unwrap();
@ -745,7 +724,7 @@ async fn post_call_increments_count_in_memory_now_and_reach_redis_with_the_batch
.unwrap();
assert_eq!(
limits
.local(&Entry::Counter(increment.counter.clone()))
.get_local(&Entry::Counter(increment.counter.clone()))
.unwrap(),
Some(Local::Number(Number::Int(-40)))
);
@ -762,7 +741,7 @@ async fn post_call_increments_count_in_memory_now_and_reach_redis_with_the_batch
})
.await;
assert_eq!(written, "60");
let Ttl::Expires(ttl) = raw.ttl("{end_user:e1}:tokens").await.unwrap() else {
let Ttl::Expires(ttl) = raw.send(commands::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
@ -786,7 +765,7 @@ async fn increment_keeps_an_existing_ttl_and_counts_in_memory(
amount: Number::Int(5),
ttl: Some(Duration::from_secs(60)),
};
raw.set("{organization:o1}:tokens", "10", Duration::from_secs(600))
raw.send(commands::set("{organization:o1}:tokens", "10", Duration::from_secs(600)))
.await
.unwrap();
@ -794,13 +773,13 @@ async fn increment_keeps_an_existing_ttl_and_counts_in_memory(
limits.increment(&increment).await.unwrap(),
Number::Float(15.0)
);
let Ttl::Expires(ttl) = raw.ttl("{organization:o1}:tokens").await.unwrap() else {
let Ttl::Expires(ttl) = raw.send(commands::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()))
.get_local(&Entry::Counter(increment.counter.clone()))
.unwrap(),
Some(Local::Number(Number::Int(5)))
);
@ -821,9 +800,8 @@ async fn increment_answers_with_memory_while_redis_is_down() {
assert_eq!(limits.increment(&increment).await.unwrap(), Number::Int(14));
assert!(
limits
.sliding_window(None, std::slice::from_ref(&increment.counter), 1_000, 60)
.sliding_window(RedisSend::Now, std::slice::from_ref(&increment.counter), 1_000, 60)
.unwrap()
.resolve()
.await
.is_err()
);
@ -831,17 +809,19 @@ async fn increment_answers_with_memory_while_redis_is_down() {
#[tokio::test]
async fn scripts_need_redis() {
let limits = RateLimits::new(None, None, REDIS_TTL);
let limits = Arc::new(RateLimits::new(None, None));
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
assert!(
limits
.sliding_window(
None,
.sliding_window(RedisSend::Now,
&[counter(&scope("api_key", "k"), Kind::Requests)],
0,
60
)
.is_err()
.is_none()
);
assert!(!limits.token_increment_post_call(&request, Vec::new()));
assert!(!limits.parallel_release_post_call(&request, &[scope("api_key", "k")], "slot"));
assert_eq!(
limits.slot_groups(&[scope("a", "1"), scope("b", "2")], false),
vec![vec![0, 1]]
@ -855,23 +835,22 @@ async fn the_limiter_script_and_the_armed_cooldown_read_share_one_pipeline(
) {
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 cooldowns = std::sync::Arc::new(RouterCooldowns::new(Some(remote(client.clone())), COOLDOWN_POLICY));
let limits = RateLimits::new(Some(remote(client)), None);
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let armed = cooldowns
.declare_active(&request, &["dep-a".to_owned()])
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let admission = limits
.sliding_window(
Some(&request),
.sliding_window(RedisSend::PreCall(&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();
let counts = admission.await.unwrap();
let active = armed.await.unwrap();
assert_eq!(counts[0].count, 1);
assert!(active.is_empty());
@ -908,7 +887,7 @@ async fn an_increment_without_a_ttl_takes_the_redis_default_like_async_increment
.await
.unwrap();
let Ttl::Expires(ttl) = raw.ttl("{api_key:k}:tokens").await.unwrap() else {
let Ttl::Expires(ttl) = raw.send(commands::ttl("{api_key:k}:tokens")).await.unwrap() else {
panic!("the counter has no TTL");
};
assert!(within(ttl, 60), "{ttl:?}");
@ -950,9 +929,16 @@ async fn a_cold_script_cache_costs_the_evalshas_redis_py_sends(
}];
let reply = limits
.parallel_acquire(on_request.then_some(&request), &gauges, SLOT_TTL, "slot-1")
.parallel_acquire(
match on_request {
true => RedisSend::PreCall(&request),
false => RedisSend::Now,
},
&gauges,
SLOT_TTL,
"slot-1",
)
.unwrap()
.resolve()
.await
.unwrap();

27
litellm-rust/crates/redis-store/tests/request.rs
View file

@ -2,14 +2,15 @@ mod support;
use std::time::Duration;
use litellm_redis_client::{Amount, Operation};
use litellm_redis_client::{Amount, Expiry, Operation};
use litellm_redis_client::commands;
use litellm_redis_store::RequestRedis;
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, connect_observed};
use tokio::{runtime::Handle, sync::oneshot};
const TTL: Option<Duration> = Some(Duration::from_secs(60));
const TTL: Option<Expiry> = Some(Expiry::KeepExisting(Duration::from_secs(60)));
// A batch is one pipeline on a standalone Redis and one call per increment on a cluster, like RedisBatch
fn increments(topology: Topology, count: usize) -> Vec<Operation> {
@ -30,11 +31,11 @@ async fn pre_declarations_on_one_client_share_one_pipeline_and_another_client_ge
let (other, other_trips) = connect_observed(&deployment);
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let first = request.pre(&client).incr_by_float("a", Amount::Int(1), TTL);
let first = request.pre(&client).add(commands::incr_by_float("a", Amount::Int(1), TTL));
let second = request
.pre(&client.clone())
.incr_by_float("b", Amount::Int(2), TTL);
let third = request.pre(&other).incr_by_float("c", Amount::Int(3), TTL);
.add(commands::incr_by_float("b", Amount::Int(2), TTL));
let third = request.pre(&other).add(commands::incr_by_float("c", Amount::Int(3), TTL));
request.flush_pre().await;
assert_eq!(trips.seen(), increments(topology, 2));
@ -63,7 +64,7 @@ async fn a_post_declaration_waits_for_flush_post_and_flush_pre_leaves_it_alone(
let total = tokio::spawn(
request
.post(&client)
.incr_by_float("spend", Amount::Float(2.5), TTL)
.add(commands::incr_by_float("spend", Amount::Float(2.5), TTL))
.settled(),
);
request.flush_pre().await;
@ -71,7 +72,7 @@ async fn a_post_declaration_waits_for_flush_post_and_flush_pre_leaves_it_alone(
assert!(!total.is_finished());
assert!(trips.seen().is_empty());
assert_eq!(raw.get::<String>("spend").await.unwrap(), None);
assert_eq!(raw.send(commands::get::<String>("spend")).await.unwrap(), None);
request.flush_post().await;
@ -89,7 +90,7 @@ async fn flush_post_returns_after_the_follow_ups_of_the_ops_it_sent(
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let total = request
.post(&client)
.incr_by_float("spend", Amount::Float(2.5), TTL);
.add(commands::incr_by_float("spend", Amount::Float(2.5), TTL));
let (mirror, mut mirrored) = oneshot::channel();
request.spawn(async move {
mirror.send(total.settled().await.unwrap()).unwrap();
@ -111,10 +112,10 @@ async fn a_post_batch_nobody_flushes_goes_out_at_the_deadline(
let first = request
.post(&client)
.incr_by_float("a", Amount::Int(1), TTL);
.add(commands::incr_by_float("a", Amount::Int(1), TTL));
let second = request
.post(&client)
.incr_by_float("b", Amount::Int(2), TTL);
.add(commands::incr_by_float("b", Amount::Int(2), TTL));
let settled = tokio::time::timeout(Duration::from_secs(5), async {
(first.settled().await, second.settled().await)
})
@ -139,7 +140,7 @@ async fn a_post_declaration_from_a_thread_without_a_runtime_goes_out_at_the_dead
move || {
request
.post(&client)
.incr_by_float("spend", Amount::Float(2.5), TTL)
.add(commands::incr_by_float("spend", Amount::Float(2.5), TTL))
}
})
.join()
@ -163,12 +164,12 @@ async fn a_post_declaration_after_a_flush_arms_a_new_deadline(
let early = request
.post(&client)
.incr_by_float("a", Amount::Int(1), TTL);
.add(commands::incr_by_float("a", Amount::Int(1), TTL));
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(200)).await;
let late = request
.post(&client)
.incr_by_float("a", Amount::Int(1), TTL);
.add(commands::incr_by_float("a", Amount::Int(1), TTL));
tokio::time::sleep(Duration::from_millis(200)).await;
let flushed_before_its_deadline = trips.seen().len() > 1;
let late = tokio::time::timeout(Duration::from_secs(5), late.settled())

101
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
View file

@ -3,15 +3,21 @@ mod support;
use std::time::Duration;
use litellm_redis_client::{Operation, Ttl};
use litellm_redis_client::commands;
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::{POLICY, RouterCooldowns},
router_cooldowns::{Cooldown, POLICY, RouterCooldowns},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{ALLOWED_FAILS, connect, cooldowns, within};
use support::{ALLOWED_FAILS, connect, cooldowns, remote, within};
use tokio::runtime::Handle;
async fn read_active(store: &std::sync::Arc<RouterCooldowns>, model_ids: &[String]) -> Vec<(String, Cooldown)> {
store.active(RedisSend::Now, model_ids).unwrap().await.unwrap()
}
/// A clock fixed just after the fixture timestamps, so every fixture cooldown is still active.
fn fixture_clock() -> litellm_redis_store::Clock {
std::sync::Arc::new(|| Duration::from_secs_f64(1_730_000_000.75))
@ -24,15 +30,15 @@ async fn put_writes_mains_str_dict_bytes_with_whole_second_ttl(
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
let store = std::sync::Arc::new(RouterCooldowns::new(Some(remote(connect(&deployment))), POLICY));
for case in cooldowns() {
let ttl = Duration::from_secs_f64(case.ttl_seconds);
store.put(case.model_id, &case.cooldown, ttl).await.unwrap();
let key = case.key;
let written: Option<String> = raw.get(key).await.unwrap();
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
let written: Option<String> = raw.send(commands::get(key)).await.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(written.as_deref(), Some(case.str_dict));
@ -48,19 +54,19 @@ async fn reads_cooldowns_python_wrote_as_str_dict_or_json(
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(connect(&deployment))), POLICY, fixture_clock()));
for case in cooldowns() {
let bytes = match format {
"json" => case.json,
_ => case.str_dict,
};
raw.set(case.key, bytes, Duration::from_secs(60))
raw.send(commands::set(case.key, bytes, Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(
store.active(&[case.model_id.to_owned()]).await.unwrap(),
read_active(&store, &[case.model_id.to_owned()]).await,
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
@ -74,13 +80,13 @@ async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum(
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let late = std::sync::Arc::new(|| Duration::from_secs(1_730_000_100));
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, late);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(connect(&deployment))), POLICY, late));
let models = cooldowns()
.iter()
.map(|case| case.model_id.to_owned())
.collect::<Vec<_>>();
for case in cooldowns() {
raw.set(case.key, case.str_dict, Duration::from_secs(60))
raw.send(commands::set(case.key, case.str_dict, Duration::from_secs(60)))
.await
.unwrap();
}
@ -89,7 +95,7 @@ async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum(
store.minimum(&models).await.unwrap(),
Some(Number::Float(0.5))
);
assert_eq!(store.active(&models).await.unwrap(), vec![]);
assert_eq!(read_active(&store, &models).await, vec![]);
}
#[rstest]
@ -99,12 +105,12 @@ async fn allowed_fails_count_across_routers_and_keep_the_first_ttl(
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
let second = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
let first = std::sync::Arc::new(RouterCooldowns::new(Some(remote(connect(&deployment))), POLICY));
let second = std::sync::Arc::new(RouterCooldowns::new(Some(remote(connect(&deployment))), POLICY));
for case in ALLOWED_FAILS {
let key = case.key;
raw.del(&[key]).await.unwrap();
raw.send(commands::del(&[key])).await.unwrap();
let suffix = case.suffix;
let counts = [
@ -117,7 +123,7 @@ async fn allowed_fails_count_across_routers_and_keep_the_first_ttl(
.await
.unwrap(),
];
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
@ -133,8 +139,8 @@ async fn python_written_fail_counts_continue_in_rust(
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
raw.incr_by("deployment:dep-a:allowed_fails", 4, Duration::from_secs(60))
let store = std::sync::Arc::new(RouterCooldowns::new(Some(remote(connect(&deployment))), POLICY));
raw.send(commands::incr_by("deployment:dep-a:allowed_fails", 4, Duration::from_secs(60)))
.await
.unwrap();
@ -150,7 +156,7 @@ async fn python_written_fail_counts_continue_in_rust(
#[rstest]
#[tokio::test]
async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
let store = RouterCooldowns::with_clock(None, POLICY, fixture_clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, fixture_clock()));
let [case, _] = cooldowns();
store
@ -170,7 +176,7 @@ async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
store.flush_local_counters().unwrap();
assert_eq!(
store.active(&[case.model_id.to_owned()]).await.unwrap(),
read_active(&store, &[case.model_id.to_owned()]).await,
vec![(case.model_id.to_owned(), case.cooldown)]
);
assert_eq!(fails, [1, 2]);
@ -187,7 +193,7 @@ async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
#[tokio::test]
async fn a_redis_outage_serves_memory_and_counts_fails_locally() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(connect(&deployment))), POLICY, fixture_clock()));
let [case, _] = cooldowns();
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
@ -195,10 +201,7 @@ async fn a_redis_outage_serves_memory_and_counts_fails_locally() {
.unwrap();
deployment.stop().await.unwrap();
let active = store
.active(&[case.model_id.to_owned(), "dep-missing".to_owned()])
.await
.unwrap();
let active = read_active(&store, &[case.model_id.to_owned(), "dep-missing".to_owned()]).await;
let fails = store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
@ -235,16 +238,16 @@ async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh(
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let time = StepClock::at(1_730_000_000.75);
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, time.clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(connect(&deployment))), POLICY, time.clock()));
let [case, _] = cooldowns();
let before = store.active(&[case.model_id.to_owned()]).await.unwrap();
raw.set(case.key, case.str_dict, Duration::from_secs(60))
let before = read_active(&store, &[case.model_id.to_owned()]).await;
raw.send(commands::set(case.key, case.str_dict, Duration::from_secs(60)))
.await
.unwrap();
let throttled = store.active(&[case.model_id.to_owned()]).await.unwrap();
let throttled = read_active(&store, &[case.model_id.to_owned()]).await;
time.advance(Duration::from_secs(1));
let refreshed = store.active(&[case.model_id.to_owned()]).await.unwrap();
let refreshed = read_active(&store, &[case.model_id.to_owned()]).await;
assert_eq!(before, vec![]);
assert_eq!(throttled, vec![]);
@ -269,19 +272,19 @@ async fn a_breaker_refusal_does_not_hold_back_the_next_read() {
);
let time = StepClock::at(1_730_000_000.75);
let guarded = support::connect_guarded(&deployment, breaker);
let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, time.clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(guarded.clone())), POLICY, time.clock()));
let [case, _] = cooldowns();
deployment.stop().await.unwrap();
store.active(&[case.model_id.to_owned()]).await.unwrap();
read_active(&store, &[case.model_id.to_owned()]).await;
time.advance(Duration::from_secs(1));
let refused = store.active(&[case.model_id.to_owned()]).await.unwrap();
let refused = read_active(&store, &[case.model_id.to_owned()]).await;
deployment.resume().await.unwrap();
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
raw.set(case.key, case.str_dict, Duration::from_secs(60))
raw.send(commands::set(case.key, case.str_dict, Duration::from_secs(60)))
.await
.unwrap();
let recovered = store.active(&[case.model_id.to_owned()]).await.unwrap();
let recovered = read_active(&store, &[case.model_id.to_owned()]).await;
assert_eq!(refused, vec![]);
assert_eq!(
@ -298,19 +301,19 @@ async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
let store = RouterCooldowns::with_clock(Some(client), POLICY, fixture_clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(client)), POLICY, fixture_clock()));
let models = cooldowns().map(|case| case.model_id.to_owned());
for case in cooldowns() {
raw.set(case.key, case.str_dict, Duration::from_secs(60))
raw.send(commands::set(case.key, case.str_dict, Duration::from_secs(60)))
.await
.unwrap();
}
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let declared = store.declare_active(&request, &models).unwrap();
let declared = store.active(RedisSend::PreCall(&request), &models).unwrap();
assert!(trips.seen().is_empty());
request.flush_pre().await;
let active = store.resolve_active(declared).await.unwrap();
let active = declared.await.unwrap();
let one_mget = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
@ -323,7 +326,7 @@ async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_
.map(|case| (case.model_id.to_owned(), case.cooldown))
.to_vec()
);
assert_eq!(store.active(&models).await.unwrap(), active);
assert_eq!(read_active(&store, &models).await, active);
assert_eq!(trips.seen().len(), 1, "the declared read backfilled memory");
}
@ -333,12 +336,12 @@ async fn a_cooldown_recorded_after_the_declaration_wins_over_the_redis_miss(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let store = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(connect(&deployment))), POLICY, fixture_clock()));
let [case, _] = cooldowns();
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let declared = store
.declare_active(&request, &[case.model_id.to_owned()])
.active(RedisSend::PreCall(&request), &[case.model_id.to_owned()])
.unwrap();
request.flush_pre().await;
store
@ -347,7 +350,7 @@ async fn a_cooldown_recorded_after_the_declaration_wins_over_the_redis_miss(
.unwrap();
assert_eq!(
store.resolve_active(declared).await.unwrap(),
declared.await.unwrap(),
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
@ -366,30 +369,30 @@ async fn a_breaker_refusal_of_a_declared_read_does_not_hold_back_the_next_read()
Box::new(|_| {}),
);
let guarded = support::connect_guarded(&deployment, breaker);
let store = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, fixture_clock());
let store = std::sync::Arc::new(RouterCooldowns::with_clock(Some(remote(guarded.clone())), POLICY, fixture_clock()));
let [case, _] = cooldowns();
let models = [case.model_id.to_owned()];
deployment.stop().await.unwrap();
let opened = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let failed = store
.declare_active(&opened, &["dep-other".to_owned()])
.active(RedisSend::PreCall(&opened), &["dep-other".to_owned()])
.unwrap();
opened.flush_pre().await;
store.resolve_active(failed).await.unwrap();
failed.await.unwrap();
let refused_request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let refused = store.declare_active(&refused_request, &models).unwrap();
let refused = store.active(RedisSend::PreCall(&refused_request), &models).unwrap();
refused_request.flush_pre().await;
let refused = store.resolve_active(refused).await.unwrap();
let refused = refused.await.unwrap();
deployment.resume().await.unwrap();
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
raw.set(case.key, case.str_dict, Duration::from_secs(60))
raw.send(commands::set(case.key, case.str_dict, Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(refused, vec![]);
assert_eq!(
store.active(&models).await.unwrap(),
read_active(&store, &models).await,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}

67
litellm-rust/crates/redis-store/tests/router_usage.rs
View file

@ -3,13 +3,14 @@ mod support;
use std::time::Duration;
use litellm_redis_client::{Amount, Expiry, Operation, Ttl};
use litellm_redis_client::commands;
use litellm_redis_store::{
Number,
router_usage::{COUNTER_TTL, Counter, Kind, POLICY, RouterUsage},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, usage, within};
use support::{connect, remote, usage, within};
fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] {
[Kind::Tpm, Kind::Rpm].map(|kind| Counter {
@ -42,10 +43,10 @@ async fn reads_counters_python_incremented_under_mains_keys(
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let cases = usage();
for (case, delta) in cases.iter().zip([7.75, 3.0]) {
raw.incr_by_float(case.key, Amount::Float(delta), Expiry::Refresh(COUNTER_TTL))
raw.send(commands::incr_by_float(case.key, Amount::Float(delta), Some(Expiry::Refresh(COUNTER_TTL))))
.await
.unwrap();
}
@ -66,8 +67,8 @@ async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let writer = RouterUsage::new(Some(connect(&deployment)), POLICY);
let reader = RouterUsage::new(Some(connect(&deployment)), POLICY);
let writer = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let reader = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let [case, _] = usage();
let model = case.counter.model.as_deref();
let [tpm_key, rpm_key] =
@ -77,17 +78,17 @@ async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_
.increment_deployment("dep", model, "12-34", Number::Float(5.5), 1)
.await
.unwrap();
raw.expire(&tpm_key, Duration::from_secs(5)).await.unwrap();
raw.send(commands::expire(&tpm_key, Duration::from_secs(5))).await.unwrap();
writer
.increment_deployment("dep", model, "12-34", Number::Int(2), 1)
.await
.unwrap();
let written: [Option<String>; 2] = [
raw.get(&tpm_key).await.unwrap(),
raw.get(&rpm_key).await.unwrap(),
raw.send(commands::get(&tpm_key)).await.unwrap(),
raw.send(commands::get(&rpm_key)).await.unwrap(),
];
let Ttl::Expires(remaining) = raw.ttl(&tpm_key).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&tpm_key)).await.unwrap() else {
panic!("{tpm_key} has no TTL");
};
assert_eq!(written, [Some("7.5".to_owned()), Some("2".to_owned())]);
@ -109,14 +110,14 @@ async fn a_request_increment_keeps_the_window_python_started(
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let [_, case] = usage();
let key = case.key;
raw.incr_by_float(
raw.send(commands::incr_by_float(
key,
Amount::Float(1.0),
Expiry::KeepExisting(Duration::from_secs(20)),
)
Some(Expiry::KeepExisting(Duration::from_secs(20))),
))
.await
.unwrap();
@ -124,7 +125,7 @@ async fn a_request_increment_keeps_the_window_python_started(
.increment(&case.counter, Number::Int(1), COUNTER_TTL)
.await
.unwrap();
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
@ -156,11 +157,35 @@ async fn without_redis_counters_accumulate_in_memory_until_flushed() {
assert_eq!(store.get_many(&counters).await.unwrap(), vec![None, None]);
}
#[rstest]
#[tokio::test]
async fn a_late_attach_keeps_memory_and_later_attaches_are_ignored(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let attached = RedisDeployment::start(topology).await.unwrap();
let ignored = RedisDeployment::start(topology).await.unwrap();
let store = RouterUsage::new(None, POLICY);
let [case, _] = usage();
store.increment(&case.counter, Number::Int(2), COUNTER_TTL).await.unwrap();
assert!(store.attach(remote(connect(&attached))));
assert!(!store.attach(remote(connect(&ignored))));
let total = store.increment(&case.counter, Number::Int(3), COUNTER_TTL).await.unwrap();
assert_eq!(total, Number::Float(3.0));
assert_eq!(
store.get_many(std::slice::from_ref(&case.counter)).await.unwrap(),
vec![Some(Number::Int(5))]
);
let ignored_total: Option<String> = connect(&ignored).send(commands::get(case.key)).await.unwrap();
assert_eq!(ignored_total, None);
}
#[rstest]
#[tokio::test]
async fn a_failed_pipeline_keeps_the_memory_counts_as_integers() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
deployment.stop().await.unwrap();
store
@ -184,11 +209,11 @@ async fn a_deployment_increment_is_its_own_pipeline_and_mirrors_the_redis_totals
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
let store = RouterUsage::new(Some(client), POLICY);
let store = RouterUsage::new(Some(remote(client)), POLICY);
let counters = tpm_and_rpm(Some("model"));
let [tpm_key, rpm_key] =
["tpm", "rpm"].map(|kind| format!("global_router:dep:model:{kind}:12-34"));
raw.incr_by_float(&tpm_key, Amount::Float(10.0), Expiry::Refresh(COUNTER_TTL))
raw.send(commands::incr_by_float(&tpm_key, Amount::Float(10.0), Some(Expiry::Refresh(COUNTER_TTL))))
.await
.unwrap();
@ -203,11 +228,11 @@ async fn a_deployment_increment_is_its_own_pipeline_and_mirrors_the_redis_totals
vec![Some(Number::Int(15)), Some(Number::Int(1))]
);
let written: [Option<String>; 2] = [
raw.get(&tpm_key).await.unwrap(),
raw.get(&rpm_key).await.unwrap(),
raw.send(commands::get(&tpm_key)).await.unwrap(),
raw.send(commands::get(&rpm_key)).await.unwrap(),
];
assert_eq!(written, [Some("15.5".to_owned()), Some("1".to_owned())]);
let Ttl::Expires(remaining) = raw.ttl(&rpm_key).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&rpm_key)).await.unwrap() else {
panic!("{rpm_key} has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
@ -217,7 +242,7 @@ async fn a_deployment_increment_is_its_own_pipeline_and_mirrors_the_redis_totals
#[tokio::test]
async fn a_failed_deployment_pipeline_keeps_the_memory_counts_as_integers() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = RouterUsage::new(Some(connect(&deployment)), POLICY);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let counters = tpm_and_rpm(Some("model"));
deployment.stop().await.unwrap();

62
litellm-rust/crates/redis-store/tests/spend_counters.rs
View file

@ -3,6 +3,8 @@ mod support;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_client::commands;
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::{self, RouterCooldowns},
@ -10,7 +12,7 @@ use litellm_redis_store::{
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, connect_observed, within};
use support::{connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const TTL: Duration = Duration::from_secs(60);
@ -44,7 +46,7 @@ fn team(id: &str) -> SpendCounter {
}
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
Arc::new(SpendCounters::new(client, TTL))
Arc::new(SpendCounters::new(client.map(remote)))
}
/// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, which checks the same keys and
@ -142,7 +144,7 @@ async fn reads_values_python_wrote_and_a_clean_miss_is_a_redis_answer(
.map(|index| team(&format!("t{index}")))
.collect::<Vec<_>>();
for ((bytes, _), counter) in READS.iter().zip(&counters) {
raw.set(&counter.key(), *bytes, TTL).await.unwrap();
raw.send(commands::set(&counter.key(), *bytes, TTL)).await.unwrap();
}
let expected = READS
.iter()
@ -152,7 +154,7 @@ async fn reads_values_python_wrote_and_a_clean_miss_is_a_redis_answer(
for (counter, value) in counters.iter().zip(&expected) {
assert_eq!(store.read(counter).await.unwrap(), Reading::Redis(*value));
}
assert_eq!(store.read_many(&counters).await, Some(expected));
assert_eq!(store.read_many(RedisSend::Now, &counters).unwrap().await.unwrap(), Some(expected));
assert_eq!(
store.read(&team("absent")).await.unwrap(),
Reading::Redis(None)
@ -172,15 +174,15 @@ async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl(
for (spend, bytes) in SETS {
store.set(&counter, spend, None).await.unwrap();
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
raw.send(commands::get::<String>(&counter.key())).await.unwrap(),
Some(bytes.to_owned())
);
}
for (spend, bytes) in SET_MAXES {
raw.del(&[&counter.key()]).await.unwrap();
raw.send(commands::del(&[&counter.key()])).await.unwrap();
let repaired = store.set_max(&counter, spend).await.unwrap();
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
raw.send(commands::get::<String>(&counter.key())).await.unwrap(),
Some(bytes.to_owned())
);
assert_eq!(repaired, Some(Number::Float(spend.as_f64())));
@ -189,7 +191,7 @@ async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl(
.set(&counter, Number::Int(5), Some(Duration::from_secs(30)))
.await
.unwrap();
let Ttl::Expires(remaining) = raw.ttl(&counter.key()).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
panic!("a reset counter has no TTL");
};
assert!(within(remaining, 30), "{remaining:?}");
@ -215,10 +217,10 @@ async fn set_max_only_ever_raises_a_counter(
assert_eq!(lower, Some(Number::Float(9.5)));
assert_eq!(higher, Some(Number::Float(12.0)));
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
raw.send(commands::get::<String>(&counter.key())).await.unwrap(),
Some("12.0".to_owned())
);
let Ttl::Expires(remaining) = raw.ttl(&counter.key()).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
panic!("a repaired counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
@ -243,7 +245,7 @@ async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory(
.increment(&first, Number::Float(1.25))
.await
.unwrap();
raw.expire(&first.key(), Duration::from_secs(5))
raw.send(commands::expire(&first.key(), Duration::from_secs(5)))
.await
.unwrap();
@ -257,12 +259,12 @@ async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory(
assert_eq!(totals, [Number::Float(1.75), Number::Float(2.0)]);
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
let Ttl::Expires(remaining) = raw.ttl(&first.key()).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&first.key())).await.unwrap() else {
panic!("an incremented counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
raw.get::<String>(&first.key()).await.unwrap(),
raw.send(commands::get::<String>(&first.key())).await.unwrap(),
Some("1.75".to_owned())
);
assert_eq!(store.get_local(&first).unwrap(), Some(Number::Float(1.75)));
@ -293,7 +295,7 @@ async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_an
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
assert_eq!(mirrored, Number::Float(3.5));
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
raw.send(commands::get::<String>(&counter.key())).await.unwrap(),
Some("3.5".to_owned())
);
}
@ -314,7 +316,7 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
request.flush_pre().await;
assert!(trips.seen().is_empty());
assert_eq!(raw.get::<String>(&counter.key()).await.unwrap(), None);
assert_eq!(raw.send(commands::get::<String>(&counter.key())).await.unwrap(), None);
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(100)).await;
@ -325,7 +327,7 @@ async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
};
assert_eq!(trips.seen(), [increment]);
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
raw.send(commands::get::<String>(&counter.key())).await.unwrap(),
Some("3.5".to_owned())
);
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Float(3.5)));
@ -339,7 +341,7 @@ async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = RouterCooldowns::new(Some(client.clone()), router_cooldowns::POLICY);
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(Some(remote(client.clone())), router_cooldowns::POLICY));
let spend = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counters = [
@ -349,13 +351,13 @@ async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
},
team("t"),
];
raw.set(&counters[0].key(), "4.5", TTL).await.unwrap();
raw.send(commands::set(&counters[0].key(), "4.5", TTL)).await.unwrap();
let active = cooldowns
.declare_active(&request, &["dep-a".to_owned()])
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let budget = spend
.declare_read_many_pre_call(&request, &counters)
.read_many(RedisSend::PreCall(&request), &counters)
.unwrap();
request.flush_pre().await;
@ -364,9 +366,9 @@ async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
Topology::Cluster => vec![Operation::Command("MGET"); 2],
};
assert_eq!(trips.seen(), expected);
assert_eq!(cooldowns.resolve_active(active).await.unwrap(), vec![]);
assert_eq!(active.await.unwrap(), vec![]);
assert_eq!(
budget.resolve().await,
budget.await.unwrap(),
Some(vec![Some(Number::Float(4.5)), None])
);
assert_eq!(trips.seen().len(), expected.len());
@ -394,10 +396,10 @@ async fn a_seed_lets_the_first_pod_win_and_others_answer_with_its_value(
assert_eq!([won, lost], [Number::Float(10.0), Number::Float(10.0)]);
assert_eq!(
raw.get::<String>(&counter.key()).await.unwrap(),
raw.send(commands::get::<String>(&counter.key())).await.unwrap(),
Some("10.0".to_owned())
);
let Ttl::Expires(remaining) = raw.ttl(&counter.key()).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counter.key())).await.unwrap() else {
panic!("a seeded counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
@ -419,12 +421,12 @@ async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_hel
for counter in &counters {
store.increment(counter, Number::Int(1)).await.unwrap();
}
raw.expire(&counters[0].key(), Duration::from_secs(5))
raw.send(commands::expire(&counters[0].key(), Duration::from_secs(5)))
.await
.unwrap();
assert!(store.refresh_ttl(&counters[0]).await.unwrap());
let Ttl::Expires(remaining) = raw.ttl(&counters[0].key()).await.unwrap() else {
let Ttl::Expires(remaining) = raw.send(commands::ttl(&counters[0].key())).await.unwrap() else {
panic!("a refreshed counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
@ -432,7 +434,7 @@ async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_hel
store.delete(&counters).await.unwrap();
for counter in &counters {
assert_eq!(raw.get::<String>(&counter.key()).await.unwrap(), None);
assert_eq!(raw.send(commands::get::<String>(&counter.key())).await.unwrap(), None);
assert_eq!(store.get_local(counter).unwrap(), None);
}
assert!(!store.refresh_ttl(&counters[0]).await.unwrap());
@ -454,7 +456,7 @@ async fn a_redis_outage_reads_memory_and_a_failed_increment_drops_the_counter()
assert_eq!(read, Reading::Memory(Some(Number::Float(2.0))));
assert!(increment.is_err());
assert_eq!(store.get_local(&failed).unwrap(), None);
assert_eq!(store.read_many(&[kept]).await, None);
assert_eq!(store.read_many(RedisSend::Now, &[kept]).unwrap().await.unwrap(), None);
}
#[rstest]
@ -503,7 +505,7 @@ async fn without_redis_counters_live_in_memory() {
assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]));
assert!(
store
.declare_read_many_pre_call(&request, &[counter])
.read_many(RedisSend::PreCall(&request), &[counter])
.is_none()
);
}
@ -527,7 +529,7 @@ async fn spend_counters_send_the_commands_main_sends(
.increment(&counters[0], Number::Float(1.5))
.await
.unwrap();
store.read_many(&counters).await.unwrap();
store.read_many(RedisSend::Now, &counters).unwrap().await.unwrap().unwrap();
store
.set_max(&counters[0], Number::Float(9.0))
.await

10
litellm-rust/crates/redis-store/tests/support/mod.rs
View file

@ -11,7 +11,7 @@ use litellm_redis_client::{
RedisEvents, RedisUrl, Target, Timeouts,
};
use litellm_redis_store::{
Number,
Number, Remote,
router_cooldowns::Cooldown,
router_usage::{Counter, Kind},
};
@ -21,6 +21,14 @@ pub fn connect(deployment: &RedisDeployment) -> RedisClient {
connect_with(deployment, None, None)
}
/// `client` with `RedisCache`'s default TTL when `litellm.default_redis_ttl` is unset.
pub fn remote(client: RedisClient) -> Remote {
Remote {
client,
default_ttl: Duration::from_secs(60),
}
}
/// Every round trip a client made, as the client reports it.
#[derive(Default)]
pub struct RoundTrips(Mutex<Vec<Operation>>);

63
litellm/caching/redis_batch.py
View file

@ -19,8 +19,9 @@ from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence
from contextvars import ContextVar, Token
from dataclasses import dataclass, field
from datetime import timedelta
from enum import Enum, auto
from types import MappingProxyType, TracebackType
from typing import TYPE_CHECKING, Final, Generic, Protocol, TypeAlias, TypeVar
from typing import TYPE_CHECKING, Final, Generic, Literal, Protocol, TypeAlias, TypeVar
from litellm._internal_context import (
REDIS_FAMILIES_METADATA_KEY,
@ -39,7 +40,7 @@ from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.types.services import ServiceTypes
if TYPE_CHECKING:
from litellm.rust_bridge._native import NativeRequestRedis
from litellm.rust_bridge._native import NativeRequestRedis, NativeSend
_T = TypeVar("_T")
_T_co: Final = TypeVar("_T_co", covariant=True)
@ -478,6 +479,30 @@ def _native_request(post_call_deadline: float) -> NativeRequestRedis:
return NativeRequestRedis(post_call_deadline)
class Batch(Enum):
"""Which of the request's batches a store op rides: none, so it goes out on its own when awaited, the pre-call
batch of reads the request needs before answering, or the post-call batch of writes nobody waits on. Outside a
request every op goes out on its own."""
NONE = auto()
PRE_CALL = auto()
POST_CALL = auto()
RequestPhase: TypeAlias = Literal[Batch.PRE_CALL, Batch.POST_CALL]
NativeSendFactory: TypeAlias = Callable[[RequestPhase, "NativeRequestRedis"], "NativeSend"]
def _native_send(phase: RequestPhase, request: NativeRequestRedis) -> NativeSend:
from litellm.rust_bridge._native import NativeSend
match phase:
case Batch.PRE_CALL:
return NativeSend.PreCall(request)
case Batch.POST_CALL:
return NativeSend.PostCall(request)
_open_post_call: Final[weakref.WeakSet[RequestRedisBatches]] = weakref.WeakSet()
"""Requests whose post-call batch still holds declared ops, so a shutdown can send them before Redis goes away."""
@ -498,6 +523,7 @@ class RequestRedisBatches:
"_native",
"_native_factory",
"_native_flush_hooks",
"_native_send_factory",
"_native_misses",
"_post_call",
"post_call_deadline",
@ -508,11 +534,13 @@ class RequestRedisBatches:
self,
post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS,
native_factory: NativeRequestFactory = _native_request,
native_send_factory: NativeSendFactory = _native_send,
) -> None:
self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
self.post_call_deadline: Final = post_call_deadline
self._native_factory: Final = native_factory
self._native_send_factory: Final = native_send_factory
self._native: NativeRequestRedis | None = None
self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry
self._native_misses: frozenset[str] = frozenset()
@ -553,11 +581,20 @@ class RequestRedisBatches:
_open_post_call.add(self)
return self.native()
def native_send(self, phase: RequestPhase) -> NativeSend:
"""The ``NativeSend`` for ``phase``; a post-call one registers the request with the shutdown drain."""
match phase:
case Batch.PRE_CALL:
return self._native_send_factory(phase, self.native())
case Batch.POST_CALL:
return self._native_send_factory(phase, self.native_post_call())
def native_result(self, pending: NativePending[_T]) -> NativeBatchResult[_T]:
return NativeBatchResult(self, pending)
def add_native_flush_hook(self, hook: Callable[[], None]) -> None:
"""Called before every native pre-call flush so late readers can declare into the same trip."""
def add_pre_call_flush_hook(self, hook: Callable[[], None]) -> None:
"""Called before every pre-call flush so late readers can declare into the same trip. The hook runs with this
request active, so a ``Batch.PRE_CALL`` op it declares rides this request whoever triggered the flush."""
self._native_flush_hooks.append(hook)
def note_native_misses(self, keys: Sequence[str]) -> None:
@ -571,8 +608,12 @@ class RequestRedisBatches:
return key in self._native_misses
async def flush_native_pre(self) -> None:
for hook in self._native_flush_hooks:
hook()
token: Final = _active_request_batches.set(self)
try:
for hook in self._native_flush_hooks:
hook()
finally:
_active_request_batches.reset(token)
if self._native is not None:
await self._native.flush_pre()
@ -581,7 +622,7 @@ class RequestRedisBatches:
self._deadline_flush = asyncio.ensure_future(self.flush_post_call())
async def flush_all(self) -> None:
"""Send whatever is still declared (write-backs nobody awaits) before the request scope closes."""
"""Batch whatever is still declared (write-backs nobody awaits) before the request scope closes."""
await asyncio.gather(
*(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre()
)
@ -643,23 +684,27 @@ async def drain_post_call_redis_batches() -> None:
class request_redis_batch_scope:
"""Redis reads declared inside share one pipeline per backend; nested scopes join the outer one."""
__slots__ = ("_native_factory", "_post_call_deadline", "_token")
__slots__ = ("_native_factory", "_native_send_factory", "_post_call_deadline", "_token")
def __init__(
self,
post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS,
native_factory: NativeRequestFactory = _native_request,
native_send_factory: NativeSendFactory = _native_send,
) -> None:
self._token: Token[RequestRedisBatches | None] | None = None
self._post_call_deadline: Final = post_call_deadline
self._native_factory: Final = native_factory
self._native_send_factory: Final = native_send_factory
def __enter__(self) -> RequestRedisBatches:
outer: Final = _active_request_batches.get()
if outer is not None:
return outer
batches: Final = RequestRedisBatches(
post_call_deadline=self._post_call_deadline, native_factory=self._native_factory
post_call_deadline=self._post_call_deadline,
native_factory=self._native_factory,
native_send_factory=self._native_send_factory,
)
self._token = _active_request_batches.set(batches)
return batches

7
litellm/proxy/auth/auth_object_prefetch.py
View file

@ -14,6 +14,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
from litellm._internal_context import service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_batch import Batch
from litellm.constants import DEFAULT_IN_MEMORY_TTL, REGISTRY_ERROR_NEGATIVE_CACHE_TTL
from litellm.models.organization import LiteLLM_OrganizationTable
from litellm.models.team import LiteLLM_TeamTableCachedObj
@ -222,7 +223,7 @@ async def _fill_from_redis(
before_read()
try:
with service_target(AUTH_OBJECTS_TARGET):
found: Final = await shared.read_missing(sorted(entry.cache_key for entry in entries))
found: Final = await shared.read_missing(sorted(entry.cache_key for entry in entries), batch=Batch.PRE_CALL)
except Exception as e: # noqa: BLE001 # the DB fill below takes over, as it does after a failed MGET
verbose_proxy_logger.debug("auth prefetch Redis read failed, filling from the database: %s", e)
return
@ -272,7 +273,7 @@ async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: U
shared: Final = cache.shared
if shared is not None:
with service_target(AUTH_OBJECTS_TARGET):
await shared.write_back(payloads)
await shared.write_back(payloads, batch=Batch.PRE_CALL)
return
for cache_key, payload, ttl in payloads:
cache.set_local(cache_key, payload, ttl)
@ -339,7 +340,7 @@ async def prefetch_identity_keys(
return
before_read()
with service_target(AUTH_OBJECTS_TARGET):
found: Final = await shared.read_missing(sorted(missing))
found: Final = await shared.read_missing(sorted(missing), batch=Batch.PRE_CALL)
management_ttl: Final = get_management_object_ttl(user_api_key_cache)
except Exception as e: # noqa: BLE001 # warm-up only; the getters read Redis and the database on their own
verbose_proxy_logger.warning("auth identity prefetch skipped, falling back to per-key lookups: %s", e)

4
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -353,7 +353,7 @@ class UserApiKeyCache(DualCache):
ttl: Final = _ttl(kwargs)
if local_only:
shared.set_local(key, payload, ttl)
elif not pipelined or shared.set_pre_call(key, payload, ttl) is None:
elif not pipelined or shared.declare_set(key, payload, ttl) is None:
await shared.set(key, payload, ttl)
except Exception as e: # noqa: BLE001 # DualCache.async_set_cache swallows a failed write
verbose_proxy_logger.error("UserApiKeyCache.async_set_cache failed for key=%r: %s", key, e)
@ -378,7 +378,7 @@ class UserApiKeyCache(DualCache):
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
if self._shared is not None:
return self._shared.delete_pre_call(key)
return self._shared.declare_delete(key)
if is_user_key_cache_key(key):
return await self.key_object_cache.async_delete_cache_pre_call(key)
return await super().async_delete_cache_pre_call(key)

43
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -34,6 +34,7 @@ from typing_extensions import NotRequired, ReadOnly
from litellm import DualCache
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_batch import Batch
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import (
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE,
@ -1210,7 +1211,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
async def _read_counter_values_from_redis(self, counters: Sequence[RateLimitCounter]) -> list[CounterState]:
groups: Final = self._slot_groups(counters)
group_values: Final = [
await self.rate_limits.read_counters(tuple(counters[index] for index in group)) for group in groups
await self.rate_limits.read_counters(tuple(counters[index] for index in group), batch=Batch.NONE)
for group in groups
]
by_index: Final = {
index: (values[offset * 2], values[offset * 2 + 1])
@ -1256,14 +1258,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
slot_groups: Final = self._slot_groups(counters)
groups: Final = tuple(tuple(counters[index] for index in group) for group in slot_groups)
states_by_group: Final[list[Sequence[CounterState]]] = [] # mutable-ok: one entry appended per group
batched: Final = self.rate_limits.has_pre_call_batch
pipelined: Final = tuple(
self.rate_limits.declare_sliding_window(group, now_int, self.window_size) for group in groups
self.rate_limits.sliding_window(group, now_int, self.window_size, batch=Batch.PRE_CALL) if batched else None
for group in groups
)
for index, (group, group_result) in enumerate(zip(groups, pipelined)):
try:
counts = (
await self.rate_limits.sliding_window(group, now_int, self.window_size)
await self.rate_limits.sliding_window(group, now_int, self.window_size, batch=Batch.NONE)
if group_result is None
else await group_result
)
@ -1505,7 +1509,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
counts = [
max(0, value)
for value in await self.rate_limits.parallel_count(
gauge_scopes, PARALLEL_REQUEST_SLOT_TTL_SECONDS
gauge_scopes, PARALLEL_REQUEST_SLOT_TTL_SECONDS, batch=Batch.NONE
)
]
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror unless fail-closed rejects
@ -1539,6 +1543,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
tuple((gauge["scope"], gauge["limit"]) for gauge in gauges),
PARALLEL_REQUEST_SLOT_TTL_SECONDS,
slot_id,
batch=Batch.NONE,
)
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500
self._reject_if_rate_limit_unverifiable("parallel_acquire_script", e)
@ -1652,7 +1657,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return
if self.rate_limits.has_redis:
try:
remaining: Final = await self.rate_limits.parallel_release(gauges, slot_id)
remaining: Final = await self.rate_limits.parallel_release(gauges, slot_id, batch=Batch.NONE)
self._mirror_released_parallel_slots(gauges, remaining)
return
except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500
@ -1678,7 +1683,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if not self.rate_limits.has_post_call_batch or not gauges or not slot_id:
return False
await self._release_parallel_request_slots_in_memory(gauges, slot_id)
return await self.rate_limits.parallel_release_post_call(gauges, slot_id)
return self.rate_limits.parallel_release_post_call(gauges, slot_id)
def _mirror_released_parallel_slots(
self, gauges: Sequence[RateLimitScope], remaining_by_gauge: Sequence[int]
@ -1849,16 +1854,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop
reply: AtomicReply
pipelined: Final = tuple(
self.rate_limits.declare_check_and_increment(counters) for counters, _meta in descriptor_groups
)
batched: Final = tuple(result for result in pipelined if result is not None)
if len(batched) == len(descriptor_groups):
if self.rate_limits.has_pre_call_batch:
batched: Final = tuple(
self.rate_limits.check_and_increment(counters, batch=Batch.PRE_CALL)
for counters, _meta in descriptor_groups
)
return await self._settle_pipelined_descriptor_groups(descriptor_groups, batched, parent_otel_span)
for counters, meta in descriptor_groups:
try:
reply = await self.rate_limits.check_and_increment(counters)
reply = await self.rate_limits.check_and_increment(counters, batch=Batch.NONE)
except Exception as e:
# Lua failure (timeout, OOM, network partition) leaves Redis
# state ambiguous. Refund any prior groups so Redis returns
@ -2002,7 +2007,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return
for refund in refunds:
try:
await self.rate_limits.guarded_increment(refund.counter, refund.window_start, -refund.increment, 0)
await self.rate_limits.guarded_increment(
refund.counter, refund.window_start, -refund.increment, 0, batch=Batch.NONE
)
except Exception as e: # noqa: BLE001 # best-effort rollback, the rejection already decided the request
log_redis_failure(
verbose_proxy_logger,
@ -4094,9 +4101,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
operation.amount,
operation.ttl or 0,
)
if await self.rate_limits.token_increment_post_call(group_operations):
if self.rate_limits.token_increment_post_call(group_operations):
continue
await self.rate_limits.token_increment(group_operations)
await self.rate_limits.token_increment(group_operations, batch=Batch.NONE)
@with_service_target("rate_limits")
async def async_increment_tokens_with_ttl_preservation(
@ -4151,7 +4158,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if self.rate_limits.has_redis:
try:
await self.rate_limits.guarded_increment(
increment.counter, operation.expected_window_start, increment.amount, increment.ttl
increment.counter,
operation.expected_window_start,
increment.amount,
increment.ttl,
batch=Batch.NONE,
)
continue
except Exception as e: # noqa: BLE001 # Redis failures use the plain increment fallback

30
litellm/proxy/spend_tracking/spend_counter_batch.py
View file

@ -1,14 +1,14 @@
"""One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each."""
import asyncio
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Awaitable, Iterator, Mapping, Sequence
from contextvars import ContextVar, Token
from dataclasses import dataclass
from types import MappingProxyType, TracebackType
from typing import Final
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_batch import NativeBatchResult, RequestRedisBatches, active_request_redis_batches
from litellm.caching.redis_batch import Batch, RequestRedisBatches, active_request_redis_batches
from litellm.proxy._types import UserAPIKeyAuth
from litellm.redis.stores.spend_counters import (
SpendCounterStore,
@ -28,8 +28,7 @@ _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
@dataclass(frozen=True, slots=True)
class _InflightRead:
counters: tuple[str, ...]
result: NativeBatchResult[Sequence[float | None] | None]
result: Awaitable[Mapping[str, float | None] | None]
@dataclass(frozen=True, slots=True)
@ -43,9 +42,9 @@ class SpendCounterBatch:
``async_batch_get_cache`` maps a clean miss to ``None`` and drops keys only when Redis failed, so an absent
key means "read it yourself" and a present ``None`` is an authoritative miss.
Inside a ``request_redis_batch_scope`` the read rides the request's native pre-call pipeline instead: the
native flush hook declares whatever is bound but unread, so whichever native read the request awaits first
(the armed cooldown read, or this batch's own) carries the spend counters in the same round trip."""
Inside a ``request_redis_batch_scope`` the read rides the request's pre-call pipeline instead: the pre-call
flush hook declares whatever is bound but unread, so whichever pre-call read the request awaits first (the
armed cooldown read, or this batch's own) carries the spend counters in the same round trip."""
__slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_request", "_store")
@ -59,7 +58,7 @@ class SpendCounterBatch:
self._inflight: Final[list[_InflightRead]] = [] # mutable-ok: drained by _load
self._request: Final[RequestRedisBatches | None] = active_request_redis_batches()
if self._request is not None:
self._request.add_native_flush_hook(self._declare_pending)
self._request.add_pre_call_flush_hook(self._declare_pending)
@property
def counter_keys(self) -> frozenset[str]:
@ -116,32 +115,29 @@ class SpendCounterBatch:
return self._loaded
def _declare_pending(self) -> None:
"""Native flush hook: put every bound-but-unread counter on the native pre-call pipeline about to go out."""
"""Pre-call flush hook: put every bound-but-unread counter on the pre-call pipeline about to go out."""
if self._request is None or not self._open:
return
pending: Final = tuple(sorted(self._keys - self._fetched))
if not pending:
return
result: Final = self._store.declare_read_many(self._request, pending)
if result is None:
if not pending or not self._store.has_redis:
return
self._fetched = self._fetched | frozenset(pending)
self._inflight.append(_InflightRead(counters=pending, result=result))
self._inflight.append(_InflightRead(result=self._store.read_many(pending, batch=Batch.PRE_CALL)))
async def _collect_inflight(self) -> None:
reads: Final = tuple(self._inflight)
self._inflight.clear()
for read in reads:
try:
values: Sequence[float | None] | None = await read.result
values: Mapping[str, float | None] | None = await read.result
except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback
verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e)
continue
if values is not None:
self._loaded = MappingProxyType({**dict(zip(read.counters, values, strict=True)), **self._loaded})
self._loaded = MappingProxyType({**values, **self._loaded})
async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]:
fetched: Final = await self._store.read_many(sorted(keys))
fetched: Final = await self._store.read_many(sorted(keys), batch=Batch.NONE)
return _NO_VALUES if fetched is None else fetched

37
litellm/redis/stores/_route.py Normal file
View file

@ -0,0 +1,37 @@
from __future__ import annotations
from collections.abc import Awaitable
from dataclasses import dataclass
from typing import Final, TypeVar
from litellm.caching.redis_batch import Batch, NativePending, RequestRedisBatches, active_request_redis_batches
from litellm.rust_bridge._native import NativeRequestRedis, NativeSend
_T = TypeVar("_T")
@dataclass(frozen=True, slots=True)
class Route:
"""The batch a native store op rides, and how its reply is awaited. A pre-call read awaits through
``native_result``, so the pre-call flush hooks run first and late readers join the same trip."""
send: NativeSend
request: RequestRedisBatches | None
def result(self, pending: NativePending[_T]) -> Awaitable[_T]:
return pending.resolve() if self.request is None else self.request.native_result(pending)
def route(batch: Batch) -> Route:
request: Final = active_request_redis_batches()
match batch:
case Batch.NONE:
return Route(NativeSend.Now(), None)
case Batch.PRE_CALL | Batch.POST_CALL:
return Route(NativeSend.Now(), None) if request is None else Route(request.native_send(batch), request)
def post_call_request() -> NativeRequestRedis | None:
"""The current request's native batches for a post-call write nobody awaits, None outside a request."""
request: Final = active_request_redis_batches()
return None if request is None else request.native_post_call()

192
litellm/redis/stores/auth_cache.py
View file

@ -1,192 +0,0 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import Final, Protocol, TypeAlias
from litellm.caching.redis_batch import (
NativeBatchResult,
NativePending,
active_request_redis_batches,
)
from litellm.rust_bridge._native import NativeAuthCache, NativeRedisClient, NativeRequestRedis
WriteBack: TypeAlias = tuple[str, object, float | None]
@dataclass(frozen=True, slots=True)
class AuthCacheSettings:
"""``UserApiKeyCache``'s memory configuration: each partition's size and ``default_in_memory_ttl``."""
memory_entries: int
memory_ttl: float | None
class AuthCacheHandle(Protocol):
def get(self, key: str) -> Awaitable[object]: ...
def get_sync(self, key: str) -> object: ...
def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Awaitable[None]: ...
def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
def set_pre_call(
self, request: NativeRequestRedis, key: str, value: object, ttl_seconds: float | None = None
) -> NativePending[None] | None: ...
def delete_pre_call(self, request: NativeRequestRedis, key: str) -> NativePending[None] | None: ...
def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Awaitable[None]: ...
def write_back(self, request: NativeRequestRedis | None, entries: Sequence[WriteBack]) -> Awaitable[None]: ...
def delete(self, keys: Sequence[str]) -> Awaitable[None]: ...
def delete_sync(self, keys: Sequence[str]) -> None: ...
def declare_get_many(
self, request: NativeRequestRedis, keys: Sequence[str]
) -> NativePending[Sequence[object]] | None: ...
def get_many(self, keys: Sequence[str]) -> Awaitable[Sequence[object] | None]: ...
def increment_if_present(self, key: str, amount: float) -> Awaitable[float | None]: ...
def get_local(self, key: str) -> object: ...
def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
def delete_local(self, key: str) -> None: ...
def flush_local(self) -> None: ...
def len_local(self) -> tuple[int, int]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, AuthCacheSettings, int], AuthCacheHandle]
def native_handle(client: NativeRedisClient | None, settings: AuthCacheSettings, redis_ttl: int) -> AuthCacheHandle:
return NativeAuthCache(
client,
memory_entries=settings.memory_entries,
memory_ttl_seconds=settings.memory_ttl,
redis_ttl_seconds=redis_ttl,
)
class AuthCacheStore:
"""The auth object cache ``enable_redis_auth_cache`` shares across workers: memory partitions in front of Redis,
all native. Keys are ``UserApiKeyCache``'s logical cache keys; values are plain data (what ``json.loads``
returns), so callers validate models themselves. Redis calls are reported under the caller's service target.
Inside a request, a key a prefetch read found absent in Redis is answered as absent by ``get`` until something
writes it, so per-key getters after a prefetch cost no round trip."""
def __init__(
self,
client: NativeRedisClient | None,
settings: AuthCacheSettings,
*,
redis_ttl: int,
handle: HandleFactory = native_handle,
) -> None:
self._client: Final = client
self._redis_ttl: Final = redis_ttl
self._build: Final = handle
self._settings = settings
self._handle = handle(client, settings, redis_ttl)
@property
def settings(self) -> AuthCacheSettings:
return self._settings
def reconfigure(self, settings: AuthCacheSettings) -> None:
"""A changed memory configuration starts from an empty memory tier, as a resized partition would."""
if settings == self._settings:
return
self._settings = settings
self._handle = self._build(self._client, settings, self._redis_ttl)
async def get(self, key: str) -> object:
request: Final = active_request_redis_batches()
if request is not None and request.native_read_as_missing(key) and self._handle.get_local(key) is None:
return None
return await self._handle.get(key)
def get_sync(self, key: str) -> object:
return self._handle.get_sync(key)
async def set(self, key: str, value: object, ttl: float | None) -> None:
self._forget_miss(key)
await self._handle.set(key, value, ttl)
def set_sync(self, key: str, value: object, ttl: float | None) -> None:
self._forget_miss(key)
self._handle.set_sync(key, value, ttl)
def set_pre_call(self, key: str, value: object, ttl: float | None) -> NativeBatchResult[None] | None:
"""Memory now and the SET on the request's native pre-call batch, sent with its next flush; None outside a
request, so the caller writes through."""
request: Final = active_request_redis_batches()
if request is None:
return None
pending: Final = self._handle.set_pre_call(request.native(), key, value, ttl)
if pending is None:
return None
request.forget_native_miss(key)
return request.native_result(pending)
def delete_pre_call(self, key: str) -> NativeBatchResult[None] | None:
request: Final = active_request_redis_batches()
if request is None:
return None
pending: Final = self._handle.delete_pre_call(request.native(), key)
if pending is None:
return None
request.note_native_misses((key,))
return request.native_result(pending)
async def set_many(self, entries: Sequence[tuple[str, object]], ttl: float | None) -> None:
for key, _value in entries:
self._forget_miss(key)
await self._handle.set_many(entries, ttl)
async def write_back(self, entries: Sequence[WriteBack]) -> None:
"""Memory now; inside a request the SETs ride its native pre-call batch, sent with the next flush, else they
go out in a round trip of their own. An entry without a TTL keeps ``default_in_memory_ttl`` in memory and the
connection's default TTL in Redis."""
request: Final = active_request_redis_batches()
for key, _value, _ttl in entries:
self._forget_miss(key)
await self._handle.write_back(None if request is None else request.native(), entries)
async def delete(self, keys: Sequence[str]) -> None:
"""Memory, then Redis in one round trip; raises when Redis failed, with memory already dropped."""
await self._handle.delete(keys)
def delete_sync(self, keys: Sequence[str]) -> None:
self._handle.delete_sync(keys)
async def read_missing(self, keys: Sequence[str]) -> dict[str, object]:
"""One read of ``keys`` (already missing in memory): on the request's native pre-call batch inside a request,
noting the keys Redis lacks, else a round trip of its own. Absent keys map to None; raises when Redis failed."""
request: Final = active_request_redis_batches()
if request is None:
values: Final = await self._handle.get_many(keys)
return {} if values is None else dict(zip(keys, values, strict=True))
pending: Final = self._handle.declare_get_many(request.native(), keys)
if pending is None:
return {}
found: Final = dict(zip(keys, await request.native_result(pending), strict=True))
request.note_native_misses(tuple(key for key, value in found.items() if value is None))
return found
async def increment_if_present(self, key: str, amount: float) -> float | None:
return await self._handle.increment_if_present(key, amount)
def get_local(self, key: str) -> object:
return self._handle.get_local(key)
def set_local(self, key: str, value: object, ttl: float | None) -> None:
self._handle.set_local(key, value, ttl)
def delete_local(self, key: str) -> None:
self._handle.delete_local(key)
def flush_local(self) -> None:
self._handle.flush_local()
def len_local(self) -> tuple[int, int]:
"""(management objects, user-key objects) held in memory."""
return self._handle.len_local()
@staticmethod
def _forget_miss(key: str) -> None:
request: Final = active_request_redis_batches()
if request is not None:
request.forget_native_miss(key)

145
litellm/redis/stores/auth_cache/__init__.py Normal file
View file

@ -0,0 +1,145 @@
from __future__ import annotations
from collections.abc import Awaitable, Sequence
from typing import Final
from litellm.caching.redis_batch import Batch, active_request_redis_batches
from litellm.redis.stores._route import route
from litellm.redis.stores.auth_cache.handle import AuthCacheHandle, HandleFactory, native_handle
from litellm.redis.stores.auth_cache.types import AuthCacheSettings, WriteBack
from litellm.rust_bridge._native import NativeRedisClient
__all__ = [
"AuthCacheHandle",
"AuthCacheSettings",
"AuthCacheStore",
"HandleFactory",
"WriteBack",
"native_handle",
]
class AuthCacheStore:
"""The auth object cache ``enable_redis_auth_cache`` shares across workers: memory partitions in front of Redis,
all native. Keys are ``UserApiKeyCache``'s logical cache keys; values are plain data (what ``json.loads``
returns), so callers validate models themselves. Redis calls are reported under the caller's service target.
Inside a request, a key a prefetch read found absent in Redis is answered as absent by ``get`` until something
writes it, so per-key getters after a prefetch cost no round trip."""
def __init__(
self,
client: NativeRedisClient | None,
settings: AuthCacheSettings,
*,
redis_ttl: int,
handle: HandleFactory = native_handle,
) -> None:
self._client: Final = client
self._redis_ttl: Final = redis_ttl
self._build: Final = handle
self._settings = settings
self._handle = handle(client, settings, redis_ttl)
@property
def settings(self) -> AuthCacheSettings:
return self._settings
def reconfigure(self, settings: AuthCacheSettings) -> None:
"""A changed memory configuration starts from an empty memory tier, as a resized partition would."""
if settings == self._settings:
return
self._settings = settings
self._handle = self._build(self._client, settings, self._redis_ttl)
async def get(self, key: str) -> object:
request: Final = active_request_redis_batches()
if request is not None and request.native_read_as_missing(key) and self._handle.get_local(key) is None:
return None
return await self._handle.get(key)
def get_sync(self, key: str) -> object:
return self._handle.get_sync(key)
async def set(self, key: str, value: object, ttl: float | None) -> None:
self._forget_miss(key)
await self._handle.set(key, value, ttl)
def set_sync(self, key: str, value: object, ttl: float | None) -> None:
self._forget_miss(key)
self._handle.set_sync(key, value, ttl)
def declare_set(self, key: str, value: object, ttl: float | None) -> Awaitable[None] | None:
"""Memory now and the SET on the request's pre-call batch, sent with its next flush; None outside a request,
so the caller writes through."""
target: Final = route(Batch.PRE_CALL)
if target.request is None:
return None
pending: Final = self._handle.declare_set(key, value, ttl, target.send)
target.request.forget_native_miss(key)
return target.result(pending)
def declare_delete(self, key: str) -> Awaitable[None] | None:
target: Final = route(Batch.PRE_CALL)
if target.request is None:
return None
pending: Final = self._handle.declare_delete(key, target.send)
target.request.note_native_misses((key,))
return target.result(pending)
async def set_many(self, entries: Sequence[tuple[str, object]], ttl: float | None) -> None:
for key, _value in entries:
self._forget_miss(key)
await self._handle.set_many(entries, ttl)
async def write_back(self, entries: Sequence[WriteBack], *, batch: Batch) -> None:
"""Memory now; the SETs ride the batch ``batch`` picks, so a pre-call write inside a request goes out with the
next flush, else in a round trip of its own. An entry without a TTL keeps ``default_in_memory_ttl`` in memory and the
connection's default TTL in Redis."""
for key, _value, _ttl in entries:
self._forget_miss(key)
await self._handle.write_back(entries, route(batch).send)
async def delete(self, keys: Sequence[str]) -> None:
"""Memory, then Redis in one round trip; raises when Redis failed, with memory already dropped."""
await self._handle.delete(keys)
def delete_sync(self, keys: Sequence[str]) -> None:
self._handle.delete_sync(keys)
async def read_missing(self, keys: Sequence[str], *, batch: Batch) -> dict[str, object]:
"""One read of ``keys`` (already missing in memory) on the batch ``batch`` picks, noting the keys Redis lacks
when it rode a request. Absent keys map to None; raises when Redis failed."""
target: Final = route(batch)
pending: Final = self._handle.get_many(keys, target.send)
if pending is None:
return {}
found: Final = dict(zip(keys, await target.result(pending), strict=True))
if target.request is not None:
target.request.note_native_misses(tuple(key for key, value in found.items() if value is None))
return found
async def increment_if_present(self, key: str, amount: float) -> float | None:
return await self._handle.increment_if_present(key, amount)
def get_local(self, key: str) -> object:
return self._handle.get_local(key)
def set_local(self, key: str, value: object, ttl: float | None) -> None:
self._handle.set_local(key, value, ttl)
def delete_local(self, key: str) -> None:
self._handle.delete_local(key)
def flush_local(self) -> None:
self._handle.flush_local()
def len_local(self) -> tuple[int, int]:
"""(management objects, user-key objects) held in memory."""
return self._handle.len_local()
@staticmethod
def _forget_miss(key: str) -> None:
request: Final = active_request_redis_batches()
if request is not None:
request.forget_native_miss(key)

40
litellm/redis/stores/auth_cache/handle.py Normal file
View file

@ -0,0 +1,40 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Sequence
from typing import Protocol, TypeAlias
from litellm.caching.redis_batch import NativePending
from litellm.redis.stores.auth_cache.types import AuthCacheSettings, WriteBack
from litellm.rust_bridge._native import NativeAuthCache, NativeRedisClient, NativeSend
class AuthCacheHandle(Protocol):
def get(self, key: str) -> Awaitable[object]: ...
def get_sync(self, key: str) -> object: ...
def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Awaitable[None]: ...
def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ...
def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ...
def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Awaitable[None]: ...
def write_back(self, entries: Sequence[WriteBack], send: NativeSend) -> Awaitable[None]: ...
def delete(self, keys: Sequence[str]) -> Awaitable[None]: ...
def delete_sync(self, keys: Sequence[str]) -> None: ...
def get_many(self, keys: Sequence[str], send: NativeSend) -> NativePending[Sequence[object]] | None: ...
def increment_if_present(self, key: str, amount: float) -> Awaitable[float | None]: ...
def get_local(self, key: str) -> object: ...
def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
def delete_local(self, key: str) -> None: ...
def flush_local(self) -> None: ...
def len_local(self) -> tuple[int, int]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, AuthCacheSettings, int], AuthCacheHandle]
def native_handle(client: NativeRedisClient | None, settings: AuthCacheSettings, redis_ttl: int) -> AuthCacheHandle:
return NativeAuthCache(
client,
memory_entries=settings.memory_entries,
memory_ttl_seconds=settings.memory_ttl,
redis_ttl_seconds=redis_ttl,
)

14
litellm/redis/stores/auth_cache/types.py Normal file
View file

@ -0,0 +1,14 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import TypeAlias
WriteBack: TypeAlias = tuple[str, object, float | None]
@dataclass(frozen=True, slots=True)
class AuthCacheSettings:
"""``UserApiKeyCache``'s memory configuration: each partition's size and ``default_in_memory_ttl``."""
memory_entries: int
memory_ttl: float | None

404
litellm/redis/stores/rate_limits.py
View file

@ -1,404 +0,0 @@
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())))

252
litellm/redis/stores/rate_limits/__init__.py Normal file
View file

@ -0,0 +1,252 @@
from __future__ import annotations
from collections.abc import Awaitable, Sequence
from typing import Final, Never, TypeVar
from litellm.caching.redis_batch import Batch, NativePending, active_request_redis_batches
from litellm.redis.stores._route import Route, post_call_request, route
from litellm.redis.stores.rate_limits.handle import HandleFactory, RateLimitsHandle
from litellm.redis.stores.rate_limits.types import (
AcquireReply,
AtomicApplied,
AtomicCounter,
AtomicOverLimit,
AtomicReply,
CounterIncrement,
CounterKind,
GaugeOverLimit,
GaugesAcquired,
GuardedReply,
LocalValue,
NativeAtomicCounter,
NativeCounter,
NativeEntry,
NativeIncrement,
NativeOverLimit,
NativeScope,
RateLimitCounter,
RateLimitEntry,
RateLimitGauge,
RateLimitScope,
RateLimitWindow,
WindowCount,
native_atomic,
native_counter,
native_entry,
native_increment,
native_scope,
)
from litellm.rust_bridge._native import NativeRateLimits, NativeRedisClient
__all__ = [
"AcquireReply",
"AtomicApplied",
"AtomicCounter",
"AtomicOverLimit",
"AtomicReply",
"CounterIncrement",
"CounterKind",
"GaugeOverLimit",
"GaugesAcquired",
"GuardedReply",
"HandleFactory",
"LocalValue",
"NativeAtomicCounter",
"NativeCounter",
"NativeEntry",
"NativeIncrement",
"NativeScope",
"RateLimitCounter",
"RateLimitEntry",
"RateLimitGauge",
"RateLimitScope",
"RateLimitStore",
"RateLimitWindow",
"RateLimitsHandle",
"WindowCount",
]
_T = TypeVar("_T")
async def _no_redis() -> Never:
raise RuntimeError("the rate limiter has no Redis to run its scripts on")
def _script(pending: NativePending[_T] | None, target: Route) -> Awaitable[_T]:
"""A script's reply; one declared without Redis raises where it is awaited, like a failed script."""
return _no_redis() if pending is None else target.result(pending)
async def _window_counts(reply: Awaitable[list[tuple[str, int]]]) -> tuple[WindowCount, ...]:
return tuple(WindowCount(start, count) for start, count in await reply)
async def _values(reply: Awaitable[list[str | None]]) -> tuple[str | None, ...]:
return tuple(await reply)
async def _counts(reply: Awaitable[list[int]]) -> tuple[int, ...]:
return tuple(await reply)
async def _totals(reply: Awaitable[list[float]]) -> tuple[float, ...]:
return tuple(await reply)
async def _guarded(reply: Awaitable[tuple[bool, int]]) -> GuardedReply:
applied, count = await reply
return GuardedReply(applied=applied, count=count)
async def _atomic_reply(reply: Awaitable[tuple[NativeOverLimit | None, list[tuple[int, int]]]]) -> AtomicReply:
over_limit, applied = await reply
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(reply: Awaitable[tuple[NativeOverLimit | None, list[int]]]) -> AcquireReply:
over_limit, in_flight = await reply
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. Each script is declared on the batch its ``Batch`` picks when
called, so several declared before any is awaited share one trip; awaiting the result sends it."""
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_pre_call_batch(self) -> bool:
"""Whether ``Batch.PRE_CALL`` scripts share the request's pre-call trip rather than each go out on its own."""
return self.has_redis and active_request_redis_batches() is not None
@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(native_scope, scopes)), force))
def local(self, entries: Sequence[RateLimitEntry]) -> tuple[LocalValue | None, ...]:
return tuple(self._handle.get_local(native_entry(entry)) for entry in entries)
def set_local(self, entry: RateLimitEntry, value: LocalValue, ttl: float | None) -> None:
self._handle.set_local(native_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 = native_entry(entry)
local: Final = self._handle.get_local(native)
if local is not None or not self.has_redis:
return local
return await self._handle.get(native, backfill_ttl_seconds=backfill_ttl)
def sliding_window(
self, counters: Sequence[RateLimitCounter], now: int, window_size: int, *, batch: Batch
) -> Awaitable[tuple[WindowCount, ...]]:
target: Final = route(batch)
native: Final = tuple(map(native_counter, counters))
return _window_counts(_script(self._handle.sliding_window(native, now, window_size, target.send), target))
def read_counters(self, counters: Sequence[RateLimitCounter], *, batch: Batch) -> Awaitable[tuple[str | None, ...]]:
"""Each counter's window start then its count, as ``BATCH_COUNTER_READ_SCRIPT`` reads them."""
target: Final = route(batch)
native: Final = tuple(map(native_counter, counters))
return _values(_script(self._handle.read_counters(native, target.send), target))
def check_and_increment(self, counters: Sequence[AtomicCounter], *, batch: Batch) -> Awaitable[AtomicReply]:
target: Final = route(batch)
native: Final = tuple(map(native_atomic, counters))
return _atomic_reply(_script(self._handle.check_and_increment(native, target.send), target))
def guarded_increment(
self,
counter: RateLimitCounter,
expected_window_start: str,
amount: float,
ttl: int | None,
*,
batch: Batch,
) -> Awaitable[GuardedReply]:
target: Final = route(batch)
native: Final = native_counter(counter)
return _guarded(
_script(self._handle.guarded_increment(native, expected_window_start, amount, ttl, target.send), target)
)
def parallel_acquire(
self, gauges: Sequence[tuple[RateLimitScope, int]], slot_ttl: int, slot_id: str, *, batch: Batch
) -> Awaitable[AcquireReply]:
target: Final = route(batch)
native: Final = tuple((native_scope(scope), limit) for scope, limit in gauges)
return _acquire_reply(_script(self._handle.parallel_acquire(native, slot_ttl, slot_id, target.send), target))
def parallel_release(
self, gauges: Sequence[RateLimitScope], slot_id: str, *, batch: Batch
) -> Awaitable[tuple[int, ...]]:
target: Final = route(batch)
native: Final = tuple(map(native_scope, gauges))
return _counts(_script(self._handle.parallel_release(native, slot_id, target.send), target))
def parallel_count(
self, gauges: Sequence[RateLimitScope], slot_ttl: int, *, batch: Batch
) -> Awaitable[tuple[int, ...]]:
target: Final = route(batch)
native: Final = tuple(map(native_scope, gauges))
return _counts(_script(self._handle.parallel_count(native, slot_ttl, target.send), target))
def token_increment(self, increments: Sequence[CounterIncrement], *, batch: Batch) -> Awaitable[tuple[float, ...]]:
target: Final = route(batch)
native: Final = tuple(map(native_increment, increments))
return _totals(_script(self._handle.token_increment(native, target.send), target))
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 = post_call_request() if self.has_redis else None
if request is None:
return False
return self._handle.token_increment_post_call(request, tuple(map(native_increment, increments)))
def parallel_release_post_call(self, gauges: Sequence[RateLimitScope], slot_id: str) -> bool:
request: Final = post_call_request() if self.has_redis else None
if request is None:
return False
return self._handle.parallel_release_post_call(request, tuple(map(native_scope, gauges)), slot_id)
async def increment(self, increment: CounterIncrement) -> int | float:
"""``DualCache.async_increment_cache``: memory, then Redis keeping an existing TTL."""
return await self._handle.increment(native_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(native_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."""
native: Final = tuple(map(native_increment, increments))
request: Final = post_call_request() if self.has_redis else None
if request is None:
await self._handle.increment_many(native)
return
self._handle.increment_many_post_call(request, native)

57
litellm/redis/stores/rate_limits/handle.py Normal file
View file

@ -0,0 +1,57 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Sequence
from typing import Protocol, TypeAlias
from litellm.caching.redis_batch import NativePending
from litellm.redis.stores.rate_limits.types import (
LocalValue,
NativeAtomicCounter,
NativeCounter,
NativeEntry,
NativeIncrement,
NativeOverLimit,
NativeScope,
)
from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend
class RateLimitsHandle(Protocol):
@property
def has_redis(self) -> bool: ...
def slot_groups(self, scopes: Sequence[NativeScope], force: bool) -> list[list[int]]: ...
def get_local(self, entry: NativeEntry) -> 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, send: NativeSend
) -> NativePending[list[tuple[str, int]]] | None: ...
def read_counters(self, counters: Sequence[NativeCounter], send: NativeSend) -> NativePending[list[str | None]] | None: ...
def check_and_increment(
self, counters: Sequence[NativeAtomicCounter], send: NativeSend
) -> NativePending[tuple[NativeOverLimit | None, list[tuple[int, int]]]] | None: ...
def guarded_increment(
self, counter: NativeCounter, expected_window_start: str, amount: float, ttl_seconds: int | None, send: NativeSend
) -> NativePending[tuple[bool, int]] | None: ...
def parallel_acquire(
self, gauges: Sequence[tuple[NativeScope, int]], slot_ttl_seconds: int, slot_id: str, send: NativeSend
) -> NativePending[tuple[NativeOverLimit | None, list[int]]] | None: ...
def parallel_release(
self, gauges: Sequence[NativeScope], slot_id: str, send: NativeSend
) -> NativePending[list[int]] | None: ...
def parallel_count(
self, gauges: Sequence[NativeScope], slot_ttl_seconds: int, send: NativeSend
) -> NativePending[list[int]] | None: ...
def token_increment(
self, increments: Sequence[NativeIncrement], send: NativeSend
) -> NativePending[list[float]] | None: ...
def token_increment_post_call(self, request: NativeRequestRedis, increments: Sequence[NativeIncrement]) -> bool: ...
def parallel_release_post_call(
self, request: NativeRequestRedis, gauges: Sequence[NativeScope], slot_id: str
) -> bool: ...
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]) -> None: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, float | None, int], RateLimitsHandle]

145
litellm/redis/stores/rate_limits/types.py Normal file
View file

@ -0,0 +1,145 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Literal, TypeAlias
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
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]
NativeOverLimit: TypeAlias = tuple[int, int, int]
def native_scope(scope: RateLimitScope) -> NativeScope:
return (scope.descriptor_key, scope.descriptor_value)
def native_counter(counter: RateLimitCounter) -> NativeCounter:
return (counter.scope.descriptor_key, counter.scope.descriptor_value, counter.kind)
def native_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 native_increment(increment: CounterIncrement) -> NativeIncrement:
return (native_counter(increment.counter), increment.amount, increment.ttl)
def native_atomic(counter: AtomicCounter) -> NativeAtomicCounter:
return (native_counter(counter.counter), counter.limit, counter.increment, counter.ttl, counter.window_size)

108
litellm/redis/stores/router_cooldowns.py
View file

@ -1,108 +0,0 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
if TYPE_CHECKING:
from opentelemetry.trace import Span
from typing_extensions import ReadOnly, TypedDict
from litellm._internal_context import service_target
from litellm.caching.redis_batch import NativeBatchResult, NativePending, RequestRedisBatches
from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeRouterCooldowns
class CooldownCacheValue(TypedDict):
exception_received: ReadOnly[str]
status_code: ReadOnly[str]
timestamp: ReadOnly[float]
cooldown_time: ReadOnly[float]
ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]]
class RouterCooldowns(Protocol):
def attach_redis(self, client: NativeRedisClient | None) -> None: ...
def flush_local_counters(self) -> None: ...
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ...
async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
def declare_active(
self, request: RequestRedisBatches, model_ids: Sequence[str]
) -> NativeBatchResult[ActiveCooldowns] | None: ...
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ...
def evict_local(self, model_id: str) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
class CooldownsHandle(Protocol):
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
def active(self, model_ids: Sequence[str]) -> Sequence[tuple[str, CooldownCacheValue]]: ...
def aactive(self, model_ids: Sequence[str]) -> Awaitable[Sequence[tuple[str, CooldownCacheValue]]]: ...
def declare_active(
self, request: NativeRequestRedis, model_ids: Sequence[str]
) -> NativePending[ActiveCooldowns]: ...
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
def evict_local(self, model_id: str) -> None: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None], CooldownsHandle]
class RouterCooldownStore:
def __init__(self, client: NativeRedisClient | None, *, handle: HandleFactory = NativeRouterCooldowns) -> None:
self._connected = client is not None
self._build: Final = handle
self._handle = handle(client)
def attach_redis(self, client: NativeRedisClient | None) -> None:
if client is None or self._connected:
return
self._connected = True
self._handle = self._build(client)
def flush_local_counters(self) -> None:
self._handle.flush_local_counters()
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None:
with service_target(ROUTER_COOLDOWNS_TARGET):
self._handle.put(model_id, value, ttl)
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int:
with service_target(ROUTER_COOLDOWNS_TARGET):
return self._handle.increment_allowed_fails(model_id, suffix, ttl)
async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
with service_target(ROUTER_COOLDOWNS_TARGET):
pending: Final = self._handle.aactive(model_ids)
return tuple(await pending)
def declare_active(
self, request: RequestRedisBatches, model_ids: Sequence[str]
) -> NativeBatchResult[ActiveCooldowns] | None:
"""``active`` declared on the request's native pre-call batch; None when there is no Redis to read."""
if not self._connected:
return None
with service_target(ROUTER_COOLDOWNS_TARGET):
pending: Final = self._handle.declare_active(request.native(), model_ids)
return request.native_result(pending)
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
with service_target(ROUTER_COOLDOWNS_TARGET):
return tuple(self._handle.active(model_ids))
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
with service_target(ROUTER_COOLDOWNS_TARGET):
return self._handle.minimum(model_ids) or default
def evict_local(self, model_id: str) -> None:
self._handle.evict_local(model_id)
def snapshot(self) -> Mapping[str, CooldownCacheValue]:
return self._handle.snapshot()

81
litellm/redis/stores/router_cooldowns/__init__.py Normal file
View file

@ -0,0 +1,81 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm._internal_context import with_service_target
from litellm.caching.redis_batch import Batch
from litellm.redis.stores._route import route
from litellm.redis.stores.router_cooldowns.handle import HandleFactory, RouterCooldownsHandle
from litellm.redis.stores.router_cooldowns.types import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, CooldownCacheValue
from litellm.rust_bridge._native import NativeRedisClient, NativeRouterCooldowns
__all__ = [
"ROUTER_COOLDOWNS_TARGET",
"ActiveCooldowns",
"CooldownCacheValue",
"HandleFactory",
"RouterCooldownStore",
"RouterCooldowns",
"RouterCooldownsHandle",
]
class RouterCooldowns(Protocol):
def attach_redis(self, client: NativeRedisClient | None) -> None: ...
def flush_local_counters(self) -> None: ...
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ...
def active(self, model_ids: Sequence[str], *, batch: Batch) -> Awaitable[ActiveCooldowns]: ...
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ...
def evict_local(self, model_id: str) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
async def _tuple(pending: Awaitable[ActiveCooldowns]) -> ActiveCooldowns:
return tuple(await pending)
class RouterCooldownStore:
def __init__(self, client: NativeRedisClient | None, *, handle: HandleFactory = NativeRouterCooldowns) -> None:
self._handle: Final = handle(client)
def attach_redis(self, client: NativeRedisClient | None) -> None:
"""Adds the Redis tier once, keeping this router's memory, as ``DualCache.attach_redis_cache`` does."""
if client is not None:
self._handle.attach(client)
def flush_local_counters(self) -> None:
self._handle.flush_local_counters()
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None:
self._handle.put(model_id, value, ttl)
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int:
return self._handle.increment_allowed_fails(model_id, suffix, ttl)
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def active(self, model_ids: Sequence[str], *, batch: Batch) -> Awaitable[ActiveCooldowns]:
"""Declared now on the batch ``batch`` picks, so several reads can share one trip before any is awaited."""
target: Final = route(batch)
return _tuple(target.result(self._handle.active(model_ids, target.send)))
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
return tuple(self._handle.active_sync(model_ids))
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
return self._handle.minimum(model_ids) or default
def evict_local(self, model_id: str) -> None:
self._handle.delete_local(model_id)
def snapshot(self) -> Mapping[str, CooldownCacheValue]:
return self._handle.snapshot()

23
litellm/redis/stores/router_cooldowns/handle.py Normal file
View file

@ -0,0 +1,23 @@
from __future__ import annotations
from collections.abc import Callable, Mapping, Sequence
from typing import Protocol, TypeAlias
from litellm.caching.redis_batch import NativePending
from litellm.redis.stores.router_cooldowns.types import ActiveCooldowns, CooldownCacheValue
from litellm.rust_bridge._native import NativeRedisClient, NativeSend
class RouterCooldownsHandle(Protocol):
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool: ...
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
def active(self, model_ids: Sequence[str], send: NativeSend) -> NativePending[ActiveCooldowns]: ...
def active_sync(self, model_ids: Sequence[str]) -> Sequence[tuple[str, CooldownCacheValue]]: ...
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
def delete_local(self, model_id: str) -> None: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None], RouterCooldownsHandle]

17
litellm/redis/stores/router_cooldowns/types.py Normal file
View file

@ -0,0 +1,17 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Final, TypeAlias
from typing_extensions import ReadOnly, TypedDict
class CooldownCacheValue(TypedDict):
exception_received: ReadOnly[str]
status_code: ReadOnly[str]
timestamp: ReadOnly[float]
cooldown_time: ReadOnly[float]
ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]]

74
litellm/redis/stores/router_usage.py → litellm/redis/stores/router_usage/__init__.py
View file

@ -1,25 +1,32 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm._internal_context import with_service_target
from litellm.redis.stores.router_usage.handle import HandleFactory, RouterUsageHandle
from litellm.redis.stores.router_usage.types import (
ROUTER_USAGE_TARGET,
CounterKind,
NativeCounter,
UsageCounter,
native_counter,
)
from litellm.rust_bridge._native import NativeRedisClient, NativeRouterUsage
ROUTER_USAGE_TARGET: Final = "router_usage"
CounterKind: TypeAlias = Literal["tpm", "rpm"]
@dataclass(frozen=True, slots=True)
class UsageCounter:
deployment_id: str
model: str | None
minute: str
kind: CounterKind
__all__ = [
"ROUTER_USAGE_TARGET",
"CounterKind",
"HandleFactory",
"NativeCounter",
"RouterUsage",
"RouterUsageHandle",
"RouterUsageStore",
"UsageCounter",
]
class RouterUsage(Protocol):
@ -41,42 +48,14 @@ class RouterUsage(Protocol):
def snapshot(self) -> Mapping[str, float]: ...
NativeCounter: TypeAlias = tuple[str, str | None, str, CounterKind]
class UsageHandle(Protocol):
def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ...
def increment_deployment(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> Awaitable[None]: ...
def get_many(self, counters: Sequence[NativeCounter]) -> Awaitable[Sequence[float | None]]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None], UsageHandle]
def _native(counter: UsageCounter) -> NativeCounter:
return (counter.deployment_id, counter.model, counter.minute, counter.kind)
class RouterUsageStore:
def __init__(self, client: NativeRedisClient | None, *, handle: HandleFactory = NativeRouterUsage) -> None:
self._connected = client is not None
self._build: Final = handle
self._handle = handle(client)
self._handle: Final = handle(client)
def attach_redis(self, client: NativeRedisClient | None) -> None:
if client is None or self._connected:
return
self._connected = True
self._handle = self._build(client)
"""Adds the Redis tier once, keeping this router's memory counts, as ``DualCache.attach_redis_cache`` does."""
if client is not None:
self._handle.attach(client)
def flush_local_counters(self) -> None:
self._handle.flush_local_counters()
@ -85,7 +64,7 @@ class RouterUsageStore:
async def async_increment(
self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None
) -> float | None:
return await self._handle.increment(_native(counter), value, ttl)
return await self._handle.increment(native_counter(counter), value, ttl)
@with_service_target(ROUTER_USAGE_TARGET)
async def increment_deployment(
@ -99,8 +78,9 @@ class RouterUsageStore:
) -> None:
await self._handle.increment_deployment(deployment_id, model, minute, tokens, requests)
@with_service_target(ROUTER_USAGE_TARGET)
async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None:
return await self._handle.get_many(tuple(map(_native, counters)))
return await self._handle.get_many(tuple(map(native_counter, counters)))
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()

26
litellm/redis/stores/router_usage/handle.py Normal file
View file

@ -0,0 +1,26 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Protocol, TypeAlias
from litellm.redis.stores.router_usage.types import NativeCounter
from litellm.rust_bridge._native import NativeRedisClient
class RouterUsageHandle(Protocol):
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool: ...
def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ...
def increment_deployment(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> Awaitable[None]: ...
def get_many(self, counters: Sequence[NativeCounter]) -> Awaitable[Sequence[float | None]]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None], RouterUsageHandle]

22
litellm/redis/stores/router_usage/types.py Normal file
View file

@ -0,0 +1,22 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, Literal, TypeAlias
ROUTER_USAGE_TARGET: Final = "router_usage"
CounterKind: TypeAlias = Literal["tpm", "rpm"]
@dataclass(frozen=True, slots=True)
class UsageCounter:
deployment_id: str
model: str | None
minute: str
kind: CounterKind
NativeCounter: TypeAlias = tuple[str, str | None, str, CounterKind]
def native_counter(counter: UsageCounter) -> NativeCounter:
return (counter.deployment_id, counter.model, counter.minute, counter.kind)

225
litellm/redis/stores/spend_counters.py
View file

@ -1,225 +0,0 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final, Literal, Protocol, TypeAlias
from litellm._internal_context import service_target, with_service_target
from litellm.caching.redis_batch import (
NativeBatchResult,
NativePending,
RequestRedisBatches,
active_request_redis_batches,
)
from litellm.rust_bridge._native import (
NativeRedisClient,
NativeRedisStoreError,
NativeRequestRedis,
NativeSpendCounters,
)
SPEND_COUNTERS_TARGET: Final = "spend_counters"
SPEND_COUNTER_PREFIX: Final = "spend:"
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
SpendIncrement: TypeAlias = tuple[str, float]
SpendReading: TypeAlias = tuple[float | None, bool]
"""(value, authoritative): authoritative when Redis answered, a clean miss included."""
def key_counter(token: str) -> str:
return f"spend:key:{token}"
def key_window_counter(token: str, budget_duration: str) -> str:
return f"spend:key:{token}:window:{budget_duration}"
def team_counter(team_id: str) -> str:
return f"spend:team:{team_id}"
def team_window_counter(team_id: str, budget_duration: str) -> str:
return f"spend:team:{team_id}:window:{budget_duration}"
def team_member_counter(user_id: str, team_id: str) -> str:
return f"spend:team_member:{user_id}:{team_id}"
def user_counter(user_id: str) -> str:
return f"spend:user:{user_id}"
def end_user_counter(end_user_id: str) -> str:
return f"{END_USER_COUNTER_PREFIX}{end_user_id}"
def org_counter(org_id: str) -> str:
return f"spend:org:{org_id}"
def tag_counter(tag: str) -> str:
return f"spend:tag:{tag}"
def project_counter(project_id: str) -> str:
return f"spend:project:{project_id}"
def model_access_group_counter(access_group: str) -> str:
return f"spend:model_access_group:{access_group}"
def shadow_eval_counter(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
def is_unavailable(error: BaseException) -> bool:
"""Redis timed out or its breaker is open, the failures a post-call increment drops."""
return isinstance(error, NativeRedisStoreError) and error.args[:1] in (("timeout",), ("breaker_open",))
class SpendCountersHandle(Protocol):
def read(self, counter: str) -> Awaitable[tuple[Literal["redis", "memory"], float | None]]: ...
def read_many(self, counters: Sequence[str]) -> Awaitable[Sequence[float | None] | None]: ...
def declare_read_many(
self, request: NativeRequestRedis, counters: Sequence[str]
) -> NativePending[Sequence[float | None] | None] | None: ...
def increment(self, counter: str, amount: float) -> Awaitable[float]: ...
def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ...
def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ...
def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ...
def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ...
def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ...
def delete(self, counters: Sequence[str]) -> Awaitable[None]: ...
def refresh_ttl(self, counter: str) -> Awaitable[bool]: ...
def get_local(self, counter: str) -> float | None: ...
def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None: ...
def delete_local(self, counter: str) -> None: ...
def increment_local(self, counter: str, amount: float) -> float: ...
def flush_local(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, int], SpendCountersHandle]
class SpendCounterStore:
"""The ``spend:*`` budget counters: Redis shared by every worker, and this worker's memory copy that answers
only when Redis cannot. Counters are named by the builders above."""
def __init__(
self, client: NativeRedisClient | None, *, redis_ttl: int = 60, handle: HandleFactory = NativeSpendCounters
) -> None:
self._build: Final = handle
self._handle = handle(client, redis_ttl)
self._redis_ttl = None if client is None else redis_ttl
def attach_redis(self, client: NativeRedisClient | None, *, redis_ttl: int) -> None:
if client is None or self._redis_ttl is not None:
return
self._handle = self._build(client, redis_ttl)
self._redis_ttl = redis_ttl
@property
def redis_ttl(self) -> int | None:
"""The TTL every Redis write gives a counter, None when there is no Redis."""
return self._redis_ttl
@property
def has_redis(self) -> bool:
return self._redis_ttl is not None
@with_service_target(SPEND_COUNTERS_TARGET)
async def read(self, counter: str) -> SpendReading:
source, value = await self._handle.read(counter)
return value, source == "redis"
@with_service_target(SPEND_COUNTERS_TARGET)
async def read_many(self, counters: Sequence[str]) -> Mapping[str, float | None] | None:
"""None when there is no Redis or the read failed, so each caller reads on its own."""
values: Final = await self._handle.read_many(counters)
return None if values is None else dict(zip(counters, values, strict=True))
def declare_read_many(
self, request: RequestRedisBatches, counters: Sequence[str]
) -> NativeBatchResult[Sequence[float | None] | None] | None:
"""``read_many`` declared on the request's native pre-call batch; None when there is no Redis."""
if not self.has_redis:
return None
with service_target(SPEND_COUNTERS_TARGET):
pending: Final = self._handle.declare_read_many(request.native(), counters)
return None if pending is None else request.native_result(pending)
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment(self, counter: str, amount: float) -> float:
"""Raises when the increment may not have landed, after dropping the counter everywhere."""
return await self._handle.increment(counter, amount)
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_many(self, increments: Sequence[SpendIncrement]) -> tuple[float, ...]:
return tuple(await self._handle.increment_many(increments))
def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool:
"""Declares the increments on the request's post-call batch, False when there is no Redis or request.
Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere."""
request: Final = active_request_redis_batches()
if request is None or not self.has_redis or not increments:
return False
with service_target(SPEND_COUNTERS_TARGET):
return self._handle.increment_many_post_call(request.native_post_call(), increments)
@with_service_target(SPEND_COUNTERS_TARGET)
async def set_max(self, counter: str, spend: float) -> float | None:
"""Raise memory and Redis to ``spend``, never lowering either; what Redis holds after, None without Redis."""
return await self._handle.set_max(counter, spend)
@with_service_target(SPEND_COUNTERS_TARGET)
async def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> float:
"""Seed a cold counter with SET NX so one worker wins, answering with what the counter then holds."""
return await self._handle.seed(counter, spend, increment_if_lost=increment_if_lost)
@with_service_target(SPEND_COUNTERS_TARGET)
async def set(self, counter: str, spend: float, ttl: float | None = None) -> None:
"""Memory, then Redis; raises when Redis failed, with memory already written."""
await self._handle.set(counter, spend, ttl)
@with_service_target(SPEND_COUNTERS_TARGET)
async def delete(self, counters: Sequence[str]) -> None:
"""Memory, then Redis in one round trip; raises when Redis failed, with memory already dropped."""
await self._handle.delete(counters)
@with_service_target(SPEND_COUNTERS_TARGET)
async def refresh_ttl(self, counter: str) -> bool:
return await self._handle.refresh_ttl(counter)
def get_local(self, counter: str) -> float | None:
return self._handle.get_local(counter)
def set_local(self, counter: str, spend: float, ttl: float | None = None) -> None:
self._handle.set_local(counter, spend, ttl)
def delete_local(self, counter: str) -> None:
self._handle.delete_local(counter)
def increment_local(self, counter: str, amount: float) -> float:
return self._handle.increment_local(counter, amount)
def flush_local(self) -> None:
self._handle.flush_local()
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()
def apply_broadcast(self, cache_key: str, new_value: float | None, ttl: float | None) -> None:
"""A cross-worker auth cache invalidation: a SET writes this worker's copy, a delete drops it. Keys outside
the spend counter keyspace are not this store's."""
if not cache_key.startswith(SPEND_COUNTER_PREFIX):
return
try:
if new_value is None:
self.delete_local(cache_key)
else:
self.set_local(cache_key, new_value, ttl)
except ValueError:
return

173
litellm/redis/stores/spend_counters/__init__.py Normal file
View file

@ -0,0 +1,173 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Final
from litellm._internal_context import with_service_target
from litellm.caching.redis_batch import Batch
from litellm.redis.stores._route import post_call_request, route
from litellm.redis.stores.spend_counters.handle import HandleFactory, SpendCountersHandle
from litellm.redis.stores.spend_counters.types import (
END_USER_COUNTER_PREFIX,
SPEND_COUNTER_PREFIX,
SPEND_COUNTERS_TARGET,
SpendIncrement,
SpendReading,
end_user_counter,
is_unavailable,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
shadow_eval_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.rust_bridge._native import NativeRedisClient, NativeSpendCounters
__all__ = [
"END_USER_COUNTER_PREFIX",
"SPEND_COUNTERS_TARGET",
"SPEND_COUNTER_PREFIX",
"HandleFactory",
"SpendCounterStore",
"SpendCountersHandle",
"SpendIncrement",
"SpendReading",
"end_user_counter",
"is_unavailable",
"key_counter",
"key_window_counter",
"model_access_group_counter",
"org_counter",
"project_counter",
"shadow_eval_counter",
"tag_counter",
"team_counter",
"team_member_counter",
"team_window_counter",
"user_counter",
]
async def _by_counter(
counters: Sequence[str], values: Awaitable[Sequence[float | None] | None] | None
) -> Mapping[str, float | None] | None:
read: Final = None if values is None else await values
return None if read is None else dict(zip(counters, read, strict=True))
class SpendCounterStore:
"""The ``spend:*`` budget counters: Redis shared by every worker, and this worker's memory copy that answers
only when Redis cannot. Counters are named by the builders above."""
def __init__(
self, client: NativeRedisClient | None, *, redis_ttl: int = 60, handle: HandleFactory = NativeSpendCounters
) -> None:
self._handle: Final = handle(client, redis_ttl)
self._redis_ttl = None if client is None else redis_ttl
def attach_redis(self, client: NativeRedisClient | None, *, redis_ttl: int) -> None:
"""Adds the Redis tier once, keeping the memory copies, as ``DualCache.attach_redis_cache`` does."""
if client is None or not self._handle.attach(client, redis_ttl):
return
self._redis_ttl = redis_ttl
@property
def redis_ttl(self) -> int | None:
"""The TTL every Redis write gives a counter, None when there is no Redis."""
return self._redis_ttl
@property
def has_redis(self) -> bool:
return self._redis_ttl is not None
@with_service_target(SPEND_COUNTERS_TARGET)
async def read(self, counter: str) -> SpendReading:
source, value = await self._handle.read(counter)
return value, source == "redis"
@with_service_target(SPEND_COUNTERS_TARGET)
def read_many(self, counters: Sequence[str], *, batch: Batch) -> Awaitable[Mapping[str, float | None] | None]:
"""One MGET declared now on the batch ``batch`` picks. None when there is no Redis or the read failed, so
each caller reads on its own."""
target: Final = route(batch)
pending: Final = self._handle.read_many(counters, target.send)
return _by_counter(counters, None if pending is None else target.result(pending))
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment(self, counter: str, amount: float) -> float:
"""Raises when the increment may not have landed, after dropping the counter everywhere."""
return await self._handle.increment(counter, amount)
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_many(self, increments: Sequence[SpendIncrement]) -> tuple[float, ...]:
return tuple(await self._handle.increment_many(increments))
@with_service_target(SPEND_COUNTERS_TARGET)
def increment_many_post_call(self, increments: Sequence[SpendIncrement]) -> bool:
"""Declares the increments on the request's post-call batch, False when there is no Redis or request.
Each total lands in memory once the batch goes out; a counter whose increment failed is dropped everywhere."""
request: Final = post_call_request() if self.has_redis and increments else None
if request is None:
return False
return self._handle.increment_many_post_call(request, increments)
@with_service_target(SPEND_COUNTERS_TARGET)
async def set_max(self, counter: str, spend: float) -> float | None:
"""Raise memory and Redis to ``spend``, never lowering either; what Redis holds after, None without Redis."""
return await self._handle.set_max(counter, spend)
@with_service_target(SPEND_COUNTERS_TARGET)
async def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> float:
"""Seed a cold counter with SET NX so one worker wins, answering with what the counter then holds."""
return await self._handle.seed(counter, spend, increment_if_lost=increment_if_lost)
@with_service_target(SPEND_COUNTERS_TARGET)
async def set(self, counter: str, spend: float, ttl: float | None = None) -> None:
"""Memory, then Redis; raises when Redis failed, with memory already written."""
await self._handle.set(counter, spend, ttl)
@with_service_target(SPEND_COUNTERS_TARGET)
async def delete(self, counters: Sequence[str]) -> None:
"""Memory, then Redis in one round trip; raises when Redis failed, with memory already dropped."""
await self._handle.delete(counters)
@with_service_target(SPEND_COUNTERS_TARGET)
async def refresh_ttl(self, counter: str) -> bool:
return await self._handle.refresh_ttl(counter)
def get_local(self, counter: str) -> float | None:
return self._handle.get_local(counter)
def set_local(self, counter: str, spend: float, ttl: float | None = None) -> None:
self._handle.set_local(counter, spend, ttl)
def delete_local(self, counter: str) -> None:
self._handle.delete_local(counter)
def increment_local(self, counter: str, amount: float) -> float:
return self._handle.increment_local(counter, amount)
def flush_local(self) -> None:
self._handle.flush_local()
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()
def apply_broadcast(self, cache_key: str, new_value: float | None, ttl: float | None) -> None:
"""A cross-worker auth cache invalidation: a SET writes this worker's copy, a delete drops it. Keys outside
the spend counter keyspace are not this store's."""
if not cache_key.startswith(SPEND_COUNTER_PREFIX):
return
try:
if new_value is None:
self.delete_local(cache_key)
else:
self.set_local(cache_key, new_value, ttl)
except ValueError:
return

33
litellm/redis/stores/spend_counters/handle.py Normal file
View file

@ -0,0 +1,33 @@
from __future__ import annotations
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Literal, Protocol, TypeAlias
from litellm.caching.redis_batch import NativePending
from litellm.redis.stores.spend_counters.types import SpendIncrement
from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend
class SpendCountersHandle(Protocol):
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool: ...
def read(self, counter: str) -> Awaitable[tuple[Literal["redis", "memory"], float | None]]: ...
def read_many(
self, counters: Sequence[str], send: NativeSend
) -> NativePending[Sequence[float | None] | None] | None: ...
def increment(self, counter: str, amount: float) -> Awaitable[float]: ...
def increment_many(self, increments: Sequence[SpendIncrement]) -> Awaitable[Sequence[float]]: ...
def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[SpendIncrement]) -> bool: ...
def set_max(self, counter: str, spend: float) -> Awaitable[float | None]: ...
def seed(self, counter: str, spend: float, *, increment_if_lost: bool) -> Awaitable[float]: ...
def set(self, counter: str, spend: float, ttl_seconds: float | None = None) -> Awaitable[None]: ...
def delete(self, counters: Sequence[str]) -> Awaitable[None]: ...
def refresh_ttl(self, counter: str) -> Awaitable[bool]: ...
def get_local(self, counter: str) -> float | None: ...
def set_local(self, counter: str, spend: float, ttl_seconds: float | None = None) -> None: ...
def delete_local(self, counter: str) -> None: ...
def increment_local(self, counter: str, amount: float) -> float: ...
def flush_local(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None, int], SpendCountersHandle]

66
litellm/redis/stores/spend_counters/types.py Normal file
View file

@ -0,0 +1,66 @@
from __future__ import annotations
from typing import Final, TypeAlias
from litellm.rust_bridge._native import NativeRedisStoreError
SPEND_COUNTERS_TARGET: Final = "spend_counters"
SPEND_COUNTER_PREFIX: Final = "spend:"
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
SpendIncrement: TypeAlias = tuple[str, float]
SpendReading: TypeAlias = tuple[float | None, bool]
"""(value, authoritative): authoritative when Redis answered, a clean miss included."""
def key_counter(token: str) -> str:
return f"spend:key:{token}"
def key_window_counter(token: str, budget_duration: str) -> str:
return f"spend:key:{token}:window:{budget_duration}"
def team_counter(team_id: str) -> str:
return f"spend:team:{team_id}"
def team_window_counter(team_id: str, budget_duration: str) -> str:
return f"spend:team:{team_id}:window:{budget_duration}"
def team_member_counter(user_id: str, team_id: str) -> str:
return f"spend:team_member:{user_id}:{team_id}"
def user_counter(user_id: str) -> str:
return f"spend:user:{user_id}"
def end_user_counter(end_user_id: str) -> str:
return f"{END_USER_COUNTER_PREFIX}{end_user_id}"
def org_counter(org_id: str) -> str:
return f"spend:org:{org_id}"
def tag_counter(tag: str) -> str:
return f"spend:tag:{tag}"
def project_counter(project_id: str) -> str:
return f"spend:project:{project_id}"
def model_access_group_counter(access_group: str) -> str:
return f"spend:model_access_group:{access_group}"
def shadow_eval_counter(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
def is_unavailable(error: BaseException) -> bool:
"""Redis timed out or its breaker is open, the failures a post-call increment drops."""
return isinstance(error, NativeRedisStoreError) and error.args[:1] in (("timeout",), ("breaker_open",))

3
litellm/router_utils/cooldown_cache.py
View file

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm.caching.redis_batch import Batch
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.redis.stores.router_cooldowns import (
ROUTER_COOLDOWNS_TARGET,
@ -43,7 +44,7 @@ class CooldownCache:
async def async_get_active_cooldowns(
self, model_ids: Sequence[str], parent_otel_span: Span | None
) -> list[tuple[str, CooldownCacheValue]]:
return list(await self.cooldown_store.active(model_ids, parent_otel_span))
return list(await self.cooldown_store.active(model_ids, batch=Batch.NONE))
def get_min_cooldown(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> float:
return self.cooldown_store.minimum(model_ids, self.default_cooldown_time, parent_otel_span)

9
litellm/router_utils/routing_prefetch.py
View file

@ -11,7 +11,7 @@ if TYPE_CHECKING:
from litellm._internal_context import service_target
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
from litellm.caching.redis_batch import Batch, BatchResult, active_request_redis_batches
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, RouterCooldowns
_PREFETCH_SLOT: Final = "routing_read"
@ -86,9 +86,8 @@ def arm_cooldowns(store: RouterCooldowns, model_ids: Sequence[str]) -> None:
request: Final = active_request_redis_batches()
if request is None or _COOLDOWNS_SLOT in request.prefetched:
return
result: Final = store.declare_active(request, model_ids)
if result is not None:
request.prefetched[_COOLDOWNS_SLOT] = ArmedCooldowns(frozenset(model_ids), result)
result: Final = store.active(model_ids, batch=Batch.PRE_CALL)
request.prefetched[_COOLDOWNS_SLOT] = ArmedCooldowns(frozenset(model_ids), result)
def armed() -> bool:
@ -120,7 +119,7 @@ async def read_cooldowns(
prefetched: Final = await _armed_cooldowns(model_ids)
if prefetched is not None:
return prefetched
return await store.active(model_ids, parent_otel_span)
return await store.active(model_ids, batch=Batch.NONE)
def _take(read: TierRead) -> ArmedRead | None:

112
litellm/rust_bridge/_native.pyi
View file

@ -57,24 +57,38 @@ class NativeRequestRedis:
class NativePending(Generic[_T_co]):
def resolve(self) -> Future[_T_co]: ...
class NativeSend:
@final
class Now(NativeSend):
def __new__(cls) -> NativeSend.Now: ...
@final
class PreCall(NativeSend):
def __new__(cls, request: NativeRequestRedis, /) -> NativeSend.PreCall: ...
@final
class PostCall(NativeSend):
def __new__(cls, request: NativeRequestRedis, /) -> NativeSend.PostCall: ...
@final
class NativeRouterCooldowns:
def __new__(cls, client: NativeRedisClient | None, /) -> NativeRouterCooldowns: ...
def __new__(cls, client: NativeRedisClient | None, redis_ttl_seconds: int = 60, /) -> NativeRouterCooldowns: ...
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool: ...
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
def active(self, model_ids: Sequence[str]) -> list[tuple[str, _NativeCooldown]]: ...
def aactive(self, model_ids: Sequence[str]) -> Future[list[tuple[str, _NativeCooldown]]]: ...
def declare_active(
self, request: NativeRequestRedis, model_ids: Sequence[str]
def active(
self, model_ids: Sequence[str], send: NativeSend
) -> NativePending[list[tuple[str, _NativeCooldown]]]: ...
def active_sync(self, model_ids: Sequence[str]) -> list[tuple[str, _NativeCooldown]]: ...
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
def evict_local(self, model_id: str) -> None: ...
def delete_local(self, model_id: str) -> None: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> dict[str, _NativeCooldown]: ...
@final
class NativeRouterUsage:
def __new__(cls, client: NativeRedisClient | None, /) -> NativeRouterUsage: ...
def __new__(cls, client: NativeRedisClient | None, redis_ttl_seconds: int = 60, /) -> NativeRouterUsage: ...
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool: ...
def increment(
self, counter: tuple[str, str | None, str, Literal["tpm", "rpm"]], value: float, ttl_seconds: int
) -> Future[float]: ...
@ -85,8 +99,6 @@ 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"]]]
@ -100,6 +112,7 @@ _RateLimitEntry = tuple[Literal["window", "requests", "tokens", "max_parallel_re
_RateLimitLocal = int | float | str | dict[str, float]
_RateLimitIncrement = tuple[_RateLimitCounter, int | float, int | None]
_RateLimitAtomicCounter = tuple[_RateLimitCounter, int, int, int, int]
_RateLimitOverLimit = tuple[int, int, int]
@final
class NativeRateLimits:
@ -110,64 +123,54 @@ class NativeRateLimits:
redis_ttl_seconds: int = 60,
/,
) -> NativeRateLimits: ...
@property
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 slot_groups(self, scopes: Sequence[_RateLimitScope], force: bool) -> list[list[int]]: ...
def get_local(self, entry: _RateLimitEntry) -> _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,
) -> NativePending[object]: ...
self, counters: Sequence[_RateLimitCounter], now: int, window_size: int, send: NativeSend
) -> NativePending[list[tuple[str, int]]] | None: ...
def read_counters(
self, counters_: Sequence[_RateLimitCounter], *, request: NativeRequestRedis | None = None
) -> NativePending[object]: ...
self, counters: Sequence[_RateLimitCounter], send: NativeSend
) -> NativePending[list[str | None]] | None: ...
def check_and_increment(
self, counters_: Sequence[_RateLimitAtomicCounter], *, request: NativeRequestRedis | None = None
) -> NativePending[object]: ...
self, counters: Sequence[_RateLimitAtomicCounter], send: NativeSend
) -> NativePending[tuple[_RateLimitOverLimit | None, list[tuple[int, int]]]] | None: ...
def guarded_increment(
self,
counter: _RateLimitCounter,
expected_window_start: str,
amount: float,
ttl_seconds: int | None,
*,
request: NativeRequestRedis | None = None,
) -> NativePending[object]: ...
send: NativeSend,
) -> NativePending[tuple[bool, int]] | None: ...
def parallel_acquire(
self,
gauges: Sequence[tuple[_RateLimitScope, int]],
slot_ttl_seconds: int,
slot_id: str,
*,
request: NativeRequestRedis | None = None,
) -> NativePending[object]: ...
self, gauges: Sequence[tuple[_RateLimitScope, int]], slot_ttl_seconds: int, slot_id: str, send: NativeSend
) -> NativePending[tuple[_RateLimitOverLimit | None, list[int]]] | None: ...
def parallel_release(
self, gauges: Sequence[_RateLimitScope], slot_id: str, *, request: NativeRequestRedis | None = None
) -> NativePending[object]: ...
self, gauges: Sequence[_RateLimitScope], slot_id: str, send: NativeSend
) -> NativePending[list[int]] | None: ...
def parallel_count(
self, gauges: Sequence[_RateLimitScope], slot_ttl_seconds: int, *, request: NativeRequestRedis | None = None
) -> NativePending[object]: ...
self, gauges: Sequence[_RateLimitScope], slot_ttl_seconds: int, send: NativeSend
) -> NativePending[list[int]] | None: ...
def token_increment(
self, increments_: Sequence[_RateLimitIncrement], *, request: NativeRequestRedis | None = None
) -> NativePending[object]: ...
self, increments: Sequence[_RateLimitIncrement], send: NativeSend
) -> NativePending[list[float]] | None: ...
def token_increment_post_call(
self, request: NativeRequestRedis, increments_: Sequence[_RateLimitIncrement]
) -> Future[None]: ...
self, request: NativeRequestRedis, increments: Sequence[_RateLimitIncrement]
) -> bool: ...
def parallel_release_post_call(
self, request: NativeRequestRedis, gauges: Sequence[_RateLimitScope], slot_id: str
) -> Future[None]: ...
) -> bool: ...
def increment(self, increment: _RateLimitIncrement) -> Future[int | float]: ...
def increment_many(self, increments_: Sequence[_RateLimitIncrement]) -> Future[None]: ...
def increment_many(self, increments: Sequence[_RateLimitIncrement]) -> Future[None]: ...
def increment_many_post_call(
self, request: NativeRequestRedis, increments_: Sequence[_RateLimitIncrement]
) -> Future[None]: ...
self, request: NativeRequestRedis, increments: Sequence[_RateLimitIncrement]
) -> None: ...
@final
class NativeAuthCache:
@ -186,20 +189,13 @@ class NativeAuthCache:
def get_sync(self, key: str) -> object: ...
def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ...
def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
def set_pre_call(
self, request: NativeRequestRedis, key: str, value: object, ttl_seconds: float | None = None
) -> NativePending[None] | None: ...
def delete_pre_call(self, request: NativeRequestRedis, key: str) -> NativePending[None] | None: ...
def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ...
def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ...
def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ...
def write_back(
self, request: NativeRequestRedis | None, entries: Sequence[tuple[str, object, float | None]]
) -> Future[None]: ...
def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ...
def delete(self, keys: Sequence[str]) -> Future[None]: ...
def delete_sync(self, keys: Sequence[str]) -> None: ...
def declare_get_many(
self, request: NativeRequestRedis, keys: Sequence[str]
) -> NativePending[list[object]] | None: ...
def get_many(self, keys: Sequence[str]) -> Future[list[object] | None]: ...
def get_many(self, keys: Sequence[str], send: NativeSend) -> NativePending[list[object]] | None: ...
def increment_if_present(self, key: str, amount: float) -> Future[float | None]: ...
def get_local(self, key: str) -> object: ...
def set_local(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ...
@ -210,10 +206,10 @@ class NativeAuthCache:
@final
class NativeSpendCounters:
def __new__(cls, client: NativeRedisClient | None, redis_ttl_seconds: int = 60, /) -> NativeSpendCounters: ...
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool: ...
def read(self, counter: str) -> Future[tuple[Literal["redis", "memory"], float | None]]: ...
def read_many(self, counters: Sequence[str]) -> Future[list[float | None] | None]: ...
def declare_read_many(
self, request: NativeRequestRedis, counters: Sequence[str]
def read_many(
self, counters: Sequence[str], send: NativeSend
) -> NativePending[list[float | None] | None] | None: ...
def increment(self, counter: str, amount: float) -> Future[float]: ...
def increment_many(self, increments: Sequence[tuple[str, float]]) -> Future[list[float]]: ...

407
redis-api-design.md Normal file
View file

@ -0,0 +1,407 @@
# Native Redis API design
This covers the layering for the Rust Redis stack behind the Python proxy, from how a Redis command is written down in `redis-client` up to how Python store wrappers pick a request phase. It came out of reviewing the PyO3 store bindings in the `litellm_router_stores_rust` stack (#49525, #49820, #49834, #49821)
## Problem
### The bridge carries knowledge it does not need
`python-bridge/src/redis/store/*.rs` was meant to be pure PyO3 conversion, but each family file also spelled store knowledge by hand. `rate_limits.rs` matched `"requests"`, `"tokens"`, `"window"` and `"max_parallel_requests"`, `router_usage.rs` matched `"tpm"` and `"rpm"`, `router_cooldowns.rs` named every `CooldownCacheValue` field, and `spend_counters.rs` and `auth_cache.rs` each mapped Redis client errors to kind strings. Those strings are Redis key suffixes and value shapes that `redis-store` already owns and spelled a second time
### Every store picks its own request phase
Whether an op goes out alone, rides the request's pre-call batch, or rides its post-call batch is a request concern, yet each store decides it inside its own method names
| Store | How it picks | Methods |
|---|---|---|
| `router_cooldowns` | always pre-call | `declare_active` |
| `spend_counters` | phase in the name | `declare_read_many_pre_call`, `increment_many_post_call` |
| `auth_cache` | phase in the name | `set_pre_call`, `delete_pre_call`, `declare_get_many` |
| `rate_limits` | `Option<&RequestRedis>` plus a private `Route`, and separate `*_post_call` methods | 8 script methods, 3 post-call methods |
The Rust stores call `request.pre(client)` or `request.post(client)` in 13 places. The Python wrappers in `litellm/redis/stores/*.py` repeat it one layer up, each looking up `active_request_redis_batches()`, checking for Redis and a request, picking `native()` or `native_post_call()`, and wrapping the result in `native_result`, with `rate_limits.py` adding its own `_declare` and `_post_call_request`. Moving an op between phases means a new method in Rust, the bridge, the stub and Python, and the names disagree across stores
`Route` (`Batch` or `Direct`) is shared infrastructure but lives in `rate_limits/scripts.rs`, so no other store can send an op alone, and `Pending<T>` there carries a second `Reply::Declared | Reply::Direct` enum just to resolve a reply
A pre-call read has to be awaited through `NativeBatchResult` so the native flush hooks run first and late readers such as the routing prefetch join the same trip. Every store has to remember that, and a bare `pending.resolve()` silently drops them
### Commands are written three times in redis-client
`commands.rs` builds each `redis::Cmd` once, but the surfaces around it repeat every command
| Surface | What it repeats |
|---|---|
| `RedisClient` (`client.rs`) | an async method per command applying the namespace, querying and decoding |
| `Pipeline` (`pipeline.rs`) | a builder per command applying the namespace and wrapping an `Op { entry, key, idempotent }` |
| `RedisBatch` (`batch.rs`) | a `BatchOp` variant per command, a match in `fn pipeline()` mapping it to `Pipeline`, and a match in `run_alone` mapping it to `RedisClient` |
One new batchable command touches five places
## Layers
| Layer | Owns | Does not know about |
|---|---|---|
| `redis-client` | commands as values, reply decoding, `RedisClient`, `Pipeline`, `RedisBatch` and its alone mode | requests, phases, keyspaces |
| request layer, `RequestRedis` and `Send` (in `redis-store` for now) | a request's pre and post batches, the post-call deadline, `spawn`, the shutdown drain, and turning a phase plus a client into a batch | keyspaces, Python |
| store modules (`router_cooldowns`, `router_usage`, `spend_counters`, `auth_cache`, `rate_limits`) | keys, value codecs, memory tiers, scripts, post-call follow-ups, Python value shapes through `FromValue` and `IntoValue` | which phase they run in |
| `python-bridge` | PyO3 conversion, `NativePending`, `NativeSend`, exception types | store logic |
| Python (`litellm/caching/redis_batch.py`, `litellm/redis/stores/*.py`) | `Phase`, `NativeTarget`, coexistence with stores still on redis-py | Rust internals |
`RequestRedis` depends only on `redis-client` and a Tokio handle, so it could move out of `redis-store`. It stays out of `redis-client`, whose `AGENTS.md` keeps it free of domain knowledge, and a request with pre and post phases is a host lifecycle concept. Keep it in `redis-store`, which both hosts already depend on, and move it to its own crate only when something besides the stores needs request batches
## redis-client
### One command value
```rust
pub struct Command<T> {
plan: Plan,
key: String,
idempotent: bool,
decode: fn(Value) -> Result<T, Error>,
}
enum Plan {
Commands(Vec<Cmd>),
Script(ScriptCall),
IncrKeepingTtl { amount: Amount, ttl: Duration },
}
```
Each command is one constructor in `commands.rs`
```rust
pub fn get<T: FromRedisValue>(key: &str) -> Command<Option<T>>;
pub fn mget<T: FromRedisValue>(keys: &[&str]) -> Command<Vec<Option<T>>>;
pub fn set(key: &str, value: impl ToSingleRedisArg, ttl: Duration) -> Command<()>;
pub fn del(keys: &[&str]) -> Command<usize>;
pub fn expire(key: &str, ttl: Duration) -> Command<bool>;
pub fn incr_by_float(key: &str, amount: Amount, expiry: Expiry) -> Command<f64>;
pub fn eval<T: FromRedisValue>(script: &Script, keys: &[&str], args: ...) -> Command<T>;
```
The constructor arguments are the wire encode, through redis-rs `ToRedisArgs` (`Amount` writes an `int` with `str()` and a `float` with `repr()`, like redis-py). `decode` is the matching reply decode. Both stay generic, so `Command` has no separate `encode`
### Composing a store codec
```rust
impl<T> Command<T> {
pub fn map<U>(self, decode: fn(T) -> Result<U, Error>) -> Command<U>;
}
```
A store attaches its keyspace codec without a new type, and the same `Command<Cooldown>` works on every executor
```rust
commands::get::<Vec<u8>>(&key).map(Cooldown::decode)
batch.add(commands::set(&key, cooldown.encode(), ttl))
```
### Three executors, one entry point each
```rust
impl RedisClient {
pub async fn send<T>(&self, command: Command<T>) -> Result<T, Error>;
}
impl Pipeline {
pub fn add<T>(self, command: Command<T>) -> Self;
}
impl RedisBatch {
pub fn add<T>(&self, command: Command<T>) -> Declared<T>;
}
```
The executor applies its client's namespace to `key` once. `BatchOp`, `batch.rs`'s `fn pipeline()`, `run_alone` and `Pipeline`'s private `Op` go away, and a new command is one constructor. Call sites change from `client.get::<String>(key).await` to `client.send(commands::get::<String>(key)).await` and from `batch.mget(&keys)` to `batch.add(commands::mget(&keys))`. `subscribe` returns a stream and stays a method
### Alone versus pipelined behavior is explicit
A command run alone follows the `RedisCache` method behind it, and a pipelined one follows Python's pipeline helper, which is why `Plan` keeps the two cases that differ
| Command | Alone | Pipelined |
|---|---|---|
| `incr_by_float` with `Expiry::KeepExisting` | `INCRBYFLOAT`, read `TTL`, `EXPIRE` only a key with none, like `async_increment` | lowered to `Refresh`, `INCRBYFLOAT` plus `EXPIRE`, like `RedisBatch.increment` |
| `eval` | `EVALSHA`, and on `NOSCRIPT` a `SCRIPT LOAD` then `EVALSHA`, like redis-py's `Script` (commit `91ddc9f032f`) | `EVALSHA` with the pipeline's retry |
Each command with both forms gets a contract test against the commands Redis receives, so the forms never drift from Python's
### A batch that runs every op alone
```rust
impl RedisBatch {
pub fn alone(client: RedisClient) -> Self;
}
```
`flush` already runs each op alone on a cluster. The alone flag takes that path everywhere, so a fresh alone batch is "send it now" and the first `resolve()` sends it. `Declared<T>` becomes the one reply handle for every route, which removes `rate_limits::scripts`'s `Route` and `Reply` enums and reduces its `Pending<T>` to nothing once `Command::map` carries the decode
## Request layer
```rust
pub enum Send<'a> {
Now,
PreCall(&'a RequestRedis),
PostCall(&'a RequestRedis),
}
impl Send<'_> {
pub fn batch(&self, client: &RedisClient) -> RedisBatch {
match self {
Self::Now => RedisBatch::alone(client.clone()),
Self::PreCall(request) => request.pre(client),
Self::PostCall(request) => request.post(client),
}
}
}
```
The caller decides the phase, because only it knows when it needs the answer. The store owns its client, which is optional and can differ per store, so the caller passes `Send` and the store calls `send.batch(&client)`, the one place the request and the client meet
`RequestRedis` keeps its current API, `pre`, `post`, `spawn`, `flush_pre` and `flush_post`. The pre-call batch goes out at the first `resolve()` or at `flush_pre()`. The post-call batch goes out at `flush_post()` or `post_call_deadline` after its first declaration, and `flush_post` waits for the follow-ups spawned on it so the shutdown drain keeps fallback writes in flight
## Stores
### Each op exists once
```rust
impl RateLimits {
pub fn token_increment(&self, send: Send<'_>, increments: &[Increment]) -> Result<Declared<Vec<f64>>, Error> {
let batch = send.batch(self.redis.as_ref().ok_or(Error::NoRedis)?);
Ok(scripts::token_increment(&batch, increments))
}
}
impl RouterCooldowns {
pub fn active(&self, send: Send<'_>, model_ids: &[String]) -> Result<Declared<Vec<(String, Cooldown)>>, Error>;
}
impl SpendCounters {
pub fn read_many(&self, send: Send<'_>, counters: &[SpendCounter]) -> Declared<Option<Vec<Option<Number>>>>;
}
impl AuthCache {
pub fn get_many(&self, send: Send<'_>, keys: &[String]) -> Declared<Option<Vec<Option<Value>>>>;
}
```
### Post-call follow-ups stay with the store
Nobody resolves a post-call write, so work after it settles runs on `request.spawn`. A post-call method exists only where that follow-up exists, as a thin wrapper over the shared op
```rust
pub fn token_increment_post_call(self: &Arc<Self>, request: &RequestRedis, increments: Vec<Increment>) -> Result<(), Error> {
let declared = self.token_increment(Send::PostCall(request), &increments)?;
let store = Arc::clone(self);
request.spawn(async move {
if declared.settled().await.is_err() {
let _ = store.increment_many(&increments).await;
}
});
Ok(())
}
```
A plain post-call write is `op(Send::PostCall(request), ...)` with the `Declared` dropped
### Spellings and shapes live with the store
Key suffixes and wire spellings derive from strum on the store's enums, per `litellm-rust/.agents/skills/rust-string-enums/SKILL.md`, so there is one spelling per value
```rust
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)]
pub enum Kind {
#[strum(serialize = "requests")]
Requests,
#[strum(serialize = "tokens")]
Tokens,
}
```
The same goes for `router_usage::Kind` (`tpm`, `rpm`), `spend_counters::Reading` (`redis`, `memory`), and `ErrorKind` (`breaker_open`, `timeout`, `transport`, `rejected`, `decode`, `encode`, `memory`), with `Error::kind()` on each store error and `From<&litellm_redis_client::Error> for ErrorKind`. `rate_limits::Entry::new(suffix, scope)` reads the suffix Python names an entry by, and the window and gauge suffixes are constants shared with the key builders
### Python value shapes
When the Python side of a store type is plain data, the store type converts through `litellm_python_compat::Value`, so the bridge never names a field
```rust
pub trait FromValue: Sized {
fn from_value(value: &Value) -> Result<Self, DecodeError>;
}
pub trait IntoValue {
fn into_value(self) -> Value;
}
```
`Cooldown::to_value` and `from_value` are the first case, and `encode` and `decode` are built on them (`repr(&self.to_value())` and `Self::from_value(&python_value(raw)?)`). `Counter`, `Scope`, `Entry`, `Increment` and `AtomicCounter` implement the same pair for the tuples Python sends, so each family's Python wire shape is a `redis-store` contract that parity tests can pin without pyo3. Unlike a store trait, this pair has a generic consumer, the bridge's `Py<T>`
Going through `Value` costs one small intermediate allocation over extracting a tuple directly, a few times per request at most, which readability outranks here
## python-bridge
### What stays
The bridge keeps only what pyo3 and the orphan rule force on it
1. `NativePending`, the one resolve-once handle for a declared op
2. `NativeSend`, the Python face of `Send`
3. `Py<T>` over `FromValue` and `IntoValue`, plus `PyValue`, `PyNumber` (the int versus float check) and the few non-data shapes such as `PyLocal`
4. Reply views such as the rate limiter's `(over_limit, applied)` tuples, which are Python-specific
5. `NativeRedisStoreError(kind, message)`, built once by `kind_error(error.kind(), &error)`
6. The `redis_async` and `redis_sync` helpers, which attach the `ServiceLogging` labels and run one scoped op
7. Host settings projection (`policy.rs`)
Each pyclass keeps its own explicit `#[pymethods]` list, with no store trait and no macro writing methods
### NativePending
```rust
#[pyclass(frozen, name = "NativePending")]
pub(crate) struct NativePending {
resolve: Mutex<Option<Resolve>>,
call_type: &'static str,
}
```
`resolve()` takes the future once, flushing the op's batch if it has not gone out, and raises `RuntimeError` on a second call. `call_type` is the `ServiceLogging` label, attached at resolve time where the traffic happens. Rust futures are lazy, so building the future at declare time sends nothing. It replaced `NativeDeclaredCooldowns`, `NativeDeclaredSpendCounters`, `NativeDeclaredAuthReads`, `NativeDeclaredAuthWrite` and `NativeRateLimitPending`
### NativeSend
`redis-store` stays free of pyo3, so the enum mirroring `Send` lives in the bridge. Each variant carries exactly what it needs, so there are no phase strings and no phase without a request
```rust
#[pyclass(frozen, name = "NativeSend")]
pub(crate) enum NativeSend {
Now(),
PreCall(Py<NativeRequestRedis>),
PostCall(Py<NativeRequestRedis>),
}
impl NativeSend {
fn send<'a>(&'a self, py: Python<'a>) -> Send<'a> {
match self {
Self::Now() => Send::Now,
Self::PreCall(request) => Send::PreCall(&request.get().0),
Self::PostCall(request) => Send::PostCall(&request.get().0),
}
}
}
```
`NativeRequestRedis` is `frozen`, so `Py::get()` borrows its `RequestRedis` without a clone. `Now()` needs parentheses because PyO3 complex enums do not mix unit and tuple variants. Every store method takes `send: &NativeSend` and returns `NativePending`
```rust
fn token_increment(&self, py: Python<'_>, increments_: Vec<Py<Increment>>, send: &NativeSend) -> PyResult<NativePending> {
pending("token_increment_script", self.0.token_increment(send.send(py), &increments(increments_)), std::convert::identity)
}
```
Stub
```python
@final
class NativeSend:
@final
class Now(NativeSend):
def __new__(cls) -> NativeSend.Now: ...
@final
class PreCall(NativeSend):
def __new__(cls, request: NativeRequestRedis, /) -> NativeSend.PreCall: ...
@final
class PostCall(NativeSend):
def __new__(cls, request: NativeRequestRedis, /) -> NativeSend.PostCall: ...
```
## Python
### Phase and NativeTarget in redis_batch.py
```python
class Phase(Enum):
PRE_CALL = auto()
POST_CALL = auto()
@dataclass(frozen=True, slots=True)
class NativeTarget:
send: NativeSend
batches: RequestRedisBatches | None
def result(self, pending: NativePending[_T]) -> Awaitable[_T]:
return self.batches.native_result(pending) if self.batches else pending.resolve()
def native_target(phase: Phase) -> NativeTarget:
batches: Final = active_request_redis_batches()
if batches is None:
return NativeTarget(NativeSend.Now(), None)
match phase:
case Phase.PRE_CALL:
return NativeTarget(NativeSend.PreCall(batches.native()), batches)
case Phase.POST_CALL:
return NativeTarget(NativeSend.PostCall(batches.native_post_call()), batches)
```
`result()` makes the flush-hook rule part of the API. A pre-call read always goes through `native_result`, which awaits `flush_native_pre()` and its hooks before `resolve()`, so no store can drop late readers. `native()` and `native_post_call()` return the same `NativeRequestRedis`, and `POST_CALL` keeps calling `native_post_call()` because it registers the request with the shutdown drain
### Python store ops
```python
def check_and_increment(self, counters: Sequence[AtomicCounter], target: NativeTarget) -> Awaitable[AtomicReply]:
return _atomic_reply(target.result(self._handle.check_and_increment(tuple(map(_atomic, counters)), target.send)))
```
This drops the `declare_x`, `x` and `x_post_call` triples, the per-store `_declare` and `_post_call_request`, and the `request=` keyword
### App code
```python
target = native_target(Phase.PRE_CALL)
window, checked, spend, cooldowns = await asyncio.gather(
rate_limits.sliding_window(counters, now, window_size, target),
rate_limits.check_and_increment(atomic_counters, target),
spend_counters.read_many(spend_keys, target),
router_cooldowns.active(model_ids, target),
)
```
Inside a request the first await flushes, so the four reads share one trip. Outside one the target is `Now()` and each read goes alone, and the calling code is the same. Post-call writes keep their current call sites, such as `token_increment_post_call(increments)`, and pick `Phase.POST_CALL` inside the store
| | Pre-call | Post-call |
|---|---|---|
| For | reads the request needs | writes nobody waits on |
| Python gets | `NativePending` through `NativeTarget.result` | nothing it resolves |
| Sends on | the first `resolve()` or `flush_pre()` | `flush_post()` or the deadline |
| Waits on the reply | Python, through `resolve()` | Rust, through `settled()` in a spawned task |
## Coexisting with stores still on redis-py
`RequestRedisBatches` already flushes both kinds of batch each phase, the Python `RedisBatch` per backend (`batch()` and `post_call()`) and the one `NativeRequestRedis` (`native()`), gathering both in `flush_all` and `flush_post_call`. `native_target` only touches the native half. A store on redis-py keeps `batch(redis_cache)`, and a store that migrates switches its Python wrapper from `batch()` to `native_target()`, so nothing outside `litellm/redis/stores/` changes and stores move one at a time
Until the last store on a Redis moves, a request touching both kinds sends two concurrent pipelines per phase to it, which is what already happens today. The two pipelines have no order between them, which is fine because no store reads another's keys within a batch
## Not doing
1. A `RedisStore` trait. Nothing consumes stores generically, and their ops, inputs and replies all differ. What they share is values, `Send`, `Declared<T>` and `Command<T>`
2. A `Sink` with its own command methods. It would be a fourth copy of every command, and `RedisBatch::alone` gives the same routing with none
3. Python resolving post-call writes. They are fire-and-forget by design and `flush_post` already waits for their follow-ups at shutdown
4. Keyspace codecs on `Command`. `redis-client` stays domain-agnostic, and `Command::map` composes the store's codec instead
5. Typed exception kinds for router stores and rate limits, which raise `RuntimeError` today. That is a behavior change to decide on its own
## Open questions
1. `declare_*` returns `None` today when there is no Redis, and callers fall back to their own read. With `Now()` always available, the proposal is that each op answers from memory without Redis, the way `RateLimits::get` does, which changes every `declare_*` caller
2. `Declared` from an alone batch needs its own way to run, since `Declared::resolve` currently flushes its shared batch. A boxed future polled on first await should do, but it is not prototyped
3. The three `*_post_call` rate limit methods return an already finished awaitable around synchronous work, and could return `None` instead
4. Rate limit parameters `counters_`, `scopes_` and `increments_` leak their trailing underscore into Python keyword names, and renaming the helpers fixes it
## Order of work
1. `refactor(redis-client)` below the store stack: `Command<T>`, `Command::map`, `send` and `add` on the three executors, `RedisBatch::alone`, and the contract tests for commands with both forms
2. #49525: `Send` beside `RequestRedis`, `FromValue` and `IntoValue`, the bridge's `NativeSend` and `Py<T>`, Python's `Phase`, `NativeTarget` and `native_target`, and router cooldowns and usage converted
3. #49820: spend counters converted
4. #49834: auth cache converted
5. #49821: rate limits converted, dropping `_declare`, `_post_call_request` and the `request=` keyword
Each step is a new commit on its branch, restacked bottom up. Behavior stays the same except that ops limited to pre-call today can also go alone, and the existing store tests plus one `rstest` case per `Send` variant pin it
## Already done in the stack
1. `store/mod.rs` gained `redis_async`, `redis_sync`, `NativePending` and `ttl` (#49525), and each store method makes one helper call
2. `NativeRedisStoreError(kind, message)` replaced `NativeSpendCounterError` and `NativeAuthCacheError` (#49820, #49834)
3. `rate_limits::Kind` and `Entry::new`, `router_usage::Kind`, `Cooldown::to_value` and `from_value`, `spend_counters::Reading`, and `ErrorKind` with `Error::kind()` moved the spellings and shapes into `redis-store`, with `PyValue` moved down to #49525 so cooldowns can use it. Those last commits were restacked locally and not yet pushed when this was written

20
tests/unit/caching/test_redis_batch.py
View file

@ -189,6 +189,26 @@ class FakeNativePending(Generic[_Answer]):
return self.answer()
@dataclass
class ReadyPending(Generic[_Answer]):
"""An op declared with ``NativeSend.Now()``: it answers whenever it is resolved."""
answer: Callable[[], _Answer]
async def resolve(self) -> _Answer:
return self.answer()
def fake_send(_phase: object, request: FakeNativeRequest) -> FakeNativeRequest:
"""A ``native_send_factory`` for ``FakeNatives``: the fake request stands in for its ``NativeSend``."""
return request
def fake_pending(send: object, answer: Callable[[], _Answer]) -> FakeNativePending[_Answer] | ReadyPending[_Answer]:
"""What a fake handle declares for ``send``: on the fake request's batch, or on its own for ``Now()``."""
return FakeNativePending(send, answer) if isinstance(send, FakeNativeRequest) else ReadyPending(answer)
def make(fail: Exception | None = None, namespace: str | None = None) -> tuple[FakeRedisCache, FakeClient]:
client = FakeClient(replies, fail)
return FakeRedisCache(client, namespace), client

8
tests/unit/caching/test_request_redis_batch_post_call.py
View file

@ -41,7 +41,7 @@ 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
from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache, fake_send
async def _script_outside_the_pipeline(keys: Sequence[str], args: Sequence[object]) -> object:
@ -422,7 +422,7 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa
monkeypatch.setattr(proxy_server, "spend_counter_store", spend.store())
pending = [PendingSpendIncrement("spend:key:k1", 0.5), PendingSpendIncrement("spend:team:t1", 0.5)]
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
await proxy_server._apply_spend_counter_increments(pending)
declared = list(spend.post_call)
assert natives.built[0].events == []
@ -437,7 +437,7 @@ async def test_spend_counter_increments_ride_the_native_post_call_batch(monkeypa
async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary():
natives = FakeNatives()
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
request.native_post_call()
await request.flush_all()
at_the_boundary = list(natives.built[0].events)
@ -452,7 +452,7 @@ async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_requ
async def test_native_post_call_ops_still_waiting_on_their_callbacks_are_drained_at_shutdown():
natives = FakeNatives()
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
request.native_post_call()
await request.flush_all()
await drain_post_call_redis_batches()

54
tests/unit/caching/test_request_redis_batch_pre_call.py
View file

@ -17,7 +17,7 @@ import litellm.caching.dual_cache as dual_cache_module
from litellm import Router
from litellm._internal_context import current_service_target
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_batch import active_request_redis_batches, request_redis_batch_scope
from litellm.caching.redis_batch import Batch, active_request_redis_batches, request_redis_batch_scope
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
@ -27,6 +27,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitUnverifiableError,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.redis.stores._route import route
from litellm.redis.stores.rate_limits import RateLimitScope
from litellm.redis.stores.router_cooldowns import ActiveCooldowns, RouterCooldownStore
from litellm.router_api import RoutingContext
@ -44,6 +45,7 @@ from .test_redis_batch import (
FakeNativeRequest,
FakeNatives,
FakeRedisCache,
fake_send,
)
@ -123,7 +125,7 @@ async def test_descriptor_lua_calls_share_one_native_pipeline_and_each_keeps_its
_descriptor("team", "t1", 20),
]
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
response = await limiter.atomic_check_and_increment_by_n(
descriptors=descriptors,
increments=[{"requests": 1}, {"requests": 1}, {"requests": 1}],
@ -276,13 +278,11 @@ async def test_without_a_request_scope_descriptor_groups_run_one_at_a_time_and_s
@dataclass
class _FakeCooldownsHandle:
"""A connected cooldown store whose ``declare_active`` lands on the request's fake native batch."""
"""A connected cooldown store whose ``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: ())
def active(self, model_ids: Sequence[str], send: FakeNativeRequest) -> FakeNativePending[ActiveCooldowns]:
send.events.append("declare cooldowns")
return FakeNativePending(send, lambda: ())
@pytest.mark.asyncio
@ -299,7 +299,7 @@ async def test_the_admission_script_and_the_armed_cooldown_read_go_out_in_one_na
) # pyright: ignore[reportArgumentType] # test double
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) 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)
@ -755,7 +755,7 @@ def _user_entry() -> tuple[_CacheEntry, LiteLLM_UserTable]:
async def test_auth_write_back_rides_the_native_pre_call_batch_and_the_scope_drains_what_nobody_awaited():
cache, fake = shared_cache()
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await _write_back([_user_entry()], cache)
assert cache.get_local("user-1")["user_id"] == "user-1"
assert natives.built[0].events == ["declare write_back (('user-1', 42),)"]
@ -774,7 +774,7 @@ async def test_auth_write_back_outside_a_scope_writes_through_as_before():
@pytest.mark.asyncio
async def test_a_key_the_request_read_as_absent_is_not_read_again_by_a_per_key_get():
cache, fake = shared_cache()
with request_redis_batch_scope(native_factory=FakeNatives()) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await prefetch_identity_keys(["absent-key"], cache)
assert await cache.async_get_cache("absent-key") is None
await cache.async_set_cache("absent-key", {"v": 1}, ttl=5)
@ -787,7 +787,7 @@ async def test_a_key_the_request_read_as_absent_is_not_read_again_by_a_per_key_g
async def test_management_object_writes_inside_a_request_ride_its_pipeline_and_write_through_outside():
cache, fake = shared_cache()
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await cache.async_set_cache("team_id:t1", {"team_id": "t1"}, ttl=60)
await cache.async_set_cache("hashed-key-object", {"token": "hashed-key-object"}, ttl=60)
assert cache.get_local("team_id:t1") == {"team_id": "t1"}
@ -814,7 +814,7 @@ async def test_a_team_refresh_inside_a_request_sends_its_set_and_alias_del_in_on
proxy_logging_obj.internal_usage_cache = InternalUsageCache(dual_cache=usage_cache)
team = LiteLLM_TeamTableCachedObj(team_id="t1", team_alias="alpha")
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await cache_team_object("t1", team, cache, proxy_logging_obj)
assert natives.built[0].events == [
"declare set team_id:t1 ttl=60.0",
@ -840,7 +840,7 @@ async def test_a_changed_default_in_memory_ttl_reaches_the_native_store():
async def test_identity_prefetch_is_one_read_after_which_hits_and_misses_alike_cost_no_read():
cache, fake = shared_cache(FakeAuthCache(redis={"key-hit": {"k": "key-hit"}}))
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
await prefetch_identity_keys(["key-hit", "end_user_id:eu-miss", "key-hit"], cache)
assert await cache.async_get_cache("key-hit") == {"k": "key-hit"}
assert await cache.async_get_cache("end_user_id:eu-miss") is None
@ -857,7 +857,7 @@ async def test_the_auth_prefetch_and_the_armed_cooldown_read_go_out_in_one_nativ
) # pyright: ignore[reportArgumentType] # test double
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await prefetch_identity_keys(["key-hit"], cache, before_read=lambda: arm_cooldowns(cooldowns, ["dep-a"]))
arm_cooldowns(cooldowns, ["dep-a"])
active = await read_cooldowns(cooldowns, ["dep-a"], None)
@ -875,7 +875,7 @@ async def test_a_prefetch_served_from_memory_arms_no_early_cooldown_read():
cache.set_local("key-hit", {"k": "key-hit"}, None)
armed: Final[list[str]] = [] # mutable-ok: records the arming callback
with request_redis_batch_scope(native_factory=FakeNatives()): # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
await prefetch_identity_keys(["key-hit"], cache, before_read=lambda: armed.append("armed"))
assert armed == []
@ -885,9 +885,9 @@ async def test_a_prefetch_served_from_memory_arms_no_early_cooldown_read():
async def test_awaiting_a_native_read_runs_the_native_hooks_then_flushes_the_native_pre_batch_once_per_request():
natives: Final = FakeNatives()
with request_redis_batch_scope(post_call_deadline=7, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(post_call_deadline=7, native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
native: Final = request.native()
request.add_native_flush_hook(lambda: native.events.append("hook"))
request.add_pre_call_flush_hook(lambda: native.events.append("hook"))
result: Final = request.native_result(FakeNativePending(native, lambda: "answer"))
declared_only: Final = list(native.events)
answer: Final = await result
@ -900,11 +900,27 @@ async def test_awaiting_a_native_read_runs_the_native_hooks_then_flushes_the_nat
assert native.events == ["hook", "flush_pre", "resolve", "hook", "flush_pre"]
@pytest.mark.asyncio
async def test_a_pre_call_flush_hook_declares_on_its_own_request_whichever_request_awaits_the_flush():
natives: Final = FakeNatives()
declared_on: Final[list[object]] = []
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as owner: # pyright: ignore[reportArgumentType] # test double
owner.add_pre_call_flush_hook(lambda: declared_on.append(route(Batch.PRE_CALL).send))
result: Final = owner.native_result(FakeNativePending(owner.native(), lambda: "answer"))
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as other: # pyright: ignore[reportArgumentType] # test double
other.native()
answer: Final = await result
assert answer == "answer"
assert declared_on == [natives.built[0]]
@pytest.mark.asyncio
async def test_a_request_that_no_native_store_asks_for_builds_no_native_batches():
natives: Final = FakeNatives()
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
await request.flush_all()
await request.flush_post_call()

4
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -2524,7 +2524,7 @@ async def test_refresh_cached_team_broadcasts_team_and_alias_invalidations(
async def test_refresh_cached_team_flushes_redis_set_before_broadcast(monkeypatch) -> None:
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.proxy.management_endpoints import team_endpoints
from tests.unit.caching.test_redis_batch import FakeNatives
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from tests.unit.redis.stores.auth_cache_double import shared_cache
natives: Final = FakeNatives()
@ -2537,7 +2537,7 @@ async def test_refresh_cached_team_flushes_redis_set_before_broadcast(monkeypatc
proxy_logging_obj.internal_usage_cache.dual_cache = None
monkeypatch.setattr(team_endpoints, "publish_auth_cache_invalidation", record_publish)
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
await team_endpoints.refresh_cached_team(
team_row=LiteLLM_TeamTable(team_id="team-order", team_alias=None),
user_api_key_cache=cache,

6
tests/unit/proxy/spend_tracking/test_spend_counter_batch.py
View file

@ -24,7 +24,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
)
from litellm.rust_bridge._native import NativeRedisStoreError
from litellm.caching.redis_batch import request_redis_batch_scope
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives, fake_send
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
TOKEN = UserAPIKeyAuth(token="hashed", team_id="team", user_id="user", org_id="org")
@ -92,7 +92,7 @@ async def test_counters_bound_at_auth_join_the_native_pipeline_the_armed_cooldow
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0})
with (
request_redis_batch_scope(native_factory=natives) as request, # pyright: ignore[reportArgumentType] # test double
request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request, # pyright: ignore[reportArgumentType] # test double
spend_counter_batch_scope(spend.store()),
):
bind_admission_counter_keys(TOKEN, end_user_id=None)
@ -118,7 +118,7 @@ async def test_counters_bound_after_the_native_batch_went_out_ride_the_next_one(
spend = FakeSpendCounters(redis={"spend:key:hashed": 3.0, "spend:org:org": 4.0})
with (
request_redis_batch_scope(native_factory=natives), # pyright: ignore[reportArgumentType] # test double
request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send), # pyright: ignore[reportArgumentType] # test double
spend_counter_batch_scope(spend.store(), counter_keys=frozenset({"spend:key:hashed"})),
):
batch = active_spend_counter_batch()

38
tests/unit/redis/stores/auth_cache_double.py
View file

@ -30,6 +30,11 @@ class DeclaredOp:
return self.answer()
def _on_request(send: object) -> bool:
"""A fake request stands in for a request's ``NativeSend``; anything else is ``NativeSend.Now()``."""
return getattr(send, "events", None) is not None
@dataclass
class FakeAuthCache:
"""``redis`` holds what each key decodes to; ``fail`` makes every direct Redis round trip raise."""
@ -74,17 +79,17 @@ class FakeAuthCache:
self._round_trip("SET_STR", key, ttl_seconds)
self.redis[key] = value
def set_pre_call(self, request: object, key: str, value: object, ttl_seconds: float | None = None) -> DeclaredOp:
def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: object) -> DeclaredOp:
self.set_local(key, value, ttl_seconds)
self._declare(request, f"declare set {key} ttl={ttl_seconds}")
self._declare(send, f"declare set {key} ttl={ttl_seconds}")
self.redis[key] = value
return DeclaredOp(request, lambda: None)
return DeclaredOp(send, lambda: None)
def delete_pre_call(self, request: object, key: str) -> DeclaredOp:
def declare_delete(self, key: str, send: object) -> DeclaredOp:
self.delete_local(key)
self._declare(request, f"declare del {key}")
self._declare(send, f"declare del {key}")
self.redis.pop(key, None)
return DeclaredOp(request, lambda: None)
return DeclaredOp(send, lambda: None)
async def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> None:
for key, value in entries:
@ -92,15 +97,15 @@ class FakeAuthCache:
self._round_trip("SET_MANY", tuple(key for key, _ in entries), ttl_seconds)
self.redis.update(entries)
async def write_back(self, request: object | None, entries: Sequence[WriteBack]) -> None:
async def write_back(self, entries: Sequence[WriteBack], send: object) -> None:
for key, value, ttl in entries:
self.set_local(key, value, ttl)
self.redis[key] = value
written: Final = tuple((key, ttl) for key, _value, ttl in entries)
if request is None:
if not _on_request(send):
self._round_trip("WRITE_BACK", written)
return
self._declare(request, f"declare write_back {written}")
self._declare(send, f"declare write_back {written}")
async def delete(self, keys: Sequence[str]) -> None:
self.delete_sync(keys)
@ -112,19 +117,18 @@ class FakeAuthCache:
for key in keys:
self.redis.pop(key, None)
def declare_get_many(self, request: object, keys: Sequence[str]) -> DeclaredOp:
self._declare(request, f"declare auth_cache {','.join(keys)}")
def get_many(self, keys: Sequence[str], send: object) -> DeclaredOp:
on_request: Final = _on_request(send)
self._declare(send, f"declare auth_cache {','.join(keys)}")
def answer() -> list[object]:
if self.fail is not None:
if on_request and self.fail is not None:
raise self.fail
if not on_request:
self._round_trip("MGET", tuple(keys))
return [self.redis.get(key) for key in keys]
return DeclaredOp(request, answer)
async def get_many(self, keys: Sequence[str]) -> list[object]:
self._round_trip("MGET", tuple(keys))
return [self.redis.get(key) for key in keys]
return DeclaredOp(send, answer)
async def increment_if_present(self, key: str, amount: float) -> float | None:
self._round_trip("INCRBYFLOAT_IF_PRESENT", key, amount)

144
tests/unit/redis/stores/fake_rate_limits.py
View file

@ -7,7 +7,7 @@ from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass, field
from typing import Final
from typing import Final, Generic, TypeVar
from pydantic import TypeAdapter
@ -22,6 +22,7 @@ from litellm.redis.stores.rate_limits import (
)
from litellm.rust_bridge._native import NativeRateLimits
_T = TypeVar("_T")
RawScript = Callable[..., Awaitable[object]]
RawValue = int | float | str | bytes | None
_RAW: Final = TypeAdapter(list[RawValue])
@ -40,17 +41,18 @@ 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."""
class Resolved(Generic[_T]):
"""A script call; declared with a send that logs ``events`` (a ``FakeNativeRequest`` standing in for the
request's ``NativeSend``), 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:
def __init__(self, reply: Callable[[], Awaitable[_T]], send: object, name: str) -> None:
self._reply: Final = reply
self._events: Final = _events(request)
self._events: Final = _events(send)
if self._events is not None:
self._events.append(f"declare {name}")
def resolve(self) -> Awaitable[object]:
def resolve(self) -> Awaitable[_T]:
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())
@ -97,6 +99,7 @@ class FakeRateLimits:
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))
deferred: list[asyncio.Future[None]] = field(default_factory=list)
def store(self) -> RateLimitStore:
return RateLimitStore(None, handle=lambda _client, _ttl, _redis_ttl: self)
@ -108,54 +111,55 @@ class FakeRateLimits:
raise AssertionError(f"the test gave the limiter no fake for {name}")
return _RAW.validate_python(await script(keys=list(keys), args=list(args)))
@property
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 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 get_local(self, entry: NativeEntry) -> LocalValue | None:
return self.memory.get_local(entry)
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]
return self.memory.get_local(entry)
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))]
self, counters: Sequence[NativeCounter], now: int, window_size: int, send: object
) -> Resolved[list[tuple[str, int]]]:
keys: Final = [key for counter in counters for key in (window_key(counter), counter_key(counter))]
async def reply() -> object:
async def reply() -> list[tuple[str, int]]:
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 [(_text(raw[index]) or "", _int(raw[index + 1])) for index in range(0, len(raw), 2)]
return Resolved(reply, request, "sliding_window")
return Resolved(reply, send, "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))]
def read_counters(self, counters: Sequence[NativeCounter], send: object) -> Resolved[list[str | None]]:
keys: Final = [key for counter in counters for key in (window_key(counter), counter_key(counter))]
async def reply() -> object:
async def reply() -> list[str | None]:
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")
return Resolved(reply, send, "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]
def check_and_increment(self, counters: Sequence[NativeAtomicCounter], send: object) -> Resolved[tuple[tuple[int, int, int] | None, list[tuple[int, int]]]]:
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:
async def reply() -> tuple[tuple[int, int, int] | None, list[tuple[int, int]]]:
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")
return Resolved(reply, send, "check_and_increment")
def guarded_increment(
self,
@ -163,10 +167,9 @@ class FakeRateLimits:
expected_window_start: str,
amount: float,
ttl_seconds: int | None,
*,
request: object = None,
) -> Resolved:
async def reply() -> object:
send: object,
) -> Resolved[tuple[bool, int]]:
async def reply() -> tuple[bool, int]:
raw: Final = await self._run(
self.window_guarded_token_increment_script,
"window_guarded_token_increment_script",
@ -175,74 +178,87 @@ class FakeRateLimits:
)
return _int(raw[0]) == 1, _int(raw[1])
return Resolved(reply, request, "guarded_increment")
return Resolved(reply, send, "guarded_increment")
def parallel_acquire(
self,
gauges: Sequence[tuple[NativeScope, int]],
slot_ttl_seconds: int,
slot_id: str,
*,
request: object = None,
) -> Resolved:
send: object,
) -> Resolved[tuple[tuple[int, int, int] | None, list[int]]]:
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:
async def reply() -> tuple[tuple[int, int, int] | None, list[int]]:
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")
return Resolved(reply, send, "parallel_acquire")
def parallel_release(self, gauges: Sequence[NativeScope], slot_id: str, *, request: object = None) -> Resolved:
def parallel_release(self, gauges: Sequence[NativeScope], slot_id: str, send: object) -> Resolved[list[int]]:
keys: Final = [scope_key(scope, "max_parallel_requests") for scope in gauges]
async def reply() -> object:
async def reply() -> list[int]:
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")
return Resolved(reply, send, "parallel_release")
def parallel_count(
self, gauges: Sequence[NativeScope], slot_ttl_seconds: int, *, request: object = None
) -> Resolved:
self, gauges: Sequence[NativeScope], slot_ttl_seconds: int, send: object
) -> Resolved[list[int]]:
keys: Final = [scope_key(scope, "max_parallel_requests") for scope in gauges]
async def reply() -> object:
async def reply() -> list[int]:
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")
return Resolved(reply, send, "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)]
def token_increment(self, increments: Sequence[NativeIncrement], send: object) -> Resolved[list[float]]:
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:
async def reply() -> list[float]:
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")
return Resolved(reply, send, "token_increment")
async def token_increment_post_call(self, request: object, increments_: Sequence[NativeIncrement]) -> None:
async def settle(self) -> None:
"""Runs the post-call writes declared so far, as the request's post-call flush would."""
await asyncio.gather(*self.deferred)
def token_increment_post_call(self, request: object, increments: Sequence[NativeIncrement]) -> bool:
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:
async def run() -> None:
try:
await self.token_increment(increments, None).resolve()
except Exception: # noqa: BLE001 # the store's post-call path falls back exactly like this
await self.increment_many(increments)
self.deferred.append(asyncio.ensure_future(run()))
return self.connected
def parallel_release_post_call(self, request: object, gauges: Sequence[NativeScope], slot_id: str) -> bool:
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 run() -> None:
try:
await self.parallel_release(gauges, slot_id, None).resolve()
except Exception: # noqa: BLE001 # the store's post-call path only logs this
return
self.deferred.append(asyncio.ensure_future(run()))
return self.connected
async def increment(self, increment: NativeIncrement) -> float:
total: Final = await self.memory.increment(increment)
@ -251,11 +267,11 @@ class FakeRateLimits:
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_)
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_)
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:
def increment_many_post_call(self, request: object, increments: Sequence[NativeIncrement]) -> None:
self.post_call.append("increment_many")
await self.increment_many(increments_)
self.deferred.append(asyncio.ensure_future(self.increment_many(increments)))

20
tests/unit/redis/stores/spend_counters_double.py
View file

@ -79,6 +79,13 @@ class FakeSpendCounters:
client: Final = None if self.redis is None else NativeRedisClient.connect({"host": "127.0.0.1", "port": 1})
return SpendCounterStore(client, redis_ttl=self.redis_ttl, handle=lambda _, __: self)
def attach(self, client: NativeRedisClient, redis_ttl_seconds: int = 60, /) -> bool:
if self.redis is not None:
return False
self.redis = {}
self.redis_ttl = redis_ttl_seconds
return True
def mget(self, counters: Sequence[str]) -> list[float | None] | None:
if self.redis is None or self.fail or not self.batch_reads:
return None
@ -97,16 +104,13 @@ class FakeSpendCounters:
return "redis", value
return "memory", self.memory.get(counter)
def read_many(self, counters: Sequence[str]) -> asyncio.Future[Sequence[float | None] | None]:
return _done(self.mget(counters))
def declare_read_many(self, request: object, counters: Sequence[str]) -> DeclaredRead | None:
if self.redis is None or not self.batch_reads:
def read_many(self, counters: Sequence[str], send: object) -> DeclaredRead | None:
if self.redis is None:
return None
declared: Final = DeclaredRead(tuple(counters), self, request)
self.declared.append(declared)
events: Final = getattr(request, "events", None)
declared: Final = DeclaredRead(tuple(counters), self, send)
events: Final = getattr(send, "events", None)
if isinstance(events, list):
self.declared.append(declared)
events.append("declare spend_counters")
return declared

14
tests/unit/redis/stores/test_auth_cache.py
View file

@ -5,12 +5,12 @@ from unittest.mock import MagicMock
import pytest
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_batch import Batch, request_redis_batch_scope
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.common_utils.user_api_key_cache import is_user_key_cache_key
from litellm.redis.stores.auth_cache import AuthCacheSettings, AuthCacheStore
from litellm.rust_bridge._native import NativeAuthCache
from tests.unit.caching.test_redis_batch import FakeNatives
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
from .auth_cache_double import FakeAuthCache
@ -124,13 +124,13 @@ async def test_a_key_the_request_read_found_absent_is_answered_without_another_r
fake: Final = FakeAuthCache(redis={"team_id:hit": {"team_id": "hit"}})
store: Final = AuthCacheStore(None, AuthCacheSettings(200, 60), redis_ttl=60, handle=fake.factory)
with request_redis_batch_scope(native_factory=FakeNatives()) as request: # pyright: ignore[reportArgumentType] # test double
found = await store.read_missing(["team_id:absent", "team_id:hit"])
with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double
found = await store.read_missing(["team_id:absent", "team_id:hit"], batch=Batch.PRE_CALL)
absent = await store.get("team_id:absent")
await store.set("team_id:absent", {"team_id": "now"}, None)
store.delete_local("team_id:absent")
refilled = await store.get("team_id:absent")
store.delete_pre_call("team_id:hit")
store.declare_delete("team_id:hit")
deleted = await store.get("team_id:hit")
await request.flush_all()
@ -146,11 +146,11 @@ async def test_outside_a_request_a_read_is_a_round_trip_of_its_own() -> None:
fake: Final = FakeAuthCache(redis={"team_id:t1": {"team_id": "t1"}})
store: Final = AuthCacheStore(None, AuthCacheSettings(200, 60), redis_ttl=60, handle=fake.factory)
assert await store.read_missing(["team_id:t1", "team_id:t2"]) == {
assert await store.read_missing(["team_id:t1", "team_id:t2"], batch=Batch.PRE_CALL) == {
"team_id:t1": {"team_id": "t1"},
"team_id:t2": None,
}
assert store.set_pre_call("team_id:t3", {}, None) is None
assert store.declare_set("team_id:t3", {}, None) is None
assert fake.commands == [("MGET", ("team_id:t1", "team_id:t2"))]

18
tests/unit/redis/stores/test_rate_limits.py
View file

@ -4,15 +4,14 @@ from typing import Final
import pytest
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_batch import Batch, 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
from tests.unit.caching.test_redis_batch import FakeNatives, fake_send
@dataclass(frozen=True, slots=True)
@ -144,15 +143,16 @@ async def test_without_redis_increments_add_in_memory_as_python_numbers_do():
@pytest.mark.asyncio
async def test_without_redis_nothing_is_declared_on_the_request_and_post_call_writes_are_refused():
async def test_without_redis_a_script_raises_where_it_is_awaited_and_post_call_writes_are_refused():
store: Final = RateLimitStore(None)
counter: Final = RateLimitScope("api_key", "k1").counter("requests")
natives: Final = FakeNatives()
declared: Final = store.sliding_window((counter,), 0, 60, batch=Batch.NONE)
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")
with pytest.raises(RuntimeError, match="no Redis"):
await declared
with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double
assert not store.token_increment_post_call((CounterIncrement(counter, 1, 60),))
assert not store.parallel_release_post_call((counter.scope,), "slot")
assert natives.built == []

Some files were not shown because too many files have changed in this diff Show more