diff --git a/litellm/router.py b/litellm/router.py index 05c03dd69f7..f10115863b2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -179,7 +179,7 @@ from litellm.router_utils.cooldown_handlers import ( _get_cooldown_deployments, _set_cooldown_deployments, is_advisor_orchestration_failure, - is_background_response_cost_poll_failure, + is_background_response_cost_poll_not_found, is_caller_timeout_408, ) from litellm.router_utils.fallback_event_handlers import ( @@ -8142,10 +8142,10 @@ class Router: litellm_params: Final = kwargs.get("litellm_params", {}) _model_info: Final = litellm_params.get("model_info", {}) - if is_background_response_cost_poll_failure(litellm_params): + if is_background_response_cost_poll_not_found(exception, litellm_params): verbose_router_logger.debug( "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Failure came from the background response cost poll, not the deployment's health." + "Provider 404 came from the background response cost poll, not the deployment's health." ) return False diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index cda1cca497c..408ddbab34b 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -64,12 +64,9 @@ def is_advisor_orchestration_failure(exception: BaseException | None) -> bool: return bool(getattr(exception, _ADVISOR_ORCHESTRATION_FAILURE_ATTR, False)) -def is_background_response_cost_poll_failure(litellm_params: Mapping[str, object]) -> bool: - """Whether the failed call was the enterprise cost poller reading back a stored background response. - - A provider 404 there means the provider dropped the stored object, not that the deployment is unhealthy. - """ - return any( +def is_background_response_cost_poll_not_found(exception: Exception, litellm_params: Mapping[str, object]) -> bool: + """Whether a background response cost poll failed with a provider 404.""" + return getattr(exception, "status_code", None) == 404 and any( isinstance(candidate, Mapping) and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN for candidate in (litellm_params.get("metadata"), litellm_params.get("litellm_metadata")) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index e6037bebabe..184092e4096 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -8789,6 +8789,7 @@ class TestBackgroundResponseCostPollCooldown: "model_info": {"id": "dep-1"}, } ], + allowed_fails=0, ) def _cooled_down_ids(self, router): @@ -8801,16 +8802,13 @@ class TestBackgroundResponseCostPollCooldown: ) def _deployment_callback_on_failure(self, router, kwargs): - import asyncio from datetime import datetime - async def callback(): - now = datetime.now() - return router.deployment_callback_on_failure(kwargs, None, now, now) + now = datetime.now() + return router.deployment_callback_on_failure(kwargs, None, now, now) - return asyncio.run(callback()) - - def test_untagged_not_found_cools_down_deployment(self): + @pytest.mark.asyncio + async def test_untagged_not_found_cools_down_deployment(self): router = self._router() assert ( self._deployment_callback_on_failure( @@ -8856,7 +8854,33 @@ class TestBackgroundResponseCostPollCooldown: value = get_deployment_failures_for_current_minute(litellm_router_instance=router, deployment_id="dep-1") assert not value - def test_other_internal_origin_not_found_still_cools_down_deployment(self): + @pytest.mark.asyncio + async def test_cost_poll_non_404_still_cools_down_deployment(self): + from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY + from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + + router = self._router() + assert ( + self._deployment_callback_on_failure( + router, + { + "exception": litellm.InternalServerError( + message="upstream 500", llm_provider="openai", model="gpt-4.1" + ), + "litellm_params": { + "model_info": {"id": "dep-1"}, + "litellm_metadata": { + INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN + }, + }, + }, + ) + is True + ) + assert "dep-1" in self._cooled_down_ids(router) + + @pytest.mark.asyncio + async def test_other_internal_origin_not_found_still_cools_down_deployment(self): from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN