diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 941887b249c..48b6730b0c6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4558,6 +4558,8 @@ name = "litellm-redis-types" version = "0.1.0" dependencies = [ "litellm-secret-value", + "rstest", + "url", ] [[package]] diff --git a/litellm-rust/crates/circuit-breaker/AGENTS.md b/litellm-rust/crates/circuit-breaker/AGENTS.md index 52ef0ca7435..852cd72664a 100644 --- a/litellm-rust/crates/circuit-breaker/AGENTS.md +++ b/litellm-rust/crates/circuit-breaker/AGENTS.md @@ -1,13 +1,14 @@ - `fsm` is a generic state machine runtime and knows nothing about breakers. `breaker` is a generic circuit breaker and knows nothing about Redis. Deciding which errors count as failures belongs to the consumer - `Machine::step` stays pure. `&self` is config and time comes in on the event, never read inside the machine. The only clock read is `Admission`'s drop, at the call-site boundary - Transition rules are one `match (state, event)` per machine, with no macro DSL and no `_` arm that would hide a new state from the compiler -- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes, matching Python bumping `generation` only when `_state` changes. Data inside a variant (failure counters, `since`) changes without a transition +- `Shared` counts a transition, bumps the generation and calls the observer only when the state's enum variant changes. Data inside a variant (failure counters, `since`) changes without a transition - The observer runs under the lock so transitions are seen in generation order. Keep it to logging or a metric - A disabled breaker is `Option::None` at the consumer, not a flag inside the machine. `Config` comes in already resolved, reading `REDIS_CIRCUIT_BREAKER_*` env is the consumer's job -## Parity with Python `RedisCircuitBreaker` (`litellm/caching/redis_cache.py`) -- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against the Python breaker. A behavior change here needs the same check, by driving the Python class with the same steps -- Divergences, each on purpose +## The only Redis breaker +- Python's `RedisCache` admits through this breaker via `NativeRedisBreaker` (`python-bridge/src/redis/breaker.rs`), sharing one per Redis server with the native client through `redis-connect`'s `RedisHosts`, so the Python `RedisCircuitBreaker` it replaced is gone +- The rstest table `matches_the_python_redis_circuit_breaker` was checked case by case against that Python class and now pins its behavior. A behavior change here changes what every Redis caller in the proxy sees +- Divergences from the old Python class, each on purpose - monotonic `Instant` instead of `time.time()`, so a wall clock jump can't open or close the breaker - a probe `Admission` dropped without an outcome records a hard failure and reopens. Python's probe that dies on `CancelledError` (a `BaseException`) records nothing and leaves the breaker HALF_OPEN for good, and a dropped Rust future is common - `_failure_count` is dropped, it only fed a log line diff --git a/litellm-rust/crates/circuit-breaker/src/breaker.rs b/litellm-rust/crates/circuit-breaker/src/breaker.rs index c9509561b37..cf4961ef099 100644 --- a/litellm-rust/crates/circuit-breaker/src/breaker.rs +++ b/litellm-rust/crates/circuit-breaker/src/breaker.rs @@ -1,5 +1,5 @@ use std::{ - sync::Arc, + sync::{Arc, Weak}, time::{Duration, Instant}, }; @@ -128,6 +128,9 @@ impl Machine for Config { #[derive(Clone)] pub struct Breaker(Arc>); +#[derive(Clone)] +pub struct WeakBreaker(Weak>); + #[derive(Debug, PartialEq, Eq)] pub struct Rejected; @@ -165,6 +168,16 @@ impl Breaker { pub fn state(&self) -> State { self.0.state() } + + pub fn downgrade(&self) -> WeakBreaker { + WeakBreaker(Arc::downgrade(&self.0)) + } +} + +impl WeakBreaker { + pub fn upgrade(&self) -> Option { + self.0.upgrade().map(Breaker) + } } impl Admission { diff --git a/litellm-rust/crates/circuit-breaker/src/lib.rs b/litellm-rust/crates/circuit-breaker/src/lib.rs index e76228122d1..af8df68246e 100644 --- a/litellm-rust/crates/circuit-breaker/src/lib.rs +++ b/litellm-rust/crates/circuit-breaker/src/lib.rs @@ -1,5 +1,7 @@ mod breaker; mod fsm; -pub use breaker::{Admission, Breaker, Config, Event, Failure, Output, Rejected, State}; +pub use breaker::{ + Admission, Breaker, Config, Event, Failure, Output, Rejected, State, WeakBreaker, +}; pub use fsm::{Machine, Observer, Shared, Transition}; diff --git a/litellm-rust/crates/gateway/src/redis/mod.rs b/litellm-rust/crates/gateway/src/redis/mod.rs index 17b111231e7..91428a825f7 100644 --- a/litellm-rust/crates/gateway/src/redis/mod.rs +++ b/litellm-rust/crates/gateway/src/redis/mod.rs @@ -4,7 +4,7 @@ use litellm_circuit_breaker::Config as BreakerConfig; use litellm_config::{CacheSettings, Config, RedisSettings, Value}; use litellm_core_utils::settings::Lookup; use litellm_redis_client::{RedisClient, RedisEvents}; -use litellm_redis_connect::{HostServices, SharedEnvironment, TracingEvents}; +use litellm_redis_connect::{HostServices, RedisHosts, SharedEnvironment, TracingEvents}; use litellm_redis_resolve::Expected; use litellm_redis_types::Resolution; @@ -80,8 +80,9 @@ pub(crate) async fn connect( let plan = plan(config, lookup)?; let breaker = breaker_config(lookup)?; let events: Arc = Arc::new(TracingEvents::new(lookup)); + let hosts = RedisHosts::default(); let client = |params: Params| -> Result { - Ok(litellm_redis_connect::connect( + Ok(hosts.client( params, HostServices { events: Some(events.clone()), diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index dc127e41129..09b787b2c94 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -38,9 +38,9 @@ mod _native { #[pymodule_export] use crate::model_registry::bundled_model_cost_map; #[pymodule_export] - use crate::redis::client::NativeRedisClient; + use crate::redis::breaker::{NativeAdmission, NativeRedisBreaker, native_redis_breaker_states}; #[pymodule_export] - use crate::redis::client::native_redis_breaker_states; + use crate::redis::client::NativeRedisClient; #[pymodule_export] use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes}; #[pymodule_export] @@ -136,6 +136,8 @@ mod tests { "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", "NativeRedisClient", + "NativeRedisBreaker", + "NativeAdmission", "NativeRequestRedis", "NativeSend", "NativeAuthCache", diff --git a/litellm-rust/crates/python-bridge/src/redis/breaker.rs b/litellm-rust/crates/python-bridge/src/redis/breaker.rs new file mode 100644 index 00000000000..733cf1bd969 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/breaker.rs @@ -0,0 +1,134 @@ +use std::{ + sync::{Arc, Mutex, PoisonError}, + time::{Duration, Instant}, +}; + +use litellm_circuit_breaker::{Admission, Breaker, Config, Failure}; +use litellm_redis_client::RedisEvents; +use litellm_redis_types::Endpoint; +use pyo3::{exceptions::PyValueError, prelude::*}; + +use crate::redis::{ + client::{RedisSpec, resolved_endpoint}, + events::{ServiceLoggingEvents, state_label}, + hosts, +}; + +/// `litellm.redis.store_settings.RedisBreakerSettings`, filled from `litellm.constants`. +#[derive(FromPyObject)] +pub(crate) struct BreakerSettings { + failure_threshold: u32, + recovery_timeout: f64, + timeout_min_duration: f64, +} + +impl BreakerSettings { + pub(crate) fn config(&self) -> PyResult { + let seconds = |seconds: f64| { + Duration::try_from_secs_f64(seconds) + .map_err(|error| PyValueError::new_err(error.to_string())) + }; + Ok(Config { + failure_threshold: self.failure_threshold, + recovery_timeout: seconds(self.recovery_timeout)?, + timeout_min_duration: seconds(self.timeout_min_duration)?, + }) + } +} + +/// The one health signal for a Redis, shared by its `RedisCache` and its native client, so a Redis the +/// cache found wedged is fast-failed by every store too. +#[pyclass(frozen, name = "NativeRedisBreaker")] +pub(crate) struct NativeRedisBreaker(pub(crate) Breaker); + +impl NativeRedisBreaker { + fn build( + py: Python<'_>, + settings: &BreakerSettings, + endpoint: Option, + ) -> PyResult { + let config = settings.config()?; + let runtime = litellm_host_python::runtime_handle()?; + let events: Arc = + Arc::new(ServiceLoggingEvents::transitions_only(py, runtime)?); + Ok(Self(py.detach(|| { + hosts().breaker(endpoint, config, Some(events)) + }))) + } +} + +#[pymethods] +impl NativeRedisBreaker { + #[new] + fn new(py: Python<'_>, settings: BreakerSettings) -> PyResult { + Self::build(py, &settings, None) + } + + /// The breaker of the Redis server `redis_kwargs` names, which every cache and client on that server shares + /// whatever its db, credentials or namespace. The first caller's settings win. + #[staticmethod] + fn shared( + py: Python<'_>, + redis_kwargs: RedisSpec, + settings: BreakerSettings, + ) -> PyResult { + let endpoint = resolved_endpoint(py, &redis_kwargs)?; + Self::build(py, &settings, endpoint) + } + + fn admit(&self) -> Option { + self.0 + .admit(Instant::now()) + .ok() + .map(|admission| NativeAdmission(Mutex::new(Some(admission)))) + } + + fn failed(&self, timeout: bool) { + self.0.failed(failure(timeout), Instant::now()); + } + + #[getter] + fn state(&self) -> &'static str { + state_label(&self.0.state()) + } +} + +/// One admitted call. Settling it twice records nothing more, and dropping an unsettled probe reopens the breaker. +#[pyclass(frozen, name = "NativeAdmission")] +pub(crate) struct NativeAdmission(Mutex>); + +impl NativeAdmission { + fn take(&self) -> Option { + self.0.lock().unwrap_or_else(PoisonError::into_inner).take() + } +} + +#[pymethods] +impl NativeAdmission { + fn succeeded(&self) { + if let Some(admission) = self.take() { + admission.succeeded(); + } + } + + fn failed(&self, timeout: bool) { + if let Some(admission) = self.take() { + admission.failed(failure(timeout), Instant::now()); + } + } +} + +fn failure(timeout: bool) -> Failure { + match timeout { + true => Failure::Timeout, + false => Failure::Hard, + } +} + +#[pyfunction] +pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { + py.detach(|| hosts().breakers()) + .iter() + .map(|breaker| state_label(&breaker.state())) + .collect() +} diff --git a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs index 5b0e0ec5d94..989f3dc9247 100644 --- a/litellm-rust/crates/python-bridge/src/redis/client/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/client/mod.rs @@ -1,28 +1,23 @@ mod credentials; -use std::{ - collections::BTreeMap, - sync::{Arc, LazyLock, Mutex, PoisonError}, - time::Duration, -}; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; use litellm_circuit_breaker::Config as BreakerConfig; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; use litellm_python_compat::number::{self, SecondsError}; -use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient}; +use litellm_redis_client::{RedisClient, RedisEvents}; use litellm_redis_connect::{CredentialSource, HostServices, SharedEnvironment}; use litellm_redis_resolve::{Kwargs, PythonException}; -use litellm_redis_types::{ProviderHandle, Resolution}; +use litellm_redis_types::{Endpoint, ProviderHandle, Resolution}; use pyo3::{ exceptions::{PyOverflowError, PyTypeError, PyValueError}, prelude::*, - sync::MutexExt, types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyMapping, PyString, PyTuple}, }; use self::credentials::PythonCredentialSource; -use crate::redis::events::{ServiceLoggingEvents, state_label}; +use crate::redis::{breaker::BreakerSettings, events::ServiceLoggingEvents, hosts}; pyo3::import_exception!(redis.exceptions, RedisClusterException); @@ -67,42 +62,13 @@ impl<'a, 'py> FromPyObject<'a, 'py> for RedisSpec { } } -/// `litellm.redis.store_settings.RedisBreakerSettings`, which `RedisCache` fills from `litellm.constants`. -#[derive(FromPyObject)] -pub(crate) struct RedisBreaker { - #[pyo3(attribute)] - failure_threshold: u32, - #[pyo3(attribute)] - recovery_timeout: f64, - #[pyo3(attribute)] - timeout_min_duration: f64, -} - -impl RedisBreaker { - fn config(&self) -> PyResult { - let seconds = |seconds: f64| { - Duration::try_from_secs_f64(seconds) - .map_err(|error| PyValueError::new_err(error.to_string())) - }; - Ok(BreakerConfig { - failure_threshold: self.failure_threshold, - recovery_timeout: seconds(self.recovery_timeout)?, - timeout_min_duration: seconds(self.timeout_min_duration)?, - }) - } -} - /// `litellm.redis.store_settings.RedisStoreSettings`, which a `RedisCache` captures when it is built. #[derive(FromPyObject)] pub(crate) struct StoreSettings { - #[pyo3(attribute)] redis_kwargs: RedisSpec, - #[pyo3(attribute)] namespace: Option, - #[pyo3(attribute)] default_ttl: f64, - #[pyo3(attribute)] - breaker: Option, + breaker: Option, } fn is_plain(value: &Bound<'_, PyAny>) -> bool { @@ -115,22 +81,6 @@ fn is_plain(value: &Bound<'_, PyAny>) -> bool { || value.is_instance_of::() } -/// A client the process built, held weakly so a discarded `Router` drops its client and provider. -struct Shared { - resolved: Resolution, - client: WeakRedisClient, -} - -static CLIENTS: LazyLock>> = LazyLock::new(|| Mutex::new(Vec::new())); - -fn live_clients(py: Python<'_>) -> std::sync::MutexGuard<'static, Vec> { - let mut clients = CLIENTS - .lock_py_attached(py) - .unwrap_or_else(PoisonError::into_inner); - clients.retain(|shared| shared.client.upgrade().is_some()); - clients -} - /// The one client for a resolved Redis config and credential provider, so every family in this process shares its pool /// and its breaker. A config redis-py refuses building its client raises the same class here. One it only fails on /// once a command connects, where `RedisCache` logs every failed call, leaves the family on memory with an error. @@ -141,17 +91,8 @@ fn shared_client( breaker: Option, ) -> PyResult> { let environment: SharedEnvironment = Arc::new(ProcessEnvironment); - let RedisSpec { - kwargs, - credential_provider, - } = spec; - // The registry entry's client keeps the provider alive, so no live entry shares its address - let handle = credential_provider - .as_ref() - .map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64)); let logger = crate::logger::capture(py); - let environ = py.import("os")?.getattr("environ")?; - let resolved = match logger.scope(|| resolve(&kwargs, handle, &environ))? { + let resolved = match logger.scope(|| resolve_spec(py, &spec))? { Ok(resolution) => { logger.scope(|| { for ignored in &resolution.ignored { @@ -175,7 +116,8 @@ fn shared_client( events: Some(events), http: crate::http::host_client(py, litellm_http::ClientVariant::Provider)?, environment, - credential_provider: credential_provider + credential_provider: spec + .credential_provider .as_ref() .map(|provider| -> PyResult> { Ok(Arc::new(PythonCredentialSource::new( @@ -185,30 +127,14 @@ fn shared_client( }) .transpose()?, }; - let mut clients = live_clients(py); - if let Some(client) = clients - .iter() - .filter(|shared| shared.resolved == resolved) - .find_map(|shared| shared.client.upgrade()) - { - return Ok(Some(client)); - } - let connected = { + let connected = py.detach(|| { let _runtime = runtime.enter(); - litellm_redis_connect::connect(resolved.clone(), host, breaker) - }; - let client = match connected { - Ok(client) => client, - Err(error) => { - drop(clients); - return Ok(unconnected(py, &error)); - } - }; - clients.push(Shared { - resolved, - client: client.downgrade(), + hosts().client(resolved, host, breaker) }); - Ok(Some(client)) + match connected { + Ok(client) => Ok(Some(client)), + Err(error) => Ok(unconnected(py, &error)), + } } /// The Rust client for one `RedisCache`, built once and handed to every native store, so a store @@ -232,7 +158,7 @@ impl NativeRedisClient { breaker, } = settings; let default_ttl = number::seconds(default_ttl).map_err(seconds_error)?; - let breaker = breaker.as_ref().map(RedisBreaker::config).transpose()?; + let breaker = breaker.as_ref().map(BreakerSettings::config).transpose()?; Ok( shared_client(py, redis_kwargs, namespace, breaker)?.map(|client| Self { client, @@ -247,16 +173,23 @@ impl NativeRedisClient { } } -/// The breaker state of every native Redis client this process built, so the Prometheus state -/// gauge counts them next to the Python `RedisCache` breakers. -#[pyfunction] -pub(crate) fn native_redis_breaker_states(py: Python<'_>) -> Vec<&'static str> { - live_clients(py) - .iter() - .filter_map(|shared| shared.client.upgrade()) - .filter_map(|client| client.breaker_state()) - .map(|state| state_label(&state)) - .collect() +/// The Redis server `spec` names, whatever its db, credentials or namespace, or none when the Rust client cannot reach it. +pub(crate) fn resolved_endpoint(py: Python<'_>, spec: &RedisSpec) -> PyResult> { + Ok(resolve_spec(py, spec)? + .ok() + .and_then(|resolution| resolution.target.endpoint())) +} + +// A live registry entry's client keeps its provider alive, so no live entry shares its address +fn resolve_spec( + py: Python<'_>, + spec: &RedisSpec, +) -> PyResult> { + let handle = spec + .credential_provider + .as_ref() + .map(|provider| ProviderHandle::new(provider.as_ptr() as usize as u64)); + resolve(&spec.kwargs, handle, &py.import("os")?.getattr("environ")?) } /// What redis-py makes of the kwargs, with `REDIS_USERNAME` read from `environ` the way Python's Azure AD auth diff --git a/litellm-rust/crates/python-bridge/src/redis/events.rs b/litellm-rust/crates/python-bridge/src/redis/events.rs index 04b0c9dbf32..a281d78d5a6 100644 --- a/litellm-rust/crates/python-bridge/src/redis/events.rs +++ b/litellm-rust/crates/python-bridge/src/redis/events.rs @@ -54,21 +54,33 @@ pub(crate) struct ServiceLoggingEvents { } impl ServiceLoggingEvents { + const MODULE: &str = "litellm.rust_bridge.host.redis_events"; + pub(crate) fn new(py: Python<'_>, runtime: Handle) -> PyResult { - const MODULE: &str = "litellm.rust_bridge.host.redis_events"; - Self::from_module(py, &py.import(MODULE)?, runtime) + Self::from_module( + py, + &py.import(Self::MODULE)?, + runtime, + EventLoop::running(py), + ) + } + + /// For a breaker, whose transitions never run on the caller's loop, so building it never looks one up. + pub(crate) fn transitions_only(py: Python<'_>, runtime: Handle) -> PyResult { + Self::from_module(py, &py.import(Self::MODULE)?, runtime, None) } fn from_module( py: Python<'_>, module: &Bound<'_, PyModule>, runtime: Handle, + event_loop: Option, ) -> PyResult { Ok(Self { record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), context: PythonContext::capture(py)?, - event_loop: EventLoop::running(py), + event_loop, runtime, }) } @@ -226,14 +238,14 @@ def record_redis_breaker(state): let module = PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), &name, &name) .unwrap(); - let events = ServiceLoggingEvents::from_module(py, &module, Handle::current()).unwrap(); - ( - ServiceLoggingEvents { - event_loop: client_loop.map(|client_loop| stand_in_loop(py, client_loop)), - ..events - }, - module.unbind(), + let events = ServiceLoggingEvents::from_module( + py, + &module, + Handle::current(), + client_loop.map(|client_loop| stand_in_loop(py, client_loop)), ) + .unwrap(); + (events, module.unbind()) }) } diff --git a/litellm-rust/crates/python-bridge/src/redis/mod.rs b/litellm-rust/crates/python-bridge/src/redis/mod.rs index 0a54822bcd7..808f930dd20 100644 --- a/litellm-rust/crates/python-bridge/src/redis/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/mod.rs @@ -1,3 +1,15 @@ +use std::sync::LazyLock; + +use litellm_redis_connect::RedisHosts; + +pub(crate) mod breaker; pub(crate) mod client; mod events; pub(crate) mod store; + +static HOSTS: LazyLock = LazyLock::new(RedisHosts::default); + +/// The Redis servers this process talks to, shared by every `RedisCache` and native store in it. +pub(crate) fn hosts() -> &'static RedisHosts { + &HOSTS +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs index bf6ddfbeb52..ca95e1530f4 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -64,9 +64,7 @@ impl<'py> IntoPyObject<'py> for LocalValue { /// `litellm.redis.stores.rate_limits.RateLimitScope`. #[derive(FromPyObject)] struct PyScope { - #[pyo3(attribute)] descriptor_key: String, - #[pyo3(attribute)] descriptor_value: String, } @@ -82,9 +80,7 @@ impl From for Scope { /// `RateLimitCounter`. #[derive(FromPyObject)] struct PyCounter { - #[pyo3(attribute)] scope: PyScope, - #[pyo3(attribute)] kind: Spelled, } @@ -100,9 +96,7 @@ impl From for Counter { /// `RateLimitWindow`, `RateLimitCounter` or `RateLimitGauge`, told apart by the key suffix each names. #[derive(FromPyObject)] struct PyEntry { - #[pyo3(attribute)] suffix: String, - #[pyo3(attribute)] scope: PyScope, } @@ -118,11 +112,8 @@ impl TryFrom for Entry { /// `CounterIncrement`. #[derive(FromPyObject)] struct PyIncrement { - #[pyo3(attribute)] counter: PyCounter, - #[pyo3(attribute)] amount: PyNumber, - #[pyo3(attribute)] ttl: Option, } @@ -139,15 +130,10 @@ impl From for Increment { /// `AtomicCounter`. #[derive(FromPyObject)] struct PyAtomicCounter { - #[pyo3(attribute)] counter: PyCounter, - #[pyo3(attribute)] limit: i64, - #[pyo3(attribute)] increment: i64, - #[pyo3(attribute)] ttl: i64, - #[pyo3(attribute)] window_size: i64, } 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 index f9f2364da19..5d695ca7090 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_cooldowns.rs @@ -1,13 +1,10 @@ -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use litellm_redis_store::{ Send, router_cooldowns::{Cooldown, RouterCooldowns}, }; -use pyo3::{ - prelude::*, - types::{PyDict, PyTuple}, -}; +use pyo3::{prelude::*, types::PyTuple}; use super::{ NativePending, PyData, PyNumber, duration, kind_error, pending, policy::router_tier_policies, @@ -173,11 +170,13 @@ impl NativeRouterCooldowns { self.0.flush_local_counters().map_err(kind_error) } - fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { - let snapshot = PyDict::new(py); - for (key, cooldown) in self.0.snapshot().map_err(kind_error)? { - snapshot.set_item(key, PyData(cooldown))?; - } - Ok(snapshot) + fn snapshot(&self) -> PyResult>> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, cooldown)| (key, PyData(cooldown))) + .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 index bad4085cfe6..48565028de3 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -1,7 +1,7 @@ -use std::{sync::Arc, time::Duration}; +use std::{collections::HashMap, sync::Arc, time::Duration}; use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; -use pyo3::{prelude::*, types::PyDict}; +use pyo3::prelude::*; use super::{ PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote, @@ -11,13 +11,9 @@ use crate::redis::client::NativeRedisClient; /// `litellm.redis.stores.router_usage.UsageCounter`. #[derive(FromPyObject)] struct PyCounter { - #[pyo3(attribute)] deployment_id: String, - #[pyo3(attribute)] model: Option, - #[pyo3(attribute)] minute: String, - #[pyo3(attribute)] kind: Spelled, } @@ -174,11 +170,13 @@ impl NativeRouterUsage { self.0.flush_local_counters().map_err(kind_error) } - fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { - let snapshot = PyDict::new(py); - for (key, count) in self.0.snapshot().map_err(kind_error)? { - snapshot.set_item(key, PyNumber(count))?; - } - Ok(snapshot) + fn snapshot(&self) -> PyResult> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, count)| (key, PyNumber(count))) + .collect()) } } diff --git a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs index 58b77222f5e..ba016821e48 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/spend_counters.rs @@ -1,10 +1,10 @@ -use std::sync::Arc; +use std::{collections::HashMap, sync::Arc}; use litellm_redis_store::{ Number, spend_counters::{LostSeed, Reading, SpendCounter, SpendCounters}, }; -use pyo3::{prelude::*, types::PyDict}; +use pyo3::prelude::*; use super::{ NativePending, PyData, PyNumber, kind_error, pending, redis_async, remote, @@ -246,11 +246,13 @@ impl NativeSpendCounters { self.0.flush_local().map_err(kind_error) } - fn snapshot<'py>(&self, py: Python<'py>) -> PyResult> { - let snapshot = PyDict::new(py); - for (key, spend) in self.0.snapshot().map_err(kind_error)? { - snapshot.set_item(key, PyNumber(spend))?; - } - Ok(snapshot) + fn snapshot(&self) -> PyResult> { + Ok(self + .0 + .snapshot() + .map_err(kind_error)? + .into_iter() + .map(|(key, spend)| (key, PyNumber(spend))) + .collect()) } } diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index 3cf468e0d3f..d9a8b3af7e8 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -5,8 +5,8 @@ use std::{ }; use futures_util::stream::BoxStream; -use litellm_circuit_breaker::{Admission, Breaker, Failure, State}; -use redis::{Cmd, RedisError, Value}; +use litellm_circuit_breaker::{Admission, Breaker, Failure}; +use redis::{Cmd, ErrorKind, RedisError, ServerErrorKind, Value}; use crate::{ CallEvent, Command, ConnectError, Error, FailureKind, Namespace, Operation, Outcome, Pipeline, @@ -76,11 +76,6 @@ impl RedisClient { WeakRedisClient(Arc::downgrade(&self.0)) } - /// The breaker's current state, or none when the breaker is disabled. - pub fn breaker_state(&self) -> Option { - self.0.breaker.as_ref().map(Breaker::state) - } - pub async fn ping(&self) -> Result<(), Error> { self.guarded(Operation::Command("PING"), async { commands::ping() @@ -297,6 +292,15 @@ fn settle(admission: Option, error: Option<&Error>) { admission.failed(Failure::Timeout, Instant::now()) } Some(Error::Transport(_)) => admission.failed(Failure::Hard, Instant::now()), + Some(Error::Server(error)) if proves_redis_answers(error) => admission.succeeded(), Some(Error::Server(_) | Error::Decode(_) | Error::BreakerOpen) => drop(admission), } } + +// A replica refusing a write is the old primary mid-failover, which must not reset a run of failures +fn proves_redis_answers(error: &RedisError) -> bool { + !matches!( + error.kind(), + ErrorKind::InvalidClientConfig | ErrorKind::Server(ServerErrorKind::ReadOnly) + ) +} diff --git a/litellm-rust/crates/redis-client/tests/client.rs b/litellm-rust/crates/redis-client/tests/client.rs index f95023ca6db..e9c6cd7e712 100644 --- a/litellm-rust/crates/redis-client/tests/client.rs +++ b/litellm-rust/crates/redis-client/tests/client.rs @@ -824,6 +824,81 @@ async fn rejected_commands_never_open_the_breaker( ); } +#[rstest] +#[tokio::test] +async fn a_probe_redis_rejects_still_proves_redis_answers_and_closes_the_breaker( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let server = RedisDeployment::start(topology).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + client + .send(commands::set("list", "item", Duration::from_secs(60))) + .await + .unwrap(); + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + assert!( + matches!(breaker.state(), State::Open { .. }), + "{:?}", + breaker.state() + ); + tokio::time::sleep(Duration::from_millis(250)).await; + + let probe = client + .send(commands::incr_by("list", 1, Duration::from_secs(60))) + .await; + + assert!(matches!(probe, Err(Error::Server(_))), "{probe:?}"); + assert!( + matches!(breaker.state(), State::Closed { .. }), + "{:?}", + breaker.state() + ); +} + +#[rstest] +#[tokio::test] +async fn a_replica_refusing_a_write_does_not_reset_a_run_of_failures() { + let server = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let breaker = breaker(2); + let client = connect_guarded( + &server, + "", + Duration::from_millis(300), + Some(breaker.clone()), + ); + client.send(commands::get::("k")).await.unwrap(); + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + server + .cli_on_every_node(&["REPLICAOF", "127.0.0.1", "1"]) + .await + .unwrap(); + + let refused = client + .send(commands::set("k", "v", Duration::from_secs(60))) + .await; + server.freeze().await.unwrap(); + let _ = client.send(commands::get::("k")).await; + server.thaw().await.unwrap(); + + assert!(matches!(refused, Err(Error::Server(_))), "{refused:?}"); + assert!( + matches!(breaker.state(), State::Open { .. }), + "{:?}", + breaker.state() + ); +} + #[rstest] #[tokio::test] async fn an_open_breaker_fails_every_pipeline_op_without_sending( diff --git a/litellm-rust/crates/redis-connect/AGENTS.md b/litellm-rust/crates/redis-connect/AGENTS.md index 1b47804ddd8..8e9d1302eaf 100644 --- a/litellm-rust/crates/redis-connect/AGENTS.md +++ b/litellm-rust/crates/redis-connect/AGENTS.md @@ -1,5 +1,6 @@ - Turns a `litellm_redis_types::Resolution` into a `RedisClient`, the same way for every host. It depends on `litellm-redis-types`, never on `litellm-config`, so a host that resolves Redis some other way connects the same way. `redis-client` stays cloud-agnostic, so credential providers and cloud token sources are built here - The host passes only what it alone can build in `HostServices`: the event sink, the HTTP client and the environment lookup. The breaker config comes from `Resolution::breaker`, which `config` reads from `REDIS_CIRCUIT_BREAKER_*` with Python's fallback, so a bad value keeps the default instead of failing connect. Each cloud source asks its auth crate for one explicit identity, the way the Python SDK builds it: `AwsAuthService::default_chain_credentials` (botocore's default chain), `CredentialSource::application_default` (`google.auth.default`) and `AzureIdentity` (the class `_build_azure_credential` picks). LiteLLM's provider settings such as `AWS_ROLE_NAME`, `AZURE_AD_TOKEN` or `VERTEXAI_CREDENTIALS` belong to Bedrock, Azure OpenAI and Vertex and never reach those constructors. `Auth::Provider` (a handle for a Python `credential_provider` object) fails at connect with an error naming the key, since no host can supply one yet +- `RedisHosts` is the one place a host shares Redis state. Breakers are keyed by `Endpoint`, so every caller on a server sees one health signal whatever its db or credentials. A client takes only breaker settings and admits through its server's breaker, so no caller can pair a client with a breaker of its own, and clients are keyed by `Resolution` alone. Both are held weakly, so a breaker or client is freed with its last user, and the lock is held while a client connects so each is built once. A caller holding the Python GIL releases it first. Hosts own one instance (python-bridge per process, the gateway per `connect`) and never keep a registry of their own - Every expiring credential goes through the one refresh loop in `refresh.rs`. A cloud source only fetches. Never add a second loop or let a source sleep. The expiry deadline lives only in the published `Latest::Ready`, so expiring it leaves `Ready` and cannot fire again. Never keep a copy of the deadline in the loop - AWS SDK identity settings must be in the process environment. `aws-config` cannot take the host lookup, so config-only overrides for profiles, role/web-identity settings, credentials, files and metadata sources fail before connecting rather than silently using another identity. Equal process/config values are allowed. Bedrock-only settings remain unrelated to Redis - redis-rs subscribes once for the connection handshake and again for re-authentication, so a new connection sends a redundant initial AUTH. Unchanged refreshes do not trigger another AUTH. Do not skip a provider's first or alternating subscriptions, because concurrent connections and retries each need the current token. Removing the initial duplicate requires an upstream connection-level fix diff --git a/litellm-rust/crates/redis-connect/src/connect.rs b/litellm-rust/crates/redis-connect/src/connect.rs index 305aa181f2b..cf7fe2334d5 100644 --- a/litellm-rust/crates/redis-connect/src/connect.rs +++ b/litellm-rust/crates/redis-connect/src/connect.rs @@ -25,10 +25,22 @@ pub struct HostServices { // litellm/constants.py REDIS_CLUSTER_HEALTH_CHECK_INTERVAL, the default _redis.py sets for cluster clients only const CLUSTER_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(25); +/// A breaker whose transitions reach `events`, for every client that talks to one Redis. +pub(crate) fn breaker(config: BreakerConfig, events: Option>) -> Breaker { + Breaker::new( + config, + Box::new(move |transition| { + if let Some(events) = &events { + events.breaker(transition); + } + }), + ) +} + pub fn connect( resolved: Resolution, host: HostServices, - breaker: Option, + breaker: Option, ) -> Result { if rustls::crypto::CryptoProvider::get_default().is_none() { let _ = rustls::crypto::ring::default_provider().install_default(); @@ -36,17 +48,6 @@ pub fn connect( tracing::info!(target = ?resolved.target, namespace = ?resolved.namespace, "redis client configured"); let credentials = credentials(resolved.auth, &host)?; let events = host.events; - let breaker = breaker.map(|config| { - let observer = events.clone(); - Breaker::new( - config, - Box::new(move |transition| { - if let Some(events) = &observer { - events.breaker(transition); - } - }), - ) - }); let health_check = health_check(&resolved.target, resolved.health_check); Ok(RedisClient::new(RedisConfig { target: resolved.target, diff --git a/litellm-rust/crates/redis-connect/src/hosts.rs b/litellm-rust/crates/redis-connect/src/hosts.rs new file mode 100644 index 00000000000..11e69d45aff --- /dev/null +++ b/litellm-rust/crates/redis-connect/src/hosts.rs @@ -0,0 +1,241 @@ +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; + +use litellm_circuit_breaker::{Breaker, Config as BreakerConfig, WeakBreaker}; +use litellm_redis_client::{RedisClient, RedisEvents, WeakRedisClient}; +use litellm_redis_types::{Endpoint, Resolution}; + +use crate::{Error, HostServices, breaker, connect}; + +/// The Redis servers one host talks to. Every caller on a server shares its breaker, and callers with the same +/// settings share one client, which admits through its server's breaker. Both are held weakly, so each is freed +/// with its last user and the next caller builds a fresh one. +#[derive(Default)] +pub struct RedisHosts { + breakers: Mutex, WeakBreaker)>>, + clients: Mutex>, +} + +#[derive(PartialEq)] +struct ClientKey { + resolved: Resolution, + guarded: bool, +} + +impl RedisHosts { + /// The breaker of the server `endpoint` names, built from `config` by its first caller. Without an endpoint + /// the breaker is this caller's own, which `breakers` still lists. + pub fn breaker( + &self, + endpoint: Option, + config: BreakerConfig, + events: Option>, + ) -> Breaker { + let mut breakers = live(&self.breakers, WeakBreaker::upgrade); + if let Some(shared) = breakers + .iter() + .filter(|(key, _)| key.is_some() && *key == endpoint) + .find_map(|(_, weak)| weak.upgrade()) + { + return shared; + } + let built = breaker(config, events); + breakers.push((endpoint, built.downgrade())); + built + } + + /// The client for `resolved`, connected by its first caller and guarded by its server's breaker unless `breaker` + /// is none. The lock is held while it connects, so concurrent callers with the same settings connect once. + pub fn client( + &self, + resolved: Resolution, + host: HostServices, + breaker: Option, + ) -> Result { + let breaker = breaker + .map(|config| self.breaker(resolved.target.endpoint(), config, host.events.clone())); + let key = ClientKey { + resolved, + guarded: breaker.is_some(), + }; + let mut clients = live(&self.clients, WeakRedisClient::upgrade); + if let Some(shared) = clients + .iter() + .filter(|(entry, _)| *entry == key) + .find_map(|(_, weak)| weak.upgrade()) + { + return Ok(shared); + } + let client = connect(key.resolved.clone(), host, breaker)?; + clients.push((key, client.downgrade())); + Ok(client) + } + + /// Every breaker a caller still holds, for a state gauge. + pub fn breakers(&self) -> Vec { + live(&self.breakers, WeakBreaker::upgrade) + .iter() + .filter_map(|(_, weak)| weak.upgrade()) + .collect() + } +} + +fn live( + entries: &Mutex>, + upgrade: fn(&W) -> Option, +) -> MutexGuard<'_, Vec<(K, W)>> { + let mut entries = entries.lock().unwrap_or_else(PoisonError::into_inner); + entries.retain(|(_, weak)| upgrade(weak).is_some()); + entries +} + +#[cfg(test)] +mod tests { + use std::{ + sync::Arc, + time::{Duration, Instant}, + }; + + use futures_util::FutureExt; + use litellm_circuit_breaker::{Config, Failure, State}; + use litellm_redis_types::{ + Auth, Endpoint, HealthCheck, Resolution, Standalone, StaticAuth, Target, + }; + use rstest::rstest; + + use super::RedisHosts; + use crate::HostServices; + + const CONFIG: Config = Config { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(60), + timeout_min_duration: Duration::ZERO, + }; + + fn target(host: &str, db: i64) -> Target { + Target::Standalone(Standalone { + host: host.to_owned(), + port: 6379, + tls: None, + db, + }) + } + + fn endpoint(host: &str, db: i64) -> Option { + target(host, db).endpoint() + } + + fn resolved(host: &str, namespace: Option<&str>) -> Resolution { + Resolution { + target: target(host, 0), + auth: Auth::Static(StaticAuth::default()), + namespace: namespace.map(str::to_owned), + timeout: None, + connect_timeout: None, + health_check: HealthCheck::Unset, + keepalive: true, + ignored: Vec::new(), + } + } + + fn services() -> HostServices { + HostServices { + events: None, + http: litellm_http::Client::for_test(), + environment: Arc::new(|_: &str| None), + credential_provider: None, + } + } + + fn open(breaker: &litellm_circuit_breaker::Breaker) { + breaker.failed(Failure::Hard, Instant::now()); + } + + #[rstest] + fn callers_on_one_server_share_its_breaker_whatever_their_db() { + let hosts = RedisHosts::default(); + let first = hosts.breaker(endpoint("r1", 0), CONFIG, None); + let second = hosts.breaker(endpoint("r1", 1), CONFIG, None); + let other = hosts.breaker(endpoint("r2", 0), CONFIG, None); + + open(&first); + + assert!(matches!(second.state(), State::Open { .. })); + assert!(matches!(other.state(), State::Closed { .. })); + } + + #[rstest] + fn a_breaker_without_an_endpoint_is_listed_but_never_shared() { + let hosts = RedisHosts::default(); + let own = hosts.breaker(None, CONFIG, None); + let shared = hosts.breaker(endpoint("r1", 0), CONFIG, None); + + open(&own); + + assert!(matches!(shared.state(), State::Closed { .. })); + assert_eq!(hosts.breakers().len(), 2); + } + + #[rstest] + fn a_server_whose_callers_are_all_gone_gets_a_fresh_breaker() { + let hosts = RedisHosts::default(); + let first = hosts.breaker(endpoint("r1", 0), CONFIG, None); + open(&first); + drop(first); + + let again = hosts.breaker(endpoint("r1", 0), CONFIG, None); + + assert!(matches!(again.state(), State::Closed { .. })); + assert_eq!(hosts.breakers().len(), 1); + } + + #[rstest] + #[tokio::test] + async fn callers_with_the_same_settings_share_one_client() { + let hosts = RedisHosts::default(); + let first = hosts + .client(resolved("r1", Some("a")), services(), Some(CONFIG)) + .unwrap(); + + let same = hosts + .client(resolved("r1", Some("a")), services(), Some(CONFIG)) + .unwrap(); + let other_namespace = hosts + .client(resolved("r1", Some("b")), services(), Some(CONFIG)) + .unwrap(); + + assert!(first.same_client(&same)); + assert!(!first.same_client(&other_namespace)); + } + + #[rstest] + #[tokio::test] + async fn a_client_admits_through_its_servers_breaker() { + let hosts = RedisHosts::default(); + let breaker = hosts.breaker(endpoint("r1", 1), CONFIG, None); + let client = hosts + .client(resolved("r1", None), services(), Some(CONFIG)) + .unwrap(); + + open(&breaker); + + assert!(matches!( + client.ping().now_or_never(), + Some(Err(litellm_redis_client::Error::BreakerOpen)) + )); + } + + #[rstest] + #[tokio::test] + async fn a_client_without_a_breaker_is_never_shared_with_a_guarded_one() { + let hosts = RedisHosts::default(); + let guarded = hosts + .client(resolved("r1", None), services(), Some(CONFIG)) + .unwrap(); + + let unguarded = hosts + .client(resolved("r1", None), services(), None) + .unwrap(); + + assert!(!guarded.same_client(&unguarded)); + } +} diff --git a/litellm-rust/crates/redis-connect/src/lib.rs b/litellm-rust/crates/redis-connect/src/lib.rs index b3d1d411eed..7e37f721679 100644 --- a/litellm-rust/crates/redis-connect/src/lib.rs +++ b/litellm-rust/crates/redis-connect/src/lib.rs @@ -1,11 +1,14 @@ mod connect; mod error; mod events; +mod hosts; mod refresh; mod sources; +pub(crate) use connect::breaker; pub use connect::{HostServices, SharedEnvironment, connect}; pub use error::{CredentialStage, Error}; pub use events::TracingEvents; +pub use hosts::RedisHosts; pub(crate) use refresh::RefreshingCredentials; pub use refresh::{CredentialSource, Expiring, SourceFuture}; diff --git a/litellm-rust/crates/redis-types/Cargo.toml b/litellm-rust/crates/redis-types/Cargo.toml index 229325a27bd..49df066cc48 100644 --- a/litellm-rust/crates/redis-types/Cargo.toml +++ b/litellm-rust/crates/redis-types/Cargo.toml @@ -7,3 +7,7 @@ repository.workspace = true [dependencies] litellm-secret-value.workspace = true +url.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/redis-types/src/auth.rs b/litellm-rust/crates/redis-types/src/auth.rs new file mode 100644 index 00000000000..b2386cf9013 --- /dev/null +++ b/litellm-rust/crates/redis-types/src/auth.rs @@ -0,0 +1,50 @@ +use litellm_secret_value::SecretValue; + +#[derive(Clone, Debug, PartialEq)] +pub enum Auth { + Static(StaticAuth), + /// Credentials come from a host object, like a redis-py `credential_provider`. + Provider(ProviderHandle), + GcpIam { + service_account: String, + }, + AzureAd(AzureAdAuth), + AwsIam(AwsIamAuth), +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct StaticAuth { + pub username: Option, + pub password: Option, +} + +/// Names one credential provider the host holds. Two resolutions with the same handle use the +/// same provider. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct ProviderHandle(u64); + +impl ProviderHandle { + pub const fn new(id: u64) -> Self { + Self(id) + } + + pub const fn id(self) -> u64 { + self.0 + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct AzureAdAuth { + pub client_id: Option, + pub tenant_id: Option, + pub client_secret: Option, + pub username: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct AwsIamAuth { + pub user_name: String, + pub cache_name: String, + pub region: String, + pub serverless: bool, +} diff --git a/litellm-rust/crates/redis-types/src/endpoint.rs b/litellm-rust/crates/redis-types/src/endpoint.rs new file mode 100644 index 00000000000..e58bcd8d32d --- /dev/null +++ b/litellm-rust/crates/redis-types/src/endpoint.rs @@ -0,0 +1,169 @@ +use std::{collections::BTreeSet, path::PathBuf}; + +use crate::{Node, Target}; + +const DEFAULT_PORT: u16 = 6379; + +/// The Redis server a target reaches, for state that belongs to the server rather than to one +/// connection, like its circuit breaker. Targets that differ only in db, credentials or TLS reach +/// one server. Different names for one host, like `localhost` and `127.0.0.1`, stay apart. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Endpoint { + Tcp { + host: String, + port: u16, + }, + Unix(PathBuf), + Cluster(BTreeSet<(String, u16)>), + Sentinel { + monitors: BTreeSet<(String, u16)>, + service_name: String, + }, +} + +impl Target { + /// None for a url whose server cannot be read from it. + pub fn endpoint(&self) -> Option { + match self { + Target::Standalone(host) => Some(Endpoint::tcp(&host.host, host.port)), + Target::Unix(socket) => Some(Endpoint::Unix(socket.path.clone())), + Target::Url(url) => url_endpoint(url.url.expose()), + Target::Cluster(cluster) => Some(Endpoint::Cluster(nodes(&cluster.nodes))), + Target::Sentinel(sentinel) => Some(Endpoint::Sentinel { + monitors: nodes(&sentinel.monitors), + service_name: sentinel.service_name.clone(), + }), + } + } +} + +impl Endpoint { + fn tcp(host: &str, port: u16) -> Self { + Self::Tcp { + host: host.to_ascii_lowercase(), + port, + } + } +} + +fn nodes(nodes: &[Node]) -> BTreeSet<(String, u16)> { + nodes + .iter() + .map(|node| (node.host.to_ascii_lowercase(), node.port)) + .collect() +} + +// redis-py's parse_url schemes. host() rather than host_str(), which keeps an IPv6 address's brackets +fn url_endpoint(url: &str) -> Option { + let url = url::Url::parse(url).ok()?; + match url.scheme() { + "redis" | "rediss" => { + let host = match url.host()? { + url::Host::Domain(domain) => domain.to_owned(), + url::Host::Ipv4(ip) => ip.to_string(), + url::Host::Ipv6(ip) => ip.to_string(), + }; + Some(Endpoint::tcp(&host, url.port().unwrap_or(DEFAULT_PORT))) + } + "unix" => Some(Endpoint::Unix(url.path().into())), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use litellm_secret_value::SecretValue; + use rstest::rstest; + + use crate::{Cluster, Node, RedisUrl, Sentinel, Standalone, Target, TlsOptions, UnixSocket}; + + fn host(host: &str, port: u16, db: i64) -> Target { + Target::Standalone(Standalone { + host: host.to_owned(), + port, + tls: None, + db, + }) + } + + fn url(url: &str) -> Target { + Target::Url(RedisUrl { + url: SecretValue::new(url.to_owned()), + tls: None, + db: None, + username: None, + password: None, + }) + } + + fn cluster(nodes: &[(&str, u16)]) -> Target { + Target::Cluster(Cluster { + nodes: nodes + .iter() + .map(|(host, port)| Node { + host: (*host).to_owned(), + port: *port, + }) + .collect(), + tls: None, + }) + } + + fn sentinel(service_name: &str, db: i64) -> Target { + Target::Sentinel(Sentinel { + monitors: vec![Node { + host: "s1".to_owned(), + port: 26379, + }], + service_name: service_name.to_owned(), + monitor_auth: None, + tls: None, + db, + }) + } + + #[rstest] + #[case::another_db(host("r1", 6379, 0), host("r1", 6379, 1))] + #[case::host_case(host("r1", 6379, 0), host("R1", 6379, 0))] + #[case::tls( + host("r1", 6379, 0), + Target::Standalone(Standalone { + host: "r1".to_owned(), + port: 6379, + tls: Some(TlsOptions::verified()), + db: 0, + }) + )] + #[case::url_and_host(host("r1", 6379, 0), url("redis://r1:6379/0"))] + #[case::url_credentials_and_db(url("redis://r1:6379/0"), url("redis://user:pw@r1/2?db=3"))] + #[case::rediss_url(host("r1", 6380, 0), url("rediss://r1:6380"))] + #[case::ipv6_url(host("::1", 6379, 0), url("redis://[::1]:6379"))] + #[case::unix_url( + Target::Unix(UnixSocket { path: "/tmp/redis.sock".into(), db: 0 }), + url("unix:///tmp/redis.sock?db=2") + )] + #[case::cluster_node_order(cluster(&[("a", 1), ("b", 2)]), cluster(&[("b", 2), ("a", 1)]))] + #[case::sentinel_db(sentinel("mymaster", 0), sentinel("mymaster", 1))] + fn targets_on_one_server_share_an_endpoint(#[case] left: Target, #[case] right: Target) { + assert_eq!(left.endpoint(), right.endpoint()); + assert!(left.endpoint().is_some()); + } + + #[rstest] + #[case::another_port(host("r1", 6379, 0), url("redis://r1:6380"))] + #[case::another_host(host("r1", 6379, 0), host("r2", 6379, 0))] + #[case::another_cluster(cluster(&[("a", 1)]), cluster(&[("a", 1), ("b", 2)]))] + #[case::another_service(sentinel("mymaster", 0), sentinel("other", 0))] + #[case::tcp_and_unix( + host("localhost", 6379, 0), + Target::Unix(UnixSocket { path: "/tmp/redis.sock".into(), db: 0 }) + )] + fn targets_on_different_servers_do_not(#[case] left: Target, #[case] right: Target) { + assert_ne!(left.endpoint(), right.endpoint()); + } + + #[rstest] + fn a_url_with_no_server_has_no_endpoint() { + assert_eq!(url("http://r1:6379").endpoint(), None); + } +} diff --git a/litellm-rust/crates/redis-types/src/lib.rs b/litellm-rust/crates/redis-types/src/lib.rs index 0ea69edfa6a..9538f8e1e4d 100644 --- a/litellm-rust/crates/redis-types/src/lib.rs +++ b/litellm-rust/crates/redis-types/src/lib.rs @@ -1,175 +1,11 @@ -use std::{path::PathBuf, time::Duration}; +mod auth; +mod endpoint; +mod resolution; +mod target; +mod tls; -use litellm_secret_value::SecretValue; - -/// One Redis connection with every setting, environment value and secret already applied. Each -/// client library builds its client from this one value, so they cannot read a setting differently. -#[derive(Clone, Debug, PartialEq)] -pub struct Resolution { - pub target: Target, - pub auth: Auth, - pub namespace: Option, - pub timeout: Option, - pub connect_timeout: Option, - pub health_check: HealthCheck, - /// redis-py's `socket_keepalive`, on unless set false, so an idle connection a load balancer - /// dropped is noticed. - pub keepalive: bool, - pub ignored: Vec, -} - -/// redis-py's `health_check_interval`, which `_coerce_redis_kwargs_types` reads as an `int` of -/// seconds. Unset and zero differ, because a client may give an unset interval its own default. -#[derive(Clone, Copy, Debug, PartialEq)] -pub enum HealthCheck { - Unset, - Disabled, - Every(Duration), -} - -/// A setting the Rust client accepts without honoring, so the host can log what ignoring it means. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct Ignored { - /// The key, the environment variable or the url query parameter that set it. - pub setting: String, - pub consequence: &'static str, -} - -#[derive(Clone, Debug, PartialEq)] -pub enum Target { - Standalone(Standalone), - /// Like redis-py's `UnixDomainSocketConnection`, which replaces host, port and TLS. - Unix(UnixSocket), - Url(RedisUrl), - Cluster(Cluster), - Sentinel(Sentinel), -} - -#[derive(Clone, Debug, PartialEq)] -pub struct Standalone { - pub host: String, - pub port: u16, - pub tls: Option, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct UnixSocket { - pub path: PathBuf, - pub db: i64, -} - -/// A url kept verbatim, plus what redis-py's `parse_url` reads from its query that the url's own -/// grammar does not carry. -#[derive(Clone, Debug, PartialEq)] -pub struct RedisUrl { - pub url: SecretValue, - /// Set only for a `rediss://` url, from its query parameters. - pub tls: Option, - /// `?db=`, which wins over the path. - pub db: Option, - /// `?username=` and `?password=`, which the url's userinfo wins over. - pub username: Option, - pub password: Option, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct Cluster { - pub nodes: Vec, - pub tls: Option, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct Node { - pub host: String, - pub port: u16, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct Sentinel { - pub monitors: Vec, - pub service_name: String, - /// How to authenticate to the monitors, which are separate servers from the data node. - pub monitor_auth: Option, - pub tls: Option, - pub db: i64, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct TlsOptions { - pub ca_certs: Option, - /// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data` - pub ca_data: Option, - /// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path` - pub ca_path: Option, - pub certfile: Option, - pub keyfile: Option, - /// Unlocks an encrypted `keyfile`, redis-py's `ssl_password` - pub key_password: Option, - pub verify_certificate: bool, - pub check_hostname: bool, -} - -impl TlsOptions { - pub fn verified() -> Self { - Self { - ca_certs: None, - ca_data: None, - ca_path: None, - certfile: None, - keyfile: None, - key_password: None, - verify_certificate: true, - check_hostname: false, - } - } -} - -#[derive(Clone, Debug, PartialEq)] -pub enum Auth { - Static(StaticAuth), - /// Credentials come from a host object, like a redis-py `credential_provider`. - Provider(ProviderHandle), - GcpIam { - service_account: String, - }, - AzureAd(AzureAdAuth), - AwsIam(AwsIamAuth), -} - -#[derive(Clone, Debug, Default, PartialEq)] -pub struct StaticAuth { - pub username: Option, - pub password: Option, -} - -/// Names one credential provider the host holds. Two resolutions with the same handle use the -/// same provider. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub struct ProviderHandle(u64); - -impl ProviderHandle { - pub const fn new(id: u64) -> Self { - Self(id) - } - - pub const fn id(self) -> u64 { - self.0 - } -} - -#[derive(Clone, Debug, PartialEq)] -pub struct AzureAdAuth { - pub client_id: Option, - pub tenant_id: Option, - pub client_secret: Option, - pub username: Option, -} - -#[derive(Clone, Debug, PartialEq)] -pub struct AwsIamAuth { - pub user_name: String, - pub cache_name: String, - pub region: String, - pub serverless: bool, -} +pub use auth::{Auth, AwsIamAuth, AzureAdAuth, ProviderHandle, StaticAuth}; +pub use endpoint::Endpoint; +pub use resolution::{HealthCheck, Ignored, Resolution}; +pub use target::{Cluster, Node, RedisUrl, Sentinel, Standalone, Target, UnixSocket}; +pub use tls::TlsOptions; diff --git a/litellm-rust/crates/redis-types/src/resolution.rs b/litellm-rust/crates/redis-types/src/resolution.rs new file mode 100644 index 00000000000..75ffe5ba487 --- /dev/null +++ b/litellm-rust/crates/redis-types/src/resolution.rs @@ -0,0 +1,36 @@ +use std::time::Duration; + +use crate::{Auth, Target}; + +/// One Redis connection with every setting, environment value and secret already applied. Each +/// client library builds its client from this one value, so they cannot read a setting differently. +#[derive(Clone, Debug, PartialEq)] +pub struct Resolution { + pub target: Target, + pub auth: Auth, + pub namespace: Option, + pub timeout: Option, + pub connect_timeout: Option, + pub health_check: HealthCheck, + /// redis-py's `socket_keepalive`, on unless set false, so an idle connection a load balancer + /// dropped is noticed. + pub keepalive: bool, + pub ignored: Vec, +} + +/// redis-py's `health_check_interval`, which `_coerce_redis_kwargs_types` reads as an `int` of +/// seconds. Unset and zero differ, because a client may give an unset interval its own default. +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum HealthCheck { + Unset, + Disabled, + Every(Duration), +} + +/// A setting the Rust client accepts without honoring, so the host can log what ignoring it means. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Ignored { + /// The key, the environment variable or the url query parameter that set it. + pub setting: String, + pub consequence: &'static str, +} diff --git a/litellm-rust/crates/redis-types/src/target.rs b/litellm-rust/crates/redis-types/src/target.rs new file mode 100644 index 00000000000..4a6655a23a6 --- /dev/null +++ b/litellm-rust/crates/redis-types/src/target.rs @@ -0,0 +1,65 @@ +use std::path::PathBuf; + +use litellm_secret_value::SecretValue; + +use crate::{StaticAuth, TlsOptions}; + +#[derive(Clone, Debug, PartialEq)] +pub enum Target { + Standalone(Standalone), + /// Like redis-py's `UnixDomainSocketConnection`, which replaces host, port and TLS. + Unix(UnixSocket), + Url(RedisUrl), + Cluster(Cluster), + Sentinel(Sentinel), +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Standalone { + pub host: String, + pub port: u16, + pub tls: Option, + pub db: i64, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct UnixSocket { + pub path: PathBuf, + pub db: i64, +} + +/// A url kept verbatim, plus what redis-py's `parse_url` reads from its query that the url's own +/// grammar does not carry. +#[derive(Clone, Debug, PartialEq)] +pub struct RedisUrl { + pub url: SecretValue, + /// Set only for a `rediss://` url, from its query parameters. + pub tls: Option, + /// `?db=`, which wins over the path. + pub db: Option, + /// `?username=` and `?password=`, which the url's userinfo wins over. + pub username: Option, + pub password: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Cluster { + pub nodes: Vec, + pub tls: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Node { + pub host: String, + pub port: u16, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct Sentinel { + pub monitors: Vec, + pub service_name: String, + /// How to authenticate to the monitors, which are separate servers from the data node. + pub monitor_auth: Option, + pub tls: Option, + pub db: i64, +} diff --git a/litellm-rust/crates/redis-types/src/tls.rs b/litellm-rust/crates/redis-types/src/tls.rs new file mode 100644 index 00000000000..dd664c3c3c1 --- /dev/null +++ b/litellm-rust/crates/redis-types/src/tls.rs @@ -0,0 +1,33 @@ +use std::path::PathBuf; + +use litellm_secret_value::SecretValue; + +#[derive(Clone, Debug, PartialEq)] +pub struct TlsOptions { + pub ca_certs: Option, + /// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data` + pub ca_data: Option, + /// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path` + pub ca_path: Option, + pub certfile: Option, + pub keyfile: Option, + /// Unlocks an encrypted `keyfile`, redis-py's `ssl_password` + pub key_password: Option, + pub verify_certificate: bool, + pub check_hostname: bool, +} + +impl TlsOptions { + pub fn verified() -> Self { + Self { + ca_certs: None, + ca_data: None, + ca_path: None, + certfile: None, + keyfile: None, + key_password: None, + verify_certificate: true, + check_hostname: false, + } + } +} diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index ade2ac1b6ca..d58ddb9c629 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -384,7 +384,7 @@ class RedisBatch: return await pipe.execute(raise_on_error=False) try: - replies: Final = await _run_under_circuit_breaker(self.redis_cache._circuit_breaker, self.name, run) # pyright: ignore[reportPrivateUsage] # same breaker as the cache's own methods + replies: Final = await _run_under_circuit_breaker(self.redis_cache._health, self.name, run) # pyright: ignore[reportPrivateUsage] # same breaker as the cache's own methods except Exception as e: # noqa: BLE001 # each declaring caller applies its own Redis fallback log_redis_failure(verbose_logger, logging.WARNING, f"{self.name}: pipeline of {len(ops)} ops failed", e) with service_target(target): diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 4d9d26fc70c..3f047e6001d 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,7 +18,6 @@ import json import logging import threading import time -import weakref from collections import Counter from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from contextvars import ContextVar @@ -34,10 +33,6 @@ from litellm._internal_context import current_service_caller from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( DEFAULT_REDIS_MAJOR_VERSION, - REDIS_CIRCUIT_BREAKER_ENABLED, - REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, - REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, - REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, REDIS_TIMEOUT_LOG_INTERVAL, ) from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, has_attributes @@ -59,6 +54,7 @@ if TYPE_CHECKING: from redis.asyncio.cluster import ClusterPipeline from litellm._service_logger import ServiceLogging + from litellm.rust_bridge._native import NativeAdmission, NativeRedisBreaker pipeline = Pipeline cluster_pipeline = ClusterPipeline @@ -230,142 +226,6 @@ def _get_call_stack_info(num_frames: int = 2) -> str: return "unknown" -class RedisCircuitBreaker: - """ - Tracks Redis health for a RedisCache instance. - - States: - CLOSED - normal, Redis is called - OPEN - Redis is down, raise immediately (no network call) - HALF_OPEN - recovery probe: allow one request through - - Transitions: - CLOSED -> OPEN after failure_threshold consecutive hard connectivity - failures, or after an unbroken run of timeout failures - (no success or hard failure in between) that reaches - failure_threshold and spans timeout_min_duration seconds - OPEN -> HALF_OPEN after recovery_timeout seconds - HALF_OPEN -> CLOSED on success - HALF_OPEN -> OPEN on failure (resets timer) - - Timeouts are accounted separately from hard connectivity failures because the async - Redis timeout includes time waiting for the worker event loop to resume: one loop - stall makes every in-flight operation time out together, which satisfies a purely - consecutive threshold instantly even though Redis is healthy. Requiring a - timeout-only streak to also span timeout_min_duration filters such bursts while a - real outage that surfaces as timeouts still opens the breaker after that duration. - """ - - CLOSED = "closed" - OPEN = "open" - HALF_OPEN = "half_open" - - def __init__( - self, - failure_threshold: int, - recovery_timeout: int, - enabled: bool = True, - timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION, - service_logger: "ServiceLogging | None" = None, - ) -> None: - self._service_logger = service_logger - self.failure_threshold = failure_threshold - self.recovery_timeout = recovery_timeout - self.enabled = enabled - self.timeout_min_duration = timeout_min_duration - self._failure_count = 0 - self._hard_failure_count = 0 - self._timeout_count = 0 - self._timeout_streak_started_at: float | None = None - self._opened_at: float | None = None - self._state = self.CLOSED - self._generation = 0 - _LIVE_BREAKERS.add(self) - self._report(transitions_to=None, failure_class=None) - - @property - def state(self) -> str: - return self._state - - @property - def generation(self) -> int: - """Counts state transitions, so a call can tell whether the breaker moved while it ran.""" - return self._generation - - def is_open(self) -> bool: - """Returns True if Redis calls should be skipped.""" - if not self.enabled: - return False - if self._state == self.HALF_OPEN: - # Probe already in flight — fast-fail all concurrent requests. - # Only the one call that caused the OPEN→HALF_OPEN transition - # (which returned False) is the designated probe. - return True - if self._state == self.OPEN: - if time.time() - (self._opened_at or 0) > self.recovery_timeout: - self._set_state(self.HALF_OPEN) - return False # this caller is the designated probe - return True - return False - - def _should_open(self, now: float) -> bool: - if self._state == self.HALF_OPEN: - return True - if self._hard_failure_count >= self.failure_threshold: - return True - if self._timeout_count < self.failure_threshold: - return False - return now - (self._timeout_streak_started_at or now) >= self.timeout_min_duration - - def record_failure(self, is_timeout: bool = False) -> None: - if not self.enabled: - return - now: Final = time.time() - self._failure_count += 1 - if is_timeout: - self._timeout_count += 1 - if self._timeout_streak_started_at is None: - self._timeout_streak_started_at = now - else: - self._hard_failure_count += 1 - self._timeout_count = 0 - self._timeout_streak_started_at = None - self._opened_at = now - self._report(transitions_to=None, failure_class="timeout" if is_timeout else "connectivity") - if self._should_open(now): - if self._state != self.OPEN: - verbose_logger.warning( - "Redis circuit breaker OPENED after %d consecutive failures" - " (%d hard connectivity) — fast-failing Redis calls for %ds", - self._failure_count, - self._hard_failure_count, - self.recovery_timeout, - ) - self._set_state(self.OPEN) - - def record_success(self) -> None: - if not self.enabled or self._state == self.OPEN: - return - if self._state == self.HALF_OPEN: - verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered") - self._failure_count = 0 - self._hard_failure_count = 0 - self._timeout_count = 0 - self._timeout_streak_started_at = None - self._set_state(self.CLOSED) - - def _set_state(self, state: str) -> None: - if state == self._state: - return - self._state = state - self._generation += 1 - self._report(transitions_to=state, failure_class=None) - - def _report(self, transitions_to: str | None, failure_class: str | None) -> None: - if self._service_logger is not None: - self._service_logger.redis_breaker_hook(transitions_to=transitions_to, failure_class=failure_class) - - _RedisCallResult = TypeVar("_RedisCallResult") @@ -470,7 +330,22 @@ def is_redis_timeout_failure(exc: BaseException) -> bool: return any(isinstance(link, timeout_types) for link in _explicit_causes(exc)) -def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None: +@dataclass(frozen=True, slots=True) +class RedisHealth: + """The breaker a ``RedisCache`` shares with every other cache and native client on its Redis.""" + + breaker: "NativeRedisBreaker" + service_logger: "ServiceLogging" + + def failed(self, settles: "NativeRedisBreaker | NativeAdmission", exc: BaseException) -> None: + timeout: Final = is_redis_timeout_failure(exc) + settles.failed(timeout) + self.service_logger.redis_breaker_hook( + transitions_to=None, failure_class="timeout" if timeout else "connectivity" + ) + + +def _record_swallowed_redis_failure(health: RedisHealth | None, exc: BaseException) -> None: """Record a Redis failure that the calling method is about to swallow. The marker is a ContextVar rather than a counter on the breaker because breakers are @@ -479,24 +354,19 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep failure would be discarded and a Redis that is answering would still be evicted. asyncio gives each task its own copy of the context, so this is per-call. """ - if not _is_redis_health_failure(exc): + if health is None or not _is_redis_health_failure(exc): return - breaker.record_failure(is_timeout=is_redis_timeout_failure(exc)) + health.failed(health.breaker, exc) _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) -_LIVE_BREAKERS: Final[weakref.WeakSet[RedisCircuitBreaker]] = weakref.WeakSet() - - def redis_breaker_state_counts() -> Mapping[str, int]: """How many live Redis circuit breakers sit in each state, for the Prometheus services logger. - Counts the Python breakers of every live ``RedisCache`` and the breaker of every native Redis - client the router stores built, since each opens and closes on its own.""" + Each Redis has one breaker, which every ``RedisCache`` and native client on it shares.""" from litellm.rust_bridge._native import native_redis_breaker_states - python_states: Final = Counter(breaker.state for breaker in _LIVE_BREAKERS) - return MappingProxyType(dict(python_states + Counter(native_redis_breaker_states()))) + return MappingProxyType(dict(Counter(native_redis_breaker_states()))) class RedisCircuitBreakerOpenError(Exception): @@ -559,35 +429,60 @@ def log_redis_failure( @dataclass(frozen=True, slots=True) class _BreakerAdmission: + admission: "NativeAdmission" swallowed_before: int - generation: int -def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> _BreakerAdmission: - """Reject the call if the breaker is open, else record what its success may later prove.""" - if breaker.is_open(): +def _enter_circuit_breaker(health: RedisHealth | None, name: str) -> _BreakerAdmission | None: + """Reject the call if the breaker is open, else hold the admission its outcome settles.""" + if health is None: + return None + admission: Final = health.breaker.admit() + if admission is None: raise RedisCircuitBreakerOpenError(f"Redis circuit breaker is open — skipping {name}") - return _BreakerAdmission(swallowed_before=_swallowed_redis_failures.get(), generation=breaker.generation) + return _BreakerAdmission(admission=admission, swallowed_before=_swallowed_redis_failures.get()) -def _exit_circuit_breaker(breaker: RedisCircuitBreaker, admission: _BreakerAdmission) -> None: - """Record success only when nothing failed while the call ran and the breaker has not moved since. +@functools.lru_cache(maxsize=1) +def _redis_reply_error_types() -> tuple[type[Exception], type[Exception]]: + from redis.exceptions import ReadOnlyError, ResponseError + + return ResponseError, ReadOnlyError + + +def _is_redis_reply(exc: BaseException) -> bool: + """A reply that proves Redis answers. A replica refusing a write is the old primary mid-failover, so it does not.""" + try: + reply, failover = _redis_reply_error_types() + except ImportError: + return False + return isinstance(exc, reply) and not isinstance(exc, failover) + + +def _fail_circuit_breaker(health: RedisHealth | None, admission: _BreakerAdmission | None, exc: BaseException) -> None: + if health is None or admission is None: + return + if _is_redis_health_failure(exc): + health.failed(admission.admission, exc) + return + if _is_redis_reply(exc): + _exit_circuit_breaker(admission) + + +def _exit_circuit_breaker(admission: _BreakerAdmission | None) -> None: + """Record success only when nothing failed while the call ran. Several Redis methods catch their own connection errors and return a default, so a - method that returned is not on its own proof of a healthy Redis. A success also vouches - only for the breaker state that admitted the call: a call admitted before the breaker - opened, or a probe admitted before a later failure reopened it, finishes knowing nothing - about whether Redis has recovered since, so only the current probe may close the breaker. + method that returned is not on its own proof of a healthy Redis. The admission itself + ignores a success once the breaker has moved since it admitted the call. """ - if _swallowed_redis_failures.get() != admission.swallowed_before: + if admission is None or _swallowed_redis_failures.get() != admission.swallowed_before: return - if breaker.generation != admission.generation: - return - breaker.record_success() + admission.admission.succeeded() async def _run_under_circuit_breaker( - breaker: RedisCircuitBreaker, + health: RedisHealth | None, name: str, call: Callable[[], Awaitable[_RedisCallResult]], ) -> _RedisCallResult: @@ -596,31 +491,29 @@ async def _run_under_circuit_breaker( Shared by the method decorator and the Lua script executor so both feed the same health signal. """ - admission: Final = _enter_circuit_breaker(breaker, name) + admission: Final = _enter_circuit_breaker(health, name) try: result: Final = await call() except Exception as e: - if _is_redis_health_failure(e): - breaker.record_failure(is_timeout=is_redis_timeout_failure(e)) + _fail_circuit_breaker(health, admission, e) raise - _exit_circuit_breaker(breaker, admission) + _exit_circuit_breaker(admission) return result def _run_under_circuit_breaker_sync( - breaker: RedisCircuitBreaker, + health: RedisHealth | None, name: str, call: Callable[[], _RedisCallResult], ) -> _RedisCallResult: """Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path.""" - admission: Final = _enter_circuit_breaker(breaker, name) + admission: Final = _enter_circuit_breaker(health, name) try: result: Final = call() except Exception as e: - if _is_redis_health_failure(e): - breaker.record_failure(is_timeout=is_redis_timeout_failure(e)) + _fail_circuit_breaker(health, admission, e) raise - _exit_circuit_breaker(breaker, admission) + _exit_circuit_breaker(admission) return result @@ -643,7 +536,7 @@ def _redis_circuit_breaker_guard(method): return await method(self, *args, **kwargs) return await _run_under_circuit_breaker( - self._circuit_breaker, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator + self._health, # pyright: ignore[reportPrivateUsage] # the guard is RedisCache's own decorator wrapper.__name__, lambda: self._admit_async_call(guarded), # pyright: ignore[reportPrivateUsage] # same decorator ) @@ -654,7 +547,7 @@ def _redis_circuit_breaker_guard(method): def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) -> Callable[..., _RedisCallResult]: return functools.wraps(method)( lambda self, *args, **kwargs: _run_under_circuit_breaker_sync( - self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs) + self._health, method.__name__, lambda: method(self, *args, **kwargs) ) ) @@ -717,12 +610,17 @@ class RedisCache(BaseCache): except Exception: pass - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, - recovery_timeout=REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT, - enabled=REDIS_CIRCUIT_BREAKER_ENABLED, - service_logger=self.service_logger_obj, + from litellm.rust_bridge._native import NativeRedisBreaker + + store_kwargs: Final = MappingProxyType(merge_redis_kwargs(redis_kwargs)) + breaker_config: Final = breaker_settings() + self._health: Final = ( + None + if breaker_config is None + else RedisHealth(NativeRedisBreaker.shared(store_kwargs, breaker_config), self.service_logger_obj) ) + if self._health is not None: + self._health.service_logger.redis_breaker_hook(transitions_to=None, failure_class=None) self._setup_health_pings() @@ -731,10 +629,10 @@ class RedisCache(BaseCache): else: super().__init__() # defaults to 60s self.store_settings = RedisStoreSettings( - redis_kwargs=MappingProxyType(merge_redis_kwargs(redis_kwargs)), + redis_kwargs=store_kwargs, namespace=namespace, default_ttl=self.default_ttl, - breaker=breaker_settings(), + breaker=breaker_config, ) def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]: @@ -1123,9 +1021,7 @@ class RedisCache(BaseCache): litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor) return await executor(keys=keys, args=args, client=client) - return await _run_under_circuit_breaker( - self._circuit_breaker, "run_script", lambda: self._admit_async_call(execute) - ) + return await _run_under_circuit_breaker(self._health, "run_script", lambda: self._admit_async_call(execute)) return run_script @@ -1246,7 +1142,7 @@ class RedisCache(BaseCache): log_redis_failure( verbose_logger, logging.ERROR, "LiteLLM Redis Caching: async set() - Got exception from REDIS", e ) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) async def _pipeline_helper( self, @@ -1340,7 +1236,7 @@ class RedisCache(BaseCache): "LiteLLM Redis Caching: async set_cache_pipeline() - Got exception from REDIS", e, ) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) @_redis_circuit_breaker_guard async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: @@ -1385,7 +1281,7 @@ class RedisCache(BaseCache): "LiteLLM Redis Caching: async_set_cache_pipeline_with_ttls() - Got exception from REDIS", e, ) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) async def _set_cache_sadd_helper( self, @@ -1472,7 +1368,7 @@ class RedisCache(BaseCache): "LiteLLM Redis Caching: async set_cache_sadd() - Got exception from REDIS", e, ) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) @_redis_circuit_breaker_guard async def batch_cache_write(self, key, value, **kwargs): @@ -1688,7 +1584,7 @@ class RedisCache(BaseCache): log_redis_failure( verbose_logger, logging.ERROR, "litellm.caching.caching: get() - Got exception from REDIS", e ) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ @@ -1725,12 +1621,12 @@ class RedisCache(BaseCache): key_value_dict = {} _key_list: Final = [key for key in key_list if key is not None] start_time: Final = time.time() - admission: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache") + admission: Final = _enter_circuit_breaker(self._health, "batch_get_cache") try: _keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list] results: Final = self._run_redis_mget_operation(keys=_keys) - _exit_circuit_breaker(self._circuit_breaker, admission) + _exit_circuit_breaker(admission) end_time: Final = time.time() _duration: Final = end_time - start_time self.service_logger_obj.service_success_hook( @@ -1768,7 +1664,7 @@ class RedisCache(BaseCache): parent_otel_span=parent_otel_span, ) log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in batch get cache", e) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) return key_value_dict @_redis_circuit_breaker_guard @@ -1819,7 +1715,7 @@ class RedisCache(BaseCache): print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}") if raise_on_error: raise - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) @_redis_circuit_breaker_guard async def async_batch_get_cache( @@ -1893,7 +1789,7 @@ class RedisCache(BaseCache): ) ) log_redis_failure(verbose_logger, logging.ERROR, "Error occurred in async batch get cache", e) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) return key_value_dict def sync_ping(self) -> bool: @@ -2167,7 +2063,7 @@ class RedisCache(BaseCache): return ttl except Exception as e: verbose_logger.debug("Redis TTL Error: %s", e) - _record_swallowed_redis_failure(self._circuit_breaker, e) + _record_swallowed_redis_failure(self._health, e) return None @_redis_circuit_breaker_guard diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index 42a3ffac855..0e0143ace3b 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -142,10 +142,10 @@ class PrometheusServicesLogger: return metrics def _refresh_redis_breaker_gauge(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts + from litellm.caching.redis_cache import redis_breaker_state_counts counts: Final = redis_breaker_state_counts() - for state in (RedisCircuitBreaker.CLOSED, RedisCircuitBreaker.OPEN, RedisCircuitBreaker.HALF_OPEN): + for state in ("closed", "open", "half_open"): self.redis_breaker_state.labels(state).set(counts.get(state, 0)) def redis_breaker_hook(self, transitions_to: str | None, failure_class: str | None) -> None: diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 773985ad427..77af87f5eb3 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -6,7 +6,7 @@ import httpx from pydantic import JsonValue from typing_extensions import disjoint_base -from litellm.redis.store_settings import RedisStoreSettings +from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings from litellm.rust_bridge.execution import Execution from litellm.rust_bridge.public_call import NativeCall @@ -38,6 +38,22 @@ class _NativeCooldown(TypedDict): timestamp: float cooldown_time: float +@final +class NativeAdmission: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + def succeeded(self) -> None: ... + def failed(self, timeout: bool) -> None: ... + +@final +class NativeRedisBreaker: + def __new__(cls, settings: RedisBreakerSettings) -> NativeRedisBreaker: ... + @staticmethod + def shared(redis_kwargs: Mapping[str, object], settings: RedisBreakerSettings) -> NativeRedisBreaker: ... + def admit(self) -> NativeAdmission | None: ... + def failed(self, timeout: bool) -> None: ... + @property + def state(self) -> Literal["closed", "open", "half_open"]: ... + @final class NativeRedisClient: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @@ -563,6 +579,7 @@ __all__ = [ "GuardedReply", "HuggingFaceEncoding", "LocalSizes", + "NativeAdmission", "NativeAuthCache", "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", @@ -570,6 +587,7 @@ __all__ = [ "NativeDiagnosticProcessor", "NativePending", "NativeRateLimits", + "NativeRedisBreaker", "NativeRedisClient", "NativeRedisStoreError", "NativeRequestRedis", diff --git a/litellm/rust_bridge/host/redis_events.py b/litellm/rust_bridge/host/redis_events.py index 30116a1e15b..ff94378be32 100644 --- a/litellm/rust_bridge/host/redis_events.py +++ b/litellm/rust_bridge/host/redis_events.py @@ -121,4 +121,13 @@ def record_redis_call( def record_redis_breaker(state: BreakerState) -> None: + match state: + case "open": + verbose_logger.warning("Redis circuit breaker OPENED, fast-failing Redis calls") + case "half_open": + verbose_logger.info("Redis circuit breaker HALF_OPEN, probing Redis") + case "closed": + verbose_logger.info("Redis circuit breaker CLOSED, Redis recovered") + case _: + assert_never(state) _service_logging().redis_breaker_hook(transitions_to=state, failure_class=None) diff --git a/tests/unit/caching/test_dual_cache.py b/tests/unit/caching/test_dual_cache.py index 0a9eee6630a..bba0c215015 100644 --- a/tests/unit/caching/test_dual_cache.py +++ b/tests/unit/caching/test_dual_cache.py @@ -368,30 +368,13 @@ async def test_dual_cache_sync_and_async_set_cache_use_same_ttl(): assert abs(sync_expiry - async_expiry) < 1.0 -def test_circuit_breaker_opens_after_threshold(): - """Circuit opens after N consecutive Redis failures.""" - from litellm.caching.redis_cache import RedisCircuitBreaker - - cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - cb.record_failure() - - assert cb._state == "open" - - @pytest.mark.asyncio async def test_circuit_breaker_open_skips_redis(): """When circuit is open, the guard decorator raises immediately without calling the method.""" - from litellm.caching.redis_cache import ( - RedisCircuitBreaker, - _redis_circuit_breaker_guard, - ) - class FakeRedis: + class FakeRedis(_OpenBreakerRedis): def __init__(self): - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - self._circuit_breaker._state = "open" - self._circuit_breaker._opened_at = time.time() + super().__init__() self.call_count = 0 @_redis_circuit_breaker_guard @@ -406,87 +389,13 @@ async def test_circuit_breaker_open_skips_redis(): assert fr.call_count == 0 # method body never executed -def test_circuit_breaker_closes_on_recovery(): - """After recovery_timeout expires, probe is allowed and success closes the circuit.""" - from litellm.caching.redis_cache import RedisCircuitBreaker - - cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - cb._state = "open" - cb._opened_at = time.time() - 9999 # recovery timeout long expired - - # is_open() should return False to allow a probe through, and transition to HALF_OPEN - assert cb.is_open() is False - assert cb._state == "half_open" - - # Successful probe closes the circuit - cb.record_success() - assert cb._state == "closed" - - -def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): - """ - Regression test: only ONE probe gets through when the circuit transitions - OPEN → HALF_OPEN. All concurrent callers that check is_open() while the - state is already HALF_OPEN must be fast-failed (return True), not allowed - through as additional probes. - """ - from litellm.caching.redis_cache import RedisCircuitBreaker - - cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - cb._state = "open" - cb._opened_at = time.time() - 9999 # recovery timeout long expired - - # First caller: OPEN + expired → transitions to HALF_OPEN, returns False (probe) - assert cb.is_open() is False - assert cb._state == "half_open" - - # All subsequent concurrent callers: HALF_OPEN → fast-fail (return True) - for _ in range(10): - assert cb.is_open() is True, "concurrent callers should be fast-failed in HALF_OPEN" - - -def test_circuit_breaker_disabled_never_opens(): - """When disabled, failures never open the circuit and is_open() stays False.""" - from litellm.caching.redis_cache import RedisCircuitBreaker - - cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, enabled=False) - - for _ in range(100): - cb.record_failure() - - assert cb._state == "closed" - assert cb.is_open() is False - - -def test_circuit_breaker_disabled_record_success_leaves_state_untouched(): - """ - A disabled breaker must not mutate state in any state-machine method. Force - a non-default (OPEN) state and assert record_success() returns without - resetting it — the same enabled-guard contract as is_open/record_failure. - """ - from litellm.caching.redis_cache import RedisCircuitBreaker - - cb = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, enabled=False) - cb._state = "open" - cb._failure_count = 3 - - cb.record_success() - - assert cb._state == "open" - assert cb._failure_count == 3 - - @pytest.mark.asyncio async def test_circuit_breaker_disabled_guard_always_calls_method(): """A disabled breaker lets every guarded call through, even after failures.""" - from litellm.caching.redis_cache import ( - RedisCircuitBreaker, - _redis_circuit_breaker_guard, - ) class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) + self._health = None self.call_count = 0 def _admit_async_call(self, call): @@ -504,7 +413,6 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): # Every call reached the method body; the breaker never short-circuited. assert fr.call_count == 5 - assert fr._circuit_breaker.is_open() is False @pytest.mark.asyncio @@ -598,11 +506,16 @@ async def test_dual_cache_late_attach_redis_wires_writes_and_ttl_async(): class _OpenBreakerRedis: def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker + from litellm._service_logger import ServiceLogging + from litellm.caching.redis_cache import RedisHealth + from litellm.redis.store_settings import RedisBreakerSettings + from litellm.rust_bridge._native import NativeRedisBreaker - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() + breaker: Final = NativeRedisBreaker( + RedisBreakerSettings(failure_threshold=1, recovery_timeout=60.0, timeout_min_duration=5.0) + ) + breaker.failed(False) + self._health = RedisHealth(breaker, ServiceLogging()) @_redis_circuit_breaker_guard async def async_get_cache(self, key, **kwargs): diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index f0fbadd36ed..536d510f6d6 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -5,10 +5,11 @@ from __future__ import annotations import asyncio import hashlib import json +import uuid from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass, field from datetime import timedelta -from typing import Any, Generic, TypeVar +from typing import Any, Final, Generic, TypeVar import pytest from redis.exceptions import NoScriptError @@ -23,11 +24,12 @@ from litellm.caching.redis_batch import ( ) from litellm.caching.redis_cache import ( RedisCache, - RedisCircuitBreaker, + RedisHealth, _get_call_stack_info, # pyright: ignore[reportPrivateUsage] # the chain the service hook reports ) from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.redis.store_settings import RedisStoreSettings, breaker_settings +from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings +from litellm.rust_bridge._native import NativeRedisBreaker SCRIPT = "return redis.call('GET', KEYS[1])" _Answer = TypeVar("_Answer") @@ -93,15 +95,22 @@ class FakeClient: class FakeRedisCache(RedisCache): - def __init__(self, client: FakeClient, namespace: str | None = None) -> None: # super().__init__ needs a server + def __init__( + self, client: FakeClient, namespace: str | None = None, failure_threshold: int = 5 + ) -> None: # super().__init__ needs a server self.client = client self.namespace = namespace - self.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2, "db": id(self) % 1_000_000} - self.store_settings = RedisStoreSettings( - redis_kwargs=self.redis_kwargs, namespace=namespace, default_ttl=60, breaker=breaker_settings() - ) - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30) + self.redis_kwargs = {"host": f"batch-{uuid.uuid4().hex}.invalid", "port": 1, "socket_timeout": 0.2} self.service_logger_obj = ServiceLogging() + breaker: Final = RedisBreakerSettings( + failure_threshold=failure_threshold, recovery_timeout=30.0, timeout_min_duration=5.0 + ) + self._health = RedisHealth( # pyright: ignore[reportGeneralTypeIssues] # the fake skips RedisCache.__init__, which sets it + NativeRedisBreaker.shared(self.redis_kwargs, breaker), self.service_logger_obj + ) + self.store_settings = RedisStoreSettings( + redis_kwargs=self.redis_kwargs, namespace=namespace, default_ttl=60, breaker=breaker + ) self.default_ttl = None self.alone: list[tuple[str, Any]] = [] self.store: dict[str, Any] = {} @@ -211,9 +220,11 @@ def fake_pending(send: object, answer: Callable[[], _Answer]) -> FakeNativePendi return FakeNativePending(send, answer) if isinstance(send, FakeNativeRequest) else ReadyPending(answer) -def make(fail: Exception | None = None, namespace: str | None = None) -> tuple[FakeRedisCache, FakeClient]: +def make( + fail: Exception | None = None, namespace: str | None = None, failure_threshold: int = 5 +) -> tuple[FakeRedisCache, FakeClient]: client = FakeClient(replies, fail) - return FakeRedisCache(client, namespace), client + return FakeRedisCache(client, namespace, failure_threshold), client async def run_alone_script(keys: Sequence[str], args: Sequence[Any]) -> object: @@ -324,7 +335,7 @@ async def test_a_reply_an_operation_cannot_decode_fails_only_that_operation() -> @pytest.mark.asyncio async def test_pipeline_failure_fails_every_operation_and_trips_the_breaker() -> None: - cache, _client = make(fail=ConnectionError("redis down")) + cache, _client = make(fail=ConnectionError("redis down"), failure_threshold=1) batch = RedisBatch(cache) got = batch.mget(["a"]) incr = batch.increment("cnt", 1) @@ -332,7 +343,7 @@ async def test_pipeline_failure_fails_every_operation_and_trips_the_breaker() -> await got with pytest.raises(ConnectionError): await incr - assert cache._circuit_breaker._failure_count == 1 # pyright: ignore[reportPrivateUsage] + assert cache._health.breaker.state == "open" # pyright: ignore[reportPrivateUsage] # the breaker the pipeline fed @pytest.mark.asyncio diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 0fbeb05e154..d98b71f2ce3 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -2,7 +2,8 @@ import asyncio import gc import time import types -from collections.abc import Iterator +import uuid +from collections.abc import Callable, Iterator from datetime import timedelta from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -11,7 +12,9 @@ import pytest import litellm from litellm._service_logger import ServiceLogging -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError, RedisHealth +from litellm.redis.store_settings import RedisBreakerSettings +from litellm.rust_bridge._native import NativeRedisBreaker @pytest.fixture @@ -28,7 +31,7 @@ def sync_batch_redis_cache(redis_no_ping): with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point "litellm._redis.get_redis_client", return_value=MagicMock() ) as get_client: - cache = RedisCache(host="127.0.0.1", port=6379) + cache = RedisCache(host="127.0.0.1", port=_closed_port()) cache.redis_client.mget.side_effect = OSError("redis unavailable") get_client.assert_called_once() yield cache @@ -501,7 +504,8 @@ def test_sync_get_cache_failure_feeds_the_breaker_and_logs_a_well_formed_record( assert all("redis unavailable" in record.getMessage() for record in caplog.records) assert len(caplog.records) == REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - assert sync_batch_redis_cache._circuit_breaker.is_open() is True + assert sync_batch_redis_cache._health is not None + assert sync_batch_redis_cache._health.breaker.state == "open" with pytest.raises(RedisCircuitBreakerOpenError): sync_batch_redis_cache.get_cache("lit7468") @@ -527,7 +531,7 @@ async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point "litellm._redis.get_redis_client", return_value=MagicMock() ): - cache = RedisCache(host="127.0.0.1", port=6379) + cache = RedisCache(host="127.0.0.1", port=_closed_port()) with patch.object(cache, "init_async_client", return_value=failing_client): assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} @@ -544,7 +548,7 @@ def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(sto with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point "litellm._redis.get_redis_client", return_value=MagicMock() ): - cache = RedisCache(host="127.0.0.1", port=6379) + cache = RedisCache(host="127.0.0.1", port=_closed_port()) cache.redis_client.mget.return_value = [stored, None, b"0"] assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) @@ -558,7 +562,7 @@ def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[ with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point "litellm._redis.get_redis_client", return_value=failing_client ): - cache = RedisCache(host="127.0.0.1", port=6379, service_logger_obj=service_logger) + cache = RedisCache(host="127.0.0.1", port=_closed_port(), service_logger_obj=service_logger) yield cache, service_logger @@ -646,6 +650,25 @@ def test_sync_batch_get_cache_survives_a_service_callback_that_raises( cache.batch_get_cache(key_list=["lit6729"]) +_THRESHOLD: Final = 3 + + +def _health( + *, failure_threshold: int = _THRESHOLD, recovery_timeout: float = 60.0, timeout_min_duration: float = 5.0 +) -> RedisHealth: + """A breaker no other test shares, so one test opening it cannot fast-fail another.""" + return RedisHealth( + NativeRedisBreaker( + RedisBreakerSettings( + failure_threshold=failure_threshold, + recovery_timeout=recovery_timeout, + timeout_min_duration=timeout_min_duration, + ) + ), + ServiceLogging(), + ) + + def test_call_stack_info_skips_breaker_guard_frames(): """Guarded methods must still report their real callers in service-log call_type. @@ -653,14 +676,10 @@ def test_call_stack_info_skips_breaker_guard_frames(): without skipping them every guarded method logged the guard machinery instead of who actually issued the Redis call. """ - from litellm.caching.redis_cache import ( - RedisCircuitBreaker, - _get_call_stack_info, - _redis_circuit_breaker_guard_sync, - ) + from litellm.caching.redis_cache import _get_call_stack_info, _redis_circuit_breaker_guard_sync class Guarded: - _circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + _health = _health() @_redis_circuit_breaker_guard_sync def probe(self): @@ -675,7 +694,7 @@ def test_call_stack_info_skips_breaker_guard_frames(): assert caller_two() == "caller_one <- caller_two" -def test_call_stack_info_skips_guard_frames_when_deployed_without_sources(monkeypatch): +def test_call_stack_info_skips_guard_frames_when_deployed_without_sources(monkeypatch: pytest.MonkeyPatch): """Guard-frame skipping must survive a bytecode-only deployment. Shipping `.pyc` files without their `.py` sources leaves the module's `__file__` pointing @@ -684,16 +703,12 @@ def test_call_stack_info_skips_guard_frames_when_deployed_without_sources(monkey machinery instead of the real caller. """ from litellm.caching import redis_cache as redis_cache_module - from litellm.caching.redis_cache import ( - RedisCircuitBreaker, - _get_call_stack_info, - _redis_circuit_breaker_guard_sync, - ) + from litellm.caching.redis_cache import _get_call_stack_info, _redis_circuit_breaker_guard_sync monkeypatch.setattr(redis_cache_module, "__file__", redis_cache_module.__file__ + "c") class Guarded: - _circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + _health = _health() @_redis_circuit_breaker_guard_sync def probe(self): @@ -719,10 +734,11 @@ async def test_circuit_breaker_success_still_resets_the_failure_streak(): from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD cache = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) + assert cache._health is not None for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - 1): await cache.async_get_cache("lit4930") - assert cache._circuit_breaker.is_open() is False + assert cache._health.breaker.state == "closed" reachable_redis = AsyncMock() reachable_redis.get.return_value = None @@ -732,7 +748,7 @@ async def test_circuit_breaker_success_still_resets_the_failure_streak(): for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - 1): await cache.async_get_cache("lit4930") - assert cache._circuit_breaker.is_open() is False, "one success must clear the streak" + assert cache._health.breaker.state == "closed", "one success must clear the streak" @pytest.mark.asyncio @@ -770,13 +786,9 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): """ from redis.exceptions import ConnectionError as RedisConnectionError - from litellm.caching.redis_cache import ( - RedisCircuitBreaker, - _record_swallowed_redis_failure, - _run_under_circuit_breaker, - ) + from litellm.caching.redis_cache import _record_swallowed_redis_failure, _run_under_circuit_breaker - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + health = _health() # The failure has to land after both calls are already in flight, which is the only # ordering where a shared counter confuses the two. Failing before the healthy call @@ -787,7 +799,7 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): async def swallows_a_failure() -> None: await asyncio.wait_for(healthy_call_in_flight.wait(), timeout=5) - _record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable")) + _record_swallowed_redis_failure(health, RedisConnectionError("redis unreachable")) failure_recorded.set() async def succeeds_while_the_other_fails() -> str: @@ -796,16 +808,14 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure(): return "ok" await asyncio.gather( - _run_under_circuit_breaker(breaker, "failing", swallows_a_failure), - _run_under_circuit_breaker(breaker, "healthy", succeeds_while_the_other_fails), + _run_under_circuit_breaker(health, "failing", swallows_a_failure), + _run_under_circuit_breaker(health, "healthy", succeeds_while_the_other_fails), ) - rounds = breaker.failure_threshold + 1 - for _ in range(rounds): + for _ in range(_THRESHOLD + 1): await run_one_round() - assert breaker._failure_count < breaker.failure_threshold, "the healthy call must clear the streak" - assert breaker.is_open() is False, "a Redis answering every round must stay in the pool" + assert health.breaker.state == "closed", "a Redis answering every round must stay in the pool" @pytest.mark.asyncio @@ -832,22 +842,94 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker) """ import redis.exceptions - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from litellm.caching.redis_cache import _run_under_circuit_breaker - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + health = _health() raised = getattr(redis.exceptions, error)("boom") async def failing_call(): raise raised - for _ in range(breaker.failure_threshold): + for _ in range(_THRESHOLD): with pytest.raises(type(raised)): - await _run_under_circuit_breaker(breaker, "op", failing_call) + await _run_under_circuit_breaker(health, "op", failing_call) with pytest.raises(Exception, match="circuit breaker is open" if opens_breaker else "boom"): - await _run_under_circuit_breaker(breaker, "op", failing_call) + await _run_under_circuit_breaker(health, "op", failing_call) - assert breaker.is_open() is opens_breaker + assert health.breaker.state == ("open" if opens_breaker else "closed") + + +@pytest.mark.asyncio +async def test_a_recovery_probe_redis_rejects_still_closes_the_breaker(): + """A command error is Redis answering, so a probe that gets one proves Redis is back. + + Leaving that probe unsettled counted it as a failure and reopened the breaker for another + recovery timeout, so every Redis call kept fast-failing against a Redis that was serving. + """ + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import ResponseError + + from litellm.caching.redis_cache import _run_under_circuit_breaker + + health = _health(recovery_timeout=0.05) + + async def failing(error: Exception) -> None: + raise error + + for _ in range(_THRESHOLD): + with pytest.raises(RedisConnectionError): + await _run_under_circuit_breaker(health, "op", lambda: failing(RedisConnectionError("refused"))) + assert health.breaker.state == "open" + await asyncio.sleep(0.06) + + with pytest.raises(ResponseError): + await _run_under_circuit_breaker(health, "op", lambda: failing(ResponseError("WRONGTYPE"))) + + assert health.breaker.state == "closed" + + +@pytest.mark.asyncio +async def test_a_replica_refusing_a_write_does_not_reset_a_run_of_failures(): + """During a failover the old primary answers writes with READONLY between connection failures. + + Counting that reply as a success reset the streak, so the breaker never opened and every + other call kept waiting out its socket timeout. + """ + from redis.exceptions import ConnectionError as RedisConnectionError + from redis.exceptions import ReadOnlyError + + from litellm.caching.redis_cache import _run_under_circuit_breaker + + health = _health() + + async def failing(error: Exception) -> None: + raise error + + for _ in range(_THRESHOLD - 1): + with pytest.raises(RedisConnectionError): + await _run_under_circuit_breaker(health, "op", lambda: failing(RedisConnectionError("refused"))) + with pytest.raises(ReadOnlyError): + await _run_under_circuit_breaker(health, "op", lambda: failing(ReadOnlyError("READONLY"))) + with pytest.raises(RedisConnectionError): + await _run_under_circuit_breaker(health, "op", lambda: failing(RedisConnectionError("refused"))) + + assert health.breaker.state == "open" + + +def test_a_new_redis_cache_shows_its_breaker_in_the_state_gauge(monkeypatch: pytest.MonkeyPatch): + from prometheus_client import REGISTRY + + from litellm.caching.redis_cache import redis_breaker_state_counts + + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + + cache = RedisCache(host=f"gauge-{uuid.uuid4().hex}.invalid", port=6379) + + assert cache.store_settings.breaker is not None + assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( + redis_breaker_state_counts().get("closed", 0) + ) @pytest.mark.asyncio @@ -867,9 +949,9 @@ async def test_event_loop_stall_timeout_burst_keeps_breaker_closed(): from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from litellm.caching.redis_cache import _run_under_circuit_breaker - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + health = _health(timeout_min_duration=5.0) async def healthy_redis_call_with_client_timeout(): deadline = time_mod.monotonic() + 0.05 @@ -883,15 +965,15 @@ async def test_event_loop_stall_timeout_burst_keeps_breaker_closed(): time_mod.sleep(0.2) results = await asyncio.gather( - *(_run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) for _ in range(8)), + *(_run_under_circuit_breaker(health, "op", healthy_redis_call_with_client_timeout) for _ in range(8)), stall_the_loop(), return_exceptions=True, ) timeouts = [r for r in results if isinstance(r, RedisTimeoutError)] - assert len(timeouts) >= breaker.failure_threshold, "the stall must time out a full burst" + assert len(timeouts) >= _THRESHOLD, "the stall must time out a full burst" - assert breaker.is_open() is False, "a healthy Redis behind one loop stall must stay in the pool" - assert await _run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) == "ok" + assert health.breaker.state == "closed", "a healthy Redis behind one loop stall must stay in the pool" + assert await _run_under_circuit_breaker(health, "op", healthy_redis_call_with_client_timeout) == "ok" @pytest.mark.asyncio @@ -903,90 +985,46 @@ async def test_persistent_timeouts_still_open_the_breaker(): """ from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from litellm.caching.redis_cache import _run_under_circuit_breaker - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.1) + health = _health(timeout_min_duration=0.1) async def timing_out_call(): raise RedisTimeoutError("read timed out") - for _ in range(breaker.failure_threshold): + for _ in range(_THRESHOLD): with pytest.raises(RedisTimeoutError): - await _run_under_circuit_breaker(breaker, "op", timing_out_call) - assert breaker.is_open() is False, "the burst has not spanned the minimum duration yet" + await _run_under_circuit_breaker(health, "op", timing_out_call) + assert health.breaker.state == "closed", "the burst has not spanned the minimum duration yet" await asyncio.sleep(0.12) with pytest.raises(RedisTimeoutError): - await _run_under_circuit_breaker(breaker, "op", timing_out_call) + await _run_under_circuit_breaker(health, "op", timing_out_call) - assert breaker.is_open() is True + assert health.breaker.state == "open" + + +async def _eventually(condition: Callable[[], bool]) -> None: + """Breaker transitions reach ServiceLogging from the native runtime, a moment after the call that caused them.""" + deadline: Final = time.monotonic() + 5 + while not condition() and time.monotonic() < deadline: + await asyncio.sleep(0.01) + assert condition() @pytest.mark.asyncio -async def test_stale_timeout_does_not_let_sub_threshold_hard_failures_open_the_breaker(): - """Hard connectivity failures below the threshold must not open the breaker just - because an old timeout already started the streak and the duration has elapsed. - - Each class has to earn the open on its own terms: hard failures by reaching the - threshold, timeouts by reaching the threshold and spanning the minimum duration. - """ - from redis.exceptions import ConnectionError as RedisConnectionError - from redis.exceptions import TimeoutError as RedisTimeoutError - - from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure - - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) - - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) - await asyncio.sleep(0.06) - for _ in range(breaker.failure_threshold - 1): - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) - assert breaker.is_open() is False, "2 hard failures and 1 stale timeout are below both thresholds" - - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) - assert breaker.is_open() is True, "the threshold-th hard failure must still open it" - - -@pytest.mark.asyncio -async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its_own_duration(): - """A stale timeout followed by hard failures must not pre-age the duration gate: - a later short timeout burst has to span timeout_min_duration on its own. - """ - from redis.exceptions import ConnectionError as RedisConnectionError - from redis.exceptions import TimeoutError as RedisTimeoutError - - from litellm.caching.redis_cache import RedisCircuitBreaker, is_redis_timeout_failure - - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05) - - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) - await asyncio.sleep(0.06) - for _ in range(breaker.failure_threshold): - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) - assert breaker.is_open() is False, "the burst is instantaneous, so the duration gate must hold it closed" - - await asyncio.sleep(0.06) - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("read timed out"))) - assert breaker.is_open() is True, "the same run of timeouts persisting past the duration must open it" - - -@pytest.mark.asyncio -async def test_breaker_metrics_track_state_and_failure_class(monkeypatch): +async def test_breaker_metrics_track_state_and_failure_class(monkeypatch: pytest.MonkeyPatch): """Breaker accounting must be observable: failure class, transitions, and state.""" from prometheus_client import REGISTRY from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm._service_logger import ServiceLogging - from litellm.caching.redis_cache import ( - RedisCircuitBreaker, - is_redis_timeout_failure, - redis_breaker_state_counts, - ) + from litellm.caching.redis_cache import _run_under_circuit_breaker, redis_breaker_state_counts + from litellm.rust_bridge.host import redis_events monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) service_logger: Final = ServiceLogging() + monkeypatch.setattr(redis_events, "_service_logging", lambda: service_logger) def sample(name, labels=None): return REGISTRY.get_sample_value(name, labels) or 0.0 @@ -996,39 +1034,49 @@ async def test_breaker_metrics_track_state_and_failure_class(monkeypatch): hard_before: Final = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) opened_before: Final = sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) state_counts_before: Final = redis_breaker_state_counts() - open_gauge_before: Final = state_counts_before.get(RedisCircuitBreaker.OPEN, 0) - closed_gauge_before: Final = state_counts_before.get(RedisCircuitBreaker.CLOSED, 0) + open_gauge_before: Final = state_counts_before.get("open", 0) + closed_gauge_before: Final = state_counts_before.get("closed", 0) - breaker: Final = RedisCircuitBreaker( - failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0, service_logger=service_logger + health: Final = RedisHealth( + NativeRedisBreaker(RedisBreakerSettings(failure_threshold=2, recovery_timeout=0.05, timeout_min_duration=5.0)), + service_logger, ) - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisTimeoutError("t"))) - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) - breaker.record_failure(is_timeout=is_redis_timeout_failure(RedisConnectionError("refused"))) + + async def failing(error: Exception) -> None: + raise error + + for error in (RedisTimeoutError("t"), RedisConnectionError("refused"), RedisConnectionError("refused")): + with pytest.raises(type(error)): + await _run_under_circuit_breaker(health, "op", lambda error=error: failing(error)) assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) == timeout_before + 1 assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) == hard_before + 2 - assert sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) == opened_before + 1 + await _eventually( + lambda: sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) == opened_before + 1 + ) assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + 1 assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before - breaker._opened_at = time.time() - 9999 - assert breaker.is_open() is False - breaker.record_success() + async def succeeding() -> str: + return "ok" + + await asyncio.sleep(0.06) + assert await _run_under_circuit_breaker(health, "op", succeeding) == "ok" + await _eventually( + lambda: sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1 + ) assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before - assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1 -def test_prometheus_services_logger_starts_from_the_breakers_already_built(monkeypatch): +def test_prometheus_services_logger_starts_from_the_breakers_already_built(monkeypatch: pytest.MonkeyPatch): """A breaker built before prometheus_system is on shows up in the state gauge once the services logger starts.""" from prometheus_client import REGISTRY - from litellm._service_logger import ServiceLogging - from litellm.caching.redis_cache import RedisCircuitBreaker, redis_breaker_state_counts + from litellm.caching.redis_cache import redis_breaker_state_counts monkeypatch.setattr(litellm, "service_callback", []) - early: Final = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, service_logger=ServiceLogging()) - assert early.is_open() is False + early: Final = _health() + assert early.breaker.state == "closed" monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) gc.collect() @@ -1041,55 +1089,18 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke def test_sync_guard_counts_a_timeout_as_a_timeout(): from redis.exceptions import TimeoutError as RedisTimeoutError - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker_sync + from litellm.caching.redis_cache import _run_under_circuit_breaker_sync - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + health = _health(timeout_min_duration=5.0) def timing_out_call() -> str: raise RedisTimeoutError("read timed out") for _ in range(6): with pytest.raises(RedisTimeoutError): - _run_under_circuit_breaker_sync(breaker, "op", timing_out_call) + _run_under_circuit_breaker_sync(health, "op", timing_out_call) - assert breaker.is_open() is False - - -def test_success_admitted_before_the_breaker_opened_cannot_close_it(): - """A stale in-flight success must not close a breaker that opened while it ran. - - Calls admitted while the breaker was still closed finish after later failures opened it. - Recording their success unconditionally closed the breaker again, skipping the recovery - timeout and the single half-open probe, so the breaker flapped between open and closed - on every straggler while Redis was still down. - """ - from litellm.caching.redis_cache import RedisCircuitBreaker - - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - breaker.record_failure() - assert breaker._state == breaker.OPEN - - breaker.record_success() - - assert breaker._state == breaker.OPEN - assert breaker.is_open() is True - - -def test_recovery_probe_still_closes_the_breaker(): - from litellm.caching.redis_cache import RedisCircuitBreaker - - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - breaker.record_failure() - breaker._opened_at = time.time() - 9999 - assert breaker.is_open() is False - assert breaker._state == breaker.HALF_OPEN - - breaker.record_success() - - assert breaker._state == breaker.CLOSED - assert breaker.is_open() is False + assert health.breaker.state == "closed" @pytest.mark.asyncio @@ -1100,9 +1111,9 @@ async def test_stale_success_during_the_recovery_probe_leaves_the_breaker_to_the outage, so closing on the straggler's success resumed full Redis traffic before the probe had proven anything. """ - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from litellm.caching.redis_cache import _run_under_circuit_breaker - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + health = _health(recovery_timeout=0.05) stale_admitted = asyncio.Event() stale_release = asyncio.Event() probe_admitted = asyncio.Event() @@ -1118,77 +1129,94 @@ async def test_stale_success_during_the_recovery_probe_leaves_the_breaker_to_the await probe_release.wait() return "probe" - stale = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", stale_call)) + stale = asyncio.ensure_future(_run_under_circuit_breaker(health, "op", stale_call)) await stale_admitted.wait() - for _ in range(3): - breaker.record_failure() - assert breaker._state == breaker.OPEN - breaker._opened_at = time.time() - 9999 - probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", probe_call)) + for _ in range(_THRESHOLD): + health.breaker.failed(False) + assert health.breaker.state == "open" + await asyncio.sleep(0.06) + probe = asyncio.ensure_future(_run_under_circuit_breaker(health, "op", probe_call)) await probe_admitted.wait() - assert breaker._state == breaker.HALF_OPEN + assert health.breaker.state == "half_open" stale_release.set() assert await stale == "stale" - assert breaker._state == breaker.HALF_OPEN, "the straggler must not close the breaker for the probe" - assert breaker.is_open() is True + assert health.breaker.state == "half_open", "the straggler must not close the breaker for the probe" + assert health.breaker.admit() is None probe_release.set() assert await probe == "probe" - assert breaker._state == breaker.CLOSED - assert breaker.is_open() is False + assert health.breaker.state == "closed" @pytest.mark.asyncio -async def test_a_probe_overtaken_by_a_later_outage_leaves_the_breaker_to_the_new_probe(): - """A probe still in flight when a late failure reopens the breaker must not close it for the next probe. +async def test_redis_failures_through_the_cache_fast_fail_the_native_store_on_the_same_redis(): + """A Redis the cache found unreachable must fast-fail the native stores on it too, without another socket wait. - Once the breaker has reopened, only the probe admitted after that outage has reached - Redis, so the older probe's success no longer says anything about whether Redis recovered. + The cache and the native client each kept a breaker, so a management call that wrote a spend counter + through the native store after the cache's breaker had opened still waited a full socket timeout. """ - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from redis.exceptions import ConnectionError as RedisConnectionError - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - old_probe_admitted = asyncio.Event() - old_probe_release = asyncio.Event() - new_probe_admitted = asyncio.Event() - new_probe_release = asyncio.Event() + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + from litellm.redis.stores.connection import connect - async def old_probe_call() -> str: - old_probe_admitted.set() - await old_probe_release.wait() - return "old probe" + cache: Final = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) + store: Final = connect(cache).spend_counters() + unreachable: Final = AsyncMock() + unreachable.ttl.side_effect = RedisConnectionError("refused") - async def new_probe_call() -> str: - new_probe_admitted.set() - await new_probe_release.wait() - return "new probe" + with patch.object(cache, "init_async_client", return_value=unreachable): + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert await cache.async_get_ttl("lit4930") is None - for _ in range(3): - breaker.record_failure() - breaker._opened_at = time.time() - 9999 - old_probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", old_probe_call)) - await old_probe_admitted.wait() - assert breaker._state == breaker.HALF_OPEN + with pytest.raises(Exception, match="circuit breaker is open"): + await store.set("spend:key:lit4930", 1.0, 60) - breaker.record_failure() - assert breaker._state == breaker.OPEN - breaker._opened_at = time.time() - 9999 - new_probe = asyncio.ensure_future(_run_under_circuit_breaker(breaker, "op", new_probe_call)) - await new_probe_admitted.wait() - assert breaker._state == breaker.HALF_OPEN - old_probe_release.set() - assert await old_probe == "old probe" +@pytest.mark.asyncio +async def test_caches_on_one_redis_share_its_breaker_whatever_their_namespace(): + """Every cache on a Redis must learn its outage together, so the native client that the first one built, + which every later cache on that Redis reuses, fast-fails for all of them.""" + from redis.exceptions import ConnectionError as RedisConnectionError - assert breaker._state == breaker.HALF_OPEN, "the overtaken probe must not close the breaker for the new probe" - assert breaker.is_open() is True + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - new_probe_release.set() - assert await new_probe == "new probe" - assert breaker._state == breaker.CLOSED + port: Final = _closed_port() + first: Final = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=port, namespace="first") + second: Final = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=port, namespace="second") + elsewhere: Final = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), namespace="first") + unreachable: Final = AsyncMock() + unreachable.ttl.side_effect = RedisConnectionError("refused") + + with patch.object(first, "init_async_client", return_value=unreachable): + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + assert await first.async_get_ttl("lit4930") is None + + assert second._health is not None and second._health.breaker.state == "open" + assert elsewhere._health is not None and elsewhere._health.breaker.state == "closed" + + +@pytest.mark.asyncio +async def test_native_store_failures_fast_fail_the_cache_on_the_same_redis(): + """The native client's failures open the breaker the cache admits through, so neither stack relearns an outage.""" + from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + from litellm.redis.stores.connection import connect + + cache: Final = await asyncio.to_thread(RedisCache, host="127.0.0.1", port=_closed_port(), socket_timeout=0.5) + store: Final = connect(cache).spend_counters() + unreached: Final = AsyncMock() + + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD): + with pytest.raises(Exception, match=r"(?i)connect"): + await store.set("spend:key:lit4930", 1.0, 60) + + with patch.object(cache, "init_async_client", return_value=unreached): + with pytest.raises(RedisCircuitBreakerOpenError): + await cache.async_get_ttl("lit4930") + assert unreached.mock_calls == [] def test_timeouts_during_a_blip_log_once_per_interval_not_once_per_call(sync_batch_redis_cache, caplog, monkeypatch): @@ -1316,22 +1344,22 @@ async def test_pool_wait_timeout_is_a_timeout_failure_not_hard_connectivity(): from redis.asyncio import BlockingConnectionPool, Redis from redis.exceptions import ConnectionError as RedisConnectionError - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from litellm.caching.redis_cache import _run_under_circuit_breaker pool = BlockingConnectionPool(connection_class=FakeConnection, server=FakeServer(), max_connections=1, timeout=0.01) client = Redis(connection_pool=pool) - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + health = _health(timeout_min_duration=5.0) busy_connection = await pool.get_connection() try: - for _ in range(breaker.failure_threshold * 2): + for _ in range(_THRESHOLD * 2): with pytest.raises(RedisConnectionError, match="No connection available"): - await _run_under_circuit_breaker(breaker, "op", lambda: client.get("k")) + await _run_under_circuit_breaker(health, "op", lambda: client.get("k")) finally: await pool.release(busy_connection) - assert breaker.is_open() is False, "a busy pool is a timeout gated on duration, not a dead Redis" - assert await _run_under_circuit_breaker(breaker, "op", lambda: client.get("k")) is None + assert health.breaker.state == "closed", "a busy pool is a timeout gated on duration, not a dead Redis" + assert await _run_under_circuit_breaker(health, "op", lambda: client.get("k")) is None await client.aclose() @@ -1347,18 +1375,18 @@ async def test_cluster_pool_exhaustion_is_a_timeout_failure_not_hard_connectivit redis_exceptions.MaxConnectionsError, # pyright: ignore[reportAttributeAccessIssue] # old stubs ) - from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + from litellm.caching.redis_cache import _run_under_circuit_breaker - breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + health = _health(timeout_min_duration=5.0) async def pool_exhausted() -> None: raise max_connections_error() - for _ in range(breaker.failure_threshold * 2): + for _ in range(_THRESHOLD * 2): with pytest.raises(max_connections_error): - await _run_under_circuit_breaker(breaker, "op", pool_exhausted) + await _run_under_circuit_breaker(health, "op", pool_exhausted) - assert breaker.is_open() is False, "an exhausted client-side pool is a timeout, not a dead Redis" + assert health.breaker.state == "closed", "an exhausted client-side pool is a timeout, not a dead Redis" def test_timeout_classification_follows_the_explicit_cause_chain_only(): diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 2f236596960..21772593947 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -30,13 +30,13 @@ def redis_cache_double() -> "MagicMock": from unittest.mock import MagicMock from litellm.caching.redis_cache import RedisCache - from litellm.redis.store_settings import RedisStoreSettings, breaker_settings + from litellm.redis.store_settings import RedisStoreSettings redis: Final = MagicMock(spec=RedisCache) redis.redis_kwargs = {"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2} redis.namespace = None redis.store_settings = RedisStoreSettings( - redis_kwargs=redis.redis_kwargs, namespace=None, default_ttl=60, breaker=breaker_settings() + redis_kwargs=redis.redis_kwargs, namespace=None, default_ttl=60, breaker=None ) return redis diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 966e94fad80..ab8972a65da 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -422,11 +422,16 @@ async def test_no_agent_id_passes(): class _OpenBreakerRedis: def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker + from litellm._service_logger import ServiceLogging + from litellm.caching.redis_cache import RedisHealth + from litellm.redis.store_settings import RedisBreakerSettings + from litellm.rust_bridge._native import NativeRedisBreaker - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() + breaker: Final = NativeRedisBreaker( + RedisBreakerSettings(failure_threshold=1, recovery_timeout=60.0, timeout_min_duration=5.0) + ) + breaker.failed(False) + self._health = RedisHealth(breaker, ServiceLogging()) @_redis_circuit_breaker_guard async def async_get_cache(self, key, **kwargs): diff --git a/tests/unit/proxy/hooks/test_sensitive_data_routing.py b/tests/unit/proxy/hooks/test_sensitive_data_routing.py index a76cba2aed7..4c0a592810c 100644 --- a/tests/unit/proxy/hooks/test_sensitive_data_routing.py +++ b/tests/unit/proxy/hooks/test_sensitive_data_routing.py @@ -29,7 +29,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( from litellm.redis.stores.connection import RedisStores from litellm.redis.stores.sensitive_route_pins import SensitiveRoutePinStore from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings -from litellm.rust_bridge._native import NativeRedisClient +from litellm.rust_bridge._native import NativeRedisBreaker, NativeRedisClient class MockInternalUsageCache: @@ -925,19 +925,21 @@ class TestPreCallHookDeferredRouting: @pytest.mark.asyncio async def test_an_open_circuit_breaker_keeps_session_routing_in_memory_without_a_warning(caplog): + redis_kwargs = {"host": f"open-breaker-{uuid.uuid4()}.invalid", "port": 1, "socket_timeout": 0.2} + breaker = RedisBreakerSettings(failure_threshold=1, recovery_timeout=60.0, timeout_min_duration=5.0) + opened = NativeRedisBreaker.shared(redis_kwargs, breaker) + opened.failed(False) client = NativeRedisClient.connect( - RedisStoreSettings( - redis_kwargs={"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}, - namespace=f"open-breaker-{uuid.uuid4()}", - default_ttl=60, - breaker=RedisBreakerSettings(failure_threshold=1, recovery_timeout=60.0, timeout_min_duration=5.0), - ) + RedisStoreSettings(redis_kwargs=redis_kwargs, namespace=None, default_ttl=60, breaker=breaker) ) handler = PROXY_SensitiveDataRoutingHandler( internal_usage_cache=MockInternalUsageCache(), route_pins=SensitiveRoutePinStore(RedisStores(client).sensitive_route_pins()), ) - await handler.set_session_routing("opening-session", "safe-model") + for _ in range(200): + if any("circuit breaker OPENED" in record.getMessage() for record in caplog.records): + break + await asyncio.sleep(0.01) caplog.clear() with caplog.at_level(logging.DEBUG, logger="LiteLLM"): diff --git a/tests/unit/proxy/test_redis_auth_cache_flag.py b/tests/unit/proxy/test_redis_auth_cache_flag.py index f07b8db6385..d7d141b894b 100644 --- a/tests/unit/proxy/test_redis_auth_cache_flag.py +++ b/tests/unit/proxy/test_redis_auth_cache_flag.py @@ -16,7 +16,7 @@ import litellm.proxy.proxy_server as ps from litellm.caching.caching import RedisCache from litellm.caching.dual_cache import DualCache from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.redis.store_settings import RedisStoreSettings, breaker_settings +from litellm.redis.store_settings import RedisStoreSettings from litellm.redis.stores.spend_counters import SpendCounterStore from tests.unit.redis.stores.auth_cache_double import FakeAuthCache from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters @@ -40,7 +40,7 @@ class _FakeRedisCache(RedisCache): redis_kwargs={"host": "127.0.0.1", "port": 1, "socket_timeout": 0.2}, namespace=None, default_ttl=60, - breaker=breaker_settings(), + breaker=None, ) def set_cache(self, key, value, **kwargs): # type: ignore[override] diff --git a/tests/unit/redis/stores/test_connection.py b/tests/unit/redis/stores/test_connection.py index d896df6a313..bd8dfd68bf5 100644 --- a/tests/unit/redis/stores/test_connection.py +++ b/tests/unit/redis/stores/test_connection.py @@ -1,25 +1,27 @@ import gc +import logging import uuid from concurrent.futures import ThreadPoolExecutor from typing import Final +import pytest + from litellm.caching.redis_cache import RedisCache from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings -from litellm.redis.stores.connection import connect -from litellm.rust_bridge._native import NativeRedisClient, native_redis_breaker_states +from litellm.redis.stores.connection import RedisStores, connect +from litellm.redis.stores.sensitive_route_pins import SensitiveRoutePinStore +from litellm.rust_bridge._native import NativeRedisBreaker, NativeRedisClient, native_redis_breaker_states + +_BREAKER: Final = RedisBreakerSettings(failure_threshold=5, recovery_timeout=60.0, timeout_min_duration=5.0) def _settings(namespace: str) -> RedisStoreSettings: return RedisStoreSettings( - redis_kwargs={"host": "127.0.0.1", "port": 1}, - namespace=namespace, - default_ttl=60, - breaker=RedisBreakerSettings(failure_threshold=5, recovery_timeout=60.0, timeout_min_duration=5.0), + redis_kwargs={"host": "127.0.0.1", "port": 1}, namespace=namespace, default_ttl=60, breaker=None ) -def test_concurrent_first_connects_share_one_client_without_deadlocking() -> None: - before: Final = len(native_redis_breaker_states()) +def test_concurrent_first_connects_finish_without_deadlocking() -> None: namespace: Final = f"concurrent-connect-{uuid.uuid4()}" pool: Final = ThreadPoolExecutor(max_workers=8) futures: Final = tuple(pool.submit(NativeRedisClient.connect, _settings(namespace)) for _ in range(8)) @@ -29,23 +31,39 @@ def test_concurrent_first_connects_share_one_client_without_deadlocking() -> Non pool.shutdown(wait=False) assert all(client is not None for client in clients) + + +def test_concurrent_lookups_of_one_redis_share_one_breaker_without_deadlocking() -> None: + gc.collect() + before: Final = len(native_redis_breaker_states()) + redis_kwargs: Final = {"host": f"shared-breaker-{uuid.uuid4()}.invalid", "port": 6379} + pool: Final = ThreadPoolExecutor(max_workers=8) + futures: Final = tuple(pool.submit(NativeRedisBreaker.shared, redis_kwargs, _BREAKER) for _ in range(8)) + try: + breakers: Final = tuple(future.result(timeout=30) for future in futures) + finally: + pool.shutdown(wait=False) + assert len(native_redis_breaker_states()) == before + 1 + for _ in range(_BREAKER.failure_threshold): + breakers[0].failed(False) + assert {breaker.state for breaker in breakers} == {"open"} -def _breaker_states_while_connected(namespace: str) -> int: - client: Final = NativeRedisClient.connect(_settings(namespace)) - assert client is not None +def _breaker_states_while_held() -> int: + breaker: Final = NativeRedisBreaker(_BREAKER) + assert breaker.state == "closed" return len(native_redis_breaker_states()) -def test_a_discarded_client_leaves_the_process_registry() -> None: +def test_a_discarded_breaker_leaves_the_process_registry() -> None: gc.collect() before: Final = len(native_redis_breaker_states()) - while_connected: Final = _breaker_states_while_connected(f"discarded-{uuid.uuid4()}") + while_held: Final = _breaker_states_while_held() gc.collect() - assert while_connected == before + 1 + assert while_held == before + 1 assert len(native_redis_breaker_states()) == before @@ -75,3 +93,33 @@ def test_a_redis_cache_stand_in_that_skips_the_constructor_connects_nothing() -> def test_no_redis_cache_connects_nothing() -> None: assert connect(None).client is None + + +def test_caches_on_one_server_share_a_breaker_whatever_their_db_or_url() -> None: + host: Final = f"one-server-{uuid.uuid4()}.invalid" + by_host: Final = NativeRedisBreaker.shared({"host": host, "port": 6379, "db": 0}, _BREAKER) + by_url: Final = NativeRedisBreaker.shared({"url": f"redis://:pw@{host.upper()}:6379/1"}, _BREAKER) + other_port: Final = NativeRedisBreaker.shared({"host": host, "port": 6380}, _BREAKER) + + for _ in range(_BREAKER.failure_threshold): + by_host.failed(False) + + assert (by_url.state, other_port.state) == ("open", "closed") + + +@pytest.mark.asyncio +async def test_a_store_admits_through_the_breaker_its_server_opened(caplog: pytest.LogCaptureFixture) -> None: + redis_kwargs: Final = {"host": f"opened-{uuid.uuid4()}.invalid", "port": 6379} + opened: Final = NativeRedisBreaker.shared(redis_kwargs, _BREAKER) + for _ in range(_BREAKER.failure_threshold): + opened.failed(False) + client: Final = NativeRedisClient.connect( + RedisStoreSettings(redis_kwargs={**redis_kwargs, "db": 3}, namespace=None, default_ttl=60, breaker=_BREAKER) + ) + pins: Final = SensitiveRoutePinStore(RedisStores(client).sensitive_route_pins()) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM"): + await pins.set("tenant", "session", "model", 60.0) + + assert opened.state == "open" + assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) diff --git a/tests/unit/redis/stores/test_router_cooldowns.py b/tests/unit/redis/stores/test_router_cooldowns.py index 2405299cb11..f77c111d56d 100644 --- a/tests/unit/redis/stores/test_router_cooldowns.py +++ b/tests/unit/redis/stores/test_router_cooldowns.py @@ -238,8 +238,8 @@ def test_the_native_breaker_counts_in_the_redis_breaker_state_gauge() -> None: store: Final = NativeRouterCooldowns( NativeRedisClient.connect( RedisStoreSettings( - redis_kwargs={"host": "127.0.0.1", "port": 1}, - namespace=f"breaker-gauge-test-{uuid.uuid4()}", + redis_kwargs={"host": f"breaker-gauge-{uuid.uuid4()}.invalid", "port": 1}, + namespace=None, default_ttl=60, breaker=RedisBreakerSettings(failure_threshold=1, recovery_timeout=60.0, timeout_min_duration=5.0), ) 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 2bcc5c5c1d3..a64702cd286 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 @@ -1,5 +1,6 @@ import asyncio import json +import uuid from typing import Final from collections.abc import AsyncIterator, Mapping from datetime import datetime, timezone @@ -15,7 +16,11 @@ from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD, + SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -192,7 +197,7 @@ async def shared_session_router() -> AsyncIterator[tuple[litellm.Router, FakeSer patch("redis.Redis.from_url", return_value=sync_client), patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), ): - redis: Final = RedisCache(url="redis://session-outage.example:6379/0") + redis: Final = RedisCache(url=f"redis://session-outage-{uuid.uuid4().hex}.example:6379/0") router.cache.attach_redis_cache(redis) yield router, server, redis finally: @@ -223,11 +228,11 @@ async def test_redis_read_failures_and_open_breaker_preserve_session_routing( assert await _selected_session_ids(router) == ("replacement",) server.connected = False with patch("time.time", return_value=1790000000.0): - for _ in range(redis._circuit_breaker.failure_threshold + 1): + for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + 1): assert set(await _selected_session_ids(router)) == ( {"replacement"} if has_pin else {"original", "replacement", "other"} ) - assert redis._circuit_breaker.is_open() + assert redis._health is not None and redis._health.breaker.state == "open" async def test_confirmed_redis_miss_does_not_revive_a_local_pin_during_outage(