diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index 29ceb783e7e..2ecee5095b8 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -217,6 +217,7 @@ class SharedHealthCheckManager: return ( cached_results.get("healthy_endpoints", []), cached_results.get("unhealthy_endpoints", []), + {}, ) # No recent cache, try to acquire lock diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fdcbd6fba25..3ab8832acea 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2109,6 +2109,7 @@ def _schedule_background_health_check_db_save( def _write_health_state_to_router_cache( healthy_endpoints: list, unhealthy_endpoints: list, + exceptions_by_model_id: Optional[dict] = None, ) -> None: """ Write deployment health states to the router's health state cache @@ -2120,6 +2121,8 @@ def _write_health_state_to_router_cache( increment_deployment_failures_for_current_minute, ) + _exceptions: dict = exceptions_by_model_id or {} + try: if llm_router is None or not llm_router.enable_health_check_routing: return @@ -2131,7 +2134,10 @@ def _write_health_state_to_router_cache( _effective_unhealthy = [ ep for ep in unhealthy_endpoints - if getattr(ep.get("exception"), "status_code", 500) not in (429, 408) + if getattr( + _exceptions.get(ep.get("model_id")), "status_code", 500 + ) + not in (429, 408) ] states = build_deployment_health_states( @@ -2151,7 +2157,7 @@ def _write_health_state_to_router_cache( if not model_id: continue - original_exception = endpoint.get("exception") + original_exception = _exceptions.get(model_id) if original_exception is None: continue diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py index b2d298bc45d..2add9604255 100644 --- a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -282,7 +282,7 @@ class TestWriteHealthStateIntegration: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "timeout", "exception": timeout_exc}, + {"model_id": "deploy-1", "error": "timeout"}, ] healthy_endpoints = [ {"model_id": "deploy-2"}, @@ -295,6 +295,7 @@ class TestWriteHealthStateIntegration: _write_health_state_to_router_cache( healthy_endpoints=healthy_endpoints, unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, ) mock_cooldown.assert_called_once_with( litellm_router_instance=router, @@ -316,7 +317,7 @@ class TestWriteHealthStateIntegration: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "unknown failure"}, # no "exception" key + {"model_id": "deploy-1", "error": "unknown failure"}, ] with patch.object(proxy_module, "llm_router", router): @@ -326,6 +327,7 @@ class TestWriteHealthStateIntegration: _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + # no exceptions_by_model_id → cooldown should not fire ) mock_cooldown.assert_not_called() @@ -345,7 +347,7 @@ class TestWriteHealthStateIntegration: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "rate limited", "exception": rate_exc}, + {"model_id": "deploy-1", "error": "rate limited"}, ] with patch.object(proxy_module, "llm_router", router): @@ -358,6 +360,7 @@ class TestWriteHealthStateIntegration: _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, ) mock_increment.assert_called_once_with( litellm_router_instance=router, @@ -557,7 +560,7 @@ class TestHealthCheckIgnoreTransientErrors: assert getattr(rate_exc, "status_code", None) == 429 unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "rate limited", "exception": rate_exc}, + {"model_id": "deploy-1", "error": "rate limited"}, ] with patch.object(proxy_module, "llm_router", router): @@ -570,6 +573,7 @@ class TestHealthCheckIgnoreTransientErrors: _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, ) mock_cooldown.assert_not_called() mock_increment.assert_not_called() @@ -591,7 +595,7 @@ class TestHealthCheckIgnoreTransientErrors: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "timeout", "exception": timeout_exc}, + {"model_id": "deploy-1", "error": "timeout"}, ] with patch.object(proxy_module, "llm_router", router): @@ -601,6 +605,7 @@ class TestHealthCheckIgnoreTransientErrors: _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, ) mock_cooldown.assert_not_called() @@ -621,7 +626,7 @@ class TestHealthCheckIgnoreTransientErrors: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "auth failed", "exception": auth_exc}, + {"model_id": "deploy-1", "error": "auth failed"}, ] with patch.object(proxy_module, "llm_router", router): @@ -631,6 +636,7 @@ class TestHealthCheckIgnoreTransientErrors: _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": auth_exc}, ) mock_cooldown.assert_called_once() @@ -651,13 +657,14 @@ class TestHealthCheckIgnoreTransientErrors: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "rate limited", "exception": rate_exc}, + {"model_id": "deploy-1", "error": "rate limited"}, ] with patch.object(proxy_module, "llm_router", router): _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, ) # Health state cache should have NO entry for deploy-1 @@ -682,7 +689,7 @@ class TestHealthCheckIgnoreTransientErrors: ) unhealthy_endpoints = [ - {"model_id": "deploy-1", "error": "rate limited", "exception": rate_exc}, + {"model_id": "deploy-1", "error": "rate limited"}, ] with patch.object(proxy_module, "llm_router", router): @@ -692,5 +699,6 @@ class TestHealthCheckIgnoreTransientErrors: _write_health_state_to_router_cache( healthy_endpoints=[], unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, ) mock_cooldown.assert_called_once()