mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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:
parent
fc8f614a03
commit
b1b236826a
192 changed files with 17208 additions and 10154 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
265
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs
Normal file
265
litellm-rust/crates/python-bridge/src/redis/store/auth_cache.rs
Normal 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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);
|
||||
|
|
|
|||
151
litellm-rust/crates/python-bridge/src/redis/store/policy.rs
Normal file
151
litellm-rust/crates/python-bridge/src/redis/store/policy.rs
Normal 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}");
|
||||
});
|
||||
}
|
||||
}
|
||||
591
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs
Normal file
591
litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs
Normal 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));
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
8
litellm-rust/crates/redis-store/src/auth_cache/keys.rs
Normal file
8
litellm-rust/crates/redis-store/src/auth_cache/keys.rs
Normal 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))
|
||||
}
|
||||
258
litellm-rust/crates/redis-store/src/auth_cache/mod.rs
Normal file
258
litellm-rust/crates/redis-store/src/auth_cache/mod.rs
Normal 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
|
||||
}
|
||||
}
|
||||
}
|
||||
35
litellm-rust/crates/redis-store/src/auth_cache/types.rs
Normal file
35
litellm-rust/crates/redis-store/src/auth_cache/types.rs
Normal 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()
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
74
litellm-rust/crates/redis-store/src/rate_limits/keys.rs
Normal file
74
litellm-rust/crates/redis-store/src/rate_limits/keys.rs
Normal 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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
260
litellm-rust/crates/redis-store/src/rate_limits/lua.rs
Normal file
260
litellm-rust/crates/redis-store/src/rate_limits/lua.rs
Normal 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
|
||||
"#;
|
||||
304
litellm-rust/crates/redis-store/src/rate_limits/mod.rs
Normal file
304
litellm-rust/crates/redis-store/src/rate_limits/mod.rs
Normal 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()
|
||||
}
|
||||
303
litellm-rust/crates/redis-store/src/rate_limits/scripts.rs
Normal file
303
litellm-rust/crates/redis-store/src/rate_limits/scripts.rs
Normal 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")),
|
||||
}
|
||||
}
|
||||
117
litellm-rust/crates/redis-store/src/rate_limits/types.rs
Normal file
117
litellm-rust/crates/redis-store/src/rate_limits/types.rs
Normal 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",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
17
litellm-rust/crates/redis-store/src/router_cooldowns/keys.rs
Normal file
17
litellm-rust/crates/redis-store/src/router_cooldowns/keys.rs
Normal 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"),
|
||||
}
|
||||
}
|
||||
205
litellm-rust/crates/redis-store/src/router_cooldowns/mod.rs
Normal file
205
litellm-rust/crates/redis-store/src/router_cooldowns/mod.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
28
litellm-rust/crates/redis-store/src/router_usage/keys.rs
Normal file
28
litellm-rust/crates/redis-store/src/router_usage/keys.rs
Normal 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
|
||||
)
|
||||
}
|
||||
}
|
||||
110
litellm-rust/crates/redis-store/src/router_usage/mod.rs
Normal file
110
litellm-rust/crates/redis-store/src/router_usage/mod.rs
Normal 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()?)
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
pub(super) fn pin(tenant: &str, session_id: &str) -> String {
|
||||
format!("{{sensitive_route:{tenant}:{session_id}}}:model")
|
||||
}
|
||||
111
litellm-rust/crates/redis-store/src/sensitive_route_pins/mod.rs
Normal file
111
litellm-rust/crates/redis-store/src/sensitive_route_pins/mod.rs
Normal 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")
|
||||
}
|
||||
85
litellm-rust/crates/redis-store/src/spend_counters/keys.rs
Normal file
85
litellm-rust/crates/redis-store/src/spend_counters/keys.rs
Normal 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()),
|
||||
})
|
||||
}
|
||||
}
|
||||
380
litellm-rust/crates/redis-store/src/spend_counters/mod.rs
Normal file
380
litellm-rust/crates/redis-store/src/spend_counters/mod.rs
Normal 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");
|
||||
}
|
||||
}
|
||||
}
|
||||
48
litellm-rust/crates/redis-store/src/spend_counters/types.rs
Normal file
48
litellm-rust/crates/redis-store/src/spend_counters/types.rs
Normal 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()
|
||||
}
|
||||
|
|
@ -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))
|
||||
}
|
||||
|
|
|
|||
741
litellm-rust/crates/redis-store/tests/auth_cache.rs
Normal file
741
litellm-rust/crates/redis-store/tests/auth_cache.rs
Normal 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);
|
||||
}
|
||||
1083
litellm-rust/crates/redis-store/tests/rate_limits.rs
Normal file
1083
litellm-rust/crates/redis-store/tests/rate_limits.rs
Normal file
File diff suppressed because it is too large
Load diff
512
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
512
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal 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())]
|
||||
);
|
||||
}
|
||||
329
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
329
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal 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))]
|
||||
);
|
||||
}
|
||||
134
litellm-rust/crates/redis-store/tests/sensitive_route_pins.rs
Normal file
134
litellm-rust/crates/redis-store/tests/sensitive_route_pins.rs
Normal 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);
|
||||
}
|
||||
648
litellm-rust/crates/redis-store/tests/spend_counters.rs
Normal file
648
litellm-rust/crates/redis-store/tests/spend_counters.rs
Normal 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);
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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__"
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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, ...],
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
146
litellm/redis/stores/auth_cache/__init__.py
Normal file
146
litellm/redis/stores/auth_cache/__init__.py
Normal 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)
|
||||
31
litellm/redis/stores/auth_cache/handle.py
Normal file
31
litellm/redis/stores/auth_cache/handle.py
Normal 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: ...
|
||||
16
litellm/redis/stores/auth_cache/types.py
Normal file
16
litellm/redis/stores/auth_cache/types.py
Normal 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
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
192
litellm/redis/stores/rate_limits/__init__.py
Normal file
192
litellm/redis/stores/rate_limits/__init__.py
Normal 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)
|
||||
68
litellm/redis/stores/rate_limits/handle.py
Normal file
68
litellm/redis/stores/rate_limits/handle.py
Normal 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: ...
|
||||
107
litellm/redis/stores/rate_limits/types.py
Normal file
107
litellm/redis/stores/rate_limits/types.py
Normal 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",
|
||||
]
|
||||
93
litellm/redis/stores/router_cooldowns/__init__.py
Normal file
93
litellm/redis/stores/router_cooldowns/__init__.py
Normal 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()
|
||||
26
litellm/redis/stores/router_cooldowns/handle.py
Normal file
26
litellm/redis/stores/router_cooldowns/handle.py
Normal 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]: ...
|
||||
17
litellm/redis/stores/router_cooldowns/types.py
Normal file
17
litellm/redis/stores/router_cooldowns/types.py
Normal 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]]
|
||||
116
litellm/redis/stores/router_usage/__init__.py
Normal file
116
litellm/redis/stores/router_usage/__init__.py
Normal 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()
|
||||
35
litellm/redis/stores/router_usage/handle.py
Normal file
35
litellm/redis/stores/router_usage/handle.py
Normal 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]: ...
|
||||
15
litellm/redis/stores/router_usage/types.py
Normal file
15
litellm/redis/stores/router_usage/types.py
Normal 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
|
||||
26
litellm/redis/stores/sensitive_route_pins/__init__.py
Normal file
26
litellm/redis/stores/sensitive_route_pins/__init__.py
Normal 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)
|
||||
9
litellm/redis/stores/sensitive_route_pins/handle.py
Normal file
9
litellm/redis/stores/sensitive_route_pins/handle.py
Normal 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]: ...
|
||||
168
litellm/redis/stores/spend_counters/__init__.py
Normal file
168
litellm/redis/stores/spend_counters/__init__.py
Normal 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
|
||||
34
litellm/redis/stores/spend_counters/handle.py
Normal file
34
litellm/redis/stores/spend_counters/handle.py
Normal 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]: ...
|
||||
66
litellm/redis/stores/spend_counters/types.py
Normal file
66
litellm/redis/stores/spend_counters/types.py
Normal 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",))
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
230
litellm/router_rust/routing_stores.py
Normal file
230
litellm/router_rust/routing_stores.py
Normal 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
Loading…
Add table
Reference in a new issue