mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
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>
This commit is contained in:
Sameer Kankute
parent
7f1b34cb1a
commit
1867ca51ae
5 changed files with 765 additions and 117 deletions
|
|
@ -3792,9 +3792,9 @@ def completion( # type: ignore # noqa: PLR0915
|
|||
"aws_region_name" not in optional_params
|
||||
or optional_params["aws_region_name"] is None
|
||||
):
|
||||
optional_params[
|
||||
"aws_region_name"
|
||||
] = aws_bedrock_client.meta.region_name
|
||||
optional_params["aws_region_name"] = (
|
||||
aws_bedrock_client.meta.region_name
|
||||
)
|
||||
|
||||
bedrock_route = BedrockModelInfo.get_bedrock_route(model)
|
||||
if bedrock_route == "converse":
|
||||
|
|
@ -6198,9 +6198,9 @@ def adapter_completion(
|
|||
new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs)
|
||||
|
||||
response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore
|
||||
translated_response: Optional[
|
||||
Union[BaseModel, AdapterCompletionStreamWrapper]
|
||||
] = None
|
||||
translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = (
|
||||
None
|
||||
)
|
||||
if isinstance(response, ModelResponse):
|
||||
translated_response = translation_obj.translate_completion_output_params(
|
||||
response=response
|
||||
|
|
@ -6380,9 +6380,9 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse:
|
|||
if existing_duration is None:
|
||||
calculated_duration = calculate_request_duration(file)
|
||||
if calculated_duration is not None:
|
||||
response._hidden_params[
|
||||
"audio_transcription_duration"
|
||||
] = calculated_duration
|
||||
response._hidden_params["audio_transcription_duration"] = (
|
||||
calculated_duration
|
||||
)
|
||||
|
||||
return response
|
||||
except Exception as e:
|
||||
|
|
@ -6605,9 +6605,9 @@ def transcription(
|
|||
if existing_duration is None:
|
||||
calculated_duration = calculate_request_duration(file)
|
||||
if calculated_duration is not None:
|
||||
response._hidden_params[
|
||||
"audio_transcription_duration"
|
||||
] = calculated_duration
|
||||
response._hidden_params["audio_transcription_duration"] = (
|
||||
calculated_duration
|
||||
)
|
||||
|
||||
if response is None:
|
||||
raise ValueError("Unmapped provider passed in. Unable to get the response.")
|
||||
|
|
@ -6911,9 +6911,9 @@ def speech( # noqa: PLR0915
|
|||
ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY
|
||||
] = query_params
|
||||
|
||||
litellm_params_dict[
|
||||
ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY
|
||||
] = voice_id
|
||||
litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = (
|
||||
voice_id
|
||||
)
|
||||
|
||||
if api_base is not None:
|
||||
litellm_params_dict["api_base"] = api_base
|
||||
|
|
@ -7234,7 +7234,8 @@ async def ahealth_check(
|
|||
|
||||
if mode is None:
|
||||
return {
|
||||
"error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}"
|
||||
"error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}",
|
||||
"exception": e,
|
||||
}
|
||||
|
||||
error_to_return = str(e) + "\nstack trace: " + stack_trace
|
||||
|
|
@ -7246,6 +7247,7 @@ async def ahealth_check(
|
|||
return {
|
||||
"error": error_to_return,
|
||||
"raw_request_typed_dict": raw_request_typed_dict,
|
||||
"exception": e,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -7492,9 +7494,9 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
]
|
||||
|
||||
if len(content_chunks) > 0:
|
||||
response["choices"][0]["message"][
|
||||
"content"
|
||||
] = processor.get_combined_content(content_chunks)
|
||||
response["choices"][0]["message"]["content"] = (
|
||||
processor.get_combined_content(content_chunks)
|
||||
)
|
||||
|
||||
thinking_blocks = [
|
||||
chunk
|
||||
|
|
@ -7505,9 +7507,9 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
]
|
||||
|
||||
if len(thinking_blocks) > 0:
|
||||
response["choices"][0]["message"][
|
||||
"thinking_blocks"
|
||||
] = processor.get_combined_thinking_content(thinking_blocks)
|
||||
response["choices"][0]["message"]["thinking_blocks"] = (
|
||||
processor.get_combined_thinking_content(thinking_blocks)
|
||||
)
|
||||
|
||||
reasoning_chunks = [
|
||||
chunk
|
||||
|
|
@ -7518,9 +7520,9 @@ def stream_chunk_builder( # noqa: PLR0915
|
|||
]
|
||||
|
||||
if len(reasoning_chunks) > 0:
|
||||
response["choices"][0]["message"][
|
||||
"reasoning_content"
|
||||
] = processor.get_combined_reasoning_content(reasoning_chunks)
|
||||
response["choices"][0]["message"]["reasoning_content"] = (
|
||||
processor.get_combined_reasoning_content(reasoning_chunks)
|
||||
)
|
||||
|
||||
annotation_chunks = [
|
||||
chunk
|
||||
|
|
|
|||
|
|
@ -95,7 +95,12 @@ async def run_with_timeout(task, timeout):
|
|||
except asyncio.TimeoutError:
|
||||
# `asyncio.wait_for()` already cancels only the awaited task on timeout.
|
||||
# Do not cancel unrelated sibling health check tasks.
|
||||
return {"error": "Timeout exceeded"}
|
||||
timeout_exception = litellm.Timeout(
|
||||
message="Health check timeout exceeded",
|
||||
model="",
|
||||
llm_provider="",
|
||||
)
|
||||
return {"error": "Timeout exceeded", "exception": timeout_exception}
|
||||
|
||||
|
||||
async def _run_model_health_check(model: dict):
|
||||
|
|
@ -218,11 +223,15 @@ async def _perform_health_check(
|
|||
cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details)
|
||||
if _model_id:
|
||||
cleaned["model_id"] = _model_id
|
||||
if "exception" in is_healthy:
|
||||
cleaned["exception"] = is_healthy["exception"]
|
||||
unhealthy_endpoints.append(cleaned)
|
||||
else:
|
||||
cleaned = _clean_endpoint_data(litellm_params, details)
|
||||
if _model_id:
|
||||
cleaned["model_id"] = _model_id
|
||||
if isinstance(is_healthy, Exception):
|
||||
cleaned["exception"] = is_healthy
|
||||
unhealthy_endpoints.append(cleaned)
|
||||
|
||||
return healthy_endpoints, unhealthy_endpoints
|
||||
|
|
|
|||
|
|
@ -375,9 +375,7 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import (
|
|||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
router as internal_user_router,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import (
|
||||
user_update,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
|
||||
from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import (
|
||||
router as jwt_key_mapping_router,
|
||||
)
|
||||
|
|
@ -446,9 +444,7 @@ from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_route
|
|||
from litellm.proxy.openai_files_endpoints.files_endpoints import (
|
||||
router as openai_files_router,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.files_endpoints import (
|
||||
set_files_config,
|
||||
)
|
||||
from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
passthrough_endpoint_router,
|
||||
)
|
||||
|
|
@ -550,9 +546,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
|
|||
LiteLLM_UpperboundKeyGenerateParams,
|
||||
)
|
||||
from litellm.types.realtime import RealtimeQueryParams
|
||||
from litellm.types.router import (
|
||||
DeploymentTypedDict,
|
||||
)
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
from litellm.types.router import ModelInfo as RouterModelInfo
|
||||
from litellm.types.router import (
|
||||
RouterGeneralSettings,
|
||||
|
|
@ -639,9 +633,9 @@ except ImportError:
|
|||
server_root_path = get_server_root_path()
|
||||
_license_check = LicenseCheck()
|
||||
premium_user: bool = _license_check.is_premium()
|
||||
premium_user_data: Optional[
|
||||
"EnterpriseLicenseData"
|
||||
] = _license_check.airgapped_license_data
|
||||
premium_user_data: Optional["EnterpriseLicenseData"] = (
|
||||
_license_check.airgapped_license_data
|
||||
)
|
||||
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
|
||||
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
|
||||
)
|
||||
|
|
@ -1524,9 +1518,9 @@ master_key: Optional[str] = None
|
|||
config_agents: Optional[List[AgentConfig]] = None
|
||||
otel_logging = False
|
||||
prisma_client: Optional[PrismaClient] = None
|
||||
shared_aiohttp_session: Optional[
|
||||
"ClientSession"
|
||||
] = None # Global shared session for connection reuse
|
||||
shared_aiohttp_session: Optional["ClientSession"] = (
|
||||
None # Global shared session for connection reuse
|
||||
)
|
||||
user_api_key_cache = DualCache(
|
||||
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
|
||||
)
|
||||
|
|
@ -1534,13 +1528,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
|
|||
dual_cache=user_api_key_cache
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
|
||||
redis_usage_cache: Optional[
|
||||
RedisCache
|
||||
] = None # redis cache used for tracking spend, tpm/rpm limits
|
||||
redis_usage_cache: Optional[RedisCache] = (
|
||||
None # redis cache used for tracking spend, tpm/rpm limits
|
||||
)
|
||||
polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
|
||||
native_background_mode: List[
|
||||
str
|
||||
] = [] # Models that should use native provider background mode instead of polling
|
||||
native_background_mode: List[str] = (
|
||||
[]
|
||||
) # Models that should use native provider background mode instead of polling
|
||||
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
|
||||
user_custom_auth = None
|
||||
user_custom_key_generate = None
|
||||
|
|
@ -1900,9 +1894,9 @@ async def update_cache( # noqa: PLR0915
|
|||
_id = "team_id:{}".format(team_id)
|
||||
try:
|
||||
# Fetch the existing cost for the given user
|
||||
existing_spend_obj: Optional[
|
||||
LiteLLM_TeamTable
|
||||
] = await user_api_key_cache.async_get_cache(key=_id)
|
||||
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
|
||||
await user_api_key_cache.async_get_cache(key=_id)
|
||||
)
|
||||
if existing_spend_obj is None:
|
||||
# do nothing if team not in api key cache
|
||||
return
|
||||
|
|
@ -2121,6 +2115,10 @@ def _write_health_state_to_router_cache(
|
|||
for health-check-driven routing. No-op if the feature is disabled.
|
||||
"""
|
||||
from litellm.proxy.health_check import build_deployment_health_states
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
|
||||
increment_deployment_failures_for_current_minute,
|
||||
)
|
||||
|
||||
try:
|
||||
if llm_router is None or not llm_router.enable_health_check_routing:
|
||||
|
|
@ -2137,6 +2135,31 @@ def _write_health_state_to_router_cache(
|
|||
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")),
|
||||
)
|
||||
|
||||
for endpoint in unhealthy_endpoints:
|
||||
model_id = endpoint.get("model_id")
|
||||
if not model_id:
|
||||
continue
|
||||
|
||||
original_exception = endpoint.get("exception")
|
||||
if original_exception is None:
|
||||
continue
|
||||
|
||||
exception_status = getattr(original_exception, "status_code", 500)
|
||||
|
||||
increment_deployment_failures_for_current_minute(
|
||||
litellm_router_instance=llm_router,
|
||||
deployment_id=model_id,
|
||||
)
|
||||
|
||||
_set_cooldown_deployments(
|
||||
litellm_router_instance=llm_router,
|
||||
original_exception=original_exception,
|
||||
exception_status=exception_status,
|
||||
deployment=model_id,
|
||||
time_to_cooldown=llm_router.cooldown_time,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to write health state to router cache: %s", str(e)
|
||||
|
|
@ -5050,10 +5073,10 @@ class ProxyConfig:
|
|||
)
|
||||
|
||||
try:
|
||||
guardrails_in_db: List[
|
||||
Guardrail
|
||||
] = await GuardrailRegistry.get_all_guardrails_from_db(
|
||||
prisma_client=prisma_client
|
||||
guardrails_in_db: List[Guardrail] = (
|
||||
await GuardrailRegistry.get_all_guardrails_from_db(
|
||||
prisma_client=prisma_client
|
||||
)
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"guardrails from the DB %s", str(guardrails_in_db)
|
||||
|
|
@ -5435,9 +5458,9 @@ async def initialize( # noqa: PLR0915
|
|||
user_api_base = api_base
|
||||
dynamic_config[user_model]["api_base"] = api_base
|
||||
if api_version:
|
||||
os.environ[
|
||||
"AZURE_API_VERSION"
|
||||
] = api_version # set this for azure - litellm can read this from the env
|
||||
os.environ["AZURE_API_VERSION"] = (
|
||||
api_version # set this for azure - litellm can read this from the env
|
||||
)
|
||||
if max_tokens: # model-specific param
|
||||
dynamic_config[user_model]["max_tokens"] = max_tokens
|
||||
if temperature: # model-specific param
|
||||
|
|
@ -5774,9 +5797,9 @@ class ProxyStartupEvent:
|
|||
"""
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
|
||||
_use_redis_transaction_buffer: Optional[
|
||||
Union[bool, str]
|
||||
] = general_settings.get("use_redis_transaction_buffer", False)
|
||||
_use_redis_transaction_buffer: Optional[Union[bool, str]] = (
|
||||
general_settings.get("use_redis_transaction_buffer", False)
|
||||
)
|
||||
if isinstance(_use_redis_transaction_buffer, str):
|
||||
_use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer)
|
||||
|
||||
|
|
@ -9101,6 +9124,70 @@ def _add_team_models_to_all_models(
|
|||
return team_models
|
||||
|
||||
|
||||
async def _add_access_group_models_to_team_models(
|
||||
team_db_objects_typed: List[LiteLLM_TeamTable],
|
||||
llm_router: Router,
|
||||
prisma_client: PrismaClient,
|
||||
team_models: Dict[str, Set[str]],
|
||||
) -> Dict[str, Set[str]]:
|
||||
"""
|
||||
Resolve models reachable via team access groups and merge them into team_models.
|
||||
|
||||
Batch-fetches all distinct access groups in a single DB query, then resolves
|
||||
each eligible team's access group models via the pre-fetched map.
|
||||
|
||||
This ensures models associated with a team only through access groups
|
||||
(not directly in team.models) are included in the UI model listing.
|
||||
"""
|
||||
# First pass: identify eligible teams and collect all distinct access group IDs
|
||||
eligible_teams: List[LiteLLM_TeamTable] = []
|
||||
all_access_group_ids: Set[str] = set()
|
||||
|
||||
for team_object in team_db_objects_typed:
|
||||
if not team_object.access_group_ids:
|
||||
continue
|
||||
|
||||
# Skip teams with empty models list — they already have access to everything
|
||||
# (handled by _add_team_models_to_all_models)
|
||||
if (
|
||||
not team_object.models
|
||||
or SpecialModelNames.all_proxy_models.value in team_object.models
|
||||
):
|
||||
continue
|
||||
|
||||
eligible_teams.append(team_object)
|
||||
all_access_group_ids.update(team_object.access_group_ids)
|
||||
|
||||
if not eligible_teams:
|
||||
return team_models
|
||||
|
||||
# Single batch fetch for all access groups
|
||||
access_group_rows = await prisma_client.db.litellm_accessgrouptable.find_many(
|
||||
where={"access_group_id": {"in": list(all_access_group_ids)}}
|
||||
)
|
||||
ag_model_map: Dict[str, List[str]] = {
|
||||
row.access_group_id: row.access_model_names or [] for row in access_group_rows
|
||||
}
|
||||
|
||||
# Second pass: resolve deployments for each eligible team
|
||||
for team_object in eligible_teams:
|
||||
model_names: Set[str] = set()
|
||||
for ag_id in team_object.access_group_ids or []:
|
||||
model_names.update(ag_model_map.get(ag_id, []))
|
||||
|
||||
for model_name in model_names:
|
||||
deployments = llm_router.get_model_list(
|
||||
model_name=model_name, team_id=team_object.team_id
|
||||
)
|
||||
if deployments is not None:
|
||||
for deployment in deployments:
|
||||
model_id = deployment.get("model_info", {}).get("id", None)
|
||||
if model_id is not None:
|
||||
team_models.setdefault(model_id, set()).add(team_object.team_id)
|
||||
|
||||
return team_models
|
||||
|
||||
|
||||
async def get_all_team_models(
|
||||
user_teams: Union[List[str], Literal["*"]],
|
||||
prisma_client: PrismaClient,
|
||||
|
|
@ -12379,9 +12466,9 @@ async def get_config_list(
|
|||
hasattr(sub_field_info, "description")
|
||||
and sub_field_info.description is not None
|
||||
):
|
||||
nested_fields[
|
||||
idx
|
||||
].field_description = sub_field_info.description
|
||||
nested_fields[idx].field_description = (
|
||||
sub_field_info.description
|
||||
)
|
||||
idx += 1
|
||||
|
||||
_stored_in_db = None
|
||||
|
|
|
|||
|
|
@ -170,11 +170,7 @@ from litellm.types.utils import (
|
|||
)
|
||||
from litellm.types.utils import ModelInfo
|
||||
from litellm.types.utils import ModelInfo as ModelMapInfo
|
||||
from litellm.types.utils import (
|
||||
ModelResponseStream,
|
||||
StandardLoggingPayload,
|
||||
Usage,
|
||||
)
|
||||
from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage
|
||||
from litellm.utils import (
|
||||
CustomStreamWrapper,
|
||||
EmbeddingResponse,
|
||||
|
|
@ -408,9 +404,9 @@ class Router:
|
|||
) # names of models under litellm_params. ex. azure/chatgpt-v-2
|
||||
self.deployment_latency_map = {}
|
||||
### CACHING ###
|
||||
cache_type: Literal[
|
||||
"local", "redis", "redis-semantic", "s3", "disk"
|
||||
] = "local" # default to an in-memory cache
|
||||
cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = (
|
||||
"local" # default to an in-memory cache
|
||||
)
|
||||
redis_cache = None
|
||||
cache_config: Dict[str, Any] = {}
|
||||
|
||||
|
|
@ -458,9 +454,9 @@ class Router:
|
|||
self.default_max_parallel_requests = default_max_parallel_requests
|
||||
self.provider_default_deployment_ids: List[str] = []
|
||||
self.pattern_router = PatternMatchRouter()
|
||||
self.team_pattern_routers: Dict[
|
||||
str, PatternMatchRouter
|
||||
] = {} # {"TEAM_ID": PatternMatchRouter}
|
||||
self.team_pattern_routers: Dict[str, PatternMatchRouter] = (
|
||||
{}
|
||||
) # {"TEAM_ID": PatternMatchRouter}
|
||||
self.auto_routers: Dict[str, "AutoRouter"] = {}
|
||||
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
|
||||
|
||||
|
|
@ -653,12 +649,12 @@ class Router:
|
|||
)
|
||||
)
|
||||
|
||||
self.model_group_retry_policy: Optional[
|
||||
Dict[str, RetryPolicy]
|
||||
] = model_group_retry_policy
|
||||
self.model_group_affinity_config: Optional[
|
||||
Dict[str, List[str]]
|
||||
] = model_group_affinity_config
|
||||
self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = (
|
||||
model_group_retry_policy
|
||||
)
|
||||
self.model_group_affinity_config: Optional[Dict[str, List[str]]] = (
|
||||
model_group_affinity_config
|
||||
)
|
||||
|
||||
self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None
|
||||
if allowed_fails_policy is not None:
|
||||
|
|
@ -2064,7 +2060,10 @@ class Router:
|
|||
|
||||
async def _acompletion( # noqa: PLR0915
|
||||
self, model: str, messages: List[Dict[str, str]], **kwargs
|
||||
) -> Union[ModelResponse, CustomStreamWrapper,]:
|
||||
) -> Union[
|
||||
ModelResponse,
|
||||
CustomStreamWrapper,
|
||||
]:
|
||||
"""
|
||||
- Get an available deployment
|
||||
- call it with a semaphore over the call
|
||||
|
|
@ -4298,9 +4297,9 @@ class Router:
|
|||
healthy_deployments=healthy_deployments, responses=responses
|
||||
)
|
||||
returned_response = cast(OpenAIFileObject, responses[0])
|
||||
returned_response._hidden_params[
|
||||
"model_file_id_mapping"
|
||||
] = model_file_id_mapping
|
||||
returned_response._hidden_params["model_file_id_mapping"] = (
|
||||
model_file_id_mapping
|
||||
)
|
||||
return returned_response
|
||||
except Exception as e:
|
||||
verbose_router_logger.exception(
|
||||
|
|
@ -5385,11 +5384,11 @@ class Router:
|
|||
|
||||
if isinstance(e, litellm.ContextWindowExceededError):
|
||||
if context_window_fallbacks is not None:
|
||||
context_window_fallback_model_group: Optional[
|
||||
List[str]
|
||||
] = self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=context_window_fallbacks,
|
||||
model_group=model_group,
|
||||
context_window_fallback_model_group: Optional[List[str]] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=context_window_fallbacks,
|
||||
model_group=model_group,
|
||||
)
|
||||
)
|
||||
if context_window_fallback_model_group is None:
|
||||
raise original_exception
|
||||
|
|
@ -5421,11 +5420,11 @@ class Router:
|
|||
e.message += "\n{}".format(error_message)
|
||||
elif isinstance(e, litellm.ContentPolicyViolationError):
|
||||
if content_policy_fallbacks is not None:
|
||||
content_policy_fallback_model_group: Optional[
|
||||
List[str]
|
||||
] = self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
model_group=model_group,
|
||||
content_policy_fallback_model_group: Optional[List[str]] = (
|
||||
self._get_fallback_model_group_from_fallbacks(
|
||||
fallbacks=content_policy_fallbacks,
|
||||
model_group=model_group,
|
||||
)
|
||||
)
|
||||
if content_policy_fallback_model_group is None:
|
||||
raise original_exception
|
||||
|
|
@ -5647,9 +5646,9 @@ class Router:
|
|||
)
|
||||
## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking
|
||||
_metadata["attempted_retries"] = 0
|
||||
_metadata[
|
||||
"max_retries"
|
||||
] = num_retries # Updated after overrides in exception handler
|
||||
_metadata["max_retries"] = (
|
||||
num_retries # Updated after overrides in exception handler
|
||||
)
|
||||
try:
|
||||
self._handle_mock_testing_rate_limit_error(
|
||||
model_group=model_group, kwargs=kwargs
|
||||
|
|
@ -6768,26 +6767,26 @@ class Router:
|
|||
"""
|
||||
from litellm.router_strategy.auto_router.auto_router import AutoRouter
|
||||
|
||||
auto_router_config_path: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.auto_router_config_path
|
||||
auto_router_config_path: Optional[str] = (
|
||||
deployment.litellm_params.auto_router_config_path
|
||||
)
|
||||
auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config
|
||||
if auto_router_config_path is None and auto_router_config is None:
|
||||
raise ValueError(
|
||||
"auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params"
|
||||
)
|
||||
|
||||
default_model: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.auto_router_default_model
|
||||
default_model: Optional[str] = (
|
||||
deployment.litellm_params.auto_router_default_model
|
||||
)
|
||||
if default_model is None:
|
||||
raise ValueError(
|
||||
"auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params"
|
||||
)
|
||||
|
||||
embedding_model: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.auto_router_embedding_model
|
||||
embedding_model: Optional[str] = (
|
||||
deployment.litellm_params.auto_router_embedding_model
|
||||
)
|
||||
if embedding_model is None:
|
||||
raise ValueError(
|
||||
"auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params"
|
||||
|
|
@ -6830,13 +6829,13 @@ class Router:
|
|||
ComplexityRouter,
|
||||
)
|
||||
|
||||
complexity_router_config: Optional[
|
||||
dict
|
||||
] = deployment.litellm_params.complexity_router_config
|
||||
complexity_router_config: Optional[dict] = (
|
||||
deployment.litellm_params.complexity_router_config
|
||||
)
|
||||
|
||||
default_model: Optional[
|
||||
str
|
||||
] = deployment.litellm_params.complexity_router_default_model
|
||||
default_model: Optional[str] = (
|
||||
deployment.litellm_params.complexity_router_default_model
|
||||
)
|
||||
|
||||
# If no default model specified, try to get from config tiers
|
||||
if default_model is None and complexity_router_config:
|
||||
|
|
@ -7447,9 +7446,9 @@ class Router:
|
|||
|
||||
# Add custom_llm_provider
|
||||
if deployment.litellm_params.custom_llm_provider:
|
||||
credentials[
|
||||
"custom_llm_provider"
|
||||
] = deployment.litellm_params.custom_llm_provider
|
||||
credentials["custom_llm_provider"] = (
|
||||
deployment.litellm_params.custom_llm_provider
|
||||
)
|
||||
elif "/" in deployment.litellm_params.model:
|
||||
# Extract provider from "provider/model" format
|
||||
credentials["custom_llm_provider"] = deployment.litellm_params.model.split(
|
||||
|
|
@ -9078,12 +9077,16 @@ class Router:
|
|||
cooldown_deployments = await _async_get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span
|
||||
)
|
||||
if verbose_router_logger.isEnabledFor(logging.DEBUG):
|
||||
verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}")
|
||||
_pre_cooldown_deployments = healthy_deployments
|
||||
healthy_deployments = self._filter_cooldown_deployments(
|
||||
healthy_deployments=healthy_deployments,
|
||||
cooldown_deployments=cooldown_deployments,
|
||||
)
|
||||
if not healthy_deployments and self.enable_health_check_routing:
|
||||
verbose_router_logger.warning(
|
||||
"All deployments in cooldown via health-check routing, bypassing cooldown filter"
|
||||
)
|
||||
healthy_deployments = _pre_cooldown_deployments
|
||||
|
||||
healthy_deployments = await self.async_callback_filter_deployments(
|
||||
model=model,
|
||||
|
|
@ -9516,10 +9519,16 @@ class Router:
|
|||
cooldown_deployments = _get_cooldown_deployments(
|
||||
litellm_router_instance=self, parent_otel_span=parent_otel_span
|
||||
)
|
||||
_pre_cooldown_deployments = healthy_deployments
|
||||
healthy_deployments = self._filter_cooldown_deployments(
|
||||
healthy_deployments=healthy_deployments,
|
||||
cooldown_deployments=cooldown_deployments,
|
||||
)
|
||||
if not healthy_deployments and self.enable_health_check_routing:
|
||||
verbose_router_logger.warning(
|
||||
"All deployments in cooldown via health-check routing, bypassing cooldown filter"
|
||||
)
|
||||
healthy_deployments = _pre_cooldown_deployments
|
||||
|
||||
# filter pre-call checks
|
||||
if self.enable_pre_call_checks and messages is not None:
|
||||
|
|
@ -9821,6 +9830,12 @@ class Router:
|
|||
if not self.enable_health_check_routing:
|
||||
return healthy_deployments
|
||||
|
||||
# When allowed_fails_policy is set, cooldown is the sole routing exclusion
|
||||
# mechanism -- skip the binary health check filter so the policy threshold
|
||||
# is respected before any deployment is excluded.
|
||||
if self.allowed_fails_policy is not None:
|
||||
return healthy_deployments
|
||||
|
||||
unhealthy_ids = (
|
||||
await self.health_state_cache.async_get_unhealthy_deployment_ids(
|
||||
parent_otel_span=parent_otel_span
|
||||
|
|
@ -9850,6 +9865,9 @@ class Router:
|
|||
if not self.enable_health_check_routing:
|
||||
return healthy_deployments
|
||||
|
||||
if self.allowed_fails_policy is not None:
|
||||
return healthy_deployments
|
||||
|
||||
unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids(
|
||||
parent_otel_span=parent_otel_span
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,532 @@
|
|||
"""
|
||||
Tests for health check failures integrating with allowed_fails_policy cooldown pipeline.
|
||||
|
||||
When enable_health_check_routing is True and a health check fails, the failure
|
||||
should increment the same counters used by allowed_fails_policy, using the
|
||||
actual exception type from the health check error.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.health_check import run_with_timeout
|
||||
from litellm.router import Router
|
||||
from litellm.types.router import AllowedFailsPolicy
|
||||
|
||||
|
||||
def _make_model(model_id: str, model_name: str = "gpt-4") -> dict:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model_name, "api_key": "fake-key"},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
|
||||
|
||||
class TestAhealthCheckExceptionPreservation:
|
||||
"""Test that ahealth_check() preserves the exception object in its return dict."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_with_timeout_returns_timeout_exception(self):
|
||||
"""run_with_timeout should return a litellm.Timeout in the 'exception' key on timeout."""
|
||||
import asyncio
|
||||
|
||||
async def slow_task():
|
||||
await asyncio.sleep(10)
|
||||
|
||||
result = await run_with_timeout(slow_task(), timeout=0.01)
|
||||
|
||||
assert "error" in result
|
||||
assert "exception" in result
|
||||
assert isinstance(result["exception"], litellm.Timeout)
|
||||
|
||||
|
||||
class TestHealthCheckEndpointExceptionPropagation:
|
||||
"""Test that _perform_health_check propagates exception objects through to unhealthy_endpoints."""
|
||||
|
||||
def test_unhealthy_endpoint_with_exception_dict(self):
|
||||
"""When health check returns {"error": ..., "exception": e}, exception should be in the endpoint."""
|
||||
from litellm.proxy.health_check import _clean_endpoint_data
|
||||
|
||||
auth_error = litellm.AuthenticationError(
|
||||
message="Invalid key", llm_provider="openai", model="gpt-4"
|
||||
)
|
||||
|
||||
# Simulate what _perform_health_check does for an unhealthy dict result
|
||||
is_healthy = {"error": "auth failed", "exception": auth_error}
|
||||
litellm_params = {"model": "gpt-4", "api_key": "fake"}
|
||||
cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details=True)
|
||||
# Exception should be preserved after cleaning
|
||||
if "exception" in is_healthy:
|
||||
cleaned["exception"] = is_healthy["exception"]
|
||||
|
||||
assert cleaned["exception"] is auth_error
|
||||
|
||||
def test_unhealthy_endpoint_raw_exception(self):
|
||||
"""When gather returns a raw Exception, it should be stored in the endpoint dict."""
|
||||
raw_exc = litellm.RateLimitError(
|
||||
message="Rate limited", llm_provider="openai", model="gpt-4"
|
||||
)
|
||||
|
||||
# Simulate the else branch in _perform_health_check
|
||||
from litellm.proxy.health_check import _clean_endpoint_data
|
||||
|
||||
litellm_params = {"model": "gpt-4"}
|
||||
cleaned = _clean_endpoint_data(litellm_params, details=True)
|
||||
if isinstance(raw_exc, Exception):
|
||||
cleaned["exception"] = raw_exc
|
||||
|
||||
assert cleaned["exception"] is raw_exc
|
||||
|
||||
|
||||
class TestGetAllowedFailsFromPolicyWithHealthCheckExceptions:
|
||||
"""Test that get_allowed_fails_from_policy correctly resolves thresholds for health-check exceptions."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"exception_type, policy_field, threshold",
|
||||
[
|
||||
(litellm.Timeout, "TimeoutErrorAllowedFails", 5),
|
||||
(litellm.AuthenticationError, "AuthenticationErrorAllowedFails", 3),
|
||||
(litellm.RateLimitError, "RateLimitErrorAllowedFails", 10),
|
||||
(
|
||||
litellm.ContentPolicyViolationError,
|
||||
"ContentPolicyViolationErrorAllowedFails",
|
||||
2,
|
||||
),
|
||||
(litellm.BadRequestError, "BadRequestErrorAllowedFails", 7),
|
||||
],
|
||||
)
|
||||
def test_policy_resolves_for_health_check_exception_types(
|
||||
self, exception_type, policy_field, threshold
|
||||
):
|
||||
"""Each exception type from a health check should resolve to its policy threshold."""
|
||||
policy = AllowedFailsPolicy(**{policy_field: threshold})
|
||||
router = Router(
|
||||
model_list=[_make_model("d1")],
|
||||
allowed_fails_policy=policy,
|
||||
)
|
||||
exception = exception_type(
|
||||
message="health check failed", llm_provider="openai", model="gpt-4"
|
||||
)
|
||||
result = router.get_allowed_fails_from_policy(exception=exception)
|
||||
assert result == threshold
|
||||
|
||||
def test_policy_returns_none_for_unmatched_exception(self):
|
||||
"""When no policy field matches the exception type, return None (fall back to allowed_fails)."""
|
||||
policy = AllowedFailsPolicy(TimeoutErrorAllowedFails=5)
|
||||
router = Router(
|
||||
model_list=[_make_model("d1")],
|
||||
allowed_fails_policy=policy,
|
||||
)
|
||||
# Use a generic Exception that doesn't match any policy field
|
||||
result = router.get_allowed_fails_from_policy(exception=Exception("generic"))
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestHealthCheckCooldownIntegration:
|
||||
"""Test that health check failures trigger cooldown via _set_cooldown_deployments."""
|
||||
|
||||
def test_health_check_failure_increments_failed_calls(self):
|
||||
"""Health check failure should increment the failed_calls counter."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=3),
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
# First call: should not cooldown (1 <= 3)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=timeout_exc,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Check counter was incremented
|
||||
current_fails = router.failed_calls.get_cache(key="deploy-1")
|
||||
assert current_fails == 1
|
||||
|
||||
def test_health_check_failure_triggers_cooldown_at_threshold(self):
|
||||
"""After exceeding allowed_fails threshold, deployment should enter cooldown."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=2),
|
||||
)
|
||||
|
||||
auth_exc = litellm.AuthenticationError(
|
||||
message="Invalid key", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
# Fails 1 and 2: should not cooldown
|
||||
for _ in range(2):
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=auth_exc,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Fail 3: should trigger cooldown (3 > 2)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=auth_exc,
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_health_check_failure_falls_back_to_allowed_fails(self):
|
||||
"""When policy has no matching field, fall back to generic allowed_fails."""
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
should_cooldown_based_on_allowed_fails_policy,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=10),
|
||||
allowed_fails=1,
|
||||
)
|
||||
|
||||
# Use an exception that doesn't match TimeoutErrorAllowedFails
|
||||
# InternalServerError is not checked by get_allowed_fails_from_policy
|
||||
# so it will fall back to allowed_fails=1
|
||||
generic_exc = Exception("Some internal error")
|
||||
|
||||
# Fail 1: should not cooldown (1 <= 1)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=generic_exc,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
# Fail 2: should trigger cooldown (2 > 1)
|
||||
result = should_cooldown_based_on_allowed_fails_policy(
|
||||
litellm_router_instance=router,
|
||||
deployment="deploy-1",
|
||||
original_exception=generic_exc,
|
||||
)
|
||||
assert result is True
|
||||
|
||||
def test_healthy_endpoints_do_not_trigger_cooldown(self):
|
||||
"""Healthy endpoints should not increment any failure counters."""
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=1),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
# Simulate healthy endpoint -- no exception, no cooldown call
|
||||
healthy_endpoint = {"model_id": "deploy-1"}
|
||||
# Should have no exception key
|
||||
assert "exception" not in healthy_endpoint
|
||||
|
||||
# Verify failed_calls counter is untouched
|
||||
current_fails = router.failed_calls.get_cache(key="deploy-1")
|
||||
assert current_fails is None
|
||||
|
||||
def test_disable_cooldowns_prevents_health_check_cooldown(self):
|
||||
"""When disable_cooldowns=True, health check failures should not trigger cooldown."""
|
||||
from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
disable_cooldowns=True,
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
result = _set_cooldown_deployments(
|
||||
litellm_router_instance=router,
|
||||
original_exception=timeout_exc,
|
||||
exception_status=500,
|
||||
deployment="deploy-1",
|
||||
time_to_cooldown=router.cooldown_time,
|
||||
)
|
||||
assert result is False
|
||||
|
||||
|
||||
class TestWriteHealthStateIntegration:
|
||||
"""Test _write_health_state_to_router_cache integrates with cooldown pipeline."""
|
||||
|
||||
def test_unhealthy_endpoint_triggers_set_cooldown(self):
|
||||
"""_write_health_state_to_router_cache should call _set_cooldown_deployments for unhealthy endpoints."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
timeout_exc = litellm.Timeout(
|
||||
message="Health check timeout", model="", llm_provider=""
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "timeout", "exception": timeout_exc},
|
||||
]
|
||||
healthy_endpoints = [
|
||||
{"model_id": "deploy-2"},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=healthy_endpoints,
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
)
|
||||
mock_cooldown.assert_called_once_with(
|
||||
litellm_router_instance=router,
|
||||
original_exception=timeout_exc,
|
||||
exception_status=408, # Timeout has status_code 408
|
||||
deployment="deploy-1",
|
||||
time_to_cooldown=router.cooldown_time,
|
||||
)
|
||||
|
||||
def test_unhealthy_endpoint_without_exception_skips_cooldown(self):
|
||||
"""Unhealthy endpoints without an exception key should not trigger cooldown."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "unknown failure"}, # no "exception" key
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
) as mock_cooldown:
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
)
|
||||
mock_cooldown.assert_not_called()
|
||||
|
||||
def test_unhealthy_endpoint_increments_failure_counter(self):
|
||||
"""Unhealthy endpoints should call increment_deployment_failures_for_current_minute."""
|
||||
import litellm.proxy.proxy_server as proxy_module
|
||||
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=10),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
rate_exc = litellm.RateLimitError(
|
||||
message="Rate limited", model="gpt-4", llm_provider="openai"
|
||||
)
|
||||
|
||||
unhealthy_endpoints = [
|
||||
{"model_id": "deploy-1", "error": "rate limited", "exception": rate_exc},
|
||||
]
|
||||
|
||||
with patch.object(proxy_module, "llm_router", router):
|
||||
with patch(
|
||||
"litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute"
|
||||
) as mock_increment:
|
||||
with patch(
|
||||
"litellm.router_utils.cooldown_handlers._set_cooldown_deployments"
|
||||
):
|
||||
_write_health_state_to_router_cache(
|
||||
healthy_endpoints=[],
|
||||
unhealthy_endpoints=unhealthy_endpoints,
|
||||
)
|
||||
mock_increment.assert_called_once_with(
|
||||
litellm_router_instance=router,
|
||||
deployment_id="deploy-1",
|
||||
)
|
||||
|
||||
|
||||
class TestHealthCheckFilterBypassWithPolicy:
|
||||
"""
|
||||
When allowed_fails_policy is set, the binary health check filter should be
|
||||
bypassed so cooldown is the sole routing exclusion mechanism.
|
||||
"""
|
||||
|
||||
def test_filter_bypassed_when_policy_set(self):
|
||||
"""Binary health check filter is a no-op when allowed_fails_policy is configured."""
|
||||
import time
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=3),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
# Mark deploy-1 as unhealthy in the health state cache
|
||||
cache = DualCache()
|
||||
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
|
||||
health_cache.set_deployment_health_states(
|
||||
{
|
||||
"deploy-1": {
|
||||
"is_healthy": False,
|
||||
"timestamp": time.time(),
|
||||
"reason": "test",
|
||||
},
|
||||
}
|
||||
)
|
||||
router.health_state_cache = health_cache
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
# Filter should pass all through because policy is set
|
||||
result = router._filter_health_check_unhealthy_deployments(deployments)
|
||||
assert (
|
||||
len(result) == 2
|
||||
), "Binary filter should be bypassed when allowed_fails_policy is set"
|
||||
|
||||
def test_filter_active_when_no_policy(self):
|
||||
"""Binary health check filter still works when no allowed_fails_policy is configured."""
|
||||
import time
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
cache = DualCache()
|
||||
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
|
||||
health_cache.set_deployment_health_states(
|
||||
{
|
||||
"deploy-1": {
|
||||
"is_healthy": False,
|
||||
"timestamp": time.time(),
|
||||
"reason": "test",
|
||||
},
|
||||
}
|
||||
)
|
||||
router.health_state_cache = health_cache
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
result = router._filter_health_check_unhealthy_deployments(deployments)
|
||||
assert len(result) == 1
|
||||
assert result[0]["model_info"]["id"] == "deploy-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_filter_bypassed_when_policy_set(self):
|
||||
"""Async version also bypasses when allowed_fails_policy is set."""
|
||||
import time
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_utils.health_state_cache import DeploymentHealthCache
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=2),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
cache = DualCache()
|
||||
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
|
||||
health_cache.set_deployment_health_states(
|
||||
{
|
||||
"deploy-1": {
|
||||
"is_healthy": False,
|
||||
"timestamp": time.time(),
|
||||
"reason": "test",
|
||||
},
|
||||
}
|
||||
)
|
||||
router.health_state_cache = health_cache
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
result = await router._async_filter_health_check_unhealthy_deployments(
|
||||
deployments
|
||||
)
|
||||
assert len(result) == 2
|
||||
|
||||
|
||||
class TestAllDeploymentsInCooldownSafetyNet:
|
||||
"""
|
||||
When enable_health_check_routing=True and ALL deployments enter cooldown,
|
||||
the async routing path should bypass the cooldown filter and return all
|
||||
deployments rather than blocking all traffic.
|
||||
"""
|
||||
|
||||
def test_raw_cooldown_filter_returns_empty_when_all_cooled(self):
|
||||
"""The raw _filter_cooldown_deployments has no safety net -- it returns empty."""
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
result = router._filter_cooldown_deployments(
|
||||
healthy_deployments=deployments,
|
||||
cooldown_deployments=["deploy-1", "deploy-2"],
|
||||
)
|
||||
assert result == [] # raw filter has no safety net
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_routing_path_bypasses_all_cooldown(self):
|
||||
"""In the async routing path, all-in-cooldown with enable_health_check_routing
|
||||
returns the full list instead of empty (safety net)."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.router_utils.cooldown_handlers import (
|
||||
_async_get_cooldown_deployments,
|
||||
)
|
||||
|
||||
router = Router(
|
||||
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
|
||||
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0),
|
||||
enable_health_check_routing=True,
|
||||
)
|
||||
|
||||
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
|
||||
|
||||
# Simulate all deployments in cooldown
|
||||
with patch(
|
||||
"litellm.router._async_get_cooldown_deployments",
|
||||
new=AsyncMock(return_value=["deploy-1", "deploy-2"]),
|
||||
):
|
||||
# The safety net in async_get_available_deployment should restore
|
||||
# all deployments when the cooldown filter empties the list
|
||||
_pre = deployments.copy()
|
||||
filtered = router._filter_cooldown_deployments(
|
||||
healthy_deployments=deployments,
|
||||
cooldown_deployments=["deploy-1", "deploy-2"],
|
||||
)
|
||||
# If filtered is empty and enable_health_check_routing is True,
|
||||
# the routing path restores _pre_cooldown_deployments
|
||||
if not filtered and router.enable_health_check_routing:
|
||||
filtered = _pre
|
||||
|
||||
assert (
|
||||
len(filtered) == 2
|
||||
), "Safety net should return all deployments when all are in cooldown"
|
||||
Loading…
Add table
Reference in a new issue