feat(redis): serve the proxy's common Redis call with Rust stores (#50170)

* feat(redis): route the router, rate limiters and auth through the native store families

* style(redis): satisfy the python lint gates

Format the rate limit store, sort its __all__ and give the routing read batch and two tests concrete list and dict types.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(router): give the registry's default usage store a memory-only handle

RouterUsageStore now takes a handle, so the default passed None and any usage read before RustCore swapped in its store raised AttributeError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* test(redis): build stores from native handles and wait for router tasks

Tests still used the pre-refactor store constructors (Store(None), handle= factories, NativeIncrement, tuple scopes). The router outcome tests counted 20 loop ticks, too few now that the allowed-fails increment resolves on the native runtime, so settle() waits for the scheduled tasks instead.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
yujonglee 2026-11-06 15:01:11 +09:00 • committed by GitHub
parent fc8f614a03
commit b1b236826a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
192 changed files with 17208 additions and 10154 deletions

View file

@ -4381,6 +4381,7 @@ dependencies = [
"criterion",
"fancy-regex 0.19.2",
"futures-util",
"indexmap 2.14.0",
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
@ -5773,6 +5774,7 @@ version = "0.29.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4688ddedf473e32662b9b067670129a8afb8c18e351482c70d62ba4a88171e8b"
dependencies = [
"indexmap 2.14.0",
"libc",
"once_cell",
"portable-atomic",

View file

@ -60,7 +60,8 @@ litellm-router-core.workspace = true
litellm-host-python.workspace = true
litellm-model-catalog.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
indexmap = "2.14.0"
pyo3 = { workspace = true, features = ["indexmap"] }
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true

View file

@ -42,9 +42,24 @@ mod _native {
#[pymodule_export]
use crate::redis::client::native_redis_breaker_states;
#[pymodule_export]
use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes};
#[pymodule_export]
use crate::redis::store::rate_limits::{
NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired,
PyGuardedReply, PyWindowCount,
};
#[pymodule_export]
use crate::redis::store::request::{NativeRequestRedis, NativeSend};
#[pymodule_export]
use crate::redis::store::{NativePending, NativeRedisStoreError};
use crate::redis::store::router_usage::NativeRouterUsage;
#[pymodule_export]
use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins;
#[pymodule_export]
use crate::redis::store::spend_counters::{NativeSpendCounters, PySpendRead};
#[pymodule_export]
use crate::redis::store::{
NativePending, NativeRedisStoreError, router_cooldowns::NativeRouterCooldowns,
};
#[pymodule_export]
use crate::router::{RouterCall, RouterCoreHandle};
#[pymodule_export]
@ -123,8 +138,22 @@ mod tests {
"NativeRedisClient",
"NativeRequestRedis",
"NativeSend",
"NativeAuthCache",
"NativeRateLimits",
"WindowCount",
"AtomicApplied",
"AtomicOverLimit",
"GuardedReply",
"GaugesAcquired",
"GaugeOverLimit",
"NativePending",
"NativeRouterCooldowns",
"NativeRouterUsage",
"NativeSensitiveRoutePins",
"NativeRedisStoreError",
"NativeSpendCounters",
"SpendRead",
"LocalSizes",
"native_redis_breaker_states",
"NativeClickHouseSpendConfig",
"NativeClickHouseSpendStorage",

View file

@ -13,6 +13,8 @@ pub(crate) enum PythonSettings {
UrlPolicy,
#[strum(serialize = "provider_defaults")]
ProviderDefaults,
#[strum(serialize = "router_tiers")]
RouterTiers,
#[strum(serialize = "secret_manager")]
SecretManager,
#[strum(serialize = "secret_manager_binding")]

View file

@ -18,7 +18,7 @@ use pyo3::{
exceptions::{PyOverflowError, PyTypeError, PyValueError},
prelude::*,
sync::MutexExt,
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyTuple},
types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple},
};
use self::credentials::PythonCredentialSource;
@ -40,16 +40,18 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
fn extract(kwargs: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let py = kwargs.py();
let source = PyDict::new(py);
let mapping = kwargs.cast::<PyMapping>()?;
source.update(&mapping)?;
let plain = PyDict::new(py);
let mut host_objects = Vec::new();
for (name, value) in kwargs.cast::<PyDict>()?.iter() {
for (name, value) in source.iter() {
match value.is_none() || is_plain(&value) {
true => plain.set_item(&name, value)?,
false => host_objects.push(name.extract::<String>()?),
}
}
let credential_provider = kwargs
.cast::<PyDict>()?
let credential_provider = source
.get_item("credential_provider")?
.filter(|provider| !provider.is_none() && !is_plain(provider))
.map(Bound::unbind);
@ -65,27 +67,44 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec {
}
}
/// `litellm.constants`' `REDIS_CIRCUIT_BREAKER_*`, which Python reads from the environment once.
pub(crate) struct RedisBreaker(BreakerConfig);
/// `litellm.redis.store_settings.RedisBreakerSettings`, which `RedisCache` fills from `litellm.constants`.
#[derive(FromPyObject)]
pub(crate) struct RedisBreaker {
#[pyo3(attribute)]
failure_threshold: u32,
#[pyo3(attribute)]
recovery_timeout: f64,
#[pyo3(attribute)]
timeout_min_duration: f64,
}
impl<'a, 'py> FromPyObject<'a, 'py> for RedisBreaker {
type Error = PyErr;
fn extract(breaker: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let (failure_threshold, recovery_timeout, timeout_min_duration): (u32, f64, f64) =
breaker.extract()?;
impl RedisBreaker {
fn config(&self) -> PyResult<BreakerConfig> {
let seconds = |seconds: f64| {
Duration::try_from_secs_f64(seconds)
.map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))
.map_err(|error| PyValueError::new_err(error.to_string()))
};
Ok(Self(BreakerConfig {
failure_threshold,
recovery_timeout: seconds(recovery_timeout)?,
timeout_min_duration: seconds(timeout_min_duration)?,
}))
Ok(BreakerConfig {
failure_threshold: self.failure_threshold,
recovery_timeout: seconds(self.recovery_timeout)?,
timeout_min_duration: seconds(self.timeout_min_duration)?,
})
}
}
/// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built.
#[derive(FromPyObject)]
pub(crate) struct StoreSettings {
#[pyo3(attribute)]
redis_kwargs: RedisSpec,
#[pyo3(attribute)]
namespace: Option<String>,
#[pyo3(attribute)]
default_ttl: f64,
#[pyo3(attribute)]
breaker: Option<RedisBreaker>,
}
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
value.is_instance_of::<PyString>()
|| value.is_instance_of::<PyBool>()
@ -119,7 +138,7 @@ fn shared_client(
py: Python<'_>,
spec: RedisSpec,
namespace: Option<String>,
breaker: Option<RedisBreaker>,
breaker: Option<BreakerConfig>,
) -> PyResult<Option<RedisClient>> {
let environment: SharedEnvironment = Arc::new(ProcessEnvironment);
let RedisSpec {
@ -176,7 +195,7 @@ fn shared_client(
}
let connected = {
let _runtime = runtime.enter();
litellm_redis_connect::connect(resolved.clone(), host, breaker.map(|breaker| breaker.0))
litellm_redis_connect::connect(resolved.clone(), host, breaker)
};
let client = match connected {
Ok(client) => client,
@ -205,15 +224,15 @@ pub(crate) struct NativeRedisClient {
impl NativeRedisClient {
/// `None` when redis-py would fail only once a command connects, which leaves every store given it on memory.
#[staticmethod]
#[pyo3(signature = (redis_kwargs, *, default_ttl_seconds, namespace = None, breaker = None))]
fn connect(
py: Python<'_>,
redis_kwargs: RedisSpec,
default_ttl_seconds: f64,
namespace: Option<String>,
breaker: Option<RedisBreaker>,
) -> PyResult<Option<Self>> {
let default_ttl = number::seconds(default_ttl_seconds).map_err(seconds_error)?;
fn connect(py: Python<'_>, settings: StoreSettings) -> PyResult<Option<Self>> {
let StoreSettings {
redis_kwargs,
namespace,
default_ttl,
breaker,
} = settings;
let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?;
let breaker = breaker.as_ref().map(RedisBreaker::config).transpose()?;
Ok(
shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self {
client,

View file

@ -1,11 +1,3 @@
#[expect(
dead_code,
reason = "the native stores that use these helpers land in the next PR"
)]
pub(crate) mod client;
mod events;
#[expect(
dead_code,
reason = "the native stores that use these helpers land in the next PR"
)]
pub(crate) mod store;

View file

@ -0,0 +1,265 @@
use std::sync::Arc;
use litellm_python_compat::Value;
use litellm_redis_store::auth_cache::{AuthCache, Settings};
use pyo3::prelude::*;
use super::{
NativePending, PyNumber, kind_error, pending, redis_async, redis_sync, remote,
request::{NativeRequestRedis, NativeSend},
ttl,
value::PyValue,
};
use crate::redis::client::NativeRedisClient;
fn values(values: Vec<Option<Value>>) -> Vec<Option<PyValue>> {
values.into_iter().map(|value| value.map(PyValue)).collect()
}
/// `UserApiKeyCache` with `enable_redis_auth_cache` on. The sync methods back `DualCache`'s sync
/// API and block with the GIL released, like redis-py's sync client.
#[pyclass(frozen, name = "NativeAuthCache")]
pub(crate) struct NativeAuthCache(Arc<AuthCache>);
#[pymethods]
impl NativeAuthCache {
#[new]
#[pyo3(signature = (client, *, memory_entries, memory_ttl_seconds = None))]
fn new(
client: Option<PyRef<'_, NativeRedisClient>>,
memory_entries: usize,
memory_ttl_seconds: Option<f64>,
) -> PyResult<Self> {
Ok(Self(Arc::new(AuthCache::new(
client.map(|client| remote(&client)),
Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds)?,
},
))))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (key, request = None))]
fn get<'py>(
&self,
py: Python<'py>,
key: String,
request: Option<&NativeRequestRedis>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let request = request.map(|request| request.0.clone());
redis_async(py, "async_get_cache", kind_error, async move {
store
.get(request.as_ref(), &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", kind_error, async move {
store.get(None, &key).await.map(|value| value.map(PyValue))
})
}
/// `update_cache_ttl` and `update_in_memory_max_size`: both partitions keep what they hold.
#[pyo3(signature = (*, memory_entries, memory_ttl_seconds = None))]
fn reconfigure(&self, memory_entries: usize, memory_ttl_seconds: Option<f64>) -> PyResult<()> {
self.0.reconfigure(Settings {
memory_entries,
memory_ttl: ttl(memory_ttl_seconds)?,
});
Ok(())
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&key, value.0, ttl).await
})
}
#[pyo3(signature = (key, value, ttl_seconds = None))]
fn set_sync(
&self,
py: Python<'_>,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_sync(py, "set_cache", kind_error, async move {
store.set_str(&key, value.0, ttl).await
})
}
/// Memory now and the `SET` on the batch `send` picks.
#[pyo3(signature = (key, value, ttl_seconds, send))]
fn declare_set(
&self,
key: String,
value: PyValue,
ttl_seconds: Option<f64>,
send: &NativeSend,
) -> PyResult<NativePending> {
let written = self
.0
.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 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_delete_cache",
kind_error,
deleted,
std::convert::identity,
))
}
#[pyo3(signature = (entries, ttl_seconds = None))]
fn set_many<'py>(
&self,
py: Python<'py>,
entries: Vec<(String, PyValue)>,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let entries = entries
.into_iter()
.map(|(key, value)| (key, value.0))
.collect();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache_pipeline", kind_error, async move {
store.set_many(entries, ttl).await
})
}
/// 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>,
entries: Vec<(String, PyValue, Option<f64>)>,
send: &NativeSend,
) -> PyResult<Bound<'py, PyAny>> {
let entries = entries
.into_iter()
.map(|(key, value, seconds)| Ok((key, value.0, ttl(seconds)?)))
.collect::<PyResult<_>>()?;
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", 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", kind_error, async move {
store.delete(&keys).await
})
}
/// One `MGET` of keys the caller found missing in memory, on the batch `send` picks; `None`
/// without Redis.
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>(
&self,
py: Python<'py>,
key: String,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_if_present", 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(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(kind_error)
}
fn delete_local(&self, key: &str) -> PyResult<()> {
self.0.delete_local(key).map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
self.0.flush_local().map_err(kind_error)
}
fn len_local(&self) -> PyResult<PyLocalSizes> {
let (objects, key_objects) = self.0.len_local().map_err(kind_error)?;
Ok(PyLocalSizes {
objects,
key_objects,
})
}
}
/// Entries the management-object and the user-key partitions hold in memory.
#[pyclass(frozen, eq, get_all, name = "LocalSizes")]
#[derive(PartialEq)]
pub(crate) struct PyLocalSizes {
objects: usize,
key_objects: usize,
}
#[pymethods]
impl PyLocalSizes {
#[new]
fn new(objects: usize, key_objects: usize) -> Self {
Self {
objects,
key_objects,
}
}
}

View file

@ -1,4 +1,11 @@
pub(crate) mod auth_cache;
mod policy;
pub(crate) mod rate_limits;
pub(crate) mod request;
pub(crate) mod router_cooldowns;
pub(crate) mod router_usage;
pub(crate) mod sensitive_route_pins;
pub(crate) mod spend_counters;
mod value;
use value::PyValue;
@ -159,6 +166,7 @@ where
}
/// A Python `int` or `float` across the bridge, keeping which one it is.
#[derive(Clone, Copy, PartialEq)]
pub(crate) struct PyNumber(Number);
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
@ -188,6 +196,21 @@ impl<'py> IntoPyObject<'py> for PyNumber {
}
}
/// A store enum read from its wire spelling, so the bridge names no variant.
pub(crate) struct Spelled<T>(pub(crate) T);
impl<'a, 'py, T: std::str::FromStr> FromPyObject<'a, 'py> for Spelled<T> {
type Error = PyErr;
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
let spelling: String = value.extract()?;
spelling
.parse()
.map(Self)
.map_err(|_| PyValueError::new_err(format!("unknown spelling {spelling:?}")))
}
}
/// 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);

View file

@ -0,0 +1,151 @@
use std::time::Duration;
use litellm_redis_store::TierPolicy;
use pyo3::prelude::*;
use crate::{
coercion::{Field, FieldSpec, ProjectionError},
python_settings::{PythonSettings, Snapshot},
};
const COOLDOWN_REFRESH: FieldSpec<Duration> =
FieldSpec::new("cooldown_redis_read_interval", seconds);
const USAGE_REFRESH: FieldSpec<Duration> = FieldSpec::new("usage_redis_read_interval", seconds);
const TRACKED_READS: FieldSpec<usize> = FieldSpec::new("max_redis_batch_cache_size", count);
const MEMORY_TTL: FieldSpec<Option<Duration>> =
FieldSpec::new("default_in_memory_ttl", optional_seconds);
/// The tier values this Python process runs with, from its environment and `litellm` globals.
#[derive(Debug)]
pub(crate) struct RouterTierPolicies {
pub(crate) cooldowns: TierPolicy,
pub(crate) usage: TierPolicy,
}
pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult<RouterTierPolicies> {
Ok(project(&PythonSettings::RouterTiers.read(py)?)?)
}
fn project(snapshot: &Snapshot<'_>) -> Result<RouterTierPolicies, ProjectionError> {
let tracked_reads = snapshot.read(&TRACKED_READS)?;
let default_in_memory_ttl = snapshot.read(&MEMORY_TTL)?;
Ok(RouterTierPolicies {
cooldowns: TierPolicy {
refresh_after: snapshot.read(&COOLDOWN_REFRESH)?,
tracked_reads,
default_in_memory_ttl,
},
usage: TierPolicy {
refresh_after: snapshot.read(&USAGE_REFRESH)?,
tracked_reads,
default_in_memory_ttl,
},
})
}
fn seconds(field: &Field<'_>) -> Result<Duration, ProjectionError> {
match field.value().extract::<f64>() {
Ok(seconds) if seconds.is_finite() && seconds >= 0.0 => {
Ok(Duration::from_secs_f64(seconds))
}
_ => Err(field.invalid("a non-negative number of seconds")),
}
}
fn optional_seconds(field: &Field<'_>) -> Result<Option<Duration>, ProjectionError> {
match field.value().is_none() {
true => Ok(None),
false => seconds(field).map(Some),
}
}
fn count(field: &Field<'_>) -> Result<usize, ProjectionError> {
field
.value()
.extract::<usize>()
.map_err(|_| field.invalid("a non-negative integer"))
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use pyo3::{prelude::*, types::PyDict};
use super::project;
use crate::python_settings::PythonSettings;
fn projected(py: Python<'_>, fields: &str) -> Result<super::RouterTierPolicies, String> {
let locals = PyDict::new(py);
let source = std::ffi::CString::new(format!(
"import types\nsettings = types.SimpleNamespace({fields})"
))
.unwrap();
py.run(&source, Some(&locals), Some(&locals)).unwrap();
let snapshot =
PythonSettings::RouterTiers.snapshot(locals.get_item("settings").unwrap().unwrap());
project(&snapshot).map_err(|error| PyErr::from(error).to_string())
}
#[test]
fn the_host_values_are_each_familys_refresh_and_tracked_reads() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=2.5, usage_redis_read_interval=30, max_redis_batch_cache_size=42, default_in_memory_ttl=None",
)
.unwrap();
assert_eq!(
policies.cooldowns.refresh_after,
Duration::from_millis(2500)
);
assert_eq!(policies.usage.refresh_after, Duration::from_secs(30));
assert_eq!(policies.cooldowns.tracked_reads, 42);
assert_eq!(policies.usage.tracked_reads, 42);
assert_eq!(policies.cooldowns.default_in_memory_ttl, None);
});
}
#[test]
fn a_set_default_in_memory_ttl_reaches_both_familys_backfill() {
Python::initialize();
Python::attach(|py| {
let policies = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=7.5",
)
.unwrap();
assert_eq!(
policies.cooldowns.default_in_memory_ttl,
Some(Duration::from_millis(7500))
);
assert_eq!(
policies.usage.default_in_memory_ttl,
Some(Duration::from_millis(7500))
);
});
}
#[test]
fn a_value_python_would_not_run_with_is_a_configuration_error() {
Python::initialize();
Python::attach(|py| {
let negative = projected(
py,
"cooldown_redis_read_interval=-1, usage_redis_read_interval=10, max_redis_batch_cache_size=1000, default_in_memory_ttl=None",
)
.unwrap_err();
let text = projected(
py,
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many', default_in_memory_ttl=None",
)
.unwrap_err();
assert!(
negative.contains("cooldown_redis_read_interval"),
"{negative}"
);
assert!(text.contains("max_redis_batch_cache_size"), "{text}");
});
}
}

View file

@ -0,0 +1,591 @@
use std::{sync::Arc, time::Duration};
use indexmap::IndexMap;
use litellm_redis_store::{
Pending,
rate_limits::{
AcquireReply, Applied, AtomicCounter, AtomicReply, Counter, Entry, Gauge, GuardedIncrement,
GuardedReply, Increment, Kind, Local, RateLimits, Scope, WindowCount,
},
};
use pyo3::{exceptions::PyValueError, prelude::*};
use super::{
NativePending, PyNumber, Spelled, kind_error, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl,
};
use crate::redis::client::NativeRedisClient;
/// The v3 rate limiter's keyspace. Memory reads and writes are synchronous like `InMemoryCache`;
/// every Redis call is a future.
#[pyclass(frozen, name = "NativeRateLimits")]
pub(crate) struct NativeRateLimits(Arc<RateLimits>);
/// An `int` or `float` counter, a `str` window start, or a `{slot_id: acquired_at}` registry.
#[derive(FromPyObject)]
enum PyLocal {
Number(PyNumber),
Text(String),
Slots(IndexMap<String, f64>),
}
impl From<PyLocal> for Local {
fn from(local: PyLocal) -> Self {
match local {
PyLocal::Number(number) => Self::Number(number.0),
PyLocal::Text(text) => Self::Text(text),
PyLocal::Slots(slots) => Self::Slots(slots.into_iter().collect()),
}
}
}
/// What the memory tier holds, back as the Python value it was written as.
struct LocalValue(Local);
impl<'py> IntoPyObject<'py> for LocalValue {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
Ok(match self.0 {
Local::Number(number) => PyNumber(number).into_pyobject(py)?,
Local::Text(text) => text.into_pyobject(py)?.into_any(),
Local::Slots(slots) => slots
.into_iter()
.collect::<IndexMap<_, _>>()
.into_pyobject(py)?
.into_any(),
})
}
}
/// `litellm.redis.stores.rate_limits.RateLimitScope`.
#[derive(FromPyObject)]
struct PyScope {
#[pyo3(attribute)]
descriptor_key: String,
#[pyo3(attribute)]
descriptor_value: String,
}
impl From<PyScope> for Scope {
fn from(scope: PyScope) -> Self {
Self {
descriptor: scope.descriptor_key,
value: scope.descriptor_value,
}
}
}
/// `RateLimitCounter`.
#[derive(FromPyObject)]
struct PyCounter {
#[pyo3(attribute)]
scope: PyScope,
#[pyo3(attribute)]
kind: Spelled<Kind>,
}
impl From<PyCounter> for Counter {
fn from(counter: PyCounter) -> Self {
Self {
scope: counter.scope.into(),
kind: counter.kind.0,
}
}
}
/// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names.
#[derive(FromPyObject)]
struct PyEntry {
#[pyo3(attribute)]
suffix: String,
#[pyo3(attribute)]
scope: PyScope,
}
impl TryFrom<PyEntry> for Entry {
type Error = PyErr;
fn try_from(entry: PyEntry) -> PyResult<Self> {
Self::new(&entry.suffix, entry.scope.into())
.map_err(|_| PyValueError::new_err(format!("unknown key suffix {:?}", entry.suffix)))
}
}
/// `CounterIncrement`.
#[derive(FromPyObject)]
struct PyIncrement {
#[pyo3(attribute)]
counter: PyCounter,
#[pyo3(attribute)]
amount: PyNumber,
#[pyo3(attribute)]
ttl: Option<u64>,
}
impl From<PyIncrement> for Increment {
fn from(increment: PyIncrement) -> Self {
Self {
counter: increment.counter.into(),
amount: increment.amount.0,
ttl: increment.ttl.map(Duration::from_secs),
}
}
}
/// `AtomicCounter`.
#[derive(FromPyObject)]
struct PyAtomicCounter {
#[pyo3(attribute)]
counter: PyCounter,
#[pyo3(attribute)]
limit: i64,
#[pyo3(attribute)]
increment: i64,
#[pyo3(attribute)]
ttl: i64,
#[pyo3(attribute)]
window_size: i64,
}
impl From<PyAtomicCounter> for AtomicCounter {
fn from(counter: PyAtomicCounter) -> Self {
Self {
counter: counter.counter.into(),
limit: counter.limit,
increment: counter.increment,
ttl_seconds: counter.ttl,
window_size_seconds: counter.window_size,
}
}
}
fn all<P, T: From<P>>(records: Vec<P>) -> Vec<T> {
records.into_iter().map(T::from).collect()
}
#[pyclass(frozen, eq, get_all, name = "WindowCount")]
#[derive(PartialEq)]
pub(crate) struct PyWindowCount {
window_start: String,
count: i64,
}
#[pymethods]
impl PyWindowCount {
#[new]
fn new(window_start: String, count: i64) -> Self {
Self {
window_start,
count,
}
}
}
/// Per counter, the count after the increment and the start of the window it landed in.
#[pyclass(frozen, eq, get_all, name = "AtomicApplied")]
#[derive(PartialEq)]
pub(crate) struct PyAtomicApplied {
counters: Vec<(i64, i64)>,
}
#[pymethods]
impl PyAtomicApplied {
#[new]
fn new(counters: Vec<(i64, i64)>) -> Self {
Self { counters }
}
}
#[pyclass(frozen, eq, get_all, name = "AtomicOverLimit")]
#[derive(PartialEq)]
pub(crate) struct PyAtomicOverLimit {
index: usize,
current: i64,
limit: i64,
}
#[pymethods]
impl PyAtomicOverLimit {
#[new]
fn new(index: usize, current: i64, limit: i64) -> Self {
Self {
index,
current,
limit,
}
}
}
#[derive(IntoPyObject)]
enum PyAtomicReply {
Applied(PyAtomicApplied),
OverLimit(PyAtomicOverLimit),
}
#[pyclass(frozen, eq, get_all, name = "GuardedReply")]
#[derive(PartialEq)]
pub(crate) struct PyGuardedReply {
applied: bool,
count: i64,
}
#[pymethods]
impl PyGuardedReply {
#[new]
fn new(applied: bool, count: i64) -> Self {
Self { applied, count }
}
}
#[pyclass(frozen, eq, get_all, name = "GaugesAcquired")]
#[derive(PartialEq)]
pub(crate) struct PyGaugesAcquired {
in_flight: Vec<i64>,
}
#[pymethods]
impl PyGaugesAcquired {
#[new]
fn new(in_flight: Vec<i64>) -> Self {
Self { in_flight }
}
}
#[pyclass(frozen, eq, get_all, name = "GaugeOverLimit")]
#[derive(PartialEq)]
pub(crate) struct PyGaugeOverLimit {
index: usize,
in_flight: i64,
limit: i64,
}
#[pymethods]
impl PyGaugeOverLimit {
#[new]
fn new(index: usize, in_flight: i64, limit: i64) -> Self {
Self {
index,
in_flight,
limit,
}
}
}
#[derive(IntoPyObject)]
enum PyAcquireReply {
Acquired(PyGaugesAcquired),
OverLimit(PyGaugeOverLimit),
}
fn window_counts(counts: Vec<WindowCount>) -> Vec<PyWindowCount> {
counts
.into_iter()
.map(
|WindowCount {
window_start,
count,
}| PyWindowCount {
window_start,
count,
},
)
.collect()
}
fn atomic(reply: AtomicReply) -> PyAtomicReply {
match reply {
AtomicReply::Applied(applied) => PyAtomicReply::Applied(PyAtomicApplied {
counters: applied
.into_iter()
.map(
|Applied {
count,
window_start,
}| (count, window_start),
)
.collect(),
}),
AtomicReply::OverLimit {
index,
current,
limit,
} => PyAtomicReply::OverLimit(PyAtomicOverLimit {
index,
current,
limit,
}),
}
}
fn guarded(GuardedReply { applied, count }: GuardedReply) -> PyGuardedReply {
PyGuardedReply { applied, count }
}
fn acquire(reply: AcquireReply) -> PyAcquireReply {
match reply {
AcquireReply::Acquired(in_flight) => {
PyAcquireReply::Acquired(PyGaugesAcquired { in_flight })
}
AcquireReply::OverLimit {
index,
in_flight,
limit,
} => PyAcquireReply::OverLimit(PyGaugeOverLimit {
index,
in_flight,
limit,
}),
}
}
/// 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,
{
Some(super::pending(call_type, kind_error, declared?, reply))
}
#[pymethods]
impl NativeRateLimits {
#[new]
#[pyo3(signature = (client, backfill_ttl_seconds = None))]
fn new(
client: Option<PyRef<'_, NativeRedisClient>>,
backfill_ttl_seconds: Option<f64>,
) -> PyResult<Self> {
Ok(Self(Arc::new(RateLimits::new(
client.map(|client| remote(&client)),
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(&all(scopes), force)
}
fn get_local(&self, entry: PyEntry) -> PyResult<Option<LocalValue>> {
Ok(self
.0
.get_local(&entry.try_into()?)
.map_err(kind_error)?
.map(LocalValue))
}
#[pyo3(signature = (entry, value, ttl_seconds))]
fn set_local(&self, entry: PyEntry, value: PyLocal, ttl_seconds: Option<f64>) -> PyResult<()> {
self.0
.set_local(&entry.try_into()?, value.into(), ttl(ttl_seconds)?)
.map_err(kind_error)
}
/// `request` answers a key its pre-call reads already found absent without another `GET`.
#[pyo3(signature = (entry, *, backfill_ttl_seconds = None, request = None))]
fn get<'py>(
&self,
py: Python<'py>,
entry: PyEntry,
backfill_ttl_seconds: Option<f64>,
request: Option<&NativeRequestRedis>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let entry = Entry::try_from(entry)?;
let backfill_ttl = ttl(backfill_ttl_seconds)?;
let request = request.map(|request| request.0.clone());
redis_async(py, "async_get_cache", kind_error, async move {
store
.get(request.as_ref(), &entry, backfill_ttl)
.await
.map(|local| local.map(LocalValue))
})
}
fn sliding_window(
&self,
counters: Vec<PyCounter>,
now: i64,
window_size: i64,
send: &NativeSend,
) -> Option<NativePending> {
script(
"batch_rate_limiter_script",
self.0
.sliding_window(send.send(), &all(counters), now, window_size),
window_counts,
)
}
fn read_counters(&self, counters: Vec<PyCounter>, send: &NativeSend) -> Option<NativePending> {
script(
"batch_counter_read_script",
self.0.read_counters(send.send(), &all(counters)),
std::convert::identity,
)
}
fn check_and_increment(
&self,
counters: Vec<PyAtomicCounter>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"check_and_increment_by_n_script",
self.0.check_and_increment(send.send(), &all(counters)),
atomic,
)
}
fn guarded_increment(
&self,
counter: PyCounter,
expected_window_start: String,
amount: PyNumber,
ttl_seconds: Option<u64>,
send: &NativeSend,
) -> Option<NativePending> {
let increment = GuardedIncrement {
counter: counter.into(),
expected_window_start,
amount: amount.0,
ttl: ttl_seconds.map(Duration::from_secs),
};
script(
"window_guarded_token_increment_script",
self.0.guarded_increment(send.send(), &increment),
guarded,
)
}
fn parallel_acquire(
&self,
gauges: Vec<(PyScope, i64)>,
slot_ttl_seconds: u64,
slot_id: &str,
send: &NativeSend,
) -> Option<NativePending> {
let gauges: Vec<Gauge> = gauges
.into_iter()
.map(|(scope, limit)| Gauge {
scope: scope.into(),
limit,
})
.collect();
script(
"parallel_acquire_script",
self.0.parallel_acquire(
send.send(),
&gauges,
Duration::from_secs(slot_ttl_seconds),
slot_id,
),
acquire,
)
}
fn parallel_release(
&self,
gauges: Vec<PyScope>,
slot_id: &str,
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_release_script",
self.0.parallel_release(send.send(), &all(gauges), slot_id),
std::convert::identity,
)
}
fn parallel_count(
&self,
gauges: Vec<PyScope>,
slot_ttl_seconds: u64,
send: &NativeSend,
) -> Option<NativePending> {
script(
"parallel_count_script",
self.0.parallel_count(
send.send(),
&all(gauges),
Duration::from_secs(slot_ttl_seconds),
),
std::convert::identity,
)
}
fn token_increment(
&self,
increments: Vec<PyIncrement>,
send: &NativeSend,
) -> Option<NativePending> {
script(
"token_increment_script",
self.0.token_increment(send.send(), &all(increments)),
std::convert::identity,
)
}
/// `False` without Redis.
fn token_increment_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<PyIncrement>,
) -> bool {
self.0
.token_increment_post_call(&request.0, all(increments))
}
/// `False` without Redis.
fn parallel_release_post_call(
&self,
request: &NativeRequestRedis,
gauges: Vec<PyScope>,
slot_id: &str,
) -> bool {
self.0
.parallel_release_post_call(&request.0, &all(gauges), slot_id)
}
fn increment<'py>(
&self,
py: Python<'py>,
increment: PyIncrement,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increment = Increment::from(increment);
redis_async(py, "async_increment", kind_error, async move {
Ok(store.increment(&increment).await.map(PyNumber))
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
increments: Vec<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let increments: Vec<Increment> = all(increments);
redis_async(py, "async_increment_pipeline", kind_error, async move {
store.increment_many(&increments).await;
Ok(())
})
}
fn increment_many_post_call(&self, request: &NativeRequestRedis, increments: Vec<PyIncrement>) {
self.0
.increment_many_post_call(&request.0, &all(increments));
}
}

View file

@ -0,0 +1,183 @@
use std::sync::Arc;
use litellm_redis_store::{
Send,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use super::{
NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies,
redis_async, redis_sync, remote, request::NativeSend,
};
use crate::redis::client::NativeRedisClient;
/// One router's cooldown store. The router's sync completion path calls `put`, `active_sync`,
/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py; the
/// `async_` twins return awaitables for async callers.
#[pyclass(frozen, name = "NativeRouterCooldowns")]
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
/// `[(model_id, cooldown), ...]` the way `RedisCache` hands the router its batch read.
struct PyActive(Vec<(String, Cooldown)>);
impl<'py> IntoPyObject<'py> for PyActive {
type Target = PyAny;
type Output = Bound<'py, PyAny>;
type Error = PyErr;
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
active_to_py(py, self.0)?
.into_pyobject(py)
.map(Bound::into_any)
}
}
fn active_to_py(
py: Python<'_>,
active: Vec<(String, Cooldown)>,
) -> PyResult<Vec<Bound<'_, PyTuple>>> {
active
.into_iter()
.map(|(model_id, cooldown)| {
PyTuple::new(
py,
[
model_id.into_pyobject(py)?.into_any(),
PyData(cooldown).into_pyobject(py)?,
],
)
})
.collect()
}
#[pymethods]
impl NativeRouterCooldowns {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let remote = client.map(|client| remote(&client));
let policy = router_tier_policies(py)?.cooldowns;
Ok(Self(Arc::new(RouterCooldowns::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn put(
&self,
py: Python<'_>,
model_id: String,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<()> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_sync(py, "set_cache", kind_error, async move {
store.put_str(&model_id, &cooldown.0, ttl).await
})
}
fn increment_allowed_fails(
&self,
py: Python<'_>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<i64> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_sync(py, "increment_cache", kind_error, async move {
store
.increment_allowed_fails_int(&model_id, suffix.as_deref(), ttl)
.await
})
}
fn async_put<'py>(
&self,
py: Python<'py>,
model_id: String,
cooldown: PyData<Cooldown>,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.put(&model_id, &cooldown.0, ttl).await
})
}
fn async_increment_allowed_fails<'py>(
&self,
py: Python<'py>,
model_id: String,
suffix: Option<String>,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_increment", kind_error, async move {
store
.increment_allowed_fails(&model_id, suffix.as_deref(), ttl)
.await
})
}
/// 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(kind_error)?;
Ok(pending(
"async_batch_get_cache",
kind_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 active = self.0.active(Send::Now, &model_ids).map_err(kind_error)?;
let active = redis_sync(py, "batch_get_cache", kind_error, active)?;
active_to_py(py, active)
}
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
redis_sync(py, "batch_get_cache", kind_error, async move {
store
.minimum(&model_ids)
.await
.map(|minimum| minimum.map(PyNumber))
})
}
fn delete_local(&self, model_id: &str) -> PyResult<()> {
self.0.delete_local(model_id).map_err(kind_error)
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(kind_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(kind_error)? {
snapshot.set_item(key, PyData(cooldown))?;
}
Ok(snapshot)
}
}

View file

@ -0,0 +1,184 @@
use std::{sync::Arc, time::Duration};
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
use pyo3::{prelude::*, types::PyDict};
use super::{
PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote,
};
use crate::redis::client::NativeRedisClient;
/// `litellm.redis.stores.router_usage.UsageCounter`.
#[derive(FromPyObject)]
struct PyCounter {
#[pyo3(attribute)]
deployment_id: String,
#[pyo3(attribute)]
model: Option<String>,
#[pyo3(attribute)]
minute: String,
#[pyo3(attribute)]
kind: Spelled<Kind>,
}
impl From<PyCounter> for Counter {
fn from(counter: PyCounter) -> Self {
Self {
deployment_id: counter.deployment_id,
model: counter.model,
minute: counter.minute,
kind: counter.kind.0,
}
}
}
fn counters(counters: Vec<PyCounter>) -> Vec<Counter> {
counters.into_iter().map(Counter::from).collect()
}
/// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the
/// Rust router's sync calls, which have no event loop to await on.
#[pyclass(frozen, name = "NativeRouterUsage")]
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
#[pymethods]
impl NativeRouterUsage {
#[new]
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
let remote = client.map(|client| remote(&client));
let policy = router_tier_policies(py)?.usage;
Ok(Self(Arc::new(RouterUsage::new(remote, policy))))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
fn increment<'py>(
&self,
py: Python<'py>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_async(py, "async_increment", kind_error, async move {
Ok(store
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
.await
.map(PyNumber))
})
}
/// `_increment_deployment`'s `async_increment_pipeline`, always on its own round trip.
fn increment_deployment<'py>(
&self,
py: Python<'py>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment_pipeline", kind_error, async move {
store
.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
)
.await
})
}
fn get_many<'py>(
&self,
py: Python<'py>,
counters: Vec<PyCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_async(py, "async_batch_get_cache", kind_error, async move {
let read = store.get_many(&counters).await?;
Ok(read
.into_iter()
.map(|count| count.map(PyNumber))
.collect::<Vec<_>>())
})
}
fn increment_sync(
&self,
py: Python<'_>,
counter: PyCounter,
value: PyNumber,
ttl_seconds: u64,
) -> PyResult<Option<PyNumber>> {
let store = self.0.clone();
let counter = Counter::from(counter);
redis_sync(py, "increment_cache", kind_error, async move {
Ok(store
.increment(&counter, value.0, Duration::from_secs(ttl_seconds))
.await
.map(PyNumber))
})
}
fn increment_deployment_sync(
&self,
py: Python<'_>,
deployment_id: String,
model: Option<String>,
minute: String,
tokens: PyNumber,
requests: i64,
) -> PyResult<()> {
let store = self.0.clone();
redis_sync(py, "increment_pipeline", kind_error, async move {
store
.increment_deployment(
&deployment_id,
model.as_deref(),
&minute,
tokens.0,
requests,
)
.await
})
}
fn get_many_sync(
&self,
py: Python<'_>,
counters: Vec<PyCounter>,
) -> PyResult<Vec<Option<PyNumber>>> {
let store = self.0.clone();
let counters = self::counters(counters);
redis_sync(py, "batch_get_cache", kind_error, async move {
let read = store.get_many(&counters).await?;
Ok(read.into_iter().map(|count| count.map(PyNumber)).collect())
})
}
fn flush_local_counters(&self) -> PyResult<()> {
self.0.flush_local_counters().map_err(kind_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(kind_error)? {
snapshot.set_item(key, PyNumber(count))?;
}
Ok(snapshot)
}
}

View file

@ -0,0 +1,49 @@
use std::sync::Arc;
use litellm_redis_store::sensitive_route_pins::SensitiveRoutePins;
use pyo3::prelude::*;
use super::{duration, kind_error, redis_async, remote};
use crate::redis::client::NativeRedisClient;
#[pyclass(frozen, name = "NativeSensitiveRoutePins")]
pub(crate) struct NativeSensitiveRoutePins(Arc<SensitiveRoutePins>);
#[pymethods]
impl NativeSensitiveRoutePins {
#[new]
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
Self(Arc::new(SensitiveRoutePins::new(
client.map(|client| remote(&client)),
)))
}
fn get<'py>(
&self,
py: Python<'py>,
tenant: String,
session_id: String,
default_ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let default_ttl = duration(default_ttl_seconds)?;
redis_async(py, "async_get_cache", kind_error, async move {
store.get(&tenant, &session_id, default_ttl).await
})
}
fn set<'py>(
&self,
py: Python<'py>,
tenant: String,
session_id: String,
model: String,
ttl_seconds: f64,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = duration(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&tenant, &session_id, &model, ttl).await
})
}
}

View file

@ -0,0 +1,256 @@
use std::sync::Arc;
use litellm_redis_store::{
Number,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
};
use pyo3::{prelude::*, types::PyDict};
use super::{
NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote,
request::{NativeRequestRedis, NativeSend},
ttl, unwrap_all,
};
use crate::redis::client::NativeRedisClient;
/// A counter's value and whether Redis answered it, a clean miss included.
#[pyclass(frozen, eq, get_all, name = "SpendRead")]
#[derive(PartialEq)]
pub(crate) struct PySpendRead {
value: Option<PyNumber>,
from_redis: bool,
}
#[pymethods]
impl PySpendRead {
#[new]
fn new(value: Option<PyNumber>, from_redis: bool) -> Self {
Self { value, from_redis }
}
}
fn increments(increments: Vec<(PyData<SpendCounter>, PyNumber)>) -> Vec<(SpendCounter, Number)> {
increments
.into_iter()
.map(|(counter, amount)| (counter.0, amount.0))
.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.
#[pyclass(frozen, name = "NativeSpendCounters")]
pub(crate) struct NativeSpendCounters(Arc<SpendCounters>);
#[pymethods]
impl NativeSpendCounters {
#[new]
fn new(client: Option<PyRef<'_, NativeRedisClient>>) -> Self {
Self(Arc::new(SpendCounters::new(
client.map(|client| remote(&client)),
)))
}
/// Adds the Redis tier and keeps memory; `false` when one was already attached.
fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool {
self.0.attach(remote(&client))
}
#[getter]
fn has_redis(&self) -> bool {
self.0.has_redis()
}
#[getter]
fn redis_ttl_seconds(&self) -> Option<f64> {
self.0.redis_ttl().map(|ttl| ttl.as_secs_f64())
}
/// `("redis" | "memory", value)`: where the answer came from, a Redis miss included.
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", kind_error, async move {
let reading = store.read(&counter.0).await?;
Ok(PySpendRead {
value: reading.value().map(PyNumber),
from_redis: matches!(reading, Reading::Redis(_)),
})
})
}
/// 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: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_increment", kind_error, async move {
store.increment(&counter.0, amount.0).await.map(PyNumber)
})
}
fn increment_many<'py>(
&self,
py: Python<'py>,
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", kind_error, async move {
let totals = store.increment_many(&increments).await?;
Ok(totals.into_iter().map(PyNumber).collect::<Vec<_>>())
})
}
/// Declares the increments on the request's post-call batch; `False` when there is no Redis.
fn increment_many_post_call(
&self,
request: &NativeRequestRedis,
increments: Vec<(PyData<SpendCounter>, PyNumber)>,
) -> bool {
self.0
.increment_many_post_call(&request.0, &self::increments(increments))
}
fn set_max<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_set_max", kind_error, async move {
store
.set_max(&counter.0, spend.0)
.await
.map(|stored| stored.map(PyNumber))
})
}
#[pyo3(signature = (counter, spend, *, increment_if_lost))]
fn seed<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
increment_if_lost: bool,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let lost = if increment_if_lost {
LostSeed::Increment
} else {
LostSeed::Answer
};
redis_async(py, "async_set_cache", kind_error, async move {
store.seed(&counter.0, spend.0, lost).await.map(PyNumber)
})
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let ttl = ttl(ttl_seconds)?;
redis_async(py, "async_set_cache", kind_error, async move {
store.set(&counter.0, spend.0, ttl).await
})
}
fn delete<'py>(
&self,
py: Python<'py>,
counters: Vec<PyData<SpendCounter>>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
let counters = unwrap_all(counters);
redis_async(py, "async_delete_cache", kind_error, async move {
store.delete(&counters).await
})
}
fn refresh_ttl<'py>(
&self,
py: Python<'py>,
counter: PyData<SpendCounter>,
) -> PyResult<Bound<'py, PyAny>> {
let store = self.0.clone();
redis_async(py, "async_refresh_ttl", kind_error, async move {
store.refresh_ttl(&counter.0).await
})
}
fn get_local(&self, counter: PyData<SpendCounter>) -> PyResult<Option<PyNumber>> {
Ok(self
.0
.get_local(&counter.0)
.map_err(kind_error)?
.map(PyNumber))
}
#[pyo3(signature = (counter, spend, ttl_seconds = None))]
fn set_local(
&self,
counter: PyData<SpendCounter>,
spend: PyNumber,
ttl_seconds: Option<f64>,
) -> PyResult<()> {
self.0
.set_local(&counter.0, spend.0, ttl(ttl_seconds)?)
.map_err(kind_error)
}
fn delete_local(&self, counter: PyData<SpendCounter>) -> PyResult<()> {
self.0.delete_local(&counter.0).map_err(kind_error)
}
fn increment_local(
&self,
counter: PyData<SpendCounter>,
amount: PyNumber,
) -> PyResult<PyNumber> {
self.0
.increment_local(&counter.0, amount.0)
.map(PyNumber)
.map_err(kind_error)
}
fn flush_local(&self) -> PyResult<()> {
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(kind_error)? {
snapshot.set_item(key, PyNumber(spend))?;
}
Ok(snapshot)
}
}

View file

@ -1,11 +1,11 @@
- 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` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`)
- 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` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). 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` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) 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 family holds one `DualCache` per `DualCache` its Python held on `main`, and each `DualCache` method is the Python method it names, so a family op that was a `DualCache` call is one call here. Where `main`'s Python reached past `DualCache` to `in_memory_cache` or `redis_cache`, the family does the same through `memory()` and `redis()`, which keeps that reach visible. Caches a family's Python attached to one `RedisCache` share one `Arc<RemoteSlot>`. Only a family that batch-reads builds its cache with `DualCache::batch_reading` from a `TierPolicy` the host passes in, and the Python host fills it from 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. 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
- Key suffixes and wire spellings derive from strum on the family's enums. A value the family stores in Redis implements `FromValue` / `IntoValue` over `litellm_python_compat::Value`. An argument or reply is a typed record, never a positional tuple: a frozen Python dataclass read in `python-bridge` by a `#[derive(FromPyObject)]` struct with `#[pyo3(attribute)]` fields and converted into the family's type with `From`, and a reply returned as a frozen `#[pyclass]`. The bridge reads an enum through its strum spelling, never a hand-written match
- 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. 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

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

View file

@ -0,0 +1,258 @@
mod keys;
mod types;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_python_compat::{Value, json};
use litellm_redis_client::{Declared, RedisBatch, commands};
pub use self::keys::is_key_object;
pub use self::types::Settings;
use crate::{
Error, Number, Pending, RequestRedis, Send,
dual_cache::DualCache,
memory::{Clock, lock, system_clock},
remote::{Remote, RemoteSlot},
request::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: `UserApiKeyCache`'s two
/// `DualCache` 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: DualCache<Value>,
objects: DualCache<Value>,
redis: Arc<RemoteSlot>,
settings: Mutex<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 {
let redis = Arc::new(RemoteSlot::new(remote));
let partition = || DualCache::new(redis.clone(), settings.memory(), now.clone());
Self {
key_objects: partition(),
objects: partition(),
redis,
settings: Mutex::new(settings),
}
}
pub fn has_redis(&self) -> bool {
self.redis.get().is_some()
}
pub fn settings(&self) -> Settings {
*lock(&self.settings)
}
/// `UserApiKeyCache.update_cache_ttl` and `update_in_memory_max_size` on both partitions, which
/// keep what they hold.
pub fn reconfigure(&self, settings: Settings) {
*lock(&self.settings) = settings;
self.key_objects.reconfigure(settings.memory());
self.objects.reconfigure(settings.memory());
}
pub fn get_local(&self, key: &str) -> Result<Option<Value>, Error> {
Ok(self.partition(key).memory().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> {
let partition = self.partition(key);
Ok(partition
.memory()
.set(key, value, ttl.or(partition.default_in_memory_ttl()))?)
}
pub fn delete_local(&self, key: &str) -> Result<(), Error> {
Ok(self.partition(key).memory().delete(key)?)
}
pub fn flush_local(&self) -> Result<(), Error> {
self.key_objects.memory().flush()?;
Ok(self.objects.memory().flush()?)
}
/// Entries held by the management-object and the user-key partitions.
pub fn len_local(&self) -> Result<(usize, usize), Error> {
Ok((
self.objects.memory().len()?,
self.key_objects.memory().len()?,
))
}
/// `DualCache.async_get_cache`. A key `request` already read as missing from Redis is not read
/// again.
pub async fn get(
&self,
request: Option<&RequestRedis>,
key: &str,
) -> Result<Option<Value>, Error> {
self.partition(key).get(request, key, None).await
}
/// `DualCache.async_set_cache`, writing `json.dumps(value)`.
pub async fn set(&self, key: &str, value: Value, ttl: Option<Duration>) -> Result<(), Error> {
self.partition(key).set(key, value, ttl).await
}
/// `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> {
self.partition(key).set_str(key, value, ttl).await
}
/// `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> {
self.partition(key).declare_set(send, key, value, ttl)
}
/// `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.partition(key).declare_delete(send, 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<(), Error> {
let (key_objects, others): (Vec<_>, Vec<_>) =
entries.into_iter().partition(|(key, _)| is_key_object(key));
self.key_objects.set_pipeline(key_objects, ttl).await?;
self.objects.set_pipeline(others, ttl).await
}
/// The auth prefetch write-back: memory first, an entry without a TTL taking
/// `default_in_memory_ttl`, then `RedisCache.async_set_cache_pipeline_with_ttls`, which encodes
/// every value before it sends anything and gives an entry without a TTL the connection's
/// default. On a request batch the `SET`s are declared and nothing is left to wait for.
/// [`Send::Now`] is one pipeline of its own rather than one round trip per entry.
pub fn write_back(
&self,
send: Send<'_>,
entries: Vec<(String, Value, Option<Duration>)>,
) -> Result<Pending<()>, Error> {
let encoded = self
.redis
.get()
.map(|_| {
entries
.iter()
.map(|(_, value, _)| json::dumps(value))
.collect::<Result<Vec<_>, _>>()
})
.transpose();
let writes = entries
.into_iter()
.map(|(key, value, ttl)| {
self.set_local(&key, value, ttl)?;
Ok((key, ttl))
})
.collect::<Result<Vec<_>, Error>>()?;
let (Some(remote), Some(encoded)) = (self.redis.get(), encoded?) else {
return Ok(ready(()));
};
let sets = writes.iter().zip(encoded).map(|((key, ttl), value)| {
commands::set(key, value, ttl.unwrap_or(remote.default_ttl))
});
let Send::Now = send else {
let batch = send.batch(&remote.client);
sets.for_each(|set| drop(batch.add(set)));
return Ok(ready(()));
};
let pipeline = RedisBatch::pipelined(remote.client.clone());
let written: Vec<_> = sets.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`: the key-object partition and the rest as two
/// concurrent `DualCache.async_delete_cache_keys`. Both run even when one fails.
pub async fn delete(&self, keys: &[String]) -> Result<(), Error> {
let (key_objects, others): (Vec<String>, Vec<String>) =
keys.iter().cloned().partition(|key| is_key_object(key));
let (key_objects, others) = tokio::join!(
self.key_objects.delete_keys(&key_objects),
self.objects.delete_keys(&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, noting the keys Redis lacks as the request's misses. `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 client = self.redis.client()?;
let misses = send.misses(client);
let read = send.batch(client).add(commands::mget::<Vec<u8>>(
&keys.iter().map(String::as_str).collect::<Vec<_>>(),
));
let keys = keys.to_vec();
Some(Box::pin(async move {
let raw = read.resolve().await?;
if let Some(misses) = misses {
misses.note(
keys.into_iter()
.zip(&raw)
.filter(|(_, raw)| raw.is_none())
.map(|(key, _)| key),
);
}
raw.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.redis.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 partition(&self, key: &str) -> &DualCache<Value> {
if is_key_object(key) {
&self.key_objects
} else {
&self.objects
}
}
}

View file

@ -0,0 +1,35 @@
use std::time::Duration;
use litellm_python_compat::Value;
use crate::{DecodeError, codec::python_value, dual_cache::MemorySettings};
/// `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,
}
}
}
impl Settings {
pub(super) fn memory(self) -> MemorySettings {
MemorySettings {
entries: Some(self.memory_entries),
ttl: self.memory_ttl,
}
}
}
pub(super) fn decode(raw: Option<Vec<u8>>) -> Result<Option<Value>, DecodeError> {
raw.as_deref().map(python_value).transpose()
}

View file

@ -1,23 +1,17 @@
#![expect(
dead_code,
reason = "the store families built on these tiers and codecs land in the next PR"
)]
#![cfg_attr(
not(test),
expect(
unused_imports,
reason = "the store families built on these tiers and codecs land in the next PR"
)
)]
pub mod auth_cache;
mod codec;
mod dual_cache;
mod error;
mod failure_log;
mod memory;
mod number;
pub mod rate_limits;
mod remote;
mod request;
pub mod router_cooldowns;
pub mod router_usage;
pub mod sensitive_route_pins;
pub mod spend_counters;
mod value;
pub use codec::DecodeError;

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

View file

@ -0,0 +1,260 @@
//! The limiter's Lua, byte for byte as `main`'s Python registers it, so both share one script SHA.
pub(super) const BATCH_RATE_LIMITER: &str = r#"
local results = {}
local now = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
-- Process each window/counter pair
for i = 1, #KEYS, 2 do
local window_key = KEYS[i]
local counter_key = KEYS[i + 1]
local increment_value = 1
-- Check if window exists and is valid
local window_start = redis.call('GET', window_key)
if not window_start or (now - tonumber(window_start)) >= window_size then
-- Reset window and counter
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment_value)
redis.call('EXPIRE', window_key, window_size)
redis.call('EXPIRE', counter_key, window_size)
table.insert(results, tostring(now)) -- window_start
table.insert(results, increment_value) -- counter
else
local counter = redis.call('INCR', counter_key)
-- This happens when window_key exists but counter_key doesn't (e.g., tokens key
-- created after requests key when both share the same window_key)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 then
redis.call('EXPIRE', counter_key, window_size)
end
table.insert(results, window_start) -- window_start
table.insert(results, counter) -- counter
end
end
return results
"#;
pub(super) const BATCH_COUNTER_READ: &str = r#"return redis.call('MGET', unpack(KEYS))"#;
pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#"
-- Atomic check-and-increment-by-N across one or more descriptors.
-- All-or-nothing: if any descriptor would exceed its limit, no counter is
-- modified.
--
-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied
-- timestamp so that window resets are deterministic across replicas with
-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the
-- TOCTOU window across multi-replica deployments.
--
-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor.
-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]:
-- ARGV[(i-1)*4 + 1] = limit
-- ARGV[(i-1)*4 + 2] = increment
-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets)
-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length)
--
-- Return on success:
-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... }
-- Return on over-limit: { 1, descriptor_index, current_counter, limit }
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local descriptor_count = #KEYS / 2
local reset_windows = {}
-- Pass 1: read state, validate. Abort without writing if any over limit.
local descriptor_state = {}
for i = 1, descriptor_count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local arg_base = (i - 1) * 4 + 1
local limit = tonumber(ARGV[arg_base])
local increment = tonumber(ARGV[arg_base + 1])
local window_size = tonumber(ARGV[arg_base + 3])
local window_start = redis.call('GET', window_key)
local window_expired = (not window_start) or
((now - tonumber(window_start)) >= window_size)
local current_counter
if window_expired then
current_counter = 0
else
current_counter = tonumber(redis.call('GET', counter_key) or 0)
end
local blocked
if increment > 0 then
blocked = current_counter + increment > limit
else
blocked = current_counter >= limit
end
if blocked then
return { 1, i, current_counter, limit }
end
descriptor_state[i] = { window_expired, current_counter, window_start }
end
-- Pass 2: all checks passed. Apply increments.
local results = { 0 }
for i = 1, descriptor_count do
local window_key = KEYS[(i - 1) * 2 + 1]
local counter_key = KEYS[(i - 1) * 2 + 2]
local arg_base = (i - 1) * 4 + 1
local increment = tonumber(ARGV[arg_base + 1])
local ttl = tonumber(ARGV[arg_base + 2])
local window_size = tonumber(ARGV[arg_base + 3])
local window_expired = descriptor_state[i][1]
local active_window_start
if window_expired then
active_window_start = now
if not reset_windows[window_key] then
local prefix = string.sub(window_key, 1, -(#':window') - 1)
redis.call('DEL', prefix .. ':requests', prefix .. ':tokens')
reset_windows[window_key] = true
end
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment)
redis.call('EXPIRE', window_key, window_size)
if ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, increment)
else
active_window_start = tonumber(descriptor_state[i][3])
local new_counter = redis.call('INCRBY', counter_key, increment)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, new_counter)
end
table.insert(results, active_window_start)
end
return results
"#;
pub(super) const WINDOW_GUARDED_TOKEN_INCREMENT: &str = r#"
local results = {}
for i = 1, #KEYS, 2 do
local window_key = KEYS[i]
local counter_key = KEYS[i + 1]
local arg_base = ((i - 1) / 2) * 3 + 1
local expected_window_start = ARGV[arg_base]
local increment = tonumber(ARGV[arg_base + 1])
local ttl = tonumber(ARGV[arg_base + 2])
local active_window_start = redis.call('GET', window_key)
if active_window_start and active_window_start == expected_window_start then
local new_counter = redis.call('INCRBY', counter_key, increment)
local current_ttl = redis.call('TTL', counter_key)
if current_ttl == -1 and ttl > 0 then
redis.call('EXPIRE', counter_key, ttl)
end
table.insert(results, 1)
table.insert(results, new_counter)
else
table.insert(results, 0)
table.insert(results, tonumber(redis.call('GET', counter_key) or 0))
end
end
return results
"#;
pub(super) const PARALLEL_ACQUIRE: &str = r#"
-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge.
-- Each gauge key is a sorted set of per-request slot ids scored by acquire
-- time (Redis server clock). In-flight requests are counted by ZCARD after
-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM
-- counters the gauge is never reset while requests are in flight, a
-- rejected request never occupies a slot, and a slot leaked by a crashed
-- worker self-heals after the slot TTL even under continuous traffic.
--
-- KEYS: one gauge zset key per descriptor.
-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id).
-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }.
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
for i = 1, #KEYS do
local limit = tonumber(ARGV[(i - 1) * 3 + 1])
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl)
local in_flight = redis.call('ZCARD', KEYS[i])
if in_flight + 1 > limit then
return { 1, i, in_flight, limit }
end
end
local results = { 0 }
for i = 1, #KEYS do
local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2])
local slot_id = ARGV[(i - 1) * 3 + 3]
redis.call('ZADD', KEYS[i], now, slot_id)
redis.call('EXPIRE', KEYS[i], slot_ttl)
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"#;
pub(super) const PARALLEL_RELEASE: &str = r#"
-- Release one slot per gauge key by removing this request's slot id.
-- ZREM of an absent member (or key) is a no-op, so a release without a
-- matching acquire (proxy-side rejection, double-fired callback, slot
-- already expired) can never free a slot owned by another request.
-- KEYS: gauge zset keys. ARGV: per-key slot_id.
-- Returns the remaining in-flight count per key.
local results = {}
for i = 1, #KEYS do
redis.call('ZREM', KEYS[i], ARGV[i])
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"#;
pub(super) const PARALLEL_COUNT: &str = r#"
-- Read the current in-flight count per gauge key (prunes expired slots
-- first so leaked slots do not inflate the reading).
-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds.
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local results = {}
for i = 1, #KEYS do
redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i]))
table.insert(results, redis.call('ZCARD', KEYS[i]))
end
return results
"#;
pub(super) const TOKEN_INCREMENT: &str = r#"
local results = {}
-- Process each key/increment_value/ttl triplet
for i = 1, #KEYS do
local key = KEYS[i]
local increment_value = tonumber(ARGV[i * 2 - 1])
local ttl_seconds = tonumber(ARGV[i * 2])
-- Increment the value
local new_value = redis.call('INCRBYFLOAT', key, increment_value)
-- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL
-- ttl_seconds can be 0 (no TTL) or positive (set TTL)
if ttl_seconds and ttl_seconds > 0 then
local current_ttl = redis.call('TTL', key)
if current_ttl == -1 then
redis.call('EXPIRE', key, ttl_seconds)
end
end
table.insert(results, new_value)
end
return results
"#;

View file

@ -0,0 +1,304 @@
mod keys;
mod lua;
mod scripts;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Declared, RedisBatch, RedisClient, Slot};
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,
dual_cache::{DualCache, MemorySettings, PipelineIncrement},
memory::{Clock, 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 {
cache: DualCache<Local>,
}
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 {
cache: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings {
entries: None,
ttl: backfill_ttl,
},
now,
),
}
}
pub fn has_redis(&self) -> bool {
self.cache.redis().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
.cache
.redis()
.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.cache.memory().get(&entry.key())?)
}
pub fn set_local(
&self,
entry: &Entry,
value: Local,
ttl: Option<Duration>,
) -> Result<(), Error> {
Ok(self.cache.memory().set(&entry.key(), value, ttl)?)
}
/// `DualCache.async_get_cache`, backfilling memory for `backfill_ttl` or the store's default.
/// A key `request` already read as missing from Redis is not read again.
pub async fn get(
&self,
request: Option<&RequestRedis>,
entry: &Entry,
backfill_ttl: Option<Duration>,
) -> Result<Option<Local>, Error> {
self.cache.get(request, &entry.key(), 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;
};
crate::failure_log::redis_failure!(
&error,
"TTL preservation failed, falling back to regular pipeline"
);
store.increment_many(&increments).await;
});
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 {
crate::failure_log::redis_failure!(
&error,
"parallel_release_script failed, the slot stays released in memory only"
);
}
});
true
}
/// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an
/// entry memory cannot add to with `None`.
pub async fn increment(&self, increment: &Increment) -> Option<Number> {
self.cache
.increment(&increment.counter.key(), increment.amount, increment.ttl)
.await
}
/// `DualCache.async_increment_cache_pipeline`, whose failures only log.
pub async fn increment_many(&self, increments: &[Increment]) {
self.cache.increment_pipeline(&pipeline(increments)).await;
}
/// `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]) {
self.cache
.increment_pipeline_post_call(request, &pipeline(increments));
}
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))
}
fn batch(&self, send: Send<'_>) -> Option<RedisBatch> {
self.cache.redis().batch(send)
}
}
fn pipeline(increments: &[Increment]) -> Vec<PipelineIncrement> {
increments
.iter()
.map(|increment| PipelineIncrement {
key: increment.counter.key(),
amount: increment.amount,
ttl: increment.ttl,
})
.collect()
}

View file

@ -0,0 +1,303 @@
use std::{sync::LazyLock, time::Duration};
use litellm_python_compat::repr::float_repr;
use litellm_redis_client::{
Command, Declared, Error, RedisBatch, Script, Value,
commands::{self, unexpected_reply},
};
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> =
LazyLock::new(|| Script::new(lua::BATCH_RATE_LIMITER));
static BATCH_COUNTER_READ: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::BATCH_COUNTER_READ));
static CHECK_AND_INCREMENT_BY_N: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::CHECK_AND_INCREMENT_BY_N));
static WINDOW_GUARDED_TOKEN_INCREMENT: LazyLock<Script> =
LazyLock::new(|| Script::new(lua::WINDOW_GUARDED_TOKEN_INCREMENT));
static PARALLEL_ACQUIRE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_ACQUIRE));
static PARALLEL_RELEASE: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_RELEASE));
static PARALLEL_COUNT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::PARALLEL_COUNT));
static TOKEN_INCREMENT: LazyLock<Script> = LazyLock::new(|| Script::new(lua::TOKEN_INCREMENT));
pub(super) fn sliding_window(
batch: &RedisBatch,
counters: &[Counter],
now: i64,
window_size: i64,
) -> Declared<Vec<WindowCount>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
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(
batch: &RedisBatch,
counters: &[Counter],
) -> Declared<Vec<Option<String>>> {
let keys = counters
.iter()
.flat_map(|counter| [counter.scope.window_key(), counter.key()])
.collect::<Vec<_>>();
batch.add(eval(&BATCH_COUNTER_READ, &keys, Vec::new()).map(optional_texts))
}
pub(super) fn check_and_increment(
batch: &RedisBatch,
counters: &[AtomicCounter],
) -> Declared<AtomicReply> {
let keys = counters
.iter()
.flat_map(|counter| [counter.counter.scope.window_key(), counter.counter.key()])
.collect::<Vec<_>>();
let args = counters
.iter()
.flat_map(|counter| {
[
counter.limit,
counter.increment,
counter.ttl_seconds,
counter.window_size_seconds,
]
.map(|arg| arg.to_string())
})
.collect();
batch.add(eval(&CHECK_AND_INCREMENT_BY_N, &keys, args).map(atomic_reply))
}
pub(super) fn guarded_increment(
batch: &RedisBatch,
increment: &GuardedIncrement,
) -> Declared<GuardedReply> {
let keys = [
increment.counter.scope.window_key(),
increment.counter.key(),
];
let args = vec![
increment.expected_window_start.clone(),
python_arg(increment.amount),
ttl_arg(increment.ttl),
];
batch.add(eval(&WINDOW_GUARDED_TOKEN_INCREMENT, &keys, args).map(guarded_reply))
}
pub(super) fn parallel_acquire(
batch: &RedisBatch,
gauges: &[Gauge],
slot_ttl: Duration,
slot_id: &str,
) -> Declared<AcquireReply> {
let keys = gauges
.iter()
.map(|gauge| gauge.scope.gauge_key())
.collect::<Vec<_>>();
let args = gauges
.iter()
.flat_map(|gauge| {
[
gauge.limit.to_string(),
slot_ttl.as_secs().to_string(),
slot_id.to_owned(),
]
})
.collect();
batch.add(eval(&PARALLEL_ACQUIRE, &keys, args).map(acquire_reply))
}
pub(super) fn parallel_release(
batch: &RedisBatch,
gauges: &[Scope],
slot_id: &str,
) -> Declared<Vec<i64>> {
let keys = gauges.iter().map(Scope::gauge_key).collect::<Vec<_>>();
let args = gauges.iter().map(|_| slot_id.to_owned()).collect();
batch.add(eval(&PARALLEL_RELEASE, &keys, args).map(integers))
}
pub(super) fn parallel_count(
batch: &RedisBatch,
gauges: &[Scope],
slot_ttl: Duration,
) -> 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();
batch.add(eval(&PARALLEL_COUNT, &keys, args).map(integers))
}
pub(super) fn token_increment(batch: &RedisBatch, increments: &[Increment]) -> Declared<Vec<f64>> {
let keys = increments
.iter()
.map(|increment| increment.counter.key())
.collect::<Vec<_>>();
let args = increments
.iter()
.flat_map(|increment| [python_arg(increment.amount), ttl_arg(increment.ttl)])
.collect();
batch.add(eval(&TOKEN_INCREMENT, &keys, args).map(floats))
}
/// `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`.
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.
fn python_arg(amount: Number) -> String {
match amount {
Number::Int(value) => value.to_string(),
Number::Float(value) => float_repr(value),
}
}
/// `ttl or 0`: the scripts read a zero TTL as "leave the expiry alone".
fn ttl_arg(ttl: Option<Duration>) -> String {
ttl.map_or(0, |ttl| ttl.as_secs()).to_string()
}
fn items(reply: Value) -> Result<Vec<Value>, Error> {
match reply {
Value::Array(items) => Ok(items),
_ => Err(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(unexpected_reply(
"rate limiter script reply holds a non-integer where a count belongs",
)),
}
}
fn text(value: &Value) -> Result<Option<String>, Error> {
match value {
Value::Nil => Ok(None),
Value::BulkString(bytes) => String::from_utf8(bytes.clone()).map(Some).map_err(|_| {
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(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_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(|_| unexpected_reply("rate limiter script reply names an index below 1"))
}
fn window_counts(reply: Value) -> Result<Vec<WindowCount>, Error> {
items(reply)?
.chunks(2)
.map(|pair| match pair {
[window_start, count] => Ok(WindowCount {
window_start: text(window_start)?.ok_or_else(|| {
unexpected_reply("rate limiter script reply holds a nil window start")
})?,
count: integer(count)?,
}),
_ => Err(unexpected_reply(
"rate limiter script reply has an odd length",
)),
})
.collect()
}
fn optional_texts(reply: Value) -> Result<Vec<Option<String>>, Error> {
items(reply)?.iter().map(text).collect()
}
fn integers(reply: Value) -> Result<Vec<i64>, Error> {
items(reply)?.iter().map(integer).collect()
}
fn floats(reply: Value) -> Result<Vec<f64>, Error> {
items(reply)?.into_iter().map(float).collect()
}
fn atomic_reply(reply: Value) -> Result<AtomicReply, Error> {
let items = items(reply)?;
match items.split_first() {
Some((Value::Int(0), rest)) => rest
.chunks(2)
.map(|pair| match pair {
[count, window_start] => Ok(Applied {
count: integer(count)?,
window_start: integer(window_start)?,
}),
_ => Err(unexpected_reply(
"rate limiter script reply has an odd length",
)),
})
.collect::<Result<_, _>>()
.map(AtomicReply::Applied),
Some((Value::Int(1), [index_value, current, limit])) => Ok(AtomicReply::OverLimit {
index: index(index_value)?,
current: integer(current)?,
limit: integer(limit)?,
}),
_ => Err(unexpected_reply("rate limiter script reply has no status")),
}
}
fn guarded_reply(reply: Value) -> Result<GuardedReply, Error> {
match items(reply)?.as_slice() {
[applied, count] => Ok(GuardedReply {
applied: integer(applied)? == 1,
count: integer(count)?,
}),
_ => Err(unexpected_reply(
"rate limiter script reply is not an (applied, count) pair",
)),
}
}
fn acquire_reply(reply: Value) -> Result<AcquireReply, Error> {
let items = items(reply)?;
match items.split_first() {
Some((Value::Int(0), rest)) => rest
.iter()
.map(integer)
.collect::<Result<_, _>>()
.map(AcquireReply::Acquired),
Some((Value::Int(1), [index_value, in_flight, limit])) => Ok(AcquireReply::OverLimit {
index: index(index_value)?,
in_flight: integer(in_flight)?,
limit: integer(limit)?,
}),
_ => Err(unexpected_reply("rate limiter script reply has no status")),
}
}

View file

@ -0,0 +1,117 @@
use std::time::Duration;
use litellm_python_compat::Value;
use super::keys::{Counter, Scope};
use crate::{DecodeError, FromValue, Number};
/// `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,
},
}
/// 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)>),
}
impl From<Number> for Local {
fn from(number: Number) -> Self {
Self::Number(number)
}
}
impl crate::memory::Counter for Local {
fn as_number(&self) -> Option<Number> {
match self {
Self::Number(number) => Some(*number),
Self::Text(_) | Self::Slots(_) => None,
}
}
}
/// A `RedisPipelineIncrementOperation`.
#[derive(Clone, Debug, PartialEq)]
pub struct Increment {
pub counter: Counter,
pub amount: Number,
pub ttl: Option<Duration>,
}
/// A rate-limit value as `RedisCache._get_cache_logic` reads it.
impl FromValue for Local {
fn from_value(value: &Value) -> Result<Self, DecodeError> {
match value {
Value::Str(text) => Ok(Self::Text(text.clone())),
value => Number::from_value(value)
.map(Local::Number)
.ok_or(DecodeError::Shape(
"a rate-limit value is a number or a string",
)),
}
}
}

View file

@ -0,0 +1,17 @@
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"),
}
}

View file

@ -0,0 +1,205 @@
mod keys;
mod types;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::{Amount, Command, Expiry, commands};
pub use self::types::Cooldown;
use crate::{
Error, Number, Pending, Send,
dual_cache::{BatchReads, DualCache, TierPolicy},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
};
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: DualCache<Cooldown, BatchReads>,
fails: DualCache<i64>,
redis: Arc<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 {
let redis = Arc::new(RemoteSlot::new(remote));
Self {
fails: DualCache::new(redis.clone(), policy.memory(), now.clone()),
cooldowns: DualCache::batch_reading(redis.clone(), policy, now),
redis,
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.redis.attach(remote)
}
pub fn has_redis(&self) -> bool {
self.redis.get().is_some()
}
/// `DualCache.async_set_cache`: memory first, then a `SET` of `json.dumps(value)` with
/// `EX int(ttl)` whose failure only logs.
pub async fn put(
&self,
model_id: &str,
cooldown: &Cooldown,
ttl: Duration,
) -> Result<(), Error> {
self.cooldowns
.set(&keys::cooldown(model_id), cooldown.clone(), Some(ttl))
.await
}
/// `DualCache.set_cache`: memory first, then a `SET` of `str(value)` with `EX int(ttl)` whose
/// failure only logs.
pub async fn put_str(
&self,
model_id: &str,
cooldown: &Cooldown,
ttl: Duration,
) -> Result<(), Error> {
self.cooldowns
.set_str(&keys::cooldown(model_id), cooldown.clone(), Some(ttl))
.await
}
/// `DualCache.async_increment_cache`: the fleet-wide fail count through `INCRBYFLOAT`. 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 command =
commands::incr_by_float(&key, Amount::Int(1), Some(Expiry::KeepExisting(ttl)));
self.count_fail(&key, ttl, command, Number::Float).await
}
/// `DualCache.increment_cache`: the fleet-wide fail count through `INCRBY`. A Redis failure
/// falls back to this router's own count.
pub async fn increment_allowed_fails_int(
&self,
model_id: &str,
suffix: Option<&str>,
ttl: Duration,
) -> Result<i64, Error> {
let key = keys::allowed_fails(model_id, suffix);
let command = commands::incr_by(&key, 1, ttl);
self.count_fail(&key, ttl, command, Number::Int).await
}
async fn count_fail<T>(
&self,
key: &str,
ttl: Duration,
command: Command<T>,
total: fn(T) -> Number,
) -> Result<i64, Error> {
let local = self
.fails
.memory()
.update(key, Some(ttl), |count| Ok(count.unwrap_or(0) + 1))?;
let Some(client) = self.redis.client() else {
return Ok(local);
};
match client.send(command).await {
Ok(count) => Ok(total(count).truncate()),
Err(error) => {
crate::failure_log::redis_failure!(
&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_batch_get(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_batch_get(read).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
.batch_get(&keys::cooldowns(model_ids))
.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.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, 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)
}
}

View file

@ -0,0 +1,61 @@
use std::time::Duration;
use litellm_python_compat::Value;
use crate::{DecodeError, FromValue, IntoValue, Number};
#[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()
}
}
/// 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")?,
})
}
}

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

View file

@ -0,0 +1,110 @@
mod keys;
use std::{sync::Arc, time::Duration};
pub use self::keys::{Counter, Kind};
use crate::{
Error, Number,
dual_cache::{BatchReads, DualCache, PipelineIncrement, TierPolicy},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
};
/// `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: DualCache<Number, BatchReads>,
}
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: DualCache::batch_reading(Arc::new(RemoteSlot::new(remote)), policy, now),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.usage.redis().attach(remote)
}
pub fn has_redis(&self) -> bool {
self.usage.redis().get().is_some()
}
/// `DualCache.async_increment_cache`. A Redis failure answers with the memory count, and an
/// entry memory cannot add to with `None`.
pub async fn increment(
&self,
counter: &Counter,
value: Number,
ttl: Duration,
) -> Option<Number> {
self.usage.increment(&counter.key(), value, Some(ttl)).await
}
/// One deployment's tokens and requests, `async_increment_pipeline` on its own round trip, with
/// the totals mirrored back into memory as `int`s.
pub async fn increment_deployment(
&self,
deployment_id: &str,
model: Option<&str>,
minute: &str,
tokens: Number,
requests: i64,
) -> Result<(), Error> {
let increment = |kind, amount| PipelineIncrement {
key: Counter {
deployment_id: deployment_id.to_owned(),
model: model.map(str::to_owned),
minute: minute.to_owned(),
kind,
}
.key(),
amount,
ttl: Some(COUNTER_TTL),
};
let increments = [
increment(Kind::Tpm, tokens),
increment(Kind::Rpm, Number::Int(requests)),
];
let Some(totals) = self.usage.increment_pipeline(&increments).await else {
return Ok(());
};
if self.usage.redis().get().is_none() {
return Ok(());
}
increments
.iter()
.zip(totals)
.try_for_each(|(increment, total)| {
self.usage.memory().set(
&increment.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<_>>();
self.usage.batch_get(&keys).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()?)
}
}

View file

@ -0,0 +1,3 @@
pub(super) fn pin(tenant: &str, session_id: &str) -> String {
format!("{{sensitive_route:{tenant}:{session_id}}}:model")
}

View file

@ -0,0 +1,111 @@
mod keys;
use std::{sync::Arc, time::Duration};
use litellm_python_compat::{Value, json};
use litellm_redis_client::{Ttl, commands};
use crate::{
DecodeError, Error,
codec::python_value,
dual_cache::{DualCache, MemorySettings},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
value::text,
};
/// The model each session is pinned to: memory per worker, Redis shared by every pod. Like
/// `SensitiveDataRoutingHandler` on `main`, it reaches past `DualCache` to Redis first.
pub struct SensitiveRoutePins {
pins: DualCache<String>,
}
impl SensitiveRoutePins {
pub fn new(remote: Option<Remote>) -> Self {
Self::with_clock(remote, system_clock())
}
pub fn with_clock(remote: Option<Remote>, now: Clock) -> Self {
Self {
pins: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings::default(),
now,
),
}
}
/// Redis first, kept in memory for the key's remaining TTL. A Redis miss or failure serves memory.
pub async fn get(
&self,
tenant: &str,
session_id: &str,
default_ttl: Duration,
) -> Result<Option<String>, Error> {
let key = keys::pin(tenant, session_id);
if let Some(model) = self.read_remote(&key, default_ttl).await? {
return Ok(Some(model));
}
Ok(self.pins.memory().get(&key)?)
}
/// `SET ... EX` whose failure only logs, then memory.
pub async fn set(
&self,
tenant: &str,
session_id: &str,
model: &str,
ttl: Duration,
) -> Result<(), Error> {
let key = keys::pin(tenant, session_id);
if let Some(client) = self.pins.redis().client() {
let encoded = json::dumps(&Value::Str(model.to_owned()))?;
if let Err(error) = client.send(commands::set(&key, encoded, ttl)).await {
crate::failure_log::redis_failure!(
&error,
"writing a sensitive route pin to Redis failed"
);
}
}
self.pins.memory().set(&key, model.to_owned(), Some(ttl))?;
Ok(())
}
/// `GET`, then on a hit its own `TTL` trip, as `RedisCache.async_get_cache` and `async_get_ttl`
/// make them. A key without a TTL, or a failed `TTL`, keeps `default_ttl`.
async fn read_remote(&self, key: &str, default_ttl: Duration) -> Result<Option<String>, Error> {
let Some(client) = self.pins.redis().client() else {
return Ok(None);
};
let read = client
.send(commands::get::<Vec<u8>>(key))
.await
.map_err(Error::from)
.and_then(|raw| raw.as_deref().map(decode).transpose().map_err(Error::from));
let model = match read {
Ok(Some(model)) => model,
Ok(None) => return Ok(None),
Err(error) => {
crate::failure_log::redis_failure!(
&error,
"reading a sensitive route pin from Redis failed, serving memory"
);
return Ok(None);
}
};
let ttl = match client.send(commands::ttl(key)).await {
Ok(Ttl::Expires(remaining)) => remaining,
Ok(Ttl::Missing | Ttl::Persistent) => default_ttl,
Err(error) => {
tracing::debug!(%error, "Redis TTL Error");
default_ttl
}
};
self.pins.memory().set(key, model.clone(), Some(ttl))?;
Ok(Some(model))
}
}
fn decode(raw: &[u8]) -> Result<String, DecodeError> {
text(&python_value(raw)?, "a pinned model name")
}

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

View file

@ -0,0 +1,380 @@
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,
dual_cache::{DualCache, MemorySettings},
memory::{Clock, system_clock},
remote::{Remote, RemoteSlot},
};
/// `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. Like `spend_counter_cache`'s callers on `main`, most ops reach past
/// `DualCache` to its Redis and its memory.
pub struct SpendCounters {
cache: DualCache<Number>,
}
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 {
cache: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings::default(),
now,
),
}
}
/// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached.
pub fn attach(&self, remote: Remote) -> bool {
self.cache.redis().attach(remote)
}
pub fn has_redis(&self) -> bool {
self.cache.redis().get().is_some()
}
/// The TTL every Redis write gives a counter, `None` without Redis.
pub fn redis_ttl(&self) -> Option<Duration> {
self.cache.redis().get().map(|remote| remote.default_ttl)
}
/// `read_spend_counter_cache_value`: Redis first. `RedisCache.async_get_cache` swallows an unreachable
/// `GET` as a miss, so an open breaker, a transport error, or a value `float()` refuses reads memory instead.
pub async fn read(&self, counter: &SpendCounter) -> Result<Reading, Error> {
let Some(client) = self.cache.redis().client() else {
return Ok(Reading::Memory(self.get_local(counter)?));
};
match client.send(commands::get::<Vec<u8>>(&counter.key())).await {
Ok(raw) => match types::decode(raw) {
Ok(value) => Ok(Reading::Redis(value)),
Err(error) => {
tracing::debug!(%error, "spend counter Redis value did not decode, falling back to memory");
Ok(Reading::Memory(self.get_local(counter)?))
}
},
Err(litellm_redis_client::Error::BreakerOpen) => {
Ok(Reading::Memory(self.get_local(counter)?))
}
Err(error) => {
crate::failure_log::redis_failure!(&error, "spend counter Redis read failed");
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.cache.redis().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.cache.redis().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.cache.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.cache.redis().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.cache.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.cache.redis().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.cache.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();
crate::failure_log::redis_failure!(
&error,
&format!("spend counter {key} 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.cache.memory().guard();
if self
.cache
.memory()
.get(&key)?
.is_none_or(|cached| cached.as_f64() < spend.as_f64())
{
self.cache.memory().set(&key, spend, None)?;
}
}
let Some(remote) = self.cache.redis().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.cache.redis().get()
else {
return Ok(self.cache.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.cache.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.cache.memory().set(&key, spend, ttl)?;
let Some(remote) = self.cache.redis().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();
self.cache.delete_keys(&keys).await
}
/// `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.cache.redis().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.cache.memory().get(&counter.key())?)
}
pub fn set_local(
&self,
counter: &SpendCounter,
spend: Number,
ttl: Option<Duration>,
) -> Result<(), Error> {
Ok(self.cache.memory().set(&counter.key(), spend, ttl)?)
}
pub fn delete_local(&self, counter: &SpendCounter) -> Result<(), Error> {
Ok(self.cache.memory().delete(&counter.key())?)
}
/// `DualCache.async_increment_cache(local_only=True)`, which passes `InMemoryCache` no TTL.
pub fn increment_local(&self, counter: &SpendCounter, amount: Number) -> Result<Number, Error> {
Ok(self
.cache
.memory()
.increment(&counter.key(), amount, None)?)
}
pub fn flush_local(&self) -> Result<(), Error> {
Ok(self.cache.memory().flush()?)
}
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, Error> {
Ok(self.cache.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");
}
}
}

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

View file

@ -2,12 +2,12 @@ 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.
/// A store value read from the plain Python data it is held as.
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.
/// A store value written as the plain Python data it is held as.
pub trait IntoValue {
fn into_value(self) -> Value;
}
@ -36,20 +36,6 @@ impl IntoValue for Number {
}
}
/// 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()),
@ -57,23 +43,6 @@ pub(crate) fn text(value: &Value, shape: &'static str) -> Result<String, DecodeE
}
}
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))
}

View file

@ -0,0 +1,741 @@
mod support;
use std::time::Duration;
use litellm_python_compat::Value;
use litellm_redis_client::commands;
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, Remote, RequestRedis,
auth_cache::{self, AuthCache, Settings},
router_cooldowns::RouterCooldowns,
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const LONG_DEADLINE: Duration = Duration::from_secs(60);
const KEY_OBJECT: &str = "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b";
const TEAM_JSON: &str = concat!(
r#"{"team_id": "team-\u00e9", "models": ["gpt-5", "claude-opus-5"], "max_budget": null, "#,
r#""spend": 0.30000000000000004, "blocked": false, "tpm_limit": 10000000000, "#,
r#""metadata": {"tags": [], "quoted": "it's \"x\""}}"#,
);
const TEAM_STR: &str = concat!(
"{'team_id': 'team-é', 'models': ['gpt-5', 'claude-opus-5'], 'max_budget': None, ",
"'spend': 0.30000000000000004, 'blocked': False, 'tpm_limit': 10000000000, ",
r#"'metadata': {'tags': [], 'quoted': 'it\'s "x"'}}"#,
);
fn dict<const N: usize>(entries: [(&str, Value); N]) -> Value {
Value::Dict(
entries
.into_iter()
.map(|(key, value)| (Value::from(key), value))
.collect(),
)
}
fn strings(items: &[&str]) -> Value {
Value::List(items.iter().copied().map(Value::from).collect())
}
fn team_record() -> Value {
dict([
("team_id", Value::from("team-é")),
("models", strings(&["gpt-5", "claude-opus-5"])),
("max_budget", Value::None),
("spend", Value::Float(0.30000000000000004)),
("blocked", Value::Bool(false)),
("tpm_limit", Value::from(10000000000)),
(
"metadata",
dict([
("tags", strings(&[])),
("quoted", Value::from("it's \"x\"")),
]),
),
])
}
/// The writes `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own
/// `json.dumps` and `str`.
fn written() -> [(Value, &'static str, &'static str); 7] {
[
(team_record(), TEAM_JSON, TEAM_STR),
(
dict([
("token", Value::from(KEY_OBJECT)),
("spend", Value::Float(1e-05)),
("rpm_limit", Value::from(5)),
("expires", Value::from("2026-11-03T00:00:00Z")),
]),
concat!(
r#"{"token": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b", "spend": 1e-05, "#,
r#""rpm_limit": 5, "expires": "2026-11-03T00:00:00Z"}"#,
),
concat!(
"{'token': '88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b', 'spend': 1e-05, ",
"'rpm_limit': 5, 'expires': '2026-11-03T00:00:00Z'}",
),
),
(
Value::from("__no_team_membership__"),
"\"__no_team_membership__\"",
"__no_team_membership__",
),
(
strings(&["mcp-a", "mcp-b"]),
r#"["mcp-a", "mcp-b"]"#,
"['mcp-a', 'mcp-b']",
),
(Value::Float(12.5), "12.5", "12.5"),
(Value::from(7), "7", "7"),
(Value::Bool(true), "true", "True"),
]
}
/// The reads `tests/unit/redis/stores/test_auth_cache.py` checks against Python's own
/// `RedisCache._get_cache_logic`.
fn read_back() -> [(&'static str, Value); 5] {
[
(TEAM_JSON, team_record()),
(TEAM_STR, team_record()),
("['mcp-a', 'mcp-b']", strings(&["mcp-a", "mcp-b"])),
(
"{'user_id': 'u1', 'teams': ['t1'], 'active': True, 'budget': None}",
dict([
("user_id", Value::from("u1")),
("teams", strings(&["t1"])),
("active", Value::Bool(true)),
("budget", Value::None),
]),
),
("\"__NO_MAPPING__\"", Value::from("__NO_MAPPING__")),
]
}
fn settings(memory_ttl: Option<u64>) -> Settings {
Settings {
memory_ttl: memory_ttl.map(Duration::from_secs),
..Settings::default()
}
}
fn store(client: Option<RedisClient>) -> AuthCache {
AuthCache::new(client.map(remote), settings(Some(60)))
}
fn team(id: &str) -> Value {
Value::Dict(vec![(Value::from("team_id"), Value::from(id))])
}
async fn redis_ttl(raw: &RedisClient, key: &str) -> Duration {
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
remaining
}
#[rstest]
#[case(KEY_OBJECT, true)]
#[case(
"88DC28D0F030C55ED4AB77ED8FAF098196CB1C05DF778539800C9F1243FE6B4B",
false
)]
#[case(
"88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4",
false
)]
#[case(
"team_id:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
false
)]
#[case("team_membership:user-1:team-1", false)]
#[case("end_user_restricted_registry", false)]
fn only_a_bare_lowercase_sha256_digest_is_a_user_key_object(
#[case] key: &str,
#[case] key_object: bool,
) {
assert_eq!(auth_cache::is_key_object(key), key_object, "{key}");
}
#[rstest]
#[tokio::test]
async fn writes_mains_json_and_str_bytes(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
for (index, (value, json, printed)) in written().into_iter().enumerate() {
let key = format!("team_id:t{index}");
store.set(&key, value.clone(), None).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&key))
.await
.unwrap()
.as_deref(),
Some(json)
);
store.set_str(&key, value, None).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&key))
.await
.unwrap()
.as_deref(),
Some(printed)
);
}
}
#[rstest]
#[tokio::test]
async fn reads_every_encoding_main_writes_and_backfills_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
for (index, (bytes, expected)) in read_back().into_iter().enumerate() {
let key = format!("user_id:u{index}");
raw.send(commands::set(&key, bytes, Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(store.get(None, &key).await.unwrap(), Some(expected.clone()));
raw.send(commands::del(&[&key])).await.unwrap();
assert_eq!(
store.get(None, &key).await.unwrap(),
Some(expected),
"{key}"
);
}
assert_eq!(store.get(None, "user_id:absent").await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn ttls_follow_dual_cache_and_the_write_back_rule(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = AuthCache::new(
Some(Remote {
client: connect(&deployment),
default_ttl: Duration::from_secs(45),
}),
Settings {
memory_ttl: Some(Duration::from_secs(20)),
..Settings::default()
},
);
store
.set(
"team_id:given",
team("a"),
Some(Duration::from_secs_f64(30.9)),
)
.await
.unwrap();
store.set("team_id:default", team("b"), None).await.unwrap();
store
.write_back(
RedisSend::Now,
vec![
("user-1".to_owned(), team("c"), None),
("user-2".to_owned(), team("d"), Some(Duration::from_secs(5))),
],
)
.unwrap()
.await
.unwrap();
assert!(within(redis_ttl(&raw, "team_id:given").await, 30));
assert!(within(redis_ttl(&raw, "team_id:default").await, 20));
assert!(within(redis_ttl(&raw, "user-1").await, 45));
assert!(within(redis_ttl(&raw, "user-2").await, 5));
assert_eq!(store.get_local("user-1").unwrap(), Some(team("c")));
}
#[tokio::test]
async fn user_key_objects_survive_churn_in_the_management_object_partition() {
let store = AuthCache::new(
None,
Settings {
memory_entries: 2,
..Settings::default()
},
);
store.set(KEY_OBJECT, team("key"), None).await.unwrap();
for index in 0..5 {
store
.set(&format!("team_id:t{index}"), team("t"), None)
.await
.unwrap();
}
assert_eq!(
store.get(None, KEY_OBJECT).await.unwrap(),
Some(team("key"))
);
assert_eq!(store.len_local().unwrap(), (2, 1));
}
#[rstest]
#[tokio::test]
async fn a_value_json_refuses_stays_in_memory_and_never_reaches_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let members = Value::Set(vec![Value::from("a")]);
store
.set("team_id:set", members.clone(), None)
.await
.unwrap();
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(
raw.send(commands::get::<String>("team_id:set"))
.await
.unwrap(),
None
);
}
// RedisCache._pipeline_helper runs json.dumps inside the pipeline's try, so one value it refuses
// aborts every SET of that pipeline, while DualCache has already written them all to memory
#[tokio::test]
async fn a_value_json_refuses_keeps_its_whole_pipeline_out_of_redis() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let members = Value::Set(vec![Value::from("a")]);
store
.set_many(
vec![
("team_id:a".to_owned(), team("a")),
("team_id:set".to_owned(), members.clone()),
],
None,
)
.await
.unwrap();
assert_eq!(store.get_local("team_id:a").unwrap(), Some(team("a")));
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(
raw.send(commands::get::<String>("team_id:a"))
.await
.unwrap(),
None
);
}
// auth_object_prefetch writes every entry to memory before
// RedisCache.async_set_cache_pipeline_with_ttls encodes them, and that encodes all before sending
#[tokio::test]
async fn a_write_back_value_json_refuses_fails_after_memory_and_before_redis() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let members = Value::Set(vec![Value::from("a")]);
let written = store.write_back(
RedisSend::Now,
vec![
("team_id:set".to_owned(), members.clone(), None),
("team_id:a".to_owned(), team("a"), None),
],
);
assert!(written.is_err());
assert_eq!(store.get_local("team_id:set").unwrap(), Some(members));
assert_eq!(store.get_local("team_id:a").unwrap(), Some(team("a")));
assert_eq!(
raw.send(commands::get::<String>("team_id:a"))
.await
.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn pre_call_writes_and_reads_ride_one_pipeline_with_the_cooldown_read(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(client.clone())),
COOLDOWN_POLICY,
));
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
raw.send(commands::set(
KEY_OBJECT,
"{\"token\": \"t\"}",
Duration::from_secs(60),
))
.await
.unwrap();
raw.send(commands::set(
"team_alias:alpha",
"\"stale\"",
Duration::from_secs(60),
))
.await
.unwrap();
let reads = store
.get_many(
RedisSend::PreCall(&request),
&[KEY_OBJECT.to_owned(), "end_user_id:absent".to_owned()],
)
.unwrap();
let active = cooldowns
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let written = store
.declare_set(RedisSend::PreCall(&request), "team_id:t1", team("t1"), None)
.unwrap();
let deleted = store
.declare_delete(RedisSend::PreCall(&request), "team_alias:alpha")
.unwrap();
assert!(trips.seen().is_empty());
assert_eq!(store.get_local("team_id:t1").unwrap(), Some(team("t1")));
request.flush_pre().await;
let mut seen = trips.seen();
seen.sort_by_key(|operation| format!("{operation:?}"));
let expected = match topology {
Topology::Standalone => vec![Operation::Pipeline { ops: 4 }],
Topology::Cluster => vec![
Operation::Command("DEL"),
Operation::Command("MGET"),
Operation::Command("MGET"),
Operation::Command("SET"),
],
};
assert_eq!(seen, expected);
assert_eq!(
reads.await.unwrap(),
vec![
Some(Value::Dict(vec![(Value::from("token"), Value::from("t"))])),
None
]
);
assert_eq!(active.await.unwrap(), vec![]);
written.await.unwrap();
deleted.await.unwrap();
assert_eq!(
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.send(commands::get::<String>("team_alias:alpha"))
.await
.unwrap(),
None
);
assert_eq!(trips.seen().len(), expected.len());
}
#[rstest]
#[tokio::test]
async fn a_write_back_on_the_request_waits_for_the_next_flush(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
store
.write_back(
RedisSend::PreCall(&request),
vec![(
"user-1".to_owned(),
team("u"),
Some(Duration::from_secs(42)),
)],
)
.unwrap()
.await
.unwrap();
assert!(trips.seen().is_empty());
assert_eq!(store.get_local("user-1").unwrap(), Some(team("u")));
request.flush_pre().await;
let write = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("SET"),
};
assert_eq!(trips.seen(), [write]);
assert!(within(redis_ttl(&raw, "user-1").await, 42));
}
#[rstest]
#[tokio::test]
async fn an_eviction_broadcast_drops_only_this_pods_copy_and_the_next_read_refills_from_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let this_pod = store(Some(connect(&deployment)));
let other_pod = store(Some(connect(&deployment)));
this_pod.set("team_id:t1", team("old"), None).await.unwrap();
other_pod
.set("team_id:t1", team("new"), None)
.await
.unwrap();
assert_eq!(
this_pod.get(None, "team_id:t1").await.unwrap(),
Some(team("old"))
);
this_pod.delete_local("team_id:t1").unwrap();
assert_eq!(
this_pod.get(None, "team_id:t1").await.unwrap(),
Some(team("new"))
);
assert!(
raw.send(commands::get::<String>("team_id:t1"))
.await
.unwrap()
.is_some()
);
}
// RedisBatch.read_as_missing: once the request's MGET found a key absent, or the request deleted it,
// a GET of it is answered without a round trip until the request sets it again
#[rstest]
#[tokio::test]
async fn a_key_the_request_read_as_missing_is_not_read_again(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let other_pod_writes = |key: &'static str| {
raw.send(commands::set(
key,
TEAM_JSON.to_owned(),
Duration::from_secs(60),
))
};
let read = store
.get_many(RedisSend::PreCall(&request), &["team_id:read".to_owned()])
.unwrap();
assert_eq!(read.await.unwrap(), vec![None]);
store
.declare_delete(RedisSend::PreCall(&request), "team_id:deleted")
.unwrap()
.await
.unwrap();
other_pod_writes("team_id:read").await.unwrap();
other_pod_writes("team_id:deleted").await.unwrap();
assert_eq!(
store.get(Some(&request), "team_id:read").await.unwrap(),
None
);
assert_eq!(
store.get(Some(&request), "team_id:deleted").await.unwrap(),
None
);
assert_eq!(
store.get(None, "team_id:read").await.unwrap(),
Some(team_record())
);
store
.declare_set(
RedisSend::PreCall(&request),
"team_id:deleted",
team("t"),
None,
)
.unwrap()
.await
.unwrap();
store.delete_local("team_id:deleted").unwrap();
assert_eq!(
store.get(Some(&request), "team_id:deleted").await.unwrap(),
Some(team("t"))
);
}
#[rstest]
#[tokio::test]
async fn delete_drops_memory_and_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let keys = [KEY_OBJECT.to_owned(), "team_id:t1".to_owned()];
store
.set_many(
keys.iter().map(|key| (key.clone(), team("t"))).collect(),
None,
)
.await
.unwrap();
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.send(commands::get::<String>(key)).await.unwrap(), None);
}
}
#[rstest]
#[tokio::test]
async fn increment_if_present_only_raises_an_entry_redis_holds(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
assert_eq!(
store
.increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5))
.await
.unwrap(),
None
);
assert_eq!(
raw.send(commands::get::<String>("litellm-proxy-budget:spend"))
.await
.unwrap(),
None
);
store
.set("litellm-proxy-budget:spend", Value::Float(1.5), None)
.await
.unwrap();
assert_eq!(
store
.increment_if_present("litellm-proxy-budget:spend", Number::Float(2.5))
.await
.unwrap(),
Some(4.0)
);
}
#[tokio::test]
async fn without_redis_every_operation_is_memory_only() {
let store = store(None);
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
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
.get_many(RedisSend::PreCall(&request), &["team_id:t1".to_owned()])
.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))
.await
.unwrap(),
None
);
assert_eq!(
store.get(None, "team_id:t1").await.unwrap(),
Some(team("t"))
);
assert_eq!(store.get_local("team_id:t2").unwrap(), None);
}
// UserApiKeyCache writes the key-object partition, then the rest, and deletes both partitions
// concurrently, while async_increment_if_present sends its script as a plain EVAL
#[rstest]
#[tokio::test]
async fn the_auth_cache_sends_the_commands_user_api_key_cache_sends(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
deployment
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
.await
.unwrap();
store
.set_many(
vec![
(KEY_OBJECT.to_owned(), team("t0")),
("team_id:t1".to_owned(), team("t1")),
],
None,
)
.await
.unwrap();
let written = trips.seen();
store
.delete(&[KEY_OBJECT.to_owned(), "team_id:t1".to_owned()])
.await
.unwrap();
store
.increment_if_present("litellm-proxy-budget:spend", Number::Int(5))
.await
.unwrap();
assert_eq!(written, vec![Operation::Pipeline { ops: 1 }; 2]);
assert_eq!(
trips.seen()[written.len()..written.len() + 2],
[Operation::Command("DEL"), Operation::Command("DEL")]
);
assert_eq!(support::calls(&deployment, "eval").await, 1);
assert_eq!(support::calls(&deployment, "evalsha").await, 0);
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,512 @@
mod support;
use std::time::Duration;
use litellm_redis_client::commands;
use litellm_redis_client::{Operation, Ttl};
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::{Cooldown, RouterCooldowns},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{ALLOWED_FAILS, COOLDOWN_POLICY as POLICY, 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))
}
#[rstest]
#[tokio::test]
async fn put_writes_the_bytes_mains_set_path_writes_with_whole_second_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
#[values("json", "str_dict")] format: &str,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
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);
let (written_by_put, expected) = match format {
"json" => (
store.put(case.model_id, &case.cooldown, ttl).await,
case.json,
),
_ => (
store.put_str(case.model_id, &case.cooldown, ttl).await,
case.str_dict,
),
};
written_by_put.unwrap();
let key = case.key;
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(expected));
assert!(within(remaining, case.redis_ttl_seconds), "{remaining:?}");
}
}
#[rstest]
#[tokio::test]
async fn reads_cooldowns_python_wrote_as_str_dict_or_json(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
#[values("str_dict", "json")] format: &str,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
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.send(commands::set(case.key, bytes, Duration::from_secs(60)))
.await
.unwrap();
assert_eq!(
read_active(&store, &[case.model_id.to_owned()]).await,
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
}
#[rstest]
#[tokio::test]
async fn an_ended_cooldown_is_not_active_but_still_counts_toward_the_minimum(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
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 = 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.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
}
assert_eq!(
store.minimum(&models).await.unwrap(),
Some(Number::Float(0.5))
);
assert_eq!(read_active(&store, &models).await, vec![]);
}
#[rstest]
#[tokio::test]
async fn allowed_fails_count_across_async_and_sync_routers_and_keep_the_first_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
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.send(commands::del(&[key])).await.unwrap();
let suffix = case.suffix;
let counts = [
first
.increment_allowed_fails("dep-a", suffix, Duration::from_secs_f64(30.9))
.await
.unwrap(),
second
.increment_allowed_fails_int("dep-a", suffix, Duration::from_secs(120))
.await
.unwrap(),
];
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(counts, [1, 2], "{key}");
assert!(within(remaining, 30), "{remaining:?}");
}
}
#[rstest]
#[case::async_increment_cache_sends_incrbyfloat("INCRBYFLOAT")]
#[case::increment_cache_sends_incrby("INCRBY")]
#[tokio::test]
async fn allowed_fails_send_the_increment_mains_path_sends(
#[case] command: &'static str,
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let (client, trips) = support::connect_observed(&deployment);
let store = std::sync::Arc::new(RouterCooldowns::new(Some(remote(client)), POLICY));
let ttl = Duration::from_secs(60);
let count = match command {
"INCRBYFLOAT" => store.increment_allowed_fails("dep-a", None, ttl).await,
_ => store.increment_allowed_fails_int("dep-a", None, ttl).await,
};
assert_eq!(count.unwrap(), 1);
assert_eq!(trips.seen().first(), Some(&Operation::Command(command)));
}
#[rstest]
#[tokio::test]
async fn python_written_fail_counts_continue_in_rust(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
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();
assert_eq!(
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
5
);
}
#[rstest]
#[tokio::test]
async fn a_router_without_redis_keeps_its_cooldowns_and_counts_in_memory() {
let store = std::sync::Arc::new(RouterCooldowns::with_clock(None, POLICY, fixture_clock()));
let [case, _] = cooldowns();
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
let fails = [
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
];
store.flush_local_counters().unwrap();
assert_eq!(
read_active(&store, &[case.model_id.to_owned()]).await,
vec![(case.model_id.to_owned(), case.cooldown)]
);
assert_eq!(fails, [1, 2]);
assert_eq!(
store
.increment_allowed_fails("dep-a", None, Duration::from_secs(60))
.await
.unwrap(),
1
);
}
#[rstest]
#[tokio::test]
async fn a_redis_outage_serves_memory_and_counts_fails_locally() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
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))
.await
.unwrap();
deployment.stop().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
.unwrap();
assert_eq!(active, vec![(case.model_id.to_owned(), case.cooldown)]);
assert_eq!(fails, 1);
}
struct StepClock(std::sync::Mutex<Duration>);
impl StepClock {
fn at(seconds: f64) -> std::sync::Arc<Self> {
std::sync::Arc::new(Self(std::sync::Mutex::new(Duration::from_secs_f64(
seconds,
))))
}
fn advance(&self, by: Duration) {
*self.0.lock().unwrap() += by;
}
fn clock(self: &std::sync::Arc<Self>) -> litellm_redis_store::Clock {
let this = self.clone();
std::sync::Arc::new(move || *this.0.lock().unwrap())
}
}
#[rstest]
#[tokio::test]
async fn a_miss_is_read_from_redis_again_only_after_the_one_second_refresh(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let time = StepClock::at(1_730_000_000.75);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(connect(&deployment))),
POLICY,
time.clock(),
));
let [case, _] = cooldowns();
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 = read_active(&store, &[case.model_id.to_owned()]).await;
time.advance(Duration::from_secs(1));
let refreshed = read_active(&store, &[case.model_id.to_owned()]).await;
assert_eq!(before, vec![]);
assert_eq!(throttled, vec![]);
assert_eq!(
refreshed,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}
#[rstest]
#[tokio::test]
async fn a_breaker_refusal_does_not_hold_back_the_next_read() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(200),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let time = StepClock::at(1_730_000_000.75);
let guarded = support::connect_guarded(&deployment, breaker);
let store = std::sync::Arc::new(RouterCooldowns::with_clock(
Some(remote(guarded.clone())),
POLICY,
time.clock(),
));
let [case, _] = cooldowns();
deployment.stop().await.unwrap();
read_active(&store, &[case.model_id.to_owned()]).await;
time.advance(Duration::from_secs(1));
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.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
let recovered = read_active(&store, &[case.model_id.to_owned()]).await;
assert_eq!(refused, vec![]);
assert_eq!(
recovered,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}
#[rstest]
#[tokio::test]
async fn a_declared_active_read_rides_the_request_batch_and_answers_what_active_does(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
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.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.active(RedisSend::PreCall(&request), &models).unwrap();
assert!(trips.seen().is_empty());
request.flush_pre().await;
let active = declared.await.unwrap();
let one_mget = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("MGET"),
};
assert_eq!(trips.seen(), [one_mget]);
assert_eq!(
active,
cooldowns()
.map(|case| (case.model_id.to_owned(), case.cooldown))
.to_vec()
);
assert_eq!(read_active(&store, &models).await, active);
assert_eq!(trips.seen().len(), 1, "the declared read backfilled memory");
}
#[rstest]
#[tokio::test]
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 = 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
.active(RedisSend::PreCall(&request), &[case.model_id.to_owned()])
.unwrap();
request.flush_pre().await;
store
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
.await
.unwrap();
assert_eq!(
declared.await.unwrap(),
vec![(case.model_id.to_owned(), case.cooldown)]
);
}
#[rstest]
#[tokio::test]
async fn a_breaker_refusal_of_a_declared_read_does_not_hold_back_the_next_read() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_millis(200),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let guarded = support::connect_guarded(&deployment, breaker);
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
.active(RedisSend::PreCall(&opened), &["dep-other".to_owned()])
.unwrap();
opened.flush_pre().await;
failed.await.unwrap();
let refused_request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let refused = store
.active(RedisSend::PreCall(&refused_request), &models)
.unwrap();
refused_request.flush_pre().await;
let refused = refused.await.unwrap();
deployment.resume().await.unwrap();
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
raw.send(commands::set(
case.key,
case.str_dict,
Duration::from_secs(60),
))
.await
.unwrap();
assert_eq!(refused, vec![]);
assert_eq!(
read_active(&store, &models).await,
vec![(case.model_id.to_owned(), case.cooldown.clone())]
);
}

View file

@ -0,0 +1,329 @@
mod support;
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_redis_client::commands;
use litellm_redis_client::{Amount, Expiry, Operation, Ttl};
use litellm_redis_store::{
Number, TierPolicy,
router_usage::{COUNTER_TTL, Counter, Kind, RouterUsage},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{USAGE_POLICY as POLICY, connect, remote, usage, within};
fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] {
[Kind::Tpm, Kind::Rpm].map(|kind| Counter {
deployment_id: "dep".to_owned(),
model: model.map(str::to_owned),
minute: "12-34".to_owned(),
kind,
})
}
#[rstest]
#[case::tpm(Kind::Tpm, "tpm")]
#[case::rpm(Kind::Rpm, "rpm")]
fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) {
assert_eq!(<&str>::from(kind), segment);
assert_eq!(segment.parse::<Kind>().unwrap(), kind);
}
#[rstest]
#[case::unknown("tokens")]
#[case::wrong_case("TPM")]
fn other_text_is_not_a_counter_kind(#[case] segment: &str) {
assert!(segment.parse::<Kind>().is_err());
}
// DualCache._backfill_kwargs: a router's batch read keeps a Redis value in memory for
// litellm.default_in_memory_ttl when one is set, else for InMemoryCache's own 600 s
#[rstest]
#[case::configured(Some(Duration::from_secs(5)), None)]
#[case::unset(None, Some(Number::Int(3)))]
#[tokio::test]
async fn a_batch_read_backfills_memory_for_default_in_memory_ttl(
#[case] default_in_memory_ttl: Option<Duration>,
#[case] after_six_seconds: Option<Number>,
) {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let seconds = Arc::new(AtomicU64::new(1_000));
let now = {
let seconds = seconds.clone();
Arc::new(move || Duration::from_secs(seconds.load(Ordering::SeqCst)))
};
let policy = TierPolicy {
refresh_after: Duration::from_secs(3600),
default_in_memory_ttl,
..POLICY
};
let store = RouterUsage::with_clock(Some(remote(connect(&deployment))), policy, now);
let [case, _] = usage();
raw.send(commands::incr_by_float(
case.key,
Amount::Float(3.0),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
store
.get_many(std::slice::from_ref(&case.counter))
.await
.unwrap();
raw.send(commands::del(&[case.key])).await.unwrap();
seconds.fetch_add(6, Ordering::SeqCst);
let read = store.get_many(&[case.counter]).await.unwrap();
assert_eq!(read, vec![after_six_seconds]);
}
#[rstest]
#[tokio::test]
async fn reads_counters_python_incremented_under_mains_keys(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let cases = usage();
for (case, delta) in cases.iter().zip([7.75, 3.0]) {
raw.send(commands::incr_by_float(
case.key,
Amount::Float(delta),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
}
assert_eq!(
store
.get_many(&cases.map(|case| case.counter))
.await
.unwrap(),
vec![Some(Number::Float(7.75)), Some(Number::Int(3))]
);
}
#[rstest]
#[tokio::test]
async fn deployment_increments_write_incrbyfloat_totals_with_a_refreshed_minute_window(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
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] =
["tpm", "rpm"].map(|kind| case.key.replace(":tpm:", &format!(":{kind}:")));
writer
.increment_deployment("dep", model, "12-34", Number::Float(5.5), 1)
.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.send(commands::get(&tpm_key)).await.unwrap(),
raw.send(commands::get(&rpm_key)).await.unwrap(),
];
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())]);
assert!(within(remaining, 60), "{remaining:?}");
assert_eq!(
reader.get_many(&tpm_and_rpm(model)).await.unwrap(),
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
);
assert_eq!(
writer.get_many(&tpm_and_rpm(model)).await.unwrap(),
vec![Some(Number::Int(7)), Some(Number::Int(2))]
);
}
#[rstest]
#[tokio::test]
async fn a_request_increment_keeps_the_window_python_started(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY);
let [_, case] = usage();
let key = case.key;
raw.send(commands::incr_by_float(
key,
Amount::Float(1.0),
Some(Expiry::KeepExisting(Duration::from_secs(20))),
))
.await
.unwrap();
let total = store
.increment(&case.counter, Number::Int(1), COUNTER_TTL)
.await
.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else {
panic!("{key} has no TTL");
};
assert_eq!(total, Number::Float(2.0));
assert!(within(remaining, 20), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn without_redis_counters_accumulate_in_memory_until_flushed() {
let store = RouterUsage::new(None, POLICY);
let counters = tpm_and_rpm(Some("model"));
store
.increment_deployment("dep", Some("model"), "12-34", Number::Int(2), 1)
.await
.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
.await
.unwrap();
let counted = store.get_many(&counters).await.unwrap();
store.flush_local_counters().unwrap();
assert_eq!(
counted,
vec![Some(Number::Float(7.5)), Some(Number::Int(2))]
);
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(remote(connect(&deployment))), POLICY);
deployment.stop().await.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
.await
.unwrap();
assert_eq!(
store.get_many(&tpm_and_rpm(Some("model"))).await.unwrap(),
vec![Some(Number::Int(5)), Some(Number::Int(1))]
);
}
// main's _increment_deployment sends async_increment_pipeline on its own round trip, never on the
// request's post-call pipeline
#[rstest]
#[tokio::test]
async fn a_deployment_increment_is_its_own_pipeline_and_mirrors_the_redis_totals(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
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.send(commands::incr_by_float(
&tpm_key,
Amount::Float(10.0),
Some(Expiry::Refresh(COUNTER_TTL)),
))
.await
.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.5), 1)
.await
.unwrap();
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
assert_eq!(
store.get_many(&counters).await.unwrap(),
vec![Some(Number::Int(15)), Some(Number::Int(1))]
);
let written: [Option<String>; 2] = [
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.send(commands::ttl(&rpm_key)).await.unwrap() else {
panic!("{rpm_key} has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
}
#[rstest]
#[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(remote(connect(&deployment))), POLICY);
let counters = tpm_and_rpm(Some("model"));
deployment.stop().await.unwrap();
store
.increment_deployment("dep", Some("model"), "12-34", Number::Float(5.75), 1)
.await
.unwrap();
assert_eq!(
store.get_many(&counters).await.unwrap(),
[Some(Number::Int(5)), Some(Number::Int(1))]
);
}

View file

@ -0,0 +1,134 @@
mod support;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_redis_client::{Ttl, commands};
use litellm_redis_store::{Clock, sensitive_route_pins::SensitiveRoutePins};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, remote, within};
const PIN_KEY: &str = "{sensitive_route:tenant-a:session-1}:model";
const HOUR: Duration = Duration::from_secs(3600);
struct StepClock(Mutex<Duration>);
impl StepClock {
fn start() -> Arc<Self> {
Arc::new(Self(Mutex::new(Duration::from_secs(1_000))))
}
fn advance(&self, by: Duration) {
*self.0.lock().unwrap() += by;
}
fn clock(self: &Arc<Self>) -> Clock {
let this = self.clone();
Arc::new(move || *this.0.lock().unwrap())
}
}
#[rstest]
#[tokio::test]
async fn a_pin_one_pod_writes_is_read_by_another_under_mains_key_and_encoding(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let writer = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
let reader = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
writer
.set("tenant-a", "session-1", "on-prem-model", HOUR)
.await
.unwrap();
let written: Option<String> = raw.send(commands::get(PIN_KEY)).await.unwrap();
let Ttl::Expires(remaining) = raw.send(commands::ttl(PIN_KEY)).await.unwrap() else {
panic!("{PIN_KEY} has no TTL");
};
assert_eq!(written.as_deref(), Some("\"on-prem-model\""));
assert!(within(remaining, 3600), "{remaining:?}");
assert_eq!(
reader.get("tenant-a", "session-1", HOUR).await.unwrap(),
Some("on-prem-model".to_owned())
);
assert_eq!(
reader.get("tenant-b", "session-1", HOUR).await.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn a_pin_read_from_redis_stays_in_memory_for_its_remaining_ttl_only() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let time = StepClock::start();
let reader = SensitiveRoutePins::with_clock(Some(remote(connect(&deployment))), time.clock());
raw.send(commands::set(
PIN_KEY,
"\"on-prem-model\"",
Duration::from_secs(30),
))
.await
.unwrap();
let read = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
deployment.stop().await.unwrap();
time.advance(Duration::from_secs(20));
let before_expiry = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
time.advance(Duration::from_secs(15));
let after_expiry = reader.get("tenant-a", "session-1", HOUR).await.unwrap();
assert_eq!(read, Some("on-prem-model".to_owned()));
assert_eq!(before_expiry, read);
assert_eq!(after_expiry, None);
}
#[rstest]
#[tokio::test]
async fn a_redis_outage_keeps_pinning_sessions_in_memory() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = SensitiveRoutePins::new(Some(remote(connect(&deployment))));
store
.set("tenant-a", "session-1", "on-prem-model", HOUR)
.await
.unwrap();
deployment.stop().await.unwrap();
let pinned_before = store.get("tenant-a", "session-1", HOUR).await.unwrap();
store
.set("tenant-a", "session-2", "other-model", HOUR)
.await
.unwrap();
let pinned_during = store.get("tenant-a", "session-2", HOUR).await.unwrap();
assert_eq!(pinned_before, Some("on-prem-model".to_owned()));
assert_eq!(pinned_during, Some("other-model".to_owned()));
}
#[tokio::test]
async fn without_redis_pins_live_in_memory_until_they_expire() {
let time = StepClock::start();
let store = SensitiveRoutePins::with_clock(None, time.clock());
store
.set(
"tenant-a",
"session-1",
"on-prem-model",
Duration::from_secs(10),
)
.await
.unwrap();
let pinned = store.get("tenant-a", "session-1", HOUR).await.unwrap();
time.advance(Duration::from_secs(11));
let expired = store.get("tenant-a", "session-1", HOUR).await.unwrap();
assert_eq!(pinned, Some("on-prem-model".to_owned()));
assert_eq!(expired, None);
}

View file

@ -0,0 +1,648 @@
mod support;
use std::{sync::Arc, time::Duration};
use litellm_redis_client::commands;
use litellm_redis_client::{Operation, RedisClient, Ttl};
use litellm_redis_store::Send as RedisSend;
use litellm_redis_store::{
Number, RequestRedis,
router_cooldowns::RouterCooldowns,
spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{COOLDOWN_POLICY, connect, connect_observed, remote, within};
use tokio::runtime::Handle;
const TTL: Duration = Duration::from_secs(60);
const LONG_DEADLINE: Duration = Duration::from_secs(60);
const READS: [(&str, Number); 4] = [
("7.25", Number::Float(7.25)),
("3", Number::Int(3)),
("0.0", Number::Float(0.0)),
("1e-05", Number::Float(1e-05)),
];
const SETS: [(Number, &str); 4] = [
(Number::Float(0.0), "0.0"),
(Number::Int(12), "12"),
(Number::Float(1e-05), "1e-05"),
(Number::Float(0.30000000000000004), "0.30000000000000004"),
];
const SET_MAXES: [(Number, &str); 3] = [
(Number::Float(2.5), "2.5"),
(Number::Int(100), "100"),
(Number::Float(1e-05), "1e-05"),
];
fn team(id: &str) -> SpendCounter {
SpendCounter::Team {
team_id: id.to_owned(),
window: None,
}
}
fn store(client: Option<RedisClient>) -> Arc<SpendCounters> {
Arc::new(SpendCounters::new(client.map(remote)))
}
/// Mirrored by `tests/unit/redis/stores/test_spend_counters.py`, which checks the same keys and
/// bytes against Python's own key builders and `RedisCache` codec.
#[rstest]
#[case::redis(Reading::Redis(Some(Number::Int(3))), "redis", Some(Number::Int(3)))]
#[case::redis_miss(Reading::Redis(None), "redis", None)]
#[case::memory(
Reading::Memory(Some(Number::Float(2.5))),
"memory",
Some(Number::Float(2.5))
)]
fn a_reading_names_where_it_was_answered(
#[case] reading: Reading,
#[case] source: &str,
#[case] value: Option<Number>,
) {
assert_eq!(<&str>::from(reading), source);
assert_eq!(reading.value(), value);
}
#[rstest]
#[case(
SpendCounter::Key {
token: "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b".to_owned(),
window: None,
},
"spend:key:88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b"
)]
#[case(
SpendCounter::Key {
token: "88dc28d0".to_owned(),
window: Some("30d".to_owned()),
},
"spend:key:88dc28d0:window:30d"
)]
#[case(
SpendCounter::Team {
team_id: "team-a".to_owned(),
window: None,
},
"spend:team:team-a"
)]
#[case(
SpendCounter::Team {
team_id: "team-a".to_owned(),
window: Some("1h".to_owned()),
},
"spend:team:team-a:window:1h"
)]
#[case(
SpendCounter::TeamMember {
user_id: "user:with:colons".to_owned(),
team_id: "team-a".to_owned(),
},
"spend:team_member:user:with:colons:team-a"
)]
#[case(SpendCounter::User("user@example.com".to_owned()), "spend:user:user@example.com")]
#[case(SpendCounter::EndUser("customer é 1".to_owned()), "spend:end_user:customer é 1")]
#[case(SpendCounter::Org("org-1".to_owned()), "spend:org:org-1")]
#[case(SpendCounter::Tag("prod:eu".to_owned()), "spend:tag:prod:eu")]
#[case(SpendCounter::Project("proj-1".to_owned()), "spend:project:proj-1")]
#[case(
SpendCounter::ModelAccessGroup("gpt-group".to_owned()),
"spend:model_access_group:gpt-group"
)]
#[case(SpendCounter::ShadowEval("job-7".to_owned()), "spend:shadow_eval:job-7")]
fn every_counter_formats_and_parses_back_to_mains_key(
#[case] counter: SpendCounter,
#[case] key: &str,
) {
assert_eq!(counter.key(), key);
assert_eq!(SpendCounter::parse(key).unwrap(), counter, "{key}");
}
#[rstest]
#[case("spend_db_floor:spend:key:abc")]
#[case("virtual_key_spend:abc:gpt:1d")]
#[case("spend:unknown:abc")]
#[case("spend:team_member:no-team")]
#[case("spend:")]
fn keys_outside_the_spend_counter_keyspace_are_refused(#[case] key: &str) {
assert!(SpendCounter::parse(key).is_err(), "{key}");
}
#[rstest]
#[tokio::test]
async fn reads_values_python_wrote_and_a_clean_miss_is_a_redis_answer(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counters = (0..READS.len())
.map(|index| team(&format!("t{index}")))
.collect::<Vec<_>>();
for ((bytes, _), counter) in READS.iter().zip(&counters) {
raw.send(commands::set(&counter.key(), *bytes, TTL))
.await
.unwrap();
}
let expected = READS
.iter()
.map(|(_, spend)| Some(*spend))
.collect::<Vec<_>>();
for (counter, value) in counters.iter().zip(&expected) {
assert_eq!(store.read(counter).await.unwrap(), Reading::Redis(*value));
}
assert_eq!(
store
.read_many(RedisSend::Now, &counters)
.unwrap()
.await
.unwrap(),
Some(expected)
);
assert_eq!(
store.read(&team("absent")).await.unwrap(),
Reading::Redis(None)
);
}
#[rstest]
#[tokio::test]
async fn resets_and_repairs_write_mains_bytes_with_the_redis_ttl(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counter = team("t");
for (spend, bytes) in SETS {
store.set(&counter, spend, None).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some(bytes.to_owned())
);
}
for (spend, bytes) in SET_MAXES {
raw.send(commands::del(&[&counter.key()])).await.unwrap();
let repaired = store.set_max(&counter, spend).await.unwrap();
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some(bytes.to_owned())
);
assert_eq!(repaired, Some(Number::Float(spend.as_f64())));
}
store
.set(&counter, Number::Int(5), Some(Duration::from_secs(30)))
.await
.unwrap();
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:?}");
}
#[rstest]
#[tokio::test]
async fn set_max_only_ever_raises_a_counter(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counter = SpendCounter::Key {
token: "hashed".to_owned(),
window: None,
};
store.increment(&counter, Number::Float(9.5)).await.unwrap();
let lower = store.set_max(&counter, Number::Float(4.0)).await.unwrap();
let higher = store.set_max(&counter, Number::Float(12.0)).await.unwrap();
assert_eq!(lower, Some(Number::Float(9.5)));
assert_eq!(higher, Some(Number::Float(12.0)));
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("12.0".to_owned())
);
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:?}");
assert_eq!(
store.get_local(&counter).unwrap(),
Some(Number::Float(12.0))
);
}
#[rstest]
#[tokio::test]
async fn increments_refresh_the_ttl_and_mirror_the_shared_total_into_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let other_pod = store(Some(connect(&deployment)));
let store = store(Some(client));
let [first, second] = [team("a"), team("b")];
other_pod
.increment(&first, Number::Float(1.25))
.await
.unwrap();
raw.send(commands::expire(&first.key(), Duration::from_secs(5)))
.await
.unwrap();
let totals = store
.increment_many(&[
(first.clone(), Number::Float(0.5)),
(second.clone(), Number::Int(2)),
])
.await
.unwrap();
assert_eq!(totals, [Number::Float(1.75), Number::Float(2.0)]);
assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]);
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.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)));
assert_eq!(store.get_local(&second).unwrap(), Some(Number::Float(2.0)));
}
#[rstest]
#[tokio::test]
async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_in_redis_and_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
let declared = std::thread::spawn({
let (store, request, counter) = (Arc::clone(&store), request.clone(), counter.clone());
move || store.increment_many_post_call(&request, &[(counter, Number::Float(3.5))])
})
.join()
.expect("a Python thread has no Tokio context, and declaring there must not panic");
request.flush_post().await;
assert!(declared);
let mirrored =
litellm_redis_testing::eventually(async || store.get_local(&counter).unwrap()).await;
assert_eq!(mirrored, Number::Float(3.5));
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("3.5".to_owned())
);
}
#[rstest]
#[tokio::test]
async fn post_call_increments_wait_for_the_post_batch_then_land_in_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let store = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(3.5))]));
request.flush_pre().await;
assert!(trips.seen().is_empty());
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
None
);
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(100)).await;
let increment = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Command("INCRBYFLOAT"),
};
assert_eq!(trips.seen(), [increment]);
assert_eq!(
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)));
}
#[rstest]
#[tokio::test]
async fn the_budget_read_rides_the_same_pre_call_pipeline_as_the_cooldown_read(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = connect_observed(&deployment);
let cooldowns = std::sync::Arc::new(RouterCooldowns::new(
Some(remote(client.clone())),
COOLDOWN_POLICY,
));
let spend = store(Some(client));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counters = [
SpendCounter::Key {
token: "hashed".to_owned(),
window: None,
},
team("t"),
];
raw.send(commands::set(&counters[0].key(), "4.5", TTL))
.await
.unwrap();
let active = cooldowns
.active(RedisSend::PreCall(&request), &["dep-a".to_owned()])
.unwrap();
let budget = spend
.read_many(RedisSend::PreCall(&request), &counters)
.unwrap();
request.flush_pre().await;
let expected = match topology {
Topology::Standalone => vec![Operation::Pipeline { ops: 2 }],
Topology::Cluster => vec![Operation::Command("MGET"); 2],
};
assert_eq!(trips.seen(), expected);
assert_eq!(active.await.unwrap(), vec![]);
assert_eq!(
budget.await.unwrap(),
Some(vec![Some(Number::Float(4.5)), None])
);
assert_eq!(trips.seen().len(), expected.len());
}
#[rstest]
#[tokio::test]
async fn a_seed_lets_the_first_pod_win_and_others_answer_with_its_value(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = store(Some(connect(&deployment)));
let second = store(Some(connect(&deployment)));
let counter = SpendCounter::Org("org".to_owned());
let won = first
.seed(&counter, Number::Float(10.0), LostSeed::Answer)
.await
.unwrap();
let lost = second
.seed(&counter, Number::Float(99.0), LostSeed::Increment)
.await
.unwrap();
assert_eq!([won, lost], [Number::Float(10.0), Number::Float(10.0)]);
assert_eq!(
raw.send(commands::get::<String>(&counter.key()))
.await
.unwrap(),
Some("10.0".to_owned())
);
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:?}");
assert_eq!(
second.get_local(&counter).unwrap(),
Some(Number::Float(10.0))
);
}
#[rstest]
#[tokio::test]
async fn delete_drops_memory_and_redis_and_refresh_ttl_reports_whether_redis_held_it(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let store = store(Some(connect(&deployment)));
let counters = [team("a"), team("b")];
for counter in &counters {
store.increment(counter, Number::Int(1)).await.unwrap();
}
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.send(commands::ttl(&counters[0].key())).await.unwrap() else {
panic!("a refreshed counter has no TTL");
};
assert!(within(remaining, 60), "{remaining:?}");
store.delete(&counters).await.unwrap();
for counter in &counters {
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());
}
// main's read_spend_counter_cache_value: RedisCache.async_get_cache swallows a failed GET and
// returns None, an authoritative miss. Only the breaker guard's refusal raises into the fallback
// that reads this pod's memory
#[rstest]
#[tokio::test]
async fn a_failed_read_serves_the_local_counter_and_a_failed_increment_drops_the_counter() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = store(Some(connect(&deployment)));
let [kept, failed] = [team("kept"), team("failed")];
store.increment(&kept, Number::Float(2.0)).await.unwrap();
store.increment(&failed, Number::Float(3.0)).await.unwrap();
deployment.stop().await.unwrap();
let read = store.read(&kept).await.unwrap();
let increment = store.increment(&failed, Number::Float(1.0)).await;
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(RedisSend::Now, &[kept])
.unwrap()
.await
.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn a_read_with_an_open_breaker_serves_the_local_counter() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let breaker = litellm_circuit_breaker::Breaker::new(
litellm_circuit_breaker::Config {
failure_threshold: 1,
recovery_timeout: Duration::from_secs(60),
..litellm_circuit_breaker::Config::default()
},
Box::new(|_| {}),
);
let store = store(Some(support::connect_guarded(&deployment, breaker)));
let counter = team("t");
store.increment(&counter, Number::Float(2.0)).await.unwrap();
deployment.stop().await.unwrap();
let failed = store.read(&counter).await.unwrap();
let refused = store.read(&counter).await.unwrap();
assert_eq!(failed, Reading::Memory(Some(Number::Float(2.0))));
assert_eq!(refused, Reading::Memory(Some(Number::Float(2.0))));
}
// increment_in_memory calls async_increment_cache(local_only=True), which hands InMemoryCache no
// TTL, so the entry takes InMemoryCache's own default rather than spend_counter_cache's 60 s
#[test]
fn a_local_increment_outlives_the_default_in_memory_ttl() {
let seconds = Arc::new(std::sync::atomic::AtomicU64::new(1_000));
let now = {
let seconds = seconds.clone();
Arc::new(move || Duration::from_secs(seconds.load(std::sync::atomic::Ordering::SeqCst)))
};
let store = SpendCounters::with_clock(None, now);
let counter = team("t");
store.increment_local(&counter, Number::Int(1)).unwrap();
seconds.fetch_add(61, std::sync::atomic::Ordering::SeqCst);
assert_eq!(store.get_local(&counter).unwrap(), Some(Number::Int(1)));
}
#[rstest]
#[tokio::test]
async fn a_failed_post_call_increment_drops_the_counter_from_memory() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let store = store(Some(connect(&deployment)));
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let counter = team("t");
store.increment(&counter, Number::Float(1.0)).await.unwrap();
assert!(store.increment_many_post_call(&request, &[(counter.clone(), Number::Float(1.0))]));
deployment.stop().await.unwrap();
request.flush_post().await;
tokio::time::sleep(Duration::from_millis(200)).await;
assert_eq!(store.get_local(&counter).unwrap(), None);
}
#[tokio::test]
async fn without_redis_counters_live_in_memory() {
let store = store(None);
let counter = team("t");
let request = RequestRedis::new(Handle::current(), LONG_DEADLINE);
let totals = store
.increment_many(&[
(counter.clone(), Number::Int(2)),
(counter.clone(), Number::Float(0.5)),
])
.await
.unwrap();
let seeded = store
.seed(&counter, Number::Float(1.0), LostSeed::Answer)
.await
.unwrap();
let raised = store.set_max(&counter, Number::Float(7.0)).await.unwrap();
assert_eq!(totals, [Number::Int(2), Number::Float(2.5)]);
assert_eq!(seeded, Number::Float(2.5));
assert_eq!(raised, None);
assert_eq!(
store.read(&counter).await.unwrap(),
Reading::Memory(Some(Number::Float(7.0)))
);
assert!(!store.increment_many_post_call(&request, &[(counter.clone(), Number::Int(1))]));
assert!(
store
.read_many(RedisSend::PreCall(&request), &[counter])
.is_none()
);
}
// main's RedisCache: async_increment(refresh_ttl=True), MGET for a batched read, async_set_max as
// a plain EVAL, SET EX NX for a seed and one multi-key DEL per chunk
#[rstest]
#[tokio::test]
async fn spend_counters_send_the_commands_main_sends(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let store = store(Some(connect(&deployment)));
let counters = [team("a"), team("b")];
deployment
.cli_on_every_node(&["CONFIG", "RESETSTAT"])
.await
.unwrap();
store
.increment(&counters[0], Number::Float(1.5))
.await
.unwrap();
store
.read_many(RedisSend::Now, &counters)
.unwrap()
.await
.unwrap()
.unwrap();
store
.set_max(&counters[0], Number::Float(9.0))
.await
.unwrap();
store
.seed(&counters[1], Number::Float(2.0), LostSeed::Answer)
.await
.unwrap();
store.delete(&counters).await.unwrap();
// commandstats also counts the GET, SET and EXPIRE the set_max script runs inside Redis
let calls = async |command| support::calls(&deployment, command).await;
assert_eq!(calls("incrbyfloat").await, 1);
assert_eq!(calls("expire").await, 2);
assert_eq!(calls("eval").await, 1);
assert_eq!(calls("evalsha").await, 0);
assert_eq!(calls("set").await, 2);
assert!(calls("mget").await > 0);
assert_eq!(calls("get").await, 1);
assert!(calls("del").await > 0);
}

View file

@ -10,9 +10,27 @@ use litellm_redis_client::{
CallEvent, Cluster, Credentials, Namespace, Node, Operation, RedisClient, RedisConfig,
RedisEvents, RedisUrl, Target, Timeouts, current_tag,
};
use litellm_redis_store::Remote;
use litellm_redis_store::{
Number, Remote, TierPolicy,
router_cooldowns::Cooldown,
router_usage::{Counter, Kind},
};
use litellm_redis_testing::RedisDeployment;
/// `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const COOLDOWN_POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(1),
tracked_reads: 1000,
default_in_memory_ttl: None,
};
/// `DualCache`'s default `redis_batch_cache_expiry` and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const USAGE_POLICY: TierPolicy = TierPolicy {
refresh_after: Duration::from_secs(10),
tracked_reads: 1000,
default_in_memory_ttl: None,
};
pub fn connect(deployment: &RedisDeployment) -> RedisClient {
connect_with(deployment, None, None)
}
@ -116,6 +134,103 @@ fn connect_with(
.unwrap()
}
/// What main's `RedisCache` writes for a cooldown: `str(value)` with `EX int(ttl)`, read back with
/// `json.loads` then `ast.literal_eval` (`litellm/caching/redis_cache.py` `set_cache` and
/// `_get_cache_logic`). Keys come from main's `CooldownCache`. The same literals sit in
/// `tests/unit/redis/stores/test_router_cooldowns.py`, checked against main's codec.
pub struct CooldownCase {
pub model_id: &'static str,
pub cooldown: Cooldown,
pub ttl_seconds: f64,
pub key: &'static str,
pub redis_ttl_seconds: u64,
pub str_dict: &'static str,
pub json: &'static str,
}
pub fn cooldowns() -> [CooldownCase; 2] {
[
CooldownCase {
model_id: "dep-a",
cooldown: Cooldown {
exception_received: "RateLimitError: limited é 'quoted'".to_owned(),
status_code: "429".to_owned(),
timestamp: Number::Float(1_730_000_000.25),
cooldown_time: Number::Int(60),
},
ttl_seconds: 90.5,
key: "deployment:dep-a:cooldown",
redis_ttl_seconds: 90,
str_dict: "{'exception_received': \"RateLimitError: limited é 'quoted'\", 'status_code': '429', 'timestamp': 1730000000.25, 'cooldown_time': 60}",
json: "{\"exception_received\": \"RateLimitError: limited \\u00e9 'quoted'\", \"status_code\": \"429\", \"timestamp\": 1730000000.25, \"cooldown_time\": 60}",
},
CooldownCase {
model_id: "dep-b",
cooldown: Cooldown {
exception_received: "Timeout".to_owned(),
status_code: "408".to_owned(),
timestamp: Number::Float(1_730_000_000.5),
cooldown_time: Number::Float(0.5),
},
ttl_seconds: 30.0,
key: "deployment:dep-b:cooldown",
redis_ttl_seconds: 30,
str_dict: "{'exception_received': 'Timeout', 'status_code': '408', 'timestamp': 1730000000.5, 'cooldown_time': 0.5}",
json: "{\"exception_received\": \"Timeout\", \"status_code\": \"408\", \"timestamp\": 1730000000.5, \"cooldown_time\": 0.5}",
},
]
}
/// main's allowed-fails keys for `dep-a`: an empty suffix is the same key as none.
pub struct AllowedFailsCase {
pub suffix: Option<&'static str>,
pub key: &'static str,
}
pub const ALLOWED_FAILS: [AllowedFailsCase; 3] = [
AllowedFailsCase {
suffix: None,
key: "deployment:dep-a:allowed_fails",
},
AllowedFailsCase {
suffix: Some(""),
key: "deployment:dep-a:allowed_fails",
},
AllowedFailsCase {
suffix: Some("timeout"),
key: "deployment:dep-a:allowed_fails:timeout",
},
];
/// main's `RouterCacheEnum` usage keys, including the `None` a missing model prints into the key.
pub struct UsageCase {
pub counter: Counter,
pub key: &'static str,
}
pub fn usage() -> [UsageCase; 2] {
[
UsageCase {
counter: Counter {
deployment_id: "dep".to_owned(),
model: Some("anthropic/claude-sonnet-5-5".to_owned()),
minute: "12-34".to_owned(),
kind: Kind::Tpm,
},
key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34",
},
UsageCase {
counter: Counter {
deployment_id: "dep".to_owned(),
model: None,
minute: "09-05".to_owned(),
kind: Kind::Rpm,
},
key: "global_router:dep:None:rpm:09-05",
},
]
}
/// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back.
pub fn within(ttl: Duration, whole_seconds: u64) -> bool {
ttl <= Duration::from_secs(whole_seconds)

View file

@ -105,8 +105,15 @@ def current_service_target() -> str | None:
return _service_target.get()
def with_service_target(target: str) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
"""Run every call of the decorated function, coroutine functions included, under ``service_target(target)``."""
def with_service_target(
target: str, *, preserve_existing: bool = False
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
"""Run calls under the target, or the inherited target when ``preserve_existing`` is true."""
def call_target() -> str:
if preserve_existing:
return current_service_target() or target
return target
def decorate(fn: Callable[_P, _R]) -> Callable[_P, _R]:
if inspect.iscoroutinefunction(fn):
@ -116,14 +123,14 @@ def with_service_target(target: str) -> Callable[[Callable[_P, _R]], Callable[_P
@wraps(fn)
async def run_async(*args: _P.args, **kwargs: _P.kwargs) -> object:
with service_target(target):
with service_target(call_target()):
return await awaitable_fn(*args, **kwargs)
return cast("Callable[_P, _R]", run_async) # cast-ok: same coroutine-returning signature as ``fn``
@wraps(fn)
def run(*args: _P.args, **kwargs: _P.kwargs) -> _R:
with service_target(target):
with service_target(call_target()):
return fn(*args, **kwargs)
return run

View file

@ -518,7 +518,7 @@ def _or_secrets(value: object, read: Callable[[str], object], *names: str) -> ob
return reduce(lambda found, name: found or read(name), names, value)
def _merge_redis_kwargs(
def merge_redis_kwargs(
env_overrides: Mapping[str, object],
) -> dict[str, object]: # mutable-ok: _get_redis_client_logic swaps settings for client objects in place
overrides: Final = {
@ -580,7 +580,7 @@ def _get_redis_client_logic(**env_overrides):
"""
Common functionality across sync + async redis client implementations
"""
redis_kwargs: Final = _merge_redis_kwargs(env_overrides)
redis_kwargs: Final = merge_redis_kwargs(env_overrides)
_startup_nodes: Final = redis_kwargs.get("startup_nodes") or None

View file

@ -12,14 +12,14 @@ import asyncio
import itertools
import logging
import time
from collections.abc import Mapping, Sequence
from collections.abc import Awaitable, Mapping, Sequence
from dataclasses import dataclass
from threading import Lock
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS
from litellm.types.caching import RedisPipelineIncrementOperation
from .base_cache import BaseCache
@ -91,7 +91,9 @@ class DualCache(BaseCache):
self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size)
self._last_redis_batch_access_time_lock = Lock()
self.redis_batch_cache_expiry = (
default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10
default_redis_batch_cache_expiry
or litellm.default_redis_batch_cache_expiry
or DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS
)
self.default_in_memory_ttl = default_in_memory_ttl or litellm.default_in_memory_ttl
self.default_redis_ttl = default_redis_ttl or litellm.default_redis_ttl
@ -497,7 +499,7 @@ class DualCache(BaseCache):
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)
return None if batch is None else await self._set_on_batch(batch, key, value, ttl)
async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None:
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
"""Memory now, the Redis DEL on the request's pipeline; None when no pipeline is open, so the caller
takes its direct path."""
batch: Final = None if self.redis_cache is None else active_request_redis_batch(self.redis_cache)

View file

@ -622,9 +622,11 @@ class RequestRedisBatches:
self._deadline_flush = asyncio.ensure_future(self.flush_post_call())
async def flush_all(self) -> None:
"""Batch 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. The pre-call
flush hooks are skipped: they declare reads for later readers, and none are left."""
await asyncio.gather(
*(batch.flush() for batch in self._batches.values() if batch.pending), self.flush_native_pre()
*(batch.flush() for batch in self._batches.values() if batch.pending),
*(() if self._native is None else (self._native.flush_pre(),)),
)
async def flush_post_call(self) -> None:

View file

@ -42,7 +42,7 @@ from litellm.constants import (
)
from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.redis.store_settings import RedisStoreSettings
from litellm.redis.store_settings import RedisStoreSettings, breaker_settings
from litellm.types.caching import (
RedisPipelineIncrementOperation,
RedisPipelineLpopOperation,
@ -661,7 +661,7 @@ class RedisCache(BaseCache):
):
from litellm._service_logger import ServiceLogging
from .._redis import get_redis_client, get_redis_connection_pool
from .._redis import get_redis_client, get_redis_connection_pool, merge_redis_kwargs
redis_kwargs: Final = {}
if host is not None:
@ -715,9 +715,10 @@ class RedisCache(BaseCache):
else:
super().__init__() # defaults to 60s
self.store_settings = RedisStoreSettings(
redis_kwargs=MappingProxyType(dict(redis_kwargs)),
redis_kwargs=MappingProxyType(merge_redis_kwargs(redis_kwargs)),
namespace=namespace,
default_ttl=self.default_ttl,
breaker=breaker_settings(),
)
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:

View file

@ -159,6 +159,7 @@ DEFAULT_ALLOWED_FAILS: Final = get_env_int("DEFAULT_ALLOWED_FAILS", 3)
DEFAULT_REDIS_SYNC_INTERVAL: Final = get_env_int("DEFAULT_REDIS_SYNC_INTERVAL", 1)
DEFAULT_COOLDOWN_TIME_SECONDS: Final = get_env_int("DEFAULT_COOLDOWN_TIME_SECONDS", 5)
DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS: Final = get_env_float("DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 1.0)
DEFAULT_REDIS_BATCH_CACHE_EXPIRY_SECONDS: Final = 10
DEFAULT_REPLICATE_POLLING_RETRIES: Final = get_env_int("DEFAULT_REPLICATE_POLLING_RETRIES", 5)
DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = get_env_int("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)
DEFAULT_IMAGE_TOKEN_COUNT: Final = get_env_int("DEFAULT_IMAGE_TOKEN_COUNT", 250)

View file

@ -40,6 +40,9 @@ from litellm.litellm_core_utils.llm_judge import (
)
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.redis.stores.spend_counters import (
shadow_eval_counter,
)
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.management_endpoints.auto_router_endpoints import ShadowEvalDirection
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
@ -587,7 +590,7 @@ def _leg_eval_spend(sums: Mapping[str, object]) -> float:
def _job_spend_counter_key(job_id: str) -> str:
return f"spend:shadow_eval:{job_id}"
return shadow_eval_counter(job_id)
async def _job_spend_from_counter(counter_key: str, fallback_spend: float, max_budget: float) -> float:

View file

@ -133,10 +133,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
get_management_object_ttl,
model_access_group_cache_key,
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
object_permission_cache_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
tag_registry_cache_key,
team_membership_auth_cache_key,
@ -151,8 +149,21 @@ from litellm.proxy.guardrails.tool_name_extraction import (
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
end_user_counter,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
from litellm.repositories.organization_repository import OrganizationRepository
@ -1426,7 +1437,7 @@ async def common_checks(
user_budget: Final = user_object.max_budget
user_spend: Final = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
counter_key=user_counter(user_object.user_id),
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
@ -1974,7 +1985,7 @@ async def check_end_user_budget(
from litellm.proxy.proxy_server import get_current_spend
end_user_spend: Final = await get_current_spend(
counter_key=f"spend:end_user:{end_user_obj.user_id}",
counter_key=end_user_counter(end_user_obj.user_id),
fallback_spend=end_user_obj.spend or 0.0,
max_budget=end_user_budget,
fallback_authoritative=True,
@ -3187,7 +3198,7 @@ async def cache_team_object(
usage_cache: Final = None if proxy_logging_obj is None else proxy_logging_obj.internal_usage_cache.dual_cache
# On a shared Redis the write below replaces the team entry and the alias DEL below removes the alias entry
# for both caches, so the usage cache only has its own memory to clear.
redis_shared: Final = usage_cache is not None and usage_cache.redis_cache is user_api_key_cache.redis_cache
redis_shared: Final = usage_cache is not None and usage_cache.redis_cache is user_api_key_cache.redis_connection
await _invalidate_usage_cache_entry(usage_cache, key, redis_shared=redis_shared, stale="team object")
@ -3328,37 +3339,40 @@ async def invalidate_team_member_spend_state(
await await_inflight_team_membership_load(user_id=user_id, team_id=team_id)
if new_spend is not None:
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import (
SPEND_DB_FLOOR_CACHE_TTL_SECONDS,
spend_counter_cache,
spend_counter_store,
)
spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}"
spend_counter_key: Final = team_member_counter(user_id, team_id)
spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}"
spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_store.set(spend_counter_key, new_spend, ttl=60)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
"read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
spend_counter_key,
e,
)
try:
await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60)
except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up
await spend_counter_store.delete((spend_counter_key,))
except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
verbose_proxy_logger.warning(
"Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next "
"read reseeds from the DB rather than keeping the stale pre-reset value authoritative",
"Failed to delete stale spend counter %s in Redis after a failed reset write",
spend_counter_key,
e,
exc_info=True,
)
try:
await spend_counter_cache.redis_cache.async_delete_cache(key=spend_counter_key)
except Exception: # noqa: BLE001 # stale value now authoritative in Redis; surface instead of reporting success
verbose_proxy_logger.warning(
"Failed to delete stale spend counter %s in Redis after a failed reset write",
spend_counter_key,
exc_info=True,
)
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={
"error": "Spend was reset in the database, but Redis is unreachable and still "
"holds the pre-reset counter. Retry once Redis is reachable."
},
) from e
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail={
"error": "Spend was reset in the database, but Redis is unreachable and still "
"holds the pre-reset counter. Retry once Redis is reachable."
},
) from e
spend_counter_store.set_local(spend_counter_key, new_spend, ttl=60)
spend_counter_cache.in_memory_cache.set_cache(
key=spend_db_floor_key,
@ -6722,7 +6736,7 @@ async def virtual_key_max_budget_check(
from litellm.proxy.proxy_server import get_current_spend
fallback_spend: Final = valid_token.spend or 0.0
counter_key: Final = f"spend:key:{valid_token.token}"
counter_key: Final = key_counter(str(valid_token.token))
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
spend: Final = await get_current_spend(
@ -6811,7 +6825,7 @@ async def _virtual_key_multi_budget_check(
for window in valid_token.budget_limits:
w: dict = window if isinstance(window, dict) else window.model_dump()
counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}"
counter_key = key_window_counter(str(valid_token.token), str(w["budget_duration"]))
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
@ -7118,7 +7132,7 @@ async def _check_team_member_budget(
from litellm.proxy.proxy_server import get_current_spend
team_member_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
counter_key=team_member_counter(valid_token.user_id, team_object.team_id),
fallback_spend=team_member_spend,
max_budget=team_member_budget,
)
@ -7230,7 +7244,7 @@ async def team_max_budget_check(
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
spend: Final = await get_current_spend(
counter_key=f"spend:team:{team_object.team_id}",
counter_key=team_counter(team_object.team_id),
fallback_spend=team_object.spend or 0.0,
max_budget=team_object.max_budget,
)
@ -7283,7 +7297,7 @@ async def _team_multi_budget_check(
for window in team_object.budget_limits:
w: dict = window if isinstance(window, dict) else window.model_dump()
counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}"
counter_key = team_window_counter(team_object.team_id, str(w["budget_duration"]))
window_spend = await get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
@ -7402,7 +7416,7 @@ async def _project_max_budget_check(
from litellm.proxy.proxy_server import get_current_spend
project_spend: Final = await get_current_spend(
counter_key=project_spend_counter_key(project_object.project_id),
counter_key=project_counter(project_object.project_id),
fallback_spend=project_object.spend or 0.0,
max_budget=max_budget,
)
@ -7609,7 +7623,7 @@ async def organization_max_budget_check(
from litellm.proxy.proxy_server import get_current_spend
org_spend: Final = await get_current_spend(
counter_key=f"spend:org:{org_id}",
counter_key=org_counter(org_id),
fallback_spend=org_table.spend or 0.0,
max_budget=org_max_budget,
)
@ -7699,7 +7713,7 @@ async def tag_max_budget_check_for_tags(
from litellm.proxy.proxy_server import get_current_spend
tag_spend = await get_current_spend(
counter_key=f"spend:tag:{tag_name}",
counter_key=tag_counter(tag_name),
fallback_spend=tag_object.spend or 0.0,
max_budget=tag_object.litellm_budget_table.max_budget,
fallback_authoritative=True,
@ -7754,7 +7768,7 @@ async def _model_access_group_max_budget_check(
continue
group_spend = await get_current_spend(
counter_key=model_access_group_spend_counter_key(group),
counter_key=model_access_group_counter(group),
fallback_spend=budget.spend,
max_budget=budget.max_budget,
fallback_authoritative=True,

View file

@ -5,17 +5,16 @@ the readers and the fallback, so enforcement never depends on this running."""
from __future__ import annotations
import time
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, Protocol, TypeAlias
from typing import Final, Literal, TypeAlias
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 active_request_redis_batch
from litellm.caching.redis_cache import RedisCache
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
@ -32,6 +31,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
)
from litellm.proxy.db.db_span import db_span
from litellm.proxy.utils import PrismaClient
from litellm.redis.stores.auth_cache import AuthCacheStore
_RowKind: TypeAlias = Literal["user_row", "team_row", "membership_row", "organization_row", "project_row"]
@ -162,11 +162,6 @@ class AuthObjectRefs:
)
class _InMemoryCache(Protocol):
def get_cache(self, key: str) -> object: ...
def set_cache(self, key: str, value: object, *, ttl: float | None = ...) -> None: ...
@dataclass(frozen=True, slots=True)
class _CacheEntry:
cache_key: str
@ -211,39 +206,32 @@ def _entries(refs: AuthObjectRefs, cache: UserApiKeyCache) -> tuple[_CacheEntry,
return tuple(_iter_entries(refs, get_management_object_ttl(cache)))
def _missing_in_memory(entries: Sequence[_CacheEntry], memory: _InMemoryCache) -> tuple[_CacheEntry, ...]:
return tuple(entry for entry in entries if memory.get_cache(key=entry.cache_key) is None)
def _missing_in_memory(entries: Sequence[_CacheEntry], cache: UserApiKeyCache) -> tuple[_CacheEntry, ...]:
return tuple(entry for entry in entries if cache.get_local(entry.cache_key) is None)
def _set_in_memory(memory: _InMemoryCache, cache_key: str, value: object, ttl: float | None) -> None:
if ttl is None:
memory.set_cache(key=cache_key, value=value)
else:
memory.set_cache(key=cache_key, value=value, ttl=ttl)
def _nothing_to_arm() -> None:
return None
async def _read_redis_rows(keys: list[str], redis_cache: RedisCache) -> Mapping[str, object]:
"""On the request pipeline when one is open; a failed pipeline reads as a miss, like ``async_batch_get_cache``."""
batch: Final = active_request_redis_batch(redis_cache)
with service_target(AUTH_OBJECTS_TARGET):
if batch is None:
return await redis_cache.async_batch_get_cache(key_list=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API
try:
return await batch.mget(keys)
except Exception as e: # noqa: BLE001 # the DB fill below takes over, as it does after a failed MGET today
verbose_proxy_logger.debug("auth prefetch Redis read failed, filling from the database: %s", e)
return MappingProxyType({})
async def _fill_from_redis(entries: Sequence[_CacheEntry], redis_cache: RedisCache, memory: _InMemoryCache) -> None:
async def _fill_from_redis(
entries: Sequence[_CacheEntry], shared: AuthCacheStore, cache: UserApiKeyCache, before_read: Callable[[], None]
) -> None:
"""On the request's native pipeline when one is open; a failed read is a miss, like ``async_batch_get_cache``."""
if not entries:
return
found: Final = _RowValues.validate_python(
await _read_redis_rows(sorted(entry.cache_key for entry in entries), redis_cache)
)
before_read()
try:
with service_target(AUTH_OBJECTS_TARGET):
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
for entry, value in ((entry, found.get(entry.cache_key)) for entry in entries):
if value is not None:
_set_in_memory(memory, entry.cache_key, value, entry.ttl)
shared.set_local(
entry.cache_key, value, cache.in_memory_cache.default_ttl if entry.ttl is None else entry.ttl
)
def _validate_row(
@ -282,18 +270,13 @@ async def _write_back(entries: Sequence[tuple[_CacheEntry, BaseModel]], cache: U
(entry.cache_key, CacheCodec.serialize(value, model_type=entry.model_type), entry.ttl)
for entry, value in entries
)
memory: Final[_InMemoryCache] = cache.in_memory_cache
for cache_key, payload, ttl in payloads:
_set_in_memory(memory, cache_key, payload, cache.default_in_memory_ttl if ttl is None else ttl)
if cache.redis_cache is None:
shared: Final = cache.shared
if shared is not None:
with service_target(AUTH_OBJECTS_TARGET):
await shared.write_back(payloads, batch=Batch.PRE_CALL)
return
batch: Final = active_request_redis_batch(cache.redis_cache)
with service_target(AUTH_OBJECTS_TARGET):
if batch is None:
await cache.redis_cache.async_set_cache_pipeline_with_ttls(payloads)
return
for cache_key, payload, ttl in payloads: # rides the request's next round trip; the scope drains leftovers
batch.set(cache_key, payload, ttl)
for cache_key, payload, ttl in payloads:
cache.set_local(cache_key, payload, ttl)
async def _fill_from_db(
@ -318,16 +301,18 @@ async def prefetch_auth_objects(
refs: AuthObjectRefs,
user_api_key_cache: UserApiKeyCache,
prisma_client: PrismaClient | None,
before_read: Callable[[], None] = _nothing_to_arm,
) -> None:
"""Best effort: any failure leaves the per-object getters to fetch as before."""
"""Best effort: any failure leaves the per-object getters to fetch as before. ``before_read`` runs right before
a Redis read goes out, so a caller can declare other reads into the same round trip."""
try:
memory: Final[_InMemoryCache] = user_api_key_cache.in_memory_cache
missing: Final = _missing_in_memory(_entries(refs, user_api_key_cache), memory)
if user_api_key_cache.redis_cache is not None:
await _fill_from_redis(missing, user_api_key_cache.redis_cache, memory)
missing: Final = _missing_in_memory(_entries(refs, user_api_key_cache), user_api_key_cache)
shared: Final = user_api_key_cache.shared
if shared is not None:
await _fill_from_redis(missing, shared, user_api_key_cache, before_read)
if prisma_client is None:
return
await _fill_from_db(refs, _missing_in_memory(missing, memory), user_api_key_cache, prisma_client)
await _fill_from_db(refs, _missing_in_memory(missing, user_api_key_cache), user_api_key_cache, prisma_client)
except Exception as e: # noqa: BLE001 # warm-up only; the getters enforce and fail closed on their own
verbose_proxy_logger.warning("auth prefetch skipped, falling back to per-object lookups: %s", e)
@ -337,28 +322,29 @@ def _identity_memory_ttl(value: object, management_ttl: float) -> float:
return min(REGISTRY_ERROR_NEGATIVE_CACHE_TTL, management_ttl) if isinstance(value, str) else management_ttl
async def prefetch_identity_keys(cache_keys: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None:
async def prefetch_identity_keys(
cache_keys: Sequence[str],
user_api_key_cache: UserApiKeyCache,
before_read: Callable[[], None] = _nothing_to_arm,
) -> None:
"""Warm the entries auth reads before it knows the key's owners (the key object, the end user and the two
registries) in one MGET on the request pipeline. Keys the MGET finds absent stay noted on the pipeline, so the
registries) in one read on the request's native pipeline. Keys it finds absent stay noted on the request, so the
per-key getters that follow go to the database without a GET of their own. Best effort, like the
owner prefetch: the getters read and enforce on their own."""
try:
redis_cache: Final = user_api_key_cache.redis_cache
if redis_cache is None:
shared: Final = user_api_key_cache.shared
if shared is None:
return
missing: Final = tuple(
key
for key in dict.fromkeys(cache_keys)
if user_api_key_cache.in_memory_cache_for(key).get_cache(key=key) is None
)
missing: Final = tuple(key for key in dict.fromkeys(cache_keys) if shared.get_local(key) is None)
if not missing:
return
found: Final = _RowValues.validate_python(await _read_redis_rows(sorted(missing), redis_cache))
before_read()
with service_target(AUTH_OBJECTS_TARGET):
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)
return
for key, value in ((key, found.get(key)) for key in missing):
if value is not None:
memory: _InMemoryCache = user_api_key_cache.in_memory_cache_for(key)
_set_in_memory(memory, key, value, _identity_memory_ttl(value, management_ttl))
shared.set_local(key, value, _identity_memory_ttl(value, management_ttl))

View file

@ -44,6 +44,10 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
_is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_model_cost_zero, # pyright: ignore[reportPrivateUsage] # the zero-cost predicate the auth-time budget checks use; no public equivalent
)
from litellm.redis.stores.spend_counters import (
key_counter,
user_counter,
)
from litellm.router_api import RegistryReader, RoutingContext
from litellm.types.llms.base import LiteLLMBaseModel
@ -133,7 +137,7 @@ async def is_token_within_budget_for_model(
key_budget: Final = valid_token.max_budget
if key_budget is not None and valid_token.token is not None:
key_spend: Final = await _counter_spend(
counter_key=f"spend:key:{valid_token.token}",
counter_key=key_counter(valid_token.token),
fallback_spend=valid_token.spend or 0.0,
max_budget=key_budget,
)
@ -149,7 +153,7 @@ async def is_token_within_budget_for_model(
and (valid_token.team_id is None or _applies_user_budget_to_team_keys())
):
user_spend: Final = await _counter_spend(
counter_key=f"spend:user:{valid_token.user_id}",
counter_key=user_counter(valid_token.user_id),
fallback_spend=valid_token.user_spend or 0.0,
max_budget=user_budget,
)

View file

@ -24,10 +24,9 @@ from pydantic import TypeAdapter
from starlette.exceptions import WebSocketException
import litellm
from litellm._internal_context import service_target
from litellm._internal_context import service_target, set_request_parent_otel_span
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
CLIENT_REQUESTED_MODEL_SCOPE_KEY,
GLOBAL_PROXY_SPEND_CACHE_KEY,
@ -182,10 +181,19 @@ from litellm.proxy.utils import (
ProxyLogging,
normalize_route_for_root_path,
)
from litellm.redis.stores.spend_counters import (
SpendCounterStore,
team_member_counter,
)
from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository
from litellm.repositories.user_repository import UserRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.router_api import fallbacks_disabled, references_provider_scoped_resource, request_fallback_chain
from litellm.router_api import (
RoutingContext,
fallbacks_disabled,
references_provider_scoped_resource,
request_fallback_chain,
)
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
@ -2051,7 +2059,8 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
if open_telemetry_logger is None:
return
if getattr(request.state, "parent_otel_span", None) is not None:
if (existing_span := getattr(request.state, "parent_otel_span", None)) is not None:
set_request_parent_otel_span(existing_span)
return
parent_otel_span: Final = open_telemetry_logger.create_litellm_proxy_request_started_span(
start_time=start_time,
@ -2067,6 +2076,7 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None:
http_route=get_request_route_template(request),
)
request.state.parent_otel_span = parent_otel_span
set_request_parent_otel_span(parent_otel_span)
async def _read_request_body_deferring_parse_failure(
@ -2648,6 +2658,7 @@ async def user_api_key_auth_builder(
await prefetch_identity_keys(
_identity_cache_keys(api_key, end_user_id=end_user_id, key_is_resolved=valid_token is not None),
user_api_key_cache=user_api_key_cache,
before_read=arm_routing_cooldowns,
)
if end_user_id:
try:
@ -3097,7 +3108,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
team_member_spend = valid_token.team_member_spend
if valid_token.user_id is not None and valid_token.team_id is not None:
team_member_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}",
counter_key=team_member_counter(valid_token.user_id, valid_token.team_id),
fallback_spend=team_member_spend,
max_budget=team_member_budget,
)
@ -4266,10 +4277,10 @@ async def _authorize_authenticated_request(
return None
def _spend_counter_redis_cache() -> RedisCache | None:
from litellm.proxy.proxy_server import spend_counter_cache
def _spend_counter_store() -> SpendCounterStore:
from litellm.proxy.proxy_server import spend_counter_store
return spend_counter_cache.redis_cache
return spend_counter_store
def _identity_cache_keys(api_key: str, *, end_user_id: str | None, key_is_resolved: bool) -> tuple[str, ...]:
@ -4300,9 +4311,20 @@ async def _prefetch_referenced_auth_objects(
refs=AuthObjectRefs.from_token(valid_token),
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
before_read=arm_routing_cooldowns,
)
def arm_routing_cooldowns() -> None:
"""Declares the cooldown read routing takes later into the auth prefetch's native round trip, so a request whose
auth objects come from Redis pays one round trip for both. Armed only when a prefetch reads Redis, so a request
served from memory, or denied, reads no cooldowns early."""
from litellm.proxy.proxy_server import llm_router # noqa: PLC0415 # runtime global
if llm_router is not None:
RoutingContext(llm_router).arm_cooldown_prefetch()
def _seed_request_destinations(user_api_key_dict: UserAPIKeyAuth, request: Request | None = None) -> None:
"""Anchor the OTLP destinations this key or team overrides its traces to.
@ -4372,7 +4394,7 @@ async def user_api_key_auth(
with (
phase_span(f"auth {route}"),
service_target(AUTH_OBJECTS_TARGET),
spend_counter_batch_scope(_spend_counter_redis_cache()),
spend_counter_batch_scope(_spend_counter_store()),
):
try:
user_api_key_auth_obj: Final = await user_api_key_auth_builder(

View file

@ -3,7 +3,7 @@ import json
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Final, Protocol
from pydantic import TypeAdapter
@ -19,6 +19,13 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # leg
)
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
class SpendCounterCopies(Protocol):
"""This worker's memory copies of the spend counters, which a broadcast SET or delete reaches too."""
def apply_broadcast(self, cache_key: str, new_value: float | None, ttl: float | None) -> None: ...
if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
@ -183,7 +190,7 @@ def evict_local(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache"
def _drop_from_memory(cache_key: str, user_api_key_cache: "UserApiKeyCache") -> None:
user_api_key_cache.delete_from_memory(cache_key)
user_api_key_cache.delete_local(cache_key)
@with_service_target(AUTH_OBJECTS_TARGET)
@ -232,17 +239,19 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
class AuthCacheInvalidationSubscriber:
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache")
__slots__ = ("_additional_in_memory_caches", "_redis_cache", "_spend_counters", "_task", "_user_api_key_cache")
def __init__(
self,
redis_cache: "RedisCache",
user_api_key_cache: "UserApiKeyCache",
additional_in_memory_caches: Sequence["InMemoryCache"] = (),
spend_counters: SpendCounterCopies | None = None,
) -> None:
self._redis_cache = redis_cache
self._user_api_key_cache = user_api_key_cache
self._additional_in_memory_caches = tuple(additional_in_memory_caches)
self._spend_counters: Final = spend_counters
self._task: asyncio.Task[None] | None = None
def start(self) -> None:
@ -297,6 +306,8 @@ class AuthCacheInvalidationSubscriber:
parsed: Final = _message_from_data(data)
if parsed is None:
return
if self._spend_counters is not None:
self._spend_counters.apply_broadcast(parsed.cache_key, parsed.new_value, parsed.ttl)
if parsed.new_value is not None:
for additional_cache in self._additional_in_memory_caches:
additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl)

View file

@ -16,7 +16,8 @@ are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` f
from __future__ import annotations
from typing import Any, TypeVar
from collections.abc import Mapping, Sequence
from typing import Any, TypeAlias, TypeVar
from pydantic import BaseModel, ValidationError
@ -24,6 +25,9 @@ from litellm._logging import verbose_proxy_logger
T = TypeVar("T", bound=BaseModel)
CacheData: TypeAlias = "str | int | float | bool | None | Sequence[CacheData] | Mapping[str, CacheData]"
CacheValue: TypeAlias = "BaseModel | CacheData"
class CacheCodec:
"""

View file

@ -15,7 +15,6 @@ from typing_extensions import assert_never
import litellm
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import (
GLOBAL_PROXY_SPEND_CACHE_KEY,
LITELLM_PROXY_BUDGET_NAME,
@ -42,17 +41,29 @@ from litellm.proxy.common_utils.user_api_key_cache import (
AUTH_OBJECTS_TARGET,
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
from litellm.proxy.db.db_span import db_span, db_spanned
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
SpendCounterStore,
end_user_counter,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import PrismaBatch, SpendLinkedTable
from litellm.repositories.project_repository import ProjectRepository
@ -157,7 +168,7 @@ def _row_carried_spend(row: _BudgetLinkedRow, caps: Mapping[str, float]) -> floa
def _team_membership_counter_key(row: _TeamMembershipRow) -> str:
return f"spend:team_member:{row.user_id}:{row.team_id}"
return team_member_counter(row.user_id, row.team_id)
def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]:
@ -165,7 +176,7 @@ def _team_membership_cache_keys(row: _TeamMembershipRow) -> tuple[str, ...]:
def _key_counter_key(row: _KeyRow) -> str:
return f"spend:key:{row.token}"
return key_counter(row.token)
def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]:
@ -173,7 +184,7 @@ def _key_cache_keys(row: _KeyRow) -> tuple[str, ...]:
def _org_counter_key(row: _OrgRow) -> str:
return f"spend:org:{row.organization_id}"
return org_counter(row.organization_id)
def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]:
@ -184,7 +195,7 @@ def _org_cache_keys(row: _OrgRow) -> tuple[str, ...]:
def _tag_counter_key(row: _TagRow) -> str:
return f"spend:tag:{row.tag_name}"
return tag_counter(row.tag_name)
def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]:
@ -192,7 +203,7 @@ def _tag_cache_keys(row: _TagRow) -> tuple[str, ...]:
def _model_access_group_counter_key(row: _ModelAccessGroupRow) -> str:
return model_access_group_spend_counter_key(row.access_group_name)
return model_access_group_counter(row.access_group_name)
def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]:
@ -200,7 +211,7 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
def _project_counter_key(row: _ProjectRow) -> str:
return project_spend_counter_key(row.project_id)
return project_counter(row.project_id)
def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
@ -208,7 +219,7 @@ def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
def _enduser_counter_key(row: _EndUserRow) -> str:
return f"spend:end_user:{row.user_id}"
return end_user_counter(row.user_id)
def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]:
@ -427,7 +438,7 @@ class _WindowSource:
table: str
id_column: str
entity_type: Litellm_EntityType
counter_prefix: str
window_counter: Callable[[str, str], str]
log_subject: str
retry_subject: str
write: Callable[[PrismaClient, str, str], Awaitable[None]]
@ -452,7 +463,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
table="LiteLLM_VerificationToken",
id_column="token",
entity_type=Litellm_EntityType.KEY,
counter_prefix="spend:key",
window_counter=key_window_counter,
log_subject="keys",
retry_subject="key",
write=_write_key_windows,
@ -461,7 +472,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = (
table="LiteLLM_TeamTable",
id_column="team_id",
entity_type=Litellm_EntityType.TEAM,
counter_prefix="spend:team",
window_counter=team_window_counter,
log_subject="teams",
retry_subject="team",
write=_write_team_windows,
@ -607,22 +618,16 @@ class ResetBudgetJob:
commit opens a window where get_current_spend reads 0 from Redis
while the DB still holds the pre-reset value, allowing bypass.
"""
try:
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. "
"Budget may be over-enforced until counter expires.",
counter_key,
redis_err,
)
try:
await spend_counter_store.delete((counter_key,))
except Exception as e:
verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e)
verbose_proxy_logger.warning(
"Failed to reset spend counter %s in Redis: %s. Budget may be over-enforced until counter expires.",
counter_key,
e,
)
@staticmethod
async def _invalidate_global_proxy_spend_cache() -> None:
@ -671,11 +676,11 @@ class ResetBudgetJob:
if not keys:
return
try:
from litellm.proxy.proxy_server import spend_counter_cache, user_api_key_cache
from litellm.proxy.proxy_server import spend_counter_store, user_api_key_cache
match cache:
case "spend counter":
await spend_counter_cache.async_delete_cache_keys(keys)
await spend_counter_store.delete(keys)
case "user_api_key_cache":
await user_api_key_cache.async_delete_cache_keys(keys)
case _:
@ -1176,7 +1181,7 @@ class ResetBudgetJob:
for k in updated_keys:
token = getattr(k.row, "token", None)
if token:
await self._invalidate_spend_counter(f"spend:key:{token}")
await self._invalidate_spend_counter(key_counter(str(token)))
end_time = time.time()
outcome: Final = _ChunkOutcome(
@ -1287,7 +1292,7 @@ class ResetBudgetJob:
for u in updated_users:
user_id = getattr(u.row, "user_id", None)
if user_id:
await self._invalidate_spend_counter(f"spend:user:{user_id}")
await self._invalidate_spend_counter(user_counter(str(user_id)))
if user_id == LITELLM_PROXY_BUDGET_NAME:
await self._invalidate_global_proxy_spend_cache()
@ -1402,7 +1407,7 @@ class ResetBudgetJob:
for t in updated_teams:
team_id = getattr(t.row, "team_id", None)
if team_id:
await self._invalidate_spend_counter(f"spend:team:{team_id}")
await self._invalidate_spend_counter(team_counter(str(team_id)))
end_time = time.time()
outcome: Final = _ChunkOutcome(
@ -1465,7 +1470,7 @@ class ResetBudgetJob:
async def _reset_expired_window(
window: dict,
counter_key: str,
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
now: datetime,
reset_settings: BudgetResetSettings,
prisma_client: PrismaClient,
@ -1479,13 +1484,11 @@ class ResetBudgetJob:
reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None)
if reset_at > now:
return False
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache)
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_value)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value)
except Exception as redis_err:
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_store)
try:
await spend_counter_store.set(counter_key, new_value)
except Exception as redis_err:
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
budget_duration: Final = window["budget_duration"]
next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings)
window["reset_at"] = next_reset_at.isoformat()
@ -1539,7 +1542,7 @@ class ResetBudgetJob:
@staticmethod
@with_service_target(SPEND_COUNTERS_TARGET)
async def _window_carried_spend(
window: Mapping[str, object], counter_key: str, spend_counter_cache: DualCache
window: Mapping[str, object], counter_key: str, spend_counter_store: SpendCounterStore
) -> float:
"""Per-window spend lives only in the counter, so the carried overage is
read from it before the reset overwrites it."""
@ -1550,11 +1553,11 @@ class ResetBudgetJob:
if cap is None:
return 0.0
try:
current: Final = await spend_counter_cache.async_get_cache(key=counter_key)
current, _ = await spend_counter_store.read(counter_key)
except Exception as e: # noqa: BLE001 # an unreadable counter falls back to a plain zero reset
verbose_proxy_logger.warning("Failed to read spend counter %s for rollover: %s", counter_key, e)
return 0.0
if not isinstance(current, (int, float)):
if current is None:
return 0.0
return _carried_spend(float(current), cap)
@ -1564,12 +1567,12 @@ class ResetBudgetJob:
reset_at <= now. Only the expired windows are reset; other windows are untouched.
"""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
now: Final = datetime.utcnow()
for source in _WINDOW_SOURCES:
try:
await self._reset_windows_for(source=source, now=now, spend_counter_cache=spend_counter_cache)
await self._reset_windows_for(source=source, now=now, spend_counter_store=spend_counter_store)
except Exception as e:
verbose_proxy_logger.exception("Failed to reset budget windows for %s: %s", source.log_subject, e)
@ -1577,7 +1580,7 @@ class ResetBudgetJob:
self,
source: _WindowSource,
now: datetime,
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
) -> None:
"""Walk one table's windowed rows a page at a time, to the end.
@ -1599,7 +1602,7 @@ class ResetBudgetJob:
source=source,
cursor=cursor,
now=now,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
)
if next_cursor is None:
return
@ -1610,7 +1613,7 @@ class ResetBudgetJob:
source: _WindowSource,
cursor: str,
now: datetime,
spend_counter_cache: DualCache,
spend_counter_store: SpendCounterStore,
) -> str | None:
"""Reset one page of windows; return the next cursor, or None when drained."""
rows: Final = await self._with_db_retry(
@ -1629,11 +1632,11 @@ class ResetBudgetJob:
windows: list[dict[str, object]] = raw if isinstance(raw, list) else json.loads(raw)
changed = False
for window in windows:
counter_key = f"{source.counter_prefix}:{row_id}:window:{window['budget_duration']}"
counter_key = source.window_counter(row_id, str(window["budget_duration"]))
if await ResetBudgetJob._reset_expired_window(
window,
counter_key,
spend_counter_cache,
spend_counter_store,
now,
self.reset_settings,
prisma_client=self.prisma_client,

View file

@ -2,23 +2,23 @@ from __future__ import annotations
import asyncio
import re
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar, cast, overload
from pydantic import BaseModel
from pydantic import BaseModel, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec, CacheValue
from litellm.redis.stores.auth_cache import AuthCacheSettings, AuthCacheStore
from litellm.redis.stores.connection import connect
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm.caching.redis_batch import BatchResult
T = TypeVar("T", bound=BaseModel)
from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa: E402, PLC0414 # public re-export
@ -26,6 +26,12 @@ from litellm.constants import AUTH_OBJECTS_TARGET as AUTH_OBJECTS_TARGET # noqa
_HASHED_TOKEN_CACHE_KEY: Final = re.compile(r"[0-9a-f]{64}")
_BUDGET_WINDOWS_ABSENT_PREFIX: Final = "budget_windows_absent:"
AuthCacheFactory: TypeAlias = Callable[[RedisCache, AuthCacheSettings], AuthCacheStore]
def _connected_auth_cache(redis_cache: RedisCache, settings: AuthCacheSettings) -> AuthCacheStore:
return AuthCacheStore(connect(redis_cache).auth_cache(settings))
def is_user_key_cache_key(key: str) -> bool:
"""Only user-key objects are cached under a bare ``hash_token`` digest; every other object uses a prefixed key."""
@ -33,13 +39,18 @@ def is_user_key_cache_key(key: str) -> bool:
_PIPELINED_SET_OPTIONS: Final = frozenset(("ttl",))
_TTL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
def _ttl(kwargs: Mapping[str, object]) -> float | None:
return _TTL.validate_python(kwargs.get("ttl"))
class UserApiKeyCache(DualCache):
"""
DualCache wrapper for UserAPIKeyAuth-like payloads.
DualCache for UserAPIKeyAuth-like payloads.
Stores a Redis-safe JSON payload in BOTH in-memory and Redis to avoid
Stores a Redis-safe JSON payload in memory and, with ``enable_redis_auth_cache``, in Redis, to avoid
"memory returns BaseModel, Redis returns dict" format drift.
When ``model_type`` is provided:
@ -53,12 +64,12 @@ class UserApiKeyCache(DualCache):
When ``model_type`` is omitted, the interface behaves like ``DualCache``:
raw cached payload is returned (dict/str/etc.).
``async_set_cache_pipeline`` applies the same untyped Codec pass as omitting
``model_type`` on ``async_set_cache`` (so ``BaseModel`` rows are dumped before Redis).
User-key objects (see ``is_user_key_cache_key``) live in their own in-memory partition,
``key_object_cache``, so churn in the other management objects cannot evict them. Both
partitions share the same Redis backend and TTL settings.
``key_object_cache``, so churn in the other management objects cannot evict them.
Without Redis both partitions are this class's Python memory. ``attach_redis_cache`` moves the whole
keyspace, both memory partitions included, onto the native ``AuthCacheStore`` (``shared``); redis-py never
reads or writes it, and ``redis_cache`` stays None.
``get_cache`` / ``async_get_cache`` overloads and implementations must be contiguous
(no other methods in between) so mypy resolves ``@overload`` + implementation correctly.
@ -72,16 +83,15 @@ class UserApiKeyCache(DualCache):
default_redis_ttl: float | None = None,
key_object_in_memory_cache: InMemoryCache | None = None,
budget_windows_absent_in_memory_cache: InMemoryCache | None = None,
auth_cache: AuthCacheFactory = _connected_auth_cache,
) -> None:
super().__init__(
in_memory_cache=in_memory_cache,
redis_cache=redis_cache,
default_in_memory_ttl=default_in_memory_ttl,
default_redis_ttl=default_redis_ttl,
)
self.key_object_cache: Final = DualCache(
in_memory_cache=key_object_in_memory_cache or InMemoryCache(),
redis_cache=redis_cache,
default_in_memory_ttl=default_in_memory_ttl,
default_redis_ttl=default_redis_ttl,
)
@ -90,14 +100,33 @@ class UserApiKeyCache(DualCache):
or InMemoryCache(default_ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL),
default_in_memory_ttl=default_in_memory_ttl,
)
self._auth_cache: Final = auth_cache
self._redis_connection: RedisCache | None = None
self._shared: AuthCacheStore | None = None
self.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl)
def in_memory_cache_for(self, key: str) -> InMemoryCache:
return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache
@property
def shared(self) -> AuthCacheStore | None:
"""The shared store holding the keyspace once ``enable_redis_auth_cache`` attached Redis."""
return self._shared
def delete_from_memory(self, key: str) -> None:
self.in_memory_cache_for(key).delete_cache(key)
if key.startswith(_BUDGET_WINDOWS_ABSENT_PREFIX):
self.budget_windows_absent_cache.delete_cache(key)
@property
def has_redis(self) -> bool:
return self._shared is not None
@property
def redis_connection(self) -> RedisCache | None:
"""The Redis the auth cache was attached to, for features that share it under keys of their own."""
return self._redis_connection
def _settings(self) -> AuthCacheSettings:
return AuthCacheSettings(
memory_entries=self.in_memory_cache.max_size_in_memory, memory_ttl=self.default_in_memory_ttl
)
def _python_only(self, method: str) -> None:
if self._shared is not None:
raise NotImplementedError(f"UserApiKeyCache.{method} is not part of the shared auth cache keyspace")
def update_cache_ttl(self, default_in_memory_ttl: float | None, default_redis_ttl: float | None) -> None:
super().update_cache_ttl(default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=default_redis_ttl)
@ -107,17 +136,54 @@ class UserApiKeyCache(DualCache):
self.budget_windows_absent_cache.update_cache_ttl(
default_in_memory_ttl=default_in_memory_ttl, default_redis_ttl=None
)
if self._shared is not None:
self._shared.reconfigure(self._settings())
def update_in_memory_max_size(self, max_size: int | None) -> None:
super().update_in_memory_max_size(max_size)
self.key_object_cache.update_in_memory_max_size(max_size)
self.budget_windows_absent_cache.update_in_memory_max_size(max_size)
if self._shared is not None:
self._shared.reconfigure(self._settings())
def attach_redis_cache(
self, redis_cache: RedisCache | None = None, *, default_redis_ttl: float | None = None
) -> None:
super().attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl)
self.key_object_cache.attach_redis_cache(redis_cache, default_redis_ttl=default_redis_ttl)
if redis_cache is None or self._shared is not None:
return
if default_redis_ttl is not None:
self.update_cache_ttl(default_in_memory_ttl=None, default_redis_ttl=default_redis_ttl)
self._redis_connection = redis_cache
self._shared = self._auth_cache(redis_cache, self._settings())
def get_local(self, key: str) -> object:
"""This worker's memory copy alone."""
if self._shared is not None:
return self._shared.get_local(key)
return self._python_memory(key).get_cache(key=key)
def set_local(self, key: str, value: object, ttl: float | None) -> None:
"""A memory write alone, with ``default_in_memory_ttl`` when ``ttl`` is None."""
if self._shared is not None:
self._shared.set_local(key, value, ttl)
return
effective_ttl: Final = self.default_in_memory_ttl if ttl is None else ttl
if effective_ttl is None:
self._python_memory(key).set_cache(key=key, value=value)
return
self._python_memory(key).set_cache(key=key, value=value, ttl=effective_ttl)
def delete_local(self, key: str) -> None:
"""Drops this worker's memory copy alone, as a cross-worker eviction broadcast does."""
if key.startswith(_BUDGET_WINDOWS_ABSENT_PREFIX):
self.budget_windows_absent_cache.delete_cache(key)
if self._shared is not None:
self._shared.delete_local(key)
return
self._python_memory(key).delete_cache(key)
def _python_memory(self, key: str) -> InMemoryCache:
return self.key_object_cache.in_memory_cache if is_user_key_cache_key(key) else self.in_memory_cache
@overload
def get_cache(
@ -150,24 +216,17 @@ class UserApiKeyCache(DualCache):
) -> object:
if model_type is None and "model_type" in kwargs:
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
shared: Final = self._shared
cached: Final = (
self.key_object_cache.get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs)
self._get_shared_sync(shared, key, local_only)
if shared is not None
else self.key_object_cache.get_cache(
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
)
if is_user_key_cache_key(key)
else super().get_cache(key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs)
)
if model_type is None:
return cached
if cached is None:
return None
decoded: Final = CacheCodec.deserialize(cached, model_type=model_type)
if decoded is None:
verbose_proxy_logger.error(
"UserApiKeyCache.get_cache failed to deserialize cached value for key=%r model_type=%s",
key,
getattr(model_type, "__name__", str(model_type)),
)
return None
return decoded
return _decoded(cached, key, model_type, "get_cache")
@overload
async def async_get_cache(
@ -200,8 +259,11 @@ class UserApiKeyCache(DualCache):
) -> object:
if model_type is None and "model_type" in kwargs:
model_type = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
shared: Final = self._shared
cached: Final = (
await self.key_object_cache.async_get_cache(
await self._get_shared(shared, key, local_only)
if shared is not None
else await self.key_object_cache.async_get_cache(
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
)
if is_user_key_cache_key(key)
@ -209,28 +271,44 @@ class UserApiKeyCache(DualCache):
key=key, parent_otel_span=parent_otel_span, local_only=local_only, **kwargs
)
)
if model_type is None:
return cached
if cached is None:
return None
decoded: Final = CacheCodec.deserialize(cached, model_type=model_type)
if decoded is None:
verbose_proxy_logger.error(
"UserApiKeyCache.async_get_cache failed to deserialize cached value for key=%r model_type=%s",
key,
getattr(model_type, "__name__", str(model_type)),
)
return None
return decoded
return _decoded(cached, key, model_type, "async_get_cache")
def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
@staticmethod
def _get_shared_sync(shared: AuthCacheStore, key: str, local_only: bool) -> object:
try:
return shared.get_local(key) if local_only else shared.get_sync(key)
except Exception as e: # noqa: BLE001 # DualCache.get_cache reads a failure as a miss
verbose_proxy_logger.error("UserApiKeyCache.get_cache failed for key=%r: %s", key, e)
return None
@staticmethod
async def _get_shared(shared: AuthCacheStore, key: str, local_only: bool) -> object:
try:
return shared.get_local(key) if local_only else await shared.get(key)
except Exception as e: # noqa: BLE001 # DualCache.async_get_cache reads a failure as a miss
verbose_proxy_logger.error("UserApiKeyCache.async_get_cache failed for key=%r: %s", key, e)
return None
def set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
payload: Final[object] = CacheCodec.serialize(value, model_type=model_type)
shared: Final = self._shared
if shared is not None:
if key is None:
return None
try:
if local_only:
shared.set_local(key, payload, _ttl(kwargs))
else:
shared.set_sync(key, payload, _ttl(kwargs))
except Exception as e: # noqa: BLE001 # DualCache.set_cache swallows a failed write
verbose_proxy_logger.error("UserApiKeyCache.set_cache failed for key=%r: %s", key, e)
return None
if key is not None and is_user_key_cache_key(key):
return self.key_object_cache.set_cache(key=key, value=payload, local_only=local_only, **kwargs)
return super().set_cache(key=key, value=payload, local_only=local_only, **kwargs)
async def async_set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
async def async_set_cache(self, key: str | None, value: CacheValue, local_only: bool = False, **kwargs: object):
"""Inside a request the Redis SET rides the request's pipeline (memory is written at once); anywhere
else, or with options the pipeline does not carry, it goes to Redis directly as before."""
model_type: Final = cast(type[BaseModel] | None, kwargs.pop("model_type", None))
@ -242,6 +320,11 @@ class UserApiKeyCache(DualCache):
and kwargs.keys() <= _PIPELINED_SET_OPTIONS
and (ttl is None or isinstance(ttl, (int, float)))
)
shared: Final = self._shared
if shared is not None:
if key is not None:
await self._set_shared(shared, key, payload, local_only=local_only, pipelined=pipelined, kwargs=kwargs)
return None
if key is not None and is_user_key_cache_key(key):
if pipelined and await self.key_object_cache.async_set_cache_pre_call(key, payload, ttl) is not None:
return None
@ -250,19 +333,46 @@ class UserApiKeyCache(DualCache):
return None
return await super().async_set_cache(key=key, value=payload, local_only=local_only, **kwargs)
@staticmethod
async def _set_shared(
shared: AuthCacheStore,
key: str,
payload: object,
*,
local_only: bool,
pipelined: bool,
kwargs: Mapping[str, object],
) -> None:
try:
ttl: Final = _ttl(kwargs)
if local_only:
shared.set_local(key, payload, ttl)
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)
def delete_cache(self, key: str) -> None:
if self._shared is not None:
self._shared.delete_sync((key,))
return
if is_user_key_cache_key(key):
self.key_object_cache.delete_cache(key)
return
super().delete_cache(key)
async def async_delete_cache(self, key: str) -> None:
if self._shared is not None:
await self._shared.delete((key,))
return
if is_user_key_cache_key(key):
await self.key_object_cache.async_delete_cache(key)
return
await super().async_delete_cache(key)
async def async_delete_cache_pre_call(self, key: str) -> BatchResult[None] | None:
async def async_delete_cache_pre_call(self, key: str) -> Awaitable[None] | None:
if self._shared is not None:
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)
@ -274,6 +384,10 @@ class UserApiKeyCache(DualCache):
Both partitions are cleared even when one raises, because a caller
batching these has already committed the rows they cache.
"""
if self._shared is not None:
if keys:
await self._shared.delete(keys)
return
key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key))
other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key))
outcomes: Final = await asyncio.gather(
@ -286,9 +400,12 @@ class UserApiKeyCache(DualCache):
raise failed[0]
def flush_cache(self) -> None:
self.budget_windows_absent_cache.flush_cache()
if self._shared is not None:
self._shared.flush_local()
return
super().flush_cache()
self.key_object_cache.in_memory_cache.flush_cache()
self.budget_windows_absent_cache.flush_cache()
async def async_set_cache_pipeline(
self, cache_list: Sequence[tuple[str, object]], local_only: bool = False, **kwargs: object
@ -298,6 +415,10 @@ class UserApiKeyCache(DualCache):
``model_type``: ``BaseModel`` values become JSON-safe dicts; dicts/scalars unchanged.
"""
normalized: Final = tuple((key, CacheCodec.serialize(value, model_type=None)) for key, value in cache_list)
shared: Final = self._shared
if shared is not None:
await self._set_many_shared(shared, normalized, local_only=local_only, ttl=_ttl(kwargs))
return
key_object_entries: Final = tuple(entry for entry in normalized if is_user_key_cache_key(entry[0]))
other_entries: Final = tuple(entry for entry in normalized if not is_user_key_cache_key(entry[0]))
if key_object_entries:
@ -307,6 +428,67 @@ class UserApiKeyCache(DualCache):
if other_entries:
await super().async_set_cache_pipeline(cache_list=other_entries, local_only=local_only, **kwargs)
@staticmethod
async def _set_many_shared(
shared: AuthCacheStore, entries: Sequence[tuple[str, object]], *, local_only: bool, ttl: float | None
) -> None:
try:
if not local_only:
await shared.set_many(entries, ttl)
return
for key, value in entries:
shared.set_local(key, value, ttl)
except Exception as e: # noqa: BLE001 # DualCache.async_set_cache_pipeline swallows a failed write
verbose_proxy_logger.error("UserApiKeyCache.async_set_cache_pipeline failed: %s", e)
async def increment_if_present(self, key: str, amount: float) -> float | None:
"""``INCRBYFLOAT`` of an entry Redis still holds; None without Redis or when Redis lacks the entry. Raises
when Redis failed."""
if self._shared is None:
return None
return await self._shared.increment_if_present(key, amount)
async def async_batch_get_cache(
self,
keys: list[str],
parent_otel_span: Span | None = None,
local_only: bool = False,
throttle_redis: bool = True,
**kwargs: object, # kwargs-ok: forward the base cache options in Python-only mode
) -> list[object | None] | None:
self._python_only("async_batch_get_cache")
return await super().async_batch_get_cache(
keys, parent_otel_span=parent_otel_span, local_only=local_only, throttle_redis=throttle_redis, **kwargs
)
async def async_increment_cache(
self,
key: str,
value: float,
parent_otel_span: Span | None = None,
local_only: bool = False,
refresh_ttl: bool = False,
**kwargs: object, # kwargs-ok: forward the base cache options in Python-only mode
) -> float | None:
self._python_only("async_increment_cache")
return await super().async_increment_cache(
key, value, parent_otel_span=parent_otel_span, local_only=local_only, refresh_ttl=refresh_ttl, **kwargs
)
def _decoded(cached: object, key: str, model_type: type[BaseModel] | None, method: str) -> object:
if model_type is None or cached is None:
return cached
decoded: Final = CacheCodec.deserialize(cached, model_type=model_type)
if decoded is None:
verbose_proxy_logger.error(
"UserApiKeyCache.%s failed to deserialize cached value for key=%r model_type=%s",
method,
key,
getattr(model_type, "__name__", str(model_type)),
)
return decoded
#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row,
#: so a human without an entitlement costs no DB read per request. Lives beside the key builder
@ -360,25 +542,10 @@ def model_access_group_registry_cache_key() -> str:
return "model_access_group_registry"
def model_access_group_spend_counter_key(access_group_name: str) -> str:
"""Spend counter key for one model access group; shared so its four owners cannot drift.
The reservation path writes it up front, the cost callback writes it after the call, auth
reads it to enforce ``max_budget``, and the reset job clears it on rollover. A copy that
drifts in any one of them silently resets or reads a counter nobody else touches, which shows
up as a budget that never trips or never resets.
"""
return f"spend:model_access_group:{access_group_name}"
def project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
def project_spend_counter_key(project_id: str) -> str:
return f"spend:project:{project_id}"
#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds
#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch.
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__"

View file

@ -20,10 +20,10 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.constants import (
GITHUB_COPILOT_DEFAULT_HOST,
GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS,
@ -203,7 +203,7 @@ async def _try_cache_set(token_cache: "RedisCache", key: str, value: str, nx: bo
@with_service_target("user_provider_connections")
async def invalidate_user_provider_credential_cache(
cache: DualCache,
cache: "UserApiKeyCache",
user_id: str,
credential_name: str,
) -> bool:
@ -211,7 +211,7 @@ async def invalidate_user_provider_credential_cache(
before the disconnect must not fill the token back in behind it (fills are
set-if-absent). Returns False when Redis is attached and the write fails, so
the disconnect can refuse to drop the row while a stale token might linger."""
token_cache: Final = cache.redis_cache
token_cache: Final = cache.redis_connection
if token_cache is None:
return True
try:
@ -235,13 +235,13 @@ async def invalidate_user_provider_credential_cache(
@with_service_target("user_provider_connections")
async def drop_user_provider_credential_cache(
cache: DualCache,
cache: "UserApiKeyCache",
user_id: str,
credential_name: str,
) -> None:
"""Connect path: delete the key outright so the next read fills fresh (a
tombstone would linger for the TTL and hide the new connection)."""
token_cache: Final = cache.redis_cache
token_cache: Final = cache.redis_connection
if token_cache is None:
return
try:
@ -252,7 +252,7 @@ async def drop_user_provider_credential_cache(
@with_service_target("user_provider_connections")
async def set_user_provider_credential_cache(
cache: DualCache,
cache: "UserApiKeyCache",
user_id: str,
credential_name: str,
payload: GithubCopilotUserConnectionPayload,
@ -262,7 +262,7 @@ async def set_user_provider_credential_cache(
not-connected marker. The overwrite is verified by decoding the key back
(async_set_cache swallows write errors); on a mismatch the key is deleted,
and False means a stale entry survived both attempts."""
token_cache: Final = cache.redis_cache
token_cache: Final = cache.redis_connection
if token_cache is None:
return True
key: Final = _cache_key(user_id, credential_name)
@ -296,7 +296,7 @@ async def set_user_provider_credential_cache(
@with_service_target("user_provider_connections")
async def aget_user_provider_credentials(
prisma_client: "PrismaClient",
cache: DualCache,
cache: "UserApiKeyCache",
user_id: str,
credential_names: Sequence[str],
) -> Mapping[str, Mapping[str, str]]:
@ -306,7 +306,7 @@ async def aget_user_provider_credentials(
plaintext connection data only ever lives in the returned mapping. Redis is the only
cache layer used: without it there is no shared invalidation, so a worker
serving a stale local entry after a disconnect is worse than a DB read."""
token_cache: Final = cache.redis_cache
token_cache: Final = cache.redis_connection
names: Final = tuple(dict.fromkeys(credential_names))
if not names:
return {}

View file

@ -26,10 +26,10 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate
from litellm.proxy.spend_tracking.spend_counter_batch import (
SPEND_COUNTERS_TARGET,
read_batched_spend_counter,
record_spend_counter_value,
)
from litellm.redis.stores.spend_counters import END_USER_COUNTER_PREFIX, SPEND_COUNTERS_TARGET, SpendCounterStore
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
@ -47,7 +47,6 @@ from litellm.repositories.verification_token_repository import (
if TYPE_CHECKING:
from prisma.types import LiteLLM_EndUserTableWhereUniqueInput
from litellm.caching.dual_cache import DualCache
from litellm.proxy.utils import PrismaClient
@ -58,8 +57,6 @@ _WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType(
}
)
END_USER_COUNTER_PREFIX: Final = "spend:end_user:"
_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{
"Key": "api_key",
@ -114,13 +111,11 @@ class SpendCounterReseed:
@staticmethod
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None:
async def increment_in_memory(spend_counter_store: SpendCounterStore, counter_key: str, increment: float) -> float:
"""Apply local deltas after an in-flight reseed establishes the spend balance."""
lock: Final = await SpendCounterReseed._get_lock(counter_key)
async with lock:
return await spend_counter_cache.async_increment_cache(
key=counter_key, value=increment, local_only=True, refresh_ttl=True
)
return spend_counter_store.increment_local(counter_key, increment)
@staticmethod
async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
@ -222,7 +217,7 @@ class SpendCounterReseed:
@with_service_target(SPEND_COUNTERS_TARGET)
async def coalesced(
prisma_client: Optional["PrismaClient"],
spend_counter_cache: "DualCache",
spend_counter_store: SpendCounterStore,
counter_key: str,
require_cache_warm: bool = False,
) -> float | None:
@ -236,24 +231,9 @@ class SpendCounterReseed:
"""
lock: Final = await SpendCounterReseed._get_lock(counter_key)
async with lock:
batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
if batched is not None and batched[0] is not None:
return batched[0]
# Re-check after acquiring the lock. Skip in-memory on a clean
# Redis miss - in-memory is per-pod-stale.
redis_clean_miss = batched is not None
if spend_counter_cache.redis_cache is not None and not redis_clean_miss:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
cached: Final = await SpendCounterReseed._read_cached(spend_counter_store, counter_key)
if cached is not None:
return cached
db_spend: Final = await SpendCounterReseed.from_db(prisma_client, counter_key)
if db_spend is None:
@ -264,29 +244,8 @@ class SpendCounterReseed:
# Redis key with db_spend; concurrent seeders read the winner's
# value. INCRBYFLOAT-of-db_spend from N pods would multiply the
# counter (N x db_spend) and trigger spurious budget alerts.
current_value: float = float(db_spend)
try:
if spend_counter_cache.redis_cache is not None:
seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=db_spend,
nx=True,
)
if seeded:
current_value = float(db_spend)
else:
cached: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
current_value = float(cached) if cached is not None else float(db_spend)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
record_spend_counter_value(counter_key, current_value)
else:
cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
seeded_spend: Final = max(db_spend, float(cached_spend)) if cached_spend is not None else db_spend
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
return seeded_spend
seeded: Final = await spend_counter_store.seed(counter_key, db_spend, increment_if_lost=False)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.coalesced: failed to warm counter %s",
@ -294,7 +253,20 @@ class SpendCounterReseed:
)
if require_cache_warm:
raise
return current_value
return float(db_spend)
if spend_counter_store.has_redis:
record_spend_counter_value(counter_key, float(seeded))
return float(seeded)
@staticmethod
async def _read_cached(spend_counter_store: SpendCounterStore, counter_key: str) -> float | None:
"""Re-check after acquiring the lock. A clean Redis miss, in the request's MGET or a GET, skips memory,
which only holds this pod's writes."""
batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
if batched is not None:
return batched[0]
value, _ = await spend_counter_store.read(counter_key)
return None if value is None else float(value)
@staticmethod
async def window_from_table(
@ -425,7 +397,7 @@ class SpendCounterReseed:
@with_service_target(SPEND_COUNTERS_TARGET)
async def coalesced_window(
prisma_client: Optional["PrismaClient"],
spend_counter_cache: "DualCache",
spend_counter_store: SpendCounterStore,
counter_key: str,
entity_type: str,
entity_id: str,
@ -434,22 +406,9 @@ class SpendCounterReseed:
) -> float | None:
lock: Final = await SpendCounterReseed._get_lock(counter_key)
async with lock:
batched: Final = await SpendCounterReseed._read_active_batch(counter_key)
if batched is not None and batched[0] is not None:
return batched[0]
redis_clean_miss = batched is not None
if spend_counter_cache.redis_cache is not None and not redis_clean_miss:
try:
val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if val is not None:
return float(val)
redis_clean_miss = True
except Exception:
pass
if not redis_clean_miss:
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
if val is not None:
return float(val)
cached: Final = await SpendCounterReseed._read_cached(spend_counter_store, counter_key)
if cached is not None:
return cached
window_spend: Final = await SpendCounterReseed.window_from_db(
prisma_client=prisma_client,
@ -461,39 +420,13 @@ class SpendCounterReseed:
if window_spend is None:
return None
try:
if spend_counter_cache.redis_cache is not None:
seeded: Final = await spend_counter_cache.redis_cache.async_set_cache(
key=counter_key,
value=window_spend,
nx=True,
)
if seeded:
current_value = window_spend
else:
current_cached_value = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
if current_cached_value is None:
current_value = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=window_spend,
)
else:
current_value = float(current_cached_value)
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
record_spend_counter_value(counter_key, float(current_value))
else:
cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
seeded_spend: Final = (
max(window_spend, float(cached_spend)) if cached_spend is not None else window_spend
)
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend)
return seeded_spend
seeded: Final = await spend_counter_store.seed(counter_key, window_spend, increment_if_lost=True)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.coalesced_window: failed to warm counter %s",
counter_key,
)
raise
return current_value
if spend_counter_store.has_redis:
record_spend_counter_value(counter_key, float(seeded))
return float(seeded)

View file

@ -24,6 +24,7 @@ from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS
from litellm.proxy._types import UserAPIKeyAuth
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET
from litellm.types.llms.base import LiteLLMBaseModel
if TYPE_CHECKING:
@ -224,7 +225,7 @@ class BatchEnqueuedTokenStore:
def _record_key(batch_id: str) -> str:
return f"batch_enqueued_token_reservation:{batch_id}"
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def reserve(
self,
tokens: int,
@ -329,7 +330,7 @@ class BatchEnqueuedTokenStore:
tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token, reserved_at_monotonic=started
)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def refund(
self,
reservation: BatchEnqueuedTokenReservation,
@ -368,7 +369,7 @@ class BatchEnqueuedTokenStore:
"Redis enqueued-token refund failed; leaked increments expire with the TTL: %s", str(e)
)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def save_reservation(
self,
batch_id: str,
@ -401,7 +402,7 @@ class BatchEnqueuedTokenStore:
local_only=True,
)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def pop_reservation(
self,
batch_id: str,

View file

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

View file

@ -21,6 +21,7 @@ from litellm.proxy.hooks.rate_limiter_utils import (
convert_priority_to_percent,
resolve_llm_provider_for_rate_limit,
)
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral, LLMResponseTypes
from litellm.utils import get_utc_datetime
@ -38,7 +39,7 @@ class DynamicRateLimiterCache:
self.ttl = 60 # 1 min ttl
self.time_fn = time_fn
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_get_cache(self, model: str) -> int | None:
dt: Final = self.time_fn()
current_minute: Final = dt.strftime("%H-%M")
@ -49,7 +50,7 @@ class DynamicRateLimiterCache:
response = len(_response)
return response
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_set_cache_sadd(self, model: str, value: list):
"""
Add value to set.
@ -85,7 +86,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
def update_variables(self, llm_router: Router) -> None:
self.llm_router = llm_router
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def check_available_usage(
self, model: str, priority: str | None = None
) -> tuple[int | None, int | None, int | None, int | None, int | None]:
@ -183,7 +184,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
)
return None, None, None, None, None
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -239,7 +240,7 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
)
return None
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> LLMResponseTypes | None:

View file

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

View file

@ -23,6 +23,7 @@ from litellm.proxy.auth.budget_throttle import throttled_limit
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit, socket_session_usage
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET
from litellm.types.utils import CallTypes, Usage
if TYPE_CHECKING:
@ -88,7 +89,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
except Exception:
pass
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def check_key_in_limits(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -230,7 +231,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
descriptor_key=descriptor_key,
)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def get_all_cache_objects(
self,
current_global_requests: str | None,
@ -273,7 +274,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
request_count_end_user_id=results[5],
)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -517,7 +518,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
) # don't block execution for cache updates
)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None:
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
@ -718,7 +719,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
except Exception as e:
self.print_verbose(e)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
try:
self.print_verbose("Inside Max Parallel Request Failure Hook")
@ -789,7 +790,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
except Exception as e:
verbose_proxy_logger.exception("Inside Parallel Request Limiter: An exception occurred - %s", e)
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_release_max_parallel_requests_on_disconnect(self, user_api_key_dict: UserAPIKeyAuth) -> None:
"""Release the slot ``async_pre_call_hook`` took for a connection whose per-turn callbacks released none."""
from litellm.proxy.proxy_server import (
@ -851,7 +852,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger):
verbose_proxy_logger.debug("Parallel Request Limiter: Error getting user object", str(e))
return None
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def async_post_call_success_hook(self, data: dict, user_api_key_dict: UserAPIKeyAuth, response) -> None:
"""
Retrieve the key's remaining rate limits.

File diff suppressed because it is too large Load diff

View file

@ -1294,7 +1294,7 @@ async def _update_database_and_spend_counters(
"""The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then
spans the database write and the counter update, so the post-call counters are read with a single MGET after the
write and their increments leave in a single pipeline."""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
from litellm.proxy.spend_tracking.budget_reservation import get_reserved_counter_keys
if budget_reservation is not None:
@ -1313,7 +1313,7 @@ async def _update_database_and_spend_counters(
model_access_groups=model_access_groups,
project_id=project_id,
)
with spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys):
with spend_counter_batch_scope(spend_counter_store, counter_keys=counter_keys):
return await _update_database_and_spend_counters_in_batch(
proxy_logging_obj=proxy_logging_obj,
spend_writer=spend_writer,

View file

@ -7,22 +7,19 @@ this hook manages:
2. Checking incoming requests for existing routing overrides
3. Applying sticky routing so all subsequent requests in a session go to the same model
Works across multiple proxy instances via DualCache (in-memory + Redis).
Works across multiple proxy instances via the native sensitive route pin store (in-memory + Redis).
"""
import logging
import os
from typing import TYPE_CHECKING, Any, Final
from pydantic import TypeAdapter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import log_redis_failure
from litellm.integrations.custom_guardrail import get_session_id_from_request_data
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.redis.stores.connection import connect
from litellm.redis.stores.sensitive_route_pins import SensitiveRoutePinStore
if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
@ -32,9 +29,7 @@ else:
InternalUsageCache = Any
SENSITIVE_ROUTING_CACHE_PREFIX: Final = "sensitive_route"
DEFAULT_SENSITIVE_ROUTING_TTL: Final = 3600
_CACHED_ROUTE: Final = TypeAdapter(object)
class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
@ -46,8 +41,10 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
model field if a routing override exists for the session.
"""
def __init__(self, internal_usage_cache: InternalUsageCache):
self.internal_usage_cache = internal_usage_cache
def __init__(self, internal_usage_cache: InternalUsageCache, route_pins: SensitiveRoutePinStore | None = None):
self.route_pins: Final = route_pins or SensitiveRoutePinStore(
connect(internal_usage_cache.dual_cache.redis_cache).sensitive_route_pins()
)
self.ttl = int(
os.getenv(
"LITELLM_SENSITIVE_ROUTING_TTL",
@ -55,9 +52,6 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
)
)
def _make_cache_key(self, session_id: str, tenant: str) -> str:
return f"{{{SENSITIVE_ROUTING_CACHE_PREFIX}:{tenant}:{session_id}}}:model"
@staticmethod
def _resolve_tenant(user_api_key_dict: UserAPIKeyAuth | None) -> str:
"""
@ -83,47 +77,10 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
]
return "|".join(principal) if principal else "default"
@with_service_target("sensitive_route_pins")
async def _get_routed_model(self, session_id: str, user_api_key_dict: UserAPIKeyAuth | None) -> str | None:
"""Get the model this session should be routed to, if any."""
cache_key: Final = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict))
return await self.route_pins.get(self._resolve_tenant(user_api_key_dict), session_id, self.ttl)
if self.internal_usage_cache.dual_cache.redis_cache is not None:
try:
redis_route: Final = _CACHED_ROUTE.validate_python(
await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache(key=cache_key)
)
if redis_route is not None:
routed_model: Final = str(redis_route)
remaining_ttl = await self.internal_usage_cache.dual_cache.redis_cache.async_get_ttl(key=cache_key)
await self.internal_usage_cache.async_set_cache(
key=cache_key,
value=routed_model,
ttl=remaining_ttl if remaining_ttl is not None else self.ttl,
litellm_parent_otel_span=None,
local_only=True,
)
return routed_model
except Exception as e:
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"SensitiveDataRoutingHandler: Redis GET failed, falling back to in-memory",
e,
)
local_route: Final = _CACHED_ROUTE.validate_python(
await self.internal_usage_cache.async_get_cache(
key=cache_key,
litellm_parent_otel_span=None,
local_only=True,
)
)
if local_route is not None:
return str(local_route)
return None
@with_service_target("sensitive_route_pins")
async def set_session_routing(
self,
session_id: str,
@ -138,8 +95,6 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
route the session to a specific model. The override is scoped to the
requesting principal so sessions from different tenants cannot collide.
"""
cache_key: Final = self._make_cache_key(session_id, self._resolve_tenant(user_api_key_dict))
verbose_proxy_logger.info(
"SensitiveDataRoutingHandler: Setting session routing session_id=%s model=%s guardrail=%s ttl=%s",
session_id,
@ -147,31 +102,8 @@ class _PROXY_SensitiveDataRoutingHandler(CustomLogger):
guardrail_name,
self.ttl,
)
await self.route_pins.set(self._resolve_tenant(user_api_key_dict), session_id, model, self.ttl)
if self.internal_usage_cache.dual_cache.redis_cache is not None:
try:
await self.internal_usage_cache.dual_cache.redis_cache.async_set_cache(
key=cache_key,
value=model,
ttl=self.ttl,
)
except Exception as e:
log_redis_failure(
verbose_proxy_logger,
logging.WARNING,
"SensitiveDataRoutingHandler: Redis SET failed, falling back to in-memory",
e,
)
await self.internal_usage_cache.async_set_cache(
key=cache_key,
value=model,
ttl=self.ttl,
litellm_parent_otel_span=None,
local_only=True,
)
@with_service_target("sensitive_route_pins")
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,

View file

@ -86,6 +86,9 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa:
from litellm.proxy.management_helpers.team_roster_sync import user_tx_db
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import handle_exception_on_proxy, hash_password
from litellm.redis.stores.spend_counters import (
user_counter,
)
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import TableActions
from litellm.repositories.table_repositories import (
@ -1477,7 +1480,7 @@ async def _invalidate_user_spend_counter_if_changed(
if non_default_values.get("spend") is not None:
from litellm.proxy.proxy_server import invalidate_spend_counter
await invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}")
await invalidate_spend_counter(counter_key=user_counter(str(non_default_values["user_id"])))
def _clears_object_permission(user_request: UpdateUserRequest) -> bool:

View file

@ -30,7 +30,7 @@ from pydantic import TypeAdapter
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._internal_context import service_target, with_service_target
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.caching.dual_cache import DualCache
@ -157,7 +157,6 @@ from litellm.proxy.mcp_server.outbound_credentials.sso_assertion_store import (
)
from litellm.proxy.search_endpoints.search_tool_registry import rotate_search_tools_master_key
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # legacy module exports
_is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
is_master_key,
@ -171,6 +170,11 @@ from litellm.proxy.utils import (
is_valid_api_key,
require_enterprise_license,
)
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
key_counter,
key_window_counter,
)
from litellm.repositories.base_repository import BaseRepository, record_to_dict
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
@ -4263,21 +4267,18 @@ async def update_key_fn(
)
if data.spend is not None:
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
counter_key: Final = f"spend:key:{hash_token_if_needed(key)}"
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
with service_target(SPEND_COUNTERS_TARGET):
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=data.spend, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis after key spend update: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
counter_key: Final = key_counter(hash_token_if_needed(key))
try:
await spend_counter_store.set(counter_key, data.spend, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis after key spend update: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
asyncio.create_task(
KeyManagementEventHooks.async_key_updated_hook(
@ -4901,7 +4902,7 @@ async def _budget_window_usage(
if not isinstance(duration, str) or not duration:
return None
spend: Final = await get_current_spend(
counter_key=f"spend:key:{api_key_hash}:window:{duration}",
counter_key=key_window_counter(api_key_hash, duration),
fallback_spend=0.0,
max_budget=_window_max_budget(window),
window_entity_type="Key",
@ -6822,19 +6823,16 @@ async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: f
Without the broadcast, a worker that already cached the pre-reset key object
or floor marker keeps enforcing against it until its own TTL expires.
"""
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache
from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache, spend_counter_store
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=value, ttl=60)
if spend_counter_cache.redis_cache is not None:
try:
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=value, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. "
"Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
try:
await spend_counter_store.set(counter_key, value, ttl=60)
except Exception as redis_err:
verbose_proxy_logger.warning(
"Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.",
counter_key,
redis_err,
)
floor_key: Final = f"spend_db_floor:{counter_key}"
spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS)
@ -6944,7 +6942,7 @@ async def _reset_key_budget_windows(
for window in reset_windows:
duration = window.get("budget_duration")
if isinstance(duration, str) and duration:
counter_key = f"spend:key:{hashed_api_key}:window:{duration}"
counter_key = key_window_counter(hashed_api_key, duration)
await _set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0)
@ -7018,7 +7016,7 @@ async def reset_key_spend_fn(
# budget_limits windows, so get_current_spend() returns the correct
# amount for every enforcement check immediately instead of the stale
# pre-reset value.
_counter_key: Final = f"spend:key:{hashed_api_key}"
_counter_key: Final = key_counter(hashed_api_key)
await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to)
await _reset_key_budget_windows(
prisma_client=prisma_client,

View file

@ -31,6 +31,9 @@ from litellm.proxy.management_endpoints.common_daily_activity import (
get_daily_activity,
)
from litellm.proxy.management_helpers.utils import handle_budget_for_entity
from litellm.redis.stores.spend_counters import (
tag_counter,
)
from litellm.repositories.daily_activity_buckets import api_key_from_bucket
from litellm.repositories.model_repository import ModelRepository
from litellm.repositories.table_repositories import (
@ -319,7 +322,7 @@ async def new_tag(
)
await _evict_tag_cache_keys((tag_cache_key(tag.name), tag_registry_cache_key()))
await invalidate_spend_counter(counter_key=f"spend:tag:{tag.name}")
await invalidate_spend_counter(counter_key=tag_counter(tag.name))
# Update models with new tag
if tag.models:

View file

@ -121,7 +121,7 @@ async def _post_token_endpoint(url: str, form: dict[str, str], headers: dict[str
def _redis_cache_is_available() -> bool:
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415
return user_api_key_cache.redis_cache is not None
return user_api_key_cache.has_redis
def _runtime_backend_and_coordinator() -> tuple[TokenCacheBackend | None, RefreshCoordinator | None, bool]:

View file

@ -28,7 +28,7 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator imp
def runtime_refresh_coordinator() -> RefreshCoordinator | None:
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # runtime global
redis_cache: Final = user_api_key_cache.redis_cache
redis_cache: Final = user_api_key_cache.redis_connection
if redis_cache is None:
return None
# The Redis client from init_async_client() is only partially typed; the lock validates every

View file

@ -16,6 +16,7 @@ from litellm.constants import (
MAX_FILE_DOWNLOADS_PER_MINUTE_KEY,
)
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET
if TYPE_CHECKING:
from litellm.proxy.utils import InternalUsageCache
@ -149,7 +150,7 @@ async def _increment_all_or_none(
return exceeded
@with_service_target("rate_limits")
@with_service_target(RATE_LIMITS_TARGET)
async def consume_file_usage(
cache: "InternalUsageCache",
limits: tuple[ScopedFileUsageLimit, ...],

View file

@ -206,7 +206,6 @@ from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
tuning_limit_violation,
)
from litellm.rust_bridge.catalog import Target
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import (
PRICING_OVERRIDES_KEY,
CallTypes,
@ -333,10 +332,7 @@ from litellm._logging import (
verbose_router_logger,
)
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
drain_post_call_redis_batches,
)
from litellm.caching.redis_batch import drain_post_call_redis_batches
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.constants import (
@ -620,9 +616,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
get_management_object_ttl,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.common_utils.validation_error_body import public_validation_errors
@ -887,6 +881,23 @@ from litellm.proxy.telemetry.middleware import TelemetryMiddleware
from litellm.proxy.telemetry.runtime import TelemetryRuntime
from litellm.proxy.telemetry.settings import describe_errors as describe_telemetry_errors
from litellm.proxy.telemetry.settings import load_settings as load_telemetry_settings
from litellm.redis.stores.connection import connect
from litellm.redis.stores.spend_counters import (
SPEND_COUNTERS_TARGET,
SpendCounterStore,
end_user_counter,
is_unavailable,
key_counter,
key_window_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
team_window_counter,
user_counter,
)
from litellm.router import router_arguments
from litellm.router_config import build_router_config
@ -973,7 +984,6 @@ from litellm.proxy.spend_tracking.spend_capture_rate import (
run_scheduled_spend_capture_rate_check,
)
from litellm.proxy.spend_tracking.spend_counter_batch import (
SPEND_COUNTERS_TARGET,
PendingSpendIncrement,
active_spend_counter_batch,
forget_spend_counter,
@ -3148,6 +3158,7 @@ user_api_key_cache: UserApiKeyCache = UserApiKeyCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
spend_counter_cache: Final = DualCache(default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value)
spend_counter_store: Final = SpendCounterStore(connect(None).spend_counters())
cli_sso_session_cache: Final = DualCache(default_in_memory_ttl=CLI_SSO_SESSION_TTL_SECONDS)
saml_sso_cache: Final = DualCache(default_in_memory_ttl=SAML_AUTHN_REQUEST_TTL_SECONDS)
model_max_budget_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=spend_counter_cache)
@ -3458,28 +3469,18 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
in-memory copy is guarded by a read-compare-write with no await in between,
so it is atomic within the worker.
"""
cached: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
needs_update = True
if cached is not None:
try:
needs_update = float(cached) < db_spend
except (TypeError, ValueError):
needs_update = True
if needs_update:
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
if spend_counter_cache.redis_cache is not None:
forget_spend_counter(counter_key)
try:
repaired: Final = await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend)
except Exception:
verbose_proxy_logger.debug(
"Unable to repair stale spend counter %s in Redis",
counter_key,
exc_info=True,
)
return
if repaired is not None:
record_spend_counter_value(counter_key, repaired)
forget_spend_counter(counter_key)
try:
repaired: Final = await spend_counter_store.set_max(counter_key, db_spend)
except Exception:
verbose_proxy_logger.debug(
"Unable to repair stale spend counter %s in Redis",
counter_key,
exc_info=True,
)
return
if repaired is not None:
record_spend_counter_value(counter_key, repaired)
async def reseed_spend_counter_from_db(counter_key: str) -> bool:
@ -3576,19 +3577,8 @@ async def read_spend_counter_cache_value(counter_key: str) -> tuple[float | None
if batched is not None:
return batched
if spend_counter_cache.redis_cache is not None:
try:
redis_val: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key)
return (float(redis_val) if redis_val is not None else None), True
except Exception as e:
verbose_proxy_logger.debug(
"get_current_spend: Redis read failed for %s, falling back to in-memory: %s",
counter_key,
e,
)
in_memory_val: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
return (float(in_memory_val) if in_memory_val is not None else None), False
value, authoritative = await spend_counter_store.read(counter_key)
return (float(value) if value is not None else None), authoritative
async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float) -> tuple[float, bool]:
@ -3603,7 +3593,7 @@ async def _read_spend_counter_estimate(counter_key: str, fallback_spend: float)
# Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
db_spend: Final = await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
counter_key=counter_key,
)
if db_spend is not None:
@ -3630,7 +3620,7 @@ async def increment_spend_counters(
"""
Atomically increment spend counters for budget enforcement.
Uses spend_counter_cache (DualCache with Redis backend when available)
Uses spend_counter_store (Redis when available, this worker's memory behind it)
so counters are shared across all pods. Budget check functions read
from these counters via get_current_spend() (Redis-first).
@ -3638,7 +3628,7 @@ async def increment_spend_counters(
updated before the next request's auth check runs.
"""
with spend_counter_batch_scope(
spend_counter_cache.redis_cache,
spend_counter_store,
counter_keys=post_call_counter_keys(
token=token,
team_id=team_id,
@ -3805,7 +3795,11 @@ async def _increment_spend_counters_batched(
window_reset_at: Final = (
window.get("reset_at") if isinstance(window, dict) else getattr(window, "reset_at", None)
)
window_counter: Final = f"spend:{entity_type.value}:{entity_id}:window:{duration}"
window_counter: Final = (
key_window_counter(entity_id, str(duration))
if entity_type is Litellm_EntityType.KEY
else team_window_counter(entity_id, str(duration))
)
window_start: Final = get_budget_window_start(window)
pending_window: Final = (
await _prepare_window_spend_counter_increment(
@ -3868,7 +3862,7 @@ async def _increment_spend_counters_batched(
hashed_token: Final = (
hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token
)
key_counter_key: Final = f"spend:key:{hashed_token}"
key_counter_key: Final = key_counter(hashed_token)
key_pending: Final[tuple[PendingSpendIncrement, ...]] = (
()
if key_counter_key in reserved_counter_keys
@ -3884,7 +3878,7 @@ async def _increment_spend_counters_batched(
return key_pending + await _cached_window_increments(Litellm_EntityType.KEY, hashed_token, hashed_token)
async def _team_scope(scope_team_id: str) -> tuple[_SpendScopeItem, ...]:
team_counter_key: Final = f"spend:team:{scope_team_id}"
team_counter_key: Final = team_counter(scope_team_id)
team_pending: Final[tuple[PendingSpendIncrement, ...]] = (
()
if team_counter_key in reserved_counter_keys
@ -3904,7 +3898,7 @@ async def _increment_spend_counters_batched(
async def _team_member_scope(
scope_user_id: str, scope_team_id: str
) -> tuple[PendingSpendIncrement | BaseException, ...]:
team_member_counter_key: Final = f"spend:team_member:{scope_user_id}:{scope_team_id}"
team_member_counter_key: Final = team_member_counter(scope_user_id, scope_team_id)
if team_member_counter_key in reserved_counter_keys:
return ()
return (
@ -3916,7 +3910,7 @@ async def _increment_spend_counters_batched(
)
async def _user_scope(scope_user_id: str) -> tuple[PendingSpendIncrement | BaseException, ...]:
user_counter_key: Final = f"spend:user:{scope_user_id}"
user_counter_key: Final = user_counter(scope_user_id)
if user_counter_key in reserved_counter_keys:
return ()
return (
@ -4079,7 +4073,7 @@ async def _prepare_end_user_and_tag_spend_increments(
coro
for coro in (
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:end_user:{end_user_id}",
counter_key=end_user_counter(end_user_id),
source_cache_key=end_user_cache_key(end_user_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4088,7 +4082,7 @@ async def _prepare_end_user_and_tag_spend_increments(
else None,
*(
_prepare_unreserved_spend_counter_increment(
counter_key=f"spend:tag:{tag_name}",
counter_key=tag_counter(tag_name),
source_cache_key=tag_cache_key(tag_name),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4124,7 +4118,7 @@ async def _prepare_model_access_group_spend_increments(
results: Final = await asyncio.gather(
*(
_prepare_unreserved_spend_counter_increment(
counter_key=model_access_group_spend_counter_key(group),
counter_key=model_access_group_counter(group),
source_cache_key=model_access_group_cache_key(group),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4145,7 +4139,7 @@ async def _prepare_org_spend_increment(
return ()
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=f"spend:org:{org_id}",
counter_key=org_counter(org_id),
source_cache_key=[f"org_id:{org_id}:with_budget", f"org_id:{org_id}"],
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4162,7 +4156,7 @@ async def _prepare_project_spend_increment(
return ()
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=project_spend_counter_key(project_id),
counter_key=project_counter(project_id),
source_cache_key=project_cache_key(project_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
@ -4197,7 +4191,7 @@ async def _prepare_spend_counter_increment(
pipelined Redis call.
On first access per pod:
1. Check spend_counter_cache (in-memory -> Redis via DualCache)
1. Check spend_counter_store (Redis, then this worker's memory)
2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`.
Falls back to the cached object's `.spend` via user_api_key_cache
only if prisma is unavailable, since that value can lag the flusher.
@ -4302,7 +4296,7 @@ async def ensure_spend_counter_initialized(
# Shares the per-counter lock with get_current_spend.
db_spend: Final = await SpendCounterReseed.coalesced(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
counter_key=counter_key,
require_cache_warm=True,
)
@ -4343,7 +4337,7 @@ async def ensure_window_spend_counter_initialized(
window_spend: Final = await SpendCounterReseed.coalesced_window(
prisma_client=prisma_client,
spend_counter_cache=spend_counter_cache,
spend_counter_store=spend_counter_store,
counter_key=counter_key,
entity_type=entity_type,
entity_id=entity_id,
@ -4369,29 +4363,13 @@ async def _is_spend_counter_cache_warm(counter_key: str) -> bool:
batched_value, _ = batched
if batched_value is None:
return False
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=batched_value)
spend_counter_store.set_local(counter_key, batched_value)
return True
if spend_counter_cache.redis_cache is not None:
try:
current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache(
key=counter_key,
)
if current_value is None:
return False
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
return True
except Exception as e:
verbose_proxy_logger.debug(
"Unable to read Redis spend counter %s before initialization, falling back to in-memory: %s",
counter_key,
e,
)
return spend_counter_cache.in_memory_cache.get_cache(key=counter_key) is not None
current_value, from_redis = await spend_counter_store.read(counter_key)
if from_redis and current_value is not None:
spend_counter_store.set_local(counter_key, current_value)
return current_value is not None
async def increment_spend_counter(counter_key: str, increment: float):
@ -4403,10 +4381,10 @@ async def increment_spend_counter(counter_key: str, increment: float):
@with_service_target(SPEND_COUNTERS_TARGET)
async def refresh_spend_counter_ttl(counter_key: str) -> bool:
if spend_counter_cache.redis_cache is None:
if not spend_counter_store.has_redis:
return False
try:
return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key)
return await spend_counter_store.refresh_ttl(counter_key)
except Exception as e:
verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e)
return False
@ -4414,25 +4392,17 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool:
@with_service_target(SPEND_COUNTERS_TARGET)
async def increment_spend_counter_cache(counter_key: str, increment: float) -> float | None:
if spend_counter_cache.redis_cache is not None:
if spend_counter_store.has_redis:
try:
current_value: Final = await spend_counter_cache.redis_cache.async_increment(
key=counter_key,
value=increment,
refresh_ttl=True,
)
current_value: Final = await spend_counter_store.increment(counter_key, increment)
except Exception:
await invalidate_spend_counter(counter_key=counter_key)
forget_spend_counter(counter_key)
raise
spend_counter_cache.in_memory_cache.set_cache(
key=counter_key,
value=current_value,
)
record_spend_counter_value(counter_key, float(current_value))
return current_value
return await SpendCounterReseed.increment_in_memory(
spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment
spend_counter_store=spend_counter_store, counter_key=counter_key, increment=increment
)
@ -4442,10 +4412,8 @@ _increment_spend_counter_cache: Final = increment_spend_counter_cache
@with_service_target(SPEND_COUNTERS_TARGET)
async def invalidate_spend_counter(counter_key: str) -> None:
forget_spend_counter(counter_key)
spend_counter_cache.in_memory_cache.delete_cache(key=counter_key)
if spend_counter_cache.redis_cache is not None:
with contextlib.suppress(Exception):
await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key)
with contextlib.suppress(Exception):
await spend_counter_store.delete((counter_key,))
_invalidate_spend_counter: Final = invalidate_spend_counter
@ -4457,51 +4425,22 @@ async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncremen
try:
await increment_spend_counters_pipeline(pending=pending)
except Exception as e:
if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e):
if isinstance(e, RedisCircuitBreakerOpenError) or is_redis_timeout_failure(e) or is_unavailable(e):
return
raise
def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> bool:
"""Post-call increments ride the request's post-call pipeline with the other counters. Each counter's
new value lands in memory when the pipeline settles; a failed one is invalidated so no reader trusts a
"""Post-call increments ride the request's native post-call pipeline with the other counters. Each counter's
new value lands in memory when the pipeline settles; a failed one is dropped everywhere so no reader trusts a
counter whose increment may not have applied, as ``increment_spend_counters_pipeline`` does."""
redis_cache: Final = spend_counter_cache.redis_cache
if redis_cache is None or not pending:
return False
batch: Final = active_post_call_redis_batch(redis_cache)
if batch is None:
return False
ttl: Final = redis_cache.get_ttl()
with service_target(SPEND_COUNTERS_TARGET):
for item in pending:
batch.increment(item.counter_key, item.increment, ttl).on_settled(_settle_spend_counter_increment(item))
return True
def _settle_spend_counter_increment(item: PendingSpendIncrement) -> Callable[[asyncio.Future[float]], Awaitable[None]]:
async def settle(future: asyncio.Future[float]) -> None:
if not future.cancelled() and future.exception() is None:
current_value: Final = float(future.result())
spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
record_spend_counter_value(item.counter_key, current_value)
return
if future.cancelled():
if spend_counter_cache.in_memory_cache.get_cache(key=item.counter_key) is not None:
spend_counter_cache.in_memory_cache.increment_cache(key=item.counter_key, value=item.increment)
return
verbose_proxy_logger.warning(
"Spend counter %s increment did not land in the post-call pipeline; invalidating it", item.counter_key
)
await invalidate_spend_counter(counter_key=item.counter_key)
return settle
return spend_counter_store.increment_many_post_call(tuple((item.counter_key, item.increment) for item in pending))
async def increment_spend_counters_pipeline(pending: Sequence[PendingSpendIncrement]) -> tuple[float | None, ...]:
"""One INCRBYFLOAT+EXPIRE pipeline for every pending counter, returning each counter's new value in order; on
failure every counter is invalidated before the error propagates, so no caller can read a half-applied batch."""
if spend_counter_cache.redis_cache is None:
if not spend_counter_store.has_redis:
return await run_spend_counter_pipeline(pending=pending)
try:
return await run_spend_counter_pipeline(pending=pending)
@ -4516,26 +4455,21 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) -
happens to counters whose increment may or may not have landed when the pipeline fails."""
if not pending:
return ()
redis_cache: Final = spend_counter_cache.redis_cache
if redis_cache is None:
if not spend_counter_store.has_redis:
return tuple(
[
await SpendCounterReseed.increment_in_memory(
spend_counter_cache=spend_counter_cache, counter_key=item.counter_key, increment=item.increment
spend_counter_store=spend_counter_store, counter_key=item.counter_key, increment=item.increment
)
for item in pending
]
)
ttl: Final = redis_cache.get_ttl()
increment_list: Final = [
RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl)
for item in pending
]
results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list)
for item, current_value in zip(pending, results or ()):
spend_counter_cache.in_memory_cache.set_cache(key=item.counter_key, value=current_value)
record_spend_counter_value(item.counter_key, float(current_value))
return tuple(float(current_value) for current_value in results or ())
totals: Final = await spend_counter_store.increment_many(
tuple((item.counter_key, item.increment) for item in pending)
)
for item, current_value in zip(pending, totals, strict=True):
record_spend_counter_value(item.counter_key, current_value)
return totals
_CachedSpendObjectT = TypeVar(
@ -4558,15 +4492,9 @@ def _add_spend_to_locally_cached_object(
)
async def _add_spend_to_shared_global_proxy_spend(
cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None
) -> float | None:
if cache.redis_cache is None:
return None
async def _add_spend_to_shared_global_proxy_spend(cache: UserApiKeyCache, response_cost: float) -> float | None:
try:
return await cache.redis_cache.async_increment_if_present(
key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=response_cost, parent_otel_span=parent_otel_span
)
return await cache.increment_if_present(GLOBAL_PROXY_SPEND_CACHE_KEY, response_cost)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to add spend to the global proxy spend in Redis. "
@ -4577,11 +4505,9 @@ async def _add_spend_to_shared_global_proxy_spend(
return None
async def _add_spend_to_global_proxy_spend(
cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None
) -> None:
async def _add_spend_to_global_proxy_spend(cache: UserApiKeyCache, response_cost: float) -> None:
local_before: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True)
shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost, parent_otel_span)
shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost)
local_total: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True)
if not isinstance(local_total, (int, float)):
return
@ -4697,7 +4623,7 @@ async def update_cache(
if litellm.max_budget <= 0:
return
with service_target(SPEND_COUNTERS_TARGET):
await _add_spend_to_global_proxy_spend(target, response_cost, parent_otel_span)
await _add_spend_to_global_proxy_spend(target, response_cost)
def run_ollama_serve():
@ -5420,6 +5346,7 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache:
redis_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
spend_counter_store.attach_redis(connect(redis_cache))
cli_sso_session_cache.attach_redis_cache(
redis_cache,
default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS,
@ -9059,6 +8986,7 @@ class ProxyConfig:
redis_cache=redis_cache,
user_api_key_cache=user_api_key_cache,
additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache),
spend_counters=spend_counter_store,
)
self.auth_cache_invalidation_subscriber = subscriber
subscriber.start()

View file

@ -35,9 +35,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
team_membership_reservation_cache_key,
)
@ -50,6 +48,17 @@ from litellm.proxy.spend_tracking.input_tokens import (
)
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement, spend_counter_batch_scope
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.redis.stores.spend_counters import (
end_user_counter,
key_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
user_counter,
)
from litellm.router import Router
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
from litellm.types.router import DeploymentTypedDict
@ -128,15 +137,16 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c
"""A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL
while the request is in flight so a request longer than the TTL does not drop its
reservation and admit concurrent requests against the DB floor on any worker."""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
if spend_counter_cache.redis_cache is None or not counter_keys:
redis_ttl: Final = spend_counter_store.redis_ttl
if redis_ttl is None or not counter_keys:
return
task: Final = asyncio.create_task(
_renew_reservation_lease(
budget_reservation=budget_reservation,
counter_keys=counter_keys,
interval=spend_counter_cache.redis_cache.default_ttl / 2,
interval=redis_ttl / 2,
request_task=asyncio.current_task(),
)
)
@ -174,7 +184,7 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token:
"""
if valid_token is None:
return False
return counter_key == f"spend:key:{valid_token.token}" and should_throttle_budget_exceeded(valid_token)
return counter_key == key_counter(str(valid_token.token)) and should_throttle_budget_exceeded(valid_token)
async def _apply_over_budget_reservation_policy(
@ -489,7 +499,7 @@ async def _get_budget_counters(
if valid_token.max_budget is not None and valid_token.max_budget > 0:
counters.append(
_BudgetCounter(
counter_key=f"spend:key:{valid_token.token}",
counter_key=key_counter(valid_token.token),
source_cache_key=valid_token.token,
max_budget=float(valid_token.max_budget),
fallback_spend=float(valid_token.spend or 0.0),
@ -499,7 +509,7 @@ async def _get_budget_counters(
)
counters.extend(
_get_budget_limit_counters(
entity_prefix=f"spend:key:{valid_token.token}",
entity_prefix=key_counter(valid_token.token),
entity_type="Key",
entity_id=valid_token.token,
budget_limits=valid_token.budget_limits,
@ -512,7 +522,7 @@ async def _get_budget_counters(
if team_object.max_budget is not None and team_object.max_budget > 0:
counters.append(
_BudgetCounter(
counter_key=f"spend:team:{team_id}",
counter_key=team_counter(team_id),
source_cache_key=f"team_id:{team_id}",
max_budget=float(team_object.max_budget),
fallback_spend=float(team_object.spend or 0.0),
@ -522,7 +532,7 @@ async def _get_budget_counters(
)
counters.extend(
_get_budget_limit_counters(
entity_prefix=f"spend:team:{team_id}",
entity_prefix=team_counter(team_id),
entity_type="Team",
entity_id=team_id,
budget_limits=team_object.budget_limits,
@ -540,7 +550,7 @@ async def _get_budget_counters(
):
counters.append(
_BudgetCounter(
counter_key=f"spend:user:{user_object.user_id}",
counter_key=user_counter(user_object.user_id),
source_cache_key=user_object.user_id,
max_budget=float(user_object.max_budget),
fallback_spend=float(user_object.spend or 0.0),
@ -623,7 +633,7 @@ async def _get_end_user_budget_counter(
return None
return _BudgetCounter(
counter_key=f"spend:end_user:{end_user_id}",
counter_key=end_user_counter(end_user_id),
source_cache_key=source_cache_key,
max_budget=max_budget,
fallback_spend=fallback_spend,
@ -663,7 +673,7 @@ async def _get_tag_budget_counters(
continue
counters.append(
_BudgetCounter(
counter_key=f"spend:tag:{tag_name}",
counter_key=tag_counter(tag_name),
source_cache_key=tag_cache_key(tag_name),
max_budget=max_budget,
fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0,
@ -705,7 +715,7 @@ def _model_access_group_counter(group: str, budget: ModelAccessGroupBudget | Non
if budget is None or budget.max_budget is None or budget.max_budget <= 0:
return None
return _BudgetCounter(
counter_key=model_access_group_spend_counter_key(group),
counter_key=model_access_group_counter(group),
source_cache_key=model_access_group_cache_key(group),
max_budget=budget.max_budget,
fallback_spend=budget.spend,
@ -769,7 +779,7 @@ async def _get_team_member_budget_counter(
team_member_spend = cast(LiteLLM_TeamMembership, team_membership).spend if team_membership is not None else 0.0
return _BudgetCounter(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}",
counter_key=team_member_counter(valid_token.user_id, team_object.team_id),
source_cache_key=membership_cache_key,
max_budget=float(team_member_budget),
fallback_spend=float(team_member_spend or 0.0),
@ -808,7 +818,7 @@ async def _get_org_budget_counter(
org_spend: Final = _to_float(_get_value(org_table, "spend")) or 0.0
return _BudgetCounter(
counter_key=f"spend:org:{org_id}",
counter_key=org_counter(org_id),
source_cache_key=f"org_id:{org_id}:with_budget",
max_budget=org_max_budget,
fallback_spend=org_spend,
@ -841,7 +851,7 @@ async def _get_project_budget_counter(
return None
return _BudgetCounter(
counter_key=project_spend_counter_key(valid_token.project_id),
counter_key=project_counter(valid_token.project_id),
source_cache_key=source_cache_key,
max_budget=project_max_budget,
fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0,
@ -1096,9 +1106,9 @@ async def _get_current_counter_value(counter: _BudgetCounter) -> float:
def _counters_batch_scope(counter_keys: frozenset[str]) -> spend_counter_batch_scope:
"""Each counter is read once, then written, so one MGET up front serves every read in the loop."""
from litellm.proxy.proxy_server import spend_counter_cache
from litellm.proxy.proxy_server import spend_counter_store
return spend_counter_batch_scope(spend_counter_cache.redis_cache, counter_keys=counter_keys)
return spend_counter_batch_scope(spend_counter_store, counter_keys=counter_keys)
@dataclass(frozen=True, slots=True)

View file

@ -1,27 +1,34 @@
"""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 pydantic import TypeAdapter
from litellm._internal_context import service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_batch import BatchResult, RedisBatch, active_request_redis_batch
from litellm.caching.redis_cache import RedisCache
from litellm.caching.redis_batch import Batch, RequestRedisBatches, active_request_redis_batches
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_spend_counter_key,
project_spend_counter_key,
from litellm.redis.stores.spend_counters import (
SpendCounterStore,
end_user_counter,
key_counter,
model_access_group_counter,
org_counter,
project_counter,
tag_counter,
team_counter,
team_member_counter,
user_counter,
)
_CounterValues: Final = TypeAdapter(dict[str, float | None])
_NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
SPEND_COUNTERS_TARGET: Final = "spend_counters"
@dataclass(frozen=True, slots=True)
class _InflightRead:
result: Awaitable[Mapping[str, float | None] | None]
@dataclass(frozen=True, slots=True)
@ -35,23 +42,23 @@ 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 MGET rides the request's pipeline instead: the batch's flush
hook declares whatever is bound but unread, so whoever flushes first (the auth object prefetch, usually)
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", "_redis_cache", "_request_batch")
__slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_request", "_store")
def __init__(self, redis_cache: RedisCache) -> None:
self._redis_cache: Final = redis_cache
def __init__(self, store: SpendCounterStore) -> None:
self._store: Final = store
self._lock: Final = asyncio.Lock()
self._open = True
self._keys: frozenset[str] = frozenset()
self._fetched: frozenset[str] = frozenset()
self._loaded: Mapping[str, float | None] = _NO_VALUES
self._inflight: Final[list[BatchResult[Mapping[str, object]]]] = [] # mutable-ok: drained by _load
self._request_batch: Final[RedisBatch | None] = active_request_redis_batch(redis_cache)
if self._request_batch is not None:
self._request_batch.add_flush_hook(self._declare_pending)
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_pre_call_flush_hook(self._declare_pending)
@property
def counter_keys(self) -> frozenset[str]:
@ -96,7 +103,7 @@ class SpendCounterBatch:
async def _load(self) -> Mapping[str, float | None]:
async with self._lock:
if self._request_batch is not None:
if self._request is not None:
self._declare_pending()
await self._collect_inflight()
return self._loaded
@ -108,34 +115,30 @@ class SpendCounterBatch:
return self._loaded
def _declare_pending(self) -> None:
"""Flush hook: put every bound-but-unread counter on the request pipeline that is about to go out."""
if self._request_batch is None or not self._open:
"""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 = self._keys - self._fetched
if pending:
self._fetched = self._fetched | pending
with service_target(SPEND_COUNTERS_TARGET):
self._inflight.append(self._request_batch.mget(sorted(pending)))
pending: Final = tuple(sorted(self._keys - self._fetched))
if not pending or not self._store.has_redis:
return
self._fetched = self._fetched | frozenset(pending)
self._inflight.append(_InflightRead(result=self._store.read_many(pending, batch=Batch.PRE_CALL)))
async def _collect_inflight(self) -> None:
results: Final = tuple(self._inflight)
reads: Final = tuple(self._inflight)
self._inflight.clear()
for result in results:
for read in reads:
try:
fetched: Mapping[str, float | None] = _CounterValues.validate_python(await 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
self._loaded = MappingProxyType({**fetched, **self._loaded})
if values is not None:
self._loaded = MappingProxyType({**values, **self._loaded})
async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]:
try:
with service_target(SPEND_COUNTERS_TARGET):
values: Final = await self._redis_cache.async_batch_get_cache(key_list=sorted(keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped cache API
return _CounterValues.validate_python(values) # pyright: ignore[reportUnknownArgumentType] # untyped cache API
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)
return _NO_VALUES
fetched: Final = await self._store.read_many(sorted(keys), batch=Batch.NONE)
return _NO_VALUES if fetched is None else fetched
_active_batch: Final[ContextVar[SpendCounterBatch | None]] = ContextVar("spend_counter_batch", default=None)
@ -149,21 +152,21 @@ class spend_counter_batch_scope:
"""Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it.
Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET."""
__slots__ = ("_counter_keys", "_redis_cache", "_token")
__slots__ = ("_counter_keys", "_store", "_token")
def __init__(self, redis_cache: RedisCache | None, counter_keys: frozenset[str] = frozenset()) -> None:
self._redis_cache: Final = redis_cache
def __init__(self, store: SpendCounterStore, counter_keys: frozenset[str] = frozenset()) -> None:
self._store: Final = store
self._counter_keys: Final = counter_keys
self._token: Token[SpendCounterBatch | None] | None = None
def __enter__(self) -> None:
if self._redis_cache is None:
if not self._store.has_redis:
return
outer: Final = _active_batch.get()
if outer is not None and outer.is_open:
outer.bind(self._counter_keys)
return
batch: Final = SpendCounterBatch(self._redis_cache)
batch: Final = SpendCounterBatch(self._store)
batch.bind(self._counter_keys)
self._token = _active_batch.set(batch)
@ -191,19 +194,19 @@ def _iter_entity_counter_keys(
"""Only string ids name a counter; anything else (None, or an unresolved placeholder in synthetic
logging payloads) simply has no counter to bind."""
if isinstance(token, str):
yield f"spend:key:{token}"
yield key_counter(token)
if isinstance(team_id, str):
yield f"spend:team:{team_id}"
yield team_counter(team_id)
if isinstance(user_id, str):
yield f"spend:team_member:{user_id}:{team_id}"
yield team_member_counter(user_id, team_id)
if isinstance(user_id, str):
yield f"spend:user:{user_id}"
yield user_counter(user_id)
if isinstance(end_user_id, str):
yield f"spend:end_user:{end_user_id}"
yield end_user_counter(end_user_id)
if isinstance(org_id, str):
yield f"spend:org:{org_id}"
yield org_counter(org_id)
if isinstance(project_id, str):
yield project_spend_counter_key(project_id)
yield project_counter(project_id)
def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]:
@ -240,11 +243,9 @@ def post_call_counter_keys(
end_user_id=end_user_id,
)
)
tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str))
tag_keys: Final = frozenset(tag_counter(tag) for tag in tags or () if tag and isinstance(tag, str))
group_keys: Final = frozenset(
model_access_group_spend_counter_key(group)
for group in model_access_groups or ()
if group and isinstance(group, str)
model_access_group_counter(group) for group in model_access_groups or () if group and isinstance(group, str)
)
return entity_keys | tag_keys | group_keys

View file

@ -1,6 +1,20 @@
from collections.abc import Mapping
from dataclasses import dataclass
from litellm.constants import (
REDIS_CIRCUIT_BREAKER_ENABLED,
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
)
@dataclass(frozen=True, slots=True)
class RedisBreakerSettings:
failure_threshold: int
recovery_timeout: float
timeout_min_duration: float
@dataclass(frozen=True, slots=True)
class RedisStoreSettings:
@ -9,3 +23,14 @@ class RedisStoreSettings:
redis_kwargs: Mapping[str, object]
namespace: str | None
default_ttl: int
breaker: RedisBreakerSettings | None
def breaker_settings() -> RedisBreakerSettings | None:
if not REDIS_CIRCUIT_BREAKER_ENABLED:
return None
return RedisBreakerSettings(
failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
timeout_min_duration=REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
)

View file

@ -1,14 +1,45 @@
from __future__ import annotations
from collections.abc import Awaitable, Iterable
from collections.abc import Awaitable, Callable, Generator, Iterable
from dataclasses import dataclass
from typing import Final, TypeVar
from typing import Final, Generic, TypeVar, final
from litellm._internal_context import current_service_target, service_target
from litellm.caching.redis_batch import Batch, NativePending, RequestRedisBatches, active_request_redis_batches
from litellm.rust_bridge._native import NativeRequestRedis, NativeSend
_T: Final = TypeVar("_T")
_T = TypeVar("_T")
_U = TypeVar("_U")
@final
class Reply(Generic[_T]):
"""A store reply that builds its coroutine only once awaited, so one declared on a batch and never awaited,
such as an armed prefetch the request no longer needs, leaves no coroutine behind to warn when collected."""
__slots__ = ("_start",)
def __init__(self, start: Callable[[], Awaitable[_T]]) -> None:
self._start: Final = start
def __await__(self) -> Generator[object, None, _T]:
return self._start().__await__()
def map(self, convert: Callable[[_T], _U]) -> Reply[_U]:
return Reply(lambda: _converted(self, convert))
async def _converted(reply: Reply[_T], convert: Callable[[_T], _U]) -> _U:
return convert(await reply)
async def _ready(value: _T) -> _T:
return value
def ready(value: _T) -> Reply[_T]:
"""A reply already known, awaited without yielding to the event loop."""
return Reply(lambda: _ready(value))
@dataclass(frozen=True, slots=True)
@ -20,10 +51,13 @@ class Route:
send: NativeSend
request: RequestRedisBatches | None
def result(self, pending: NativePending[_T]) -> Awaitable[_T]:
if self.request is None:
return pending.resolve()
return _resolved_under(current_service_target(), self.request.native_result(pending))
def result(self, pending: NativePending[_T]) -> Reply[_T]:
request: Final = self.request
if request is None:
resolved: Final = pending.resolve()
return Reply(lambda: resolved)
target: Final = current_service_target()
return Reply(lambda: _resolved_under(target, request.native_result(pending)))
async def _resolved_under(target: str | None, reply: Awaitable[_T]) -> _T:
@ -31,8 +65,8 @@ async def _resolved_under(target: str | None, reply: Awaitable[_T]) -> _T:
return await reply
async def as_tuple(reply: Awaitable[Iterable[_T]]) -> tuple[_T, ...]:
return tuple(await reply)
def as_tuple(reply: Reply[Iterable[_T]]) -> Reply[tuple[_T, ...]]:
return reply.map(tuple)
def route(batch: Batch) -> Route:

View file

@ -0,0 +1,146 @@
from __future__ import annotations
import math
from collections.abc import Awaitable, Mapping, Sequence
from types import MappingProxyType
from typing import Final
from litellm._internal_context import with_service_target
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
from litellm.redis.stores.auth_cache.types import AUTH_CACHE_TARGET, AuthCacheSettings, WriteBack
from litellm.rust_bridge._native import LocalSizes
__all__ = [
"AUTH_CACHE_TARGET",
"AuthCacheSettings",
"AuthCacheStore",
"WriteBack",
]
def _redis_ttl(ttl: float | None) -> float | None:
"""``RedisCache.get_ttl``: a NaN TTL, which ``int()`` refuses, takes the default TTL."""
return None if ttl is not None and math.isnan(ttl) else 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, handle: AuthCacheHandle) -> None:
self._handle: Final = handle
def reconfigure(self, settings: AuthCacheSettings) -> None:
"""Both memory partitions keep what they hold, as ``update_cache_ttl`` and ``update_in_memory_max_size`` do."""
self._handle.reconfigure(memory_entries=settings.memory_entries, memory_ttl_seconds=settings.memory_ttl)
@with_service_target(AUTH_CACHE_TARGET)
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)
@with_service_target(AUTH_CACHE_TARGET)
def get_sync(self, key: str) -> object:
return self._handle.get_sync(key)
@with_service_target(AUTH_CACHE_TARGET)
async def set(self, key: str, value: object, ttl: float | None) -> None:
self._forget_miss(key)
await self._handle.set(key, value, _redis_ttl(ttl))
@with_service_target(AUTH_CACHE_TARGET)
def set_sync(self, key: str, value: object, ttl: float | None) -> None:
self._forget_miss(key)
self._handle.set_sync(key, value, _redis_ttl(ttl))
@with_service_target(AUTH_CACHE_TARGET)
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, _redis_ttl(ttl), target.send)
target.request.forget_native_miss(key)
return target.result(pending)
@with_service_target(AUTH_CACHE_TARGET)
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)
@with_service_target(AUTH_CACHE_TARGET)
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, _redis_ttl(ttl))
@with_service_target(AUTH_CACHE_TARGET)
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(
tuple((key, value, _redis_ttl(ttl)) for key, value, ttl in entries), route(batch).send
)
@with_service_target(AUTH_CACHE_TARGET)
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)
@with_service_target(AUTH_CACHE_TARGET)
def delete_sync(self, keys: Sequence[str]) -> None:
self._handle.delete_sync(keys)
@with_service_target(AUTH_CACHE_TARGET)
async def read_missing(self, keys: Sequence[str], *, batch: Batch) -> Mapping[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 MappingProxyType({})
found: Final = MappingProxyType(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
@with_service_target(AUTH_CACHE_TARGET, preserve_existing=True)
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) -> LocalSizes:
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)

View file

@ -0,0 +1,31 @@
from __future__ import annotations
from collections.abc import Awaitable, Sequence
from typing import Protocol
from litellm.caching.redis_batch import NativePending
from litellm.redis.stores.auth_cache.types import WriteBack
from litellm.rust_bridge._native import LocalSizes, 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) -> LocalSizes: ...
def reconfigure(self, *, memory_entries: int, memory_ttl_seconds: float | None = None) -> None: ...

View file

@ -0,0 +1,16 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Final, TypeAlias
AUTH_CACHE_TARGET: Final = "auth_cache"
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

View file

@ -3,55 +3,48 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from litellm.constants import (
REDIS_CIRCUIT_BREAKER_ENABLED,
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
from litellm.rust_bridge._native import (
NativeAuthCache,
NativeRateLimits,
NativeRedisClient,
NativeRouterCooldowns,
NativeRouterUsage,
NativeSensitiveRoutePins,
NativeSpendCounters,
)
from litellm.rust_bridge._native import NativeRedisClient
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
_KWARGS: Final = TypeAdapter(dict[str, object])
from litellm.redis.stores.auth_cache.types import AuthCacheSettings
@dataclass(frozen=True, slots=True)
class RedisConnection:
"""A Redis connection the stores share. Pass one to every store that should ride the same client."""
class RedisStores:
"""Builds every store family's handle on one shared client. Without a client the handles run memory-only."""
client: NativeRedisClient
client: NativeRedisClient | None
def router_cooldowns(self) -> NativeRouterCooldowns:
return NativeRouterCooldowns(self.client)
def router_usage(self) -> NativeRouterUsage:
return NativeRouterUsage(self.client)
def spend_counters(self) -> NativeSpendCounters:
return NativeSpendCounters(self.client)
def rate_limits(self, backfill_ttl: float | None = None) -> NativeRateLimits:
return NativeRateLimits(self.client, backfill_ttl)
def auth_cache(self, settings: AuthCacheSettings) -> NativeAuthCache:
return NativeAuthCache(
self.client, memory_entries=settings.memory_entries, memory_ttl_seconds=settings.memory_ttl
)
def sensitive_route_pins(self) -> NativeSensitiveRoutePins:
return NativeSensitiveRoutePins(self.client)
def _breaker() -> tuple[int, float, float] | None:
if not REDIS_CIRCUIT_BREAKER_ENABLED:
return None
return (
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
float(REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT),
REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
)
def connect(redis_cache: RedisCache | None) -> RedisConnection | None:
def connect(redis_cache: RedisCache | None) -> RedisStores:
settings: Final = None if redis_cache is None else redis_cache.store_settings
if settings is None:
return None
from litellm._redis import ( # redis-py, which importing litellm must not require
_merge_redis_kwargs, # pyright: ignore[reportPrivateUsage] # Reuse the shared internal kwargs normalizer
)
client: Final = NativeRedisClient.connect(
_merge_redis_kwargs(_KWARGS.validate_python(settings.redis_kwargs)),
default_ttl_seconds=settings.default_ttl,
namespace=settings.namespace,
breaker=_breaker(),
)
return None if client is None else RedisConnection(client)
def client_of(connection: RedisConnection | None) -> NativeRedisClient | None:
return None if connection is None else connection.client
return RedisStores(None if settings is None else NativeRedisClient.connect(settings))

View file

@ -0,0 +1,192 @@
from __future__ import annotations
from collections.abc import Awaitable, Sequence
from typing import Final, TypeVar
from typing_extensions import Never
from litellm._internal_context import with_service_target
from litellm.caching.redis_batch import Batch, NativePending, active_request_redis_batches
from litellm.redis.stores._route import Reply, Route, as_tuple, post_call_request, route
from litellm.redis.stores.rate_limits.handle import RateLimitsHandle
from litellm.redis.stores.rate_limits.types import (
RATE_LIMITS_TARGET,
AcquireReply,
AtomicApplied,
AtomicCounter,
AtomicOverLimit,
AtomicReply,
CounterIncrement,
CounterKind,
GaugeOverLimit,
GaugesAcquired,
GuardedReply,
LocalValue,
RateLimitCounter,
RateLimitEntry,
RateLimitGauge,
RateLimitScope,
RateLimitWindow,
WindowCount,
)
__all__ = [
"RATE_LIMITS_TARGET",
"AcquireReply",
"AtomicApplied",
"AtomicCounter",
"AtomicOverLimit",
"AtomicReply",
"CounterIncrement",
"CounterKind",
"GaugeOverLimit",
"GaugesAcquired",
"GuardedReply",
"LocalValue",
"RateLimitCounter",
"RateLimitEntry",
"RateLimitGauge",
"RateLimitScope",
"RateLimitStore",
"RateLimitWindow",
"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) -> Reply[_T]:
"""A script's reply; one declared without Redis raises where it is awaited, like a failed script."""
return Reply(_no_redis) if pending is None else target.result(pending)
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, handle: RateLimitsHandle) -> None:
self._handle: Final = handle
@property
def has_redis(self) -> bool:
return self._handle.has_redis
@property
def has_request_batch(self) -> bool:
"""Whether ``Batch.PRE_CALL`` scripts share the request's pre-call trip and ``*_post_call`` writes ride its
post-call batch, rather than each go out on its own or 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(scopes, force))
def local(self, entries: Sequence[RateLimitEntry]) -> tuple[LocalValue | None, ...]:
return tuple(self._handle.get_local(entry) for entry in entries)
def set_local(self, entry: RateLimitEntry, value: LocalValue, ttl: float | None) -> None:
self._handle.set_local(entry, value, ttl)
@with_service_target(RATE_LIMITS_TARGET)
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."""
local: Final = self._handle.get_local(entry)
if local is not None or not self.has_redis:
return local
return await self._handle.get(entry, backfill_ttl_seconds=backfill_ttl)
@with_service_target(RATE_LIMITS_TARGET)
def sliding_window(
self, counters: Sequence[RateLimitCounter], now: int, window_size: int, *, batch: Batch
) -> Awaitable[tuple[WindowCount, ...]]:
target: Final = route(batch)
return as_tuple(_script(self._handle.sliding_window(counters, now, window_size, target.send), target))
@with_service_target(RATE_LIMITS_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)
return as_tuple(_script(self._handle.read_counters(counters, target.send), target))
@with_service_target(RATE_LIMITS_TARGET)
def check_and_increment(self, counters: Sequence[AtomicCounter], *, batch: Batch) -> Awaitable[AtomicReply]:
target: Final = route(batch)
return _script(self._handle.check_and_increment(counters, target.send), target)
@with_service_target(RATE_LIMITS_TARGET)
def guarded_increment(
self,
counter: RateLimitCounter,
expected_window_start: str,
amount: float,
ttl: int | None,
*,
batch: Batch,
) -> Awaitable[GuardedReply]:
target: Final = route(batch)
return _script(self._handle.guarded_increment(counter, expected_window_start, amount, ttl, target.send), target)
@with_service_target(RATE_LIMITS_TARGET)
def parallel_acquire(
self, gauges: Sequence[tuple[RateLimitScope, int]], slot_ttl: int, slot_id: str, *, batch: Batch
) -> Awaitable[AcquireReply]:
target: Final = route(batch)
return _script(self._handle.parallel_acquire(gauges, slot_ttl, slot_id, target.send), target)
@with_service_target(RATE_LIMITS_TARGET)
def parallel_release(
self, gauges: Sequence[RateLimitScope], slot_id: str, *, batch: Batch
) -> Awaitable[tuple[int, ...]]:
target: Final = route(batch)
return as_tuple(_script(self._handle.parallel_release(gauges, slot_id, target.send), target))
@with_service_target(RATE_LIMITS_TARGET)
def parallel_count(
self, gauges: Sequence[RateLimitScope], slot_ttl: int, *, batch: Batch
) -> Awaitable[tuple[int, ...]]:
target: Final = route(batch)
return as_tuple(_script(self._handle.parallel_count(gauges, slot_ttl, target.send), target))
@with_service_target(RATE_LIMITS_TARGET)
def token_increment(self, increments: Sequence[CounterIncrement], *, batch: Batch) -> Awaitable[tuple[float, ...]]:
target: Final = route(batch)
return as_tuple(_script(self._handle.token_increment(increments, target.send), target))
@with_service_target(RATE_LIMITS_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, increments)
@with_service_target(RATE_LIMITS_TARGET)
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, gauges, slot_id)
@with_service_target(RATE_LIMITS_TARGET)
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)
@with_service_target(RATE_LIMITS_TARGET)
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(increments)
@with_service_target(RATE_LIMITS_TARGET)
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 = post_call_request() if self.has_redis else None
if request is None:
await self._handle.increment_many(increments)
return
self._handle.increment_many_post_call(request, increments)

View file

@ -0,0 +1,68 @@
from __future__ import annotations
from collections.abc import Awaitable, Sequence
from typing import Protocol
from litellm.caching.redis_batch import NativePending
from litellm.redis.stores.rate_limits.types import (
AcquireReply,
AtomicCounter,
AtomicReply,
CounterIncrement,
GuardedReply,
LocalValue,
RateLimitCounter,
RateLimitEntry,
RateLimitScope,
WindowCount,
)
from litellm.rust_bridge._native import NativeRequestRedis, NativeSend
class RateLimitsHandle(Protocol):
@property
def has_redis(self) -> bool: ...
def slot_groups(self, scopes: Sequence[RateLimitScope], force: bool) -> list[list[int]]: ...
def get_local(self, entry: RateLimitEntry) -> LocalValue | None: ...
def set_local(self, entry: RateLimitEntry, value: LocalValue, ttl_seconds: float | None) -> None: ...
def get(
self, entry: RateLimitEntry, *, backfill_ttl_seconds: float | None = None
) -> Awaitable[LocalValue | None]: ...
def sliding_window(
self, counters: Sequence[RateLimitCounter], now: int, window_size: int, send: NativeSend
) -> NativePending[list[WindowCount]] | None: ...
def read_counters(
self, counters: Sequence[RateLimitCounter], send: NativeSend
) -> NativePending[list[str | None]] | None: ...
def check_and_increment(
self, counters: Sequence[AtomicCounter], send: NativeSend
) -> NativePending[AtomicReply] | None: ...
def guarded_increment(
self,
counter: RateLimitCounter,
expected_window_start: str,
amount: float,
ttl_seconds: int | None,
send: NativeSend,
) -> NativePending[GuardedReply] | None: ...
def parallel_acquire(
self, gauges: Sequence[tuple[RateLimitScope, int]], slot_ttl_seconds: int, slot_id: str, send: NativeSend
) -> NativePending[AcquireReply] | None: ...
def parallel_release(
self, gauges: Sequence[RateLimitScope], slot_id: str, send: NativeSend
) -> NativePending[list[int]] | None: ...
def parallel_count(
self, gauges: Sequence[RateLimitScope], slot_ttl_seconds: int, send: NativeSend
) -> NativePending[list[int]] | None: ...
def token_increment(
self, increments: Sequence[CounterIncrement], send: NativeSend
) -> NativePending[list[float]] | None: ...
def token_increment_post_call(
self, request: NativeRequestRedis, increments: Sequence[CounterIncrement]
) -> bool: ...
def parallel_release_post_call(
self, request: NativeRequestRedis, gauges: Sequence[RateLimitScope], slot_id: str
) -> bool: ...
def increment(self, increment: CounterIncrement) -> Awaitable[int | float]: ...
def increment_many(self, increments: Sequence[CounterIncrement]) -> Awaitable[None]: ...
def increment_many_post_call(self, request: NativeRequestRedis, increments: Sequence[CounterIncrement]) -> None: ...

View file

@ -0,0 +1,107 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import ClassVar, Final, Literal, TypeAlias
from litellm.rust_bridge._native import (
AtomicApplied,
AtomicOverLimit,
GaugeOverLimit,
GaugesAcquired,
GuardedReply,
WindowCount,
)
RATE_LIMITS_TARGET: Final = "rate_limits"
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
suffix: ClassVar[str] = "window"
@dataclass(frozen=True, slots=True, order=True)
class RateLimitCounter:
scope: RateLimitScope
kind: CounterKind
@property
def window(self) -> RateLimitWindow:
return RateLimitWindow(self.scope)
@property
def suffix(self) -> str:
return self.kind
@dataclass(frozen=True, slots=True)
class RateLimitGauge:
scope: RateLimitScope
suffix: ClassVar[str] = "max_parallel_requests"
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
AtomicReply: TypeAlias = AtomicApplied | AtomicOverLimit
AcquireReply: TypeAlias = GaugesAcquired | GaugeOverLimit
__all__ = [
"RATE_LIMITS_TARGET",
"AcquireReply",
"AtomicApplied",
"AtomicCounter",
"AtomicOverLimit",
"AtomicReply",
"CounterIncrement",
"CounterKind",
"GaugeOverLimit",
"GaugesAcquired",
"GuardedReply",
"LocalValue",
"RateLimitCounter",
"RateLimitEntry",
"RateLimitGauge",
"RateLimitScope",
"RateLimitWindow",
"WindowCount",
]

View file

@ -0,0 +1,93 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol
from litellm._internal_context import with_service_target
from litellm.caching.redis_batch import Batch
from litellm.redis.stores._route import as_tuple, ready, route
from litellm.redis.stores.router_cooldowns.handle import RouterCooldownsHandle
from litellm.redis.stores.router_cooldowns.types import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, CooldownCacheValue
if TYPE_CHECKING:
from litellm.redis.stores.connection import RedisStores
__all__ = [
"ROUTER_COOLDOWNS_TARGET",
"ActiveCooldowns",
"CooldownCacheValue",
"RouterCooldownStore",
"RouterCooldowns",
]
class RouterCooldowns(Protocol):
def attach_redis(self, redis: RedisStores) -> None: ...
@property
def has_redis(self) -> bool: ...
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 async_put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ...
async def async_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]) -> ActiveCooldowns: ...
def minimum(self, model_ids: Sequence[str], default: float) -> float: ...
def evict_local(self, model_id: str) -> None: ...
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
class RouterCooldownStore:
def __init__(self, handle: RouterCooldownsHandle) -> None:
self._handle: Final = handle
@property
def has_redis(self) -> bool:
return self._handle.has_redis
def attach_redis(self, redis: RedisStores) -> None:
"""Adds the Redis tier once, keeping this router's memory, as ``DualCache.attach_redis_cache`` does."""
if redis.client is not None:
self._handle.attach(redis.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)
async def async_put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None:
await self._handle.async_put(model_id, value, ttl)
@with_service_target(ROUTER_COOLDOWNS_TARGET)
async def async_increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int:
return await self._handle.async_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.
Without Redis it is answered from memory without suspending, as ``DualCache``'s memory read was."""
if not self.has_redis:
return ready(self.active_sync(model_ids))
target: Final = route(batch)
return as_tuple(target.result(self._handle.active(model_ids, target.send)))
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def active_sync(self, model_ids: Sequence[str]) -> ActiveCooldowns:
return tuple(self._handle.active_sync(model_ids))
@with_service_target(ROUTER_COOLDOWNS_TARGET)
def minimum(self, model_ids: Sequence[str], default: float) -> 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()

View file

@ -0,0 +1,26 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Protocol
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):
@property
def has_redis(self) -> bool: ...
def attach(self, client: NativeRedisClient, /) -> 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 async_put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> Awaitable[None]: ...
def async_increment_allowed_fails(
self, model_id: str, suffix: str | None, ttl_seconds: float
) -> Awaitable[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]: ...

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

View file

@ -0,0 +1,116 @@
from __future__ import annotations
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Final, Protocol
from litellm._internal_context import with_service_target
from litellm.redis.stores.router_usage.handle import RouterUsageHandle
from litellm.redis.stores.router_usage.types import (
ROUTER_USAGE_TARGET,
CounterKind,
UsageCounter,
)
if TYPE_CHECKING:
from litellm.redis.stores.connection import RedisStores
__all__ = [
"ROUTER_USAGE_TARGET",
"CounterKind",
"RouterUsage",
"RouterUsageStore",
"UsageCounter",
]
class RouterUsage(Protocol):
def attach_redis(self, redis: RedisStores) -> None: ...
def flush_local_counters(self) -> None: ...
async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float: ...
async def increment_deployment(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> None: ...
async def get_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ...
def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float: ...
def increment_deployment_sync(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> None: ...
def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ...
def snapshot(self) -> Mapping[str, float]: ...
class RouterUsageStore:
def __init__(self, handle: RouterUsageHandle) -> None:
self._handle: Final = handle
@property
def has_redis(self) -> bool:
return self._handle.has_redis
def attach_redis(self, redis: RedisStores) -> None:
"""Adds the Redis tier once, keeping this router's memory counts, as ``DualCache.attach_redis_cache`` does."""
if redis.client is not None:
self._handle.attach(redis.client)
def flush_local_counters(self) -> None:
self._handle.flush_local_counters()
@with_service_target(ROUTER_USAGE_TARGET)
async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float:
"""Without Redis every async method answers from memory without suspending, as ``DualCache``'s did."""
if not self.has_redis:
return self.increment_sync(counter, value, ttl)
return await self._handle.increment(counter, value, ttl)
@with_service_target(ROUTER_USAGE_TARGET)
async def increment_deployment(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> None:
if not self.has_redis:
self.increment_deployment_sync(deployment_id, model, minute, tokens, requests)
return
await self._handle.increment_deployment(deployment_id, model, minute, tokens, requests)
@with_service_target(ROUTER_USAGE_TARGET)
async def get_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]:
if not self.has_redis:
return self.get_many_sync(counters)
return await self._handle.get_many(counters)
@with_service_target(ROUTER_USAGE_TARGET)
def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float:
"""``increment`` for a sync call, which has no event loop to await on."""
return self._handle.increment_sync(counter, value, ttl)
@with_service_target(ROUTER_USAGE_TARGET)
def increment_deployment_sync(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> None:
self._handle.increment_deployment_sync(deployment_id, model, minute, tokens, requests)
@with_service_target(ROUTER_USAGE_TARGET)
def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]:
return self._handle.get_many_sync(counters)
def snapshot(self) -> Mapping[str, float]:
return self._handle.snapshot()

View file

@ -0,0 +1,35 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Protocol
from litellm.redis.stores.router_usage.types import UsageCounter
from litellm.rust_bridge._native import NativeRedisClient
class RouterUsageHandle(Protocol):
@property
def has_redis(self) -> bool: ...
def attach(self, client: NativeRedisClient, /) -> bool: ...
def increment(self, counter: UsageCounter, 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[UsageCounter]) -> Awaitable[Sequence[float | None]]: ...
def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: ...
def increment_deployment_sync(
self,
deployment_id: str,
model: str | None,
minute: str,
tokens: float,
requests: int,
) -> None: ...
def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ...
def flush_local_counters(self) -> None: ...
def snapshot(self) -> Mapping[str, float]: ...

View file

@ -0,0 +1,15 @@
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

View file

@ -0,0 +1,26 @@
from __future__ import annotations
from typing import Final
from litellm._internal_context import with_service_target
from litellm.redis.stores.sensitive_route_pins.handle import SensitiveRoutePinsHandle
__all__ = [
"SENSITIVE_ROUTE_PINS_TARGET",
"SensitiveRoutePinStore",
]
SENSITIVE_ROUTE_PINS_TARGET: Final = "sensitive_route_pins"
class SensitiveRoutePinStore:
def __init__(self, handle: SensitiveRoutePinsHandle) -> None:
self._handle: Final = handle
@with_service_target(SENSITIVE_ROUTE_PINS_TARGET)
async def get(self, tenant: str, session_id: str, default_ttl: float) -> str | None:
return await self._handle.get(tenant, session_id, default_ttl)
@with_service_target(SENSITIVE_ROUTE_PINS_TARGET)
async def set(self, tenant: str, session_id: str, model: str, ttl: float) -> None:
await self._handle.set(tenant, session_id, model, ttl)

View file

@ -0,0 +1,9 @@
from __future__ import annotations
from collections.abc import Awaitable
from typing import Protocol
class SensitiveRoutePinsHandle(Protocol):
def get(self, tenant: str, session_id: str, default_ttl_seconds: float) -> Awaitable[str | None]: ...
def set(self, tenant: str, session_id: str, model: str, ttl_seconds: float) -> Awaitable[None]: ...

View file

@ -0,0 +1,168 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, 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, ready, route
from litellm.redis.stores.spend_counters.handle import 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,
)
if TYPE_CHECKING:
from litellm.redis.stores.connection import RedisStores
__all__ = [
"END_USER_COUNTER_PREFIX",
"SPEND_COUNTERS_TARGET",
"SPEND_COUNTER_PREFIX",
"SpendCounterStore",
"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",
]
def _by_counter(counters: Sequence[str], read: Sequence[float | None] | None) -> Mapping[str, float | None] | None:
return None if read is None else MappingProxyType(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 in ``types``."""
def __init__(self, handle: SpendCountersHandle) -> None:
self._handle: Final = handle
def attach_redis(self, redis: RedisStores) -> None:
"""Adds the Redis tier once, keeping the memory copies, as ``DualCache.attach_redis_cache`` does."""
if redis.client is not None:
self._handle.attach(redis.client)
@property
def redis_ttl(self) -> float | None:
"""The TTL every Redis write gives a counter, None when there is no Redis."""
return self._handle.redis_ttl_seconds
@property
def has_redis(self) -> bool:
return self._handle.has_redis
@with_service_target(SPEND_COUNTERS_TARGET)
async def read(self, counter: str) -> SpendReading:
reading: Final = await self._handle.read(counter)
return reading.value, reading.from_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 (ready(None) if pending is None else target.result(pending)).map(partial(_by_counter, counters))
@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

View file

@ -0,0 +1,34 @@
from __future__ import annotations
from collections.abc import Awaitable, Mapping, Sequence
from typing import Protocol
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, SpendRead
class SpendCountersHandle(Protocol):
@property
def has_redis(self) -> bool: ...
@property
def redis_ttl_seconds(self) -> float | None: ...
def attach(self, client: NativeRedisClient, /) -> bool: ...
def read(self, counter: str) -> Awaitable[SpendRead]: ...
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]: ...

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

View file

@ -1659,6 +1659,11 @@ class RoutingContext:
) -> None:
return self._router._update_kwargs_before_fallbacks(model, kwargs, metadata_variable_name)
def arm_cooldown_prefetch(self) -> None:
from litellm.router_utils.routing_read_batch import RoutingPrefetch
RoutingPrefetch.arm_cooldowns(self)
def arm_read_prefetch(
self,
model: str,

View file

@ -31,6 +31,8 @@ if TYPE_CHECKING:
from litellm.caching.dual_cache import DualCache
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.redis.stores.connection import RedisStores
from litellm.redis.stores.router_usage import RouterUsage
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIFileObject
from litellm.types.rerank import RerankResponse
@ -147,6 +149,12 @@ class PythonHost:
return DualCache(redis_cache=redis_cache_for(config), in_memory_cache=InMemoryCache())
def router_usage(self, redis: "RedisStores") -> "RouterUsage":
"""The per-minute tpm/rpm counters a router counts its deployments' calls in, on ``redis``."""
from litellm.redis.stores.router_usage import RouterUsageStore
return RouterUsageStore(redis.router_usage())
def redis_cache_for(config: RouterConfig) -> "RedisCache | None":
"""The Redis cache the router's ``redis_*`` and ``cache_kwargs`` arguments describe, or None."""

View file

@ -16,6 +16,7 @@ Settings are resolved from the ``RouterConfig`` and the host's process settings;
core does not serve raises ``RustCoreUnsupported`` (an ``AttributeError``).
"""
import asyncio
import contextlib
import dataclasses
import inspect
@ -30,6 +31,8 @@ from time import time as unix_time
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, TypeVar, final, overload
from pydantic import TypeAdapter
from litellm._internal_context import service_target
from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_TARGET
from litellm.integrations.custom_guardrail import is_guardrail_intervention
@ -46,7 +49,6 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import (
from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.litellm_core_utils.core_helpers import (
as_object_list,
as_str_mapping,
get_parent_otel_span_from_kwargs,
without,
@ -61,6 +63,9 @@ from litellm.litellm_core_utils.routing_common_utils import (
)
from litellm.litellm_core_utils.routing_groups import validate_routing_strategy
from litellm.litellm_core_utils.secret_redaction import redact_string
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, RouterCooldownStore
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage
from litellm.router_api import (
ROUTER_INIT_ARGS,
SETTING_FIELDS,
@ -108,6 +113,7 @@ from litellm.router_rust.invoke.streaming import (
shape_attempt_kwargs,
)
from litellm.router_rust.invoke.walk import FallbackOutcomes, WalkMetadata, await_all, log_retry_into
from litellm.router_rust.routing_stores import RoutingStores
from litellm.router_rust.steps.access import caller_fields, selection_refusal
from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks
from litellm.router_rust.steps.plugins import SYNC_PATH_MESSAGE, candidate_models, run_plugins
@ -204,7 +210,6 @@ from litellm.router_rust.surface.settings import (
settings_view,
)
from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve
from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET
from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure
from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint
from litellm.router_utils.kubernetes_pod_discovery import (
@ -215,7 +220,7 @@ from litellm.router_utils.kubernetes_pod_discovery import (
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure
from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata
from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingReadBatch
from litellm.router_utils.routing_read_batch import RoutingReadBatch
from litellm.types.router import (
AllowedFailsPolicy,
Deployment,
@ -1385,8 +1390,10 @@ class RustCore(
self._legacy: Final = legacy
self._health_view: RustHealth | None = None
self._cache: DualCache = host.routing_cache(config)
redis: Final = connect(self._cache.redis_cache)
self._usage: RouterUsage = host.router_usage(redis)
self._cooldowns: CooldownCache = CooldownCache(
cache=self._cache,
store=RouterCooldownStore(redis.router_cooldowns()),
default_cooldown_time=self._cooldown_given, # pyright: ignore[reportArgumentType] # the constructor's value, as the Python core passes it
)
self._health_checks: DeploymentHealthCache = resolved.health_routing.health_checks(self._cache)
@ -1414,6 +1421,7 @@ class RustCore(
on_change=self._sync_registry,
)
self._hosted._core_cache = self._cache # pyright: ignore[reportPrivateUsage] # usage reads come from the cache the Rust core writes
self._hosted._core_usage_store = self._usage # pyright: ignore[reportPrivateUsage] # tpm/rpm reads come from the store the Rust core writes
ctx: Final = RoutingContext(self) # pyright: ignore[reportArgumentType] # the strategies read registry(), settings(), health() and select()
self._ctx: Final = ctx
strategies: Final = PreRoutingStrategies(
@ -2015,6 +2023,9 @@ class RustCore(
if changes.redis_cache is not None:
# `_core_update_redis_cache`: the routing state's cache gains the Redis cache.
self._cache.attach_redis_cache(changes.redis_cache)
redis: Final = connect(changes.redis_cache)
self._usage.attach_redis(redis)
self._cooldowns.cooldown_store.attach_redis(redis)
self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache)
self._prerouter.attach_redis_cache(changes.redis_cache)
return ApplyResult(
@ -2254,6 +2265,8 @@ class RustCore(
def flush_cache(self) -> None:
"""Drop the host's response cache and the routing state."""
self._host.set_global_cache(None)
self._usage.flush_local_counters()
self._cooldowns.cooldown_store.flush_local_counters()
self._cache.flush_cache()
self._checks.prompt_caching_pin_cache.in_memory_cache.flush_cache()
@ -3446,99 +3459,41 @@ class RustCore(
# -- cache steps -----------------------------------------------------------------------------
def _store_for(self, key: str) -> "DualCache":
"""Cooldowns live in the cooldown store; every other routing key in the routing cache."""
return self._cooldowns.cooldown_store if key.endswith(":cooldown") else self._cache
async def _batch_read(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list
"""One batch read, as the Python core makes it: the keys memory cannot serve, at most once per
read interval. Cooldowns are ``async_get_cooldown_deployments``' read, through the request's
armed prefetch when there is one."""
if not keys or not keys[0].endswith(":cooldown"):
read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above
store: Final = self._cooldowns.cooldown_store
"""One batch read: cooldowns and tpm/rpm usage from their stores, every other key from the routing cache."""
stored: Final = await self._stored_with_usage(keys)
if stored is not None:
return stored
read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above
async def _stored_with_usage(self, keys: Sequence[str]) -> tuple[object, ...] | None:
"""A usage-based-routing-v2 pick reads its tpm/rpm counters alongside its cooldowns (``_usage_read_scope``)."""
usage_read: Final = _USAGE_READ.get()
if usage_read is not None:
return self._corrected_cooldowns(keys, await self._batch_read_with_usage(keys, *usage_read))
with service_target(ROUTER_COOLDOWNS_TARGET):
prefetched: Final = await RoutingReadBatch._read_prefetched(((store, keys),)) # pyright: ignore[reportPrivateUsage] # the batch's own prefetch read
results: Final[object] = (
prefetched[0] if prefetched is not None else await store.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType] # see above
)
return self._corrected_cooldowns(keys, results)
async def _batch_read_with_usage(
self,
keys: list[str],
batch: RoutingReadBatch,
deployments: list[Mapping[str, object]], # mutable-ok: the read's key and deployment lists
) -> object:
"""The cooldown read with the usage-based-routing-v2 counters of ``deployments`` in one MGET
(``RoutingReadBatch.async_get_cooldown_deployments``), the counters kept on ``batch``."""
from litellm.caching.dual_cache import DualCache
from litellm.router_strategy.lowest_tpm_rpm_v2 import PrefetchedUsage
from litellm.router_utils.routing_read_batch import (
_routing_read_target, # pyright: ignore[reportPrivateUsage] # the batch's own target name
)
selector: Final = batch.usage_selector
assert selector is not None # set by `_usage_read_scope`
usage_keys: Final = list(itertools.chain(*selector.usage_counter_keys(deployments))) # pyright: ignore[reportArgumentType] # deployment dicts
reads: Final = ((self._cooldowns.cooldown_store, keys), (selector.router_cache, usage_keys))
with service_target(_routing_read_target(keys, usage_keys)):
results: Final = await RoutingReadBatch._read_prefetched(
reads
) or await DualCache.async_batch_get_cache_shared( # pyright: ignore[reportPrivateUsage] # see above
reads, parent_otel_span=None
)
usage_values: Final = results[1]
batch.prefetched_usage = PrefetchedUsage(
keys=frozenset(usage_keys),
values=None if usage_values is None else MappingProxyType(dict(zip(usage_keys, usage_values))),
)
return results[0]
if usage_read is None or not keys or not keys[0].endswith(":cooldown"):
return await self._stores.batch_get(keys)
batch, deployments = usage_read
stored, _ = await asyncio.gather(self._stores.batch_get(keys), batch.read_usage(deployments, None))
return stored
def _batch_read_sync(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list
"""``_batch_read`` for a sync pick: ``get_active_cooldowns``' read."""
if not keys or not keys[0].endswith(":cooldown"):
read: Final[object] = self._cache.batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above
with service_target(ROUTER_COOLDOWNS_TARGET):
results: Final = as_object_list(self._cooldowns.cooldown_store.batch_get_cache(keys=keys)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above
return self._corrected_cooldowns(keys, results)
"""``_batch_read`` for a sync pick."""
stored: Final = self._stores.batch_get_sync(keys)
if stored is not None:
return stored
read: Final[object] = self._cache.batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above
def _corrected_cooldowns(self, keys: Sequence[str], results: object) -> Sequence[object]:
"""The cooldowns a batch read found, once ``active_cooldowns_from_results`` corrected the in-memory
TTL of each it promoted from Redis and dropped each that expired."""
if not isinstance(results, list):
return [None] * len(keys)
ids: Final = [key.removeprefix("deployment:").removesuffix(":cooldown") for key in keys]
_ = self._cooldowns.active_cooldowns_from_results(ids, results) # pyright: ignore[reportUnknownArgumentType] # see above
return results # pyright: ignore[reportUnknownVariableType] # see above
async def _increment_usage(self, keys: Sequence[object], amounts: Sequence[object], ttl: float | None) -> None:
"""``_increment_deployment_usage``'s write: one pipeline, then this worker's memory holds what
Redis answered, so its usage reads see what the other workers counted."""
from litellm.types.caching import RedisPipelineIncrementOperation
operations: Final = [
RedisPipelineIncrementOperation(
key=str(key), increment_value=_amount(amount), ttl=None if ttl is None else int(ttl)
)
for key, amount in zip(keys, amounts)
]
async def _increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> None:
try:
values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations)
counted: Final = await self._stores.increment_usage(keys, amounts)
except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router
from litellm._logging import verbose_router_logger
verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e)
return
if values is None or self._cache.redis_cache is None:
return
for operation, value in zip(operations, values):
await self._cache.async_set_cache(operation["key"], int(value), local_only=True, ttl=operation["ttl"]) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
if not counted:
raise _not_usage(keys)
async def _answer_cache(self, step: Mapping[str, object]) -> object:
with service_target(_cache_target(str(step.get("key")))):
@ -3548,32 +3503,33 @@ class RustCore(
key: Final = str(step.get("key"))
ttl: Final = _ttl(step)
local: Final = step.get("local") is True
stored: Final = None if local else await self._stores.answer(step.get("op"), key, step.get("value"), ttl)
if stored is not None:
return _plain(stored.value)
match step.get("op"):
case "get":
return _plain(await self._store_for(key).async_get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
return _plain(await self._cache.async_get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case "batch_get":
return _plain(await self._batch_read([str(k) for k in _as_list(step.get("keys"))]))
case "set":
store: Final = self._store_for(key)
if ttl is None:
await store.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
await self._cache.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
else:
await store.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
await self._cache.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
return None
case "increment":
by: Final = _amount(step.get("value"))
counter: Final = self._store_for(key)
total: Final = (
await counter.async_increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType] # see above
await self._cache.async_increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType] # see above
if ttl is None
else await counter.async_increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
else await self._cache.async_increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
)
return by if total is None else total
case "increment_usage":
await self._increment_usage(_as_list(step.get("keys")), _as_list(step.get("value")), ttl)
await self._increment_usage(_as_list(step.get("keys")), _as_list(step.get("value")))
return None
case "delete":
await self._store_for(key).async_delete_cache(key)
await self._cache.async_delete_cache(key)
return None
case other:
raise RuntimeError(f"router-core asked for an unknown cache op {other!r}")
@ -3594,6 +3550,16 @@ class RustCore(
except StopIteration as finished:
return finished.value # pyright: ignore[reportAny] # the coroutine's result
@property
def _stores(self) -> RoutingStores:
"""The stores router-core's cache steps reach, over the cooldowns as last assigned."""
return RoutingStores(self._cooldowns, self._usage)
def io_token_usage_sync(self, model_group: str) -> tuple[int | None, int | None]:
"""``get_model_group_io_token_usage``, whose counters stay in the routing cache, on a sync call."""
read: Final = self.run_routing_cache_sync(lambda: self._hosted.get_model_group_io_token_usage(model_group))
return _IO_TOKEN_USAGE.validate_python(read)
def _answer_cache_sync(self, step: Mapping[str, object]) -> object:
with service_target(_cache_target(str(step.get("key")))):
return self._answer_cache_sync_in_target(step)
@ -3602,34 +3568,36 @@ class RustCore(
key: Final = str(step.get("key"))
ttl: Final = _ttl(step)
local: Final = step.get("local") is True
match step.get("op"):
op: Final = step.get("op")
stored: Final = None if local else self._stores.answer_sync(op, key, step.get("value"), ttl)
if stored is not None:
return _plain(stored.value)
match op:
case "get":
return _plain(self._store_for(key).get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
return _plain(self._cache.get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped
case "batch_get":
return _plain(self._batch_read_sync([str(k) for k in _as_list(step.get("keys"))]))
case "set":
store: Final = self._store_for(key)
if ttl is None:
store.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
self._cache.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped
else:
store.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
self._cache.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above
return None
case "increment":
by: Final = _amount(step.get("value"))
counter: Final = self._store_for(key)
# DualCache.increment_cache is annotated int but adds floats as given.
return (
counter.increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
self._cache.increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
if ttl is None
else counter.increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
else self._cache.increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above
)
case "increment_usage":
keys: Final = _as_list(step.get("keys"))
amounts: Final = _as_list(step.get("value"))
self.run_routing_cache_sync(lambda: self._increment_usage(keys, amounts, ttl))
if not self._stores.increment_usage_sync(keys, _as_list(step.get("value"))):
raise _not_usage(keys)
return None
case "delete":
self._store_for(key).delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped
self._cache.delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped
return None
case other:
raise RuntimeError(f"router-core asked for an unknown cache op {other!r}")
@ -3989,6 +3957,13 @@ def _seconds(value: object, default: float) -> float | int:
return int(default) if float(default).is_integer() else default
_IO_TOKEN_USAGE: Final = TypeAdapter(tuple[int | None, int | None])
def _not_usage(keys: Sequence[object]) -> RuntimeError:
return RuntimeError(f"router-core asked to add usage to keys that are not one deployment's tpm/rpm: {keys!r}")
def _ttl(step: Mapping[str, object]) -> float | None:
ttl: Final = step.get("ttl")
return ttl if isinstance(ttl, int | float) else None
@ -5390,11 +5365,10 @@ class _Driver:
name: Final = self._served[0]
if name is not None:
core: Final = self._core
remaining: Final = as_str_mapping(
core.run_routing_cache_sync(lambda: core.health().usage(name, kind="remaining"))
remaining: Final = core._hosted._core_get_remaining_model_group_usage_sync( # pyright: ignore[reportPrivateUsage] # the core reads its own host registry
name, core.io_token_usage_sync
)
if remaining is not None:
apply_remaining_usage_headers(ensure_response_additional_headers(response), dict(remaining))
apply_remaining_usage_headers(ensure_response_additional_headers(response), remaining)
counted: Final = self._count_response(response, step)
self.log_outcomes_sync()
return counted

View file

@ -0,0 +1,230 @@
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Final, TypeAlias
from pydantic import TypeAdapter
from litellm._internal_context import service_target, with_service_target
from litellm.caching.redis_batch import Batch
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter
from litellm.router_utils import routing_prefetch
from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
__all__ = ("AllowedFailsKey", "CooldownKey", "Handled", "RoutingStores", "UsageKey", "parse_key")
_COOLDOWN_VALUE: Final = TypeAdapter(CooldownCacheValue)
def _cooldown_value(value: object) -> CooldownCacheValue:
"""The cooldown router-core wrote, its whole-second ``cooldown_time`` kept an ``int`` as it serialized it."""
validated: Final = _COOLDOWN_VALUE.validate_python(value)
written: Final = as_str_mapping(value)
cooldown_time: Final = None if written is None else written.get("cooldown_time")
if isinstance(cooldown_time, int) and not isinstance(cooldown_time, bool):
return {**validated, "cooldown_time": cooldown_time}
return validated
@dataclass(frozen=True, slots=True)
class CooldownKey:
model_id: str
@dataclass(frozen=True, slots=True)
class AllowedFailsKey:
model_id: str
suffix: str | None
@dataclass(frozen=True, slots=True)
class UsageKey:
counter: UsageCounter
RoutingKey: TypeAlias = CooldownKey | AllowedFailsKey | UsageKey
@dataclass(frozen=True, slots=True)
class Handled:
value: object
def _deployment_key(rest: str) -> RoutingKey | None:
if rest.endswith(":cooldown"):
return CooldownKey(rest.removesuffix(":cooldown"))
model_id, marker, suffix = rest.partition(":allowed_fails")
if not marker or (suffix and not suffix.startswith(":")):
return None
return AllowedFailsKey(model_id, suffix.removeprefix(":") or None)
def _usage_key(rest: str) -> UsageKey | None:
parts: Final = rest.rsplit(":", 2)
if len(parts) != 3:
return None
head, kind, minute = parts
model_id, separator, model = head.partition(":")
if not separator:
return None
match kind:
case "tpm" | "rpm":
return UsageKey(UsageCounter(model_id, model, minute, kind))
case _:
return None
def parse_key(key: str) -> RoutingKey | None:
"""The store family a router-core key belongs to. Each store rebuilds the same key from what is parsed
here, so an id or model holding a ``:`` still names the key router-core asked for."""
if key.startswith("deployment:"):
return _deployment_key(key.removeprefix("deployment:"))
if key.startswith("global_router:"):
return _usage_key(key.removeprefix("global_router:"))
return None
def _service_target(key: RoutingKey | None) -> str | None:
if isinstance(key, CooldownKey | AllowedFailsKey):
return ROUTER_COOLDOWNS_TARGET
if isinstance(key, UsageKey):
return ROUTER_USAGE_TARGET
return None
def _cooldown_ids(keys: Sequence[str]) -> tuple[str, ...] | None:
ids: Final = tuple(parsed.model_id for parsed in map(parse_key, keys) if isinstance(parsed, CooldownKey))
return ids if keys and len(ids) == len(keys) else None
def _usage_counters(keys: Sequence[object]) -> tuple[UsageCounter, ...] | None:
counters: Final = tuple(
parsed.counter for parsed in (parse_key(str(key)) for key in keys) if isinstance(parsed, UsageKey)
)
return counters if keys and len(counters) == len(keys) else None
def _aligned(model_ids: Sequence[str], active: Sequence[tuple[str, CooldownCacheValue]]) -> tuple[object, ...]:
found: Final = dict(active)
return tuple(found.get(model_id) for model_id in model_ids)
def _one_deployment(tpm: UsageCounter, rpm: UsageCounter) -> bool:
return (tpm.kind, rpm.kind) == ("tpm", "rpm") and (tpm.deployment_id, tpm.model, tpm.minute) == (
rpm.deployment_id,
rpm.model,
rpm.minute,
)
def _deployment_increment(keys: Sequence[object], amounts: Sequence[object]) -> tuple[UsageCounter, float, int] | None:
"""``_increment_deployment_usage``'s tpm then rpm increment of one deployment."""
counters: Final = _usage_counters(keys)
if counters is None or len(counters) != 2 or len(amounts) != 2:
return None
tpm, rpm = counters
tokens, requests = amounts
if not _one_deployment(tpm, rpm) or not isinstance(tokens, int | float) or not isinstance(requests, int | float):
return None
return tpm, tokens, int(requests)
def _by(value: object) -> float:
return value if isinstance(value, int | float) and not isinstance(value, bool) else 1
class RoutingStores:
"""Router-core's cache steps for cooldowns, allowed fails and tpm/rpm usage, answered by the native stores
the Python core uses. ``None`` leaves the step to the routing cache."""
def __init__(self, cooldowns: CooldownCache, usage: RouterUsage) -> None:
self._cooldowns: Final = cooldowns
self._usage: Final = usage
def answer_sync(self, op: object, key: str, value: object, ttl: float | None) -> Handled | None:
store: Final = self._cooldowns.cooldown_store
parsed: Final = parse_key(key)
target: Final = _service_target(parsed)
if target is None:
return None
with service_target(target):
match op, parsed:
case "get", CooldownKey(model_id):
return Handled(_aligned((model_id,), store.active_sync((model_id,)))[0])
case "set", CooldownKey(model_id):
store.put(model_id, _cooldown_value(value), ttl or 0.0)
return Handled(None)
case "delete", CooldownKey(model_id):
store.evict_local(model_id)
return Handled(None)
case "increment", AllowedFailsKey(model_id, suffix):
return Handled(store.increment_allowed_fails(model_id, suffix, ttl or 0.0))
case "increment", UsageKey(counter):
return Handled(self._usage.increment_sync(counter, _by(value), int(ttl or 0)))
case _:
return None
async def answer(self, op: object, key: str, value: object, ttl: float | None) -> Handled | None:
parsed: Final = parse_key(key)
target: Final = _service_target(parsed)
if target is None:
return None
store: Final = self._cooldowns.cooldown_store
with service_target(target):
match op, parsed:
case "get", CooldownKey(model_id):
return Handled(_aligned((model_id,), await store.active((model_id,), batch=Batch.NONE))[0])
case "set", CooldownKey(model_id):
await store.async_put(model_id, _cooldown_value(value), ttl or 0.0)
return Handled(None)
case "increment", AllowedFailsKey(model_id, suffix):
return Handled(await store.async_increment_allowed_fails(model_id, suffix, ttl or 0.0))
case "increment", UsageKey(counter):
by: Final = _by(value)
return Handled(await self._usage.increment(counter, by, int(ttl or 0)))
case _:
return self.answer_sync(op, key, value, ttl)
def batch_get_sync(self, keys: Sequence[str]) -> tuple[object, ...] | None:
model_ids: Final = _cooldown_ids(keys)
if model_ids is not None:
with service_target(ROUTER_COOLDOWNS_TARGET):
return _aligned(model_ids, self._cooldowns.cooldown_store.active_sync(model_ids))
counters: Final = _usage_counters(keys)
if counters is None:
return None
with service_target(ROUTER_USAGE_TARGET):
return tuple(self._usage.get_many_sync(counters))
async def batch_get(self, keys: Sequence[str]) -> tuple[object, ...] | None:
model_ids: Final = _cooldown_ids(keys)
if model_ids is not None:
store: Final = self._cooldowns.cooldown_store
with service_target(ROUTER_COOLDOWNS_TARGET):
return _aligned(model_ids, await routing_prefetch.read_cooldowns(store, model_ids, None))
counters: Final = _usage_counters(keys)
if counters is None:
return None
with service_target(ROUTER_USAGE_TARGET):
return tuple(await self._usage.get_many(counters))
@with_service_target(ROUTER_USAGE_TARGET)
def increment_usage_sync(self, keys: Sequence[object], amounts: Sequence[object]) -> bool:
found: Final = _deployment_increment(keys, amounts)
if found is None:
return False
tpm, tokens, requests = found
self._usage.increment_deployment_sync(tpm.deployment_id, tpm.model, tpm.minute, tokens, requests)
return True
@with_service_target(ROUTER_USAGE_TARGET)
async def increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> bool:
found: Final = _deployment_increment(keys, amounts)
if found is None:
return False
tpm, tokens, requests = found
await self._usage.increment_deployment(tpm.deployment_id, tpm.model, tpm.minute, tokens, requests)
return True

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