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.
This commit is contained in:
Yucheng Zhu 2026-08-19 14:29:20 -07:00
parent 6d54a4eedb
commit 792f120aef
3 changed files with 110 additions and 0 deletions

View file

@ -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(

View file

@ -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()

View file

@ -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