diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index a0b69ecb0bf..b696ed980bf 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1834,6 +1834,10 @@ class ProxyBaseLLMRequestProcessing: llm_router: Router | None, ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.common_utils.request_pressure_metrics import ( + clear_request_shed_marker, + mark_request_shed_by_proxy, + ) try: return await self.common_processing_pre_call_logic( @@ -1876,6 +1880,11 @@ class ProxyBaseLLMRequestProcessing: if fallback_model == original_model: continue self.data["model"] = fallback_model + # This rejection is being recovered from, not returned, so it + # must not leave the request marked as shed by this proxy. A + # provider 429 on the fallback would otherwise be counted as + # this pod shedding load. + clear_request_shed_marker() try: return await self.common_processing_pre_call_logic( request=request, @@ -1900,6 +1909,9 @@ class ProxyBaseLLMRequestProcessing: raise self.data["model"] = original_model + # Every fallback was rate limited too, so the original rejection is + # what the client gets and the mark has to be restored. + mark_request_shed_by_proxy() raise original_exc def _resolve_fallback_models( diff --git a/litellm/proxy/common_utils/request_pressure_metrics.py b/litellm/proxy/common_utils/request_pressure_metrics.py index ec9ef3036e6..3b2c8d462f9 100644 --- a/litellm/proxy/common_utils/request_pressure_metrics.py +++ b/litellm/proxy/common_utils/request_pressure_metrics.py @@ -34,6 +34,17 @@ def mark_request_shed_by_proxy() -> None: proxy_shed_request.set(True) +def clear_request_shed_marker() -> None: + """Undo the mark because the rejection is being recovered from, not returned. + + The mark is set when the error is constructed, which is the only point every + raise site passes through. A caught rejection that falls back to another + model would otherwise leave the request marked, and a provider 429 on that + fallback would be counted as this pod shedding load. + """ + proxy_shed_request.set(False) + + def was_request_shed_by_proxy() -> bool: return proxy_shed_request.get() diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 355c6d27eb2..3825cc72857 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4955,6 +4955,93 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == fallback_model assert call_count == 2 + @pytest.mark.asyncio + async def test_a_recovered_rate_limit_does_not_leave_the_request_marked_as_shed(self): + """The mark is set when the error is constructed, so a rejection that is + recovered from by falling back would otherwise stay marked and make a + provider 429 on the fallback count as this pod shedding load.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.common_utils.request_pressure_metrics import ( + clear_request_shed_marker, + was_request_shed_by_proxy, + ) + + clear_request_shed_marker() + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + marked_during_fallback = [] + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded for gpt-4") + marked_during_fallback.append(was_request_shed_by_proxy()) + return processor.data, MagicMock() + + mock_router = MagicMock() + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert marked_during_fallback == [False], "the fallback attempt must not inherit the recovered rejection's mark" + assert not was_request_shed_by_proxy() + + @pytest.mark.asyncio + async def test_a_rejection_returned_to_the_client_stays_marked_as_shed(self): + """When every fallback is rate limited too, the original rejection is what + the client gets, so it must still count as this pod shedding.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.proxy.common_utils.request_pressure_metrics import ( + clear_request_shed_marker, + was_request_shed_by_proxy, + ) + + clear_request_shed_marker() + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + raise ProxyRateLimitError(detail="TPM limit exceeded") + + mock_router = MagicMock() + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + with pytest.raises(ProxyRateLimitError): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None), + version=None, + proxy_config=MagicMock(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert was_request_shed_by_proxy() + @pytest.mark.asyncio async def test_raises_when_no_fallbacks_configured(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError