diff --git a/litellm/main.py b/litellm/main.py index eace9c630ba..8f14b1865f7 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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 diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 3e05ee3c484..aeabe492f54 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e2925251dae..4155716a85b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/router.py b/litellm/router.py index 0a3cd1c9c23..427a7cf197f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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 ) diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py new file mode 100644 index 00000000000..7aca9b79653 --- /dev/null +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -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"