From e20f07d178ab93e8bac4034f93be03a140f53513 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 13 Nov 2026 12:22:28 +0900 Subject: [PATCH] feat(redis): serve the health checks' deployment states from a native store (#50357) * 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. * refactor(redis): type the deployment health states in the store The health store kept a raw Value from Redis to Python and back, so Python re-checked the shape by hand and router-core only got a typed record after a Python pass. The family now has a types.rs with HealthState and HealthStates, the bridge converts through PyData, and the Python store returns DeploymentHealthStateValue mappings Reads keep main's leniency for records another writer left: extra fields are ignored, a falsy flag is unhealthy and a missing one healthy, a timestamp that is no number reads as 0, and a record that is no dict is dropped alone. A rewrite writes only is_healthy, timestamp and reason --- litellm/redis/stores/connection.py | 4 + .../stores/deployment_health/__init__.py | 47 +++++ .../redis/stores/deployment_health/handle.py | 17 ++ .../redis/stores/deployment_health/types.py | 15 ++ litellm/router_rust/core.py | 6 +- litellm/router_rust/routing_stores.py | 5 +- litellm/router_rust/surface/health_routing.py | 6 +- litellm/router_utils/health_state_cache.py | 116 ++++-------- litellm/rust/crates/python-bridge/src/lib.rs | 3 + .../src/redis/store/deployment_health.rs | 72 +++++++ .../python-bridge/src/redis/store/mod.rs | 1 + .../python-bridge/src/redis/store/policy.rs | 8 +- litellm/rust/crates/redis-store/AGENTS.md | 2 +- .../redis-store/src/deployment_health/keys.rs | 2 + .../redis-store/src/deployment_health/mod.rs | 78 ++++++++ .../src/deployment_health/types.rs | 80 ++++++++ litellm/rust/crates/redis-store/src/lib.rs | 1 + .../redis-store/tests/deployment_health.rs | 178 ++++++++++++++++++ litellm/rust_bridge/_native.pyi | 17 ++ .../redis/stores/test_deployment_health.py | 88 +++++++++ tests/unit/router_rust/test_routing_stores.py | 31 ++- .../router_utils/test_health_state_cache.py | 156 +++++++++------ .../core/test_parity_rulings_2026_11.py | 2 +- 23 files changed, 769 insertions(+), 166 deletions(-) create mode 100644 litellm/redis/stores/deployment_health/__init__.py create mode 100644 litellm/redis/stores/deployment_health/handle.py create mode 100644 litellm/redis/stores/deployment_health/types.py 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/src/deployment_health/types.rs create mode 100644 litellm/rust/crates/redis-store/tests/deployment_health.rs create mode 100644 tests/unit/redis/stores/test_deployment_health.py 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..027cd3d5251 --- /dev/null +++ b/litellm/redis/stores/deployment_health/__init__.py @@ -0,0 +1,47 @@ +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 +from litellm.redis.stores.deployment_health.types import DeploymentHealthStates, DeploymentHealthStateValue + +if TYPE_CHECKING: + from litellm.redis.stores.connection import RedisStores + +__all__ = ["DeploymentHealthStateValue", "DeploymentHealthStates", "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) -> DeploymentHealthStates | None: + return await self._handle.get() + + @with_service_target(HEALTH_CHECKS_TARGET) + def get_sync(self) -> DeploymentHealthStates | None: + return self._handle.get_sync() + + @with_service_target(HEALTH_CHECKS_TARGET) + def get_shared_sync(self) -> DeploymentHealthStates | None: + """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: DeploymentHealthStates, 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..680032c23db --- /dev/null +++ b/litellm/redis/stores/deployment_health/handle.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from collections.abc import Awaitable +from typing import Protocol + +from litellm.redis.stores.deployment_health.types import DeploymentHealthStates +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[DeploymentHealthStates | None]: ... + def get_sync(self) -> DeploymentHealthStates | None: ... + def get_shared_sync(self) -> DeploymentHealthStates | None: ... + def set_sync(self, states: DeploymentHealthStates, ttl_seconds: float) -> None: ... diff --git a/litellm/redis/stores/deployment_health/types.py b/litellm/redis/stores/deployment_health/types.py new file mode 100644 index 00000000000..b43de912472 --- /dev/null +++ b/litellm/redis/stores/deployment_health/types.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +from collections.abc import Mapping +from typing import TypeAlias + +from typing_extensions import ReadOnly, TypedDict + + +class DeploymentHealthStateValue(TypedDict): + is_healthy: ReadOnly[bool] + timestamp: ReadOnly[float] + reason: ReadOnly[str] + + +DeploymentHealthStates: TypeAlias = Mapping[str, DeploymentHealthStateValue] diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index e151bae80fb..a689b05a485 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -76,6 +76,7 @@ from litellm.litellm_core_utils.total_deadline import ( enforce_total_timeout_scope, ) 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 ( @@ -1486,7 +1487,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 @@ -2127,6 +2130,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/routing_stores.py b/litellm/router_rust/routing_stores.py index 706866a9d05..9d120c0fea4 100644 --- a/litellm/router_rust/routing_stores.py +++ b/litellm/router_rust/routing_stores.py @@ -10,6 +10,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_router_logger from litellm.constants import HEALTH_CHECKS_TARGET from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.redis.stores.deployment_health import DeploymentHealthStates from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter from litellm.redis.stores.router_usage.types import CounterKind @@ -140,8 +141,8 @@ def service_target_of(step: StoreStep | LocalStep) -> str: class HealthStates(Protocol): - async def async_get_states(self) -> object: ... - def get_states(self) -> object: ... + async def async_get_states(self) -> DeploymentHealthStates | None: ... + def get_states(self) -> DeploymentHealthStates | None: ... def _aligned(model_ids: Sequence[str], active: Sequence[tuple[str, CooldownCacheValue]]) -> tuple[object, ...]: 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 47ac0abeaff..0f90fbbe1d8 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. @@ -9,15 +9,10 @@ import time from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final -from pydantic import TypeAdapter from typing_extensions import ReadOnly, 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.litellm_core_utils.core_helpers import as_str_mapping +from litellm.redis.stores.deployment_health import DeploymentHealthStates, DeploymentHealthStore if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -27,26 +22,6 @@ else: Span = Any -_CACHED_HEALTH_SNAPSHOT: Final = TypeAdapter(object) - - -class DeploymentHealthStateValue(TypedDict): - is_healthy: bool - timestamp: float - 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. @@ -56,14 +31,11 @@ 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: + def set_deployment_health_states(self, states: DeploymentHealthStates) -> None: """Merge the given states into the shared cache entry, pruning expired ones. Merging instead of replacing lets writers probing different deployment @@ -76,56 +48,46 @@ 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) - existing: Final = raw if isinstance(raw, dict) else {} + shared: Final = self.store.get_shared_sync() + existing: Final = shared if shared is not None else self.store.get_sync() or {} expiry_seconds: Final = self.staleness_threshold * 1.5 now: Final = time.time() - merged: Final = { - model_id: state - 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(_unexpired({**existing, **states}, now - expiry_seconds), int(expiry_seconds)) except Exception as e: verbose_logger.error( "DeploymentHealthCache::set_deployment_health_states - Exception: %s", str(e), ) - def _extract_unhealthy_ids(self, raw: object) -> set[str]: - """Given raw cache value, return set of non-stale unhealthy deployment IDs.""" - if not raw or not isinstance(raw, dict): - return set() + def _extract_unhealthy_ids(self, states: DeploymentHealthStates | None) -> set[str]: + """The non-stale unhealthy deployment IDs.""" now: Final = time.time() return { model_id - for model_id, state in raw.items() - if isinstance(state, dict) - and not state.get("is_healthy", True) - and (now - _timestamp_of(state)) < self.staleness_threshold # pyright: ignore[reportUnknownArgumentType] # a cached record + for model_id, state in (states or {}).items() + if not state["is_healthy"] and now - state["timestamp"] < self.staleness_threshold } - @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 + async def async_get_states(self) -> DeploymentHealthStates | None: + """Every deployment's state, as the router core's health filter reads it; a failed read is None.""" + 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: + def get_states(self) -> DeploymentHealthStates | None: """``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 = _CACHED_HEALTH_SNAPSHOT.validate_python(await self.cache.async_get_cache(key=self.CACHE_KEY)) - return self._extract_unhealthy_ids(raw) + return self._extract_unhealthy_ids(await self.store.get()) except Exception as e: verbose_logger.debug( "DeploymentHealthCache::async_get_unhealthy_deployment_ids - Exception: %s", @@ -133,12 +95,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 = _CACHED_HEALTH_SNAPSHOT.validate_python(self.cache.get_cache(key=self.CACHE_KEY)) - return self._extract_unhealthy_ids(raw) + return self._extract_unhealthy_ids(self.store.get_sync()) except Exception as e: verbose_logger.debug( "DeploymentHealthCache::get_unhealthy_deployment_ids - Exception: %s", @@ -147,24 +107,18 @@ class DeploymentHealthCache: return set() -def _timestamp_of(state: Mapping[str, object]) -> float: - """A health record's timestamp; one that is no number reads as 0 (stale), so a bad record - drops out alone instead of turning health routing off for every deployment.""" - timestamp: Final = state.get("timestamp", 0) - return float(timestamp) if isinstance(timestamp, int | float) and not isinstance(timestamp, bool) else 0.0 - - class DeploymentHealthCheck(TypedDict): is_healthy: ReadOnly[bool] timestamp: ReadOnly[float] -def health_states_of(raw: object) -> Mapping[str, DeploymentHealthCheck]: - """The cached states as ``_extract_unhealthy_ids`` reads them: any falsy ``is_healthy`` is unhealthy, a - record without one is healthy, and what is no record is dropped.""" - records: Final = ((model_id, as_str_mapping(state)) for model_id, state in (as_str_mapping(raw) or {}).items()) +def health_states_of(states: DeploymentHealthStates | None) -> Mapping[str, DeploymentHealthCheck]: + """The cached states as router-core reads them, which is everything but the reason.""" return { - model_id: {"is_healthy": bool(state.get("is_healthy", True)), "timestamp": _timestamp_of(state)} - for model_id, state in records - if state is not None + model_id: {"is_healthy": state["is_healthy"], "timestamp": float(state["timestamp"])} + for model_id, state in (states or {}).items() } + + +def _unexpired(states: DeploymentHealthStates, cutoff: float) -> DeploymentHealthStates: + return {model_id: state for model_id, state in states.items() if state["timestamp"] > cutoff} diff --git a/litellm/rust/crates/python-bridge/src/lib.rs b/litellm/rust/crates/python-bridge/src/lib.rs index 45ea95560e1..4a8f6917396 100644 --- a/litellm/rust/crates/python-bridge/src/lib.rs +++ b/litellm/rust/crates/python-bridge/src/lib.rs @@ -53,6 +53,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, @@ -153,6 +155,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..4c5bb740233 --- /dev/null +++ b/litellm/rust/crates/python-bridge/src/redis/store/deployment_health.rs @@ -0,0 +1,72 @@ +use std::sync::Arc; + +use litellm_redis_store::{ + Error, + deployment_health::{DeploymentHealth, HealthStates}, +}; +use pyo3::prelude::*; + +use super::{ + PyData, duration, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote, +}; +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(py: Python<'_>, client: Option>) -> PyResult { + Ok(Self(Arc::new(DeploymentHealth::new( + client.map(|client| remote(&client)), + router_tier_policies(py)?.default_in_memory_ttl, + )))) + } + + /// 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(PyData)) + }) + } + + 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(PyData)) + }) + } + + 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(PyData)) + }) + } + + fn set_sync( + &self, + py: Python<'_>, + states: PyData, + 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/python-bridge/src/redis/store/policy.rs b/litellm/rust/crates/python-bridge/src/redis/store/policy.rs index ffdf7b97a2e..506c743fad4 100644 --- a/litellm/rust/crates/python-bridge/src/redis/store/policy.rs +++ b/litellm/rust/crates/python-bridge/src/redis/store/policy.rs @@ -20,6 +20,7 @@ const MEMORY_TTL: FieldSpec> = pub(crate) struct RouterTierPolicies { pub(crate) cooldowns: TierPolicy, pub(crate) usage: TierPolicy, + pub(crate) default_in_memory_ttl: Option, } pub(crate) fn router_tier_policies(py: Python<'_>) -> PyResult { @@ -40,6 +41,7 @@ fn project(snapshot: &Snapshot<'_>) -> Result`, 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, or `None` when it returns a settlement the caller waits on, like `SpendCounterStore::increment_many_post_call` - 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..9c4530e27ed --- /dev/null +++ b/litellm/rust/crates/redis-store/src/deployment_health/mod.rs @@ -0,0 +1,78 @@ +mod keys; +mod types; + +use std::{sync::Arc, time::Duration}; + +use litellm_redis_client::commands; + +pub use self::types::{HealthState, HealthStates}; +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, default_in_memory_ttl: Option) -> Self { + Self::with_clock(remote, default_in_memory_ttl, system_clock()) + } + + pub fn with_clock( + remote: Option, + default_in_memory_ttl: Option, + now: Clock, + ) -> Self { + Self { + states: DualCache::new( + Arc::new(RemoteSlot::new(remote)), + MemorySettings { + entries: None, + ttl: default_in_memory_ttl, + }, + 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: HealthStates, ttl: Duration) -> Result<(), Error> { + self.states.set_str(keys::STATES, states, Some(ttl)).await + } +} diff --git a/litellm/rust/crates/redis-store/src/deployment_health/types.rs b/litellm/rust/crates/redis-store/src/deployment_health/types.rs new file mode 100644 index 00000000000..977a6a39a08 --- /dev/null +++ b/litellm/rust/crates/redis-store/src/deployment_health/types.rs @@ -0,0 +1,80 @@ +use litellm_python_compat::{Value, truthy::truthy}; + +use crate::{DecodeError, FromValue, IntoValue, Number}; + +/// One background health check's result, the `DeploymentHealthStateValue` dict the proxy writes. +#[derive(Clone, Debug, PartialEq)] +pub struct HealthState { + pub is_healthy: bool, + pub timestamp: Number, + pub reason: String, +} + +/// Every deployment's last health check by id, in the order the writer merged them. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct HealthStates(pub Vec<(String, HealthState)>); + +impl HealthState { + /// What main's readers take from a record: extra fields are ignored, any falsy flag is + /// unhealthy and a missing one healthy, and a timestamp that is no number reads as `0` (stale). + fn read(fields: &[(Value, Value)]) -> Self { + let field = |name: &str| { + fields + .iter() + .find(|(key, _)| matches!(key, Value::Str(key) if key == name)) + .map(|(_, value)| value) + }; + Self { + is_healthy: field("is_healthy").is_none_or(truthy), + timestamp: field("timestamp") + .and_then(Number::from_value) + .unwrap_or(Number::Float(0.0)), + reason: match field("reason") { + Some(Value::Str(reason)) => reason.clone(), + _ => String::new(), + }, + } + } +} + +impl IntoValue for HealthState { + fn into_value(self) -> Value { + let field = |name: &str, value: Value| (Value::Str(name.to_owned()), value); + Value::Dict(vec![ + field("is_healthy", Value::Bool(self.is_healthy)), + field("timestamp", self.timestamp.to_value()), + field("reason", Value::Str(self.reason)), + ]) + } +} + +/// A record that is no dict, under a key that is no string, is dropped alone. +impl FromValue for HealthStates { + fn from_value(value: &Value) -> Result { + let Value::Dict(records) = value else { + return Err(DecodeError::Shape("deployment health states are a dict")); + }; + Ok(Self( + records + .iter() + .filter_map(|record| match record { + (Value::Str(id), Value::Dict(fields)) => { + Some((id.clone(), HealthState::read(fields))) + } + _ => None, + }) + .collect(), + )) + } +} + +impl IntoValue for HealthStates { + fn into_value(self) -> Value { + Value::Dict( + self.0 + .into_iter() + .map(|(id, state)| (Value::Str(id), state.into_value())) + .collect(), + ) + } +} 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..4f3b0196456 --- /dev/null +++ b/litellm/rust/crates/redis-store/tests/deployment_health.rs @@ -0,0 +1,178 @@ +mod support; + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_redis_client::{Ttl, commands}; +use litellm_redis_store::{ + Clock, Number, + deployment_health::{DeploymentHealth, HealthState, HealthStates}, +}; +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() -> HealthStates { + states_because("timeout") +} + +fn states_because(reason: &str) -> HealthStates { + one("dep-a", false, Number::Float(1.5), reason) +} + +fn one(id: &str, is_healthy: bool, timestamp: Number, reason: &str) -> HealthStates { + HealthStates(vec![( + id.to_owned(), + HealthState { + is_healthy, + timestamp, + reason: reason.to_owned(), + }, + )]) +} + +#[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))), None); + let reader = DeploymentHealth::new(Some(remote(connect(&deployment))), None); + + 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, None); + + store.set(states(), Duration::from_secs(90)).await.unwrap(); + + assert_eq!( + (store.get().await.unwrap(), store.get_shared().await), + (Some(states()), None) + ); +} + +#[rstest] +#[tokio::test] +async fn a_snapshot_read_from_redis_is_read_again_once_the_configured_memory_ttl_passes( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let now = Arc::new(Mutex::new(Duration::from_secs(1_730_000_000))); + let clock: Clock = { + let now = now.clone(); + Arc::new(move || *now.lock().unwrap()) + }; + let writer = DeploymentHealth::new(Some(remote(connect(&deployment))), None); + let reader = DeploymentHealth::with_clock( + Some(remote(connect(&deployment))), + Some(Duration::from_secs(5)), + clock, + ); + writer.set(states(), Duration::from_secs(90)).await.unwrap(); + let first = reader.get().await.unwrap(); + writer + .set(states_because("rate limited"), Duration::from_secs(90)) + .await + .unwrap(); + + let within_ttl = reader.get().await.unwrap(); + *now.lock().unwrap() += Duration::from_secs(6); + let after_ttl = reader.get().await.unwrap(); + + assert_eq!( + (first, within_ttl, after_ttl), + ( + Some(states()), + Some(states()), + Some(states_because("rate limited")) + ) + ); +} + +#[rstest] +#[case::extra_fields_are_ignored( + "{'dep-a': {'is_healthy': False, 'timestamp': 1.5, 'reason': 'timeout', 'region': 'eu'}}", + one("dep-a", false, Number::Float(1.5), "timeout") +)] +#[case::a_falsy_flag_is_unhealthy( + "{'dep-a': {'is_healthy': 0, 'timestamp': 95, 'reason': 'timeout'}}", + one("dep-a", false, Number::Int(95), "timeout") +)] +#[case::a_missing_flag_is_healthy( + "{'dep-a': {'timestamp': 1.5}}", + one("dep-a", true, Number::Float(1.5), "") +)] +#[case::a_timestamp_that_is_no_number_is_stale( + "{'dep-a': {'is_healthy': False, 'timestamp': 'soon', 'reason': 'timeout'}}", + one("dep-a", false, Number::Float(0.0), "timeout") +)] +#[case::a_bool_timestamp_is_no_number( + "{'dep-a': {'is_healthy': False, 'timestamp': True, 'reason': 'timeout'}}", + one("dep-a", false, Number::Float(0.0), "timeout") +)] +#[case::what_is_no_record_is_dropped_alone( + "{'dep-a': 'not a record', 'dep-b': None, 'dep-c': {'is_healthy': False, 'timestamp': 1.5, 'reason': 'timeout'}}", + one("dep-c", false, Number::Float(1.5), "timeout") +)] +#[tokio::test] +async fn a_record_another_writer_left_is_read_the_way_mains_readers_read_it( + #[case] written: &str, + #[case] read: HealthStates, + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let reader = DeploymentHealth::new(Some(remote(connect(&deployment))), None); + raw.send(commands::set(STATES_KEY, written, Duration::from_secs(90))) + .await + .unwrap(); + + assert_eq!( + (reader.get_shared().await, reader.get().await.unwrap()), + (Some(read.clone()), Some(read)) + ); +} + +#[rstest] +#[tokio::test] +async fn an_entry_that_is_no_dict_reads_as_nothing_shared( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let reader = DeploymentHealth::new(Some(remote(connect(&deployment))), None); + raw.send(commands::set( + STATES_KEY, + "['dep-a']", + Duration::from_secs(90), + )) + .await + .unwrap(); + + assert_eq!( + (reader.get_shared().await, reader.get().await.unwrap()), + (None, None) + ); +} diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 22714e7f325..baeb12476ed 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -128,6 +128,11 @@ class NativeClickHouseSpendStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... +class _NativeHealthState(TypedDict): + is_healthy: bool + timestamp: float + reason: str + class _NativeCooldown(TypedDict): exception_received: str status_code: str @@ -207,6 +212,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[dict[str, _NativeHealthState] | None]: ... + def get_sync(self) -> dict[str, _NativeHealthState] | None: ... + def get_shared_sync(self) -> dict[str, _NativeHealthState] | None: ... + def set_sync(self, states: Mapping[str, object], ttl_seconds: float) -> None: ... + @final class NativeSensitiveRoutePins: def __new__(cls, client: NativeRedisClient | None) -> NativeSensitiveRoutePins: ... @@ -680,6 +696,7 @@ __all__ = [ "NativeClickHouseSpendStorage", "NativeCostMapSource", "NativeDatabase", + "NativeDeploymentHealth", "NativeDiagnosticProcessor", "NativeFetchFailed", "NativeFetchPolicy", diff --git a/tests/unit/redis/stores/test_deployment_health.py b/tests/unit/redis/stores/test_deployment_health.py new file mode 100644 index 00000000000..0ceb9a363a3 --- /dev/null +++ b/tests/unit/redis/stores/test_deployment_health.py @@ -0,0 +1,88 @@ +import ast +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from pydantic import TypeAdapter + +from litellm.redis.stores.connection import connect +from litellm.redis.stores.deployment_health import DeploymentHealthStateValue + + +@dataclass(frozen=True, slots=True) +class _WireStates: + """The health states as main's sync ``RedisCache.set_cache`` writes them (``str(value)``) and + ``RedisCache.get_cache`` reads them back (``json.loads``, then ``ast.literal_eval``). The same literals sit in + ``litellm/rust/crates/redis-store/tests/deployment_health.rs``, where real Redis is checked to hold them + after the Rust store writes.""" + + states: dict[str, DeploymentHealthStateValue] + str_dict: str + ttl_seconds: float + redis_ttl_seconds: int + + +_WIRE_STATES: Final = _WireStates( + states={"dep-a": {"is_healthy": False, "timestamp": 1.5, "reason": "timeout"}}, + str_dict="{'dep-a': {'is_healthy': False, 'timestamp': 1.5, 'reason': 'timeout'}}", + ttl_seconds=90.0, + redis_ttl_seconds=90, +) + + +def test_the_wire_case_is_what_mains_codec_writes_and_reads() -> None: + assert str(_WIRE_STATES.states) == _WIRE_STATES.str_dict + assert ast.literal_eval(_WIRE_STATES.str_dict) == _WIRE_STATES.states + assert int(_WIRE_STATES.ttl_seconds) == _WIRE_STATES.redis_ttl_seconds + + +_LEFT_RECORDS: Final = TypeAdapter(dict[str, object]) + + +@pytest.mark.parametrize( + ("written", "read"), + [ + pytest.param( + "{'dep-a': {'is_healthy': False, 'timestamp': 1.5, 'reason': 'timeout', 'region': 'eu'}}", + {"dep-a": {"is_healthy": False, "timestamp": 1.5, "reason": "timeout"}}, + id="extra_fields_are_ignored", + ), + pytest.param( + "{'dep-a': {'is_healthy': 0, 'timestamp': 95, 'reason': 'timeout'}}", + {"dep-a": {"is_healthy": False, "timestamp": 95, "reason": "timeout"}}, + id="a_falsy_flag_is_unhealthy", + ), + pytest.param( + "{'dep-a': {'timestamp': 1.5}}", + {"dep-a": {"is_healthy": True, "timestamp": 1.5, "reason": ""}}, + id="a_missing_flag_is_healthy", + ), + pytest.param( + "{'dep-a': {'is_healthy': False, 'timestamp': 'soon', 'reason': 'timeout'}}", + {"dep-a": {"is_healthy": False, "timestamp": 0.0, "reason": "timeout"}}, + id="a_timestamp_that_is_no_number_is_stale", + ), + pytest.param( + "{'dep-a': {'is_healthy': False, 'timestamp': True, 'reason': 'timeout'}}", + {"dep-a": {"is_healthy": False, "timestamp": 0.0, "reason": "timeout"}}, + id="a_bool_timestamp_is_no_number", + ), + pytest.param( + "{'dep-a': 'not a record', 'dep-b': None, 'dep-c': {'is_healthy': False, 'timestamp': 1.5, " + "'reason': 'timeout'}}", + {"dep-c": {"is_healthy": False, "timestamp": 1.5, "reason": "timeout"}}, + id="what_is_no_record_is_dropped_alone", + ), + ], +) +def test_a_record_another_writer_left_is_read_the_way_mains_readers_read_it( + written: str, read: Mapping[str, DeploymentHealthStateValue] +) -> None: + """The same cases ``litellm/rust/crates/redis-store/tests/deployment_health.rs`` reads from real Redis, decoded + here from what main's ``ast.literal_eval`` reads back.""" + store: Final = connect(None).deployment_health() + + store.set_sync(_LEFT_RECORDS.validate_python(ast.literal_eval(written)), 90.0) + + assert store.get_sync() == read diff --git a/tests/unit/router_rust/test_routing_stores.py b/tests/unit/router_rust/test_routing_stores.py index 352e2819204..97a86edc5ce 100644 --- a/tests/unit/router_rust/test_routing_stores.py +++ b/tests/unit/router_rust/test_routing_stores.py @@ -6,6 +6,7 @@ from typing import Final import pytest from litellm.redis.stores.connection import connect +from litellm.redis.stores.deployment_health import DeploymentHealthStates, DeploymentHealthStore from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore from litellm.redis.stores.router_usage import RouterUsageStore from litellm.router_rust.routing_stores import ( @@ -19,6 +20,7 @@ from litellm.router_rust.routing_stores import ( UsageStep, ) from litellm.router_utils.cooldown_cache import CooldownCache +from litellm.router_utils.health_state_cache import DeploymentHealthCache from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns _ENTRY: Final[CooldownCacheValue] = { @@ -59,22 +61,19 @@ class SlowRedis(FakeCooldowns): class NoHealthStates: - async def async_get_states(self) -> object: + async def async_get_states(self) -> DeploymentHealthStates | None: return None - def get_states(self) -> object: + def get_states(self) -> DeploymentHealthStates | None: return None -@dataclass(frozen=True) -class CachedHealthStates: - raw: object - - async def async_get_states(self) -> object: - return self.raw - - def get_states(self) -> object: - return self.raw +def _health_cache(written: Mapping[str, object] | None) -> DeploymentHealthCache: + """The native health store holding what a writer left, decoded the way it is when read back from Redis.""" + handle: Final = connect(None).deployment_health() + if written is not None: + handle.set_sync(written, 90.0) + return DeploymentHealthCache(store=DeploymentHealthStore(handle), staleness_threshold=60.0) def _stores(health: HealthStates) -> RoutingStores: @@ -140,12 +139,12 @@ async def test_an_async_cooldown_step_waiting_on_redis_lets_other_requests_run( pytest.param(None, {}, id="nothing_cached"), ], ) -async def test_health_states_reach_router_core_as_the_python_filter_reads_them( - raw: object, states: Mapping[str, object] +async def test_health_states_reach_router_core_as_mains_python_filter_read_them( + raw: Mapping[str, object] | None, states: Mapping[str, object] ) -> None: - """router-core takes each state as a strict bool and float, so the leniency of ``_extract_unhealthy_ids`` - (a falsy flag is unhealthy, a missing one healthy, a bad timestamp stale) has to be applied before it crosses.""" - stores: Final = _stores(CachedHealthStates(raw)) + """router-core takes each state as a strict bool and float, so the store's decoder applies main's leniency + (a falsy flag is unhealthy, a missing one healthy, a bad timestamp stale) before a record crosses.""" + stores: Final = _stores(_health_cache(raw)) step: Final = CACHE_STEP.validate_python({"op": "health_states"}) assert isinstance(step, HealthStatesStep) diff --git a/tests/unit/router_utils/test_health_state_cache.py b/tests/unit/router_utils/test_health_state_cache.py index aa976bb1002..b40da9194a8 100644 --- a/tests/unit/router_utils/test_health_state_cache.py +++ b/tests/unit/router_utils/test_health_state_cache.py @@ -6,19 +6,24 @@ import time import pytest -from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.redis.stores.deployment_health import DeploymentHealthStates, 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 handle() -> NativeDeploymentHealth: + return NativeDeploymentHealth(None) @pytest.fixture -def health_cache(cache): - return DeploymentHealthCache(cache=cache, staleness_threshold=60.0) +def store(handle: NativeDeploymentHealth) -> DeploymentHealthStore: + return DeploymentHealthStore(handle) + + +@pytest.fixture +def health_cache(store): + return DeploymentHealthCache(store=store, staleness_threshold=60.0) def test_set_and_get_unhealthy_ids(health_cache): @@ -101,15 +106,19 @@ def test_mixed_stale_and_fresh(health_cache): assert result == {"deploy-2"} -def test_malformed_state_entries_are_skipped(health_cache): - """Non-dict entries in the cache should be skipped safely.""" +def test_malformed_state_entries_are_skipped( + health_cache: DeploymentHealthCache, handle: NativeDeploymentHealth +) -> None: + """Records another writer left that are no dict are dropped alone when the store reads them.""" now = time.time() - states = { - "deploy-1": {"is_healthy": False, "timestamp": now, "reason": "bad"}, - "deploy-2": "not_a_dict", # malformed - "deploy-3": None, # malformed - } - health_cache.set_deployment_health_states(states) + handle.set_sync( + { + "deploy-1": {"is_healthy": False, "timestamp": now, "reason": "bad"}, + "deploy-2": "not_a_dict", + "deploy-3": None, + }, + 90.0, + ) result = health_cache.get_unhealthy_deployment_ids() assert result == {"deploy-1"} @@ -120,13 +129,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 +143,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: DeploymentHealthStates | None = 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: DeploymentHealthStates | None = None + + @property + def has_redis(self) -> bool: + return True + + def attach(self, client: NativeRedisClient) -> bool: + return False + + async def get(self) -> DeploymentHealthStates | None: + return self.get_sync() + + def get_sync(self) -> DeploymentHealthStates | None: + return self._memory if self._memory is not None else self.get_shared_sync() + + def get_shared_sync(self) -> DeploymentHealthStates | None: + return None if self._redis.fail_get else self._redis.states + + def set_sync(self, states: DeploymentHealthStates, 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 +201,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 or {}) == {"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 +223,35 @@ 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"} + assert set(redis_fake.states or {}) == {"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"} +class _UnreachableHealthHandle: + has_redis: bool = True + + def attach(self, client: NativeRedisClient) -> bool: + return True + + async def get(self) -> DeploymentHealthStates | None: + raise ConnectionError("redis down") + + def get_sync(self) -> DeploymentHealthStates | None: + raise ConnectionError("redis down") + + def get_shared_sync(self) -> DeploymentHealthStates | None: + raise ConnectionError("redis down") + + def set_sync(self, states: DeploymentHealthStates, ttl_seconds: float) -> None: + raise ConnectionError("redis down") + + +@pytest.mark.asyncio +async def test_an_unreachable_store_reads_as_no_health_state_and_drops_writes(): + cache = DeploymentHealthCache(store=DeploymentHealthStore(_UnreachableHealthHandle()), staleness_threshold=60.0) + + cache.set_deployment_health_states({"deploy-1": {"is_healthy": False, "timestamp": time.time(), "reason": "x"}}) + + assert cache.get_states() is None + assert await cache.async_get_states() is None + assert cache.get_unhealthy_deployment_ids() == set() + assert await cache.async_get_unhealthy_deployment_ids() == set() 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)