From 792f120aef98e7a111e3cff0668d432e5cc3d799 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Wed, 19 Aug 2026 14:29:20 -0700 Subject: [PATCH] fix(observability): drop the shed mark when a rejection is recovered from The mark is set when `ProxyRateLimitError` is constructed, which is the one point every raise site passes through. That also means a rejection the proxy recovers from stays marked: `_pre_call_with_fallbacks` catches the local rate limit and retries against a fallback model, so a provider 429 on that fallback was counted as this pod shedding load. That is the exact conflation the mark exists to prevent, reintroduced through the back door. The mark is now cleared before each fallback attempt, so only a rejection this proxy actually returns counts. The restore before `raise original_exc` states the postcondition at the raise site. Current control flow already satisfies it, because reaching that line means every fallback raised its own `ProxyRateLimitError` and re-marked on construction, so no test distinguishes it. It is kept so the guarantee does not depend on that incidental re-marking. --- litellm/proxy/common_request_processing.py | 12 +++ .../common_utils/request_pressure_metrics.py | 11 +++ .../proxy/test_common_request_processing.py | 87 +++++++++++++++++++ 3 files changed, 110 insertions(+) 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