litellm/tests/litellm_utils_tests/test_health_check.py
ishaan-berri 51876292a0
Litellm ishaan april4 2 (#25150)
* feat(router): integrate allowed_fails_policy into health check failures (#24988)

* feat(router): integrate allowed_fails_policy into health check failures

Health check failures now increment the same per-deployment failure
counters used by allowed_fails_policy, so users can control how many
health check failures of each error type are required before a
deployment enters cooldown.

- ahealth_check() preserves the original exception in its return dict
- run_with_timeout() returns a litellm.Timeout on health check timeout
- _perform_health_check() propagates exceptions to unhealthy endpoints
- _write_health_state_to_router_cache() calls _set_cooldown_deployments
  for each unhealthy endpoint that has an exception
- When allowed_fails_policy is set, the binary health check filter is
  bypassed so cooldown is the sole routing exclusion mechanism
- Safety net: if all deployments are in cooldown with
  enable_health_check_routing=True, the cooldown filter is bypassed

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* feat(router): add health_check_ignore_transient_errors flag

When enabled, health check failures with 429 (rate limit) or 408 (timeout)
status codes are skipped from the cooldown pipeline. These are transient
load issues, not broken deployments. Auth errors (401), 404, and 5xx errors
still increment counters and trigger cooldown as before.

Config (general_settings):
  health_check_ignore_transient_errors: true

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(router): also exclude 429/408 from health state cache when ignore_transient_errors set

The previous fix only skipped cooldown counter increments. The health state
cache was still marking 429/408 endpoints as is_healthy=False, causing the
binary health check filter to exclude them from routing.

Now, when health_check_ignore_transient_errors=True, 429/408 endpoints are
also excluded from the unhealthy list passed to build_deployment_health_states(),
so the binary filter treats them as unaffected (not unhealthy).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* docs(router): add health check driven routing guide

New standalone page covering the full health check routing feature:
allowed_fails_policy integration, health_check_ignore_transient_errors,
architecture SVG, step-by-step setup, and gotchas (TTL, AllowedFails semantics).

Replaces the inline section in health.md with a link to the new page.
Added to the Routing & Load Balancing sidebar.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(health-check-routing): fix three CI failures

- Add "exception" to ILLEGAL_DISPLAY_PARAMS in health_check.py so the
  exception object is stripped before the health endpoint serializes
  results to JSON (fixes TypeError: 'URL' object is not iterable)
- Add allowed_fails_policy = None to FakeRouter stubs in
  test_router_health_check_routing.py (fixes AttributeError)
- Add health_check_ignore_transient_errors to config_settings.md router
  settings reference table (fixes documentation test)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix litellm/tests/proxy_unit_tests/test_proxy_server.py

* fix(router): address greptile review comments

- Narrow cooldown safety-net bypass: only fires when allowed_fails_policy
  is set (cooldown is health-check driven). Without a policy, cooldowns
  are from real request failures and must not be bypassed.
- Restore cooldown deployments DEBUG log that was accidentally removed.
- Fix test_health TypeError: move exception extraction to a separate
  exceptions_by_model_id dict returned alongside endpoints, so exception
  objects never appear in the endpoint dicts that get JSON-serialized
  by the /health response.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(health-check-routing): properly isolate exceptions from health response

Return exceptions_by_model_id as a separate third value from
_perform_health_check / perform_health_check so exception objects
(which contain non-JSON-serializable httpx URL types) never appear
in the endpoint dicts that get serialized by the /health response.

Callers updated: _health_endpoints.py, shared_health_check_manager.py,
proxy_server.py background loop. All use the exceptions dict only for
cooldown integration, not for display.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(shared-health-check): fix remaining 2-value return sites and update type annotation

* fix(health-check-routing): fix P0 cooldown integration never firing

The cooldown loop was reading endpoint.get("exception") which is always
None because exceptions are now returned via exceptions_by_model_id, not
stored in endpoint dicts. Fixed to use _exceptions.get(model_id).

Also fixes the transient-error filter to use _exceptions instead of
endpoint.get("exception"), and fixes all remaining 2-value return sites
in shared_health_check_manager.py. Tests updated to pass exceptions via
exceptions_by_model_id parameter instead of endpoint dicts.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(health-check-routing): fix P1 transient-error filter broken on cache hits

When SharedHealthCheckManager returns cached results, exceptions_by_model_id
is always {} so the transient-error filter defaulted to status 500 for all
endpoints, incorrectly marking 429/408 endpoints as unhealthy.

Fix: store integer exception_status on each unhealthy endpoint dict in
_perform_health_check. _get_endpoint_exception_status() uses the live
exception object when available (direct path) and falls back to the stored
integer (cache-hit path). The integer is JSON-serializable and survives
the shared cache round-trip.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(health-check-routing): gate cooldown loop behind allowed_fails_policy

Without the policy, cooldown is not the routing exclusion mechanism.
Firing _set_cooldown_deployments for all enable_health_check_routing users
was a backwards-incompatible change — 401s would immediately cooldown
deployments that the binary filter would have recovered on the next cycle.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* revert: undo allowed_fails_policy gate on cooldown loop

Cooldown integration via health checks is intentional for all
enable_health_check_routing users, not just those with allowed_fails_policy.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(docs+tests): fix health_check_ignore_transient_errors doc section and test coverage

- Move health_check_ignore_transient_errors from router_settings to
  general_settings in config_settings.md (code reads it from general_settings)
- Remove duplicate enable_health_check_routing / health_check_staleness_threshold
  entries that were incorrectly listed under router_settings
- Replace TestHealthCheckEndpointExceptionPropagation tests with ones that
  exercise the real _perform_health_check code path via mocked ahealth_check,
  verifying exceptions appear in exceptions_by_model_id and NOT in endpoint dicts

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix(tests+docs): fix tuple unpacking and docs test failures

- Update test mocks that return (healthy, unhealthy) to return
  (healthy, unhealthy, {}) to match the new 3-value signature
- Update test unpackings of perform_shared_health_check to use
  healthy, unhealthy, _ = ...
- Add health_check_ignore_transient_errors to router_settings section
  in config_settings.md (it is a Router constructor param, so the doc
  test requires it there; it also lives in general_settings for proxy use)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Fix CodeQL errors

* fix(tests): fix 2-value unpackings of _perform_health_check in test_health_check.py

* fix(tests): fix mock _perform_health_check returning 2-tuple instead of 3

* fix team routing

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* fix: add distributed lock for key rotation job (#23364)

* fix: add distributed lock for key rotation job

* fix: address Greptile review feedback on key rotation lock (#23834)

* fix: address Greptile review feedback on key rotation lock

* fix req changes greptile

* feat(proxy): Optional on_error for guardrail pipeline (API / technical failures) (#24831)

* guardrails fallback

* docs

* docs: add LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS to environment variables reference

* fix(mypy): accept Union[Dict, Any] in _get_deployment_order and use typed list to fix min() type error

* fix(mypy): use Optional[str] for api_base in PydanticAI provider to match superclass signature

---------

Co-authored-by: Sameer Kankute <sameer@berri.ai>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Harshit Jain <48647625+Harshit28j@users.noreply.github.com>
Co-authored-by: Shivam Rawat <shivam@berri.ai>
Co-authored-by: yuneng-jiang <yuneng@berri.ai>
2026-04-04 23:09:42 +00:00

790 lines
26 KiB
Python

#### What this tests ####
# This tests if ahealth_check() actually works
import os
import sys
import pytest
from unittest.mock import AsyncMock, patch
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import asyncio
import litellm
@pytest.mark.asyncio
async def test_azure_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/gpt-4.1-mini",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
}
)
print(f"response: {response}")
assert "x-ratelimit-remaining-tokens" in response
return response
# asyncio.run(test_azure_health_check())
@pytest.mark.asyncio
async def test_text_completion_health_check():
response = await litellm.ahealth_check(
model_params={"model": "gpt-3.5-turbo-instruct"},
mode="completion",
prompt="What's the weather in SF?",
)
print(f"response: {response}")
return response
@pytest.mark.asyncio
async def test_azure_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_AI_API_KEY"),
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_version": os.getenv("AZURE_AI_API_VERSION"),
},
input=["test for litellm"],
mode="embedding",
)
print(f"response: {response}")
assert "x-ratelimit-remaining-tokens" in response
return response
@pytest.mark.asyncio
async def test_openai_img_gen_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "dall-e-3",
"api_key": os.getenv("OPENAI_API_KEY"),
},
mode="image_generation",
prompt="cute baby sea otter",
)
print(f"response: {response}")
assert isinstance(response, dict) and "error" not in response
return response
# asyncio.run(test_openai_img_gen_health_check())
@pytest.mark.skip(
reason="Azure DALL-E 3 model deployment is deprecated (410 ModelDeprecated)"
)
@pytest.mark.asyncio
async def test_azure_img_gen_health_check():
"""
Test Azure image generation health check with retry logic for transient errors.
Azure sometimes returns internal server errors which are transient and not something we can control.
"""
litellm._turn_on_debug()
max_retries = 3
retry_delay = 1 # Start with 1 second delay
for attempt in range(max_retries):
response = await litellm.ahealth_check(
model_params={
"model": "azure/dall-e-3",
"api_base": os.getenv("AZURE_AI_API_BASE"),
"api_key": os.getenv("AZURE_AI_API_KEY"),
},
mode="image_generation",
prompt="cute baby sea otter",
)
# Check if response is successful (no error)
if isinstance(response, dict) and "error" not in response:
return response
# Check if error is a transient Azure internal server error
error_str = str(response.get("error", "")).lower()
is_transient_error = (
"internalservererror" in error_str
or "internal server error" in error_str
or "internalfailure" in error_str
or "internal failure" in error_str
)
# If it's the last attempt or not a transient error, fail the test
if attempt == max_retries - 1 or not is_transient_error:
assert (
isinstance(response, dict) and "error" not in response
), f"Health check failed: {response.get('error', 'Unknown error')}"
return response
# Wait before retrying with exponential backoff
await asyncio.sleep(retry_delay)
retry_delay *= 2 # Exponential backoff
# Should not reach here, but just in case
assert False, "Health check failed after all retries"
@pytest.mark.skip(reason="AWS Suspended Account")
@pytest.mark.asyncio
async def test_sagemaker_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "sagemaker/berri-benchmarking-gpt-j-6b-fp16",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
},
mode="embedding",
input=["test from litellm"],
)
print(f"response: {response}")
assert isinstance(response, dict)
return response
# asyncio.run(test_sagemaker_embedding_health_check())
@pytest.mark.asyncio
async def test_groq_health_check():
"""
This should not fail
ensure that provider wildcard model passes health check
"""
litellm.set_verbose = True
response = await litellm.ahealth_check(
model_params={
"api_key": os.environ.get("GROQ_API_KEY"),
"model": "groq/*",
"messages": [{"role": "user", "content": "What's 1 + 1?"}],
},
mode=None,
prompt="What's 1 + 1?",
input=["test from litellm"],
)
print(f"response: {response}")
assert response == {}
return response
@pytest.mark.asyncio
async def test_cohere_rerank_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "cohere/rerank-english-v3.0",
"api_key": os.getenv("COHERE_API_KEY"),
},
mode="rerank",
prompt="Hey, how's it going",
)
assert "error" not in response
print(response)
@pytest.mark.asyncio
async def test_audio_speech_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "openai/tts-1",
"api_key": os.getenv("OPENAI_API_KEY"),
},
mode="audio_speech",
prompt="Hey",
)
assert "error" not in response
print(response)
@pytest.mark.asyncio
async def test_audio_speech_health_check_with_another_voice():
response = await litellm.ahealth_check(
model_params={
"model": "openai/tts-1",
"api_key": os.getenv("OPENAI_API_KEY"),
"health_check_voice": "en-US-JennyNeural",
},
mode="audio_speech",
prompt="Hey",
)
assert "error" not in response
print(response)
@pytest.mark.asyncio
async def test_audio_transcription_health_check():
litellm.set_verbose = True
response = await litellm.ahealth_check(
model_params={
"model": "openai/whisper-1",
"api_key": os.getenv("OPENAI_API_KEY"),
},
mode="audio_transcription",
)
print(f"response: {response}")
assert "error" not in response
print(response)
def test_update_litellm_params_for_health_check():
"""
Test if _update_litellm_params_for_health_check correctly:
1. Updates messages with a random message
2. Updates model name when health_check_model is provided
3. Updates voice when health_check_voice is provided for audio_speech mode
"""
from litellm.proxy.health_check import _update_litellm_params_for_health_check
# Test with health_check_model
model_info = {"health_check_model": "gpt-3.5-turbo"}
litellm_params = {
"model": "gpt-4",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert "messages" in updated_params
assert isinstance(updated_params["messages"], list)
assert updated_params["model"] == "gpt-3.5-turbo"
# Test without health_check_model
model_info = {}
litellm_params = {
"model": "gpt-4",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert "messages" in updated_params
assert isinstance(updated_params["messages"], list)
assert updated_params["model"] == "gpt-4"
# Test with health_check_voice for audio_speech mode
model_info = {"mode": "audio_speech", "health_check_voice": "en-US-JennyNeural"}
litellm_params = {
"model": "gpt-4",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert "voice" in updated_params
assert updated_params["voice"] == "en-US-JennyNeural"
# Test without health_check_voice for audio_speech mode
model_info = {"mode": "audio_speech"}
litellm_params = {
"model": "gpt-4",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert "voice" in updated_params
assert updated_params["voice"] == "alloy"
# Test with health_check_voice for non-audio_speech mode
model_info = {"mode": "chat", "health_check_voice": "en-US-JennyNeural"}
litellm_params = {
"model": "gpt-4",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert "voice" not in updated_params
# Test with Bedrock model with region routing - should strip bedrock/ and region/ prefix
# Issue #15807: Fixes health checks sending "region/model" as model ID to AWS
model_info = {}
litellm_params = {
"model": "bedrock/us-gov-west-1/anthropic.claude-3-7-sonnet-20250219-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "anthropic.claude-3-7-sonnet-20250219-v1:0"
# Test with Bedrock cross-region inference profile - should preserve the inference profile prefix
# AWS requires inference profile IDs like "us.anthropic.claude..." for cross-region routing
litellm_params = {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "us.anthropic.claude-haiku-4-5-20251001-v1:0"
# Test with Bedrock model without region routing - should just strip bedrock/ prefix
litellm_params = {
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "us.anthropic.claude-haiku-4-5-20251001-v1:0"
# Test that non-Bedrock models are not affected by Bedrock-specific logic
litellm_params = {
"model": "openai/gpt-4",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "openai/gpt-4" # Should remain unchanged
# Test ALL cross-region inference profile prefixes (CRIS)
cris_prefixes = ["us.", "eu.", "apac.", "jp.", "au.", "us-gov.", "global."]
for prefix in cris_prefixes:
litellm_params = {
"model": f"bedrock/{prefix}anthropic.claude-3-haiku-20240307-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(
model_info, litellm_params
)
assert (
updated_params["model"] == f"{prefix}anthropic.claude-3-haiku-20240307-v1:0"
), f"Failed to preserve CRIS prefix: {prefix}"
# Test regional + CRIS combination - region should be stripped, CRIS preserved
litellm_params = {
"model": "bedrock/us-east-2/us.anthropic.claude-3-haiku-20240307-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "us.anthropic.claude-3-haiku-20240307-v1:0"
# Test GovCloud regions
litellm_params = {
"model": "bedrock/us-gov-east-1/anthropic.claude-instant-v1",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "anthropic.claude-instant-v1"
# Test imported models with handler prefixes - handlers should be preserved
litellm_params = {
"model": "bedrock/llama/arn:aws:bedrock:us-east-1:123:imported-model/abc",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert (
updated_params["model"]
== "llama/arn:aws:bedrock:us-east-1:123:imported-model/abc"
)
litellm_params = {
"model": "bedrock/deepseek_r1/arn:aws:bedrock:us-west-2:456:imported-model/xyz",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert (
updated_params["model"]
== "deepseek_r1/arn:aws:bedrock:us-west-2:456:imported-model/xyz"
)
# Test route specifications - routes should be preserved
litellm_params = {
"model": "bedrock/converse/us.anthropic.claude-haiku-4-5-20251001-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert (
updated_params["model"]
== "converse/us.anthropic.claude-haiku-4-5-20251001-v1:0"
)
litellm_params = {
"model": "bedrock/invoke/us-west-2/anthropic.claude-instant-v1",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert updated_params["model"] == "invoke/anthropic.claude-instant-v1"
# Test ARN formats - should be preserved
litellm_params = {
"model": "bedrock/arn:aws:bedrock:eu-central-1:000:application-inference-profile/abc",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert (
updated_params["model"]
== "arn:aws:bedrock:eu-central-1:000:application-inference-profile/abc"
)
# Test edge case: region + handler + ARN
litellm_params = {
"model": "bedrock/us-west-2/llama/arn:aws:bedrock:us-east-1:123:imported-model/abc",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert (
updated_params["model"]
== "llama/arn:aws:bedrock:us-east-1:123:imported-model/abc"
)
# Test edge case: route + region + CRIS
litellm_params = {
"model": "bedrock/converse/us-west-2/eu.anthropic.claude-3-sonnet-20240229-v1:0",
"api_key": "fake_key",
}
updated_params = _update_litellm_params_for_health_check(model_info, litellm_params)
assert (
updated_params["model"] == "converse/eu.anthropic.claude-3-sonnet-20240229-v1:0"
)
@pytest.mark.asyncio
async def test_perform_health_check_filters_by_model_id():
"""
When model_id is passed, only that deployment is checked (not all deployments
that share the same model name).
"""
from litellm.proxy.health_check import perform_health_check
# Two deployments with same model_name but different ids
model_list = [
{
"model_name": "gpt-4",
"model_info": {"id": "deployment-id-1"},
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"},
},
{
"model_name": "gpt-4",
"model_info": {"id": "deployment-id-2"},
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-2"},
},
]
captured_list = []
async def mock_perform_health_check(m_list, details=True, **kwargs):
captured_list.append(m_list)
return [
{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}
], [], {}
with patch(
"litellm.proxy.health_check._perform_health_check",
side_effect=mock_perform_health_check,
):
healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check(
model_list=model_list, model_id="deployment-id-2", details=True
)
# Only one deployment (deployment-id-2) should have been passed to _perform_health_check
assert len(captured_list) == 1
assert len(captured_list[0]) == 1
assert (captured_list[0][0].get("model_info") or {}).get("id") == "deployment-id-2"
assert len(healthy_endpoints) == 1
assert healthy_endpoints[0]["api_key"] == "fake-key-2"
@pytest.mark.asyncio
async def test_perform_health_check_with_health_check_model():
"""
Test if _perform_health_check correctly uses `health_check_model` when model=`openai/*`:
1. Verifies that health_check_model overrides the original model when model=`openai/*`
2. Ensures the health check is performed with the override model
"""
from litellm.proxy.health_check import _perform_health_check
# Mock model list with health_check_model specified
model_list = [
{
"litellm_params": {"model": "openai/*", "api_key": "fake-key"},
"model_info": {
"mode": "chat",
"health_check_model": "openai/gpt-4o-mini", # Override model for health check
},
}
]
# Track which model is actually used in the health check
health_check_calls = []
async def mock_health_check(litellm_params, **kwargs):
health_check_calls.append(litellm_params["model"])
return {"status": "healthy"}
with patch("litellm.ahealth_check", side_effect=mock_health_check):
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list)
print("health check calls: ", health_check_calls)
# Verify the health check used the override model
assert health_check_calls[0] == "openai/gpt-4o-mini"
# Verify the result still shows the original model
print("healthy endpoints: ", healthy_endpoints)
assert healthy_endpoints[0]["model"] == "openai/gpt-4o-mini"
assert len(healthy_endpoints) == 1
assert len(unhealthy_endpoints) == 0
@pytest.mark.asyncio
async def test_health_check_bad_model():
from litellm.proxy.health_check import _perform_health_check
import time
model_list = [
{
"model_name": "openai-gpt-4o",
"litellm_params": {
"api_key": "sk-1234",
"api_base": "https://exampleopenaiendpoint-production.up.railway.app",
"model": "openai/my-fake-openai-endpoint",
"mock_timeout": True,
"timeout": 60,
},
"model_info": {
"id": "ca27ca2eeea2f9e38bb274ead831948a26621a3738d06f1797253f0e6c4278c0",
"db_model": False,
"health_check_timeout": 1,
},
},
]
details = None
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(
model_list, details
)
print(f"healthy_endpoints: {healthy_endpoints}")
print(f"unhealthy_endpoints: {unhealthy_endpoints}")
# Track which model is actually used in the health check
health_check_calls = []
async def mock_health_check(litellm_params, **kwargs):
health_check_calls.append(litellm_params["model"])
await asyncio.sleep(10)
return {"status": "healthy"}
with patch(
"litellm.ahealth_check", side_effect=mock_health_check
) as mock_health_check:
start_time = time.time()
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list)
end_time = time.time()
print("health check calls: ", health_check_calls)
assert len(healthy_endpoints) == 0
assert len(unhealthy_endpoints) == 1
assert (
end_time - start_time < 2
), "Health check took longer than health_check_timeout"
@pytest.mark.asyncio
async def test_health_check_respects_concurrency_limit():
from litellm.proxy.health_check import _perform_health_check
model_list = [
{"litellm_params": {"model": f"openai/gpt-4o-mini-{i}", "api_key": "fake-key"}}
for i in range(6)
]
active = 0
max_active = 0
async def mock_health_check(litellm_params, **kwargs):
nonlocal active, max_active
active += 1
max_active = max(max_active, active)
await asyncio.sleep(0.05)
active -= 1
return {"status": "healthy"}
with patch("litellm.ahealth_check", side_effect=mock_health_check):
await _perform_health_check(model_list, max_concurrency=2)
assert max_active <= 2
@pytest.mark.asyncio
async def test_health_check_creates_only_bounded_initial_tasks():
from litellm.proxy.health_check import _perform_health_check
model_list = [
{"litellm_params": {"model": f"openai/gpt-4o-mini-{i}", "api_key": "fake-key"}}
for i in range(10)
]
release_event = asyncio.Event()
create_task_call_count = 0
real_create_task = asyncio.create_task
async def mock_health_check(litellm_params, **kwargs):
await release_event.wait()
return {"status": "healthy"}
def tracked_create_task(coro):
nonlocal create_task_call_count
create_task_call_count += 1
return real_create_task(coro)
with patch("litellm.ahealth_check", side_effect=mock_health_check), patch(
"litellm.proxy.health_check.asyncio.create_task",
side_effect=tracked_create_task,
):
perform_task = real_create_task(
_perform_health_check(model_list, max_concurrency=2)
)
await asyncio.sleep(0.05)
assert create_task_call_count == 2
release_event.set()
await perform_task
@pytest.mark.asyncio
async def test_timeout_does_not_cancel_other_health_checks():
from litellm.proxy.health_check import _perform_health_check
model_list = [
{
"litellm_params": {"model": "openai/slow-model", "api_key": "fake-key"},
"model_info": {"health_check_timeout": 0.05},
},
{
"litellm_params": {"model": "openai/fast-model", "api_key": "fake-key"},
"model_info": {"health_check_timeout": 1},
},
]
async def mock_health_check(litellm_params, **kwargs):
if litellm_params["model"] == "openai/slow-model":
await asyncio.sleep(0.2)
return {"status": "healthy"}
await asyncio.sleep(0.01)
return {"status": "healthy"}
with patch("litellm.ahealth_check", side_effect=mock_health_check):
healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(
model_list, max_concurrency=1
)
healthy_models = {endpoint["model"] for endpoint in healthy_endpoints}
unhealthy_models = {endpoint["model"] for endpoint in unhealthy_endpoints}
assert "openai/fast-model" in healthy_models
assert "openai/slow-model" in unhealthy_models
@pytest.mark.asyncio
async def test_ahealth_check_ocr():
litellm._turn_on_debug()
response = await litellm.ahealth_check(
model_params={
"model": "mistral/mistral-ocr-latest",
"api_key": os.getenv("MISTRAL_API_KEY"),
},
mode="ocr",
)
print(response)
return response
@pytest.mark.asyncio
async def test_image_generation_health_check_prompt(monkeypatch):
"""Health checks should respect default and environment-configured prompts."""
import importlib
import litellm.constants as litellm_constants
import litellm.proxy.health_check as health_check
def reload_modules():
reloaded_constants = importlib.reload(litellm_constants)
reloaded_health_check = importlib.reload(health_check)
return reloaded_constants, reloaded_health_check
async def run_health_check(health_check_module):
health_check_calls = []
async def mock_health_check(litellm_params, mode=None, prompt=None, input=None):
health_check_calls.append(
{
"mode": mode,
"prompt": prompt,
"model": litellm_params.get("model"),
}
)
return {"status": "healthy"}
model_list = [
{
"litellm_params": {"model": "dall-e-3", "api_key": "fake-key"},
"model_info": {
"mode": "image_generation",
},
}
]
with patch(
"litellm.proxy.health_check.litellm.ahealth_check",
side_effect=mock_health_check,
):
await health_check_module._perform_health_check(model_list)
return health_check_calls
# Default prompt is used when env var is unset
monkeypatch.delenv("DEFAULT_HEALTH_CHECK_PROMPT", raising=False)
litellm_constants, health_check = reload_modules()
health_check_calls = await run_health_check(health_check)
assert len(health_check_calls) == 1
assert (
health_check_calls[0]["prompt"] == litellm_constants.DEFAULT_HEALTH_CHECK_PROMPT
)
# Environment override should change the prompt without code changes
override_prompt = "environment override prompt"
monkeypatch.setenv("DEFAULT_HEALTH_CHECK_PROMPT", override_prompt)
litellm_constants, health_check = reload_modules()
health_check_calls = await run_health_check(health_check)
assert len(health_check_calls) == 1
assert health_check_calls[0]["prompt"] == override_prompt
@pytest.mark.asyncio
async def test_health_check_with_custom_llm_provider():
"""
Test that ahealth_check correctly uses custom_llm_provider from model_params.
This test verifies the fix for the issue where the UI's "Test connect" button
failed with "LLM Provider NOT provided" error for OpenAI-compatible self-hosted
providers, even when a provider was selected in the dropdown.
The fix ensures that when custom_llm_provider is passed in model_params,
it's properly forwarded to get_llm_provider() to identify the correct provider.
"""
from unittest.mock import MagicMock
# Mock the completion call to avoid making real API calls
mock_response = MagicMock()
mock_response._hidden_params = {"headers": {"x-ratelimit-remaining-tokens": "1000"}}
with patch("litellm.acompletion", return_value=mock_response):
# Test with a custom model name that wouldn't be recognized without custom_llm_provider
response = await litellm.ahealth_check(
model_params={
"model": "deepseek-r1-distill-qwen-1.5B-q4",
"custom_llm_provider": "openai",
"api_base": "https://example.com/v1",
"api_key": "fake-key",
},
mode="chat",
)
# Should succeed without "LLM Provider NOT provided" error
assert "error" not in response
assert isinstance(response, dict)