From 0957d84da02d1a036539296bc839ed5b0f199395 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 15:37:43 +0900 Subject: [PATCH] feat(redis): serve the health checks' deployment states from a native store DeploymentHealthCache reads and writes its one Redis entry through a new deployment_health store family instead of the router's DualCache, like the other router families. --- litellm-rust/crates/python-bridge/src/lib.rs | 3 + .../src/redis/store/deployment_health.rs | 61 ++++++++++ .../python-bridge/src/redis/store/mod.rs | 1 + litellm-rust/crates/redis-store/AGENTS.md | 2 +- .../redis-store/src/deployment_health/keys.rs | 2 + .../redis-store/src/deployment_health/mod.rs | 70 ++++++++++++ litellm-rust/crates/redis-store/src/lib.rs | 1 + .../redis-store/tests/deployment_health.rs | 62 +++++++++++ litellm/redis/stores/connection.py | 4 + .../stores/deployment_health/__init__.py | 46 ++++++++ .../redis/stores/deployment_health/handle.py | 16 +++ litellm/router_rust/core.py | 6 +- litellm/router_rust/surface/health_routing.py | 6 +- litellm/router_utils/health_state_cache.py | 55 ++++----- litellm/rust_bridge/_native.pyi | 12 ++ .../router_utils/test_health_state_cache.py | 104 +++++++++--------- .../core/test_parity_rulings_2026_11.py | 2 +- 17 files changed, 357 insertions(+), 96 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/redis/store/deployment_health.rs create mode 100644 litellm-rust/crates/redis-store/src/deployment_health/keys.rs create mode 100644 litellm-rust/crates/redis-store/src/deployment_health/mod.rs create mode 100644 litellm-rust/crates/redis-store/tests/deployment_health.rs create mode 100644 litellm/redis/stores/deployment_health/__init__.py create mode 100644 litellm/redis/stores/deployment_health/handle.py diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index dc127e41129..02872c9b3f1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -44,6 +44,8 @@ mod _native { #[pymodule_export] use crate::redis::store::auth_cache::{NativeAuthCache, PyLocalSizes}; #[pymodule_export] + use crate::redis::store::deployment_health::NativeDeploymentHealth; + #[pymodule_export] use crate::redis::store::rate_limits::{ NativeRateLimits, PyAtomicApplied, PyAtomicOverLimit, PyGaugeOverLimit, PyGaugesAcquired, PyGuardedReply, PyWindowCount, @@ -139,6 +141,7 @@ mod tests { "NativeRequestRedis", "NativeSend", "NativeAuthCache", + "NativeDeploymentHealth", "NativeRateLimits", "WindowCount", "AtomicApplied", diff --git a/litellm-rust/crates/python-bridge/src/redis/store/deployment_health.rs b/litellm-rust/crates/python-bridge/src/redis/store/deployment_health.rs new file mode 100644 index 00000000000..b48a5bf6f62 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/redis/store/deployment_health.rs @@ -0,0 +1,61 @@ +use std::sync::Arc; + +use litellm_redis_store::{Error, deployment_health::DeploymentHealth}; +use pyo3::prelude::*; + +use super::{duration, kind_error, redis_async, redis_sync, remote, value::PyValue}; +use crate::redis::client::NativeRedisClient; + +/// The background health checks' states behind `DeploymentHealthCache`. The `_sync` methods block +/// with the GIL released, for the health check writer and the router's sync calls. +#[pyclass(frozen, name = "NativeDeploymentHealth")] +pub(crate) struct NativeDeploymentHealth(Arc); + +#[pymethods] +impl NativeDeploymentHealth { + #[new] + fn new(client: Option>) -> Self { + Self(Arc::new(DeploymentHealth::new( + client.map(|client| remote(&client)), + ))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn get<'py>(&self, py: Python<'py>) -> PyResult> { + let store = self.0.clone(); + redis_async(py, "async_get_cache", kind_error, async move { + store.get().await.map(|states| states.map(PyValue)) + }) + } + + fn get_sync(&self, py: Python<'_>) -> PyResult> { + let store = self.0.clone(); + redis_sync(py, "get_cache", kind_error, async move { + store.get().await.map(|states| states.map(PyValue)) + }) + } + + fn get_shared_sync(&self, py: Python<'_>) -> PyResult> { + let store = self.0.clone(); + redis_sync(py, "get_cache", kind_error, async move { + Ok::<_, Error>(store.get_shared().await.map(PyValue)) + }) + } + + fn set_sync(&self, py: Python<'_>, states: PyValue, ttl_seconds: f64) -> PyResult<()> { + let store = self.0.clone(); + let ttl = duration(ttl_seconds)?; + redis_sync(py, "set_cache", kind_error, async move { + store.set(states.0, ttl).await + }) + } +} diff --git a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs index d9c3aabb814..35eba7b924f 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod auth_cache; +pub(crate) mod deployment_health; mod policy; pub(crate) mod rate_limits; pub(crate) mod request; diff --git a/litellm-rust/crates/redis-store/AGENTS.md b/litellm-rust/crates/redis-store/AGENTS.md index cdd2bd7ca64..65c3886da05 100644 --- a/litellm-rust/crates/redis-store/AGENTS.md +++ b/litellm-rust/crates/redis-store/AGENTS.md @@ -1,4 +1,4 @@ -- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) +- One folder per Redis family (`router_cooldowns/`, `router_usage/`, `spend_counters/`, `auth_cache/`, `rate_limits/`, `deployment_health/`), each with `mod.rs` (the store and its ops), `keys.rs` (key types and spellings) and `types.rs` (the family's values, with the `FromValue` / `IntoValue` codecs for what it stores in Redis). It is the only place that family's keys, value codec and scripts are written, and it takes domain ids, never keys. Shared pieces live once at the crate root: `Error` and `ErrorKind` (`error.rs`), `Memory` and `Clock` (`memory.rs`), `Pending` (`request.rs`), `Remote` and its attach-once `RemoteSlot`, which knows nothing of memory (`remote.rs`), `DualCache` (`dual_cache/`) and the Python codecs (`codec.rs`, `value.rs`) - An op declared on a `Send` returns `Pending`, a lazy future built when the op is declared. An op that only means something on Redis (an `MGET` the caller falls back from, a script) returns `None` without one, and an op that replaced a `DualCache` call answers from memory. A post-call method returns `false` when there was no Redis to declare on - A store that sends a pipeline of its own uses `RedisBatch::pipelined`, never `Pipeline` and raw `Value`s - Python and Rust pods share these keys during a rolling deploy, so a module writes the bytes `main`'s Python writes and reads everything `main` reads. Python formats go through `litellm-python-compat` (`repr`, `literal_eval`, `json`), never a hand-rolled codec diff --git a/litellm-rust/crates/redis-store/src/deployment_health/keys.rs b/litellm-rust/crates/redis-store/src/deployment_health/keys.rs new file mode 100644 index 00000000000..871e49b3957 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/deployment_health/keys.rs @@ -0,0 +1,2 @@ +/// `DeploymentHealthCache.CACHE_KEY`: one entry holds every deployment's state. +pub(super) const STATES: &str = "litellm:health_check:deployment_health_state"; diff --git a/litellm-rust/crates/redis-store/src/deployment_health/mod.rs b/litellm-rust/crates/redis-store/src/deployment_health/mod.rs new file mode 100644 index 00000000000..152dc3e5797 --- /dev/null +++ b/litellm-rust/crates/redis-store/src/deployment_health/mod.rs @@ -0,0 +1,70 @@ +mod keys; + +use std::{sync::Arc, time::Duration}; + +use litellm_python_compat::Value; +use litellm_redis_client::commands; + +use crate::{ + Error, codec, + dual_cache::{DualCache, MemorySettings}, + memory::{Clock, system_clock}, + remote::{Remote, RemoteSlot}, +}; + +/// The background health checks' per-deployment states, kept by `DeploymentHealthCache` in the +/// router's `DualCache`: memory per worker, Redis shared by every pod. The writer's merge and the +/// readers' staleness rules stay with the host. +pub struct DeploymentHealth { + states: DualCache, +} + +impl DeploymentHealth { + pub fn new(remote: Option) -> Self { + Self::with_clock(remote, system_clock()) + } + + pub fn with_clock(remote: Option, now: Clock) -> Self { + Self { + states: DualCache::new( + Arc::new(RemoteSlot::new(remote)), + MemorySettings::default(), + now, + ), + } + } + + /// `DualCache.attach_redis_cache`: memory stays, and `false` means a Redis was already attached. + pub fn attach(&self, remote: Remote) -> bool { + self.states.redis().attach(remote) + } + + pub fn has_redis(&self) -> bool { + self.states.redis().get().is_some() + } + + /// `DualCache.get_cache`: memory, else Redis written back to memory. + pub async fn get(&self) -> Result, Error> { + self.states.get(None, keys::STATES, None).await + } + + /// `redis_cache.get_cache`, what every pod merged so far. No Redis, a failed read and an open + /// breaker all read as `None`. + pub async fn get_shared(&self) -> Option { + let client = self.states.redis().client()?; + client + .send(commands::get::>(keys::STATES)) + .await + .map_err(Error::from) + .and_then(|raw| Ok(raw.as_deref().map(codec::decode).transpose()?)) + .unwrap_or_else(|error| { + crate::failure_log::redis_failure!(&error, "LiteLLM Cache: exception in get_cache"); + None + }) + } + + /// `DualCache.set_cache`, whose sync `RedisCache.set_cache` writes `str(value)`. + pub async fn set(&self, states: Value, ttl: Duration) -> Result<(), Error> { + self.states.set_str(keys::STATES, states, Some(ttl)).await + } +} diff --git a/litellm-rust/crates/redis-store/src/lib.rs b/litellm-rust/crates/redis-store/src/lib.rs index 69600aac2a6..eaede611add 100644 --- a/litellm-rust/crates/redis-store/src/lib.rs +++ b/litellm-rust/crates/redis-store/src/lib.rs @@ -1,5 +1,6 @@ pub mod auth_cache; mod codec; +pub mod deployment_health; mod dual_cache; mod error; mod failure_log; diff --git a/litellm-rust/crates/redis-store/tests/deployment_health.rs b/litellm-rust/crates/redis-store/tests/deployment_health.rs new file mode 100644 index 00000000000..b0fe72493b0 --- /dev/null +++ b/litellm-rust/crates/redis-store/tests/deployment_health.rs @@ -0,0 +1,62 @@ +mod support; + +use std::time::Duration; + +use litellm_python_compat::Value; +use litellm_redis_client::{Ttl, commands}; +use litellm_redis_store::deployment_health::DeploymentHealth; +use litellm_redis_testing::{RedisDeployment, Topology}; +use rstest::rstest; +use support::{connect, remote, within}; + +const STATES_KEY: &str = "litellm:health_check:deployment_health_state"; + +fn states() -> Value { + let text = |text: &str| Value::Str(text.to_owned()); + Value::Dict(vec![( + text("dep-a"), + Value::Dict(vec![ + (text("is_healthy"), Value::Bool(false)), + (text("timestamp"), Value::Float(1.5)), + (text("reason"), text("timeout")), + ]), + )]) +} + +#[rstest] +#[tokio::test] +async fn states_one_pod_writes_are_read_by_another_under_mains_key_and_encoding( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let writer = DeploymentHealth::new(Some(remote(connect(&deployment)))); + let reader = DeploymentHealth::new(Some(remote(connect(&deployment)))); + + writer.set(states(), Duration::from_secs(90)).await.unwrap(); + + let written: Option = raw.send(commands::get(STATES_KEY)).await.unwrap(); + let Ttl::Expires(remaining) = raw.send(commands::ttl(STATES_KEY)).await.unwrap() else { + panic!("{STATES_KEY} has no TTL"); + }; + assert_eq!( + written.as_deref(), + Some("{'dep-a': {'is_healthy': False, 'timestamp': 1.5, 'reason': 'timeout'}}") + ); + assert!(within(remaining, 90), "{remaining:?}"); + assert_eq!(reader.get_shared().await, Some(states())); + assert_eq!(reader.get().await.unwrap(), Some(states())); +} + +#[rstest] +#[tokio::test] +async fn without_redis_the_states_stay_in_memory_and_nothing_is_shared() { + let store = DeploymentHealth::new(None); + + store.set(states(), Duration::from_secs(90)).await.unwrap(); + + assert_eq!( + (store.get().await.unwrap(), store.get_shared().await), + (Some(states()), None) + ); +} diff --git a/litellm/redis/stores/connection.py b/litellm/redis/stores/connection.py index 21e2cc1f6f0..e33ec80bba1 100644 --- a/litellm/redis/stores/connection.py +++ b/litellm/redis/stores/connection.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final from litellm.rust_bridge._native import ( NativeAuthCache, + NativeDeploymentHealth, NativeRateLimits, NativeRedisClient, NativeRouterCooldowns, @@ -44,6 +45,9 @@ class RedisStores: def sensitive_route_pins(self) -> NativeSensitiveRoutePins: return NativeSensitiveRoutePins(self.client) + def deployment_health(self) -> NativeDeploymentHealth: + return NativeDeploymentHealth(self.client) + def connect(redis_cache: RedisCache | None) -> RedisStores: settings: Final = None if redis_cache is None else redis_cache.store_settings diff --git a/litellm/redis/stores/deployment_health/__init__.py b/litellm/redis/stores/deployment_health/__init__.py new file mode 100644 index 00000000000..98e77e4b560 --- /dev/null +++ b/litellm/redis/stores/deployment_health/__init__.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Final + +from litellm._internal_context import with_service_target +from litellm.constants import HEALTH_CHECKS_TARGET +from litellm.redis.stores.deployment_health.handle import DeploymentHealthHandle + +if TYPE_CHECKING: + from litellm.redis.stores.connection import RedisStores + +__all__ = ["DeploymentHealthStore"] + + +class DeploymentHealthStore: + """The background health checks' per-deployment states: memory per worker, Redis shared by every pod. + The store owns the key and its encoding; merging and staleness stay with ``DeploymentHealthCache``.""" + + def __init__(self, handle: DeploymentHealthHandle) -> None: + self._handle: Final = handle + + @property + def has_redis(self) -> bool: + return self._handle.has_redis + + def attach_redis(self, redis: RedisStores) -> None: + """Adds the Redis tier once, keeping memory, as ``DualCache.attach_redis_cache`` does.""" + if redis.client is not None: + self._handle.attach(redis.client) + + @with_service_target(HEALTH_CHECKS_TARGET) + async def get(self) -> object: + return await self._handle.get() + + @with_service_target(HEALTH_CHECKS_TARGET) + def get_sync(self) -> object: + return self._handle.get_sync() + + @with_service_target(HEALTH_CHECKS_TARGET) + def get_shared_sync(self) -> object: + """What every pod merged so far, from Redis alone. No Redis, a failed read and an open breaker read as None.""" + return self._handle.get_shared_sync() + + @with_service_target(HEALTH_CHECKS_TARGET) + def set_sync(self, states: object, ttl: float) -> None: + self._handle.set_sync(states, ttl) diff --git a/litellm/redis/stores/deployment_health/handle.py b/litellm/redis/stores/deployment_health/handle.py new file mode 100644 index 00000000000..01fc78c33e8 --- /dev/null +++ b/litellm/redis/stores/deployment_health/handle.py @@ -0,0 +1,16 @@ +from __future__ import annotations + +from collections.abc import Awaitable +from typing import Protocol + +from litellm.rust_bridge._native import NativeRedisClient + + +class DeploymentHealthHandle(Protocol): + @property + def has_redis(self) -> bool: ... + def attach(self, client: NativeRedisClient) -> bool: ... + def get(self) -> Awaitable[object]: ... + def get_sync(self) -> object: ... + def get_shared_sync(self) -> object: ... + def set_sync(self, states: object, ttl_seconds: float) -> None: ... diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 471f5c8ad91..9d1cc35bf87 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -64,6 +64,7 @@ from litellm.litellm_core_utils.routing_common_utils import ( from litellm.litellm_core_utils.routing_groups import validate_routing_strategy from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.redis.stores.connection import connect +from litellm.redis.stores.deployment_health import DeploymentHealthStore from litellm.redis.stores.router_cooldowns import RouterCooldownStore from litellm.redis.stores.router_usage import RouterUsage from litellm.router_api import ( @@ -1395,7 +1396,9 @@ class RustCore( store=RouterCooldownStore(redis.router_cooldowns()), default_cooldown_time=self._cooldown_given, # pyright: ignore[reportArgumentType] # the constructor's value, as the Python core passes it ) - self._health_checks: DeploymentHealthCache = resolved.health_routing.health_checks(self._cache) + self._health_checks: DeploymentHealthCache = resolved.health_routing.health_checks( + DeploymentHealthStore(redis.deployment_health()) + ) self._version: int = 1 self._settings: RustSettings | None = None self._tag_filtering_models: tuple[int, bool] | None = None @@ -2025,6 +2028,7 @@ class RustCore( redis: Final = connect(changes.redis_cache) self._usage.attach_redis(redis) self._cooldowns.cooldown_store.attach_redis(redis) + self._health_checks.store.attach_redis(redis) self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache) self._prerouter.attach_redis_cache(changes.redis_cache) return ApplyResult( diff --git a/litellm/router_rust/surface/health_routing.py b/litellm/router_rust/surface/health_routing.py index 44a045a0d5d..6ffdaec508e 100644 --- a/litellm/router_rust/surface/health_routing.py +++ b/litellm/router_rust/surface/health_routing.py @@ -17,7 +17,7 @@ from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.types.router import AllowedFailsPolicy if TYPE_CHECKING: - from litellm.caching.caching import DualCache + from litellm.redis.stores.deployment_health import DeploymentHealthStore from litellm.router_config import RouterConfig from litellm.types.router import Deployment, DeploymentTypedDict @@ -74,8 +74,8 @@ class HealthRouting: else None, } - def health_checks(self, cache: "DualCache") -> DeploymentHealthCache: - return DeploymentHealthCache(cache=cache, staleness_threshold=self.health_check_staleness) + def health_checks(self, store: "DeploymentHealthStore") -> DeploymentHealthCache: + return DeploymentHealthCache(store=store, staleness_threshold=self.health_check_staleness) def failover_document(plain: Mapping[str, object]) -> Mapping[str, object]: diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index 30547f8002c..e5f017e72ed 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -1,5 +1,5 @@ """ -Wrapper around router cache for health-check-driven routing. +Health-check-driven routing state, kept in the native ``DeploymentHealthStore``. Stores per-deployment health state from background health checks and exposes it for router candidate filtering. @@ -12,10 +12,7 @@ from typing import TYPE_CHECKING, Any, Final from typing_extensions import TypedDict from litellm import verbose_logger -from litellm._internal_context import with_service_target -from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError -from litellm.constants import HEALTH_CHECKS_TARGET +from litellm.redis.stores.deployment_health import DeploymentHealthStore if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -31,17 +28,6 @@ class DeploymentHealthStateValue(TypedDict): reason: str -@with_service_target(HEALTH_CHECKS_TARGET) -def _read_shared_health_snapshot(cache: DualCache, key: str) -> object: - redis_cache: Final = cache.redis_cache - if redis_cache is None: - return None - try: - return redis_cache.get_cache(key) - except RedisCircuitBreakerOpenError: - return None - - class DeploymentHealthCache: """ Cache for deployment health states produced by background health checks. @@ -51,13 +37,10 @@ class DeploymentHealthCache: are treated as healthy (unknown). """ - CACHE_KEY = "litellm:health_check:deployment_health_state" - - def __init__(self, cache: DualCache, staleness_threshold: float): - self.cache = cache + def __init__(self, store: DeploymentHealthStore, staleness_threshold: float): + self.store: Final = store self.staleness_threshold = staleness_threshold - @with_service_target(HEALTH_CHECKS_TARGET) def set_deployment_health_states(self, states: dict[str, DeploymentHealthStateValue]) -> None: """Merge the given states into the shared cache entry, pruning expired ones. @@ -71,8 +54,8 @@ class DeploymentHealthCache: merge is used so peers are not erased. """ try: - redis_raw: Final = _read_shared_health_snapshot(self.cache, self.CACHE_KEY) - raw: Final = redis_raw if isinstance(redis_raw, dict) else self.cache.get_cache(key=self.CACHE_KEY) + redis_raw: Final = self.store.get_shared_sync() + raw: Final = redis_raw if isinstance(redis_raw, dict) else self.store.get_sync() existing: Final = raw if isinstance(raw, dict) else {} expiry_seconds: Final = self.staleness_threshold * 1.5 now: Final = time.time() @@ -81,11 +64,7 @@ class DeploymentHealthCache: for model_id, state in {**existing, **states}.items() if isinstance(state, dict) and (now - _timestamp_of(state)) < expiry_seconds # pyright: ignore[reportUnknownArgumentType] # a cached record } - self.cache.set_cache( - key=self.CACHE_KEY, - value=merged, - ttl=int(expiry_seconds), - ) + self.store.set_sync(merged, int(expiry_seconds)) except Exception as e: verbose_logger.error( "DeploymentHealthCache::set_deployment_health_states - Exception: %s", @@ -105,21 +84,26 @@ class DeploymentHealthCache: and (now - _timestamp_of(state)) < self.staleness_threshold # pyright: ignore[reportUnknownArgumentType] # a cached record } - @with_service_target(HEALTH_CHECKS_TARGET) async def async_get_states(self) -> object: """Every deployment's raw state, as the router core's health filter reads it; a failed read is None.""" - return await self.cache.async_get_cache(key=self.CACHE_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + try: + return await self.store.get() + except Exception as e: # noqa: BLE001 # a failed read counts as no health state, as the DualCache read it replaced swallowed errors + verbose_logger.debug("DeploymentHealthCache::async_get_states - Exception: %s", str(e)) + return None - @with_service_target(HEALTH_CHECKS_TARGET) def get_states(self) -> object: """``async_get_states`` for a sync call.""" - return self.cache.get_cache(key=self.CACHE_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + try: + return self.store.get_sync() + except Exception as e: # noqa: BLE001 # a failed read counts as no health state, as the DualCache read it replaced swallowed errors + verbose_logger.debug("DeploymentHealthCache::get_states - Exception: %s", str(e)) + return None - @with_service_target(HEALTH_CHECKS_TARGET) async def async_get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Return set of deployment IDs currently marked unhealthy and not stale.""" try: - raw: Final = await self.cache.async_get_cache(key=self.CACHE_KEY) + raw: Final = await self.store.get() return self._extract_unhealthy_ids(raw) except Exception as e: verbose_logger.debug( @@ -128,11 +112,10 @@ class DeploymentHealthCache: ) return set() - @with_service_target(HEALTH_CHECKS_TARGET) def get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Sync version: return set of deployment IDs currently marked unhealthy and not stale.""" try: - raw: Final = self.cache.get_cache(key=self.CACHE_KEY) + raw: Final = self.store.get_sync() return self._extract_unhealthy_ids(raw) except Exception as e: verbose_logger.debug( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index d0deb8000af..f4dfd69534e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -95,6 +95,17 @@ class NativeRouterCooldowns: def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, _NativeCooldown]: ... +@final +class NativeDeploymentHealth: + def __new__(cls, client: NativeRedisClient | None) -> NativeDeploymentHealth: ... + @property + def has_redis(self) -> bool: ... + def attach(self, client: NativeRedisClient) -> bool: ... + def get(self) -> Future[object]: ... + def get_sync(self) -> object: ... + def get_shared_sync(self) -> object: ... + def set_sync(self, states: object, ttl_seconds: float) -> None: ... + @final class NativeSensitiveRoutePins: def __new__(cls, client: NativeRedisClient | None) -> NativeSensitiveRoutePins: ... @@ -567,6 +578,7 @@ __all__ = [ "NativeClickHouseSpendConfig", "NativeClickHouseSpendStorage", "NativeDatabase", + "NativeDeploymentHealth", "NativeDiagnosticProcessor", "NativePending", "NativeRateLimits", diff --git a/tests/unit/router_utils/test_health_state_cache.py b/tests/unit/router_utils/test_health_state_cache.py index aa976bb1002..639b13aaf34 100644 --- a/tests/unit/router_utils/test_health_state_cache.py +++ b/tests/unit/router_utils/test_health_state_cache.py @@ -6,19 +6,19 @@ import time import pytest -from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.redis.stores.deployment_health import DeploymentHealthStore from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.rust_bridge._native import NativeDeploymentHealth, NativeRedisClient @pytest.fixture -def cache(): - return DualCache() +def store(): + return DeploymentHealthStore(NativeDeploymentHealth(None)) @pytest.fixture -def health_cache(cache): - return DeploymentHealthCache(cache=cache, staleness_threshold=60.0) +def health_cache(store): + return DeploymentHealthCache(store=store, staleness_threshold=60.0) def test_set_and_get_unhealthy_ids(health_cache): @@ -120,13 +120,11 @@ def test_set_merges_states_from_scoped_writers(health_cache): health_cache.set_deployment_health_states( {"listed-bad": {"is_healthy": False, "timestamp": now, "reason": "check_failed"}} ) - health_cache.set_deployment_health_states( - {"other-ok": {"is_healthy": True, "timestamp": now, "reason": ""}} - ) + health_cache.set_deployment_health_states({"other-ok": {"is_healthy": True, "timestamp": now, "reason": ""}}) assert health_cache.get_unhealthy_deployment_ids() == {"listed-bad"} -def test_set_prunes_expired_entries(health_cache, cache): +def test_set_prunes_expired_entries(health_cache, store): """Entries older than 1.5x the staleness threshold are dropped on write.""" expired_time = time.time() - 100 # threshold 60s, prune horizon 90s health_cache.set_deployment_health_states( @@ -136,34 +134,55 @@ def test_set_prunes_expired_entries(health_cache, cache): health_cache.set_deployment_health_states( {"fresh": {"is_healthy": False, "timestamp": now, "reason": "check_failed"}} ) - stored = cache.get_cache(key=DeploymentHealthCache.CACHE_KEY) + stored = store.get_sync() assert set(stored.keys()) == {"fresh"} -class _SharedRedisFake: - """Shared get/set key-value store standing in for the Redis layer of a DualCache.""" +class _SharedRedis: + """The Redis entry every pod shares; a failed read reads as a miss, as the store's does.""" def __init__(self): - self.store = {} + self.states = None self.fail_get = False - self.breaker_open = False - def get_cache(self, key, parent_otel_span=None, **kwargs): - if self.breaker_open: - raise RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping get_cache") - if self.fail_get: - return None # RedisCache.get_cache swallows connection errors and returns None - return self.store.get(key) - def set_cache(self, key, value, **kwargs): - self.store[key] = value +class _PodHandle: + """One pod's ``DeploymentHealthHandle``: its own memory in front of the shared Redis entry.""" + + def __init__(self, redis: _SharedRedis): + self._redis = redis + self._memory = None + + @property + def has_redis(self) -> bool: + return True + + def attach(self, client: NativeRedisClient) -> bool: + return False + + async def get(self) -> object: + return self.get_sync() + + def get_sync(self) -> object: + return self._memory if self._memory is not None else self.get_shared_sync() + + def get_shared_sync(self) -> object: + return None if self._redis.fail_get else self._redis.states + + def set_sync(self, states: object, ttl_seconds: float) -> None: + self._memory = states + self._redis.states = states + + +def _pod(redis: _SharedRedis) -> DeploymentHealthCache: + return DeploymentHealthCache(store=DeploymentHealthStore(_PodHandle(redis)), staleness_threshold=60.0) def test_scoped_writers_on_shared_redis_preserve_each_other(): """Pods with different allowlists share one Redis entry; each merge must keep the peer's scope.""" - redis_fake = _SharedRedisFake() - pod_a = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) - pod_b = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) + redis_fake = _SharedRedis() + pod_a = _pod(redis_fake) + pod_b = _pod(redis_fake) pod_a.set_deployment_health_states( {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} ) @@ -173,15 +192,15 @@ def test_scoped_writers_on_shared_redis_preserve_each_other(): pod_a.set_deployment_health_states( {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} ) - assert set(redis_fake.store[DeploymentHealthCache.CACHE_KEY]) == {"prod-bad", "internal-bad"} + assert set(redis_fake.states) == {"prod-bad", "internal-bad"} assert pod_a.get_unhealthy_deployment_ids() == {"prod-bad", "internal-bad"} def test_failed_redis_read_falls_back_to_local_copy(): """A swallowed Redis GET error must not make a writer erase peer scopes it already saw.""" - redis_fake = _SharedRedisFake() - pod_a = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) - pod_b = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) + redis_fake = _SharedRedis() + pod_a = _pod(redis_fake) + pod_b = _pod(redis_fake) pod_a.set_deployment_health_states( {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} ) @@ -195,27 +214,4 @@ def test_failed_redis_read_falls_back_to_local_copy(): pod_a.set_deployment_health_states( {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} ) - assert set(redis_fake.store[DeploymentHealthCache.CACHE_KEY]) == {"prod-bad", "internal-bad"} - - -def test_open_circuit_breaker_read_still_merges_into_local_copy(caplog): - """A read refused by the open breaker is a miss, so the merge and local write still happen quietly.""" - redis_fake = _SharedRedisFake() - pod_a = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) - pod_b = DeploymentHealthCache(cache=DualCache(redis_cache=redis_fake), staleness_threshold=60.0) - pod_a.set_deployment_health_states( - {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} - ) - pod_b.set_deployment_health_states( - {"internal-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "timeout"}} - ) - pod_a.set_deployment_health_states( - {"prod-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} - ) - redis_fake.breaker_open = True - with caplog.at_level("ERROR"): - pod_a.set_deployment_health_states( - {"prod-new-bad": {"is_healthy": False, "timestamp": time.time(), "reason": "check_failed"}} - ) - assert caplog.records == [] - assert pod_a.get_unhealthy_deployment_ids() == {"prod-bad", "internal-bad", "prod-new-bad"} + assert set(redis_fake.states) == {"prod-bad", "internal-bad"} diff --git a/tests/unit/test_router/core/test_parity_rulings_2026_11.py b/tests/unit/test_router/core/test_parity_rulings_2026_11.py index dd78224774c..08e7d1c4e3f 100644 --- a/tests/unit/test_router/core/test_parity_rulings_2026_11.py +++ b/tests/unit/test_router/core/test_parity_rulings_2026_11.py @@ -135,7 +135,7 @@ async def test_an_unreadable_health_timestamp_drops_only_that_record( } health_checks: Final = core.health().health_checks # As another writer left it in the shared cache, then merged by this router's own writer. - _sync(health_checks.cache, "set_cache")(key=health_checks.CACHE_KEY, value=states) + health_checks.store.set_sync(states, 60) await _async(core, "acompletion")(model="g", messages=MESSAGES) health_checks.set_deployment_health_states(states) # pyright: ignore[reportArgumentType] # a record with a bad timestamp, on purpose await _async(core, "acompletion")(model="g", messages=MESSAGES)