mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(redis): share one native circuit breaker between RedisCache and the Rust stores (#50473)
* fix(redis): share one native circuit breaker between RedisCache and the Rust stores RedisCache and the native Redis client each kept their own breaker, so once the cache learned the coordination Redis was wedged, a spend-counter write through the native store still paid a full socket timeout. Every Redis target now has one Rust breaker, looked up by NativeRedisBreaker.shared and injected into both the RedisCache guard and the native client. The Python FSM is gone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * style(rust): cargo fmt the shared breaker Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(redis): close the breaker when the recovery probe gets a Redis reply error A command error such as WRONGTYPE means Redis answered, but the probe was left unsettled and its drop reopened the breaker for another recovery timeout. Both the RedisCache guard and the native client now settle such a probe as a success. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(redis): pass a NativeRedisBreaker and isolate the shared breaker per test The sensitive routing test still passed RedisBreakerSettings as the store breaker, and the session affinity fixture reused one Redis URL, so a breaker an earlier test opened leaked into the next. Also make the RUF043 match pattern raw. * fix(redis): key the shared breaker by server and close the three audit gaps The breaker registry now keys on the Redis server (host and port, unix socket, cluster node set or sentinel service) instead of the full connection target, so caches that differ only in db, credentials or url spelling share one breaker. A READONLY reply from a replica mid-failover no longer counts as Redis answering, so it can't reset a run of failures. A new RedisCache refreshes the breaker state gauge, which the deleted Python breaker did at construction. A shared native client is only reused by a store that passes the same breaker. * refactor(redis): split redis-types into modules and drop redundant pyo3 attribute annotations lib.rs now only declares modules and re-exports, with resolution, target, tls, auth and endpoint each in their own file. Derived FromPyObject already reads a struct field as the attribute of the same name, so #[pyo3(attribute)] on the breaker and store settings did nothing. * refactor(redis): drop redundant pyo3 attribute annotations in the stores and let PyO3 build the cooldown snapshot Derived FromPyObject already reads each field as the attribute of the same name. The cooldown snapshot returns an IndexMap, which PyO3 turns into the same dict in the same order. * refactor(redis): return a HashMap from the three memory-tier snapshots The snapshots read cache-memory's HashMap, so there is no order to keep. PyO3 builds the same dict from a std HashMap, which replaces the hand-written PyDict loops in the spend counter and router usage stores and the IndexMap in the cooldown store. The rate-limit slot registry keeps IndexMap because it hands back a Python dict in the order Python wrote it. * refactor(redis): share Redis breakers and clients through one RedisHosts per host redis-connect's RedisHosts now owns the per-server breaker and the shared clients that python-bridge kept in two process-wide statics. A client takes 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. This drops RedisClient::guarded_by and Breaker::same. python-bridge keeps one RedisHosts per process and calls it with the GIL released. The gateway builds one per connect, so its coordination and response cache on the same server now share a breaker. RedisStoreSettings.breaker is back to breaker settings, as on main. --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
parent
210e2a075a
commit
574fcea691
42 changed files with 1487 additions and 937 deletions
2
litellm-rust/Cargo.lock
generated
2
litellm-rust/Cargo.lock
generated
|
|
@ -4558,6 +4558,8 @@ name = "litellm-redis-types"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"litellm-secret-value",
|
||||
"rstest",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -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<Breaker>::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
|
||||
|
|
|
|||
|
|
@ -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<Shared<Config>>);
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct WeakBreaker(Weak<Shared<Config>>);
|
||||
|
||||
#[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<Breaker> {
|
||||
self.0.upgrade().map(Breaker)
|
||||
}
|
||||
}
|
||||
|
||||
impl Admission {
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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<dyn RedisEvents> = Arc::new(TracingEvents::new(lookup));
|
||||
let hosts = RedisHosts::default();
|
||||
let client = |params: Params| -> Result<RedisClient, Error> {
|
||||
Ok(litellm_redis_connect::connect(
|
||||
Ok(hosts.client(
|
||||
params,
|
||||
HostServices {
|
||||
events: Some(events.clone()),
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
134
litellm-rust/crates/python-bridge/src/redis/breaker.rs
Normal file
134
litellm-rust/crates/python-bridge/src/redis/breaker.rs
Normal file
|
|
@ -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<Config> {
|
||||
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<Endpoint>,
|
||||
) -> PyResult<Self> {
|
||||
let config = settings.config()?;
|
||||
let runtime = litellm_host_python::runtime_handle()?;
|
||||
let events: Arc<dyn RedisEvents> =
|
||||
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> {
|
||||
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<Self> {
|
||||
let endpoint = resolved_endpoint(py, &redis_kwargs)?;
|
||||
Self::build(py, &settings, endpoint)
|
||||
}
|
||||
|
||||
fn admit(&self) -> Option<NativeAdmission> {
|
||||
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<Option<Admission>>);
|
||||
|
||||
impl NativeAdmission {
|
||||
fn take(&self) -> Option<Admission> {
|
||||
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()
|
||||
}
|
||||
|
|
@ -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<BreakerConfig> {
|
||||
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<String>,
|
||||
#[pyo3(attribute)]
|
||||
default_ttl: f64,
|
||||
#[pyo3(attribute)]
|
||||
breaker: Option<RedisBreaker>,
|
||||
breaker: Option<BreakerSettings>,
|
||||
}
|
||||
|
||||
fn is_plain(value: &Bound<'_, PyAny>) -> bool {
|
||||
|
|
@ -115,22 +81,6 @@ fn is_plain(value: &Bound<'_, PyAny>) -> bool {
|
|||
|| value.is_instance_of::<PyDict>()
|
||||
}
|
||||
|
||||
/// 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<Mutex<Vec<Shared>>> = LazyLock::new(|| Mutex::new(Vec::new()));
|
||||
|
||||
fn live_clients(py: Python<'_>) -> std::sync::MutexGuard<'static, Vec<Shared>> {
|
||||
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<BreakerConfig>,
|
||||
) -> PyResult<Option<RedisClient>> {
|
||||
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<Arc<dyn CredentialSource>> {
|
||||
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<Option<Endpoint>> {
|
||||
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<Result<Resolution, litellm_redis_resolve::Error>> {
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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<Self> {
|
||||
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> {
|
||||
Self::from_module(py, &py.import(Self::MODULE)?, runtime, None)
|
||||
}
|
||||
|
||||
fn from_module(
|
||||
py: Python<'_>,
|
||||
module: &Bound<'_, PyModule>,
|
||||
runtime: Handle,
|
||||
event_loop: Option<EventLoop>,
|
||||
) -> PyResult<Self> {
|
||||
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())
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<RedisHosts> = 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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<PyScope> for Scope {
|
|||
/// `RateLimitCounter`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyCounter {
|
||||
#[pyo3(attribute)]
|
||||
scope: PyScope,
|
||||
#[pyo3(attribute)]
|
||||
kind: Spelled<Kind>,
|
||||
}
|
||||
|
||||
|
|
@ -100,9 +96,7 @@ impl From<PyCounter> 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<PyEntry> for Entry {
|
|||
/// `CounterIncrement`.
|
||||
#[derive(FromPyObject)]
|
||||
struct PyIncrement {
|
||||
#[pyo3(attribute)]
|
||||
counter: PyCounter,
|
||||
#[pyo3(attribute)]
|
||||
amount: PyNumber,
|
||||
#[pyo3(attribute)]
|
||||
ttl: Option<u64>,
|
||||
}
|
||||
|
||||
|
|
@ -139,15 +130,10 @@ impl From<PyIncrement> 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,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, cooldown) in self.0.snapshot().map_err(kind_error)? {
|
||||
snapshot.set_item(key, PyData(cooldown))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
fn snapshot(&self) -> PyResult<HashMap<String, PyData<Cooldown>>> {
|
||||
Ok(self
|
||||
.0
|
||||
.snapshot()
|
||||
.map_err(kind_error)?
|
||||
.into_iter()
|
||||
.map(|(key, cooldown)| (key, PyData(cooldown)))
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
#[pyo3(attribute)]
|
||||
minute: String,
|
||||
#[pyo3(attribute)]
|
||||
kind: Spelled<Kind>,
|
||||
}
|
||||
|
||||
|
|
@ -174,11 +170,13 @@ impl NativeRouterUsage {
|
|||
self.0.flush_local_counters().map_err(kind_error)
|
||||
}
|
||||
|
||||
fn snapshot<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, count) in self.0.snapshot().map_err(kind_error)? {
|
||||
snapshot.set_item(key, PyNumber(count))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
|
||||
Ok(self
|
||||
.0
|
||||
.snapshot()
|
||||
.map_err(kind_error)?
|
||||
.into_iter()
|
||||
.map(|(key, count)| (key, PyNumber(count)))
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Bound<'py, PyDict>> {
|
||||
let snapshot = PyDict::new(py);
|
||||
for (key, spend) in self.0.snapshot().map_err(kind_error)? {
|
||||
snapshot.set_item(key, PyNumber(spend))?;
|
||||
}
|
||||
Ok(snapshot)
|
||||
fn snapshot(&self) -> PyResult<HashMap<String, PyNumber>> {
|
||||
Ok(self
|
||||
.0
|
||||
.snapshot()
|
||||
.map_err(kind_error)?
|
||||
.into_iter()
|
||||
.map(|(key, spend)| (key, PyNumber(spend)))
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<State> {
|
||||
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<Admission>, 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)
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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::<String>("k")).await;
|
||||
let _ = client.send(commands::get::<String>("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::<String>("k")).await.unwrap();
|
||||
server.freeze().await.unwrap();
|
||||
let _ = client.send(commands::get::<String>("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::<String>("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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Arc<dyn RedisEvents>>) -> 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<BreakerConfig>,
|
||||
breaker: Option<Breaker>,
|
||||
) -> Result<RedisClient, Error> {
|
||||
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,
|
||||
|
|
|
|||
241
litellm-rust/crates/redis-connect/src/hosts.rs
Normal file
241
litellm-rust/crates/redis-connect/src/hosts.rs
Normal file
|
|
@ -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<Vec<(Option<Endpoint>, WeakBreaker)>>,
|
||||
clients: Mutex<Vec<(ClientKey, WeakRedisClient)>>,
|
||||
}
|
||||
|
||||
#[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<Endpoint>,
|
||||
config: BreakerConfig,
|
||||
events: Option<Arc<dyn RedisEvents>>,
|
||||
) -> 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<BreakerConfig>,
|
||||
) -> Result<RedisClient, Error> {
|
||||
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<Breaker> {
|
||||
live(&self.breakers, WeakBreaker::upgrade)
|
||||
.iter()
|
||||
.filter_map(|(_, weak)| weak.upgrade())
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn live<K, W, V>(
|
||||
entries: &Mutex<Vec<(K, W)>>,
|
||||
upgrade: fn(&W) -> Option<V>,
|
||||
) -> 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<Endpoint> {
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -7,3 +7,7 @@ repository.workspace = true
|
|||
|
||||
[dependencies]
|
||||
litellm-secret-value.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
50
litellm-rust/crates/redis-types/src/auth.rs
Normal file
50
litellm-rust/crates/redis-types/src/auth.rs
Normal file
|
|
@ -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<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
pub tenant_id: Option<String>,
|
||||
pub client_secret: Option<SecretValue>,
|
||||
pub username: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct AwsIamAuth {
|
||||
pub user_name: String,
|
||||
pub cache_name: String,
|
||||
pub region: String,
|
||||
pub serverless: bool,
|
||||
}
|
||||
169
litellm-rust/crates/redis-types/src/endpoint.rs
Normal file
169
litellm-rust/crates/redis-types/src/endpoint.rs
Normal file
|
|
@ -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<Endpoint> {
|
||||
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<Endpoint> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<String>,
|
||||
pub timeout: Option<Duration>,
|
||||
pub connect_timeout: Option<Duration>,
|
||||
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<Ignored>,
|
||||
}
|
||||
|
||||
/// 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<TlsOptions>,
|
||||
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<TlsOptions>,
|
||||
/// `?db=`, which wins over the path.
|
||||
pub db: Option<i64>,
|
||||
/// `?username=` and `?password=`, which the url's userinfo wins over.
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Cluster {
|
||||
pub nodes: Vec<Node>,
|
||||
pub tls: Option<TlsOptions>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Node {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Sentinel {
|
||||
pub monitors: Vec<Node>,
|
||||
pub service_name: String,
|
||||
/// How to authenticate to the monitors, which are separate servers from the data node.
|
||||
pub monitor_auth: Option<StaticAuth>,
|
||||
pub tls: Option<TlsOptions>,
|
||||
pub db: i64,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct TlsOptions {
|
||||
pub ca_certs: Option<PathBuf>,
|
||||
/// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data`
|
||||
pub ca_data: Option<String>,
|
||||
/// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path`
|
||||
pub ca_path: Option<PathBuf>,
|
||||
pub certfile: Option<PathBuf>,
|
||||
pub keyfile: Option<PathBuf>,
|
||||
/// Unlocks an encrypted `keyfile`, redis-py's `ssl_password`
|
||||
pub key_password: Option<SecretValue>,
|
||||
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<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
}
|
||||
|
||||
/// 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<String>,
|
||||
pub tenant_id: Option<String>,
|
||||
pub client_secret: Option<SecretValue>,
|
||||
pub username: Option<String>,
|
||||
}
|
||||
|
||||
#[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;
|
||||
|
|
|
|||
36
litellm-rust/crates/redis-types/src/resolution.rs
Normal file
36
litellm-rust/crates/redis-types/src/resolution.rs
Normal file
|
|
@ -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<String>,
|
||||
pub timeout: Option<Duration>,
|
||||
pub connect_timeout: Option<Duration>,
|
||||
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<Ignored>,
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
65
litellm-rust/crates/redis-types/src/target.rs
Normal file
65
litellm-rust/crates/redis-types/src/target.rs
Normal file
|
|
@ -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<TlsOptions>,
|
||||
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<TlsOptions>,
|
||||
/// `?db=`, which wins over the path.
|
||||
pub db: Option<i64>,
|
||||
/// `?username=` and `?password=`, which the url's userinfo wins over.
|
||||
pub username: Option<SecretValue>,
|
||||
pub password: Option<SecretValue>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Cluster {
|
||||
pub nodes: Vec<Node>,
|
||||
pub tls: Option<TlsOptions>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Node {
|
||||
pub host: String,
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Sentinel {
|
||||
pub monitors: Vec<Node>,
|
||||
pub service_name: String,
|
||||
/// How to authenticate to the monitors, which are separate servers from the data node.
|
||||
pub monitor_auth: Option<StaticAuth>,
|
||||
pub tls: Option<TlsOptions>,
|
||||
pub db: i64,
|
||||
}
|
||||
33
litellm-rust/crates/redis-types/src/tls.rs
Normal file
33
litellm-rust/crates/redis-types/src/tls.rs
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
use std::path::PathBuf;
|
||||
|
||||
use litellm_secret_value::SecretValue;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct TlsOptions {
|
||||
pub ca_certs: Option<PathBuf>,
|
||||
/// PEM text trusted alongside `ca_certs`, redis-py's `ssl_ca_data`
|
||||
pub ca_data: Option<String>,
|
||||
/// A directory whose files are PEM certificates to trust, redis-py's `ssl_ca_path`
|
||||
pub ca_path: Option<PathBuf>,
|
||||
pub certfile: Option<PathBuf>,
|
||||
pub keyfile: Option<PathBuf>,
|
||||
/// Unlocks an encrypted `keyfile`, redis-py's `ssl_password`
|
||||
pub key_password: Option<SecretValue>,
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"):
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue