mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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:
parent
6d54a4eedb
commit
792f120aef
3 changed files with 110 additions and 0 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue