tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-10 16:22:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #27003 from BerriAI/litellm_health-endpoint-non200-on-failure

Browse Source 
fix(health): return 503 when targeted model is unhealthy or DB is disconnected
This commit is contained in:
ryan-crabbe-berri
2026-05-01 14:23:37 -07:00
committed by GitHub
parent 42122b83f5 038b180315
commit 610f79dc03
2 changed files with 508 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/health_endpoints/_health_endpoints.py
+80 -6
View File
@@ -776,6 +776,45 @@ def _strip_admin_only_fields_from_health_result(result: dict) -> dict:
return out
def _resolve_targeted_model_ids(
model_list: list, model: Optional[str], model_id: Optional[str]
) -> Optional[set]:
"""
Resolve a ``/health`` ``model`` / ``model_id`` query param to the set of
deployment IDs the response should be scoped to.
Mirrors the live-path semantics in ``perform_health_check()``: ``model``
matches either the deployment's ``model_name`` alias or its
``litellm_params.model`` provider string. ``model_id`` matches
``model_info.id``.
Both query params are validated against the supplied ``model_list``.
Callers pass an already-scoped list (filtered to the caller's allowed
models for non-admins, full list for admins), so a ``model_id`` that
isn't present resolves to an empty set rather than a single-element
set — preventing a non-admin from reading another deployment's cached
health entry by guessing its ID.
Returns ``None`` when no targeting is requested — callers should treat
that as "no filter."
"""
if not model and not model_id:
return None
target_ids: set = set()
for m in model_list:
deployment_id = (m.get("model_info") or {}).get("id")
if not deployment_id:
continue
if model_id and deployment_id == model_id:
target_ids.add(deployment_id)
continue
if model:
litellm_model = (m.get("litellm_params") or {}).get("model")
if m.get("model_name") == model or litellm_model == model:
target_ids.add(deployment_id)
return target_ids
def _filter_health_check_results_by_model_ids(
results: dict, allowed_model_ids: set
) -> dict:
@@ -926,6 +965,7 @@ async def health_endpoint(
)
is_admin = _is_proxy_admin(user_api_key_dict)
model_specific_request = bool(model or model_id)
def _post_process(result: dict) -> dict:
# api_base / api_version reveal which provider/region/internal host the
@@ -933,6 +973,12 @@ async def health_endpoint(
# still see model/model_id and the healthy/unhealthy status. We also
# set a header so non-admin clients that previously parsed those
# fields can detect the change programmatically.
# When a caller asked about a specific model/model_id and zero
# endpoints came back healthy, surface that as a 503 so monitoring
# systems can rely on the HTTP status instead of having to parse the
# body. The body shape is unchanged.
if model_specific_request and result.get("healthy_count", 0) == 0:
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
if is_admin:
return result
response.headers["Litellm-Health-Field-Notice"] = (
@@ -975,16 +1021,29 @@ async def health_endpoint(
m for m in _llm_model_list if m.get("model_name") in allowed_models
]
if use_background_health_checks:
# The cached background result covers every model. When the
# caller targets a specific model/model_id we have to narrow the
# cache to that deployment before _post_process evaluates
# healthy_count, otherwise an unhealthy "foo" combined with any
# other healthy model would still report healthy_count > 0 and
# the targeted-503 path would never fire.
targeted_ids = _resolve_targeted_model_ids(_llm_model_list, model, model_id)
if len(user_api_key_dict.models) > 0:
allowed_model_ids = {
(m.get("model_info") or {}).get("id")
for m in _llm_model_list
if (m.get("model_info") or {}).get("id")
}
filtered = _filter_health_check_results_by_model_ids(
health_check_results, allowed_model_ids
# _llm_model_list is already scoped to the caller's allowed
# model_names above, so targeted_ids is implicitly the
# intersection of "targeted" and "allowed."
filter_ids = (
targeted_ids if targeted_ids is not None else allowed_model_ids
)
if not allowed_model_ids:
filtered = _filter_health_check_results_by_model_ids(
health_check_results, filter_ids
)
if targeted_ids is None and not allowed_model_ids:
# Caller has accessible model_names but none of the
# matching deployments expose a model_info.id, so the
# cache filter (which keys on model_id) drops every
@@ -1005,6 +1064,15 @@ async def health_endpoint(
"to populate model_info.id for these models."
]
return _post_process(filtered)
if targeted_ids is not None:
# Admin caller targeting a specific model: filter the cache
# so the response (and the targeted-503 check) reflects only
# that deployment, not the global aggregate.
return _post_process(
_filter_health_check_results_by_model_ids(
health_check_results, targeted_ids
)
)
return _post_process(health_check_results)
else:
router_result = await _perform_health_check_and_save(
@@ -1384,7 +1452,7 @@ def callback_name(callback):
tags=["health"],
dependencies=[Depends(user_api_key_auth)],
)
async def health_readiness():
async def health_readiness(response: Response):
"""
Unprotected endpoint for checking if worker can receive requests
"""
@@ -1417,8 +1485,8 @@ async def health_readiness():
try:
index_info = await litellm.cache.cache._index_info()
except Exception as e:
index_info = "index does not exist - error: " + str(e)
cache_type = {"type": cache_type, "index_info": index_info}
index_info = "index does not exist - error: " + str(e) # type: ignore[assignment]
cache_type = {"type": cache_type, "index_info": index_info} # type: ignore[assignment]
# check log level
log_level_name = logging.getLevelName(verbose_logger.getEffectiveLevel())
@@ -1427,6 +1495,12 @@ async def health_readiness():
# check DB
if prisma_client is not None: # if db passed in, check if it's connected
db_health_status = await _db_health_readiness_check()
# A configured DB that is not reachable means the worker cannot
# serve requests that depend on persisted state (keys, budgets,
# spend logs). Return 503 so orchestrators take this pod out of
# rotation; "Not connected" (no DB configured at all) stays 200.
if db_health_status["status"] != "connected":
response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
return {
"status": "healthy",
"db": db_health_status["status"],
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
+428 -4
View File
@@ -927,8 +927,14 @@ async def test_health_endpoint_filters_background_cache_by_user_access():
):
from fastapi import Response
# Pass model=None, model_id=None explicitly: direct calls to the
# handler skip FastAPI's Query() resolution, so unspecified params
# would otherwise carry the Query() sentinel (which is truthy).
result = await health_endpoint(
response=Response(), user_api_key_dict=user_api_key_dict
response=Response(),
user_api_key_dict=user_api_key_dict,
model=None,
model_id=None,
)
# Sanity: the source cache had two entries before scoping; the scoping
@@ -1016,10 +1022,16 @@ async def test_health_endpoint_admin_sees_routing_fields_non_admin_does_not():
admin_response = Response()
non_admin_response = Response()
admin_result = await health_endpoint(
response=admin_response, user_api_key_dict=admin_key
response=admin_response,
user_api_key_dict=admin_key,
model=None,
model_id=None,
)
non_admin_result = await health_endpoint(
response=non_admin_response, user_api_key_dict=non_admin_key
response=non_admin_response,
user_api_key_dict=non_admin_key,
model=None,
model_id=None,
)
finally:
for p in common_patches:
@@ -1113,7 +1125,10 @@ async def test_health_endpoint_warns_when_scoped_models_lack_model_id():
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
):
result = await health_endpoint(
response=Response(), user_api_key_dict=user_api_key_dict
response=Response(),
user_api_key_dict=user_api_key_dict,
model=None,
model_id=None,
)
assert result["healthy_count"] == 0
@@ -1125,6 +1140,415 @@ async def test_health_endpoint_warns_when_scoped_models_lack_model_id():
assert any("model_info.id" in w for w in result["warnings"])
@pytest.mark.asyncio
async def test_health_endpoint_blocks_cross_scope_model_id_under_background_cache():
"""
A non-admin scoped to model-a must not be able to read model-b's cached
health entry by guessing its model_id. Before the fix,
_resolve_targeted_model_ids returned {model_id} unconditionally, so the
cache filter was driven by an unvalidated ID and the global cache
leaked id-b's entry to the caller.
"""
from fastapi import Response
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-a"},
},
{
"model_name": "model-b", # caller has no access
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-b"},
},
]
cached_results = {
"healthy_endpoints": [
{
"model": "openai/gpt-4o",
"model_id": "id-b",
"api_base": "https://leaky-internal.test",
},
],
"unhealthy_endpoints": [],
"healthy_count": 1,
"unhealthy_count": 0,
}
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-scoped",
models=["model-a"],
)
response = Response()
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
# llm_router None here means the model_id 404 lookup short-circuits;
# we patch _llm_model_list directly instead to drive the cache path.
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
):
# Calling with model="model-b" rather than model_id="id-b" because
# the model_id branch raises 404 when llm_router is None. The bug
# being verified is the same: targeted resolver must drop entries
# not in the caller's scoped model_list. With the fix, the result
# has no leaked endpoints and the targeted-503 path fires.
result = await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model="model-b",
model_id=None,
)
leaked_ids = {ep.get("model_id") for ep in result.get("healthy_endpoints", [])}
leaked_ids |= {ep.get("model_id") for ep in result.get("unhealthy_endpoints", [])}
assert (
"id-b" not in leaked_ids
), "background cache leaked an out-of-scope deployment to a scoped caller"
assert result["healthy_count"] == 0
assert response.status_code == 503
@pytest.mark.asyncio
async def test_health_endpoint_503_for_targeted_unhealthy_model_under_background_cache_admin():
"""
With background_health_checks enabled, an admin calling /health?model=foo
must get 503 when foo specifically has zero healthy endpoints — even if
other unrelated models in the cache are healthy. Without the cache-path
filter, the global healthy_count would mask the targeted failure.
"""
from fastapi import Response
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a", # the unhealthy target
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-a"},
},
{
"model_name": "model-b", # an unrelated healthy model
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-b"},
},
]
cached_results = {
"healthy_endpoints": [
{"model": "openai/gpt-4o", "model_id": "id-b"},
],
"unhealthy_endpoints": [
{"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"},
],
"healthy_count": 1,
"unhealthy_count": 1,
}
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-admin",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
response = Response()
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", True),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", cached_results),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
):
result = await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model="model-a",
model_id=None,
)
assert response.status_code == 503
# Body must be scoped to the targeted model — not the global cache.
assert result["healthy_count"] == 0
assert result["unhealthy_count"] == 1
returned_ids = {ep["model_id"] for ep in result.get("unhealthy_endpoints", [])}
assert returned_ids == {"id-a"}
@pytest.mark.asyncio
async def test_health_endpoint_returns_503_when_requested_model_has_no_healthy_endpoints():
"""
/health?model=foo must return 503 when the targeted model resolves but
has zero healthy endpoints. Body shape stays the same so existing
parsers still work; only the HTTP status changes.
"""
from fastapi import Response
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {
"model": "openai/gpt-4o",
"api_base": "https://example-a.test",
},
"model_info": {"id": "id-a"},
},
]
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
async def fake_perform(**kwargs):
return {
"healthy_endpoints": [],
"unhealthy_endpoints": [
{
"model": "openai/gpt-4o",
"model_id": "id-a",
"error": "boom",
}
],
"healthy_count": 0,
"unhealthy_count": 1,
}
response = Response()
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", {}),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
patch(
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
side_effect=fake_perform,
),
):
result = await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model="model-a",
)
assert response.status_code == 503
assert result["healthy_count"] == 0
assert result["unhealthy_count"] == 1
@pytest.mark.asyncio
async def test_health_endpoint_returns_200_when_requested_model_has_healthy_endpoints():
"""
/health?model=foo with a healthy endpoint must keep returning the
default 200. Verifies the 503 path doesn't fire when healthy_count > 0.
"""
from fastapi import Response
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-a"},
},
]
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
async def fake_perform(**kwargs):
return {
"healthy_endpoints": [{"model": "openai/gpt-4o", "model_id": "id-a"}],
"unhealthy_endpoints": [],
"healthy_count": 1,
"unhealthy_count": 0,
}
response = Response()
# Default Response() exposes status_code as None; the endpoint should
# leave it alone for the healthy path.
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", {}),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
patch(
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
side_effect=fake_perform,
),
):
await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model="model-a",
)
assert response.status_code == 200
@pytest.mark.asyncio
async def test_health_endpoint_no_model_param_returns_200_even_when_zero_healthy():
"""
The non-targeted /health (no model / model_id query) preserves the
legacy 200 behavior even when healthy_count == 0. Existing K8s probes
and dashboards depend on this; only the targeted call became 5xx-aware.
"""
from fastapi import Response
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.health_endpoints._health_endpoints import health_endpoint
full_model_list = [
{
"model_name": "model-a",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "id-a"},
},
]
user_api_key_dict = UserAPIKeyAuth(
api_key="hashed-test-key",
user_role=LitellmUserRoles.PROXY_ADMIN,
)
async def fake_perform(**kwargs):
return {
"healthy_endpoints": [],
"unhealthy_endpoints": [
{"model": "openai/gpt-4o", "model_id": "id-a", "error": "boom"}
],
"healthy_count": 0,
"unhealthy_count": 1,
}
response = Response()
with (
patch("litellm.proxy.proxy_server.llm_model_list", full_model_list),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.prisma_client", None),
patch("litellm.proxy.proxy_server.use_background_health_checks", False),
patch("litellm.proxy.proxy_server.user_model", None),
patch("litellm.proxy.proxy_server.health_check_results", {}),
patch("litellm.proxy.proxy_server.health_check_details", True),
patch("litellm.proxy.proxy_server.health_check_concurrency", 1),
patch(
"litellm.proxy.health_endpoints._health_endpoints._perform_health_check_and_save",
side_effect=fake_perform,
),
):
# Pass model=None, model_id=None explicitly: when invoked through
# FastAPI, the Query(None) defaults resolve to None, but direct
# function calls in unit tests receive Query() sentinel objects
# (which are truthy). The explicit None mirrors production routing.
await health_endpoint(
response=response,
user_api_key_dict=user_api_key_dict,
model=None,
model_id=None,
)
assert response.status_code == 200
@pytest.mark.asyncio
async def test_health_readiness_returns_503_when_db_disconnected():
"""
When a Prisma client is configured but its health_check fails, the
readiness probe should mark the worker as unhealthy via the HTTP
status — not just a body field — so K8s removes the pod from the
Service endpoints.
"""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import health_readiness
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope"))
mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope"))
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
response = Response()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await health_readiness(response=response)
assert response.status_code == 503
assert result["db"] == "disconnected"
assert result["status"] == "healthy" # body shape unchanged for back-compat
@pytest.mark.asyncio
async def test_health_readiness_returns_200_when_db_connected():
"""Happy path: connected DB keeps the legacy 200."""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import health_readiness
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "unknown",
"last_updated": datetime.now() - timedelta(seconds=60),
}
response = Response()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await health_readiness(response=response)
assert response.status_code == 200
assert result["db"] == "connected"
@pytest.mark.asyncio
async def test_health_readiness_returns_200_when_no_db_configured():
"""
`prisma_client is None` means the operator chose not to use a DB. That
is a valid configuration — the worker should still report ready. We
only flip to 503 when a DB *was* configured but is unreachable.
"""
from fastapi import Response
from litellm.proxy.health_endpoints._health_endpoints import health_readiness
response = Response()
with patch("litellm.proxy.proxy_server.prisma_client", None):
result = await health_readiness(response=response)
assert response.status_code == 200
assert result["db"] == "Not connected"
def test_clean_endpoint_data_strips_credentials_keeps_routing_fields():
"""
_clean_endpoint_data() drops credentials but leaves api_base /