From 70cad58f5e272e1017d7bae0e3c8fd43307db070 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 4 Nov 2026 10:03:03 +0900 Subject: [PATCH] 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 --- litellm-rust/crates/cache-memory/src/cache.rs | 17 + litellm-rust/crates/python-bridge/src/lib.rs | 7 + .../python-bridge/src/python_settings.rs | 2 + .../python-bridge/src/redis/store/mod.rs | 38 + .../python-bridge/src/redis/store/policy.rs | 129 ++++ .../src/redis/store/router_cooldowns.rs | 233 ++++++ .../src/redis/store/router_usage.rs | 149 ++++ litellm-rust/crates/redis-store/src/codec.rs | 20 + litellm-rust/crates/redis-store/src/lib.rs | 8 + litellm-rust/crates/redis-store/src/number.rs | 68 ++ .../redis-store/src/router_cooldowns.rs | 284 ++++++++ .../crates/redis-store/src/router_usage.rs | 262 +++++++ litellm-rust/crates/redis-store/src/tier.rs | 319 +++++++++ .../redis-store/tests/router_cooldowns.rs | 395 ++++++++++ .../crates/redis-store/tests/router_usage.rs | 277 +++++++ .../crates/redis-store/tests/support/mod.rs | 107 +++ litellm/redis/stores/_keys.py | 4 + litellm/redis/stores/router_cooldowns.py | 108 +++ litellm/redis/stores/router_usage.py | 117 +++ litellm/router_python_core.py | 70 +- litellm/router_utils/cooldown_cache.py | 251 +------ litellm/router_utils/cooldown_handlers.py | 26 +- litellm/router_utils/routing_prefetch.py | 165 +++++ litellm/router_utils/routing_read_batch.py | 217 +----- litellm/rust_bridge/_native.pyi | 50 +- litellm/rust_bridge/host/settings.py | 20 + litellm/types/router.py | 6 +- .../test_request_redis_batch_post_call.py | 30 +- .../test_request_redis_batch_pre_call.py | 673 +++++++----------- tests/unit/conftest.py | 94 ++- .../redis/stores/test_router_cooldowns.py | 260 +++++++ tests/unit/redis/stores/test_router_usage.py | 125 ++++ .../router_strategy/test_complexity_router.py | 8 +- .../test_router_tag_routing.py | 30 +- .../test_session_id_affinity.py | 48 +- .../unit/router_utils/test_cooldown_cache.py | 660 +---------------- .../router_utils/test_cooldown_handlers.py | 282 ++------ ..._health_check_allowed_fails_integration.py | 27 +- .../router_utils/test_routing_read_batch.py | 161 +++-- tests/unit/rust_bridge/host/test_settings.py | 25 +- tests/unit/test_internal_context.py | 1 + tests/unit/test_router/test_router.py | 141 ++-- .../test_router/test_router_helper_utils.py | 57 +- .../test_router_provider_endpoints.py | 95 +-- 44 files changed, 3976 insertions(+), 2090 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/policy.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs create mode 100644 litellm-rust/crates/redis-store/src/codec.rs create mode 100644 litellm-rust/crates/redis-store/src/number.rs create mode 100644 litellm-rust/crates/redis-store/src/router_cooldowns.rs create mode 100644 litellm-rust/crates/redis-store/src/router_usage.rs create mode 100644 litellm-rust/crates/redis-store/src/tier.rs create mode 100644 litellm-rust/crates/redis-store/tests/router_cooldowns.rs create mode 100644 litellm-rust/crates/redis-store/tests/router_usage.rs create mode 100644 litellm/redis/stores/_keys.py create mode 100644 litellm/redis/stores/router_cooldowns.py create mode 100644 litellm/redis/stores/router_usage.py create mode 100644 litellm/router_utils/routing_prefetch.py create mode 100644 tests/unit/redis/stores/test_router_cooldowns.py create mode 100644 tests/unit/redis/stores/test_router_usage.py diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 73fbf14b07a..b8706b5e762 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -200,6 +200,23 @@ impl InMemoryCache { .collect()) } + /// Every live entry, for callers that inspect the whole tier. + pub fn entries(&self) -> Result, 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); diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index d735f06ea2b..896bebcf52b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -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}; @@ -121,6 +125,9 @@ mod tests { "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", "NativeRequestRedis", + "NativeDeclaredCooldowns", + "NativeRouterCooldowns", + "NativeRouterUsage", "native_redis_breaker_states", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index eacec6859b5..0152fe8a5b9 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -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")] diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index 498377eb079..b33e1bd096f 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -1,6 +1,10 @@ +mod policy; pub(crate) mod request; +pub(crate) mod router_cooldowns; +pub(crate) mod router_usage; use litellm_host_python::PythonContext; +use litellm_redis_store::Number; use pyo3::{ exceptions::{PyRuntimeError, PyTypeError}, prelude::*, @@ -18,3 +22,37 @@ fn call(py: Python<'_>, call_type: &'static str) -> PyResult { 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 { + if value.is_instance_of::() && !value.is_instance_of::() { + return Ok(Self(Number::Int(value.extract()?))); + } + if value.is_instance_of::() { + 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 { + Ok(match self.0 { + Number::Int(value) => value.into_pyobject(py)?.into_any(), + Number::Float(value) => value.into_pyobject(py)?.into_any(), + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/policy.rs b/litellm-rust/crates/python-bridge/src/redis/store/policy.rs new file mode 100644 index 00000000000..ae2d52136de --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/policy.rs @@ -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 = + FieldSpec::new("cooldown_redis_read_interval", seconds); +const USAGE_REFRESH: FieldSpec = FieldSpec::new("usage_redis_read_interval", seconds); +const TRACKED_READS: FieldSpec = 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 { + Ok(project(&PythonSettings::RouterTiers.read(py)?)?) +} + +fn project(snapshot: &Snapshot<'_>) -> Result { + 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 { + match field.value().extract::() { + 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 { + field + .value() + .extract::() + .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 { + 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}"); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs new file mode 100644 index 00000000000..0e96736bab8 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs @@ -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); + +pub(crate) struct PyCooldown(Cooldown); + +impl<'a, 'py> FromPyObject<'a, 'py> for PyCooldown { + type Error = PyErr; + + fn extract(value: Borrowed<'a, 'py, PyAny>) -> PyResult { + let fields = value.cast::()?; + 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::()?.0, + cooldown_time: item("cooldown_time")?.extract::()?.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 { + 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>> { + 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>) -> PyResult { + 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, + ttl_seconds: f64, + ) -> PyResult { + 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, + ) -> PyResult>> { + 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) -> PyResult> { + 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::>()) + }) + }) + } + + fn declare_active( + &self, + request: &NativeRequestRedis, + model_ids: Vec, + ) -> PyResult { + 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) -> PyResult> { + 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> { + 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, + declared: Mutex>, +} + +#[pymethods] +impl NativeDeclaredCooldowns { + fn resolve<'py>(&self, py: Python<'py>) -> PyResult> { + 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::>()) + }) + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs new file mode 100644 index 00000000000..1af6679f453 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -0,0 +1,149 @@ +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, request::NativeRequestRedis, 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); + +/// `(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 { + let (deployment_id, model, minute, kind): (String, Option, 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>) -> PyResult { + 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> { + 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) + }) + } + + #[pyo3(signature = (deployment_id, model, minute, tokens, requests, *, request = None))] + #[expect( + clippy::too_many_arguments, + reason = "preserves the Python increment_deployment signature" + )] + fn increment_deployment<'py>( + &self, + py: Python<'py>, + deployment_id: String, + model: Option, + minute: String, + tokens: PyNumber, + requests: i64, + request: Option>, + ) -> PyResult> { + let store = self.0.clone(); + let request = request.map(|request| request.0.clone()); + let scope = call(py, "async_increment_pipeline")?; + run_async_value(py, async move { + let model = model.as_deref(); + match request { + Some(request) => store.increment_deployment_post_call( + &request, + &deployment_id, + model, + &minute, + tokens.0, + requests, + ), + None => { + scope + .scope(store.increment_deployment( + &deployment_id, + model, + &minute, + tokens.0, + requests, + )) + .await + } + } + .map_err(store_error) + }) + } + + fn get_many<'py>( + &self, + py: Python<'py>, + counters: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let counters = counters + .into_iter() + .map(|counter| counter.0) + .collect::>(); + 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::>()) + }) + } + + fn flush_local_counters(&self) -> PyResult<()> { + self.0.flush_local_counters().map_err(store_error) + } + + fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { + let snapshot = PyDict::new(py); + for (key, count) in self.0.snapshot().map_err(store_error)? { + snapshot.set_item(key, PyNumber(count))?; + } + Ok(snapshot) + } +} diff --git a/litellm-rust/crates/redis-store/src/codec.rs b/litellm-rust/crates/redis-store/src/codec.rs new file mode 100644 index 00000000000..a8662fcbe0c --- /dev/null +++ b/litellm-rust/crates/redis-store/src/codec.rs @@ -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 { + let text = std::str::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?; + match serde_json::from_str::(text) { + Ok(document) => Ok(json::from_json(document)), + Err(_) => literal::literal_eval(text).map_err(DecodeError::Unparsed), + } +} diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index f64ac797cb5..c87f54163b3 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/redis-store/src/number.rs b/litellm-rust/crates/redis-store/src/number.rs new file mode 100644 index 00000000000..65864e49bba --- /dev/null +++ b/litellm-rust/crates/redis-store/src/number.rs @@ -0,0 +1,68 @@ +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 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 { + 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) -> Option { + 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()), + } + } +} diff --git a/litellm-rust/crates/redis-store/src/router_cooldowns.rs b/litellm-rust/crates/redis-store/src/router_cooldowns.rs new file mode 100644 index 00000000000..f62129c6648 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/router_cooldowns.rs @@ -0,0 +1,284 @@ +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::{Expiry, 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 { + 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 { + 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(policy: TierPolicy, now: &Clock) -> InMemoryCache { + 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, + read: DeclaredRead, +} + +/// One router's cooldowns and allowed-fails counters: memory per router, Redis shared by every pod. +pub struct RouterCooldowns { + cooldowns: MemoryTier, RedisClient>, + fails: MemoryTier, RedisClient>, +} + +impl RouterCooldowns { + pub fn new(client: Option, policy: TierPolicy) -> Self { + Self::with_clock(client, policy, system_clock()) + } + + pub fn with_clock(client: Option, 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 { + 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, Expiry::KeepExisting(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, 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 { + self.declare(|client| request.pre(client), model_ids) + } + + pub async fn resolve_active( + &self, + declared: DeclaredCooldowns, + ) -> Result, 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::, 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, 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, MemoryError> { + self.cooldowns.memory().entries() + } + + async fn read(&self, model_ids: &[String]) -> Result>, MemoryError> { + self.cooldowns + .get_many(&keys(model_ids), Cooldown::decode) + .await + } + + fn declare( + &self, + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + model_ids: &[String], + ) -> Result { + 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 { + 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) + } +} diff --git a/litellm-rust/crates/redis-store/src/router_usage.rs b/litellm-rust/crates/redis-store/src/router_usage.rs new file mode 100644 index 00000000000..d6d410dc5e9 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/router_usage.rs @@ -0,0 +1,262 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::Error as MemoryError; +use litellm_cache_memory::InMemoryCache; +use litellm_redis_client::{Declared, Expiry, RedisBatch, RedisClient}; + +use crate::{ + DecodeError, Number, RequestRedis, + 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, + 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::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, RedisClient>); + +impl RouterUsage { + pub fn new(client: Option, policy: TierPolicy) -> Self { + Self::with_clock(client, policy, system_clock()) + } + + pub fn with_clock(client: Option, 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 { + 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.as_f64(), + 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: memory, then one pipeline that resets each + /// counter's 60 s window, 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 Some(declared) = self.declare(RedisClient::batch, increment)? else { + return Ok(()); + }; + let totals = totals(declared.totals, Declared::resolve).await; + self.mirror(&declared.keys, declared.local, totals) + } + + /// `increment_deployment` on the request's post-call batch. Memory counts at once; the Redis + /// totals reach memory when that batch goes out. + pub fn increment_deployment_post_call( + self: &Arc, + request: &RequestRedis, + 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 Some(declared) = self.declare(|client| request.post(client), increment)? else { + return Ok(()); + }; + let store = Arc::clone(self); + request.spawn(async move { + let totals = totals(declared.totals, Declared::settled).await; + if let Err(error) = store.mirror(&declared.keys, declared.local, totals) { + tracing::warn!(%error, "mirroring Redis usage totals into memory failed"); + } + }); + Ok(()) + } + + pub async fn get_many(&self, counters: &[Counter]) -> Result>, MemoryError> { + let keys = counters.iter().map(Counter::key).collect::>(); + 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, MemoryError> { + self.0.memory().entries() + } + + fn declare( + &self, + batch_for: impl FnOnce(&RedisClient) -> RedisBatch, + increment: Increment, + ) -> Result, MemoryError> { + let local = increment + .counts + .iter() + .map(|(key, value)| self.add_local(key, *value, COUNTER_TTL)) + .collect::, _>>()?; + let Some(remote) = self.0.remote() else { + return Ok(None); + }; + let batch = batch_for(remote); + let totals = increment + .counts + .iter() + .map(|(key, value)| { + batch.incr_by_float(key, value.as_f64(), Expiry::Refresh(COUNTER_TTL)) + }) + .collect(); + Ok(Some(DeclaredIncrement { + keys: increment.counts.into_iter().map(|(key, _)| key).collect(), + local, + totals, + })) + } + + fn mirror( + &self, + keys: &[String], + local: Vec, + totals: Result, litellm_redis_client::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 { + 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)), + ], + } + } +} + +struct DeclaredIncrement { + keys: Vec, + local: Vec, + totals: Vec>, +} + +async fn totals>>( + declared: Vec>, + wait: impl Fn(Declared) -> F, +) -> Result, litellm_redis_client::Error> { + let mut totals = Vec::with_capacity(declared.len()); + for total in declared { + totals.push(wait(total).await?); + } + Ok(totals) +} diff --git a/litellm-rust/crates/redis-store/src/tier.rs b/litellm-rust/crates/redis-store/src/tier.rs new file mode 100644 index 00000000000..ab3dbce51c4 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/tier.rs @@ -0,0 +1,319 @@ +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 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 { + memory: L1, + remote: Option, + policy: TierPolicy, + reads: Mutex, + updates: Mutex<()>, + now: Clock, +} + +pub fn system_clock() -> Clock { + Arc::new(|| { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + }) +} + +impl MemoryTier +where + V: Clone + Send + Sync + 'static, + L1: BaseCache + DeleteCache + FlushCache, +{ + pub fn read_through( + memory: L1, + remote: Option, + 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, + ) -> Result { + 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( + &self, + keys: &[String], + decode: impl Fn(&[u8]) -> Result, + ) -> Result>, 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 { + 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 replies = reservation + .keys + .iter() + .map(|key| batch.get::>(key)) + .collect(); + 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( + &self, + declared: DeclaredRead, + decode: impl Fn(&[u8]) -> Result, + ) -> Result>, 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::, _>>() + { + 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::>(); + 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>, Error> { + let context = ExactCacheContext::default(); + keys.iter() + .map(|key| self.memory.get_cache(key, &context)) + .collect() + } + + fn backfill(&self, key: &str, value: Option) -> Result, 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]) -> Result { + 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::>(); + 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(mutex: &Mutex) -> Result, Error> { + Ok(mutex.lock()?) +} + +async fn settle( + replies: Vec>>>, +) -> Result>>, RedisError> { + let mut fetched = Vec::with_capacity(replies.len()); + for reply in replies { + fetched.push(reply.resolve().await?); + } + Ok(fetched) +} + +/// The GETs `declare_many` put on a batch, with the refreshes they reserved. +pub struct DeclaredRead { + keys: Vec, + fetch: Option, +} + +struct Fetch { + reservation: Reservation, + replies: Vec>>>, +} + +struct Reservation { + keys: Vec, + previous: Vec>, +} + +/// 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, + times: HashMap, +} + +impl RecentReads { + fn new(capacity: usize) -> Self { + Self { + capacity, + order: VecDeque::new(), + times: HashMap::new(), + } + } + + fn get(&self, key: &str) -> Option { + 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); + } + } +} diff --git a/litellm-rust/crates/redis-store/tests/router_cooldowns.rs b/litellm-rust/crates/redis-store/tests/router_cooldowns.rs new file mode 100644 index 00000000000..d79af992ca7 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/router_cooldowns.rs @@ -0,0 +1,395 @@ +mod support; + +use std::time::Duration; + +use litellm_redis_client::{Expiry, 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 = 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::>(); + 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, + Expiry::KeepExisting(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); + +impl StepClock { + fn at(seconds: f64) -> std::sync::Arc { + 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) -> 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(); + + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + 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())] + ); +} diff --git a/litellm-rust/crates/redis-store/tests/router_usage.rs b/litellm-rust/crates/redis-store/tests/router_usage.rs new file mode 100644 index 00000000000..9ade1f0bd6d --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/router_usage.rs @@ -0,0 +1,277 @@ +mod support; + +use std::{sync::Arc, time::Duration}; + +use litellm_redis_client::{Expiry, Operation, Ttl}; +use litellm_redis_store::{ + Number, RequestRedis, + router_usage::{COUNTER_TTL, Counter, Kind, POLICY, RouterUsage}, +}; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, usage, within}; +use tokio::runtime::Handle; + +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, 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; 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, 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))] + ); +} + +#[rstest] +#[tokio::test] +async fn a_post_call_increment_from_a_thread_without_a_runtime_lands_and_mirrors( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let store = Arc::new(RouterUsage::new(Some(connect(&deployment)), POLICY)); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + std::thread::spawn({ + let (store, request) = (Arc::clone(&store), request.clone()); + move || { + store + .increment_deployment_post_call( + &request, + "dep", + Some("model"), + "12-34", + Number::Int(5), + 1, + ) + .unwrap(); + } + }) + .join() + .expect("a Python thread has no Tokio context, and declaring there must not panic"); + request.flush_post().await; + + litellm_redis_testing::eventually(async || { + (store.get_many(&tpm_and_rpm(Some("model"))).await.unwrap() + == [Some(Number::Int(5)), Some(Number::Int(1))]) + .then_some(()) + }) + .await; +} + +#[rstest] +#[tokio::test] +async fn post_call_increments_wait_for_flush_post_then_mirror_the_redis_totals_into_memory( + #[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 = Arc::new(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, 10.0, Expiry::Refresh(COUNTER_TTL)) + .await + .unwrap(); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + store + .increment_deployment_post_call( + &request, + "dep", + Some("model"), + "12-34", + Number::Float(5.5), + 1, + ) + .unwrap(); + let before_flush = store.get_many(&counters).await.unwrap(); + let unsent: Option = raw.get(&rpm_key).await.unwrap(); + let trips_before_flush = trips.seen(); + request.flush_post().await; + let mirrored = litellm_redis_testing::eventually(async || { + let counted = store.get_many(&counters).await.unwrap(); + (counted == [Some(Number::Int(15)), Some(Number::Int(1))]).then_some(counted) + }) + .await; + + assert_eq!( + before_flush, + vec![Some(Number::Float(5.5)), Some(Number::Int(1))] + ); + assert_eq!(unsent, None); + assert!(trips_before_flush.is_empty()); + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + assert_eq!(mirrored, vec![Some(Number::Int(15)), Some(Number::Int(1))]); + let written: [Option; 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_post_call_pipeline_keeps_the_memory_counts_as_integers() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let store = Arc::new(RouterUsage::new(Some(connect(&deployment)), POLICY)); + let counters = tpm_and_rpm(Some("model")); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + deployment.stop().await.unwrap(); + + store + .increment_deployment_post_call( + &request, + "dep", + Some("model"), + "12-34", + Number::Float(5.75), + 1, + ) + .unwrap(); + request.flush_post().await; + + litellm_redis_testing::eventually(async || { + (store.get_many(&counters).await.unwrap() == [Some(Number::Int(5)), Some(Number::Int(1))]) + .then_some(()) + }) + .await; +} diff --git a/litellm-rust/crates/redis-store/tests/support/mod.rs b/litellm-rust/crates/redis-store/tests/support/mod.rs index 889a596db32..7b8effbf3c1 100644 --- a/litellm-rust/crates/redis-store/tests/support/mod.rs +++ b/litellm-rust/crates/redis-store/tests/support/mod.rs @@ -10,6 +10,11 @@ use litellm_redis_client::{ CallEvent, ClusterEndpoint, ClusterNode, Credentials, Endpoint, Namespace, Operation, RedisClient, RedisConfig, RedisEvents, Timeouts, TlsOptions, }; +use litellm_redis_store::{ + Number, + router_cooldowns::Cooldown, + router_usage::{Counter, Kind}, +}; use litellm_redis_testing::RedisDeployment; pub fn connect(deployment: &RedisDeployment) -> RedisClient { @@ -79,3 +84,105 @@ 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)) +} diff --git a/litellm/redis/stores/_keys.py b/litellm/redis/stores/_keys.py new file mode 100644 index 00000000000..cd5cc285ac3 --- /dev/null +++ b/litellm/redis/stores/_keys.py @@ -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}" diff --git a/litellm/redis/stores/router_cooldowns.py b/litellm/redis/stores/router_cooldowns.py new file mode 100644 index 00000000000..64de9b58f79 --- /dev/null +++ b/litellm/redis/stores/router_cooldowns.py @@ -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() diff --git a/litellm/redis/stores/router_usage.py b/litellm/redis/stores/router_usage.py new file mode 100644 index 00000000000..04add303926 --- /dev/null +++ b/litellm/redis/stores/router_usage.py @@ -0,0 +1,117 @@ +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.caching.redis_batch import active_request_redis_batches +from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, 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, + *, + request: NativeRequestRedis | None = None, + ) -> 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: + request: Final = active_request_redis_batches() + await self._handle.increment_deployment( + deployment_id, + model, + minute, + tokens, + requests, + request=None if request is None else request.native_post_call(), + ) + + 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() diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index c63e538c998..ee5e596d3d5 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -238,6 +238,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 ( @@ -1743,6 +1746,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(redis_cache) + 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(), @@ -1848,7 +1853,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 @@ -2204,6 +2209,9 @@ class PythonCore(DeprecatedRouterMembers): and caching to just work. """ self._core_cache.attach_redis_cache(cache) + redis_client: Final = native_client(cache) + 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) @@ -10060,35 +10068,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) @@ -10305,10 +10295,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, @@ -13416,8 +13406,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 @@ -13430,28 +13420,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 @@ -17653,6 +17631,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: diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index dc21db6a601..b8012b261b6 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -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)) diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index e50d5bb47a6..b10af03562c 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -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: diff --git a/litellm/router_utils/routing_prefetch.py b/litellm/router_utils/routing_prefetch.py new file mode 100644 index 00000000000..2b5425574c7 --- /dev/null +++ b/litellm/router_utils/routing_prefetch.py @@ -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] diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 8d8d77fa546..7622991a78d 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -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] diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 5351e572f67..936b3b63b38 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -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: ... @@ -39,6 +45,45 @@ 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, + *, + request: NativeRequestRedis | None = None, + ) -> 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 @@ -228,9 +273,12 @@ __all__ = [ "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDeclaredCooldowns", "NativeDiagnosticProcessor", "NativeRedisClient", "NativeRequestRedis", + "NativeRouterCooldowns", + "NativeRouterUsage", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/litellm/rust_bridge/host/settings.py b/litellm/rust_bridge/host/settings.py index 60a58f1aaba..4932a788221 100644 --- a/litellm/rust_bridge/host/settings.py +++ b/litellm/rust_bridge/host/settings.py @@ -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 diff --git a/litellm/types/router.py b/litellm/types/router.py index 41a986837da..b23239bf4bf 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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 @@ -1290,8 +1292,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}" diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 4b73bb6dd34..b3932285220 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -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"] diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index 94f19dbdc7c..31878658049 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -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 == [] diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index ed7dd5ec146..63faf732b68 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -4,17 +4,103 @@ 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 + + redis: Final = MagicMock(spec=RedisCache) + redis.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2} + redis.namespace = None + 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", @@ -289,7 +375,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 diff --git a/tests/unit/redis/stores/test_router_cooldowns.py b/tests/unit/redis/stores/test_router_cooldowns.py new file mode 100644 index 00000000000..650d263b43b --- /dev/null +++ b/tests/unit/redis/stores/test_router_cooldowns.py @@ -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 == [] diff --git a/tests/unit/redis/stores/test_router_usage.py b/tests/unit/redis/stores/test_router_usage.py new file mode 100644 index 00000000000..fedc414e6d1 --- /dev/null +++ b/tests/unit/redis/stores/test_router_usage.py @@ -0,0 +1,125 @@ +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) + requests: list[object] = 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, + *, + request: object = None, + ) -> "asyncio.Future[None]": + self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target())) + self.requests.append(request) + 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"))) + store.attach_redis(native_client(RedisCache(host="other"))) + 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_ride_its_native_post_call_batch_and_the_shutdown_drain(): + 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) as request: # pyright: ignore[reportArgumentType] # test double + await store.increment_deployment("dep", "model", "12-34", 2, 1, None) + await drain_post_call_redis_batches() + + assert handles.built[0].requests == [None, request.native()] + assert natives.built[0].events == ["flush_post"] diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index ef7b9e485d4..4fa4cf7dd5b 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -15992,9 +15992,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" @@ -16170,9 +16168,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" diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 0c19cfc72f6..1ad6c0a3491 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -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",) diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..de57ff6c718 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -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 diff --git a/tests/unit/router_utils/test_cooldown_cache.py b/tests/unit/router_utils/test_cooldown_cache.py index 98f2b254a3f..9d52b8049ad 100644 --- a/tests/unit/router_utils/test_cooldown_cache.py +++ b/tests/unit/router_utils/test_cooldown_cache.py @@ -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 diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 71c37b5b1d5..44688207aae 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -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 diff --git a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py index 74b3a8825b3..c3b6095d4a3 100644 --- a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py @@ -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.""" diff --git a/tests/unit/router_utils/test_routing_read_batch.py b/tests/unit/router_utils/test_routing_read_batch.py index beafe1e8603..270dd6bb031 100644 --- a/tests/unit/router_utils/test_routing_read_batch.py +++ b/tests/unit/router_utils/test_routing_read_batch.py @@ -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 == [] diff --git a/tests/unit/rust_bridge/host/test_settings.py b/tests/unit/rust_bridge/host/test_settings.py index 5d3486d9a64..a73eb41bd9b 100644 --- a/tests/unit/rust_bridge/host/test_settings.py +++ b/tests/unit/rust_bridge/host/test_settings.py @@ -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: diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 3336cce9aaf..af621a29567 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -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", } ) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 4f35efdfd6e..dffadaa6383 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -105,6 +105,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 @@ -1410,17 +1411,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): """ @@ -1467,11 +1457,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( @@ -15614,11 +15606,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) @@ -15633,7 +15624,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") @@ -15643,7 +15633,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) @@ -15658,10 +15648,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) @@ -15685,10 +15674,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) @@ -15716,11 +15704,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() @@ -24307,13 +24294,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=[ { @@ -24323,42 +24304,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 @@ -24398,10 +24368,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={}) @@ -24410,7 +24376,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]]]: @@ -25176,25 +25141,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 @@ -25210,19 +25163,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(): @@ -25376,7 +25325,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) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 4964dea498f..fdb11ef01bf 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -7,7 +7,7 @@ import weakref from collections.abc import Awaitable, Callable, 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 @@ -35,6 +35,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", @@ -521,14 +525,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"), @@ -537,7 +541,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] async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: @@ -1120,7 +1124,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 @@ -1755,49 +1759,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 diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index d2331c00804..8728de27503 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -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 = {