diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 2764a6f0d4..530bea3d06 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -314,6 +314,89 @@ general_settings: health_check_details: False ``` +## Health Check Driven Routing + +By default, background health checks are observability-only — they populate the `/health` endpoint but don't affect routing. Unhealthy deployments still receive traffic until request failures trigger cooldown. + +With `enable_health_check_routing: true`, the router **excludes deployments that failed their last background health check** before selecting a candidate. This gives you proactive failover instead of reactive cooldown. + +### How it works + +1. Background health checks run on their configured interval +2. After each cycle, every deployment is marked healthy or unhealthy +3. On each incoming request, the router filters out unhealthy deployments **before** cooldown filtering and load balancing +4. If all deployments are unhealthy, the filter is bypassed (safety net — never causes a total outage) +5. If health state is stale (older than `health_check_staleness_threshold`), it is ignored + +### Quick start + +```yaml +model_list: + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY + - model_name: gpt-4 + litellm_params: + model: openai/gpt-4 + api_key: os.environ/OPENAI_API_KEY_SECONDARY + +general_settings: + background_health_checks: true + health_check_interval: 60 + enable_health_check_routing: true +``` + +### Configuration + +| Setting | Where | Default | Description | +|---------|-------|---------|-------------| +| `enable_health_check_routing` | `general_settings` | `false` | Enable/disable health-check-driven routing | +| `health_check_staleness_threshold` | `general_settings` | `health_check_interval * 2` | Seconds before health state is considered stale and ignored | +| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work | +| `health_check_interval` | `general_settings` | `300` | Seconds between health check cycles | + +### Interaction with cooldown + +Health check filtering and cooldown are **additive**. A deployment can be excluded by either mechanism: + +- **Health check filter** — proactive, runs on the configured interval, excludes deployments that failed the last check +- **Cooldown** — reactive, triggered by request failures, excludes deployments for a short TTL + +This means request failures still provide fast detection between health check intervals. + +### Staleness + +If a health check result is older than `health_check_staleness_threshold`, it is ignored and the deployment is treated as eligible. This prevents stale data from permanently excluding a deployment if the health check loop stops or slows down. + +The default staleness threshold is `health_check_interval * 2`. For a 60s interval, health state expires after 120s. + +### Example: custom staleness + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + health_check_staleness_threshold: 90 # ignore health state older than 90s +``` + +### Debugging + +Run the proxy with `--detailed_debug` and look for: + +``` +health_check_routing_state_updated healthy=3 unhealthy=1 +``` + +This is logged after each health check cycle when routing state is written. + +If the safety net triggers (all deployments unhealthy), you'll see: + +``` +All deployments marked unhealthy by health checks, bypassing health filter +``` + ## Health Check Timeout The health check timeout is set in `litellm/constants.py` and defaults to 60 seconds. diff --git a/litellm/constants.py b/litellm/constants.py index 423f01afac..252068bd7b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1402,6 +1402,9 @@ DEFAULT_SHARED_HEALTH_CHECK_TTL = int( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL = int( os.getenv("DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL", 60) ) # 1 minute - TTL for health check lock +DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER = ( + 2 # health state is stale after interval * this +) PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS = int( os.getenv("PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS", 9) ) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index a8d0e3e9af..3e05ee3c48 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -207,21 +207,65 @@ async def _perform_health_check( for is_healthy, model in zip(results, model_list): litellm_params = model["litellm_params"] + _model_id = (model.get("model_info") or {}).get("id") if isinstance(is_healthy, dict) and "error" not in is_healthy: - healthy_endpoints.append( - _clean_endpoint_data({**litellm_params, **is_healthy}, details) - ) + cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details) + if _model_id: + cleaned["model_id"] = _model_id + healthy_endpoints.append(cleaned) elif isinstance(is_healthy, dict): - unhealthy_endpoints.append( - _clean_endpoint_data({**litellm_params, **is_healthy}, details) - ) + cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details) + if _model_id: + cleaned["model_id"] = _model_id + unhealthy_endpoints.append(cleaned) else: - unhealthy_endpoints.append(_clean_endpoint_data(litellm_params, details)) + cleaned = _clean_endpoint_data(litellm_params, details) + if _model_id: + cleaned["model_id"] = _model_id + unhealthy_endpoints.append(cleaned) return healthy_endpoints, unhealthy_endpoints +def build_deployment_health_states( + healthy_endpoints: list, + unhealthy_endpoints: list, +) -> dict: + """ + Build a dict mapping deployment_id -> DeploymentHealthStateValue from + health check endpoint results. + + Each endpoint dict includes a 'model_id' field (added by _perform_health_check) + that maps back to the deployment's model_info.id. + + Used by the background health check loop to feed health state into + the router's DeploymentHealthCache for health-check-driven routing. + """ + now = time.time() + states: dict = {} + + for ep in healthy_endpoints: + model_id = ep.get("model_id") + if model_id: + states[model_id] = { + "is_healthy": True, + "timestamp": now, + "reason": "", + } + + for ep in unhealthy_endpoints: + model_id = ep.get("model_id") + if model_id: + states[model_id] = { + "is_healthy": False, + "timestamp": now, + "reason": "background_health_check_failed", + } + + return states + + def _update_litellm_params_for_health_check( model_info: dict, litellm_params: dict ) -> dict: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d3d2ceb53..28e613ef48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -480,11 +480,11 @@ from litellm.proxy.search_endpoints.search_tool_management import ( router as search_tool_management_router, ) from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router -from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, @@ -2112,6 +2112,37 @@ def _schedule_background_health_check_db_save( ) +def _write_health_state_to_router_cache( + healthy_endpoints: list, + unhealthy_endpoints: list, +) -> None: + """ + Write deployment health states to the router's health state cache + for health-check-driven routing. No-op if the feature is disabled. + """ + from litellm.proxy.health_check import build_deployment_health_states + + try: + if llm_router is None or not llm_router.enable_health_check_routing: + return + + states = build_deployment_health_states( + healthy_endpoints=healthy_endpoints, + unhealthy_endpoints=unhealthy_endpoints, + ) + if states: + llm_router.health_state_cache.set_deployment_health_states(states) + verbose_proxy_logger.debug( + "health_check_routing_state_updated healthy=%d unhealthy=%d", + sum(1 for s in states.values() if s.get("is_healthy")), + sum(1 for s in states.values() if not s.get("is_healthy")), + ) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to write health state to router cache: %s", str(e) + ) + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. @@ -2281,6 +2312,9 @@ async def _run_background_health_check(): unhealthy_endpoints, ) + # Write health state to router cache for health-check-driven routing + _write_health_state_to_router_cache(healthy_endpoints, unhealthy_endpoints) + await asyncio.sleep(health_check_interval) @@ -3048,6 +3082,8 @@ class ProxyConfig: general_settings = config.get("general_settings", {}) if general_settings is None: general_settings = {} + _enable_hc_routing = False + _hc_staleness = None if general_settings: ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### key_management_settings = general_settings.get( @@ -3227,13 +3263,21 @@ class ProxyConfig: "health_check_concurrency", None ) health_check_details = general_settings.get("health_check_details", True) + # Health-check-driven routing (opt-in, passes through to Router later) + _enable_hc_routing = general_settings.get( + "enable_health_check_routing", False + ) + _hc_staleness = general_settings.get( + "health_check_staleness_threshold", None + ) verbose_proxy_logger.info( - "background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s", + "background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s health_check_routing=%s", use_background_health_checks, use_shared_health_check, health_check_interval, health_check_concurrency, health_check_details, + _enable_hc_routing, ) ### RBAC ### @@ -3263,6 +3307,11 @@ class ProxyConfig: "cache_responses": litellm.cache is not None, # cache if user passed in cache values } + # Health-check-driven routing params (from general_settings) + if _enable_hc_routing: + router_params["enable_health_check_routing"] = True + if _hc_staleness is not None: + router_params["health_check_staleness_threshold"] = _hc_staleness ## MODEL LIST model_list = config.get("model_list", None) if model_list: diff --git a/litellm/router.py b/litellm/router.py index 5cd4f83778..6cc6bad9de 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -54,7 +54,11 @@ from litellm.caching.caching import ( RedisCache, RedisClusterCache, ) -from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE +from litellm.constants import ( + DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, + DEFAULT_MAX_LRU_CACHE_SIZE, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import ( @@ -113,6 +117,7 @@ from litellm.router_utils.handle_error import ( async_raise_no_deployment_exception, send_llm_exception_alert, ) +from litellm.router_utils.health_state_cache import DeploymentHealthCache from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) @@ -303,6 +308,8 @@ class Router: deployment_affinity_ttl_seconds: int = 3600, model_group_affinity_config: Optional[Dict[str, List[str]]] = None, ignore_invalid_deployments: bool = False, + enable_health_check_routing: bool = False, + health_check_staleness_threshold: Optional[int] = None, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -493,6 +500,13 @@ class Router: cache=self.cache, default_cooldown_time=self.cooldown_time ) self.disable_cooldowns = disable_cooldowns + self.enable_health_check_routing = enable_health_check_routing + _staleness = health_check_staleness_threshold or ( + DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER + ) + self.health_state_cache = DeploymentHealthCache( + cache=self.cache, staleness_threshold=float(_staleness) + ) self.failed_calls = ( InMemoryCache() ) # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown @@ -9154,6 +9168,14 @@ class Router: if isinstance(healthy_deployments, dict): return healthy_deployments + # Health-check-based filtering (before cooldown) + healthy_deployments = ( + await self._async_filter_health_check_unhealthy_deployments( + healthy_deployments=healthy_deployments, + parent_otel_span=parent_otel_span, + ) + ) + cooldown_deployments = await _async_get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) @@ -9585,6 +9607,13 @@ class Router: parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( request_kwargs ) + + # Health-check-based filtering (before cooldown) + healthy_deployments = self._filter_health_check_unhealthy_deployments( + healthy_deployments=healthy_deployments, + parent_otel_span=parent_otel_span, + ) + cooldown_deployments = _get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) @@ -9750,10 +9779,14 @@ class Router: llm_provider="", ) - # 4. Apply cooldown filtering + # 4. Apply health-check and cooldown filtering parent_otel_span: Optional[Span] = _get_parent_otel_span_from_kwargs( request_kwargs ) + pass_through_deployments = self._filter_health_check_unhealthy_deployments( + healthy_deployments=pass_through_deployments, + parent_otel_span=parent_otel_span, + ) cooldown_deployments = _get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) @@ -9875,6 +9908,67 @@ class Router: if deployment["model_info"]["id"] not in cooldown_set ] + async def _async_filter_health_check_unhealthy_deployments( + self, + healthy_deployments: List[Dict], + parent_otel_span: Optional[Span] = None, + ) -> List[Dict]: + """ + Filter out deployments marked unhealthy by background health checks. + No-op when enable_health_check_routing is False. + Returns all deployments if health state is unavailable, stale, or would + exclude every candidate (safety net). + """ + if not self.enable_health_check_routing: + return healthy_deployments + + unhealthy_ids = ( + await self.health_state_cache.async_get_unhealthy_deployment_ids( + parent_otel_span=parent_otel_span + ) + ) + if not unhealthy_ids: + return healthy_deployments + + filtered = [ + d for d in healthy_deployments if d["model_info"]["id"] not in unhealthy_ids + ] + + if not filtered: + verbose_router_logger.warning( + "All deployments marked unhealthy by health checks, bypassing health filter" + ) + return healthy_deployments + + return filtered + + def _filter_health_check_unhealthy_deployments( + self, + healthy_deployments: List[Dict], + parent_otel_span: Optional[Span] = None, + ) -> List[Dict]: + """Sync version of _async_filter_health_check_unhealthy_deployments.""" + if not self.enable_health_check_routing: + return healthy_deployments + + unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids( + parent_otel_span=parent_otel_span + ) + if not unhealthy_ids: + return healthy_deployments + + filtered = [ + d for d in healthy_deployments if d["model_info"]["id"] not in unhealthy_ids + ] + + if not filtered: + verbose_router_logger.warning( + "All deployments marked unhealthy by health checks, bypassing health filter" + ) + return healthy_deployments + + return filtered + def _filter_pass_through_deployments( self, healthy_deployments: List[Dict] ) -> List[Dict]: diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py new file mode 100644 index 0000000000..65b064f19d --- /dev/null +++ b/litellm/router_utils/health_state_cache.py @@ -0,0 +1,100 @@ +""" +Wrapper around router cache for health-check-driven routing. + +Stores per-deployment health state from background health checks +and exposes it for router candidate filtering. +""" + +import time +from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Union + +from typing_extensions import TypedDict + +from litellm import verbose_logger +from litellm.caching.caching import DualCache + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + Span = Union[_Span, Any] +else: + Span = Any + + +class DeploymentHealthStateValue(TypedDict): + is_healthy: bool + timestamp: float + reason: str + + +class DeploymentHealthCache: + """ + Cache for deployment health states produced by background health checks. + + Stores a single dict mapping deployment_id -> DeploymentHealthStateValue. + Staleness is enforced at read time: entries older than staleness_threshold + are treated as healthy (unknown). + """ + + CACHE_KEY = "litellm:health_check:deployment_health_state" + + def __init__(self, cache: DualCache, staleness_threshold: float): + self.cache = cache + self.staleness_threshold = staleness_threshold + + def set_deployment_health_states( + self, states: Dict[str, DeploymentHealthStateValue] + ) -> None: + """Bulk-write all deployment health states as a single cache entry.""" + try: + self.cache.set_cache( + key=self.CACHE_KEY, + value=states, + ttl=int(self.staleness_threshold * 1.5), + ) + except Exception as e: + verbose_logger.error( + "DeploymentHealthCache::set_deployment_health_states - Exception: %s", + str(e), + ) + + def _extract_unhealthy_ids(self, raw: Any) -> Set[str]: + """Given raw cache value, return set of non-stale unhealthy deployment IDs.""" + if not raw or not isinstance(raw, dict): + return set() + now = 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 - state.get("timestamp", 0)) < self.staleness_threshold + } + + async def async_get_unhealthy_deployment_ids( + self, parent_otel_span: Optional[Span] = None + ) -> Set[str]: + """Return set of deployment IDs currently marked unhealthy and not stale.""" + try: + raw = await self.cache.async_get_cache(key=self.CACHE_KEY) + return self._extract_unhealthy_ids(raw) + except Exception as e: + verbose_logger.debug( + "DeploymentHealthCache::async_get_unhealthy_deployment_ids - Exception: %s", + str(e), + ) + return set() + + def get_unhealthy_deployment_ids( + self, parent_otel_span: Optional[Span] = None + ) -> Set[str]: + """Sync version: return set of deployment IDs currently marked unhealthy and not stale.""" + try: + raw = self.cache.get_cache(key=self.CACHE_KEY) + return self._extract_unhealthy_ids(raw) + except Exception as e: + verbose_logger.debug( + "DeploymentHealthCache::get_unhealthy_deployment_ids - Exception: %s", + str(e), + ) + return set() diff --git a/tests/test_litellm/router_utils/test_health_check_routing.py b/tests/test_litellm/router_utils/test_health_check_routing.py new file mode 100644 index 0000000000..f40144b44c --- /dev/null +++ b/tests/test_litellm/router_utils/test_health_check_routing.py @@ -0,0 +1,197 @@ +""" +Tests for health-check-driven routing filter in the Router. +""" + +import time + +import pytest + +from litellm.caching.caching import DualCache +from litellm.router_utils.health_state_cache import DeploymentHealthCache + + +def _make_deployment(model_id: str, model_name: str = "gpt-4") -> dict: + """Helper to create a deployment dict for testing.""" + return { + "model_name": model_name, + "litellm_params": {"model": model_name, "api_key": "fake"}, + "model_info": {"id": model_id}, + } + + +def _make_health_cache( + unhealthy_ids: set = None, staleness_threshold: float = 60.0 +) -> DeploymentHealthCache: + """Create a health cache pre-populated with unhealthy deployment IDs.""" + cache = DualCache() + health_cache = DeploymentHealthCache( + cache=cache, staleness_threshold=staleness_threshold + ) + if unhealthy_ids: + now = time.time() + states = {} + for uid in unhealthy_ids: + states[uid] = { + "is_healthy": False, + "timestamp": now, + "reason": "test_unhealthy", + } + health_cache.set_deployment_health_states(states) + return health_cache + + +class TestFilterHealthCheckUnhealthyDeployments: + """Test the sync filter method.""" + + def _make_router_like(self, enable: bool, health_cache: DeploymentHealthCache): + """Create a minimal object that behaves like Router for filter testing.""" + + class FakeRouter: + def __init__(self): + self.enable_health_check_routing = enable + self.health_state_cache = health_cache + + # Import the actual method and bind it + from litellm.router import Router + + fake = FakeRouter() + # Use the unbound method + fake._filter_health_check_unhealthy_deployments = ( + Router._filter_health_check_unhealthy_deployments.__get__(fake, FakeRouter) + ) + return fake + + def test_filter_removes_unhealthy_deployments(self): + """Unhealthy deployments should be removed from candidates.""" + health_cache = _make_health_cache(unhealthy_ids={"deploy-2"}) + router = self._make_router_like(enable=True, health_cache=health_cache) + + deployments = [ + _make_deployment("deploy-1"), + _make_deployment("deploy-2"), + _make_deployment("deploy-3"), + ] + result = router._filter_health_check_unhealthy_deployments(deployments) + assert len(result) == 2 + assert all(d["model_info"]["id"] != "deploy-2" for d in result) + + def test_filter_noop_when_disabled(self): + """When enable_health_check_routing=False, filter should be a no-op.""" + health_cache = _make_health_cache(unhealthy_ids={"deploy-1"}) + router = self._make_router_like(enable=False, health_cache=health_cache) + + deployments = [ + _make_deployment("deploy-1"), + _make_deployment("deploy-2"), + ] + result = router._filter_health_check_unhealthy_deployments(deployments) + assert len(result) == 2 # no filtering + + def test_filter_returns_all_when_all_unhealthy(self): + """Safety net: if ALL deployments are unhealthy, return all (don't cause outage).""" + health_cache = _make_health_cache( + unhealthy_ids={"deploy-1", "deploy-2", "deploy-3"} + ) + router = self._make_router_like(enable=True, health_cache=health_cache) + + deployments = [ + _make_deployment("deploy-1"), + _make_deployment("deploy-2"), + _make_deployment("deploy-3"), + ] + result = router._filter_health_check_unhealthy_deployments(deployments) + assert len(result) == 3 # all returned, safety net + + def test_filter_returns_all_when_cache_empty(self): + """When cache is empty, all deployments should pass through.""" + health_cache = _make_health_cache() # empty + router = self._make_router_like(enable=True, health_cache=health_cache) + + deployments = [ + _make_deployment("deploy-1"), + _make_deployment("deploy-2"), + ] + result = router._filter_health_check_unhealthy_deployments(deployments) + assert len(result) == 2 + + +class TestAsyncFilterHealthCheckUnhealthyDeployments: + """Test the async filter method.""" + + def _make_router_like(self, enable: bool, health_cache: DeploymentHealthCache): + from litellm.router import Router + + class FakeRouter: + def __init__(self): + self.enable_health_check_routing = enable + self.health_state_cache = health_cache + + fake = FakeRouter() + fake._async_filter_health_check_unhealthy_deployments = ( + Router._async_filter_health_check_unhealthy_deployments.__get__( + fake, FakeRouter + ) + ) + return fake + + @pytest.mark.asyncio + async def test_async_filter_removes_unhealthy(self): + """Async version: unhealthy deployments removed.""" + health_cache = _make_health_cache(unhealthy_ids={"deploy-2"}) + router = self._make_router_like(enable=True, health_cache=health_cache) + + deployments = [ + _make_deployment("deploy-1"), + _make_deployment("deploy-2"), + _make_deployment("deploy-3"), + ] + result = await router._async_filter_health_check_unhealthy_deployments( + healthy_deployments=deployments + ) + assert len(result) == 2 + assert all(d["model_info"]["id"] != "deploy-2" for d in result) + + @pytest.mark.asyncio + async def test_async_filter_safety_net(self): + """Async version: safety net when all unhealthy.""" + health_cache = _make_health_cache(unhealthy_ids={"deploy-1", "deploy-2"}) + router = self._make_router_like(enable=True, health_cache=health_cache) + + deployments = [ + _make_deployment("deploy-1"), + _make_deployment("deploy-2"), + ] + result = await router._async_filter_health_check_unhealthy_deployments( + healthy_deployments=deployments + ) + assert len(result) == 2 # safety net + + +class TestBuildDeploymentHealthStates: + """Test the build_deployment_health_states function.""" + + def test_builds_states_from_endpoints(self): + from litellm.proxy.health_check import build_deployment_health_states + + healthy = [{"model": "gpt-4", "model_id": "deploy-1"}] + unhealthy = [{"model": "gpt-4", "model_id": "deploy-2", "error": "timeout"}] + + states = build_deployment_health_states(healthy, unhealthy) + assert states["deploy-1"]["is_healthy"] is True + assert states["deploy-2"]["is_healthy"] is False + + def test_no_model_id_skipped(self): + from litellm.proxy.health_check import build_deployment_health_states + + healthy = [{"model": "gpt-4"}] # no model_id + unhealthy = [{"model": "gpt-4", "model_id": "deploy-2"}] + + states = build_deployment_health_states(healthy, unhealthy) + assert "deploy-1" not in states + assert states["deploy-2"]["is_healthy"] is False + + def test_empty_endpoints(self): + from litellm.proxy.health_check import build_deployment_health_states + + states = build_deployment_health_states([], []) + assert states == {} diff --git a/tests/test_litellm/router_utils/test_health_state_cache.py b/tests/test_litellm/router_utils/test_health_state_cache.py new file mode 100644 index 0000000000..1af61e899b --- /dev/null +++ b/tests/test_litellm/router_utils/test_health_state_cache.py @@ -0,0 +1,113 @@ +""" +Tests for DeploymentHealthCache - the cache layer for health-check-driven routing. +""" + +import time + +import pytest + +from litellm.caching.caching import DualCache +from litellm.router_utils.health_state_cache import DeploymentHealthCache + + +@pytest.fixture +def cache(): + return DualCache() + + +@pytest.fixture +def health_cache(cache): + return DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + + +def test_set_and_get_unhealthy_ids(health_cache): + """Write states, verify unhealthy set is returned correctly.""" + now = time.time() + states = { + "deploy-1": {"is_healthy": True, "timestamp": now, "reason": ""}, + "deploy-2": {"is_healthy": False, "timestamp": now, "reason": "check_failed"}, + "deploy-3": {"is_healthy": False, "timestamp": now, "reason": "timeout"}, + } + health_cache.set_deployment_health_states(states) + result = health_cache.get_unhealthy_deployment_ids() + assert result == {"deploy-2", "deploy-3"} + + +@pytest.mark.asyncio +async def test_async_get_unhealthy_ids(health_cache): + """Async version of set and get.""" + now = time.time() + states = { + "deploy-1": {"is_healthy": True, "timestamp": now, "reason": ""}, + "deploy-2": {"is_healthy": False, "timestamp": now, "reason": "check_failed"}, + } + health_cache.set_deployment_health_states(states) + result = await health_cache.async_get_unhealthy_deployment_ids() + assert result == {"deploy-2"} + + +def test_staleness_filtering(health_cache): + """Entries older than staleness_threshold should be ignored.""" + old_time = time.time() - 120 # 2 minutes ago, threshold is 60s + states = { + "deploy-1": { + "is_healthy": False, + "timestamp": old_time, + "reason": "check_failed", + }, + } + health_cache.set_deployment_health_states(states) + result = health_cache.get_unhealthy_deployment_ids() + assert result == set() # stale entry should be ignored + + +def test_empty_cache_returns_empty_set(health_cache): + """No data in cache should return empty set.""" + result = health_cache.get_unhealthy_deployment_ids() + assert result == set() + + +def test_all_healthy_returns_empty_set(health_cache): + """All healthy deployments should return empty set.""" + now = time.time() + states = { + "deploy-1": {"is_healthy": True, "timestamp": now, "reason": ""}, + "deploy-2": {"is_healthy": True, "timestamp": now, "reason": ""}, + } + health_cache.set_deployment_health_states(states) + result = health_cache.get_unhealthy_deployment_ids() + assert result == set() + + +def test_mixed_stale_and_fresh(health_cache): + """Only fresh unhealthy entries should be returned.""" + now = time.time() + old_time = now - 120 # stale + states = { + "deploy-1": { + "is_healthy": False, + "timestamp": old_time, + "reason": "stale", + }, + "deploy-2": { + "is_healthy": False, + "timestamp": now, + "reason": "fresh", + }, + } + health_cache.set_deployment_health_states(states) + result = health_cache.get_unhealthy_deployment_ids() + assert result == {"deploy-2"} + + +def test_malformed_state_entries_are_skipped(health_cache): + """Non-dict entries in the cache should be skipped safely.""" + 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) + result = health_cache.get_unhealthy_deployment_ids() + assert result == {"deploy-1"}