mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(redis): run router cooldowns and usage on the Rust store
Router cooldowns and the global TPM/RPM counters move behind two Python stores backed by the new redis-store crate. The read-through memory tier lives in redis-store as MemoryTier, so crates/cache keeps no Redis dependency. Folds in the Python-only encapsulation step that used to be its own PR native_client takes the NativeRedisSettings a RedisCache captures when it is built, the kwargs, namespace and default TTL, instead of reading those attributes off the cache. A stand-in that skips the RedisCache constructor leaves the native stores on memory, and a test double that wants them on Redis builds the settings itself The stores send the commands main does: one MGET for the cooldown and usage reads instead of a GET per key, and a deployment's tpm and rpm increments on their own async_increment_pipeline round trip instead of the request's post-call batch.
This commit is contained in:
parent
f7a162a519
commit
c14b054bf8
51 changed files with 3898 additions and 2112 deletions
|
|
@ -200,6 +200,23 @@ impl<V: Clone> InMemoryCache<V> {
|
|||
.collect())
|
||||
}
|
||||
|
||||
/// Every live entry, for callers that inspect the whole tier.
|
||||
pub fn entries(&self) -> Result<Vec<(String, V)>, Error> {
|
||||
let now = (self.now)();
|
||||
let state = self.lock()?;
|
||||
Ok(state
|
||||
.values
|
||||
.iter()
|
||||
.filter(|(key, _)| {
|
||||
state
|
||||
.expirations
|
||||
.get(*key)
|
||||
.is_none_or(|expiration| *expiration >= now)
|
||||
})
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
pub fn delete_cache(&self, key: &str) -> Result<(), Error> {
|
||||
let mut state = self.lock()?;
|
||||
Self::remove(&mut state, key);
|
||||
|
|
|
|||
|
|
@ -43,6 +43,10 @@ mod _native {
|
|||
#[pymodule_export]
|
||||
use crate::redis::store::request::NativeRequestRedis;
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::router_cooldowns::{NativeDeclaredCooldowns, NativeRouterCooldowns};
|
||||
#[pymodule_export]
|
||||
use crate::redis::store::router_usage::NativeRouterUsage;
|
||||
#[pymodule_export]
|
||||
use crate::routes::audio_transcription::{atranscription, transcription};
|
||||
#[pymodule_export]
|
||||
use crate::routes::chat_completions::{acompletion, completion};
|
||||
|
|
@ -117,6 +121,9 @@ mod tests {
|
|||
"NativeDiagnosticProcessor",
|
||||
"NativeRedisClient",
|
||||
"NativeRequestRedis",
|
||||
"NativeDeclaredCooldowns",
|
||||
"NativeRouterCooldowns",
|
||||
"NativeRouterUsage",
|
||||
"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")]
|
||||
|
|
|
|||
|
|
@ -195,13 +195,7 @@ fn shared_client(
|
|||
/// The Rust client for one `RedisCache`, built once and handed to every native store, so a store
|
||||
/// never sees the redis-py kwargs or the namespace.
|
||||
#[pyclass(frozen, name = "NativeRedisClient")]
|
||||
pub(crate) struct NativeRedisClient(
|
||||
#[expect(
|
||||
dead_code,
|
||||
reason = "the native stores built on this client read it, starting in #49525"
|
||||
)]
|
||||
pub(crate) RedisClient,
|
||||
);
|
||||
pub(crate) struct NativeRedisClient(pub(crate) RedisClient);
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRedisClient {
|
||||
|
|
|
|||
|
|
@ -1,7 +1,15 @@
|
|||
mod policy;
|
||||
pub(crate) mod request;
|
||||
pub(crate) mod router_cooldowns;
|
||||
pub(crate) mod router_usage;
|
||||
|
||||
use litellm_host_python::PythonContext;
|
||||
use pyo3::prelude::*;
|
||||
use litellm_redis_store::Number;
|
||||
use pyo3::{
|
||||
exceptions::{PyRuntimeError, PyTypeError},
|
||||
prelude::*,
|
||||
types::{PyBool, PyFloat, PyInt},
|
||||
};
|
||||
|
||||
use crate::redis::events::RedisCall;
|
||||
|
||||
|
|
@ -14,3 +22,37 @@ fn call(py: Python<'_>, call_type: &'static str) -> PyResult<RedisCall> {
|
|||
caller: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn store_error(error: litellm_cache::Error) -> PyErr {
|
||||
PyRuntimeError::new_err(litellm_error::Report::new(&error).to_string())
|
||||
}
|
||||
|
||||
/// A Python `int` or `float` across the bridge, keeping which one it is.
|
||||
pub(crate) struct PyNumber(Number);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyNumber {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
if value.is_instance_of::<PyInt>() && !value.is_instance_of::<PyBool>() {
|
||||
return Ok(Self(Number::Int(value.extract()?)));
|
||||
}
|
||||
if value.is_instance_of::<PyFloat>() {
|
||||
return Ok(Self(Number::Float(value.extract()?)));
|
||||
}
|
||||
Err(PyTypeError::new_err("expected an int or a float"))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PyNumber {
|
||||
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 {
|
||||
Number::Int(value) => value.into_pyobject(py)?.into_any(),
|
||||
Number::Float(value) => value.into_pyobject(py)?.into_any(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
129
litellm-rust/crates/python-bridge/src/redis/store/policy.rs
Normal file
129
litellm-rust/crates/python-bridge/src/redis/store/policy.rs
Normal file
|
|
@ -0,0 +1,129 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_store::{TierPolicy, router_cooldowns, router_usage};
|
||||
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);
|
||||
|
||||
/// The tier values this Python process runs with. `POLICY` in each family holds `main`'s
|
||||
/// defaults, the host reports what its environment and `litellm` globals changed.
|
||||
#[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)?;
|
||||
Ok(RouterTierPolicies {
|
||||
cooldowns: TierPolicy {
|
||||
refresh_after: snapshot.read(&COOLDOWN_REFRESH)?,
|
||||
tracked_reads,
|
||||
..router_cooldowns::POLICY
|
||||
},
|
||||
usage: TierPolicy {
|
||||
refresh_after: snapshot.read(&USAGE_REFRESH)?,
|
||||
tracked_reads,
|
||||
..router_usage::POLICY
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
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 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_replace_only_the_refresh_and_tracked_reads_of_mains_defaults() {
|
||||
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",
|
||||
)
|
||||
.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.backfill_ttl,
|
||||
litellm_redis_store::router_cooldowns::POLICY.backfill_ttl
|
||||
);
|
||||
assert_eq!(
|
||||
policies.usage.backfill_ttl,
|
||||
litellm_redis_store::router_usage::POLICY.backfill_ttl
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[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",
|
||||
)
|
||||
.unwrap_err();
|
||||
let text = projected(
|
||||
py,
|
||||
"cooldown_redis_read_interval=1, usage_redis_read_interval=10, max_redis_batch_cache_size='many'",
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(
|
||||
negative.contains("cooldown_redis_read_interval"),
|
||||
"{negative}"
|
||||
);
|
||||
assert!(text.contains("max_redis_batch_cache_size"), "{text}");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,233 @@
|
|||
use std::{
|
||||
sync::{Arc, Mutex, PoisonError},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_redis_store::router_cooldowns::{Cooldown, DeclaredCooldowns, RouterCooldowns};
|
||||
use pyo3::{
|
||||
exceptions::PyRuntimeError,
|
||||
prelude::*,
|
||||
types::{PyDict, PyTuple},
|
||||
};
|
||||
|
||||
use super::{
|
||||
PyNumber, call, policy::router_tier_policies, request::NativeRequestRedis, store_error,
|
||||
};
|
||||
use crate::execution::{run_async_value, run_sync_value};
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
/// One router's cooldown store. The router's sync completion path calls `put`, `active`,
|
||||
/// `minimum` and `increment_allowed_fails`, so those block with the GIL released like redis-py.
|
||||
#[pyclass(frozen, name = "NativeRouterCooldowns")]
|
||||
pub(crate) struct NativeRouterCooldowns(Arc<RouterCooldowns>);
|
||||
|
||||
pub(crate) struct PyCooldown(Cooldown);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let fields = value.cast::<PyDict>()?;
|
||||
let item = |name: &str| fields.as_any().get_item(name);
|
||||
Ok(Self(Cooldown {
|
||||
exception_received: item("exception_received")?.extract()?,
|
||||
status_code: item("status_code")?.extract()?,
|
||||
timestamp: item("timestamp")?.extract::<PyNumber>()?.0,
|
||||
cooldown_time: item("cooldown_time")?.extract::<PyNumber>()?.0,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
impl<'py> IntoPyObject<'py> for PyCooldown {
|
||||
type Target = PyDict;
|
||||
type Output = Bound<'py, PyDict>;
|
||||
type Error = PyErr;
|
||||
|
||||
fn into_pyobject(self, py: Python<'py>) -> PyResult<Self::Output> {
|
||||
let fields = PyDict::new(py);
|
||||
fields.set_item("exception_received", self.0.exception_received)?;
|
||||
fields.set_item("status_code", self.0.status_code)?;
|
||||
fields.set_item("timestamp", PyNumber(self.0.timestamp))?;
|
||||
fields.set_item("cooldown_time", PyNumber(self.0.cooldown_time))?;
|
||||
Ok(fields)
|
||||
}
|
||||
}
|
||||
|
||||
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(),
|
||||
PyCooldown(cooldown).into_pyobject(py)?.into_any(),
|
||||
],
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRouterCooldowns {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
|
||||
let client = client.map(|client| client.0.clone());
|
||||
let policy = router_tier_policies(py)?.cooldowns;
|
||||
Ok(Self(Arc::new(RouterCooldowns::new(client, policy))))
|
||||
}
|
||||
|
||||
fn put(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
model_id: String,
|
||||
cooldown: PyCooldown,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<()> {
|
||||
let store = self.0.clone();
|
||||
let scope = call(py, "set_cache")?;
|
||||
run_sync_value(py, async move {
|
||||
scope
|
||||
.scope(store.put(
|
||||
&model_id,
|
||||
&cooldown.0,
|
||||
Duration::from_secs_f64(ttl_seconds.max(0.0)),
|
||||
))
|
||||
.await
|
||||
.map_err(store_error)
|
||||
})
|
||||
}
|
||||
|
||||
fn increment_allowed_fails(
|
||||
&self,
|
||||
py: Python<'_>,
|
||||
model_id: String,
|
||||
suffix: Option<String>,
|
||||
ttl_seconds: f64,
|
||||
) -> PyResult<i64> {
|
||||
let store = self.0.clone();
|
||||
let scope = call(py, "increment_cache")?;
|
||||
run_sync_value(py, async move {
|
||||
scope
|
||||
.scope(store.increment_allowed_fails(
|
||||
&model_id,
|
||||
suffix.as_deref(),
|
||||
Duration::from_secs_f64(ttl_seconds.max(0.0)),
|
||||
))
|
||||
.await
|
||||
.map_err(store_error)
|
||||
})
|
||||
}
|
||||
|
||||
fn active<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
model_ids: Vec<String>,
|
||||
) -> PyResult<Vec<Bound<'py, PyTuple>>> {
|
||||
let store = self.0.clone();
|
||||
let scope = call(py, "batch_get_cache")?;
|
||||
let active = run_sync_value(py, async move {
|
||||
scope
|
||||
.scope(store.active(&model_ids))
|
||||
.await
|
||||
.map_err(store_error)
|
||||
})?;
|
||||
active_to_py(py, active)
|
||||
}
|
||||
|
||||
fn aactive<'py>(&self, py: Python<'py>, model_ids: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let scope = call(py, "async_batch_get_cache")?;
|
||||
run_async_value(py, async move {
|
||||
let active = scope
|
||||
.scope(store.active(&model_ids))
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
Python::attach(|py| {
|
||||
Ok(active_to_py(py, active)?
|
||||
.into_iter()
|
||||
.map(Bound::unbind)
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
fn declare_active(
|
||||
&self,
|
||||
request: &NativeRequestRedis,
|
||||
model_ids: Vec<String>,
|
||||
) -> PyResult<NativeDeclaredCooldowns> {
|
||||
let declared = self
|
||||
.0
|
||||
.declare_active(&request.0, &model_ids)
|
||||
.map_err(store_error)?;
|
||||
Ok(NativeDeclaredCooldowns {
|
||||
store: self.0.clone(),
|
||||
declared: Mutex::new(Some(declared)),
|
||||
})
|
||||
}
|
||||
|
||||
fn minimum(&self, py: Python<'_>, model_ids: Vec<String>) -> PyResult<Option<PyNumber>> {
|
||||
let store = self.0.clone();
|
||||
let scope = call(py, "batch_get_cache")?;
|
||||
let minimum = run_sync_value(py, async move {
|
||||
scope
|
||||
.scope(store.minimum(&model_ids))
|
||||
.await
|
||||
.map_err(store_error)
|
||||
})?;
|
||||
Ok(minimum.map(PyNumber))
|
||||
}
|
||||
|
||||
fn evict_local(&self, model_id: &str) -> PyResult<()> {
|
||||
self.0.evict_local(model_id).map_err(store_error)
|
||||
}
|
||||
|
||||
fn flush_local_counters(&self) -> PyResult<()> {
|
||||
self.0.flush_local_counters().map_err(store_error)
|
||||
}
|
||||
|
||||
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, cooldown) in self.0.snapshot().map_err(store_error)? {
|
||||
snapshot.set_item(key, PyCooldown(cooldown))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
}
|
||||
|
||||
/// An `active` read declared on a request's pre-call batch. `resolve` answers it once.
|
||||
#[pyclass(frozen, name = "NativeDeclaredCooldowns")]
|
||||
pub(crate) struct NativeDeclaredCooldowns {
|
||||
store: Arc<RouterCooldowns>,
|
||||
declared: Mutex<Option<DeclaredCooldowns>>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeDeclaredCooldowns {
|
||||
fn resolve<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let declared = self
|
||||
.declared
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.take()
|
||||
.ok_or_else(|| PyRuntimeError::new_err("this cooldown read was already resolved"))?;
|
||||
let store = self.store.clone();
|
||||
let scope = call(py, "async_batch_get_cache")?;
|
||||
run_async_value(py, async move {
|
||||
let active = scope
|
||||
.scope(store.resolve_active(declared))
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
Python::attach(|py| {
|
||||
Ok(active_to_py(py, active)?
|
||||
.into_iter()
|
||||
.map(Bound::unbind)
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage};
|
||||
use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict};
|
||||
|
||||
use super::{PyNumber, call, policy::router_tier_policies, store_error};
|
||||
use crate::execution::run_async_value;
|
||||
use crate::redis::client::NativeRedisClient;
|
||||
|
||||
/// One router's global TPM/RPM counters. Every caller on `main` is async.
|
||||
#[pyclass(frozen, name = "NativeRouterUsage")]
|
||||
pub(crate) struct NativeRouterUsage(Arc<RouterUsage>);
|
||||
|
||||
/// `(deployment_id, model, minute, "tpm" | "rpm")`.
|
||||
pub(crate) struct PyCounter(Counter);
|
||||
|
||||
impl<'a, 'py> FromPyObject<'a, 'py> for PyCounter {
|
||||
type Error = PyErr;
|
||||
|
||||
fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
|
||||
let (deployment_id, model, minute, kind): (String, Option<String>, String, String) =
|
||||
value.extract()?;
|
||||
let kind = match kind.as_str() {
|
||||
"tpm" => Kind::Tpm,
|
||||
"rpm" => Kind::Rpm,
|
||||
_ => {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"unknown usage counter kind {kind:?}"
|
||||
)));
|
||||
}
|
||||
};
|
||||
Ok(Self(Counter {
|
||||
deployment_id,
|
||||
model,
|
||||
minute,
|
||||
kind,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
impl NativeRouterUsage {
|
||||
#[new]
|
||||
fn new(py: Python<'_>, client: Option<PyRef<'_, NativeRedisClient>>) -> PyResult<Self> {
|
||||
let client = client.map(|client| client.0.clone());
|
||||
let policy = router_tier_policies(py)?.usage;
|
||||
Ok(Self(Arc::new(RouterUsage::new(client, policy))))
|
||||
}
|
||||
|
||||
fn increment<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counter: PyCounter,
|
||||
value: PyNumber,
|
||||
ttl_seconds: u64,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let scope = call(py, "async_increment")?;
|
||||
run_async_value(py, async move {
|
||||
scope
|
||||
.scope(store.increment(&counter.0, value.0, Duration::from_secs(ttl_seconds)))
|
||||
.await
|
||||
.map(PyNumber)
|
||||
.map_err(store_error)
|
||||
})
|
||||
}
|
||||
|
||||
/// `_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();
|
||||
let scope = call(py, "async_increment_pipeline")?;
|
||||
run_async_value(py, async move {
|
||||
scope
|
||||
.scope(store.increment_deployment(
|
||||
&deployment_id,
|
||||
model.as_deref(),
|
||||
&minute,
|
||||
tokens.0,
|
||||
requests,
|
||||
))
|
||||
.await
|
||||
.map_err(store_error)
|
||||
})
|
||||
}
|
||||
|
||||
fn get_many<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
counters: Vec<PyCounter>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let store = self.0.clone();
|
||||
let counters = counters
|
||||
.into_iter()
|
||||
.map(|counter| counter.0)
|
||||
.collect::<Vec<_>>();
|
||||
let scope = call(py, "async_batch_get_cache")?;
|
||||
run_async_value(py, async move {
|
||||
let read = scope
|
||||
.scope(store.get_many(&counters))
|
||||
.await
|
||||
.map_err(store_error)?;
|
||||
Ok(read
|
||||
.into_iter()
|
||||
.map(|count| count.map(PyNumber))
|
||||
.collect::<Vec<_>>())
|
||||
})
|
||||
}
|
||||
|
||||
fn flush_local_counters(&self) -> PyResult<()> {
|
||||
self.0.flush_local_counters().map_err(store_error)
|
||||
}
|
||||
|
||||
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, count) in self.0.snapshot().map_err(store_error)? {
|
||||
snapshot.set_item(key, PyNumber(count))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
}
|
||||
}
|
||||
|
|
@ -3,5 +3,6 @@
|
|||
- A read-through family builds its own `MemoryTier` (`tier.rs`, `main`'s `DualCache` over a `RedisCache`) from a `TierPolicy` the host passes in. Its `pub const POLICY` is `main`'s defaults, and the Python host overrides what `main` reads from the environment or `litellm` globals (`litellm/rust_bridge/host/settings.py`, `router_tiers`). Callers pass a `RedisClient` (or none), never a cache. Redis-only families are plain functions over `&RedisClient`
|
||||
- Where `main` swallows a Redis failure and serves memory, the family does the same and logs. It never turns a failure into a miss that `main` would not
|
||||
- `tests/fixtures/router_wire.json` is the wire contract, checked here against real Redis and in `tests/unit/redis/stores/` against Python's own codec. Change both sides together
|
||||
- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. A family exposes a request-aware variant that declares on it (`declare_active`, `increment_deployment_post_call`) next to its direct method, which runs the same declaration on a batch of its own
|
||||
- `RequestRedis` is one request's batches: a pre-call and a post-call `RedisBatch` per client, the post one flushed by `flush_post` or its deadline. A family exposes a request-aware variant that declares on it (`declare_active`) next to its direct method, which runs the same declaration on a batch of its own. Only what main's Python puts on its request pipeline rides these batches. An op main sends on its own round trip, like a deployment's tpm and rpm increment, stays on its own round trip here
|
||||
- Every store sends the commands main's `RedisCache` sends for that operation: `MGET` where Python reads several keys at once, `SET ... EX`, and `INCRBYFLOAT` with `EXPIRE` or `TTL`. A `MONITOR` diff against main is a parity check, and any difference is a defect
|
||||
- Tests run against `litellm-redis-testing` on standalone and cluster. A test is single-topology only for outages
|
||||
|
|
|
|||
20
litellm-rust/crates/redis-store/src/codec.rs
Normal file
20
litellm-rust/crates/redis-store/src/codec.rs
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
use litellm_python_compat::{Value, json, literal};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DecodeError {
|
||||
#[error("Redis value is not UTF-8")]
|
||||
NotUtf8,
|
||||
#[error("Redis value is neither JSON nor a Python literal: {0}")]
|
||||
Unparsed(#[source] litellm_python_compat::Error),
|
||||
#[error("Redis value has the wrong shape: {0}")]
|
||||
Shape(&'static str),
|
||||
}
|
||||
|
||||
/// `RedisCache._get_cache_logic`: `json.loads`, falling back to `ast.literal_eval`.
|
||||
pub(crate) fn python_value(raw: &[u8]) -> Result<Value, DecodeError> {
|
||||
let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?;
|
||||
match serde_json::from_str::<serde_json::Value>(text) {
|
||||
Ok(document) => Ok(json::from_json(document)),
|
||||
Err(_) => literal::literal_eval(text).map_err(DecodeError::Unparsed),
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,11 @@
|
|||
mod codec;
|
||||
mod number;
|
||||
mod request;
|
||||
pub mod router_cooldowns;
|
||||
pub mod router_usage;
|
||||
mod tier;
|
||||
|
||||
pub use codec::DecodeError;
|
||||
pub use number::Number;
|
||||
pub use request::RequestRedis;
|
||||
pub use tier::{Clock, TierPolicy};
|
||||
|
|
|
|||
77
litellm-rust/crates/redis-store/src/number.rs
Normal file
77
litellm-rust/crates/redis-store/src/number.rs
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
use std::{cmp::Ordering, ops::Add};
|
||||
|
||||
use litellm_python_compat::{BigInt, Value};
|
||||
|
||||
/// A Python `int` or `float`, kept apart because `repr` and `json.dumps` write them differently.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
pub enum Number {
|
||||
Int(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl From<Number> for litellm_redis_client::Amount {
|
||||
fn from(number: Number) -> Self {
|
||||
match number {
|
||||
Number::Int(value) => Self::Int(value),
|
||||
Number::Float(value) => Self::Float(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Number {
|
||||
pub fn as_f64(self) -> f64 {
|
||||
match self {
|
||||
Self::Int(value) => value as f64,
|
||||
Self::Float(value) => value,
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `int(value)`, which truncates toward zero.
|
||||
pub fn truncate(self) -> i64 {
|
||||
match self {
|
||||
Self::Int(value) => value,
|
||||
Self::Float(value) => value.trunc() as i64,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_value(self) -> Value {
|
||||
match self {
|
||||
Self::Int(value) => Value::Int(BigInt::from(value)),
|
||||
Self::Float(value) => Value::Float(value),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_value(value: &Value) -> Option<Self> {
|
||||
match value {
|
||||
Value::Int(value) => i64::try_from(value).ok().map(Self::Int),
|
||||
Value::Float(value) => Some(Self::Float(*value)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `min` over numbers: the first of the smallest, keeping its type.
|
||||
pub fn min(values: impl IntoIterator<Item = Self>) -> Option<Self> {
|
||||
values
|
||||
.into_iter()
|
||||
.fold(None, |smallest, value| match smallest {
|
||||
Some(smallest)
|
||||
if value.as_f64().partial_cmp(&smallest.as_f64()) != Some(Ordering::Less) =>
|
||||
{
|
||||
Some(smallest)
|
||||
}
|
||||
_ => Some(value),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `+`: two ints stay an int, anything with a float is a float.
|
||||
impl Add for Number {
|
||||
type Output = Self;
|
||||
|
||||
fn add(self, other: Self) -> Self {
|
||||
match (self, other) {
|
||||
(Self::Int(left), Self::Int(right)) => Self::Int(left + right),
|
||||
(left, right) => Self::Float(left.as_f64() + right.as_f64()),
|
||||
}
|
||||
}
|
||||
}
|
||||
281
litellm-rust/crates/redis-store/src/router_cooldowns.rs
Normal file
281
litellm-rust/crates/redis-store/src/router_cooldowns.rs
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::Error as MemoryError;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_python_compat::{Value, repr::repr};
|
||||
use litellm_redis_client::{RedisBatch, RedisClient};
|
||||
|
||||
use crate::{
|
||||
DecodeError, Number, RequestRedis,
|
||||
codec::python_value,
|
||||
tier::{Clock, DeclaredRead, MemoryTier, OnBreakerOpen, TierPolicy, system_clock},
|
||||
};
|
||||
|
||||
/// `main`'s cooldown tier: `DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS` (1 s), `InMemoryCache`'s
|
||||
/// 600 s default for backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
|
||||
pub const POLICY: TierPolicy = TierPolicy {
|
||||
refresh_after: Duration::from_secs(1),
|
||||
backfill_ttl: Duration::from_secs(600),
|
||||
on_breaker_open: OnBreakerOpen::ServeMemory,
|
||||
tracked_reads: 1000,
|
||||
};
|
||||
const MEMORY_ENTRIES: usize = 200;
|
||||
const MAX_CORRECTED_TTL: Duration = Duration::from_secs(60);
|
||||
const CORRECTION_SLACK: Duration = Duration::from_secs(5);
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Cooldown {
|
||||
pub exception_received: String,
|
||||
pub status_code: String,
|
||||
pub timestamp: Number,
|
||||
pub cooldown_time: Number,
|
||||
}
|
||||
|
||||
impl Cooldown {
|
||||
fn remaining(&self, now: Duration) -> f64 {
|
||||
self.timestamp.as_f64() + self.cooldown_time.as_f64() - now.as_secs_f64()
|
||||
}
|
||||
|
||||
/// `str(dict)`, the bytes `RedisCache.set_cache` writes on `main`.
|
||||
pub fn encode(&self) -> String {
|
||||
let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value);
|
||||
repr(&Value::Dict(vec![
|
||||
field(
|
||||
"exception_received",
|
||||
Value::Str(self.exception_received.clone()),
|
||||
),
|
||||
field("status_code", Value::Str(self.status_code.clone())),
|
||||
field("timestamp", self.timestamp.to_value()),
|
||||
field("cooldown_time", self.cooldown_time.to_value()),
|
||||
]))
|
||||
}
|
||||
|
||||
/// Reads both `main`'s `str(dict)` and a JSON object, as `RedisCache._get_cache_logic` does.
|
||||
pub fn decode(raw: &[u8]) -> Result<Self, DecodeError> {
|
||||
let Value::Dict(fields) = python_value(raw)? 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")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn key(model_id: &str) -> String {
|
||||
format!("deployment:{model_id}:cooldown")
|
||||
}
|
||||
|
||||
fn keys(model_ids: &[String]) -> Vec<String> {
|
||||
model_ids.iter().map(|model_id| key(model_id)).collect()
|
||||
}
|
||||
|
||||
fn allowed_fails_key(model_id: &str, suffix: Option<&str>) -> String {
|
||||
match suffix.filter(|suffix| !suffix.is_empty()) {
|
||||
Some(suffix) => format!("deployment:{model_id}:allowed_fails:{suffix}"),
|
||||
None => format!("deployment:{model_id}:allowed_fails"),
|
||||
}
|
||||
}
|
||||
|
||||
fn memory<V: Clone>(policy: TierPolicy, now: &Clock) -> InMemoryCache<V> {
|
||||
let now = now.clone();
|
||||
InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || {
|
||||
now()
|
||||
})
|
||||
}
|
||||
|
||||
/// `BaseCache.get_ttl` sends `int(ttl)`, so Redis sees whole seconds.
|
||||
fn whole_seconds(ttl: Duration) -> Duration {
|
||||
Duration::from_secs(ttl.as_secs())
|
||||
}
|
||||
|
||||
/// An `active` read declared on a request batch, answered by `resolve_active`.
|
||||
pub struct DeclaredCooldowns {
|
||||
model_ids: Vec<String>,
|
||||
read: DeclaredRead,
|
||||
}
|
||||
|
||||
/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod.
|
||||
pub struct RouterCooldowns {
|
||||
cooldowns: MemoryTier<InMemoryCache<Cooldown>, RedisClient>,
|
||||
fails: MemoryTier<InMemoryCache<i64>, RedisClient>,
|
||||
}
|
||||
|
||||
impl RouterCooldowns {
|
||||
pub fn new(client: Option<RedisClient>, policy: TierPolicy) -> Self {
|
||||
Self::with_clock(client, policy, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(client: Option<RedisClient>, policy: TierPolicy, now: Clock) -> Self {
|
||||
Self {
|
||||
cooldowns: MemoryTier::read_through(
|
||||
memory(policy, &now),
|
||||
client.clone(),
|
||||
policy,
|
||||
now.clone(),
|
||||
),
|
||||
fails: MemoryTier::read_through(memory(policy, &now), client, policy, now),
|
||||
}
|
||||
}
|
||||
|
||||
/// Memory first, then a `SET ... EX int(ttl)` whose failure only logs, like `DualCache.set_cache`.
|
||||
pub async fn put(
|
||||
&self,
|
||||
model_id: &str,
|
||||
cooldown: &Cooldown,
|
||||
ttl: Duration,
|
||||
) -> Result<(), MemoryError> {
|
||||
let key = key(model_id);
|
||||
self.cooldowns.set_memory(&key, cooldown.clone(), ttl)?;
|
||||
if let Some(remote) = self.cooldowns.remote()
|
||||
&& let Err(error) = remote
|
||||
.set(&key, cooldown.encode(), whole_seconds(ttl))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(%error, "writing a cooldown to Redis failed");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The fleet-wide fail count. A Redis failure falls back to this router's own count.
|
||||
pub async fn increment_allowed_fails(
|
||||
&self,
|
||||
model_id: &str,
|
||||
suffix: Option<&str>,
|
||||
ttl: Duration,
|
||||
) -> Result<i64, MemoryError> {
|
||||
let key = allowed_fails_key(model_id, suffix);
|
||||
let local = self
|
||||
.fails
|
||||
.update_memory(&key, ttl, |count| count.unwrap_or(0) + 1)?;
|
||||
let Some(remote) = self.fails.remote() else {
|
||||
return Ok(local);
|
||||
};
|
||||
match remote.incr_by(&key, 1, whole_seconds(ttl)).await {
|
||||
Ok(count) => Ok(count),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "allowed_fails counter fell back to this worker's in-memory count");
|
||||
Ok(local)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The deployments still cooling down. An entry whose cooldown already ended is dropped
|
||||
/// from memory, and one memory would keep past its end is shortened, at most to 60 s.
|
||||
pub async fn active(
|
||||
&self,
|
||||
model_ids: &[String],
|
||||
) -> Result<Vec<(String, Cooldown)>, MemoryError> {
|
||||
let declared = self.declare(RedisClient::batch, model_ids)?;
|
||||
self.resolve_active(declared).await
|
||||
}
|
||||
|
||||
/// `active`, with its Redis read declared on the request's pre-call batch.
|
||||
pub fn declare_active(
|
||||
&self,
|
||||
request: &RequestRedis,
|
||||
model_ids: &[String],
|
||||
) -> Result<DeclaredCooldowns, MemoryError> {
|
||||
self.declare(|client| request.pre(client), model_ids)
|
||||
}
|
||||
|
||||
pub async fn resolve_active(
|
||||
&self,
|
||||
declared: DeclaredCooldowns,
|
||||
) -> Result<Vec<(String, Cooldown)>, MemoryError> {
|
||||
let DeclaredCooldowns { model_ids, read } = declared;
|
||||
let read = self.cooldowns.resolve_many(read, Cooldown::decode).await?;
|
||||
let now = self.cooldowns.now();
|
||||
let corrected = model_ids
|
||||
.into_iter()
|
||||
.zip(read)
|
||||
.filter_map(|(model_id, cooldown)| cooldown.map(|cooldown| (model_id, cooldown)))
|
||||
.map(|(model_id, cooldown)| {
|
||||
Ok(self
|
||||
.correct(&model_id, &cooldown, now)?
|
||||
.then_some((model_id, cooldown)))
|
||||
})
|
||||
.collect::<Result<Vec<_>, MemoryError>>()?;
|
||||
Ok(corrected.into_iter().flatten().collect())
|
||||
}
|
||||
|
||||
/// The shortest stored cooldown, ended ones included, as `main`'s `get_min_cooldown` reads it.
|
||||
pub async fn minimum(&self, model_ids: &[String]) -> Result<Option<Number>, MemoryError> {
|
||||
let read = self.read(model_ids).await?;
|
||||
Ok(Number::min(
|
||||
read.into_iter()
|
||||
.flatten()
|
||||
.map(|cooldown| cooldown.cooldown_time),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn evict_local(&self, model_id: &str) -> Result<(), MemoryError> {
|
||||
self.cooldowns.delete_memory(&key(model_id))
|
||||
}
|
||||
|
||||
pub fn flush_local_counters(&self) -> Result<(), MemoryError> {
|
||||
self.fails.flush_memory()
|
||||
}
|
||||
|
||||
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
|
||||
/// in-memory router state.
|
||||
pub fn snapshot(&self) -> Result<Vec<(String, Cooldown)>, MemoryError> {
|
||||
self.cooldowns.memory().entries()
|
||||
}
|
||||
|
||||
async fn read(&self, model_ids: &[String]) -> Result<Vec<Option<Cooldown>>, MemoryError> {
|
||||
self.cooldowns
|
||||
.get_many(&keys(model_ids), Cooldown::decode)
|
||||
.await
|
||||
}
|
||||
|
||||
fn declare(
|
||||
&self,
|
||||
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
|
||||
model_ids: &[String],
|
||||
) -> Result<DeclaredCooldowns, MemoryError> {
|
||||
Ok(DeclaredCooldowns {
|
||||
model_ids: model_ids.to_vec(),
|
||||
read: self.cooldowns.declare_many(batch_for, &keys(model_ids))?,
|
||||
})
|
||||
}
|
||||
|
||||
fn correct(
|
||||
&self,
|
||||
model_id: &str,
|
||||
cooldown: &Cooldown,
|
||||
now: Duration,
|
||||
) -> Result<bool, MemoryError> {
|
||||
let key = key(model_id);
|
||||
let remaining = cooldown.remaining(now);
|
||||
if remaining <= 0.0 {
|
||||
self.cooldowns.delete_memory(&key)?;
|
||||
return Ok(false);
|
||||
}
|
||||
let remaining = Duration::from_secs_f64(remaining);
|
||||
let memory = self.cooldowns.memory();
|
||||
if memory.expires_at(&key)? > Some(now + remaining + CORRECTION_SLACK) {
|
||||
self.cooldowns.delete_memory(&key)?;
|
||||
self.cooldowns
|
||||
.set_memory(&key, cooldown.clone(), remaining.min(MAX_CORRECTED_TTL))?;
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
}
|
||||
222
litellm-rust/crates/redis-store/src/router_usage.rs
Normal file
222
litellm-rust/crates/redis-store/src/router_usage.rs
Normal file
|
|
@ -0,0 +1,222 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use litellm_cache::Error as MemoryError;
|
||||
use litellm_cache_memory::InMemoryCache;
|
||||
use litellm_redis_client::{Error, Expiry, RedisClient, RedisError, RedisErrorKind, Value};
|
||||
|
||||
use crate::{
|
||||
DecodeError, Number,
|
||||
codec::python_value,
|
||||
tier::{Clock, MemoryTier, OnBreakerOpen, TierPolicy, system_clock},
|
||||
};
|
||||
|
||||
/// `main`'s usage tier: `DualCache`'s 10 s default refresh, `InMemoryCache`'s 600 s default for
|
||||
/// backfilled entries and `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
|
||||
pub const POLICY: TierPolicy = TierPolicy {
|
||||
refresh_after: Duration::from_secs(10),
|
||||
backfill_ttl: Duration::from_secs(600),
|
||||
on_breaker_open: OnBreakerOpen::ServeMemory,
|
||||
tracked_reads: 1000,
|
||||
};
|
||||
const MEMORY_ENTRIES: usize = 200;
|
||||
/// `RoutingArgs.ttl`, the window of a per-minute counter.
|
||||
pub const COUNTER_TTL: Duration = Duration::from_secs(60);
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Kind {
|
||||
Tpm,
|
||||
Rpm,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Counter {
|
||||
pub deployment_id: String,
|
||||
pub model: Option<String>,
|
||||
pub minute: String,
|
||||
pub kind: Kind,
|
||||
}
|
||||
|
||||
impl Counter {
|
||||
/// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`.
|
||||
fn key(&self) -> String {
|
||||
let kind = match self.kind {
|
||||
Kind::Tpm => "tpm",
|
||||
Kind::Rpm => "rpm",
|
||||
};
|
||||
let model = self.model.as_deref().unwrap_or("None");
|
||||
format!(
|
||||
"global_router:{}:{model}:{kind}:{}",
|
||||
self.deployment_id, self.minute
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn decode(raw: &[u8]) -> Result<Number, DecodeError> {
|
||||
Number::from_value(&python_value(raw)?).ok_or(DecodeError::Shape("a usage counter is a number"))
|
||||
}
|
||||
|
||||
fn whole_seconds(ttl: Duration) -> Duration {
|
||||
Duration::from_secs(ttl.as_secs())
|
||||
}
|
||||
|
||||
/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod.
|
||||
pub struct RouterUsage(MemoryTier<InMemoryCache<Number>, RedisClient>);
|
||||
|
||||
impl RouterUsage {
|
||||
pub fn new(client: Option<RedisClient>, policy: TierPolicy) -> Self {
|
||||
Self::with_clock(client, policy, system_clock())
|
||||
}
|
||||
|
||||
pub fn with_clock(client: Option<RedisClient>, policy: TierPolicy, now: Clock) -> Self {
|
||||
let clock = now.clone();
|
||||
let memory =
|
||||
InMemoryCache::with_clock(Some(MEMORY_ENTRIES), Some(policy.backfill_ttl), move || {
|
||||
clock()
|
||||
});
|
||||
Self(MemoryTier::read_through(memory, client, policy, now))
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache`: memory, then `INCRBYFLOAT` keeping an existing TTL.
|
||||
/// A Redis failure answers with the memory count.
|
||||
pub async fn increment(
|
||||
&self,
|
||||
counter: &Counter,
|
||||
value: Number,
|
||||
ttl: Duration,
|
||||
) -> Result<Number, MemoryError> {
|
||||
let key = counter.key();
|
||||
let local = self.add_local(&key, value, ttl)?;
|
||||
let Some(remote) = self.0.remote() else {
|
||||
return Ok(local);
|
||||
};
|
||||
match remote
|
||||
.incr_by_float(&key, value.into(), Expiry::KeepExisting(whole_seconds(ttl)))
|
||||
.await
|
||||
{
|
||||
Ok(total) => Ok(Number::Float(total)),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "Redis usage increment failed, falling back to the in-memory count");
|
||||
Ok(local)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One deployment's tokens and requests, `async_increment_pipeline` on its own round trip:
|
||||
/// memory, then `INCRBYFLOAT` and `EXPIRE 60` for each counter, mirrored back into memory as
|
||||
/// `int` totals.
|
||||
pub async fn increment_deployment(
|
||||
&self,
|
||||
deployment_id: &str,
|
||||
model: Option<&str>,
|
||||
minute: &str,
|
||||
tokens: Number,
|
||||
requests: i64,
|
||||
) -> Result<(), MemoryError> {
|
||||
let increment = Increment::new(deployment_id, model, minute, tokens, requests);
|
||||
let local = increment
|
||||
.counts
|
||||
.iter()
|
||||
.map(|(key, value)| self.add_local(key, *value, COUNTER_TTL))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
let Some(remote) = self.0.remote() else {
|
||||
return Ok(());
|
||||
};
|
||||
let pipeline = increment
|
||||
.counts
|
||||
.iter()
|
||||
.fold(remote.pipeline(), |pipeline, (key, value)| {
|
||||
pipeline.incr_by_float(key, (*value).into(), Some(COUNTER_TTL))
|
||||
});
|
||||
let totals = remote
|
||||
.run(&pipeline)
|
||||
.await
|
||||
.into_iter()
|
||||
.map(|total| total.and_then(float_reply))
|
||||
.collect();
|
||||
let keys: Vec<String> = increment.counts.into_iter().map(|(key, _)| key).collect();
|
||||
self.mirror(&keys, local, totals)
|
||||
}
|
||||
|
||||
pub async fn get_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, MemoryError> {
|
||||
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
|
||||
self.0.get_many(&keys, decode).await
|
||||
}
|
||||
|
||||
pub fn flush_local_counters(&self) -> Result<(), MemoryError> {
|
||||
self.0.flush_memory()
|
||||
}
|
||||
|
||||
/// The memory tier keyed as Python's `InMemoryCache` was, for the fuzz probe that reads
|
||||
/// in-memory router state.
|
||||
pub fn snapshot(&self) -> Result<Vec<(String, Number)>, MemoryError> {
|
||||
self.0.memory().entries()
|
||||
}
|
||||
|
||||
fn mirror(
|
||||
&self,
|
||||
keys: &[String],
|
||||
local: Vec<Number>,
|
||||
totals: Result<Vec<f64>, Error>,
|
||||
) -> Result<(), MemoryError> {
|
||||
let totals = match totals {
|
||||
Ok(totals) => totals.into_iter().map(Number::Float).collect(),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "Redis usage pipeline failed, keeping the in-memory counts");
|
||||
local
|
||||
}
|
||||
};
|
||||
keys.iter().zip(totals).try_for_each(|(key, total)| {
|
||||
self.0
|
||||
.set_memory(key, Number::Int(total.truncate()), COUNTER_TTL)
|
||||
})
|
||||
}
|
||||
|
||||
fn add_local(&self, key: &str, value: Number, ttl: Duration) -> Result<Number, MemoryError> {
|
||||
self.0.update_memory(key, ttl, |current| {
|
||||
current.unwrap_or(Number::Int(0)) + value
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Increment {
|
||||
counts: [(String, Number); 2],
|
||||
}
|
||||
|
||||
impl Increment {
|
||||
fn new(
|
||||
deployment_id: &str,
|
||||
model: Option<&str>,
|
||||
minute: &str,
|
||||
tokens: Number,
|
||||
requests: i64,
|
||||
) -> Self {
|
||||
let counter = |kind| Counter {
|
||||
deployment_id: deployment_id.to_owned(),
|
||||
model: model.map(str::to_owned),
|
||||
minute: minute.to_owned(),
|
||||
kind,
|
||||
};
|
||||
Self {
|
||||
counts: [
|
||||
(counter(Kind::Tpm).key(), tokens),
|
||||
(counter(Kind::Rpm).key(), Number::Int(requests)),
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn float_reply(value: Value) -> Result<f64, Error> {
|
||||
let total = match &value {
|
||||
Value::BulkString(bytes) => std::str::from_utf8(bytes)
|
||||
.ok()
|
||||
.and_then(|text| text.parse().ok()),
|
||||
Value::Double(total) => Some(*total),
|
||||
_ => None,
|
||||
};
|
||||
total.ok_or_else(|| {
|
||||
Error::Decode(RedisError::from((
|
||||
RedisErrorKind::UnexpectedReturnType,
|
||||
"INCRBYFLOAT did not answer a number",
|
||||
)))
|
||||
})
|
||||
}
|
||||
312
litellm-rust/crates/redis-store/src/tier.rs
Normal file
312
litellm-rust/crates/redis-store/src/tier.rs
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
use std::{
|
||||
collections::{HashMap, VecDeque},
|
||||
sync::{Arc, Mutex, MutexGuard},
|
||||
time::{Duration, SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use litellm_redis_client::{Declared, Error as RedisError, RedisBatch, RedisClient};
|
||||
|
||||
use litellm_cache::{BaseCache, DeleteCache, Error, ExactCacheContext, FlushCache};
|
||||
|
||||
pub type Clock = Arc<dyn Fn() -> Duration + Send + Sync>;
|
||||
|
||||
/// How a read-through family uses its memory tier, declared once by the family.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct TierPolicy {
|
||||
/// Python's `redis_batch_cache_expiry`: a key memory misses is read from Redis at most this often.
|
||||
pub refresh_after: Duration,
|
||||
/// How long a value read from Redis stays in memory.
|
||||
pub backfill_ttl: Duration,
|
||||
pub on_breaker_open: OnBreakerOpen,
|
||||
/// Python's `default_max_redis_batch_cache_size`: how many keys the refresh throttle tracks.
|
||||
pub tracked_reads: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum OnBreakerOpen {
|
||||
/// Serve what memory holds and let the next read try Redis again.
|
||||
ServeMemory,
|
||||
}
|
||||
|
||||
/// The memory tier in front of Redis that a read-through family builds for itself, as Python's
|
||||
/// `DualCache` with a `RedisCache` does.
|
||||
pub struct MemoryTier<L1, L2> {
|
||||
memory: L1,
|
||||
remote: Option<L2>,
|
||||
policy: TierPolicy,
|
||||
reads: Mutex<RecentReads>,
|
||||
updates: Mutex<()>,
|
||||
now: Clock,
|
||||
}
|
||||
|
||||
pub fn system_clock() -> Clock {
|
||||
Arc::new(|| {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
})
|
||||
}
|
||||
|
||||
impl<V, L1> MemoryTier<L1, RedisClient>
|
||||
where
|
||||
V: Clone + Send + Sync + 'static,
|
||||
L1: BaseCache<Value = V, Context = ExactCacheContext> + DeleteCache + FlushCache,
|
||||
{
|
||||
pub fn read_through(
|
||||
memory: L1,
|
||||
remote: Option<RedisClient>,
|
||||
policy: TierPolicy,
|
||||
now: Clock,
|
||||
) -> Self {
|
||||
Self {
|
||||
memory,
|
||||
remote,
|
||||
reads: Mutex::new(RecentReads::new(policy.tracked_reads)),
|
||||
policy,
|
||||
updates: Mutex::new(()),
|
||||
now,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn memory(&self) -> &L1 {
|
||||
&self.memory
|
||||
}
|
||||
|
||||
pub fn remote(&self) -> Option<&RedisClient> {
|
||||
self.remote.as_ref()
|
||||
}
|
||||
|
||||
pub fn now(&self) -> Duration {
|
||||
(self.now)()
|
||||
}
|
||||
|
||||
pub fn set_memory(&self, key: &str, value: V, ttl: Duration) -> Result<(), Error> {
|
||||
self.memory
|
||||
.set_cache(key, value, &ExactCacheContext { ttl: Some(ttl) })
|
||||
}
|
||||
|
||||
/// Read-modify-write of one memory entry, atomic against every other update of this tier.
|
||||
pub fn update_memory(
|
||||
&self,
|
||||
key: &str,
|
||||
ttl: Duration,
|
||||
update: impl FnOnce(Option<V>) -> V,
|
||||
) -> Result<V, Error> {
|
||||
let _guard = lock(&self.updates)?;
|
||||
let value = update(self.memory.get_cache(key, &ExactCacheContext::default())?);
|
||||
self.set_memory(key, value.clone(), ttl)?;
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
/// Python's `DualCache.async_batch_get_cache`: memory first, then one round trip for the misses
|
||||
/// whose refresh is due, written back to memory. A failed Redis read serves memory, as
|
||||
/// `RedisCache` swallows it, and only a breaker refusal releases the refresh it reserved.
|
||||
pub async fn get_many<E: std::fmt::Display>(
|
||||
&self,
|
||||
keys: &[String],
|
||||
decode: impl Fn(&[u8]) -> Result<V, E>,
|
||||
) -> Result<Vec<Option<V>>, Error> {
|
||||
let declared = self.declare_many(RedisClient::batch, keys)?;
|
||||
self.resolve_many(declared, decode).await
|
||||
}
|
||||
|
||||
/// The first half of `get_many`: reserves the due refreshes and declares their GETs on the
|
||||
/// batch `batch_for` picks for this tier's client, without sending anything.
|
||||
pub fn declare_many(
|
||||
&self,
|
||||
batch_for: impl FnOnce(&RedisClient) -> RedisBatch,
|
||||
keys: &[String],
|
||||
) -> Result<DeclaredRead, Error> {
|
||||
let fetch = match &self.remote {
|
||||
None => None,
|
||||
Some(remote) => {
|
||||
let reservation = self.reserve(keys, &self.read_memory(keys)?)?;
|
||||
(!reservation.keys.is_empty()).then(|| {
|
||||
let batch = batch_for(remote);
|
||||
let keys: Vec<&str> = reservation.keys.iter().map(String::as_str).collect();
|
||||
let replies = batch.mget::<Vec<u8>>(&keys);
|
||||
Fetch {
|
||||
reservation,
|
||||
replies,
|
||||
}
|
||||
})
|
||||
}
|
||||
};
|
||||
Ok(DeclaredRead {
|
||||
keys: keys.to_vec(),
|
||||
fetch,
|
||||
})
|
||||
}
|
||||
|
||||
/// The second half of `get_many`. Memory is read again here, so a value written to memory
|
||||
/// since the declaration wins over what Redis answered.
|
||||
pub async fn resolve_many<E: std::fmt::Display>(
|
||||
&self,
|
||||
declared: DeclaredRead,
|
||||
decode: impl Fn(&[u8]) -> Result<V, E>,
|
||||
) -> Result<Vec<Option<V>>, Error> {
|
||||
let DeclaredRead { keys, fetch } = declared;
|
||||
let Some(Fetch {
|
||||
reservation,
|
||||
replies,
|
||||
}) = fetch
|
||||
else {
|
||||
return self.read_memory(&keys);
|
||||
};
|
||||
let fetched = settle(replies).await;
|
||||
let local = self.read_memory(&keys)?;
|
||||
let fetched = match fetched {
|
||||
Ok(fetched) => fetched,
|
||||
Err(RedisError::BreakerOpen) => {
|
||||
self.release(reservation)?;
|
||||
return Ok(local);
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "Redis read failed, serving the memory tier");
|
||||
return Ok(local);
|
||||
}
|
||||
};
|
||||
let decoded = match fetched
|
||||
.iter()
|
||||
.map(|raw| raw.as_deref().map(&decode).transpose())
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
{
|
||||
Ok(decoded) => decoded,
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "a Redis value did not decode, serving the memory tier");
|
||||
return Ok(local);
|
||||
}
|
||||
};
|
||||
let read = reservation
|
||||
.keys
|
||||
.into_iter()
|
||||
.zip(decoded)
|
||||
.collect::<HashMap<_, _>>();
|
||||
keys.iter()
|
||||
.zip(local)
|
||||
.map(|(key, local)| match local {
|
||||
Some(value) => Ok(Some(value)),
|
||||
None => self.backfill(key, read.get(key).cloned().flatten()),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn delete_memory(&self, key: &str) -> Result<(), Error> {
|
||||
self.memory.delete_cache(key)
|
||||
}
|
||||
|
||||
pub fn flush_memory(&self) -> Result<(), Error> {
|
||||
self.memory.flush_cache()
|
||||
}
|
||||
|
||||
fn read_memory(&self, keys: &[String]) -> Result<Vec<Option<V>>, Error> {
|
||||
let context = ExactCacheContext::default();
|
||||
keys.iter()
|
||||
.map(|key| self.memory.get_cache(key, &context))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn backfill(&self, key: &str, value: Option<V>) -> Result<Option<V>, Error> {
|
||||
if let Some(value) = &value {
|
||||
self.set_memory(key, value.clone(), self.policy.backfill_ttl)?;
|
||||
}
|
||||
Ok(value)
|
||||
}
|
||||
|
||||
fn reserve(&self, keys: &[String], local: &[Option<V>]) -> Result<Reservation, Error> {
|
||||
let now = self.now();
|
||||
let mut reads = lock(&self.reads)?;
|
||||
let due = keys
|
||||
.iter()
|
||||
.zip(local)
|
||||
.filter(|(_, value)| value.is_none())
|
||||
.filter_map(|(key, _)| {
|
||||
let previous = reads.get(key);
|
||||
let due = previous
|
||||
.is_none_or(|read| now.saturating_sub(read) >= self.policy.refresh_after);
|
||||
due.then(|| {
|
||||
reads.insert(key.clone(), now);
|
||||
(key.clone(), previous)
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let (keys, previous) = due.into_iter().unzip();
|
||||
Ok(Reservation { keys, previous })
|
||||
}
|
||||
|
||||
fn release(&self, reservation: Reservation) -> Result<(), Error> {
|
||||
let mut reads = lock(&self.reads)?;
|
||||
for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) {
|
||||
match previous {
|
||||
Some(read) => reads.insert(key, read),
|
||||
None => reads.remove(&key),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn lock<T>(mutex: &Mutex<T>) -> Result<MutexGuard<'_, T>, Error> {
|
||||
Ok(mutex.lock()?)
|
||||
}
|
||||
|
||||
async fn settle(
|
||||
replies: Declared<Vec<Option<Vec<u8>>>>,
|
||||
) -> Result<Vec<Option<Vec<u8>>>, RedisError> {
|
||||
replies.resolve().await
|
||||
}
|
||||
|
||||
/// The GETs `declare_many` put on a batch, with the refreshes they reserved.
|
||||
pub struct DeclaredRead {
|
||||
keys: Vec<String>,
|
||||
fetch: Option<Fetch>,
|
||||
}
|
||||
|
||||
struct Fetch {
|
||||
reservation: Reservation,
|
||||
replies: Declared<Vec<Option<Vec<u8>>>>,
|
||||
}
|
||||
|
||||
struct Reservation {
|
||||
keys: Vec<String>,
|
||||
previous: Vec<Option<Duration>>,
|
||||
}
|
||||
|
||||
/// Python's `LimitedSizeOrderedDict`: insertion-ordered, and any write at capacity drops the
|
||||
/// oldest entry first, even when the written key is already present.
|
||||
struct RecentReads {
|
||||
capacity: usize,
|
||||
order: VecDeque<String>,
|
||||
times: HashMap<String, Duration>,
|
||||
}
|
||||
|
||||
impl RecentReads {
|
||||
fn new(capacity: usize) -> Self {
|
||||
Self {
|
||||
capacity,
|
||||
order: VecDeque::new(),
|
||||
times: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn get(&self, key: &str) -> Option<Duration> {
|
||||
self.times.get(key).copied()
|
||||
}
|
||||
|
||||
fn insert(&mut self, key: String, time: Duration) {
|
||||
if self.times.len() >= self.capacity
|
||||
&& let Some(oldest) = self.order.pop_front()
|
||||
{
|
||||
self.times.remove(&oldest);
|
||||
}
|
||||
if self.times.insert(key.clone(), time).is_none() {
|
||||
self.order.push_back(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn remove(&mut self, key: &str) {
|
||||
if self.times.remove(key).is_some() {
|
||||
self.order.retain(|tracked| tracked != key);
|
||||
}
|
||||
}
|
||||
}
|
||||
395
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
395
litellm-rust/crates/redis-store/tests/router_cooldowns.rs
Normal file
|
|
@ -0,0 +1,395 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_client::{Operation, Ttl};
|
||||
use litellm_redis_store::{
|
||||
Number, RequestRedis,
|
||||
router_cooldowns::{POLICY, RouterCooldowns},
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{ALLOWED_FAILS, connect, cooldowns, within};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
/// 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_mains_str_dict_bytes_with_whole_second_ttl(
|
||||
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
|
||||
) {
|
||||
let deployment = RedisDeployment::start(topology).await.unwrap();
|
||||
let raw = connect(&deployment);
|
||||
let store = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
|
||||
|
||||
for case in cooldowns() {
|
||||
let ttl = Duration::from_secs_f64(case.ttl_seconds);
|
||||
store.put(case.model_id, &case.cooldown, ttl).await.unwrap();
|
||||
let key = case.key;
|
||||
|
||||
let written: Option<String> = raw.get(key).await.unwrap();
|
||||
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
|
||||
panic!("{key} has no TTL");
|
||||
};
|
||||
assert_eq!(written.as_deref(), Some(case.str_dict));
|
||||
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 = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
|
||||
|
||||
for case in cooldowns() {
|
||||
let bytes = match format {
|
||||
"json" => case.json,
|
||||
_ => case.str_dict,
|
||||
};
|
||||
raw.set(case.key, bytes, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.active(&[case.model_id.to_owned()]).await.unwrap(),
|
||||
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 = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, late);
|
||||
let models = cooldowns()
|
||||
.iter()
|
||||
.map(|case| case.model_id.to_owned())
|
||||
.collect::<Vec<_>>();
|
||||
for case in cooldowns() {
|
||||
raw.set(case.key, case.str_dict, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(
|
||||
store.minimum(&models).await.unwrap(),
|
||||
Some(Number::Float(0.5))
|
||||
);
|
||||
assert_eq!(store.active(&models).await.unwrap(), vec![]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn allowed_fails_count_across_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 = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
|
||||
let second = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
|
||||
|
||||
for case in ALLOWED_FAILS {
|
||||
let key = case.key;
|
||||
raw.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("dep-a", suffix, Duration::from_secs(120))
|
||||
.await
|
||||
.unwrap(),
|
||||
];
|
||||
let Ttl::Expires(remaining) = raw.ttl(key).await.unwrap() else {
|
||||
panic!("{key} has no TTL");
|
||||
};
|
||||
|
||||
assert_eq!(counts, [1, 2], "{key}");
|
||||
assert!(within(remaining, 30), "{remaining:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[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 = RouterCooldowns::new(Some(connect(&deployment)), POLICY);
|
||||
raw.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 = 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!(
|
||||
store.active(&[case.model_id.to_owned()]).await.unwrap(),
|
||||
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 = RouterCooldowns::with_clock(Some(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 = store
|
||||
.active(&[case.model_id.to_owned(), "dep-missing".to_owned()])
|
||||
.await
|
||||
.unwrap();
|
||||
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 = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, time.clock());
|
||||
let [case, _] = cooldowns();
|
||||
|
||||
let before = store.active(&[case.model_id.to_owned()]).await.unwrap();
|
||||
raw.set(case.key, case.str_dict, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
let throttled = store.active(&[case.model_id.to_owned()]).await.unwrap();
|
||||
time.advance(Duration::from_secs(1));
|
||||
let refreshed = store.active(&[case.model_id.to_owned()]).await.unwrap();
|
||||
|
||||
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 = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, time.clock());
|
||||
let [case, _] = cooldowns();
|
||||
deployment.stop().await.unwrap();
|
||||
store.active(&[case.model_id.to_owned()]).await.unwrap();
|
||||
time.advance(Duration::from_secs(1));
|
||||
let refused = store.active(&[case.model_id.to_owned()]).await.unwrap();
|
||||
deployment.resume().await.unwrap();
|
||||
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
|
||||
raw.set(case.key, case.str_dict, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let recovered = store.active(&[case.model_id.to_owned()]).await.unwrap();
|
||||
|
||||
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 = RouterCooldowns::with_clock(Some(client), POLICY, fixture_clock());
|
||||
let models = cooldowns().map(|case| case.model_id.to_owned());
|
||||
for case in cooldowns() {
|
||||
raw.set(case.key, case.str_dict, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
|
||||
let declared = store.declare_active(&request, &models).unwrap();
|
||||
assert!(trips.seen().is_empty());
|
||||
request.flush_pre().await;
|
||||
let active = store.resolve_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!(store.active(&models).await.unwrap(), 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 = RouterCooldowns::with_clock(Some(connect(&deployment)), POLICY, fixture_clock());
|
||||
let [case, _] = cooldowns();
|
||||
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
|
||||
let declared = store
|
||||
.declare_active(&request, &[case.model_id.to_owned()])
|
||||
.unwrap();
|
||||
request.flush_pre().await;
|
||||
store
|
||||
.put(case.model_id, &case.cooldown, Duration::from_secs(90))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
store.resolve_active(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 = RouterCooldowns::with_clock(Some(guarded.clone()), POLICY, fixture_clock());
|
||||
let [case, _] = cooldowns();
|
||||
let models = [case.model_id.to_owned()];
|
||||
deployment.stop().await.unwrap();
|
||||
let opened = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
let failed = store
|
||||
.declare_active(&opened, &["dep-other".to_owned()])
|
||||
.unwrap();
|
||||
opened.flush_pre().await;
|
||||
store.resolve_active(failed).await.unwrap();
|
||||
let refused_request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
|
||||
|
||||
let refused = store.declare_active(&refused_request, &models).unwrap();
|
||||
refused_request.flush_pre().await;
|
||||
let refused = store.resolve_active(refused).await.unwrap();
|
||||
deployment.resume().await.unwrap();
|
||||
litellm_redis_testing::eventually(async || guarded.ping().await.ok()).await;
|
||||
raw.set(case.key, case.str_dict, Duration::from_secs(60))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(refused, vec![]);
|
||||
assert_eq!(
|
||||
store.active(&models).await.unwrap(),
|
||||
vec![(case.model_id.to_owned(), case.cooldown.clone())]
|
||||
);
|
||||
}
|
||||
218
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
218
litellm-rust/crates/redis-store/tests/router_usage.rs
Normal file
|
|
@ -0,0 +1,218 @@
|
|||
mod support;
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use litellm_redis_client::{Amount, Expiry, Operation, Ttl};
|
||||
use litellm_redis_store::{
|
||||
Number,
|
||||
router_usage::{COUNTER_TTL, Counter, Kind, POLICY, RouterUsage},
|
||||
};
|
||||
use litellm_redis_testing::{RedisDeployment, Topology};
|
||||
use rstest::rstest;
|
||||
use support::{connect, usage, within};
|
||||
|
||||
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]
|
||||
#[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(connect(&deployment)), POLICY);
|
||||
let cases = usage();
|
||||
for (case, delta) in cases.iter().zip([7.75, 3.0]) {
|
||||
raw.incr_by_float(case.key, Amount::Float(delta), Expiry::Refresh(COUNTER_TTL))
|
||||
.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(connect(&deployment)), POLICY);
|
||||
let reader = RouterUsage::new(Some(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.expire(&tpm_key, Duration::from_secs(5)).await.unwrap();
|
||||
writer
|
||||
.increment_deployment("dep", model, "12-34", Number::Int(2), 1)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let written: [Option<String>; 2] = [
|
||||
raw.get(&tpm_key).await.unwrap(),
|
||||
raw.get(&rpm_key).await.unwrap(),
|
||||
];
|
||||
let Ttl::Expires(remaining) = raw.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(connect(&deployment)), POLICY);
|
||||
let [_, case] = usage();
|
||||
let key = case.key;
|
||||
raw.incr_by_float(
|
||||
key,
|
||||
Amount::Float(1.0),
|
||||
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.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_failed_pipeline_keeps_the_memory_counts_as_integers() {
|
||||
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
|
||||
let store = RouterUsage::new(Some(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(client), POLICY);
|
||||
let counters = tpm_and_rpm(Some("model"));
|
||||
let [tpm_key, rpm_key] =
|
||||
["tpm", "rpm"].map(|kind| format!("global_router:dep:model:{kind}:12-34"));
|
||||
raw.incr_by_float(&tpm_key, Amount::Float(10.0), Expiry::Refresh(COUNTER_TTL))
|
||||
.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.get(&tpm_key).await.unwrap(),
|
||||
raw.get(&rpm_key).await.unwrap(),
|
||||
];
|
||||
assert_eq!(written, [Some("15.5".to_owned()), Some("1".to_owned())]);
|
||||
let Ttl::Expires(remaining) = raw.ttl(&rpm_key).await.unwrap() else {
|
||||
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(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))]
|
||||
);
|
||||
}
|
||||
|
|
@ -10,6 +10,11 @@ use litellm_redis_client::{
|
|||
CallEvent, Cluster, Credentials, Namespace, Node, Operation, RedisClient, RedisConfig,
|
||||
RedisEvents, RedisUrl, Target, Timeouts,
|
||||
};
|
||||
use litellm_redis_store::{
|
||||
Number,
|
||||
router_cooldowns::Cooldown,
|
||||
router_usage::{Counter, Kind},
|
||||
};
|
||||
use litellm_redis_testing::RedisDeployment;
|
||||
|
||||
pub fn connect(deployment: &RedisDeployment) -> RedisClient {
|
||||
|
|
@ -81,3 +86,106 @@ 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)
|
||||
&& ttl > Duration::from_secs(whole_seconds.saturating_sub(5))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -42,6 +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.native_settings import NativeRedisSettings
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineIncrementOperation,
|
||||
RedisPipelineLpopOperation,
|
||||
|
|
@ -645,6 +646,8 @@ def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) -
|
|||
class RedisCache(BaseCache):
|
||||
# if users don't provider one, use the default litellm cache
|
||||
|
||||
native_settings: NativeRedisSettings | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host=None,
|
||||
|
|
@ -711,6 +714,11 @@ class RedisCache(BaseCache):
|
|||
super().__init__(default_ttl=int(litellm.default_redis_ttl))
|
||||
else:
|
||||
super().__init__() # defaults to 60s
|
||||
self.native_settings = NativeRedisSettings(
|
||||
redis_kwargs=MappingProxyType(dict(redis_kwargs)),
|
||||
namespace=namespace,
|
||||
default_ttl=self.default_ttl,
|
||||
)
|
||||
|
||||
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
|
||||
return call()
|
||||
|
|
|
|||
11
litellm/redis/native_settings.py
Normal file
11
litellm/redis/native_settings.py
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativeRedisSettings:
|
||||
"""What a native store takes from a `RedisCache`, captured once when the cache is built."""
|
||||
|
||||
redis_kwargs: Mapping[str, object]
|
||||
namespace: str | None
|
||||
default_ttl: int
|
||||
|
|
@ -1,29 +1,20 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._redis import _merge_redis_kwargs
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
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.redis.native_settings import NativeRedisSettings
|
||||
from litellm.rust_bridge._native import NativeRedisClient
|
||||
|
||||
_KWARGS: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def _redis_kwargs(connection: RedisCache) -> Mapping[str, object]:
|
||||
return _merge_redis_kwargs(
|
||||
_KWARGS.validate_python(
|
||||
connection.redis_kwargs # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped legacy dict
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _breaker() -> tuple[int, float, float] | None:
|
||||
if not REDIS_CIRCUIT_BREAKER_ENABLED:
|
||||
return None
|
||||
|
|
@ -34,7 +25,11 @@ def _breaker() -> tuple[int, float, float] | None:
|
|||
)
|
||||
|
||||
|
||||
def native_client(connection: RedisCache | None) -> NativeRedisClient | None:
|
||||
if connection is None:
|
||||
def native_client(settings: NativeRedisSettings | None) -> NativeRedisClient | None:
|
||||
if settings is None:
|
||||
return None
|
||||
return NativeRedisClient.connect(_redis_kwargs(connection), namespace=connection.namespace, breaker=_breaker())
|
||||
return NativeRedisClient.connect(
|
||||
_merge_redis_kwargs(_KWARGS.validate_python(settings.redis_kwargs)),
|
||||
namespace=settings.namespace,
|
||||
breaker=_breaker(),
|
||||
)
|
||||
|
|
|
|||
4
litellm/redis/stores/_keys.py
Normal file
4
litellm/redis/stores/_keys.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
from typing import Final
|
||||
|
||||
ROUTER_TPM_KEY: Final = "global_router:{id}:{model}:tpm:{current_minute}"
|
||||
ROUTER_RPM_KEY: Final = "global_router:{id}:{model}:rpm:{current_minute}"
|
||||
108
litellm/redis/stores/router_cooldowns.py
Normal file
108
litellm/redis/stores/router_cooldowns.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._internal_context import service_target
|
||||
from litellm.caching.redis_batch import NativeBatchResult, NativePending, RequestRedisBatches
|
||||
from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeRouterCooldowns
|
||||
|
||||
|
||||
class CooldownCacheValue(TypedDict):
|
||||
exception_received: ReadOnly[str]
|
||||
status_code: ReadOnly[str]
|
||||
timestamp: ReadOnly[float]
|
||||
cooldown_time: ReadOnly[float]
|
||||
|
||||
|
||||
ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
|
||||
ActiveCooldowns: TypeAlias = Sequence[tuple[str, CooldownCacheValue]]
|
||||
|
||||
|
||||
class RouterCooldowns(Protocol):
|
||||
def attach_redis(self, client: NativeRedisClient | None) -> None: ...
|
||||
def flush_local_counters(self) -> None: ...
|
||||
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None: ...
|
||||
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int: ...
|
||||
async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
|
||||
def declare_active(
|
||||
self, request: RequestRedisBatches, model_ids: Sequence[str]
|
||||
) -> NativeBatchResult[ActiveCooldowns] | None: ...
|
||||
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns: ...
|
||||
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float: ...
|
||||
def evict_local(self, model_id: str) -> None: ...
|
||||
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
|
||||
|
||||
|
||||
class CooldownsHandle(Protocol):
|
||||
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
|
||||
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
|
||||
def active(self, model_ids: Sequence[str]) -> Sequence[tuple[str, CooldownCacheValue]]: ...
|
||||
def aactive(self, model_ids: Sequence[str]) -> Awaitable[Sequence[tuple[str, CooldownCacheValue]]]: ...
|
||||
def declare_active(
|
||||
self, request: NativeRequestRedis, model_ids: Sequence[str]
|
||||
) -> NativePending[ActiveCooldowns]: ...
|
||||
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
|
||||
def evict_local(self, model_id: str) -> None: ...
|
||||
def flush_local_counters(self) -> None: ...
|
||||
def snapshot(self) -> Mapping[str, CooldownCacheValue]: ...
|
||||
|
||||
|
||||
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None], CooldownsHandle]
|
||||
|
||||
|
||||
class RouterCooldownStore:
|
||||
def __init__(self, client: NativeRedisClient | None, *, handle: HandleFactory = NativeRouterCooldowns) -> None:
|
||||
self._connected = client is not None
|
||||
self._build: Final = handle
|
||||
self._handle = handle(client)
|
||||
|
||||
def attach_redis(self, client: NativeRedisClient | None) -> None:
|
||||
if client is None or self._connected:
|
||||
return
|
||||
self._connected = True
|
||||
self._handle = self._build(client)
|
||||
|
||||
def flush_local_counters(self) -> None:
|
||||
self._handle.flush_local_counters()
|
||||
|
||||
def put(self, model_id: str, value: CooldownCacheValue, ttl: float) -> None:
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
self._handle.put(model_id, value, ttl)
|
||||
|
||||
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl: float) -> int:
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
return self._handle.increment_allowed_fails(model_id, suffix, ttl)
|
||||
|
||||
async def active(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
pending: Final = self._handle.aactive(model_ids)
|
||||
return tuple(await pending)
|
||||
|
||||
def declare_active(
|
||||
self, request: RequestRedisBatches, model_ids: Sequence[str]
|
||||
) -> NativeBatchResult[ActiveCooldowns] | None:
|
||||
"""``active`` declared on the request's native pre-call batch; None when there is no Redis to read."""
|
||||
if not self._connected:
|
||||
return None
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
pending: Final = self._handle.declare_active(request.native(), model_ids)
|
||||
return request.native_result(pending)
|
||||
|
||||
def active_sync(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> ActiveCooldowns:
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
return tuple(self._handle.active(model_ids))
|
||||
|
||||
def minimum(self, model_ids: Sequence[str], default: float, parent_otel_span: Span | None) -> float:
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
return self._handle.minimum(model_ids) or default
|
||||
|
||||
def evict_local(self, model_id: str) -> None:
|
||||
self._handle.evict_local(model_id)
|
||||
|
||||
def snapshot(self) -> Mapping[str, CooldownCacheValue]:
|
||||
return self._handle.snapshot()
|
||||
106
litellm/redis/stores/router_usage.py
Normal file
106
litellm/redis/stores/router_usage.py
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm.rust_bridge._native import NativeRedisClient, NativeRouterUsage
|
||||
|
||||
ROUTER_USAGE_TARGET: Final = "router_usage"
|
||||
CounterKind: TypeAlias = Literal["tpm", "rpm"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageCounter:
|
||||
deployment_id: str
|
||||
model: str | None
|
||||
minute: str
|
||||
kind: CounterKind
|
||||
|
||||
|
||||
class RouterUsage(Protocol):
|
||||
def attach_redis(self, client: NativeRedisClient | None) -> None: ...
|
||||
def flush_local_counters(self) -> None: ...
|
||||
async def async_increment(
|
||||
self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None
|
||||
) -> float | None: ...
|
||||
async def increment_deployment(
|
||||
self,
|
||||
deployment_id: str,
|
||||
model: str | None,
|
||||
minute: str,
|
||||
tokens: float,
|
||||
requests: int,
|
||||
parent_otel_span: Span | None,
|
||||
) -> None: ...
|
||||
async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None: ...
|
||||
def snapshot(self) -> Mapping[str, float]: ...
|
||||
|
||||
|
||||
NativeCounter: TypeAlias = tuple[str, str | None, str, CounterKind]
|
||||
|
||||
|
||||
class UsageHandle(Protocol):
|
||||
def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ...
|
||||
def increment_deployment(
|
||||
self,
|
||||
deployment_id: str,
|
||||
model: str | None,
|
||||
minute: str,
|
||||
tokens: float,
|
||||
requests: int,
|
||||
) -> Awaitable[None]: ...
|
||||
def get_many(self, counters: Sequence[NativeCounter]) -> Awaitable[Sequence[float | None]]: ...
|
||||
def flush_local_counters(self) -> None: ...
|
||||
def snapshot(self) -> Mapping[str, float]: ...
|
||||
|
||||
|
||||
HandleFactory: TypeAlias = Callable[[NativeRedisClient | None], UsageHandle]
|
||||
|
||||
|
||||
def _native(counter: UsageCounter) -> NativeCounter:
|
||||
return (counter.deployment_id, counter.model, counter.minute, counter.kind)
|
||||
|
||||
|
||||
class RouterUsageStore:
|
||||
def __init__(self, client: NativeRedisClient | None, *, handle: HandleFactory = NativeRouterUsage) -> None:
|
||||
self._connected = client is not None
|
||||
self._build: Final = handle
|
||||
self._handle = handle(client)
|
||||
|
||||
def attach_redis(self, client: NativeRedisClient | None) -> None:
|
||||
if client is None or self._connected:
|
||||
return
|
||||
self._connected = True
|
||||
self._handle = self._build(client)
|
||||
|
||||
def flush_local_counters(self) -> None:
|
||||
self._handle.flush_local_counters()
|
||||
|
||||
@with_service_target(ROUTER_USAGE_TARGET)
|
||||
async def async_increment(
|
||||
self, counter: UsageCounter, value: float, ttl: int, parent_otel_span: Span | None = None
|
||||
) -> float | None:
|
||||
return await self._handle.increment(_native(counter), value, ttl)
|
||||
|
||||
@with_service_target(ROUTER_USAGE_TARGET)
|
||||
async def increment_deployment(
|
||||
self,
|
||||
deployment_id: str,
|
||||
model: str | None,
|
||||
minute: str,
|
||||
tokens: float,
|
||||
requests: int,
|
||||
parent_otel_span: Span | None,
|
||||
) -> None:
|
||||
await self._handle.increment_deployment(deployment_id, model, minute, tokens, requests)
|
||||
|
||||
async def async_batch_get(self, counters: Sequence[UsageCounter]) -> Sequence[object] | None:
|
||||
return await self._handle.get_many(tuple(map(_native, counters)))
|
||||
|
||||
def snapshot(self) -> Mapping[str, float]:
|
||||
return self._handle.snapshot()
|
||||
|
|
@ -242,6 +242,9 @@ from litellm.llms.openai_like.model_info import (
|
|||
get_openai_compatible_model_info,
|
||||
)
|
||||
from litellm.main import Chat as LitellmChat
|
||||
from litellm.redis.stores._connection import native_client
|
||||
from litellm.redis.stores.router_cooldowns import RouterCooldownStore
|
||||
from litellm.redis.stores.router_usage import RouterUsage, RouterUsageStore, UsageCounter
|
||||
from litellm.responses.bridge_check import keep_registered_mode
|
||||
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults
|
||||
from litellm.router_api import (
|
||||
|
|
@ -1768,6 +1771,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self._core_cache = DualCache(
|
||||
redis_cache=redis_cache, in_memory_cache=InMemoryCache()
|
||||
) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc.
|
||||
redis_client: Final = native_client(None if redis_cache is None else redis_cache.native_settings)
|
||||
self._core_usage_store: Final[RouterUsage] = RouterUsageStore(redis_client)
|
||||
self._claude_code_session_router_cache: DualCache = DualCache(
|
||||
redis_cache=redis_cache,
|
||||
in_memory_cache=InMemoryCache(),
|
||||
|
|
@ -1873,7 +1878,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self.explicit_allowed_fails = allowed_fails
|
||||
self._core_cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS
|
||||
self._core_cooldown_cache = CooldownCache(
|
||||
cache=self._core_cache, default_cooldown_time=self._core_cooldown_time
|
||||
store=RouterCooldownStore(redis_client), default_cooldown_time=self._core_cooldown_time
|
||||
)
|
||||
self._core_disable_cooldowns = disable_cooldowns
|
||||
self.fallback_on_unsupported_params = fallback_on_unsupported_params
|
||||
|
|
@ -2229,6 +2234,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
and caching to just work.
|
||||
"""
|
||||
self._core_cache.attach_redis_cache(cache)
|
||||
redis_client: Final = native_client(cache.native_settings)
|
||||
self._core_usage_store.attach_redis(redis_client)
|
||||
self._core_cooldown_cache.cooldown_store.attach_redis(redis_client)
|
||||
self._claude_code_session_router_cache.attach_redis_cache(cache)
|
||||
self._core_prompt_caching_pin_cache.attach_redis_cache(cache)
|
||||
|
||||
|
|
@ -10173,35 +10181,17 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
total_tokens: float,
|
||||
rpm_increment: int,
|
||||
parent_otel_span: Span | None,
|
||||
) -> str | None:
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
|
||||
) -> UsageCounter | None:
|
||||
if total_tokens <= 0 and rpm_increment <= 0:
|
||||
return None
|
||||
if not self._cached_deployment_tracks_usage(deployment_id, model_group):
|
||||
return None
|
||||
|
||||
current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock
|
||||
tpm_key: Final = RouterCacheEnum.TPM.value.format(
|
||||
id=deployment_id, current_minute=current_minute, model=deployment_name
|
||||
await self._core_usage_store.increment_deployment(
|
||||
deployment_id, deployment_name, current_minute, total_tokens, rpm_increment, parent_otel_span
|
||||
)
|
||||
rpm_key: Final = RouterCacheEnum.RPM.value.format(
|
||||
id=deployment_id, current_minute=current_minute, model=deployment_name
|
||||
)
|
||||
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [
|
||||
RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value)
|
||||
for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment))
|
||||
]
|
||||
post_increment_values: Final = await self._core_cache.async_increment_cache_pipeline(
|
||||
increment_list=pipeline_operations,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
if post_increment_values is not None and self._core_cache.redis_cache is not None:
|
||||
for operation, value in zip(pipeline_operations, post_increment_values):
|
||||
await self._core_cache.async_set_cache(
|
||||
operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value
|
||||
)
|
||||
return tpm_key
|
||||
return UsageCounter(deployment_id, deployment_name, current_minute, "tpm")
|
||||
|
||||
def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool:
|
||||
deployment_info: Final = self._core_get_deployment(model_id=deployment_id)
|
||||
|
|
@ -10431,10 +10421,10 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock
|
||||
|
||||
## RPM
|
||||
rpm_key: Final = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
|
||||
counter: Final = UsageCounter(id, deployment_name, current_minute, "rpm")
|
||||
with service_target(ROUTER_USAGE_TARGET):
|
||||
await self._core_cache.async_increment_cache(
|
||||
key=rpm_key,
|
||||
await self._core_usage_store.async_increment(
|
||||
counter=counter,
|
||||
value=1,
|
||||
parent_otel_span=parent_otel_span,
|
||||
ttl=RoutingArgs.ttl.value,
|
||||
|
|
@ -13612,8 +13602,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
"""
|
||||
dt: Final = get_utc_datetime()
|
||||
current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock
|
||||
tpm_keys: Final[list[str]] = []
|
||||
rpm_keys: Final[list[str]] = []
|
||||
tpm_keys: Final[list[UsageCounter]] = []
|
||||
rpm_keys: Final[list[UsageCounter]] = []
|
||||
|
||||
model_list: Final = self._core_get_model_list(model_name=model_group)
|
||||
if model_list is None: # no matching deployments
|
||||
|
|
@ -13626,28 +13616,16 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written.
|
||||
if id is None or litellm_model is None:
|
||||
continue
|
||||
tpm_keys.append(
|
||||
RouterCacheEnum.TPM.value.format(
|
||||
id=id,
|
||||
model=litellm_model,
|
||||
current_minute=current_minute,
|
||||
)
|
||||
)
|
||||
rpm_keys.append(
|
||||
RouterCacheEnum.RPM.value.format(
|
||||
id=id,
|
||||
model=litellm_model,
|
||||
current_minute=current_minute,
|
||||
)
|
||||
)
|
||||
tpm_keys.append(UsageCounter(id, litellm_model, current_minute, "tpm"))
|
||||
rpm_keys.append(UsageCounter(id, litellm_model, current_minute, "rpm"))
|
||||
combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys
|
||||
|
||||
combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys)
|
||||
combined_tpm_rpm_values: Final = await self._core_usage_store.async_batch_get(combined_tpm_rpm_keys)
|
||||
if combined_tpm_rpm_values is None:
|
||||
return None, None
|
||||
|
||||
tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[: len(tpm_keys)]
|
||||
rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[len(tpm_keys) :]
|
||||
tpm_usage_list: Final = combined_tpm_rpm_values[: len(tpm_keys)]
|
||||
rpm_usage_list: Final = combined_tpm_rpm_values[len(tpm_keys) :]
|
||||
|
||||
## TPM
|
||||
tpm_usage: int | None = None
|
||||
|
|
@ -17856,6 +17834,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
def flush_cache(self):
|
||||
self._host.set_global_cache(None)
|
||||
self._core_usage_store.flush_local_counters()
|
||||
self._core_cooldown_cache.cooldown_store.flush_local_counters()
|
||||
self._core_cache.flush_cache()
|
||||
session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache
|
||||
if session_in_memory_cache is not None:
|
||||
|
|
|
|||
|
|
@ -1,237 +1,54 @@
|
|||
"""
|
||||
Wrapper around router cache. Meant to handle model cooldown logic
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm._internal_context import service_target
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span as _Span
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
Span = _Span
|
||||
else:
|
||||
Span = Any
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.redis.stores.router_cooldowns import (
|
||||
ROUTER_COOLDOWNS_TARGET,
|
||||
CooldownCacheValue,
|
||||
RouterCooldowns,
|
||||
)
|
||||
|
||||
|
||||
class CooldownCacheValue(TypedDict):
|
||||
exception_received: str
|
||||
status_code: str
|
||||
timestamp: float
|
||||
cooldown_time: float
|
||||
|
||||
|
||||
# Cap on the corrected in-memory TTL set in `_corrected_active_cooldown`: re-checks the
|
||||
# real remaining cooldown against Redis at least this often, so an entry that later gets
|
||||
# deleted or extended in Redis before its original deadline is still noticed promptly.
|
||||
_MAX_CORRECTED_IN_MEMORY_TTL_SECONDS: Final = 60.0
|
||||
ROUTER_COOLDOWNS_TARGET: Final = "router_cooldowns"
|
||||
__all__: Final = ("ROUTER_COOLDOWNS_TARGET", "CooldownCache", "CooldownCacheValue")
|
||||
|
||||
|
||||
class CooldownCache:
|
||||
def __init__(
|
||||
self,
|
||||
cache: DualCache,
|
||||
default_cooldown_time: float,
|
||||
redis_read_interval_seconds: float = DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS,
|
||||
):
|
||||
self.cache = cache
|
||||
self.default_cooldown_time = default_cooldown_time
|
||||
self.in_memory_cache = InMemoryCache()
|
||||
self._cooldown_store = DualCache(
|
||||
in_memory_cache=self.in_memory_cache,
|
||||
default_redis_batch_cache_expiry=redis_read_interval_seconds,
|
||||
self, store: RouterCooldowns, default_cooldown_time: float, *, clock: Callable[[], float] | None = None
|
||||
) -> None:
|
||||
self.cooldown_store: Final = store
|
||||
self._clock: Final = clock if clock is not None else lambda: time.time()
|
||||
self.default_cooldown_time: Final = default_cooldown_time
|
||||
self.exception_masker: Final = SensitiveDataMasker(
|
||||
visible_prefix=50, visible_suffix=0, mask_char="*", mask_short_values=False
|
||||
)
|
||||
# Initialize the masker with custom settings for exception strings
|
||||
self.exception_masker = SensitiveDataMasker(
|
||||
visible_prefix=50, # Show first 50 characters
|
||||
visible_suffix=0, # Show last 0 characters
|
||||
mask_char="*", # Use * for masking
|
||||
mask_short_values=False, # Truncate long messages only; keep short ones readable
|
||||
)
|
||||
|
||||
@property
|
||||
def cooldown_store(self) -> DualCache:
|
||||
"""
|
||||
The cache cooldown entries live in, with the router's Redis attached on first use.
|
||||
|
||||
It is kept separate from the router-wide cache so that a key missing from memory is
|
||||
re-read from Redis every `redis_read_interval_seconds` rather than on the router
|
||||
cache's much longer batch interval, which is what lets a sibling replica see a
|
||||
cooldown another replica wrote, and so that unrelated router keys cannot evict a
|
||||
cooldown from the in-memory tier before it expires. Redis is attached lazily because
|
||||
the router builds its cooldown cache before it wires up the shared Redis client.
|
||||
"""
|
||||
self._cooldown_store.attach_redis_cache(self.cache.redis_cache)
|
||||
return self._cooldown_store
|
||||
|
||||
def _common_add_cooldown_logic(
|
||||
self, model_id: str, original_exception, exception_status, cooldown_time: float
|
||||
) -> tuple[str, CooldownCacheValue]:
|
||||
try:
|
||||
current_time: Final = time.time()
|
||||
cooldown_key: Final = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
# Store the cooldown information for the deployment separately
|
||||
cooldown_data: Final = CooldownCacheValue(
|
||||
exception_received=self.exception_masker.mask_value(str(original_exception)),
|
||||
status_code=str(exception_status),
|
||||
timestamp=current_time,
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
|
||||
return cooldown_key, cooldown_data
|
||||
except Exception as e:
|
||||
verbose_logger.error("CooldownCache::_common_add_cooldown_logic - Exception occurred - %s", e)
|
||||
raise e
|
||||
|
||||
def add_deployment_to_cooldown(
|
||||
self,
|
||||
model_id: str,
|
||||
original_exception: Exception,
|
||||
exception_status: int,
|
||||
cooldown_time: float | None,
|
||||
):
|
||||
try:
|
||||
#########################################################
|
||||
# get cooldown time
|
||||
# 1. If dynamic cooldown time is set for the model/deployment, use that
|
||||
# 2. If no dynamic cooldown time is set, use the default cooldown time set on CooldownCache
|
||||
_cooldown_time = cooldown_time
|
||||
if _cooldown_time is None:
|
||||
_cooldown_time = self.default_cooldown_time
|
||||
#########################################################
|
||||
|
||||
cooldown_key, cooldown_data = self._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=original_exception,
|
||||
exception_status=exception_status,
|
||||
cooldown_time=_cooldown_time,
|
||||
)
|
||||
|
||||
# Set the cache with a TTL equal to the cooldown time
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
self.cooldown_store.set_cache(
|
||||
value=cooldown_data,
|
||||
key=cooldown_key,
|
||||
ttl=_cooldown_time,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.error("CooldownCache::add_deployment_to_cooldown - Exception occurred - %s", e)
|
||||
raise e
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(maxsize=1024)
|
||||
def get_cooldown_cache_key(model_id: str) -> str:
|
||||
return "deployment:" + model_id + ":cooldown"
|
||||
|
||||
def _corrected_active_cooldown(
|
||||
self,
|
||||
key: str,
|
||||
result: Mapping[str, Any],
|
||||
current_time: float,
|
||||
) -> CooldownCacheValue | None:
|
||||
"""
|
||||
Return a CooldownCacheValue if the cooldown is still active, or None if it has expired.
|
||||
|
||||
Also corrects the in-memory TTL when DualCache promotes a Redis entry using the
|
||||
default 600s TTL instead of the true remaining cooldown time.
|
||||
"""
|
||||
cooldown_cache_value: Final = CooldownCacheValue(**result) # pyright: ignore[reportUnknownArgumentType] - result comes from an untyped cache read, not from our own code
|
||||
remaining: Final = (cooldown_cache_value["timestamp"] + cooldown_cache_value["cooldown_time"]) - current_time
|
||||
if remaining <= 0:
|
||||
self.in_memory_cache.delete_cache(key)
|
||||
return None
|
||||
current_expiry: Final = self.in_memory_cache.ttl_dict.get(key)
|
||||
if current_expiry is not None and current_expiry > current_time + remaining + 5:
|
||||
corrected_ttl: Final = min(remaining, _MAX_CORRECTED_IN_MEMORY_TTL_SECONDS)
|
||||
self.in_memory_cache.delete_cache(key)
|
||||
self.in_memory_cache.set_cache(key, result, ttl=corrected_ttl)
|
||||
return cooldown_cache_value
|
||||
self, model_id: str, original_exception: Exception, exception_status: int, cooldown_time: float | None
|
||||
) -> None:
|
||||
ttl: Final = self.default_cooldown_time if cooldown_time is None else cooldown_time
|
||||
value: Final = CooldownCacheValue(
|
||||
exception_received=self.exception_masker.mask_value(str(original_exception)),
|
||||
status_code=str(exception_status),
|
||||
timestamp=self._clock(),
|
||||
cooldown_time=ttl,
|
||||
)
|
||||
self.cooldown_store.put(model_id, value, ttl)
|
||||
|
||||
async def async_get_active_cooldowns(
|
||||
self, model_ids: list[str], parent_otel_span: Span | None
|
||||
self, model_ids: Sequence[str], parent_otel_span: Span | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]:
|
||||
# Generate the keys for the deployments
|
||||
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
|
||||
return list(await self.cooldown_store.active(model_ids, parent_otel_span))
|
||||
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
results: Final = await self.cooldown_store.async_batch_get_cache(
|
||||
keys=keys, parent_otel_span=parent_otel_span
|
||||
)
|
||||
return self.active_cooldowns_from_results(model_ids, results)
|
||||
|
||||
def active_cooldowns_from_results(
|
||||
self, model_ids: list[str], results: Sequence[object] | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]:
|
||||
"""The cooldowns still active in a `cooldown_store` batch read of `get_cooldown_cache_key(model_id)` per id."""
|
||||
active_cooldowns: Final[list[tuple[str, CooldownCacheValue]]] = []
|
||||
|
||||
if results is None or all(v is None for v in results):
|
||||
return active_cooldowns
|
||||
|
||||
current_time: Final = time.time()
|
||||
for model_id, result in zip(model_ids, results):
|
||||
if result and isinstance(result, dict):
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
cooldown_cache_value = self._corrected_active_cooldown(key, result, current_time)
|
||||
if cooldown_cache_value is not None:
|
||||
active_cooldowns.append((model_id, cooldown_cache_value))
|
||||
|
||||
return active_cooldowns
|
||||
def get_min_cooldown(self, model_ids: Sequence[str], parent_otel_span: Span | None) -> float:
|
||||
return self.cooldown_store.minimum(model_ids, self.default_cooldown_time, parent_otel_span)
|
||||
|
||||
def get_active_cooldowns(
|
||||
self, model_ids: list[str], parent_otel_span: Span | None
|
||||
self, model_ids: Sequence[str], parent_otel_span: Span | None
|
||||
) -> list[tuple[str, CooldownCacheValue]]:
|
||||
# Generate the keys for the deployments
|
||||
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
|
||||
# Retrieve the values for the keys using mget
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
|
||||
|
||||
active_cooldowns: Final = []
|
||||
current_time: Final = time.time()
|
||||
for model_id, result in zip(model_ids, results):
|
||||
if result and isinstance(result, dict):
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
cooldown_cache_value = self._corrected_active_cooldown(key, result, current_time)
|
||||
if cooldown_cache_value is not None:
|
||||
active_cooldowns.append((model_id, cooldown_cache_value))
|
||||
|
||||
return active_cooldowns
|
||||
|
||||
def get_min_cooldown(self, model_ids: list[str], parent_otel_span: Span | None) -> float:
|
||||
"""Return min cooldown time required for a group of model id's."""
|
||||
|
||||
# Generate the keys for the deployments
|
||||
keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
|
||||
|
||||
# Retrieve the values for the keys using mget
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
results: Final = self.cooldown_store.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
|
||||
|
||||
min_cooldown_time: float | None = None
|
||||
# Process the results
|
||||
for model_id, result in zip(model_ids, results):
|
||||
if result and isinstance(result, dict):
|
||||
cooldown_cache_value = CooldownCacheValue(**result)
|
||||
if min_cooldown_time is None or cooldown_cache_value["cooldown_time"] < min_cooldown_time:
|
||||
min_cooldown_time = cooldown_cache_value["cooldown_time"]
|
||||
|
||||
return min_cooldown_time or self.default_cooldown_time
|
||||
|
||||
|
||||
# Usage example:
|
||||
# cooldown_cache = CooldownCache(cache=your_cache_instance, cooldown_time=your_cooldown_time)
|
||||
# cooldown_cache.add_deployment_to_cooldown(deployment, original_exception, exception_status)
|
||||
# active_cooldowns = cooldown_cache.get_active_cooldowns()
|
||||
return list(self.cooldown_store.active_sync(model_ids, parent_otel_span))
|
||||
|
|
|
|||
|
|
@ -16,9 +16,7 @@ from typing import TYPE_CHECKING, Any, Final
|
|||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import (
|
||||
DEFAULT_COOLDOWN_TIME_SECONDS,
|
||||
DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS,
|
||||
|
|
@ -30,7 +28,7 @@ from litellm.exceptions import MidStreamFallbackError
|
|||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
|
||||
from litellm.router_deprecated import renamed_router_param
|
||||
from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCacheValue
|
||||
from litellm.router_utils.cooldown_cache import CooldownCacheValue
|
||||
from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
|
||||
from litellm.types.router import Deployment
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -748,7 +746,7 @@ def should_cooldown_based_on_allowed_fails_policy(
|
|||
When *allowed_fails_override* / *cooldown_time_override* are supplied they
|
||||
take precedence over the router-level values (used by deployment-level overrides).
|
||||
|
||||
The counter lives in the router's shared ``DualCache`` (Redis when configured), so
|
||||
The counter lives in the cooldown store (Redis when configured), so
|
||||
every worker process increments the same key and the threshold applies fleet-wide.
|
||||
|
||||
When *cache_key_suffix* is supplied the fail counter is keyed as
|
||||
|
|
@ -771,26 +769,12 @@ def should_cooldown_based_on_allowed_fails_policy(
|
|||
else (ctx.settings().cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS)
|
||||
)
|
||||
|
||||
base_key: Final = f"deployment:{deployment}:allowed_fails"
|
||||
cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key
|
||||
updated_fails: Final = _increment_allowed_fails(cache=ctx.health().store, cache_key=cache_key, ttl=cooldown_time)
|
||||
updated_fails: Final = ctx.health().cooldowns.cooldown_store.increment_allowed_fails(
|
||||
deployment, cache_key_suffix, cooldown_time
|
||||
)
|
||||
return updated_fails > allowed_fails
|
||||
|
||||
|
||||
def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> int:
|
||||
"""
|
||||
Return the fleet-wide fail count. ``DualCache.increment_cache`` bumps the in-memory tier
|
||||
before Redis and re-raises a Redis error, so a Redis outage degrades to this worker's own count.
|
||||
"""
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
try:
|
||||
return cache.increment_cache(key=cache_key, value=1, ttl=ttl)
|
||||
except Exception as e: # noqa: BLE001 # a Redis outage must not stop failing deployments from cooling down
|
||||
verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e)
|
||||
local_fails: Final = cache.get_cache(key=cache_key, local_only=True)
|
||||
return local_fails if isinstance(local_fails, int) else 0
|
||||
|
||||
|
||||
def _is_allowed_fails_set_on_router(
|
||||
ctx: "RoutingContext",
|
||||
) -> bool:
|
||||
|
|
|
|||
165
litellm/router_utils/routing_prefetch.py
Normal file
165
litellm/router_utils/routing_prefetch.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm._internal_context import service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
|
||||
from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, RouterCooldowns
|
||||
|
||||
_PREFETCH_SLOT: Final = "routing_read"
|
||||
_COOLDOWNS_SLOT: Final = "routing_cooldowns"
|
||||
ROUTER_USAGE_TARGET: Final = "router_usage"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TierRead:
|
||||
cache: DualCache
|
||||
keys: tuple[str, ...]
|
||||
|
||||
|
||||
def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead:
|
||||
return TierRead(cache, tuple(keys))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ArmedRead:
|
||||
read: TierRead
|
||||
fetched: frozenset[str]
|
||||
result: BatchResult[Mapping[str, object]]
|
||||
previous: Mapping[str, float | None]
|
||||
|
||||
def release(self) -> None:
|
||||
self.read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
|
||||
dict(self.previous)
|
||||
)
|
||||
|
||||
async def settle(self, future: asyncio.Future[Mapping[str, object]]) -> None:
|
||||
if future.cancelled() or future.exception() is not None:
|
||||
self.release()
|
||||
return
|
||||
try:
|
||||
pending: Final = await self.read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
|
||||
list(self.read.keys), local_only=True
|
||||
)
|
||||
values: Final = future.result()
|
||||
missed: Final = {
|
||||
key: values[key]
|
||||
for key, local in zip(self.read.keys, pending.result)
|
||||
if local is None and values.get(key) is not None
|
||||
}
|
||||
await self.read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # the selector owns this tier
|
||||
except Exception:
|
||||
self.release()
|
||||
raise
|
||||
|
||||
|
||||
def arm(read: TierRead) -> None:
|
||||
request: Final = active_request_redis_batches()
|
||||
redis: Final = read.cache.redis_cache
|
||||
if request is None or redis is None or _PREFETCH_SLOT in request.prefetched:
|
||||
return
|
||||
due, previous = read.cache.reserve_redis_batch_reads(read.keys)
|
||||
if not due:
|
||||
return
|
||||
with service_target(ROUTER_USAGE_TARGET):
|
||||
result: Final = request.batch(redis).mget(tuple(due))
|
||||
armed_read: Final = ArmedRead(read=read, fetched=frozenset(due), result=result, previous=previous)
|
||||
result.on_settled(armed_read.settle)
|
||||
request.prefetched[_PREFETCH_SLOT] = armed_read
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ArmedCooldowns:
|
||||
model_ids: frozenset[str]
|
||||
result: Awaitable[ActiveCooldowns]
|
||||
|
||||
|
||||
def arm_cooldowns(store: RouterCooldowns, model_ids: Sequence[str]) -> None:
|
||||
request: Final = active_request_redis_batches()
|
||||
if request is None or _COOLDOWNS_SLOT in request.prefetched:
|
||||
return
|
||||
result: Final = store.declare_active(request, model_ids)
|
||||
if result is not None:
|
||||
request.prefetched[_COOLDOWNS_SLOT] = ArmedCooldowns(frozenset(model_ids), result)
|
||||
|
||||
|
||||
def armed() -> bool:
|
||||
request: Final = active_request_redis_batches()
|
||||
return request is not None and (_PREFETCH_SLOT in request.prefetched or _COOLDOWNS_SLOT in request.prefetched)
|
||||
|
||||
|
||||
async def _armed_cooldowns(model_ids: Sequence[str]) -> ActiveCooldowns | None:
|
||||
request: Final = active_request_redis_batches()
|
||||
taken: Final = None if request is None else request.prefetched.pop(_COOLDOWNS_SLOT, None)
|
||||
if not isinstance(taken, ArmedCooldowns):
|
||||
return None
|
||||
try:
|
||||
active: Final = await taken.result
|
||||
except Exception as e: # noqa: BLE001 # the read falls back to the store's own Redis read
|
||||
verbose_router_logger.debug("cooldown prefetch failed, reading again: %s", e)
|
||||
return None
|
||||
if not taken.model_ids.issuperset(model_ids):
|
||||
return None
|
||||
wanted: Final = frozenset(model_ids)
|
||||
return tuple((model_id, cooldown) for model_id, cooldown in active if model_id in wanted)
|
||||
|
||||
|
||||
async def read_cooldowns(
|
||||
store: RouterCooldowns, model_ids: Sequence[str], parent_otel_span: Span | None
|
||||
) -> ActiveCooldowns:
|
||||
"""The armed cooldown read when it covers ``model_ids``, taken once; otherwise the store reads on its own."""
|
||||
with service_target(ROUTER_COOLDOWNS_TARGET):
|
||||
prefetched: Final = await _armed_cooldowns(model_ids)
|
||||
if prefetched is not None:
|
||||
return prefetched
|
||||
return await store.active(model_ids, parent_otel_span)
|
||||
|
||||
|
||||
def _take(read: TierRead) -> ArmedRead | None:
|
||||
request: Final = active_request_redis_batches()
|
||||
if request is None:
|
||||
return None
|
||||
taken: Final = request.prefetched.pop(_PREFETCH_SLOT, None)
|
||||
if isinstance(taken, ArmedRead) and set(taken.read.keys).issuperset(read.keys):
|
||||
return taken
|
||||
if isinstance(taken, ArmedRead):
|
||||
taken.release()
|
||||
return None
|
||||
|
||||
|
||||
async def _read_prefetched(read: TierRead) -> Sequence[object | None] | None:
|
||||
prefetch: Final = _take(read)
|
||||
if prefetch is None:
|
||||
return None
|
||||
try:
|
||||
values: Final = await prefetch.result
|
||||
except Exception as e: # noqa: BLE001 # the read falls back to the tier's own Redis read
|
||||
verbose_router_logger.debug("routing prefetch failed, reading again: %s", e)
|
||||
return None
|
||||
pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
|
||||
list(read.keys), local_only=True
|
||||
)
|
||||
if any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None):
|
||||
return None
|
||||
return await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier
|
||||
pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None}
|
||||
)
|
||||
|
||||
|
||||
async def read(read: TierRead, parent_otel_span: Span | None) -> Sequence[object | None] | None:
|
||||
with service_target(ROUTER_USAGE_TARGET):
|
||||
prefetched: Final = await _read_prefetched(read)
|
||||
if prefetched is not None:
|
||||
return prefetched
|
||||
results: Final = await DualCache.async_batch_get_cache_shared(
|
||||
((read.cache, list(read.keys)),), parent_otel_span=parent_otel_span
|
||||
)
|
||||
return results[0]
|
||||
|
|
@ -1,157 +1,48 @@
|
|||
"""
|
||||
One Redis round trip for the reads a request needs before a deployment can be picked.
|
||||
|
||||
The cooldown filter (`CooldownCache`, its own `DualCache`) and usage-based selection
|
||||
(`LowestTPMLoggingHandler_v2`, the router cache) each issue their own MGET because they live in
|
||||
different objects. `RoutingReadBatch` fetches both key sets in one
|
||||
`DualCache.async_batch_get_cache_shared` while the healthy deployments are being resolved and hands
|
||||
the usage slice to the strategy, so selection does not read again.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from collections.abc import Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._internal_context import service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.redis_batch import BatchResult, active_request_redis_batches
|
||||
from litellm.router_deprecated import renamed_router_param
|
||||
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage
|
||||
from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCache
|
||||
from litellm.router_utils import routing_prefetch
|
||||
from litellm.router_utils.routing_prefetch import ROUTER_USAGE_TARGET, TierRead, selector_read
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
||||
from litellm.router_api import RoutingContext
|
||||
|
||||
|
||||
ROUTER_COOLDOWNS_USAGE_TARGET: Final = "router_cooldowns_usage"
|
||||
ROUTER_USAGE_TARGET: Final = "router_usage"
|
||||
_PREFETCH_SLOT: Final = "routing_read"
|
||||
__all__ = ("ROUTER_USAGE_TARGET", "RoutingPrefetch", "RoutingReadBatch")
|
||||
|
||||
|
||||
def _routing_read_target(cooldown_keys: Sequence[str], usage_keys: Sequence[str]) -> str:
|
||||
if not usage_keys:
|
||||
return ROUTER_COOLDOWNS_TARGET
|
||||
if not cooldown_keys:
|
||||
return ROUTER_USAGE_TARGET
|
||||
return ROUTER_COOLDOWNS_USAGE_TARGET
|
||||
def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: list) -> list[str]:
|
||||
tpm_keys, rpm_keys = selector.usage_counter_keys(deployments)
|
||||
return tpm_keys + rpm_keys
|
||||
|
||||
|
||||
async def _backfill_prefetched_cache(
|
||||
cache: DualCache,
|
||||
due_keys: tuple[str, ...],
|
||||
values: Mapping[str, object],
|
||||
) -> None:
|
||||
cache_keys: Final = list(due_keys)
|
||||
prepare_batch_get: Final = cache._prepare_batch_get # pyright: ignore[reportPrivateUsage] # memory backfill
|
||||
pending: Final = await prepare_batch_get(cache_keys, local_only=True)
|
||||
redis_values: Final = {
|
||||
key: values[key]
|
||||
for key, local in zip(due_keys, pending.result)
|
||||
if local is None and values.get(key) is not None
|
||||
}
|
||||
apply_batch_get: Final = cache._apply_batch_get # pyright: ignore[reportPrivateUsage] # cache backfill
|
||||
await apply_batch_get(pending, redis_values)
|
||||
def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> TierRead | None:
|
||||
return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments))
|
||||
|
||||
|
||||
async def _nothing() -> Sequence[object | None] | None:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RoutingPrefetch:
|
||||
"""The cooldown and usage keys of a model group, declared on the request's Redis batch before admission
|
||||
flushes it, so the routing read rides the same round trip as the rate limiter's Lua calls."""
|
||||
|
||||
keys: frozenset[str]
|
||||
fetched: frozenset[str]
|
||||
result: BatchResult[Mapping[str, object]]
|
||||
reservations: tuple[tuple[DualCache, tuple[str, ...], dict[str, float | None]], ...]
|
||||
|
||||
def release(self) -> None:
|
||||
for cache, _, previous_access_times in self.reservations:
|
||||
cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # rollback
|
||||
previous_access_times
|
||||
)
|
||||
|
||||
async def _settle(self, future: asyncio.Future[Mapping[str, object]]) -> None:
|
||||
if future.cancelled():
|
||||
self.release()
|
||||
return
|
||||
if future.exception() is not None:
|
||||
self.release()
|
||||
return
|
||||
|
||||
values: Final = future.result()
|
||||
try:
|
||||
for cache, due_keys, _ in self.reservations:
|
||||
await _backfill_prefetched_cache(cache, due_keys, values)
|
||||
except Exception:
|
||||
self.release()
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
@renamed_router_param("litellm_router_instance")
|
||||
def arm(
|
||||
ctx: "RoutingContext",
|
||||
usage_selector: LowestTPMLoggingHandler_v2 | None,
|
||||
deployments: list,
|
||||
) -> None:
|
||||
request: Final = active_request_redis_batches()
|
||||
redis_cache: Final = ctx.health().store.redis_cache
|
||||
if request is None or redis_cache is None or _PREFETCH_SLOT in request.prefetched:
|
||||
return
|
||||
cooldown_keys: Final = tuple(
|
||||
CooldownCache.get_cooldown_cache_key(model_id) for model_id in ctx.registry().names(kind="ids")
|
||||
)
|
||||
usage_keys: Final = (
|
||||
() if usage_selector is None else tuple(itertools.chain(*usage_selector.usage_counter_keys(deployments)))
|
||||
)
|
||||
keys: Final = (*cooldown_keys, *usage_keys)
|
||||
cooldown_store: Final = ctx.health().cooldowns.cooldown_store
|
||||
cooldown_due, cooldown_previous = cooldown_store.reserve_redis_batch_reads(cooldown_keys)
|
||||
usage_cache: Final = None if usage_selector is None else usage_selector.router_cache
|
||||
usage_reservation: Final = None if usage_cache is None else usage_cache.reserve_redis_batch_reads(usage_keys)
|
||||
usage_due: Final = () if usage_reservation is None else tuple(usage_reservation[0])
|
||||
due: Final = (*cooldown_due, *usage_due)
|
||||
reservations: Final = (
|
||||
(cooldown_store, tuple(cooldown_due), cooldown_previous),
|
||||
*(
|
||||
()
|
||||
if usage_cache is None or usage_reservation is None
|
||||
else ((usage_cache, usage_due, usage_reservation[1]),)
|
||||
),
|
||||
)
|
||||
if not due:
|
||||
return
|
||||
with service_target(_routing_read_target(cooldown_due, usage_due)):
|
||||
result: Final = request.batch(redis_cache).mget(due)
|
||||
prefetch: Final = RoutingPrefetch(
|
||||
keys=frozenset(keys), fetched=frozenset(due), result=result, reservations=reservations
|
||||
)
|
||||
result.on_settled(prefetch._settle)
|
||||
request.prefetched[_PREFETCH_SLOT] = prefetch
|
||||
def arm(ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: list) -> None:
|
||||
routing_prefetch.arm_cooldowns(ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids"))
|
||||
read: Final = _selector_read(usage_selector, deployments)
|
||||
if read is not None:
|
||||
routing_prefetch.arm(read)
|
||||
|
||||
@staticmethod
|
||||
def armed() -> bool:
|
||||
request: Final = active_request_redis_batches()
|
||||
return request is not None and _PREFETCH_SLOT in request.prefetched
|
||||
|
||||
@staticmethod
|
||||
def take(needed: Sequence[str]) -> "RoutingPrefetch | None":
|
||||
"""The armed prefetch when it covers every key this read needs; taken once, so a retry reads fresh."""
|
||||
request: Final = active_request_redis_batches()
|
||||
if request is None:
|
||||
return None
|
||||
armed: Final = request.prefetched.pop(_PREFETCH_SLOT, None)
|
||||
if isinstance(armed, RoutingPrefetch) and armed.keys.issuperset(needed):
|
||||
return armed
|
||||
if isinstance(armed, RoutingPrefetch):
|
||||
armed.release()
|
||||
return None
|
||||
return routing_prefetch.armed()
|
||||
|
||||
|
||||
_active_routing_read_batch: Final[ContextVar["RoutingReadBatch | None"]] = ContextVar(
|
||||
|
|
@ -179,71 +70,25 @@ class RoutingReadBatch:
|
|||
|
||||
@staticmethod
|
||||
def for_strategy(strategy: str | None, selector: object) -> "RoutingReadBatch | None":
|
||||
"""Usage-based routing reads its counters with the cooldown state; every other strategy reads only the
|
||||
cooldown state, and only through this batch when the request armed a prefetch for it. Otherwise the
|
||||
router's plain cooldown read stays in charge."""
|
||||
if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2):
|
||||
return RoutingReadBatch(usage_selector=selector)
|
||||
return RoutingReadBatch(usage_selector=None) if RoutingPrefetch.armed() else None
|
||||
|
||||
@renamed_router_param("litellm_router_instance")
|
||||
async def async_get_cooldown_deployments(
|
||||
self,
|
||||
ctx: "RoutingContext",
|
||||
healthy_deployments: list,
|
||||
parent_otel_span: "Span | None",
|
||||
self, ctx: "RoutingContext", healthy_deployments: list, parent_otel_span: "Span | None"
|
||||
) -> list[str]:
|
||||
"""
|
||||
`_async_get_cooldown_deployments`, with the strategy's tpm/rpm counters for
|
||||
`healthy_deployments` fetched in the same MGET and kept as `prefetched_usage`.
|
||||
"""
|
||||
model_ids: Final = ctx.registry().names(kind="ids")
|
||||
cooldown_keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
|
||||
selector: Final = self.usage_selector
|
||||
usage_keys: Final = (
|
||||
() if selector is None else tuple(itertools.chain(*selector.usage_counter_keys(healthy_deployments)))
|
||||
usage_read: Final = _selector_read(self.usage_selector, healthy_deployments)
|
||||
cooldowns, selector_values = await asyncio.gather(
|
||||
routing_prefetch.read_cooldowns(
|
||||
ctx.health().cooldowns.cooldown_store, ctx.registry().names(kind="ids"), parent_otel_span
|
||||
),
|
||||
_nothing() if usage_read is None else routing_prefetch.read(usage_read, parent_otel_span),
|
||||
)
|
||||
reads: Final = (
|
||||
(ctx.health().cooldowns.cooldown_store, cooldown_keys),
|
||||
*(() if selector is None else ((selector.router_cache, list(usage_keys)),)),
|
||||
)
|
||||
with service_target(_routing_read_target(cooldown_keys, usage_keys)):
|
||||
results: Final = await self._read_prefetched(reads) or await DualCache.async_batch_get_cache_shared(
|
||||
reads, parent_otel_span=parent_otel_span
|
||||
)
|
||||
cooldown_results: Final = results[0]
|
||||
if selector is not None:
|
||||
usage_values: Final = results[1]
|
||||
if usage_read is not None:
|
||||
self.prefetched_usage = PrefetchedUsage(
|
||||
keys=frozenset(usage_keys),
|
||||
values=None if usage_values is None else MappingProxyType(dict(zip(usage_keys, usage_values))),
|
||||
keys=frozenset(usage_read.keys),
|
||||
values=None if selector_values is None else dict(zip(usage_read.keys, selector_values)),
|
||||
)
|
||||
|
||||
cooldown_models: Final = ctx.health().cooldowns.active_cooldowns_from_results(model_ids, cooldown_results)
|
||||
verbose_router_logger.debug("retrieve cooldown models: %s", cooldown_models)
|
||||
return [model_id for model_id, _ in cooldown_models]
|
||||
|
||||
@staticmethod
|
||||
async def _read_prefetched(
|
||||
reads: Sequence[tuple[DualCache, list[str]]],
|
||||
) -> list[list[object | None] | None] | None:
|
||||
"""Serve the reads from the request's armed `RoutingPrefetch`, backfilling each cache's memory tier as
|
||||
its own batch read would. None when nothing usable was armed or the prefetch failed."""
|
||||
prefetch: Final = RoutingPrefetch.take(tuple(itertools.chain.from_iterable(keys for _, keys in reads)))
|
||||
if prefetch is None:
|
||||
return None
|
||||
try:
|
||||
values: Final = await prefetch.result
|
||||
except Exception as e: # noqa: BLE001 # the shared read below applies the caches' own Redis fallback
|
||||
verbose_router_logger.debug("routing prefetch failed, reading again: %s", e)
|
||||
return None
|
||||
results: Final[list[list[object | None] | None]] = [] # mutable-ok: filled per read below
|
||||
for cache, keys in reads:
|
||||
pending = await cache._prepare_batch_get(keys, local_only=True) # pyright: ignore[reportPrivateUsage] # same two-step read as async_batch_get_cache_shared
|
||||
if any(
|
||||
key not in prefetch.fetched for key, local_value in zip(keys, pending.result) if local_value is None
|
||||
):
|
||||
return None
|
||||
missed = {key: values.get(key) for key, local in zip(keys, pending.result) if local is None}
|
||||
results.append(await cache._apply_batch_get(pending, missed)) # pyright: ignore[reportPrivateUsage] # same two-step read as async_batch_get_cache_shared
|
||||
return results
|
||||
verbose_router_logger.debug("retrieve cooldown models: %s", cooldowns)
|
||||
return [model_id for model_id, _ in cooldowns]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from asyncio import Future
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Never, final
|
||||
from typing import Literal, Never, TypedDict, final
|
||||
|
||||
import httpx
|
||||
from pydantic import JsonValue
|
||||
|
|
@ -27,6 +27,12 @@ class NativeClickHouseSpendStorage:
|
|||
def ensure_schema(self) -> Future[None]: ...
|
||||
def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ...
|
||||
|
||||
class _NativeCooldown(TypedDict):
|
||||
exception_received: str
|
||||
status_code: str
|
||||
timestamp: float
|
||||
cooldown_time: float
|
||||
|
||||
@final
|
||||
class NativeRedisClient:
|
||||
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
|
||||
|
|
@ -44,6 +50,43 @@ class NativeRequestRedis:
|
|||
def flush_pre(self) -> Future[None]: ...
|
||||
def flush_post(self) -> Future[None]: ...
|
||||
|
||||
@final
|
||||
class NativeDeclaredCooldowns:
|
||||
def resolve(self) -> Future[list[tuple[str, _NativeCooldown]]]: ...
|
||||
|
||||
@final
|
||||
class NativeRouterCooldowns:
|
||||
def __new__(cls, client: NativeRedisClient | None, /) -> NativeRouterCooldowns: ...
|
||||
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None: ...
|
||||
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int: ...
|
||||
def active(self, model_ids: Sequence[str]) -> list[tuple[str, _NativeCooldown]]: ...
|
||||
def aactive(self, model_ids: Sequence[str]) -> Future[list[tuple[str, _NativeCooldown]]]: ...
|
||||
def declare_active(self, request: NativeRequestRedis, model_ids: Sequence[str]) -> NativeDeclaredCooldowns: ...
|
||||
def minimum(self, model_ids: Sequence[str]) -> float | None: ...
|
||||
def evict_local(self, model_id: str) -> None: ...
|
||||
def flush_local_counters(self) -> None: ...
|
||||
def snapshot(self) -> dict[str, _NativeCooldown]: ...
|
||||
|
||||
@final
|
||||
class NativeRouterUsage:
|
||||
def __new__(cls, client: NativeRedisClient | None, /) -> NativeRouterUsage: ...
|
||||
def increment(
|
||||
self, counter: tuple[str, str | None, str, Literal["tpm", "rpm"]], value: float, ttl_seconds: int
|
||||
) -> Future[float]: ...
|
||||
def increment_deployment(
|
||||
self,
|
||||
deployment_id: str,
|
||||
model: str | None,
|
||||
minute: str,
|
||||
tokens: float,
|
||||
requests: int,
|
||||
) -> Future[None]: ...
|
||||
def get_many(
|
||||
self, counters: Sequence[tuple[str, str | None, str, Literal["tpm", "rpm"]]]
|
||||
) -> Future[list[float | None]]: ...
|
||||
def flush_local_counters(self) -> None: ...
|
||||
def snapshot(self) -> dict[str, float]: ...
|
||||
|
||||
@final
|
||||
class NativeDatabase:
|
||||
@staticmethod
|
||||
|
|
@ -233,9 +276,12 @@ __all__ = [
|
|||
"NativeClickHouseSpendConfig",
|
||||
"NativeClickHouseSpendStorage",
|
||||
"NativeDatabase",
|
||||
"NativeDeclaredCooldowns",
|
||||
"NativeDiagnosticProcessor",
|
||||
"NativeRedisClient",
|
||||
"NativeRequestRedis",
|
||||
"NativeRouterCooldowns",
|
||||
"NativeRouterUsage",
|
||||
"ProcessReservedForForking",
|
||||
"ResponsesWebSocketConnection",
|
||||
"RustBridgeDeclined",
|
||||
|
|
|
|||
|
|
@ -31,6 +31,15 @@ class ProviderDefaults:
|
|||
enable_azure_ad_token_refresh: object
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RouterTierSettings:
|
||||
"""The memory-over-Redis tier values `main`'s router reads from the environment and `litellm` globals."""
|
||||
|
||||
cooldown_redis_read_interval: float
|
||||
usage_redis_read_interval: float
|
||||
max_redis_batch_cache_size: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SecretManager:
|
||||
readable: bool
|
||||
|
|
@ -108,6 +117,17 @@ def secret_manager_binding() -> SecretManagerBinding:
|
|||
)
|
||||
|
||||
|
||||
def router_tiers() -> RouterTierSettings:
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS, DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
|
||||
|
||||
return RouterTierSettings(
|
||||
cooldown_redis_read_interval=DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS,
|
||||
usage_redis_read_interval=litellm.default_redis_batch_cache_expiry or 10,
|
||||
max_redis_batch_cache_size=DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE,
|
||||
)
|
||||
|
||||
|
||||
def provider_defaults() -> ProviderDefaults:
|
||||
import litellm
|
||||
|
||||
|
|
|
|||
|
|
@ -46,6 +46,8 @@ from litellm._uuid import uuid
|
|||
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
|
||||
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
|
||||
from litellm.litellm_core_utils.provider_affinity import validate_provider_affinity_header_name
|
||||
from litellm.redis.stores._keys import ROUTER_RPM_KEY as _ROUTER_RPM_KEY
|
||||
from litellm.redis.stores._keys import ROUTER_TPM_KEY as _ROUTER_TPM_KEY
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.router_weights import RouterWeights
|
||||
|
||||
|
|
@ -1291,8 +1293,8 @@ class RoutingStrategy(enum.Enum):
|
|||
|
||||
|
||||
class RouterCacheEnum(enum.Enum):
|
||||
TPM = "global_router:{id}:{model}:tpm:{current_minute}"
|
||||
RPM = "global_router:{id}:{model}:rpm:{current_minute}"
|
||||
TPM = _ROUTER_TPM_KEY
|
||||
RPM = _ROUTER_RPM_KEY
|
||||
ITPM = "global_router:{id}:{model}:itpm:{current_minute}"
|
||||
OTPM = "global_router:{id}:{model}:otpm:{current_minute}"
|
||||
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from litellm.caching.redis_cache import (
|
|||
_get_call_stack_info, # pyright: ignore[reportPrivateUsage] # the chain the service hook reports
|
||||
)
|
||||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.redis.native_settings import NativeRedisSettings
|
||||
|
||||
SCRIPT = "return redis.call('GET', KEYS[1])"
|
||||
_Answer = TypeVar("_Answer")
|
||||
|
|
@ -96,6 +97,7 @@ class FakeRedisCache(RedisCache):
|
|||
self.client = client
|
||||
self.namespace = namespace
|
||||
self.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2, "db": id(self) % 1_000_000}
|
||||
self.native_settings = NativeRedisSettings(redis_kwargs=self.redis_kwargs, namespace=namespace, default_ttl=60)
|
||||
self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30)
|
||||
self.service_logger_obj = ServiceLogging()
|
||||
self.default_ttl = None
|
||||
|
|
|
|||
|
|
@ -38,7 +38,7 @@ from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
|
|||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from .test_redis_batch import FakeClient, FakeRedisCache
|
||||
from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache
|
||||
|
||||
|
||||
async def _script_outside_the_pipeline(keys: Sequence[str], args: Sequence[object]) -> object:
|
||||
|
|
@ -514,3 +514,31 @@ async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_co
|
|||
assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it"
|
||||
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend"
|
||||
assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_post_call_ops_go_out_with_the_post_call_flush_not_the_request_boundary():
|
||||
natives = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
request.native_post_call()
|
||||
await request.flush_all()
|
||||
at_the_boundary = list(natives.built[0].events)
|
||||
await flush_post_call_redis_batches()
|
||||
await drain_post_call_redis_batches()
|
||||
|
||||
assert at_the_boundary == ["flush_pre"]
|
||||
assert natives.built[0].events == ["flush_pre", "flush_post"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_native_post_call_ops_still_waiting_on_their_callbacks_are_drained_at_shutdown():
|
||||
natives = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
request.native_post_call()
|
||||
await request.flush_all()
|
||||
await drain_post_call_redis_batches()
|
||||
await drain_post_call_redis_batches()
|
||||
|
||||
assert natives.built[0].events == ["flush_pre", "flush_post"]
|
||||
|
|
|
|||
|
|
@ -1,13 +1,13 @@
|
|||
"""One Redis pipeline per backend for the pre-call reads a request makes: rate limiter Lua groups, the
|
||||
router's cooldown and usage read, auth identity and spend counters all join the request batch."""
|
||||
router's usage counter read, auth identity and spend counters all join the request batch."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain
|
||||
from typing import Any, Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -24,21 +24,24 @@ from litellm.proxy.auth.auth_object_prefetch import _CacheEntry, _write_back, pr
|
|||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
|
||||
CHECK_AND_INCREMENT_BY_N_SCRIPT,
|
||||
PROXY_MaxParallelRequestsHandler_v3,
|
||||
RateLimitDescriptor,
|
||||
RateLimitUnverifiableError,
|
||||
PROXY_MaxParallelRequestsHandler_v3,
|
||||
)
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.router_utils.cooldown_cache import ROUTER_COOLDOWNS_TARGET, CooldownCache
|
||||
from litellm.router_utils.routing_prefetch import ArmedRead, TierRead
|
||||
from litellm.router_utils.routing_read_batch import (
|
||||
ROUTER_COOLDOWNS_USAGE_TARGET,
|
||||
ROUTER_USAGE_TARGET,
|
||||
RoutingPrefetch,
|
||||
_routing_read_target, # pyright: ignore[reportPrivateUsage] # the family rule under test
|
||||
)
|
||||
|
||||
from .test_redis_batch import FakeClient, FakeRedisCache, replies
|
||||
from .test_redis_batch import FakeClient, FakeNativePending, FakeNatives, FakeRedisCache, replies
|
||||
|
||||
|
||||
def _cooldown_key(model_id: str) -> str:
|
||||
return f"deployment:{model_id}:cooldown"
|
||||
|
||||
|
||||
_MODEL_GROUP = "claude"
|
||||
_FAR_FUTURE = 4_102_444_800.0 # 2100-01-01, a cooldown stamped then is still active
|
||||
|
|
@ -289,12 +292,76 @@ def _router(redis_cache: FakeRedisCache, routing_strategy: str = "usage-based-ro
|
|||
return router
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime",
|
||||
lambda: datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc),
|
||||
_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
|
||||
_START: Final = 1_000_000.0
|
||||
_PING: Final = [{"role": "user", "content": "ping"}]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pinned_clock(monkeypatch: pytest.MonkeyPatch) -> list[float]:
|
||||
"""Pins the usage counters' minute and the router cache's read throttle clock; tests move `now[0]`."""
|
||||
now: Final = [_START] # mutable-ok: the test advances the pinned clock
|
||||
monkeypatch.setattr("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", lambda: _PINNED_MINUTE)
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: now[0])
|
||||
return now
|
||||
|
||||
|
||||
def _counter_keys(router: Router, deployment_ids: tuple[str, ...] = ("dep-a", "dep-b")) -> frozenset[str]:
|
||||
deployments: Final = [d for d in router.model_list if d["model_info"]["id"] in deployment_ids]
|
||||
tpm, rpm = router.lowesttpm_logger_v2.usage_counter_keys(deployments)
|
||||
return frozenset(tpm + rpm)
|
||||
|
||||
|
||||
def _tpm_key(router: Router, deployment_id: str) -> str:
|
||||
deployment: Final = next(d for d in router.model_list if d["model_info"]["id"] == deployment_id)
|
||||
return router.lowesttpm_logger_v2.usage_counter_keys([deployment])[0][0]
|
||||
|
||||
|
||||
def _counter_mgets(router: Router, client: FakeClient, redis_cache: FakeRedisCache) -> tuple[frozenset[str], ...]:
|
||||
counters: Final = _counter_keys(router)
|
||||
pipelined: Final = tuple(
|
||||
frozenset(command[1:])
|
||||
for pipeline in client.pipelines
|
||||
for command in pipeline.commands
|
||||
if command[0] == "MGET" and counters.intersection(command[1:])
|
||||
)
|
||||
alone: Final = tuple(
|
||||
frozenset(keys) for command, keys in redis_cache.alone if command == "MGET" and counters.intersection(keys)
|
||||
)
|
||||
return pipelined + alone
|
||||
|
||||
|
||||
def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[str, ...]:
|
||||
pipelined: Final = tuple(
|
||||
key
|
||||
for pipeline in client.pipelines
|
||||
for command in pipeline.commands
|
||||
if command[0] == "MGET"
|
||||
for key in command[1:]
|
||||
)
|
||||
alone: Final = tuple(key for command, keys in redis_cache.alone if command == "MGET" for key in keys)
|
||||
return tuple(key for key in pipelined + alone if key.endswith(":cooldown"))
|
||||
|
||||
|
||||
def _replies_with(values: dict[str, object]) -> Callable[[tuple[Any, ...]], Any]:
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET":
|
||||
return [None if values.get(key) is None else json.dumps(values[key]) for key in command[1:]]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
return replies
|
||||
|
||||
|
||||
async def _route(router: Router) -> str:
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=_PING, request_kwargs={}
|
||||
)
|
||||
return deployment["model_info"]["id"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issues_no_read_of_its_own():
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache)
|
||||
|
|
@ -306,58 +373,40 @@ async def test_armed_routing_read_rides_the_admission_pipeline_and_routing_issue
|
|||
descriptors=[_descriptor("api_key", "k1", 10), _descriptor("team", "t1", 20)],
|
||||
increments=[{"requests": 1}, {"requests": 1}],
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
deployment = await _route(router)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
assert deployment in {"dep-a", "dep-b"}
|
||||
assert len(client.pipelines) == 1
|
||||
commands = client.pipelines[0].commands
|
||||
assert [c[0] for c in commands] == ["MGET", "EVALSHA", "EVALSHA"]
|
||||
mget_keys = set(commands[0][1:])
|
||||
assert {CooldownCache.get_cooldown_cache_key("dep-a"), CooldownCache.get_cooldown_cache_key("dep-b")} <= mget_keys
|
||||
assert any(":tpm:" in key for key in mget_keys) and any(":rpm:" in key for key in mget_keys)
|
||||
assert frozenset(commands[0][1:]) == _counter_keys(router)
|
||||
assert redis_cache.alone == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cooldown_recorded_locally_after_the_prefetch_left_still_excludes_its_deployment():
|
||||
expired = {"exception_received": "429", "status_code": "429", "timestamp": 0.0, "cooldown_time": 60}
|
||||
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET": # Redis holds a stale cooldown for dep-b and nothing for dep-a
|
||||
return [
|
||||
json.dumps(expired) if key == CooldownCache.get_cooldown_cache_key("dep-b") else None
|
||||
for key in command[1:]
|
||||
]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
client = FakeClient(replies)
|
||||
@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"])
|
||||
async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_without_a_python_cooldown_read(
|
||||
routing_strategy: str,
|
||||
):
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache)
|
||||
cooldown_store = router.cooldown_cache.cooldown_store
|
||||
assert cooldown_store.in_memory_cache is not None
|
||||
router = _router(redis_cache, routing_strategy=routing_strategy)
|
||||
router.cooldown_cache.cooldown_store.put(
|
||||
"dep-a",
|
||||
{"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60},
|
||||
60,
|
||||
)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
cooldown_store.in_memory_cache.set_cache(
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
{"exception_received": "429", "status_code": "429", "timestamp": _FAR_FUTURE, "cooldown_time": 60},
|
||||
)
|
||||
picks = {
|
||||
(
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
)["model_info"]["id"]
|
||||
for _ in range(5)
|
||||
}
|
||||
picks = {await _route(router) for _ in range(5)}
|
||||
|
||||
assert picks == {"dep-b"}
|
||||
assert len(client.pipelines) == 1
|
||||
assert _python_cooldown_reads(client, redis_cache) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_routing_reads_itself():
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
|
|
@ -366,81 +415,58 @@ async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_ro
|
|||
with request_redis_batch_scope() as request:
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
armed = request.prefetched["routing_read"]
|
||||
assert isinstance(armed, RoutingPrefetch)
|
||||
request.prefetched["routing_read"] = RoutingPrefetch(
|
||||
keys=frozenset({"other"}),
|
||||
assert isinstance(armed, ArmedRead)
|
||||
request.prefetched["routing_read"] = ArmedRead(
|
||||
read=TierRead(armed.read.cache, ("other",)),
|
||||
fetched=armed.fetched,
|
||||
result=armed.result,
|
||||
reservations=armed.reservations,
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
previous=armed.previous,
|
||||
)
|
||||
deployment = await _route(router)
|
||||
assert request.prefetched == {}
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
assert len(redis_cache.alone) == 1 # the shared cooldown+usage read, one round trip as in P1
|
||||
assert deployment in {"dep-a", "dep-b"}
|
||||
assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_prefetch_with_incomplete_usage_keys_releases_cooldown_reservations():
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_a_prefetch_with_incomplete_usage_keys_releases_its_reservations():
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
router: Final = _router(redis_cache)
|
||||
cooldown_keys: Final = frozenset(
|
||||
{
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
RoutingPrefetch.arm(RoutingContext(router), router.lowesttpm_logger_v2, router.model_list[:1])
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
deployment: Final = await _route(router)
|
||||
|
||||
fallback_cooldown_mgets: Final = tuple(
|
||||
keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys)
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
assert len(fallback_cooldown_mgets) == 1
|
||||
assert deployment in {"dep-a", "dep-b"}
|
||||
assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_prefetch_falls_back_to_the_shared_read():
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_a_failed_prefetch_falls_back_to_the_selectors_own_read():
|
||||
client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down"))
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
deployment = await _route(router)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
assert len(redis_cache.alone) == 1
|
||||
cooldown_keys: Final = frozenset(
|
||||
{
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
)
|
||||
fallback_cooldown_mgets: Final = tuple(
|
||||
keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_keys.issubset(keys)
|
||||
)
|
||||
assert len(fallback_cooldown_mgets) == 1
|
||||
assert deployment in {"dep-a", "dep-b"}
|
||||
assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_family():
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_the_armed_routing_read_is_declared_under_the_router_usage_family():
|
||||
"""The prefetch is declared before routing runs under a target of its own, so the pipeline that
|
||||
carries it renders ``redis.pipeline router_cooldowns`` instead of a bare ``redis.pipeline``."""
|
||||
carries it renders ``redis.pipeline router_usage`` instead of a bare ``redis.pipeline``."""
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
router = _router(redis_cache)
|
||||
pipeline_targets: list[str | None] = [] # mutable-ok: filled by the recording hook
|
||||
|
||||
async def record(**kwargs: object) -> None:
|
||||
|
|
@ -450,35 +476,21 @@ async def test_the_armed_routing_read_is_declared_under_the_router_cooldowns_fam
|
|||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
|
||||
await _route(router)
|
||||
async with asyncio.timeout(5):
|
||||
while not pipeline_targets:
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert pipeline_targets == [ROUTER_COOLDOWNS_TARGET]
|
||||
assert pipeline_targets == [ROUTER_USAGE_TARGET]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkeypatch):
|
||||
clock: Final = 1_000_000.0
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
|
||||
cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
|
||||
active_cooldown: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET":
|
||||
return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
client: Final = FakeClient(replies)
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_an_abandoned_prefetch_still_backfills_the_counters_it_read():
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
redis_cache.store[cooldown_key] = active_cooldown
|
||||
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
router: Final = _router(redis_cache)
|
||||
client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100})
|
||||
limiter: Final = _limiter(redis_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
|
|
@ -487,176 +499,66 @@ async def test_an_abandoned_prefetch_still_backfills_the_cooldown_it_read(monkey
|
|||
descriptors=[_descriptor("api_key", "k1", 10)],
|
||||
increments=[{"requests": 1}],
|
||||
)
|
||||
after_abandoned: Final = len(_counter_mgets(router, client, redis_cache))
|
||||
|
||||
pipeline_count: Final = len(client.pipelines)
|
||||
first_cooldown_mgets: Final = tuple(
|
||||
command
|
||||
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count])
|
||||
if command[0] == "MGET" and cooldown_key in command[1:]
|
||||
)
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None)
|
||||
deployment: Final = await _route(router)
|
||||
|
||||
second_cooldown_mgets: Final = tuple(
|
||||
command
|
||||
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:])
|
||||
if command[0] == "MGET" and cooldown_key in command[1:]
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] == "dep-b"
|
||||
assert [model_id for model_id, _ in cooldowns] == ["dep-a"]
|
||||
assert len(first_cooldown_mgets) == 1
|
||||
assert second_cooldown_mgets == ()
|
||||
assert deployment == "dep-b"
|
||||
assert after_abandoned == 1
|
||||
assert len(_counter_mgets(router, client, redis_cache)) == 1
|
||||
assert redis_cache.alone == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(monkeypatch):
|
||||
clock: Final = 1_000_000.0
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
|
||||
dep_a_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
|
||||
dep_b_key: Final = CooldownCache.get_cooldown_cache_key("dep-b")
|
||||
old_cooldown: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
newer_memory_cooldown: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE + 1,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
redis_only_cooldown: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE + 2,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET":
|
||||
return [json.dumps(redis_cache.store[key]) if key in redis_cache.store else None for key in command[1:]]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
client: Final = FakeClient(replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
redis_cache.store[dep_a_key] = old_cooldown
|
||||
redis_cache.store[dep_b_key] = redis_only_cooldown
|
||||
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
memory_cache: Final = router.cooldown_cache.cooldown_store.in_memory_cache
|
||||
assert memory_cache is not None
|
||||
|
||||
with request_redis_batch_scope() as request:
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
memory_cache.set_cache(dep_a_key, newer_memory_cooldown)
|
||||
await request.flush_all()
|
||||
|
||||
prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET")
|
||||
|
||||
assert len(prefetched_mgets) == 1
|
||||
assert frozenset(prefetched_mgets[0][1:]) == frozenset({dep_a_key, dep_b_key})
|
||||
assert memory_cache.get_cache(dep_a_key) == newer_memory_cooldown
|
||||
assert memory_cache.get_cache(dep_b_key) == redis_only_cooldown
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(monkeypatch):
|
||||
clock: Final = 1_000_000.0
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
|
||||
cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
|
||||
active_cooldown: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
mget_replies: Final = iter((ConnectionError("redis down"), None))
|
||||
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET":
|
||||
response: Final = next(mget_replies)
|
||||
if isinstance(response, Exception):
|
||||
return response
|
||||
return [json.dumps(active_cooldown) if key == cooldown_key else None for key in command[1:]]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
client: Final = FakeClient(replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
redis_cache.store[cooldown_key] = active_cooldown
|
||||
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
|
||||
with request_redis_batch_scope() as request:
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await request.flush_all()
|
||||
|
||||
pipeline_count: Final = len(client.pipelines)
|
||||
first_cooldown_mgets: Final = tuple(
|
||||
command
|
||||
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[:pipeline_count])
|
||||
if command[0] == "MGET" and cooldown_key in command[1:]
|
||||
)
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
cooldowns: Final = await router.cooldown_cache.async_get_active_cooldowns(["dep-a"], parent_otel_span=None)
|
||||
|
||||
second_cooldown_mgets: Final = tuple(
|
||||
command
|
||||
for command in chain.from_iterable(pipeline.commands for pipeline in client.pipelines[pipeline_count:])
|
||||
if command[0] == "MGET" and cooldown_key in command[1:]
|
||||
)
|
||||
|
||||
assert deployment["model_info"]["id"] == "dep-b"
|
||||
assert [model_id for model_id, _ in cooldowns] == ["dep-a"]
|
||||
assert len(first_cooldown_mgets) == 1
|
||||
assert len(second_cooldown_mgets) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cooldown_that_leaves_memory_before_routing_is_read_again(monkeypatch):
|
||||
clock: Final = 1_000_000.0
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: clock)
|
||||
cooldown_key: Final = CooldownCache.get_cooldown_cache_key("dep-a")
|
||||
active_cooldown: Final = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses():
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
redis_cache.store[cooldown_key] = active_cooldown
|
||||
router: Final = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
cooldown_store: Final = router.cooldown_cache.cooldown_store
|
||||
memory_cache: Final = cooldown_store.in_memory_cache
|
||||
assert memory_cache is not None
|
||||
memory_cache.set_cache(cooldown_key, active_cooldown)
|
||||
router: Final = _router(redis_cache)
|
||||
dep_a: Final = _tpm_key(router, "dep-a")
|
||||
dep_b: Final = _tpm_key(router, "dep-b")
|
||||
client.reply_for = _replies_with({dep_a: 5, dep_b: 7})
|
||||
usage_cache: Final = router.lowesttpm_logger_v2.router_cache
|
||||
|
||||
with request_redis_batch_scope() as request:
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
usage_cache.in_memory_cache.set_cache(dep_a, 50)
|
||||
await request.flush_all()
|
||||
|
||||
assert await usage_cache.async_batch_get_cache([dep_a, dep_b], local_only=True) == [50, 7]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(pinned_clock: list[float]):
|
||||
mget_replies: Final = iter((ConnectionError("redis down"), None))
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
router: Final = _router(redis_cache)
|
||||
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET":
|
||||
failure: Final = next(mget_replies)
|
||||
if isinstance(failure, Exception):
|
||||
return failure
|
||||
return ["100" if key == _tpm_key(router, "dep-a") else None for key in command[1:]]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
client.reply_for = replies
|
||||
|
||||
with request_redis_batch_scope() as request:
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await request.flush_all()
|
||||
memory_cache.delete_cache(cooldown_key)
|
||||
deployment: Final = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
after_failure: Final = len(_counter_mgets(router, client, redis_cache))
|
||||
|
||||
prefetched_mgets: Final = tuple(command for command in client.pipelines[0].commands if command[0] == "MGET")
|
||||
fallback_cooldown_mgets: Final = tuple(
|
||||
keys for command, keys in redis_cache.alone if command == "MGET" and cooldown_key in keys
|
||||
)
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment: Final = await _route(router)
|
||||
|
||||
assert len(prefetched_mgets) == 1
|
||||
assert prefetched_mgets[0][1:] == (CooldownCache.get_cooldown_cache_key("dep-b"),)
|
||||
assert deployment["model_info"]["id"] == "dep-b"
|
||||
assert fallback_cooldown_mgets == ((cooldown_key,),)
|
||||
assert after_failure == 1
|
||||
assert len(_counter_mgets(router, client, redis_cache)) == 2
|
||||
assert deployment == "dep-b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -668,52 +570,35 @@ async def test_arming_outside_a_request_scope_is_a_no_op():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_shuffle_prefetches_only_its_cooldown_read_into_the_admission_pipeline():
|
||||
async def test_simple_shuffle_arms_no_python_prefetch_and_reads_no_cooldown_key():
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
limiter = _limiter(redis_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
with request_redis_batch_scope() as request:
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
assert set(request.prefetched) == {"routing_cooldowns"}
|
||||
await limiter.atomic_check_and_increment_by_n(
|
||||
descriptors=[_descriptor("api_key", "k1", 10)],
|
||||
increments=[{"requests": 1}],
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
deployment = await _route(router)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
assert len(client.pipelines) == 1
|
||||
commands = client.pipelines[0].commands
|
||||
assert [c[0] for c in commands] == ["MGET", "EVALSHA"]
|
||||
assert set(commands[0][1:]) == {
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
assert deployment in {"dep-a", "dep-b"}
|
||||
assert [[c[0] for c in pipeline.commands] for pipeline in client.pipelines] == [["EVALSHA"]]
|
||||
assert redis_cache.alone == []
|
||||
|
||||
shuffle = Router(model_list=[_deployment("dep-a")], routing_strategy="simple-shuffle")
|
||||
shuffle.update_redis_cache(cache=redis_cache)
|
||||
with request_redis_batch_scope() as request:
|
||||
shuffle.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
armed = request.prefetched["routing_read"]
|
||||
assert isinstance(armed, RoutingPrefetch)
|
||||
assert armed.keys == {CooldownCache.get_cooldown_cache_key("dep-a")} # no usage counters for shuffle
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing-v2"])
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
@pytest.mark.parametrize("with_limiter", [True, False])
|
||||
async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_redis_once(
|
||||
routing_strategy: str, with_limiter: bool
|
||||
):
|
||||
async def test_requests_within_the_read_interval_read_the_counters_from_redis_once(with_limiter: bool):
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache, routing_strategy=routing_strategy)
|
||||
router = _router(redis_cache)
|
||||
limiter = _limiter(redis_cache)
|
||||
request_round_trips: list[tuple[int, int]] = []
|
||||
request_round_trips: list[tuple[int, int]] = [] # mutable-ok: one entry per simulated request
|
||||
|
||||
for _ in range(3):
|
||||
pipeline_count = len(client.pipelines)
|
||||
|
|
@ -725,35 +610,24 @@ async def test_requests_within_the_cooldown_read_interval_read_cooldowns_from_re
|
|||
descriptors=[_descriptor("api_key", "k1", 10)],
|
||||
increments=[{"requests": 1}],
|
||||
)
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
await _route(router)
|
||||
request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count))
|
||||
|
||||
pipeline_mgets = [command for pipeline in client.pipelines for command in pipeline.commands if command[0] == "MGET"]
|
||||
alone_mgets = [keys for command, keys in redis_cache.alone if command == "MGET"]
|
||||
cooldown_keys = {
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
cooldown_mgets = [command[1:] for command in pipeline_mgets if cooldown_keys.intersection(command[1:])] + [
|
||||
keys for keys in alone_mgets if cooldown_keys.intersection(keys)
|
||||
]
|
||||
|
||||
assert len(cooldown_mgets) == 1
|
||||
assert len(_counter_mgets(router, client, redis_cache)) == 1
|
||||
if not with_limiter:
|
||||
assert request_round_trips[1:] == [(0, 0), (0, 0)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_requests_share_one_cooldown_read_per_interval():
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_concurrent_requests_share_one_counter_read_per_interval():
|
||||
client: Final = FakeClient(_lua_ok_replies)
|
||||
redis_cache: Final = FakeRedisCache(client)
|
||||
router: Final = _router(redis_cache)
|
||||
first_armed: Final = asyncio.Event()
|
||||
both_armed: Final = asyncio.Event()
|
||||
|
||||
async def route_after_both_requests_arm():
|
||||
async def route_after_both_requests_arm() -> str:
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
if first_armed.is_set():
|
||||
|
|
@ -761,127 +635,57 @@ async def test_concurrent_requests_share_one_cooldown_read_per_interval():
|
|||
else:
|
||||
first_armed.set()
|
||||
await both_armed.wait()
|
||||
return await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
return await _route(router)
|
||||
|
||||
deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm())
|
||||
cooldown_keys: Final = frozenset(
|
||||
{
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
)
|
||||
cooldown_mgets: Final = tuple(
|
||||
command
|
||||
for pipeline in client.pipelines
|
||||
for command in pipeline.commands
|
||||
if command[0] == "MGET" and cooldown_keys.intersection(command[1:])
|
||||
)
|
||||
|
||||
assert all(deployment["model_info"]["id"] in {"dep-a", "dep-b"} for deployment in deployments)
|
||||
assert len(cooldown_mgets) == 1
|
||||
assert set(deployments) <= {"dep-a", "dep-b"}
|
||||
assert len(_counter_mgets(router, client, redis_cache)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_prefetch_reads_cooldowns_again_once_the_read_interval_elapses(monkeypatch):
|
||||
first_time: Final = 1_000_000.0
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time)
|
||||
active_cooldown = {
|
||||
"exception_received": "429",
|
||||
"status_code": "429",
|
||||
"timestamp": _FAR_FUTURE,
|
||||
"cooldown_time": 60,
|
||||
}
|
||||
mget_results = iter((None, active_cooldown))
|
||||
|
||||
def replies(command: tuple[Any, ...]) -> Any:
|
||||
if command[0] == "MGET":
|
||||
result = next(mget_results)
|
||||
return [
|
||||
None if result is None or key != CooldownCache.get_cooldown_cache_key("dep-a") else json.dumps(result)
|
||||
for key in command[1:]
|
||||
]
|
||||
return _lua_ok_replies(command)
|
||||
|
||||
client = FakeClient(replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache, routing_strategy="simple-shuffle")
|
||||
cooldown_store = router.cooldown_cache.cooldown_store
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
dual_cache_module.time,
|
||||
"time",
|
||||
lambda: first_time + cooldown_store.redis_batch_cache_expiry + 1,
|
||||
)
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
|
||||
cooldown_keys = {
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
cooldown_mgets = [
|
||||
command
|
||||
for pipeline in client.pipelines
|
||||
for command in pipeline.commands
|
||||
if command[0] == "MGET" and cooldown_keys.intersection(command[1:])
|
||||
]
|
||||
assert len(cooldown_mgets) == 2
|
||||
assert deployment["model_info"]["id"] == "dep-b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(monkeypatch):
|
||||
first_time: Final = 1_000_000.0
|
||||
pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time)
|
||||
monkeypatch.setattr(
|
||||
"litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime",
|
||||
lambda: pinned,
|
||||
)
|
||||
async def test_the_prefetch_reads_the_counters_again_once_the_read_interval_elapses(pinned_clock: list[float]):
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache, routing_strategy="usage-based-routing-v2")
|
||||
cooldown_store = router.cooldown_cache.cooldown_store
|
||||
router = _router(redis_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await _route(router)
|
||||
|
||||
client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100})
|
||||
pinned_clock[0] = _START + router.lowesttpm_logger_v2.router_cache.redis_batch_cache_expiry + 1
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
deployment = await _route(router)
|
||||
|
||||
assert len(_counter_mgets(router, client, redis_cache)) == 2
|
||||
assert deployment == "dep-b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(pinned_clock: list[float]):
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
router = _router(redis_cache)
|
||||
usage_cache = router.lowesttpm_logger_v2.router_cache
|
||||
time_offset = cooldown_store.redis_batch_cache_expiry + 0.5
|
||||
|
||||
assert time_offset < usage_cache.redis_batch_cache_expiry
|
||||
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
await _route(router)
|
||||
|
||||
monkeypatch.setattr(dual_cache_module.time, "time", lambda: first_time + time_offset)
|
||||
for key in _counter_keys(router, ("dep-a",)):
|
||||
usage_cache.in_memory_cache.set_cache(key, 1)
|
||||
pinned_clock[0] = _START + usage_cache.redis_batch_cache_expiry + 1
|
||||
with request_redis_batch_scope():
|
||||
router.arm_routing_read_prefetch(_MODEL_GROUP, {})
|
||||
await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, messages=[{"role": "user", "content": "ping"}], request_kwargs={}
|
||||
)
|
||||
await _route(router)
|
||||
|
||||
cooldown_keys = frozenset(
|
||||
{
|
||||
CooldownCache.get_cooldown_cache_key("dep-a"),
|
||||
CooldownCache.get_cooldown_cache_key("dep-b"),
|
||||
}
|
||||
)
|
||||
second_pipeline_mgets = tuple(command for command in client.pipelines[1].commands if command[0] == "MGET")
|
||||
|
||||
assert len(client.pipelines) == 2
|
||||
assert len(second_pipeline_mgets) == 1
|
||||
assert frozenset(second_pipeline_mgets[0][1:]) == cooldown_keys
|
||||
assert frozenset(second_pipeline_mgets[0][1:]) == _counter_keys(router, ("dep-b",))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -896,6 +700,7 @@ async def test_two_backends_flush_concurrently_one_pipeline_each():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("pinned_clock")
|
||||
async def test_a_single_lua_group_rides_the_pipeline_with_the_armed_routing_read():
|
||||
client = FakeClient(_lua_ok_replies)
|
||||
redis_cache = FakeRedisCache(client)
|
||||
|
|
@ -1080,17 +885,31 @@ async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_c
|
|||
assert cache.in_memory_cache.get_cache("end_user_id:eu-miss") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cooldown_keys", "usage_keys", "expected"),
|
||||
[
|
||||
(("cooldown:a",), (), ROUTER_COOLDOWNS_TARGET),
|
||||
((), ("usage:a",), ROUTER_USAGE_TARGET),
|
||||
(("cooldown:a",), ("usage:a",), ROUTER_COOLDOWNS_USAGE_TARGET),
|
||||
],
|
||||
)
|
||||
def test_the_routing_read_family_follows_the_keys_that_are_actually_due(
|
||||
cooldown_keys: tuple[str, ...], usage_keys: tuple[str, ...], expected: str
|
||||
):
|
||||
"""A routing MGET is ``router_cooldowns`` when only cooldown keys go out, ``router_usage`` when the
|
||||
cooldowns were already in memory and only usage counters go out, and the combined family otherwise."""
|
||||
assert _routing_read_target(cooldown_keys, usage_keys) == expected
|
||||
@pytest.mark.asyncio
|
||||
async def test_awaiting_a_native_read_runs_the_native_hooks_then_flushes_the_native_pre_batch_once_per_request():
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(post_call_deadline=7, native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
native: Final = request.native()
|
||||
request.add_native_flush_hook(lambda: native.events.append("hook"))
|
||||
result: Final = request.native_result(FakeNativePending(native, lambda: "answer"))
|
||||
declared_only: Final = list(native.events)
|
||||
answer: Final = await result
|
||||
await request.flush_all()
|
||||
|
||||
assert request.native() is native
|
||||
assert [built.post_call_deadline for built in natives.built] == [7]
|
||||
assert declared_only == []
|
||||
assert answer == "answer"
|
||||
assert native.events == ["hook", "flush_pre", "resolve", "hook", "flush_pre"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_request_that_no_native_store_asks_for_builds_no_native_batches():
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(native_factory=natives) as request: # pyright: ignore[reportArgumentType] # test double
|
||||
await request.flush_all()
|
||||
await request.flush_post_call()
|
||||
|
||||
assert natives.built == []
|
||||
|
|
|
|||
|
|
@ -4,17 +4,105 @@ import importlib
|
|||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from collections.abc import Callable, Coroutine, Iterator
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
import boto3
|
||||
import httpx
|
||||
import pytest
|
||||
from pytest_socket import enable_socket, socket_allow_hosts
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
RouterRedis: TypeAlias = tuple["RedisCache", asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]]
|
||||
RedisWriteGate: TypeAlias = tuple[asyncio.Event, asyncio.Event, asyncio.Event]
|
||||
|
||||
|
||||
def redis_cache_double() -> "MagicMock":
|
||||
"""A `RedisCache` double for router tests. The native router stores connect from its settings, which name
|
||||
a closed local port, so they serve their memory tier while the double answers the Python Redis calls."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.redis.native_settings import NativeRedisSettings
|
||||
|
||||
redis: Final = MagicMock(spec=RedisCache)
|
||||
redis.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}
|
||||
redis.namespace = None
|
||||
redis.native_settings = NativeRedisSettings(redis_kwargs=redis.redis_kwargs, namespace=None, default_ttl=60)
|
||||
return redis
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_write_gate() -> RedisWriteGate:
|
||||
released: Final = asyncio.Event()
|
||||
released.set()
|
||||
return asyncio.Event(), released, asyncio.Event()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def router_redis(
|
||||
redis_write_gate: RedisWriteGate,
|
||||
) -> AsyncIterator[RouterRedis]:
|
||||
from unittest.mock import patch
|
||||
|
||||
from fakeredis import FakeRedis, FakeServer
|
||||
from fakeredis.aioredis import FakeConnection
|
||||
from opentelemetry import trace
|
||||
from redis.asyncio import BlockingConnectionPool
|
||||
from redis.exceptions import ConnectionError
|
||||
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
commands: Final[asyncio.Queue[tuple[tuple[object, ...], str | None, str | None]]] = asyncio.Queue()
|
||||
started, released, unavailable = redis_write_gate
|
||||
|
||||
class ObservedConnection(FakeConnection):
|
||||
async def send_command(self, *args: object, **kwargs: object) -> None: # kwargs-ok: Redis transport signature
|
||||
commands.put_nowait((args, current_service_target(), getattr(trace.get_current_span(), "name", None)))
|
||||
await super().send_command(*args, **kwargs)
|
||||
|
||||
async def send_packed_command(self, command: Iterable[bytes] | bytes | str, check_health: bool = True) -> None:
|
||||
chunks: Final = (
|
||||
(command.encode(),)
|
||||
if isinstance(command, str)
|
||||
else (command,)
|
||||
if isinstance(command, bytes)
|
||||
else command
|
||||
)
|
||||
wire: Final = b"".join(chunks)
|
||||
if b"INCRBYFLOAT" in wire:
|
||||
commands.put_nowait(
|
||||
((wire,), current_service_target(), getattr(trace.get_current_span(), "name", None))
|
||||
)
|
||||
started.set()
|
||||
await released.wait()
|
||||
if unavailable.is_set():
|
||||
raise ConnectionError("scripted write outage")
|
||||
await super().send_packed_command(command, check_health=check_health)
|
||||
|
||||
server: Final = FakeServer()
|
||||
client: Final = FakeRedis(server=server)
|
||||
pool: Final = BlockingConnectionPool(connection_class=ObservedConnection, server=server)
|
||||
try:
|
||||
with (
|
||||
patch("redis.Redis.from_url", return_value=client),
|
||||
patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool),
|
||||
):
|
||||
yield RedisCache(url="redis://router-unit.example:6379/0"), commands
|
||||
finally:
|
||||
released.set()
|
||||
await pool.aclose()
|
||||
client.close()
|
||||
|
||||
|
||||
HOST_ENVIRONMENT_ALLOWLIST: Final = frozenset(
|
||||
(
|
||||
"PATH",
|
||||
|
|
@ -290,7 +378,9 @@ def isolate_litellm_globals() -> Iterator[None]:
|
|||
original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks)
|
||||
original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS}
|
||||
original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)}
|
||||
original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__}
|
||||
original_clients: Final = {
|
||||
name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__
|
||||
}
|
||||
original_loggers: Final = {
|
||||
logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters))
|
||||
for logger in ALL_LOGGERS
|
||||
|
|
|
|||
|
|
@ -4,6 +4,8 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from typing import Final
|
||||
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.redis.stores._connection import native_client
|
||||
from litellm.rust_bridge._native import NativeRedisClient, native_redis_breaker_states
|
||||
|
||||
|
||||
|
|
@ -45,3 +47,21 @@ def test_a_discarded_client_leaves_the_process_registry(monkeypatch: pytest.Monk
|
|||
|
||||
assert while_connected == before + 1
|
||||
assert len(native_redis_breaker_states()) == before
|
||||
|
||||
|
||||
def test_a_redis_cache_captures_what_native_stores_need_at_construction() -> None:
|
||||
cache: Final = RedisCache(host="127.0.0.1", port=1, namespace="ns", socket_timeout=0.2)
|
||||
|
||||
assert cache.native_settings is not None
|
||||
assert dict(cache.native_settings.redis_kwargs) == {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}
|
||||
assert cache.native_settings.namespace == "ns"
|
||||
assert cache.native_settings.default_ttl == cache.default_ttl
|
||||
|
||||
|
||||
class _UnbuiltRedisCache(RedisCache):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def test_a_redis_cache_stand_in_that_skips_the_constructor_leaves_native_stores_on_memory() -> None:
|
||||
assert native_client(_UnbuiltRedisCache().native_settings) is None
|
||||
|
|
|
|||
260
tests/unit/redis/stores/test_router_cooldowns.py
Normal file
260
tests/unit/redis/stores/test_router_cooldowns.py
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
import ast
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass, field, replace
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.redis_batch import RequestRedisBatches
|
||||
from litellm.caching.redis_cache import RedisCache, redis_breaker_state_counts
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.redis.stores._connection import native_client
|
||||
from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore
|
||||
from litellm.rust_bridge._native import NativeRedisClient, NativeRouterCooldowns
|
||||
from litellm.rust_bridge.host import settings
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNativeRequest, FakeNatives
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _WireCooldown:
|
||||
"""One cooldown as main's ``RedisCache`` writes it (``str(value)``) and reads it (``json.loads`` then
|
||||
``ast.literal_eval``). The same literals sit in ``litellm-rust/crates/redis-store/tests/support/mod.rs``,
|
||||
where real Redis is checked to hold them after the Rust store writes."""
|
||||
|
||||
cooldown: CooldownCacheValue
|
||||
str_dict: str
|
||||
json: str
|
||||
ttl_seconds: float
|
||||
redis_ttl_seconds: int
|
||||
|
||||
|
||||
_WIRE_COOLDOWNS: Final = (
|
||||
_WireCooldown(
|
||||
cooldown={
|
||||
"exception_received": "RateLimitError: limited é 'quoted'",
|
||||
"status_code": "429",
|
||||
"timestamp": 1730000000.25,
|
||||
"cooldown_time": 60,
|
||||
},
|
||||
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}',
|
||||
ttl_seconds=90.5,
|
||||
redis_ttl_seconds=90,
|
||||
),
|
||||
_WireCooldown(
|
||||
cooldown={
|
||||
"exception_received": "Timeout",
|
||||
"status_code": "408",
|
||||
"timestamp": 1730000000.5,
|
||||
"cooldown_time": 0.5,
|
||||
},
|
||||
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}',
|
||||
ttl_seconds=30.0,
|
||||
redis_ttl_seconds=30,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _entry(model_id: str = "dep") -> CooldownCacheValue:
|
||||
return {"exception_received": model_id, "status_code": "429", "timestamp": 1000.0, "cooldown_time": 120.0}
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeCooldowns:
|
||||
client: NativeRedisClient | None
|
||||
stored: dict[str, CooldownCacheValue] = field(default_factory=dict)
|
||||
fails: dict[str, int] = field(default_factory=dict)
|
||||
targets: list[str | None] = field(default_factory=list)
|
||||
smallest: float | None = None
|
||||
declared: list[tuple[str, ...]] = field(default_factory=list)
|
||||
|
||||
def put(self, model_id: str, cooldown: Mapping[str, object], ttl_seconds: float) -> None:
|
||||
self.targets.append(current_service_target())
|
||||
self.stored[model_id] = CooldownCacheValue(**cooldown) # pyright: ignore[reportArgumentType] # test double
|
||||
|
||||
def increment_allowed_fails(self, model_id: str, suffix: str | None, ttl_seconds: float) -> int:
|
||||
self.targets.append(current_service_target())
|
||||
self.fails[model_id] = self.fails.get(model_id, 0) + 1
|
||||
return self.fails[model_id]
|
||||
|
||||
def active(self, model_ids: Sequence[str]) -> list[tuple[str, CooldownCacheValue]]:
|
||||
self.targets.append(current_service_target())
|
||||
return [(model_id, self.stored[model_id]) for model_id in model_ids if model_id in self.stored]
|
||||
|
||||
def aactive(self, model_ids: Sequence[str]) -> "asyncio.Future[list[tuple[str, CooldownCacheValue]]]":
|
||||
future: Final[asyncio.Future[list[tuple[str, CooldownCacheValue]]]] = asyncio.get_running_loop().create_future()
|
||||
future.set_result(self.active(model_ids))
|
||||
return future
|
||||
|
||||
def declare_active(
|
||||
self, request: FakeNativeRequest, model_ids: Sequence[str]
|
||||
) -> FakeNativePending[list[tuple[str, CooldownCacheValue]]]:
|
||||
self.targets.append(current_service_target())
|
||||
self.declared.append(tuple(model_ids))
|
||||
return FakeNativePending(request, lambda: self.active(model_ids))
|
||||
|
||||
def minimum(self, model_ids: Sequence[str]) -> float | None:
|
||||
return self.smallest
|
||||
|
||||
def evict_local(self, model_id: str) -> None:
|
||||
self.stored.pop(model_id, None)
|
||||
|
||||
def flush_local_counters(self) -> None:
|
||||
self.fails.clear()
|
||||
|
||||
def snapshot(self) -> dict[str, CooldownCacheValue]:
|
||||
return dict(self.stored)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Handles:
|
||||
built: list[FakeCooldowns] = field(default_factory=list)
|
||||
|
||||
def __call__(self, client: NativeRedisClient | None) -> FakeCooldowns:
|
||||
handle: Final = FakeCooldowns(client)
|
||||
self.built.append(handle)
|
||||
return handle
|
||||
|
||||
|
||||
def _client(namespace: str | None = None) -> NativeRedisClient:
|
||||
client: Final = native_client(RedisCache(host="redis.internal", port=6380, namespace=namespace))
|
||||
assert client is not None
|
||||
return client
|
||||
|
||||
|
||||
def test_late_attach_rebuilds_once_and_later_attaches_are_ignored() -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterCooldownStore(None, handle=handles)
|
||||
store.put("dep", _entry(), 60.0)
|
||||
first: Final = _client()
|
||||
store.attach_redis(None)
|
||||
store.attach_redis(first)
|
||||
store.attach_redis(_client("other"))
|
||||
assert [handle.client for handle in handles.built] == [None, first]
|
||||
assert store.active_sync(["dep"], None) == ()
|
||||
|
||||
|
||||
def test_writes_and_reads_carry_the_router_cooldowns_service_target() -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterCooldownStore(_client(), handle=handles)
|
||||
store.put("dep", _entry(), 60.0)
|
||||
assert store.increment_allowed_fails("dep", "timeout", 30.0) == 1
|
||||
assert store.active_sync(["dep", "other"], None) == (("dep", _entry()),)
|
||||
assert handles.built[0].targets == ["router_cooldowns"] * 3
|
||||
assert current_service_target() is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("smallest", "expected"), ((None, 60.0), (0, 60.0), (5, 5)))
|
||||
def test_minimum_falls_back_to_the_default_when_nothing_or_zero_is_stored(
|
||||
smallest: float | None, expected: float
|
||||
) -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterCooldownStore(None, handle=handles)
|
||||
handles.built[0].smallest = smallest
|
||||
assert store.minimum(["dep"], 60.0, None) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", _WIRE_COOLDOWNS, ids=("rate-limit", "timeout"))
|
||||
def test_the_wire_cases_are_what_mains_codec_writes_and_reads(case: _WireCooldown) -> None:
|
||||
assert str(case.cooldown) == case.str_dict
|
||||
assert ast.literal_eval(case.str_dict) == case.cooldown
|
||||
assert json.loads(case.json) == case.cooldown
|
||||
assert int(case.ttl_seconds) == case.redis_ttl_seconds
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setting",
|
||||
({"ssl_ciphers": "HIGH"}, {"ssl": True, "ssl_validate_ocsp": True}, {"credential_provider": "not-a-provider"}),
|
||||
)
|
||||
def test_a_redis_setting_the_native_client_cannot_honor_never_fails_construction(
|
||||
setting: dict[str, object], caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
redis: Final = RedisCache(host="127.0.0.1", port=1, **setting)
|
||||
with caplog.at_level(logging.WARNING):
|
||||
store: Final = RouterCooldownStore(native_client(redis))
|
||||
cooldown: Final = CooldownCacheValue(
|
||||
exception_received="model", status_code="429", timestamp=time.time(), cooldown_time=120.0
|
||||
)
|
||||
store.put("model", cooldown, 60.0)
|
||||
|
||||
assert store.active_sync(["model"], None) == (("model", cooldown),)
|
||||
assert next(iter(setting)) in caplog.text
|
||||
|
||||
|
||||
class _RecordingSecrets(CustomSecretManager):
|
||||
def __init__(self, values: Mapping[str, str]) -> None:
|
||||
super().__init__(secret_manager_name="router_cooldowns_test") # pyright: ignore[reportUnknownMemberType] # CustomSecretManager takes untyped **kwargs
|
||||
self.values: Final = values
|
||||
self.asked: tuple[str, ...] = ()
|
||||
|
||||
async def async_read_secret(
|
||||
self, secret_name: str, optional_params: dict[str, object] | None = None, timeout: object = None
|
||||
) -> str | None:
|
||||
raise AssertionError("get_secret reads custom managers synchronously")
|
||||
|
||||
def sync_read_secret(
|
||||
self, secret_name: str, optional_params: dict[str, object] | None = None, timeout: object = None
|
||||
) -> str | None:
|
||||
self.asked = (*self.asked, secret_name)
|
||||
return self.values.get(secret_name)
|
||||
|
||||
|
||||
def test_the_native_client_reads_redis_secrets_through_the_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("REDIS_PASSWORD", raising=False)
|
||||
manager: Final = _RecordingSecrets({"REDIS_PASSWORD": "vault-password"})
|
||||
monkeypatch.setattr(litellm, "secret_manager_client", manager)
|
||||
monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only"))
|
||||
configured: Final = settings.secret_manager
|
||||
monkeypatch.setattr(settings, "secret_manager", lambda: replace(configured(), native=False))
|
||||
|
||||
NativeRedisClient.connect({"host": "127.0.0.1", "port": 1}, namespace="secret-manager-test")
|
||||
|
||||
assert "REDIS_PASSWORD" in manager.asked
|
||||
|
||||
|
||||
def test_the_native_breaker_counts_in_the_redis_breaker_state_gauge(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true")
|
||||
monkeypatch.setenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", "1")
|
||||
open_before: Final = redis_breaker_state_counts().get("open", 0)
|
||||
store: Final = NativeRouterCooldowns(
|
||||
NativeRedisClient.connect({"host": "127.0.0.1", "port": 1}, namespace="breaker-gauge-test")
|
||||
)
|
||||
|
||||
store.put("model", _entry("model"), 60.0)
|
||||
|
||||
assert redis_breaker_state_counts().get("open", 0) == open_before + 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_declared_active_read_answers_after_the_native_pre_flush_under_the_router_cooldowns_target() -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterCooldownStore(_client(), handle=handles)
|
||||
natives: Final = FakeNatives()
|
||||
request: Final = RequestRedisBatches(native_factory=natives) # pyright: ignore[reportArgumentType] # test double
|
||||
store.put("dep", _entry(), 60.0)
|
||||
|
||||
declared: Final = store.declare_active(request, ["dep", "other"])
|
||||
assert declared is not None
|
||||
assert natives.built[0].events == []
|
||||
|
||||
assert tuple(await declared) == (("dep", _entry()),)
|
||||
assert natives.built[0].events == ["flush_pre", "resolve"]
|
||||
assert handles.built[0].declared == [("dep", "other")]
|
||||
assert handles.built[0].targets[1] == "router_cooldowns"
|
||||
|
||||
|
||||
def test_a_store_without_redis_declares_nothing_on_the_request() -> None:
|
||||
natives: Final = FakeNatives()
|
||||
request: Final = RequestRedisBatches(native_factory=natives) # pyright: ignore[reportArgumentType] # test double
|
||||
|
||||
assert RouterCooldownStore(None, handle=Handles()).declare_active(request, ["dep"]) is None
|
||||
assert natives.built == []
|
||||
121
tests/unit/redis/stores/test_router_usage.py
Normal file
121
tests/unit/redis/stores/test_router_usage.py
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Final, TypeVar
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.redis_batch import drain_post_call_redis_batches, request_redis_batch_scope
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.redis.stores._connection import native_client
|
||||
from litellm.redis.stores.router_usage import NativeCounter, RouterUsageStore, UsageCounter
|
||||
from litellm.rust_bridge._native import NativeRedisClient
|
||||
from tests.unit.caching.test_redis_batch import FakeNatives
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def _done(value: T) -> "asyncio.Future[T]":
|
||||
future: Final[asyncio.Future[T]] = asyncio.get_running_loop().create_future()
|
||||
future.set_result(value)
|
||||
return future
|
||||
|
||||
|
||||
@dataclass
|
||||
class FakeUsage:
|
||||
connected: bool
|
||||
counts: dict[NativeCounter, float] = field(default_factory=dict)
|
||||
calls: list[tuple[str, object, str | None]] = field(default_factory=list)
|
||||
|
||||
def increment(self, counter: NativeCounter, value: float, ttl_seconds: int) -> "asyncio.Future[float]":
|
||||
self.calls.append(("increment", (counter, value, ttl_seconds), current_service_target()))
|
||||
self.counts[counter] = self.counts.get(counter, 0) + value
|
||||
return _done(self.counts[counter])
|
||||
|
||||
def increment_deployment(
|
||||
self,
|
||||
deployment_id: str,
|
||||
model: str | None,
|
||||
minute: str,
|
||||
tokens: float,
|
||||
requests: int,
|
||||
) -> "asyncio.Future[None]":
|
||||
self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target()))
|
||||
for kind, value in (("tpm", tokens), ("rpm", requests)):
|
||||
counter: NativeCounter = (deployment_id, model, minute, kind) # pyright: ignore[reportAssignmentType] # literal kinds
|
||||
self.counts[counter] = self.counts.get(counter, 0) + value
|
||||
return _done(None)
|
||||
|
||||
def get_many(self, counters: Sequence[NativeCounter]) -> "asyncio.Future[list[float | None]]":
|
||||
return _done([self.counts.get(counter) for counter in counters])
|
||||
|
||||
def flush_local_counters(self) -> None:
|
||||
self.counts.clear()
|
||||
|
||||
def snapshot(self) -> dict[str, float]:
|
||||
return {}
|
||||
|
||||
|
||||
@dataclass
|
||||
class Handles:
|
||||
built: list[FakeUsage] = field(default_factory=list)
|
||||
|
||||
def __call__(self, client: NativeRedisClient | None) -> FakeUsage:
|
||||
self.built.append(FakeUsage(client is not None))
|
||||
return self.built[-1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage_target() -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterUsageStore(None, handle=handles)
|
||||
counter: Final = UsageCounter("dep", None, "12-34", "rpm")
|
||||
|
||||
await store.async_increment(counter, 1, 60)
|
||||
await store.increment_deployment("dep", "model", "12-34", 7.5, 1, None)
|
||||
|
||||
assert handles.built[0].calls == [
|
||||
("increment", (("dep", None, "12-34", "rpm"), 1, 60), "router_usage"),
|
||||
("deployment", ("dep", "model", "12-34", 7.5, 1), "router_usage"),
|
||||
]
|
||||
assert await store.async_batch_get((counter, UsageCounter("dep", "model", "12-34", "tpm"))) == [1, 7.5]
|
||||
|
||||
|
||||
def test_late_attach_rebuilds_the_handle_once() -> None:
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterUsageStore(None, handle=handles)
|
||||
store.attach_redis(native_client(RedisCache(host="redis.internal").native_settings))
|
||||
store.attach_redis(native_client(RedisCache(host="other").native_settings))
|
||||
assert [handle.connected for handle in handles.built] == [False, True]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ("flush_cache", "reset"))
|
||||
async def test_router_reset_clears_usage_and_failure_counters(operation: str) -> None:
|
||||
router: Final = Router(model_list=[])
|
||||
usage: Final = router._core_usage_store
|
||||
cooldowns: Final = router._core_cooldown_cache.cooldown_store
|
||||
counters: Final = (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm"))
|
||||
await usage.increment_deployment("dep", "model", "12-34", 7, 1, None)
|
||||
assert await usage.async_batch_get(counters) == [7, 1]
|
||||
assert cooldowns.increment_allowed_fails("dep", None, 30.0) == 1
|
||||
getattr(router, operation)()
|
||||
assert await usage.async_batch_get(counters) == [None, None]
|
||||
assert cooldowns.increment_allowed_fails("dep", None, 30.0) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inside_a_request_deployment_increments_go_out_on_their_own_trip_like_main():
|
||||
handles: Final = Handles()
|
||||
store: Final = RouterUsageStore(None, handle=handles)
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
await store.increment_deployment("dep", "model", "12-34", 7.5, 1, None)
|
||||
with request_redis_batch_scope(post_call_deadline=60, native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
await store.increment_deployment("dep", "model", "12-34", 2, 1, None)
|
||||
await drain_post_call_redis_batches()
|
||||
|
||||
assert [call[0] for call in handles.built[0].calls] == ["deployment", "deployment"]
|
||||
assert all(not native.events for native in natives.built)
|
||||
|
|
@ -16038,9 +16038,7 @@ class TestHealthFallbackDispatch:
|
|||
{"primary-id": {"is_healthy": True, "timestamp": time.time()}}
|
||||
)
|
||||
else:
|
||||
router.cooldown_cache.cooldown_store.delete_cache(
|
||||
router.cooldown_cache.get_cooldown_cache_key("primary-id")
|
||||
)
|
||||
router.cooldown_cache.cooldown_store.evict_local("primary-id")
|
||||
recovered: Final[dict[str, object]] = {"session_id": "outage"}
|
||||
assert await self._request(router, surface, stream, recovered) == "primary"
|
||||
assert recovered["routing_decision"]["routed_model"] == "primary"
|
||||
|
|
@ -16216,9 +16214,7 @@ class TestHealthFallbackDispatch:
|
|||
{"primary-id": {"is_healthy": True, "timestamp": time.time()}}
|
||||
)
|
||||
else:
|
||||
router.cooldown_cache.cooldown_store.delete_cache(
|
||||
router.cooldown_cache.get_cooldown_cache_key("primary-id")
|
||||
)
|
||||
router.cooldown_cache.cooldown_store.evict_local("primary-id")
|
||||
recovered: Final[dict[str, object]] = {"session_id": "pinned"}
|
||||
assert await self._request(router, "chat", False, recovered) == "primary"
|
||||
assert recovered["routing_decision"]["cause"] == "session_affinity_pin"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@
|
|||
|
||||
import logging
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -1846,23 +1845,20 @@ async def test_tag_fail_open_survives_deployment_narrowing(
|
|||
) -> None:
|
||||
router: Final = _fail_open_router(route=route)
|
||||
requested_model: Final = "chat" if route == "model" else "chat-test" if route == "wildcard" else "public-chat"
|
||||
with patch("time.time", return_value=1_700_000_000.0):
|
||||
if narrowing == "cooldown":
|
||||
router.cooldown_cache.add_deployment_to_cooldown(
|
||||
model_id="opted-in",
|
||||
original_exception=RuntimeError("unavailable"),
|
||||
exception_status=500,
|
||||
cooldown_time=60,
|
||||
)
|
||||
candidates: Final = await router.async_get_healthy_deployments(
|
||||
model=requested_model,
|
||||
request_kwargs={
|
||||
"litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"},
|
||||
**(
|
||||
{"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}
|
||||
),
|
||||
},
|
||||
if narrowing == "cooldown":
|
||||
router.cooldown_cache.add_deployment_to_cooldown(
|
||||
model_id="opted-in",
|
||||
original_exception=RuntimeError("unavailable"),
|
||||
exception_status=500,
|
||||
cooldown_time=60,
|
||||
)
|
||||
candidates: Final = await router.async_get_healthy_deployments(
|
||||
model=requested_model,
|
||||
request_kwargs={
|
||||
"litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"},
|
||||
**({"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}),
|
||||
},
|
||||
)
|
||||
assert isinstance(candidates, (list, tuple))
|
||||
assert tuple(deployment["model_info"]["id"] for deployment in candidates) == ("remaining",)
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
|||
)
|
||||
|
||||
|
||||
|
||||
def _repinning_router(enabled: bool | None = True) -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
|
|
@ -100,7 +101,7 @@ async def test_session_repin_requires_opt_in_unavailable_original_and_success(
|
|||
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected}
|
||||
assert await router.cache.async_get_cache(_repin_key("caller-2")) == {"model_id": "original"}
|
||||
assert await router.cache.async_get_cache(user_key) == {"model_id": "original"}
|
||||
router.cooldown_cache.cooldown_store.flush_cache()
|
||||
router.cooldown_cache.cooldown_store.evict_local("original")
|
||||
await _record_replacement_success(router, "other")
|
||||
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected}
|
||||
finally:
|
||||
|
|
@ -129,7 +130,7 @@ async def test_concurrent_successful_replacements_keep_first_winner() -> None:
|
|||
await asyncio.gather(_record_replacement_success(router), _record_replacement_success(router, "other"))
|
||||
winner: Final = await router.cache.async_get_cache(_repin_key())
|
||||
assert winner in ({"model_id": "replacement"}, {"model_id": "other"})
|
||||
router.cooldown_cache.cooldown_store.flush_cache()
|
||||
router.cooldown_cache.cooldown_store.evict_local("original")
|
||||
await _record_replacement_success(router, "original")
|
||||
assert await router.cache.async_get_cache(_repin_key()) == winner
|
||||
finally:
|
||||
|
|
@ -320,49 +321,6 @@ async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: M
|
|||
router.discard()
|
||||
|
||||
|
||||
class _ReadCountingCache(InMemoryCache):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.read_keys: tuple[str, ...] = ()
|
||||
|
||||
def get_cache(self, key: str, **kwargs: object) -> object:
|
||||
self.read_keys += (key,)
|
||||
return super().get_cache(key, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"])
|
||||
async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None:
|
||||
router: Final = _repinning_router(case != "disabled")
|
||||
reads: Final = _ReadCountingCache()
|
||||
try:
|
||||
router.cooldown_cache.cooldown_store.in_memory_cache = reads
|
||||
callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
|
||||
if case == "group-disabled":
|
||||
callback.model_group_affinity_config = {"repin-group": []}
|
||||
await router._repin_session_after_success(
|
||||
{
|
||||
"model_info": {"id": "replacement"},
|
||||
"metadata": {
|
||||
"deployment_model_name": "repin-group",
|
||||
**({"session_id": "repin-session"} if case != "no-session" else {}),
|
||||
SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session",
|
||||
},
|
||||
},
|
||||
{"status": "completed"},
|
||||
)
|
||||
expected: Final = (
|
||||
{
|
||||
router.cooldown_cache.get_cooldown_cache_key(model_id)
|
||||
for model_id in ("original", "replacement", "other")
|
||||
}
|
||||
if case == "active"
|
||||
else set()
|
||||
)
|
||||
assert set(reads.read_keys) == expected
|
||||
finally:
|
||||
router.discard()
|
||||
|
||||
|
||||
class MockResponse:
|
||||
def __init__(self, json_data, status_code):
|
||||
self._json_data = json_data
|
||||
|
|
|
|||
|
|
@ -1,629 +1,43 @@
|
|||
"""
|
||||
Unit tests for CooldownCache exception masking functionality
|
||||
"""
|
||||
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
# Add the parent directory to the system path
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
|
||||
from litellm.redis.stores.router_cooldowns import RouterCooldownStore
|
||||
from litellm.router_utils.cooldown_cache import CooldownCache
|
||||
|
||||
|
||||
class TestCooldownCacheExceptionMasking:
|
||||
"""Test suite for CooldownCache exception masking functionality"""
|
||||
|
||||
@pytest.fixture
|
||||
def cooldown_cache(self):
|
||||
"""Create a CooldownCache instance for testing"""
|
||||
mock_dual_cache = MagicMock(spec=DualCache)
|
||||
return CooldownCache(cache=mock_dual_cache, default_cooldown_time=60.0)
|
||||
|
||||
def test_exception_masker_initialization(self, cooldown_cache):
|
||||
"""Test that the exception masker is properly initialized"""
|
||||
assert isinstance(cooldown_cache.exception_masker, SensitiveDataMasker)
|
||||
assert cooldown_cache.exception_masker.visible_prefix == 50
|
||||
assert cooldown_cache.exception_masker.visible_suffix == 0
|
||||
assert cooldown_cache.exception_masker.mask_char == "*"
|
||||
|
||||
def test_short_exception_string_not_masked(self, cooldown_cache):
|
||||
"""Test that short exception strings are not masked"""
|
||||
short_exception = "Short error"
|
||||
model_id = "test-model"
|
||||
exception_status = 500
|
||||
cooldown_time = 30.0
|
||||
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=Exception(short_exception),
|
||||
exception_status=exception_status,
|
||||
cooldown_time=cooldown_time,
|
||||
@pytest.mark.parametrize("message", ("Short error", "Rate limit " + "secret" * 100))
|
||||
@pytest.mark.parametrize("ttl", (None, 5.5, 30.0))
|
||||
def test_cooldown_masks_exception_and_preserves_status_timestamp_and_ttl(message: str, ttl: float | None) -> None:
|
||||
now: Final = time.time()
|
||||
store: Final = RouterCooldownStore(None)
|
||||
cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: now)
|
||||
cache.add_deployment_to_cooldown("dep", Exception(message), 429, ttl)
|
||||
expected_ttl: Final = 60.0 if ttl is None else ttl
|
||||
expected_message: Final = message if len(message) <= 50 else message[:50] + "*" * (len(message) - 50)
|
||||
assert cache.get_active_cooldowns(["dep"], None) == [
|
||||
(
|
||||
"dep",
|
||||
{
|
||||
"exception_received": expected_message,
|
||||
"status_code": "429",
|
||||
"timestamp": now,
|
||||
"cooldown_time": expected_ttl,
|
||||
},
|
||||
)
|
||||
|
||||
# Short exception should not be masked
|
||||
assert cooldown_data["exception_received"] == short_exception
|
||||
assert cooldown_key == f"deployment:{model_id}:cooldown"
|
||||
|
||||
def test_long_exception_string_masked(self, cooldown_cache):
|
||||
"""Test that long exception strings are properly masked"""
|
||||
# Create a long exception string that simulates prompt leakage
|
||||
long_exception = (
|
||||
"litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occurred - "
|
||||
"No deployments available for selected model, Try again in 5 seconds. "
|
||||
"Passed model=anthropic_claude_sonnet_4_v1_0. pre-call-checks=False, "
|
||||
"cooldown_list=[('deepseek_r1-eastus', {'exception_received': "
|
||||
"'litellm.RateLimitError: RateLimitError: Azure_aiException - "
|
||||
'{"error":{"code":"Invalid input","status":422,"message":"invalid input error",'
|
||||
'"details":[{"type":"model_attributes_type","loc":["body"],'
|
||||
'"msg":"Tell me a story about a dragon and a princess in a magical kingdom '
|
||||
"where the dragon is actually protecting the princess from an evil wizard "
|
||||
'who wants to steal her magical powers and use them to conquer the world"}]}'
|
||||
)
|
||||
|
||||
model_id = "test-model"
|
||||
exception_status = 429
|
||||
cooldown_time = 60.0
|
||||
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=Exception(long_exception),
|
||||
exception_status=exception_status,
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
|
||||
masked_exception = cooldown_data["exception_received"]
|
||||
|
||||
# Should start with first 50 characters
|
||||
assert masked_exception.startswith(long_exception[:50])
|
||||
|
||||
# Should contain masking characters
|
||||
assert "*" in masked_exception
|
||||
|
||||
# Should be same length (prefix + asterisks)
|
||||
assert len(masked_exception) == len(long_exception)
|
||||
|
||||
# Should not contain the sensitive prompt content
|
||||
assert "Tell me a story about a dragon" not in masked_exception
|
||||
assert "magical kingdom" not in masked_exception
|
||||
|
||||
# Should preserve the error type information at the beginning (first 50 chars)
|
||||
assert masked_exception.startswith("litellm.proxy.proxy_server._handle_llm_api_excepti")
|
||||
|
||||
def test_exception_with_api_keys_masked(self, cooldown_cache):
|
||||
"""Test that API keys in exceptions are properly masked"""
|
||||
exception_with_key = (
|
||||
"Authentication failed with api_key=sk-9876567890abcdefghijklmnopqrstuvwxyz "
|
||||
"and token=bearer_token_123456789 for model gpt-4"
|
||||
)
|
||||
|
||||
model_id = "test-model"
|
||||
exception_status = 401
|
||||
cooldown_time = 30.0
|
||||
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=Exception(exception_with_key),
|
||||
exception_status=exception_status,
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
|
||||
masked_exception = cooldown_data["exception_received"]
|
||||
|
||||
# Should mask the sensitive content while preserving structure
|
||||
assert masked_exception.startswith("Authentication failed with api_key=sk-98765678")
|
||||
assert "*" in masked_exception
|
||||
assert len(masked_exception) == len(exception_with_key)
|
||||
|
||||
def test_cooldown_data_structure(self, cooldown_cache):
|
||||
"""Test that the cooldown data structure is correctly formed"""
|
||||
exception_msg = "Test exception for structure validation"
|
||||
model_id = "test-model"
|
||||
exception_status = 500
|
||||
cooldown_time = 45.0
|
||||
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=Exception(exception_msg),
|
||||
exception_status=exception_status,
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
|
||||
# Verify cooldown data structure
|
||||
assert isinstance(cooldown_data, dict)
|
||||
assert "exception_received" in cooldown_data
|
||||
assert "status_code" in cooldown_data
|
||||
assert "timestamp" in cooldown_data
|
||||
assert "cooldown_time" in cooldown_data
|
||||
|
||||
# Verify data types
|
||||
assert isinstance(cooldown_data["exception_received"], str)
|
||||
assert isinstance(cooldown_data["status_code"], str)
|
||||
assert isinstance(cooldown_data["timestamp"], float)
|
||||
assert isinstance(cooldown_data["cooldown_time"], float)
|
||||
|
||||
# Verify values
|
||||
assert cooldown_data["status_code"] == str(exception_status)
|
||||
assert cooldown_data["cooldown_time"] == cooldown_time
|
||||
assert cooldown_data["exception_received"] == exception_msg
|
||||
|
||||
def test_exception_object_conversion(self, cooldown_cache):
|
||||
"""Test that different exception types are properly converted to strings"""
|
||||
# Test with different exception types
|
||||
exceptions = [
|
||||
ValueError("Invalid value provided"),
|
||||
KeyError("Missing required key"),
|
||||
RuntimeError("Runtime error occurred"),
|
||||
Exception("Generic exception"),
|
||||
]
|
||||
|
||||
for exc in exceptions:
|
||||
model_id = f"test-model-{exc.__class__.__name__}"
|
||||
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=exc,
|
||||
exception_status=500,
|
||||
cooldown_time=30.0,
|
||||
)
|
||||
|
||||
# Should successfully convert exception to string
|
||||
assert isinstance(cooldown_data["exception_received"], str)
|
||||
assert str(exc) == cooldown_data["exception_received"] # Short exceptions not masked
|
||||
|
||||
def test_masking_preserves_error_debugging_info(self, cooldown_cache):
|
||||
"""Test that masking preserves essential debugging information"""
|
||||
debugging_exception = (
|
||||
"RateLimitError: Rate limit exceeded for model gpt-4. "
|
||||
"Current usage: 1000 tokens/minute. Limit: 500 tokens/minute. "
|
||||
"Request details: model=gpt-4, user_id=user123, "
|
||||
"prompt='Write a comprehensive analysis of the economic implications "
|
||||
"of artificial intelligence adoption in the healthcare sector, including "
|
||||
"potential cost savings, job displacement, and regulatory challenges'"
|
||||
)
|
||||
|
||||
model_id = "gpt-4-deployment"
|
||||
exception_status = 429
|
||||
cooldown_time = 120.0
|
||||
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=Exception(debugging_exception),
|
||||
exception_status=exception_status,
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
|
||||
masked_exception = cooldown_data["exception_received"]
|
||||
|
||||
# Should preserve error type and initial debugging info (first 50 chars)
|
||||
assert masked_exception.startswith("RateLimitError: Rate limit exceeded for model gpt-")
|
||||
|
||||
# Should mask the prompt content
|
||||
assert "Write a comprehensive analysis" not in masked_exception
|
||||
assert "healthcare sector" not in masked_exception
|
||||
|
||||
# Should contain masking indicator
|
||||
assert "*" in masked_exception
|
||||
|
||||
def test_error_handling_in_common_add_cooldown_logic(self, cooldown_cache):
|
||||
"""Test error handling in the _common_add_cooldown_logic method"""
|
||||
# This test ensures that edge cases are properly handled
|
||||
model_id = "test-model"
|
||||
|
||||
# Test with None exception (edge case) - should be handled gracefully
|
||||
cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
|
||||
model_id=model_id,
|
||||
original_exception=None,
|
||||
exception_status=500,
|
||||
cooldown_time=30.0,
|
||||
)
|
||||
|
||||
# Should handle None by converting to string
|
||||
assert cooldown_data["exception_received"] == "None"
|
||||
assert cooldown_key == f"deployment:{model_id}:cooldown"
|
||||
|
||||
def test_custom_masker_settings(self):
|
||||
"""Test that custom masker settings work correctly"""
|
||||
mock_dual_cache = MagicMock(spec=DualCache)
|
||||
|
||||
# Create cooldown cache and verify default settings
|
||||
cache = CooldownCache(cache=mock_dual_cache, default_cooldown_time=60.0)
|
||||
|
||||
# Test that we can access and verify the masker configuration
|
||||
assert cache.exception_masker.visible_prefix == 50
|
||||
assert cache.exception_masker.visible_suffix == 0
|
||||
assert cache.exception_masker.mask_char == "*"
|
||||
|
||||
# Test masking behavior with these settings
|
||||
long_string = "A" * 100 # 100 character string
|
||||
masked = cache.exception_masker.mask_value(long_string)
|
||||
|
||||
# Should show first 50 characters, then all asterisks
|
||||
expected = "A" * 50 + "*" * 50
|
||||
assert masked == expected
|
||||
|
||||
|
||||
class TestCooldownCacheTTLCorrection:
|
||||
def _make_cooldown_cache(self) -> CooldownCache:
|
||||
in_memory = InMemoryCache()
|
||||
dual_cache = DualCache(in_memory_cache=in_memory)
|
||||
return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)
|
||||
|
||||
def test_expired_entry_evicted_and_not_returned(self):
|
||||
"""
|
||||
An entry with timestamp+cooldown_time in the past must be evicted from
|
||||
in-memory cache and excluded from the active cooldown list.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "expired-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
expired_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time() - 120.0,
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
|
||||
|
||||
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert active == [], "Expired cooldown entry must not appear in active cooldowns"
|
||||
assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
|
||||
|
||||
def test_active_entry_is_returned(self):
|
||||
"""
|
||||
An entry whose cooldown window has not elapsed must appear in the active list.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "active-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
active_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time(),
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, active_value, ttl=60)
|
||||
|
||||
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert len(active) == 1
|
||||
assert active[0][0] == model_id
|
||||
|
||||
def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self):
|
||||
"""
|
||||
When DualCache backfills from Redis using the default 600s TTL, the in-memory
|
||||
TTL must be corrected to min(remaining, 60) seconds.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "backfilled-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
remaining = 30.0
|
||||
value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time() - (60.0 - remaining),
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, value, ttl=600)
|
||||
|
||||
before_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
assert before_expiry is not None
|
||||
|
||||
cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
after_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
assert after_expiry is not None
|
||||
corrected_remaining = after_expiry - time.time()
|
||||
assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
|
||||
assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_expired_entry_evicted(self):
|
||||
"""
|
||||
Async path must also evict expired entries.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "async-expired"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
expired_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time() - 120.0,
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
|
||||
|
||||
active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert active == [], "Expired entry must not appear in async active cooldowns"
|
||||
assert cc.in_memory_cache.get_cache(key) is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_active_entry_is_returned(self):
|
||||
"""
|
||||
Async counterpart of test_active_entry_is_returned: an entry whose cooldown
|
||||
window has not elapsed must appear in the async active list too.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "async-active-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
active_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time(),
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, active_value, ttl=60)
|
||||
|
||||
active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert len(active) == 1
|
||||
assert active[0][0] == model_id
|
||||
|
||||
|
||||
class TestCorrectedActiveCooldown:
|
||||
def _make_cooldown_cache(self) -> CooldownCache:
|
||||
in_memory = InMemoryCache()
|
||||
dual_cache = DualCache(in_memory_cache=in_memory)
|
||||
return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)
|
||||
|
||||
def _entry(self, timestamp: float, cooldown_time: float) -> CooldownCacheValue:
|
||||
return CooldownCacheValue(
|
||||
exception_received="Rate limit",
|
||||
status_code="429",
|
||||
timestamp=timestamp,
|
||||
cooldown_time=cooldown_time,
|
||||
)
|
||||
|
||||
def test_expired_entry_returns_none_and_evicts(self):
|
||||
cc = self._make_cooldown_cache()
|
||||
key = "deployment:expired-dep:cooldown"
|
||||
entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0)
|
||||
cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)
|
||||
|
||||
result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
|
||||
|
||||
assert result is None
|
||||
assert cc.in_memory_cache.get_cache(key) is None
|
||||
|
||||
def test_active_entry_within_window_returns_value(self):
|
||||
cc = self._make_cooldown_cache()
|
||||
key = "deployment:active-dep:cooldown"
|
||||
entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
|
||||
cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
|
||||
|
||||
result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
|
||||
|
||||
assert result is not None
|
||||
assert result["status_code"] == "429"
|
||||
|
||||
def test_inflated_ttl_is_corrected(self):
|
||||
cc = self._make_cooldown_cache()
|
||||
key = "deployment:backfilled-dep:cooldown"
|
||||
remaining = 30.0
|
||||
entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0)
|
||||
cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)
|
||||
|
||||
result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
|
||||
|
||||
assert result is not None
|
||||
corrected_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
assert corrected_expiry is not None
|
||||
assert corrected_expiry - time.time() <= 60.0
|
||||
|
||||
def test_normal_ttl_not_modified(self):
|
||||
cc = self._make_cooldown_cache()
|
||||
key = "deployment:normal-dep:cooldown"
|
||||
entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
|
||||
cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
|
||||
original_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
|
||||
cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())
|
||||
|
||||
after_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
assert after_expiry == original_expiry
|
||||
|
||||
|
||||
class SharedRedisDouble:
|
||||
"""
|
||||
In-process stand-in for RedisCache, shared by several DualCache instances so that
|
||||
tests can model two proxy replicas talking to one Redis.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.store: dict = {} # mutable-ok: stands in for Redis' own mutable keyspace
|
||||
|
||||
def set_cache(self, key, value, **kwargs):
|
||||
self.store[key] = value
|
||||
|
||||
async def async_set_cache(self, key, value, **kwargs):
|
||||
self.store[key] = value
|
||||
|
||||
def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
|
||||
return {key: self.store.get(key) for key in key_list}
|
||||
|
||||
async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
|
||||
return {key: self.store.get(key) for key in key_list}
|
||||
|
||||
|
||||
class TestCooldownPropagationBetweenReplicas:
|
||||
"""
|
||||
A cooldown written by one replica has to reach its siblings quickly. The router's own
|
||||
DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads
|
||||
get their own cache with a much shorter Redis read interval.
|
||||
"""
|
||||
|
||||
def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache:
|
||||
router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis)
|
||||
if read_interval is None:
|
||||
return CooldownCache(cache=router_cache, default_cooldown_time=60.0)
|
||||
return CooldownCache(
|
||||
cache=router_cache,
|
||||
default_cooldown_time=60.0,
|
||||
redis_read_interval_seconds=read_interval,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self):
|
||||
redis = SharedRedisDouble()
|
||||
replica_a = self._make_replica(redis, read_interval=0.25)
|
||||
replica_b = self._make_replica(redis, read_interval=0.25)
|
||||
model_id = "shared-deployment"
|
||||
|
||||
assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []
|
||||
|
||||
replica_a.add_deployment_to_cooldown(
|
||||
model_id=model_id,
|
||||
original_exception=Exception("Internal server error"),
|
||||
exception_status=500,
|
||||
cooldown_time=60.0,
|
||||
)
|
||||
|
||||
time.sleep(0.3)
|
||||
|
||||
active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
|
||||
assert [model_id] == [entry[0] for entry in active], (
|
||||
"sibling replica must pick up a cooldown written by another replica within the read interval"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sibling_replica_sees_cooldown_within_default_read_interval(self):
|
||||
redis = SharedRedisDouble()
|
||||
replica_a = self._make_replica(redis)
|
||||
replica_b = self._make_replica(redis)
|
||||
model_id = "default-interval-deployment"
|
||||
|
||||
assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []
|
||||
|
||||
replica_a.add_deployment_to_cooldown(
|
||||
model_id=model_id,
|
||||
original_exception=Exception("Internal server error"),
|
||||
exception_status=500,
|
||||
cooldown_time=60.0,
|
||||
)
|
||||
|
||||
time.sleep(1.2)
|
||||
|
||||
active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
|
||||
assert [model_id] == [entry[0] for entry in active], (
|
||||
"the shipped default read interval must let a sibling replica see a cooldown about a second later"
|
||||
)
|
||||
|
||||
def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self):
|
||||
redis = SharedRedisDouble()
|
||||
replica_a = self._make_replica(redis, read_interval=0.25)
|
||||
replica_b = self._make_replica(redis, read_interval=0.25)
|
||||
model_id = "sync-shared-deployment"
|
||||
|
||||
assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == []
|
||||
|
||||
replica_a.add_deployment_to_cooldown(
|
||||
model_id=model_id,
|
||||
original_exception=Exception("Internal server error"),
|
||||
exception_status=500,
|
||||
cooldown_time=60.0,
|
||||
)
|
||||
|
||||
time.sleep(0.3)
|
||||
|
||||
active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None)
|
||||
assert [model_id] == [entry[0] for entry in active]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_attached_after_construction_is_still_used(self):
|
||||
redis = SharedRedisDouble()
|
||||
router_cache = DualCache(in_memory_cache=InMemoryCache())
|
||||
writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25)
|
||||
router_cache.attach_redis_cache(redis)
|
||||
reader = self._make_replica(redis, read_interval=0.25)
|
||||
model_id = "late-redis-deployment"
|
||||
|
||||
writer.add_deployment_to_cooldown(
|
||||
model_id=model_id,
|
||||
original_exception=Exception("Internal server error"),
|
||||
exception_status=500,
|
||||
cooldown_time=60.0,
|
||||
)
|
||||
|
||||
active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None)
|
||||
assert [model_id] == [entry[0] for entry in active], (
|
||||
"a router that wires Redis after building its cooldown cache must still publish cooldowns to it"
|
||||
)
|
||||
|
||||
|
||||
class TestCooldownSurvivesUnrelatedCacheTraffic:
|
||||
@pytest.mark.asyncio
|
||||
async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self):
|
||||
router_cache = DualCache(in_memory_cache=InMemoryCache())
|
||||
cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0)
|
||||
model_id = "busy-router-deployment"
|
||||
|
||||
cc.add_deployment_to_cooldown(
|
||||
model_id=model_id,
|
||||
original_exception=Exception("Internal server error"),
|
||||
exception_status=500,
|
||||
cooldown_time=30.0,
|
||||
)
|
||||
|
||||
for i in range(400):
|
||||
router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i})
|
||||
|
||||
active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None)
|
||||
assert [model_id] == [entry[0] for entry in active], (
|
||||
"unrelated router cache traffic must not evict a cooldown that is still running"
|
||||
)
|
||||
|
||||
|
||||
class TestCooldownStoreCallsDeclareTheirKeyFamily:
|
||||
"""Every cooldown store call runs inside ``service_target("router_cooldowns")`` so the
|
||||
Redis service spans read ``redis.set router_cooldowns`` / ``redis.mget router_cooldowns``,
|
||||
the sync paths included (the async MGET already did)."""
|
||||
|
||||
def _cooldown_cache_with_recording_store(self, seen: list[tuple[str, str | None]]) -> CooldownCache:
|
||||
cc = CooldownCache(cache=DualCache(in_memory_cache=InMemoryCache()), default_cooldown_time=60.0)
|
||||
store = MagicMock()
|
||||
|
||||
def _set_cache(**_kwargs):
|
||||
seen.append(("set", current_service_target()))
|
||||
|
||||
def _batch_get_cache(**_kwargs):
|
||||
seen.append(("mget", current_service_target()))
|
||||
return []
|
||||
|
||||
store.set_cache.side_effect = _set_cache
|
||||
store.batch_get_cache.side_effect = _batch_get_cache
|
||||
cc._cooldown_store = store
|
||||
return cc
|
||||
|
||||
def test_sync_cooldown_write_runs_under_router_cooldowns(self):
|
||||
seen: list[tuple[str, str | None]] = []
|
||||
cc = self._cooldown_cache_with_recording_store(seen)
|
||||
|
||||
cc.add_deployment_to_cooldown(
|
||||
model_id="dep-1",
|
||||
original_exception=Exception("Internal server error"),
|
||||
exception_status=500,
|
||||
cooldown_time=30.0,
|
||||
)
|
||||
|
||||
assert seen == [("set", "router_cooldowns")]
|
||||
assert current_service_target() is None
|
||||
|
||||
def test_sync_cooldown_reads_run_under_router_cooldowns(self):
|
||||
seen: list[tuple[str, str | None]] = []
|
||||
cc = self._cooldown_cache_with_recording_store(seen)
|
||||
|
||||
assert cc.get_active_cooldowns(["dep-1"], parent_otel_span=None) == []
|
||||
assert cc.get_min_cooldown(["dep-1"], parent_otel_span=None) == 60.0
|
||||
|
||||
assert seen == [("mget", "router_cooldowns"), ("mget", "router_cooldowns")]
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_and_async_router_cooldown_surfaces_share_active_state() -> None:
|
||||
now: Final = time.time()
|
||||
store: Final = RouterCooldownStore(None)
|
||||
cache: Final = CooldownCache(store, default_cooldown_time=60.0, clock=lambda: now)
|
||||
cache.add_deployment_to_cooldown("dep", Exception("unavailable"), 503, 30.0)
|
||||
expected: Final = [
|
||||
("dep", {"exception_received": "unavailable", "status_code": "503", "timestamp": now, "cooldown_time": 30.0})
|
||||
]
|
||||
assert cache.get_active_cooldowns(["absent", "dep"], None) == expected
|
||||
assert await cache.async_get_active_cooldowns(["absent", "dep"], None) == expected
|
||||
assert cache.get_min_cooldown(["dep"], None) == 30.0
|
||||
assert cache.get_min_cooldown(["absent"], None) == 60.0
|
||||
|
|
|
|||
|
|
@ -16,10 +16,11 @@ import respx
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm._internal_context import current_service_target
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.litellm_core_utils.litellm_logging import mask_api_base_credentials
|
||||
from litellm.redis.stores.router_cooldowns import RouterCooldowns, RouterCooldownStore
|
||||
from litellm.router_api import RoutingContext
|
||||
from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue
|
||||
from litellm.router_utils.cooldown_callbacks import router_cooldown_event_callback
|
||||
|
|
@ -27,7 +28,6 @@ from litellm.router_utils.cooldown_handlers import (
|
|||
_get_deployment_cooldown_policy,
|
||||
_has_explicit_allowed_fails_policy_for_exception,
|
||||
_has_fallback_group,
|
||||
_increment_allowed_fails,
|
||||
_is_cooldown_required,
|
||||
_resolve_allowed_fails_from_policy,
|
||||
_should_cooldown_based_on_deployment_policy,
|
||||
|
|
@ -74,6 +74,7 @@ def _ctx_double(
|
|||
ctx.settings.return_value.allowed_fails_policy = None
|
||||
ctx.settings.return_value.allowed_fails_for.return_value = None
|
||||
ctx.health.return_value.store = store if store is not None else DualCache(in_memory_cache=InMemoryCache())
|
||||
ctx.health.return_value.cooldowns = CooldownCache(RouterCooldownStore(None), cooldown_time)
|
||||
return ctx
|
||||
|
||||
|
||||
|
|
@ -360,85 +361,42 @@ class TestShouldCooldownBasedOnDeploymentPolicy:
|
|||
|
||||
|
||||
class TestShouldCooldownBasedOnAllowedFailsPolicy:
|
||||
def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | MagicMock | None = None) -> MagicMock:
|
||||
return _ctx_double(cooldown_time=cooldown_time, allowed_fails=0, store=cache)
|
||||
def _make_router(self, cooldown_time: float = 60.0, store: RouterCooldowns | None = None) -> MagicMock:
|
||||
ctx: Final = _ctx_double(cooldown_time=cooldown_time, allowed_fails=0)
|
||||
ctx.health.return_value.cooldowns = CooldownCache(store or RouterCooldownStore(None), cooldown_time)
|
||||
return ctx
|
||||
|
||||
def test_cooldown_time_override_zero_is_not_falsy(self):
|
||||
"""cooldown_time_override=0 must be honored; it must not fall through to the router-level value."""
|
||||
store = MagicMock()
|
||||
store.increment_cache.return_value = 1
|
||||
ctx = self._make_router(cooldown_time=60.0, cache=store)
|
||||
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
||||
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=ctx,
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
cooldown_time_override=0.0,
|
||||
)
|
||||
|
||||
increment_call = store.increment_cache.call_args
|
||||
assert increment_call is not None
|
||||
assert increment_call[1]["ttl"] == 0.0, (
|
||||
"cooldown_time_override=0 should be used as TTL, not the router-level 60.0"
|
||||
store: Final = MagicMock(spec=RouterCooldowns)
|
||||
store.increment_allowed_fails.return_value = 1
|
||||
ctx: Final = self._make_router(cooldown_time=60.0, store=store)
|
||||
exc: Final = litellm.RateLimitError("429", "openai", "gpt-4")
|
||||
assert not should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5, cooldown_time_override=0.0
|
||||
)
|
||||
store.increment_allowed_fails.assert_called_once_with("dep-1", None, 0.0)
|
||||
|
||||
def test_fail_counter_is_shared_across_router_instances(self):
|
||||
"""Two workers (two Router objects over one shared cache) must pool their failures toward allowed_fails."""
|
||||
shared_cache = DualCache(in_memory_cache=InMemoryCache())
|
||||
workers = (self._make_router(cache=shared_cache), self._make_router(cache=shared_cache))
|
||||
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
|
||||
results = [
|
||||
shared: Final = RouterCooldownStore(None)
|
||||
workers: Final = tuple(self._make_router(store=shared) for _ in range(2))
|
||||
exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
results: Final = [
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=workers[i % 2],
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
ctx=workers[i % 2], deployment="dep-1", original_exception=exc, allowed_fails_override=5
|
||||
)
|
||||
for i in range(6)
|
||||
]
|
||||
|
||||
assert results == [False, False, False, False, False, True]
|
||||
assert shared_cache.get_cache(key="deployment:dep-1:allowed_fails") == 6
|
||||
|
||||
def test_fleet_wide_count_from_redis_decides_cooldown(self):
|
||||
"""The Redis (fleet-wide) count decides, even when this process has only seen one failure."""
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.increment_cache.return_value = 6
|
||||
router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
|
||||
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=router,
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
def test_fleet_wide_count_from_the_store_decides_cooldown(self):
|
||||
store: Final = MagicMock(spec=RouterCooldowns)
|
||||
store.increment_allowed_fails.return_value = 6
|
||||
ctx: Final = self._make_router(store=store)
|
||||
exc: Final = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
assert should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=ctx, deployment="dep-1", original_exception=exc, allowed_fails_override=5
|
||||
)
|
||||
|
||||
assert result is True
|
||||
redis_cache.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0)
|
||||
|
||||
def test_redis_outage_falls_back_to_this_workers_count(self):
|
||||
"""When every Redis increment fails, the worker's own in-memory count must still cool the deployment down."""
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.increment_cache.side_effect = ConnectionError("redis down")
|
||||
router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
|
||||
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
|
||||
|
||||
results = [
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=router,
|
||||
deployment="dep-1",
|
||||
original_exception=exc,
|
||||
allowed_fails_override=5,
|
||||
)
|
||||
for _ in range(6)
|
||||
]
|
||||
|
||||
assert results == [False, False, False, False, False, True]
|
||||
assert redis_cache.increment_cache.call_count == 6
|
||||
store.increment_allowed_fails.assert_called_once_with("dep-1", None, 60.0)
|
||||
|
||||
|
||||
class TestRoutingGroupCooldownAlternatives:
|
||||
|
|
@ -586,38 +544,6 @@ class TestTeamModelCooldownAlternatives:
|
|||
)
|
||||
|
||||
|
||||
class TestIncrementAllowedFailsServiceTarget:
|
||||
def test_fail_counter_bump_declares_the_router_cooldowns_key_family(self):
|
||||
"""The allowed_fails INCR is cooldown bookkeeping, so its service span must read
|
||||
``redis.incr router_cooldowns`` rather than a bare ``redis.incr``."""
|
||||
seen: list[str | None] = []
|
||||
cache = MagicMock(spec=DualCache)
|
||||
|
||||
def _increment(**_kwargs):
|
||||
seen.append(current_service_target())
|
||||
return 2
|
||||
|
||||
cache.increment_cache.side_effect = _increment
|
||||
|
||||
assert _increment_allowed_fails(cache, "deployment:dep-1:fails", ttl=60.0) == 2
|
||||
assert seen == ["router_cooldowns"]
|
||||
assert current_service_target() is None
|
||||
|
||||
def test_in_memory_fallback_reads_under_the_same_target(self):
|
||||
seen: list[str | None] = []
|
||||
cache = MagicMock(spec=DualCache)
|
||||
cache.increment_cache.side_effect = ConnectionError("redis down")
|
||||
|
||||
def _get(**_kwargs):
|
||||
seen.append(current_service_target())
|
||||
return 4
|
||||
|
||||
cache.get_cache.side_effect = _get
|
||||
|
||||
assert _increment_allowed_fails(cache, "deployment:dep-1:fails", ttl=60.0) == 4
|
||||
assert seen == ["router_cooldowns"]
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def _vcr_outcome_gate(request, vcr):
|
||||
install_live_call_probe(request, vcr)
|
||||
|
|
@ -860,136 +786,57 @@ class TestExceptionTypeCountersTrackedIndependently:
|
|||
rate_limit_exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4")
|
||||
ise = litellm.InternalServerError("Internal error", "openai", "gpt-4")
|
||||
|
||||
for _ in range(3):
|
||||
assert [
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=RoutingContext(router),
|
||||
deployment="primary",
|
||||
original_exception=rate_limit_exc,
|
||||
allowed_fails_override=5,
|
||||
allowed_fails_override=3,
|
||||
cache_key_suffix="RateLimitError",
|
||||
)
|
||||
|
||||
rl_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
|
||||
generic_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
|
||||
|
||||
assert rl_counter == 3, "RateLimitError counter should be 3"
|
||||
assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments"
|
||||
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
for _ in range(3)
|
||||
] == [False, False, False]
|
||||
assert not should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=RoutingContext(router),
|
||||
deployment="primary",
|
||||
original_exception=ise,
|
||||
allowed_fails_override=5,
|
||||
allowed_fails_override=1,
|
||||
cache_key_suffix="generic",
|
||||
)
|
||||
assert not should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=RoutingContext(router),
|
||||
deployment="primary",
|
||||
original_exception=rate_limit_exc,
|
||||
allowed_fails_override=4,
|
||||
cache_key_suffix="RateLimitError",
|
||||
)
|
||||
assert should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=RoutingContext(router),
|
||||
deployment="primary",
|
||||
original_exception=ise,
|
||||
allowed_fails_override=1,
|
||||
cache_key_suffix="generic",
|
||||
)
|
||||
|
||||
generic_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
|
||||
rl_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
|
||||
|
||||
assert generic_counter_after == 1, "generic counter should now be 1"
|
||||
assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
||||
class TestCooldownCacheTTLCorrection:
|
||||
def _make_cooldown_cache(self) -> CooldownCache:
|
||||
in_memory = InMemoryCache()
|
||||
dual_cache = DualCache(in_memory_cache=in_memory)
|
||||
return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)
|
||||
|
||||
def test_expired_entry_evicted_and_not_returned(self):
|
||||
"""
|
||||
An entry with timestamp+cooldown_time in the past must be evicted from
|
||||
in-memory cache and excluded from the active cooldown list.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "expired-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
expired_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time() - 120.0,
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
|
||||
|
||||
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert active == [], "Expired cooldown entry must not appear in active cooldowns"
|
||||
assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"
|
||||
|
||||
def test_active_entry_is_returned(self):
|
||||
"""
|
||||
An entry whose cooldown window has not elapsed must appear in the active list.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "active-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
active_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time(),
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, active_value, ttl=60)
|
||||
|
||||
active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert len(active) == 1
|
||||
assert active[0][0] == model_id
|
||||
|
||||
def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self):
|
||||
"""
|
||||
When DualCache backfills from Redis using the default 600s TTL, the in-memory
|
||||
TTL must be corrected to min(remaining, 60) seconds.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "backfilled-deployment"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
remaining = 30.0
|
||||
value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time() - (60.0 - remaining),
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, value, ttl=600)
|
||||
|
||||
before_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
assert before_expiry is not None
|
||||
|
||||
cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
after_expiry = cc.in_memory_cache.ttl_dict.get(key)
|
||||
assert after_expiry is not None
|
||||
corrected_remaining = after_expiry - time.time()
|
||||
assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
|
||||
assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)"
|
||||
|
||||
@pytest.mark.parametrize(("age", "active"), ((120.0, False), (0.0, True)))
|
||||
@pytest.mark.parametrize("asynchronous", (False, True))
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_expired_entry_evicted(self):
|
||||
"""
|
||||
Async path must also evict expired entries.
|
||||
"""
|
||||
cc = self._make_cooldown_cache()
|
||||
model_id = "async-expired"
|
||||
key = CooldownCache.get_cooldown_cache_key(model_id)
|
||||
|
||||
expired_value: CooldownCacheValue = {
|
||||
"exception_received": "Rate limit",
|
||||
"status_code": "429",
|
||||
"timestamp": time.time() - 120.0,
|
||||
"cooldown_time": 60.0,
|
||||
}
|
||||
cc.in_memory_cache.set_cache(key, expired_value, ttl=600)
|
||||
|
||||
active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)
|
||||
|
||||
assert active == [], "Expired entry must not appear in async active cooldowns"
|
||||
assert cc.in_memory_cache.get_cache(key) is None
|
||||
async def test_only_unexpired_cooldowns_are_returned(self, age: float, active: bool, asynchronous: bool):
|
||||
store: Final = RouterCooldownStore(None)
|
||||
cc: Final = CooldownCache(store, default_cooldown_time=60.0)
|
||||
value: Final = CooldownCacheValue(
|
||||
exception_received="Rate limit", status_code="429", timestamp=time.time() - age, cooldown_time=60.0
|
||||
)
|
||||
store.put("deployment", value, 600)
|
||||
result: Final = (
|
||||
await cc.async_get_active_cooldowns(("deployment",), None)
|
||||
if asynchronous
|
||||
else cc.get_active_cooldowns(("deployment",), None)
|
||||
)
|
||||
assert result == ([("deployment", value)] if active else [])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
||||
|
|
@ -2202,7 +2049,6 @@ def test_increment_deployment_successes_for_current_minute_does_not_write_to_red
|
|||
"""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
|
||||
increment_deployment_successes_for_current_minute,
|
||||
)
|
||||
|
|
@ -2700,7 +2546,6 @@ async def test_high_traffic_cooldowns_all_healthy_deployments():
|
|||
await asyncio.sleep(0.0001)
|
||||
except litellm.InternalServerError:
|
||||
model_stats[model_id]["failures"] += 1
|
||||
pass
|
||||
except Exception as e:
|
||||
print("Failed test model stats=", model_stats)
|
||||
raise e
|
||||
|
|
@ -2794,7 +2639,6 @@ async def test_high_traffic_cooldowns_one_bad_deployment():
|
|||
await asyncio.sleep(0.0001)
|
||||
except litellm.InternalServerError:
|
||||
model_stats[model_id]["failures"] += 1
|
||||
pass
|
||||
except Exception as e:
|
||||
print("Failed test model stats=", model_stats)
|
||||
raise e
|
||||
|
|
@ -2888,10 +2732,8 @@ async def test_high_traffic_cooldowns_one_rate_limited_deployment():
|
|||
await asyncio.sleep(0.0001)
|
||||
except litellm.InternalServerError:
|
||||
model_stats[model_id]["failures"] += 1
|
||||
pass
|
||||
except litellm.RateLimitError:
|
||||
model_stats[bad_deployment_id]["failures"] += 1
|
||||
pass
|
||||
except Exception as e:
|
||||
print("Failed test model stats=", model_stats)
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ should increment the same counters used by allowed_fails_policy, using the
|
|||
actual exception type from the health check error.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
from unittest.mock import ANY, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -146,30 +147,26 @@ class TestGetAllowedFailsFromPolicyWithHealthCheckExceptions:
|
|||
class TestHealthCheckCooldownIntegration:
|
||||
"""Test that health check failures trigger cooldown via set_cooldown_deployments."""
|
||||
|
||||
def test_health_check_failure_increments_failed_calls(self):
|
||||
"""Health check failure should increment the failed_calls counter."""
|
||||
def test_health_check_failure_increments_failed_calls(self) -> None:
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
router: Final = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=3),
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(message="Health check timeout", model="gpt-4", llm_provider="openai")
|
||||
|
||||
# First call: should not cooldown (1 <= 3)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=RoutingContext(router),
|
||||
deployment="deploy-1",
|
||||
original_exception=timeout_exc,
|
||||
timeout_exc: Final = litellm.Timeout(message="Health check timeout", model="gpt-4", llm_provider="openai")
|
||||
results: Final = tuple(
|
||||
should_cooldown_based_on_allowed_fails_policy(
|
||||
ctx=RoutingContext(router),
|
||||
deployment="deploy-1",
|
||||
original_exception=timeout_exc,
|
||||
)
|
||||
for _ in range(4)
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Check counter was incremented
|
||||
current_fails = router.health().store.get_cache(key="deployment:deploy-1:allowed_fails")
|
||||
assert current_fails == 1
|
||||
assert results == (False, False, False, True)
|
||||
|
||||
def test_health_check_failure_triggers_cooldown_at_threshold(self):
|
||||
"""After exceeding allowed_fails threshold, deployment should enter cooldown."""
|
||||
|
|
|
|||
|
|
@ -1,11 +1,10 @@
|
|||
"""
|
||||
One Redis round trip per request for the router's pre-call reads.
|
||||
|
||||
Before `RoutingReadBatch`, `async_get_available_deployment` issued one MGET for the cooldown keys
|
||||
(`CooldownCache`) and a second one for the tpm/rpm counters (`LowestTPMLoggingHandler_v2`).
|
||||
The router's pre-call reads. Cooldowns go through the native store, and the tpm/rpm selector
|
||||
counters (`LowestTPMLoggingHandler_v2`) ride one MGET on the Python Redis.
|
||||
"""
|
||||
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
|
@ -13,8 +12,17 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.caching.redis_batch import request_redis_batch_scope
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.redis.stores._connection import native_client
|
||||
from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore
|
||||
from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2
|
||||
from litellm.router_utils import routing_prefetch
|
||||
from litellm.rust_bridge._native import NativeRedisClient
|
||||
from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNativeRequest, FakeNatives
|
||||
from tests.unit.conftest import redis_cache_double
|
||||
from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns
|
||||
from tests.unit.redis.stores.test_router_cooldowns import Handles as CooldownHandles
|
||||
|
||||
_MODEL_GROUP = "claude"
|
||||
_MESSAGES = [{"role": "user", "content": "ping"}]
|
||||
|
|
@ -34,7 +42,7 @@ def _redis_answering(values_by_key_prefix: dict[str, object]) -> MagicMock:
|
|||
def _mget(key_list, parent_otel_span=None):
|
||||
return {key: values_by_key_prefix.get(key.rsplit(":", 1)[0], values_by_key_prefix.get(key)) for key in key_list}
|
||||
|
||||
redis = MagicMock(spec=RedisCache)
|
||||
redis = redis_cache_double()
|
||||
redis.async_batch_get_cache = AsyncMock(side_effect=_mget)
|
||||
return redis
|
||||
|
||||
|
|
@ -55,18 +63,18 @@ def _redis_key_families(redis: MagicMock) -> list[list[str]]:
|
|||
]
|
||||
|
||||
|
||||
def _cooldown(seconds: float) -> dict:
|
||||
return {"exception_received": "429", "status_code": "429", "timestamp": time.time(), "cooldown_time": seconds}
|
||||
def _cool_down(router: Router, model_id: str, started: float) -> None:
|
||||
router.cooldown_cache.cooldown_store.put(
|
||||
model_id, {"exception_received": "429", "status_code": "429", "timestamp": started, "cooldown_time": 60}, 60
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_round_trip():
|
||||
async def test_usage_based_routing_reads_the_counters_in_one_python_redis_round_trip():
|
||||
redis = _redis_answering({})
|
||||
router = _router(redis, "usage-based-routing-v2")
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
assert _redis_key_families(redis) == [
|
||||
|
|
@ -75,10 +83,8 @@ async def test_usage_based_routing_reads_cooldowns_and_counters_in_one_redis_rou
|
|||
"dep-a:anthropic/claude-x:tpm",
|
||||
"dep-b:anthropic/claude-x:rpm",
|
||||
"dep-b:anthropic/claude-x:tpm",
|
||||
"deployment:dep-a:cooldown",
|
||||
"deployment:dep-b:cooldown",
|
||||
]
|
||||
], "cooldown state and usage counters must arrive in one MGET"
|
||||
], "the usage counters must arrive in one MGET, without cooldown keys"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -114,21 +120,19 @@ async def test_usage_based_routing_still_batches_when_the_strategy_is_a_fixed_si
|
|||
"dep-a:anthropic/claude-x:tpm",
|
||||
"dep-b:anthropic/claude-x:rpm",
|
||||
"dep-b:anthropic/claude-x:tpm",
|
||||
"deployment:dep-a:cooldown",
|
||||
"deployment:dep-b:cooldown",
|
||||
]
|
||||
], "the subclassed strategy must still get the batched read, not a second MGET"
|
||||
router.cache.async_batch_get_cache.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_shuffle_still_reads_only_cooldowns():
|
||||
async def test_simple_shuffle_reads_nothing_from_the_python_redis():
|
||||
redis = _redis_answering({})
|
||||
router = _router(redis, "simple-shuffle")
|
||||
|
||||
await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
|
||||
|
||||
assert _redis_key_families(redis) == [["deployment:dep-a:cooldown", "deployment:dep-b:cooldown"]]
|
||||
assert _redis_key_families(redis) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -153,18 +157,11 @@ async def test_batched_counters_pick_the_deployment_the_strategy_picks_reading_a
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batched_read_still_excludes_a_cooled_down_deployment():
|
||||
redis = _redis_answering(
|
||||
{
|
||||
"dep-a:anthropic/claude-x:tpm": 100,
|
||||
"dep-b:anthropic/claude-x:tpm": 10,
|
||||
"deployment:dep-b:cooldown": _cooldown(seconds=60),
|
||||
}
|
||||
)
|
||||
redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10})
|
||||
router = _router(redis, "usage-based-routing-v2")
|
||||
_cool_down(router, "dep-b", started=time.time())
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
|
||||
|
||||
assert deployment["model_info"]["id"] == "dep-a", "dep-b has the lowest tpm but is cooling down"
|
||||
assert redis.async_batch_get_cache.await_count == 1
|
||||
|
|
@ -172,25 +169,18 @@ async def test_batched_read_still_excludes_a_cooled_down_deployment():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batched_read_ignores_an_expired_cooldown():
|
||||
redis = _redis_answering(
|
||||
{
|
||||
"dep-a:anthropic/claude-x:tpm": 100,
|
||||
"dep-b:anthropic/claude-x:tpm": 10,
|
||||
"deployment:dep-b:cooldown": _cooldown(seconds=-1),
|
||||
}
|
||||
)
|
||||
redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10})
|
||||
router = _router(redis, "usage-based-routing-v2")
|
||||
_cool_down(router, "dep-b", started=time.time() - 61)
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
|
||||
|
||||
assert deployment["model_info"]["id"] == "dep-b"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did():
|
||||
redis = MagicMock(spec=RedisCache)
|
||||
redis = redis_cache_double()
|
||||
redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
|
||||
routed = _router(redis, "usage-based-routing-v2")
|
||||
alone = _router(redis, "usage-based-routing-v2")
|
||||
|
|
@ -204,17 +194,100 @@ async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did():
|
|||
|
||||
assert str(routed_error.value) == str(alone_error.value)
|
||||
assert len(routed.cache.last_redis_batch_access_time) == 0, "a failed read must not throttle the next one"
|
||||
assert len(routed.cooldown_cache.cooldown_store.last_redis_batch_access_time) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_batched_read_leaves_simple_shuffle_routing():
|
||||
redis = MagicMock(spec=RedisCache)
|
||||
redis = redis_cache_double()
|
||||
redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
|
||||
router = _router(redis, "simple-shuffle")
|
||||
|
||||
deployment = await router.async_get_available_deployment(
|
||||
model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES
|
||||
)
|
||||
deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES)
|
||||
|
||||
assert deployment["model_info"]["id"] in {"dep-a", "dep-b"}
|
||||
|
||||
|
||||
def _armed_store() -> tuple[RouterCooldownStore, CooldownHandles]:
|
||||
handles: Final = CooldownHandles()
|
||||
store: Final = RouterCooldownStore(native_client(RedisCache(host="redis.internal")), handle=handles)
|
||||
store.put("dep-a", {"exception_received": "429", "status_code": "429", "timestamp": 1.0, "cooldown_time": 60}, 60)
|
||||
return store, handles
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_routing_takes_the_armed_cooldown_read_from_the_native_batch_once_then_reads_on_its_own():
|
||||
store, handles = _armed_store()
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-a", "dep-b", "dep-c"])
|
||||
armed: Final = routing_prefetch.armed()
|
||||
first: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None)
|
||||
reads_after_first: Final = len(handles.built[0].targets)
|
||||
second: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None)
|
||||
|
||||
assert armed
|
||||
assert handles.built[0].declared == [("dep-a", "dep-b", "dep-c")]
|
||||
assert first == second == (("dep-a", store.snapshot()["dep-a"]),)
|
||||
assert natives.built[0].events == ["flush_pre", "resolve"]
|
||||
assert len(handles.built[0].targets) == reads_after_first + 1, "a second read goes to the store"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_armed_cooldown_read_that_does_not_cover_the_routing_ids_is_read_again():
|
||||
store, _ = _armed_store()
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-b"])
|
||||
active: Final = await routing_prefetch.read_cooldowns(store, ["dep-a", "dep-b"], None)
|
||||
|
||||
assert active == (("dep-a", store.snapshot()["dep-a"]),)
|
||||
assert natives.built[0].events == ["flush_pre", "resolve"]
|
||||
|
||||
|
||||
def _redis_down() -> list[tuple[str, CooldownCacheValue]]:
|
||||
raise ConnectionError("redis down")
|
||||
|
||||
|
||||
class _UnreachableCooldowns(FakeCooldowns):
|
||||
def declare_active(
|
||||
self, request: FakeNativeRequest, model_ids: Sequence[str]
|
||||
) -> FakeNativePending[list[tuple[str, CooldownCacheValue]]]:
|
||||
super().declare_active(request, model_ids)
|
||||
return FakeNativePending(request, _redis_down)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_armed_cooldown_read_falls_back_to_the_stores_own_read():
|
||||
built: Final[list[FakeCooldowns]] = []
|
||||
|
||||
def unreachable(client: NativeRedisClient | None) -> FakeCooldowns:
|
||||
built.append(_UnreachableCooldowns(client))
|
||||
return built[-1]
|
||||
|
||||
store: Final = RouterCooldownStore(native_client(RedisCache(host="redis.internal")), handle=unreachable)
|
||||
store.put("dep-a", {"exception_received": "429", "status_code": "429", "timestamp": 1.0, "cooldown_time": 60}, 60)
|
||||
|
||||
with request_redis_batch_scope(native_factory=FakeNatives()): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-a"])
|
||||
active: Final = await routing_prefetch.read_cooldowns(store, ["dep-a"], None)
|
||||
|
||||
assert built[0].declared == [("dep-a",)]
|
||||
assert active == (("dep-a", store.snapshot()["dep-a"]),)
|
||||
|
||||
|
||||
def test_arming_the_cooldown_read_outside_a_request_or_without_redis_declares_nothing():
|
||||
store, handles = _armed_store()
|
||||
offline: Final = CooldownHandles()
|
||||
natives: Final = FakeNatives()
|
||||
|
||||
routing_prefetch.arm_cooldowns(store, ["dep-a"])
|
||||
with request_redis_batch_scope(native_factory=natives): # pyright: ignore[reportArgumentType] # test double
|
||||
routing_prefetch.arm_cooldowns(RouterCooldownStore(None, handle=offline), ["dep-a"])
|
||||
armed: Final = routing_prefetch.armed()
|
||||
|
||||
assert handles.built[0].declared == []
|
||||
assert offline.built[0].declared == []
|
||||
assert not armed
|
||||
assert natives.built == []
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import httpx
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.integrations.custom_secret_manager import CustomSecretManager
|
||||
from litellm.llms.custom_httpx.http_handler import default_user_agent
|
||||
from litellm.rust_bridge.host import settings
|
||||
|
|
@ -46,6 +47,26 @@ def test_http_settings_reads_the_litellm_globals(monkeypatch: pytest.MonkeyPatch
|
|||
)
|
||||
|
||||
|
||||
def test_router_tiers_read_the_env_constants_and_the_litellm_global(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import litellm.constants
|
||||
|
||||
monkeypatch.setattr(litellm.constants, "DEFAULT_COOLDOWN_REDIS_READ_INTERVAL_SECONDS", 2.5)
|
||||
monkeypatch.setattr(litellm.constants, "DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", 42)
|
||||
monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", 30.0)
|
||||
|
||||
assert settings.router_tiers() == settings.RouterTierSettings(
|
||||
cooldown_redis_read_interval=2.5,
|
||||
usage_redis_read_interval=30.0,
|
||||
max_redis_batch_cache_size=42,
|
||||
)
|
||||
|
||||
|
||||
def test_router_tiers_fall_back_to_dual_caches_ten_second_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "default_redis_batch_cache_expiry", None)
|
||||
|
||||
assert settings.router_tiers().usage_redis_read_interval == DualCache().redis_batch_cache_expiry
|
||||
|
||||
|
||||
def test_http_settings_ignores_environment_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LITELLM_USER_AGENT", "operator/1")
|
||||
monkeypatch.setenv("SSL_VERIFY", "false")
|
||||
|
|
@ -114,9 +135,7 @@ def test_secret_manager_is_native_only_when_a_readable_system_is_configured(
|
|||
monkeypatch.setattr(litellm, "_key_management_system", system)
|
||||
monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode))
|
||||
|
||||
assert settings.secret_manager() == settings.SecretManager(
|
||||
readable=access_mode != "write_only", native=native
|
||||
)
|
||||
assert settings.secret_manager() == settings.SecretManager(readable=access_mode != "write_only", native=native)
|
||||
|
||||
|
||||
def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@ _CACHE_PARAMETER_HELPERS: Final = frozenset(
|
|||
"litellm/llms/anthropic/prompt_cache_prediction.py",
|
||||
"litellm/proxy/common_utils/cache_coordinator.py",
|
||||
"litellm/proxy/common_utils/user_api_key_cache.py",
|
||||
"litellm/router_utils/routing_prefetch.py",
|
||||
"litellm/utils.py",
|
||||
}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -108,6 +108,7 @@ from litellm.utils import _invalidate_model_cost_lowercase_map
|
|||
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
||||
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
|
||||
from tests.large_text import text
|
||||
from tests.unit.conftest import redis_cache_double
|
||||
import traceback
|
||||
import inspect
|
||||
from typing import List, Optional
|
||||
|
|
@ -1413,17 +1414,6 @@ _BATCH_NOT_FOUND = {
|
|||
}
|
||||
|
||||
|
||||
async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]:
|
||||
loop = asyncio.get_event_loop()
|
||||
deadline = loop.time() + timeout
|
||||
while loop.time() < deadline:
|
||||
keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:"))
|
||||
if keys:
|
||||
return keys
|
||||
await asyncio.sleep(0.05)
|
||||
return []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch):
|
||||
"""
|
||||
|
|
@ -1470,11 +1460,13 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m
|
|||
)
|
||||
response = await router.aretrieve_batch(batch_id=_BATCH_ID)
|
||||
payload = await collector.retrieve_batch_payload()
|
||||
usage_keys = await _router_usage_keys(router)
|
||||
requested_usage: Final = await router.get_model_group_usage(_BATCH_GROUP)
|
||||
unrelated_usage: Final = await router.get_model_group_usage(_UNRELATED_BATCH_GROUP)
|
||||
|
||||
assert response.id == _BATCH_ID
|
||||
assert payload["model_group"] == _BATCH_GROUP
|
||||
assert usage_keys == []
|
||||
assert requested_usage == (None, None)
|
||||
assert unrelated_usage == (None, None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
@ -15955,11 +15947,10 @@ class TestClaudeCodeSubagentSessionRouterBinding:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_cleanup_failure_does_not_reject_a_subagent_request(self):
|
||||
from litellm.caching.caching import RedisCache
|
||||
|
||||
router = self._router()
|
||||
del router.complexity_routers["smart-router"]
|
||||
redis_cache = MagicMock(spec=RedisCache)
|
||||
redis_cache = redis_cache_double()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="smart-router")
|
||||
redis_cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis unavailable"))
|
||||
router.update_redis_cache(cache=redis_cache)
|
||||
|
|
@ -15974,7 +15965,6 @@ class TestClaudeCodeSubagentSessionRouterBinding:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_read_failure_does_not_reject_a_subagent_request(self):
|
||||
from litellm.caching.caching import RedisCache
|
||||
|
||||
router = self._router()
|
||||
request_kwargs = self._request_kwargs(agent_id="agent-1234")
|
||||
|
|
@ -15984,7 +15974,7 @@ class TestClaudeCodeSubagentSessionRouterBinding:
|
|||
cache_key,
|
||||
"smart-router",
|
||||
)
|
||||
redis_cache = MagicMock(spec=RedisCache)
|
||||
redis_cache = redis_cache_double()
|
||||
redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open"))
|
||||
router.update_redis_cache(cache=redis_cache)
|
||||
|
||||
|
|
@ -15999,10 +15989,9 @@ class TestClaudeCodeSubagentSessionRouterBinding:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self):
|
||||
from litellm.caching.caching import RedisCache
|
||||
|
||||
router = self._router()
|
||||
redis_cache = MagicMock(spec=RedisCache)
|
||||
redis_cache = redis_cache_double()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="smart-router")
|
||||
redis_cache.async_set_cache = AsyncMock(side_effect=Exception("redis unavailable"))
|
||||
router.update_redis_cache(cache=redis_cache)
|
||||
|
|
@ -16026,10 +16015,9 @@ class TestClaudeCodeSubagentSessionRouterBinding:
|
|||
async def test_subagents_follow_the_main_threads_latest_router_across_workers(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.caching.caching import RedisCache
|
||||
|
||||
shared_binding = SimpleNamespace(value=None)
|
||||
shared_redis = MagicMock(spec=RedisCache)
|
||||
shared_redis = redis_cache_double()
|
||||
shared_redis.async_get_cache = AsyncMock(side_effect=lambda key, **_: shared_binding.value)
|
||||
shared_redis.async_set_cache = AsyncMock(
|
||||
side_effect=lambda key, value, **_: setattr(shared_binding, "value", value)
|
||||
|
|
@ -16057,11 +16045,10 @@ class TestClaudeCodeSubagentSessionRouterBinding:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_pre_routing_strategies_means_no_session_cache_traffic(self):
|
||||
from litellm.caching.caching import RedisCache
|
||||
|
||||
router = self._router()
|
||||
router.complexity_routers.clear()
|
||||
redis_cache = MagicMock(spec=RedisCache)
|
||||
redis_cache = redis_cache_double()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
redis_cache.async_set_cache = AsyncMock()
|
||||
redis_cache.async_delete_cache = AsyncMock()
|
||||
|
|
@ -24649,13 +24636,7 @@ def test_router_keeps_wif_secret_pointers_unresolved(monkeypatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_rpm_increment_declares_the_router_usage_key_family():
|
||||
"""The RPM bump a failed call still earns is router usage bookkeeping, so its Redis span
|
||||
reads ``redis.incr router_usage`` rather than a bare ``redis.incr``."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm._internal_context import current_service_target
|
||||
|
||||
async def test_a_failed_call_still_counts_one_rpm():
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
|
|
@ -24665,42 +24646,31 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family():
|
|||
}
|
||||
]
|
||||
)
|
||||
seen: list[str | None] = []
|
||||
|
||||
async def _increment(**_kwargs):
|
||||
seen.append(current_service_target())
|
||||
|
||||
with patch.object(router.cache, "async_increment_cache", new=AsyncMock(side_effect=_increment)):
|
||||
await router.async_deployment_callback_on_failure(
|
||||
kwargs={
|
||||
"call_type": "acompletion",
|
||||
"litellm_params": {
|
||||
"metadata": {"deployment": "openai/gpt-4o", "model_group": "gpt-group"},
|
||||
"model_info": {"id": "dep-1"},
|
||||
},
|
||||
await router.async_deployment_callback_on_failure(
|
||||
kwargs={
|
||||
"call_type": "acompletion",
|
||||
"litellm_params": {
|
||||
"metadata": {"deployment": "openai/gpt-4o", "model_group": "gpt-group"},
|
||||
"model_info": {"id": "dep-1"},
|
||||
},
|
||||
completion_response=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
},
|
||||
completion_response=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
assert seen == ["router_usage"]
|
||||
assert current_service_target() is None
|
||||
assert _usage_by_kind(router) == {"rpm": 1}
|
||||
|
||||
class _SpanRecordingInMemoryCache(InMemoryCache):
|
||||
"""Records the live OTel span each read runs under, so the test sees what a Redis span would nest in."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.active_span_names: list[str] = []
|
||||
def _usage_by_kind(router: Router) -> dict[str, float]:
|
||||
return {key.rsplit(":", 2)[-2]: value for key, value in router._core_usage_store.snapshot().items()}
|
||||
|
||||
async def async_batch_get_cache(self, keys, **kwargs):
|
||||
self.active_span_names.append(trace.get_current_span().name)
|
||||
return await super().async_batch_get_cache(keys, **kwargs)
|
||||
|
||||
async def async_get_cache(self, key, **kwargs):
|
||||
self.active_span_names.append(trace.get_current_span().name)
|
||||
return await super().async_get_cache(key, **kwargs)
|
||||
async def _usage_settles(router: Router, kinds: int) -> dict[str, float]:
|
||||
async with asyncio.timeout(5):
|
||||
while len(_usage_by_kind(router)) < kinds:
|
||||
await asyncio.sleep(0.01)
|
||||
return _usage_by_kind(router)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -24740,10 +24710,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo
|
|||
},
|
||||
]
|
||||
)
|
||||
recording_cache = _SpanRecordingInMemoryCache()
|
||||
router.cache.in_memory_cache = recording_cache
|
||||
router.cooldown_cache.cooldown_store.in_memory_cache = recording_cache
|
||||
|
||||
with TracerProvider().get_tracer("test").start_as_current_span("POST /v1/chat/completions") as server_span:
|
||||
deployment = await router.async_get_available_deployment(model="gpt-group", request_kwargs={})
|
||||
|
||||
|
|
@ -24752,7 +24718,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo
|
|||
assert route_span.name == "route gpt-group"
|
||||
assert route_span.parent is not None and route_span.parent.span_id == server_span.get_span_context().span_id
|
||||
assert route_span.end_time is not None
|
||||
assert recording_cache.active_span_names and set(recording_cache.active_span_names) == {"route gpt-group"}
|
||||
|
||||
|
||||
def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, str | int]]]:
|
||||
|
|
@ -25545,25 +25510,13 @@ async def test_call_router_callbacks_on_success():
|
|||
]
|
||||
)
|
||||
|
||||
with patch.object(router.cache, "async_increment_cache_pipeline", new=AsyncMock()) as mock_callback:
|
||||
await router.acompletion(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
mock_response="Hello, I'm good.",
|
||||
)
|
||||
await asyncio.sleep(1)
|
||||
assert mock_callback.call_count == 1
|
||||
response: Final = await router.acompletion(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
mock_response="Hello, I'm good.",
|
||||
)
|
||||
assert await _usage_settles(router, kinds=2) == {"tpm": response.usage.total_tokens, "rpm": 1}
|
||||
|
||||
increment_list = mock_callback.call_args_list[0].kwargs["increment_list"]
|
||||
assert len(increment_list) == 2
|
||||
|
||||
for increment in increment_list:
|
||||
if "tpm" in increment["key"]:
|
||||
assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:tpm")
|
||||
assert increment["increment_value"] == 30
|
||||
elif "rpm" in increment["key"]:
|
||||
assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm")
|
||||
assert increment["increment_value"] == 1
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown")
|
||||
@pytest.mark.serial
|
||||
|
|
@ -25579,19 +25532,15 @@ async def test_call_router_callbacks_on_failure():
|
|||
]
|
||||
)
|
||||
|
||||
with patch.object(router.cache, "async_increment_cache", new=AsyncMock()) as mock_callback:
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
mock_response="litellm.RateLimitError",
|
||||
num_retries=0,
|
||||
)
|
||||
await asyncio.sleep(3)
|
||||
print(mock_callback.call_args_list)
|
||||
assert mock_callback.call_count == 1
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(
|
||||
model="gemini/gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
||||
mock_response="litellm.RateLimitError",
|
||||
num_retries=0,
|
||||
)
|
||||
assert await _usage_settles(router, kinds=1) == {"rpm": 1}
|
||||
|
||||
assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rate_limit_error_callback():
|
||||
|
|
@ -25745,7 +25694,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result)
|
|||
def test_router_redis_cache():
|
||||
router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}])
|
||||
|
||||
redis_cache = MagicMock()
|
||||
redis_cache = redis_cache_double()
|
||||
|
||||
router.update_redis_cache(cache=redis_cache)
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import weakref
|
|||
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import MethodType
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -44,6 +44,10 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tests.unit.conftest import RouterRedis
|
||||
|
||||
|
||||
def create_standard_logging_payload() -> StandardLoggingPayload:
|
||||
return StandardLoggingPayload(
|
||||
id="test_id",
|
||||
|
|
@ -530,14 +534,14 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
|
|||
"standard_logging_object": standard_logging_payload,
|
||||
}
|
||||
|
||||
tpm_key = await router.deployment_callback_on_success(
|
||||
counter = await router.deployment_callback_on_success(
|
||||
kwargs=turn_kwargs("_aresponses_websocket"),
|
||||
completion_response=None,
|
||||
start_time=time.time(),
|
||||
end_time=time.time(),
|
||||
)
|
||||
assert tpm_key is not None
|
||||
assert await router.cache.async_get_cache(key=tpm_key) == 17
|
||||
assert counter is not None
|
||||
assert await router._core_usage_store.async_batch_get((counter,)) == [17]
|
||||
|
||||
realtime_key = await router.deployment_callback_on_success(
|
||||
kwargs=turn_kwargs("_arealtime"),
|
||||
|
|
@ -546,7 +550,7 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
|
|||
end_time=time.time(),
|
||||
)
|
||||
assert realtime_key is None
|
||||
assert await router.cache.async_get_cache(key=tpm_key) == 17
|
||||
assert await router._core_usage_store.async_batch_get((counter,)) == [17]
|
||||
|
||||
|
||||
class _DiscardingLoggingWorker:
|
||||
|
|
@ -1365,7 +1369,7 @@ async def test_deployment_callback_on_success_tracks_tpm_for_io_deployment():
|
|||
|
||||
# The IO deployment is no longer skipped: its TPM routing counter is tracked.
|
||||
assert tpm_key is not None
|
||||
assert await router.cache.async_get_cache(key=tpm_key) == 100
|
||||
assert (await router.get_model_group_usage("opus"))[0] == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -2045,49 +2049,28 @@ def test_deployment_tracks_usage_for_every_limit_source(
|
|||
assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks
|
||||
|
||||
|
||||
def _shared_redis_stub(store: dict) -> MagicMock:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
async def increment_pipeline(increment_list, **kwargs):
|
||||
for op in increment_list:
|
||||
store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"]
|
||||
return [store[op["key"]] for op in increment_list]
|
||||
|
||||
async def batch_get(keys, **kwargs):
|
||||
return {key: store.get(key) for key in keys}
|
||||
|
||||
redis_stub = MagicMock(spec=RedisCache)
|
||||
redis_stub.async_increment_pipeline = increment_pipeline
|
||||
redis_stub.async_batch_get_cache = batch_get
|
||||
return redis_stub
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("router_minute_pinned")
|
||||
async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
|
||||
store: dict = {}
|
||||
worker_a = _rpm_tpm_router("lit-3058-workers")
|
||||
worker_b = _rpm_tpm_router("lit-3058-workers")
|
||||
worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
|
||||
worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
|
||||
async def test_headers_reflect_usage_counted_by_earlier_requests():
|
||||
router = _rpm_tpm_router("lit-3058-workers")
|
||||
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
tokens_on_a = 0
|
||||
earlier_tokens = 0
|
||||
for _ in range(3):
|
||||
response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
||||
tokens_on_a += response.usage.total_tokens
|
||||
response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
||||
earlier_tokens += response.usage.total_tokens
|
||||
|
||||
response = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
||||
response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
||||
headers = _ratelimit_headers(response)
|
||||
assert headers["x-ratelimit-remaining-requests"] == 96
|
||||
assert headers["x-ratelimit-remaining-tokens"] == 1000 - tokens_on_a - response.usage.total_tokens
|
||||
assert headers["x-ratelimit-remaining-tokens"] == 1000 - earlier_tokens - response.usage.total_tokens
|
||||
|
||||
counted_tokens = tokens_on_a + response.usage.total_tokens
|
||||
counted_tokens = earlier_tokens + response.usage.total_tokens
|
||||
for _ in range(2):
|
||||
response = await worker_a.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
||||
response = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong")
|
||||
counted_tokens += response.usage.total_tokens
|
||||
|
||||
stream = await worker_b.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True)
|
||||
stream = await router.acompletion(model="gpt-5-mini", messages=messages, mock_response="pong", stream=True)
|
||||
stream_headers = _ratelimit_headers(stream)
|
||||
assert stream_headers["x-ratelimit-remaining-requests"] == 93
|
||||
assert stream_headers["x-ratelimit-remaining-tokens"] == 1000 - counted_tokens
|
||||
|
|
|
|||
|
|
@ -12,21 +12,21 @@ import litellm
|
|||
import pytest
|
||||
import respx
|
||||
from litellm import CustomLogger, Router
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
|
||||
get_deployment_successes_for_current_minute,
|
||||
)
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
from litellm.types.router import RoutingStrategy
|
||||
from litellm.types.utils import ImageResponse, ModelResponse, RerankResponse
|
||||
from litellm.types.videos.utils import encode_video_id_with_provider
|
||||
from openai import AsyncAzureOpenAI, AsyncOpenAI
|
||||
|
||||
|
||||
from tests.unit.conftest import redis_cache_double
|
||||
|
||||
EVENT_TIMEOUT_SECONDS: Final = 5
|
||||
ROUTING_SELECTIONS: Final = 40
|
||||
ROUTING_MESSAGES: Final = ({"role": "user", "content": "route this request"},)
|
||||
|
|
@ -534,86 +534,22 @@ def test_router_track_deployment_metrics_updates_observable_usage() -> None:
|
|||
assert router.cache.get_cache(key="metrics-deployment", local_only=True) == 2
|
||||
|
||||
|
||||
class _GatedIncrementCache(DualCache):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(in_memory_cache=InMemoryCache())
|
||||
self.first_increment_started = asyncio.Event()
|
||||
self.release_first_increment = asyncio.Event()
|
||||
self.deployment_success_incremented = asyncio.Event()
|
||||
self.increment_calls = 0
|
||||
|
||||
def increment_cache(self, key: str, value: int, local_only: bool = False, **kwargs: object) -> int:
|
||||
result: Final = super().increment_cache(key=key, value=value, local_only=local_only, **kwargs)
|
||||
if key.endswith(":successes"):
|
||||
self.deployment_success_incremented.set()
|
||||
return result
|
||||
|
||||
async def async_increment_cache_pipeline(
|
||||
self,
|
||||
increment_list: list[RedisPipelineIncrementOperation],
|
||||
local_only: bool = False,
|
||||
parent_otel_span: object = None,
|
||||
**kwargs: object,
|
||||
) -> list[float] | None:
|
||||
self.increment_calls += 1
|
||||
if self.increment_calls == 1:
|
||||
self.first_increment_started.set()
|
||||
await self.release_first_increment.wait()
|
||||
return await super().async_increment_cache_pipeline(
|
||||
increment_list=increment_list,
|
||||
local_only=local_only,
|
||||
parent_otel_span=parent_otel_span,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class _UnavailableIncrementCache(DualCache):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(in_memory_cache=InMemoryCache())
|
||||
self.first_increment_started = asyncio.Event()
|
||||
self.release_first_increment = asyncio.Event()
|
||||
self.deployment_success_incremented = asyncio.Event()
|
||||
self.increment_calls = 0
|
||||
|
||||
def increment_cache(self, key: str, value: int, local_only: bool = False, **kwargs: object) -> int:
|
||||
result: Final = super().increment_cache(key=key, value=value, local_only=local_only, **kwargs)
|
||||
if key.endswith(":successes"):
|
||||
self.deployment_success_incremented.set()
|
||||
return result
|
||||
|
||||
async def async_increment_cache_pipeline(
|
||||
self,
|
||||
increment_list: list[RedisPipelineIncrementOperation],
|
||||
local_only: bool = False,
|
||||
parent_otel_span: object = None,
|
||||
**kwargs: object,
|
||||
) -> list[float] | None:
|
||||
self.increment_calls += 1
|
||||
if self.increment_calls == 1:
|
||||
self.first_increment_started.set()
|
||||
await self.release_first_increment.wait()
|
||||
raise RuntimeError("cache unavailable")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("router_minute_pinned")
|
||||
async def test_router_success_callback_during_pre_header_increment_does_not_double_count() -> None:
|
||||
capture: Final = _RouterLoggingCapture("lit-3058-race")
|
||||
router: Final = _rpm_tpm_router("lit-3058-race")
|
||||
cache: Final = _GatedIncrementCache()
|
||||
router.cache = cache
|
||||
request: Final = asyncio.create_task(
|
||||
router.acompletion(
|
||||
model="gpt-5-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="pong",
|
||||
success_callback=[capture.async_log_success_event],
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.wait_for(cache.first_increment_started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
|
||||
await asyncio.wait_for(cache.deployment_success_incremented.wait(), timeout=EVENT_TIMEOUT_SECONDS)
|
||||
await capture.next_event()
|
||||
assert get_deployment_successes_for_current_minute(router, "lit-3058-race") == 1
|
||||
assert cache.increment_calls == 1
|
||||
cache.release_first_increment.set()
|
||||
response: Final = await request
|
||||
|
||||
assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
|
||||
|
|
@ -621,10 +557,10 @@ async def test_router_success_callback_during_pre_header_increment_does_not_doub
|
|||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("router_minute_pinned")
|
||||
async def test_router_failed_pre_header_increment_clears_counted_tokens_stamp() -> None:
|
||||
async def test_router_failed_redis_increment_keeps_local_usage_without_double_counting() -> None:
|
||||
capture: Final = _RouterLoggingCapture("lit-3058-fail")
|
||||
router: Final = _rpm_tpm_router("lit-3058-fail")
|
||||
cache: Final = _UnavailableIncrementCache()
|
||||
router.cache = cache
|
||||
router.update_redis_cache(redis_cache_double())
|
||||
metadata: Final[dict[str, object]] = {}
|
||||
request: Final = asyncio.create_task(
|
||||
router.acompletion(
|
||||
|
|
@ -632,21 +568,18 @@ async def test_router_failed_pre_header_increment_clears_counted_tokens_stamp()
|
|||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="pong",
|
||||
metadata=metadata,
|
||||
success_callback=[capture.async_log_success_event],
|
||||
)
|
||||
)
|
||||
|
||||
await asyncio.wait_for(cache.first_increment_started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
|
||||
await asyncio.wait_for(cache.deployment_success_incremented.wait(), timeout=EVENT_TIMEOUT_SECONDS)
|
||||
assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30
|
||||
await capture.next_event()
|
||||
assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1
|
||||
assert cache.increment_calls == 1
|
||||
cache.release_first_increment.set()
|
||||
response: Final = await request
|
||||
|
||||
assert response.usage.total_tokens == 30
|
||||
assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in metadata
|
||||
assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 100
|
||||
assert await router.get_model_group_usage("gpt-5-mini") == (None, None)
|
||||
assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == response.usage.total_tokens
|
||||
assert _ratelimit_headers(response)["x-ratelimit-remaining-requests"] == 99
|
||||
assert await router.get_model_group_usage("gpt-5-mini") == (response.usage.total_tokens, 1)
|
||||
|
||||
|
||||
ASSISTANT_RESPONSE: Final = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue