scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(router): add pre-call deadline check after backoff and fix test patching

Browse source 
Add a deadline check immediately before each provider attempt in the retry
loop so that a sleep that consumes all remaining budget cannot be followed
by a provider call. Also add early bail in both backoff paths when the
sleep duration >= remaining budget (avoids a useless sleep before the
inevitable rejection).

Fix the sleep-patching pattern in deadline tests: PythonHost uses __slots__
which prevents instance-level attribute assignment, so patch the class
method instead. Restore the accidentally-dropped
test_litellm_timeout_takes_priority_over_stainless_timeout function header,
and add metadata kwarg required by _record_failed_attempt.

Three new tests assert the exact invariant the reviewer asked for: zero
additional provider calls when the first-backoff sleep exhausts the
deadline, one retry when the in-loop sleep nearly exhausts it, and a
positive control confirming retries do proceed with sufficient budget.
This commit is contained in:
KJyothirmayi-2025sl93025 2026-11-04 18:42:48 +05:30
parent a5c339168d
commit 3d9adc3229
2 changed files with 146 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/router_python_core.py
View file

@ -9658,13 +9658,15 @@ class PythonCore(DeprecatedRouterMembers):
if request_deadline is not None:
_remaining: Final = request_deadline - time.monotonic()
if _remaining <= 0:
if _remaining <= 0 or retry_after >= _remaining:
raise original_exception
await self._host.sleep(min(retry_after, _remaining))
else:
await self._host.sleep(retry_after)
for current_attempt in range(num_retries):
if request_deadline is not None and request_deadline - time.monotonic() <= 0:
raise original_exception
try:
# Update retry tracking metadata before each retry attempt
_metadata["attempted_retries"] = current_attempt + 1
@ -9741,7 +9743,7 @@ class PythonCore(DeprecatedRouterMembers):
)
if request_deadline is not None:
_loop_remaining: Final = request_deadline - time.monotonic()
if _loop_remaining <= 0:
if _loop_remaining <= 0 or _sleep_duration >= _loop_remaining:
raise e
await self._host.sleep(min(_sleep_duration, _loop_remaining))
else:

146
tests/unit/test_router/test_router_timeout.py
View file

@ -15,6 +15,7 @@ import respx
import litellm
from litellm import Router
from litellm.router_python_host import PythonHost
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.router_utils.cooldown_handlers import is_caller_timeout_408
@ -744,7 +745,7 @@ def test_request_deadline_not_set_when_no_timeout_header() -> None:
async def test_retry_skips_sleep_when_deadline_already_expired() -> None:
sleep_calls: Final[list[float]] = []
async def fake_sleep(duration: float) -> None:
async def fake_sleep(_host: object, duration: float) -> None:
sleep_calls.append(duration)
router: Final = Router(
@ -763,13 +764,14 @@ async def test_retry_skips_sleep_when_deadline_already_expired() -> None:
error.litellm_response_headers = httpx.Headers({"retry-after": "60"})
with (
patch.object(router._host, "sleep", new=fake_sleep), # type: ignore[attr-defined]
patch.object(PythonHost, "sleep", new=fake_sleep),
pytest.raises(litellm.RateLimitError),
):
await router.async_function_with_retries(
original_function=AsyncMock(side_effect=error),
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
metadata={"model_group": "gpt-3.5-turbo"},
request_deadline=time.monotonic() - 1.0,
fallbacks=None,
context_window_fallbacks=None,
@ -783,7 +785,7 @@ async def test_retry_skips_sleep_when_deadline_already_expired() -> None:
async def test_retry_caps_sleep_to_remaining_budget() -> None:
sleep_calls: Final[list[float]] = []
async def fake_sleep(duration: float) -> None:
async def fake_sleep(_host: object, duration: float) -> None:
sleep_calls.append(duration)
router: Final = Router(
@ -804,13 +806,14 @@ async def test_retry_caps_sleep_to_remaining_budget() -> None:
deadline: Final = time.monotonic() + 2.0
with (
patch.object(router._host, "sleep", new=fake_sleep), # type: ignore[attr-defined]
patch.object(PythonHost, "sleep", new=fake_sleep),
pytest.raises(litellm.RateLimitError),
):
await router.async_function_with_retries(
original_function=AsyncMock(side_effect=error),
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
metadata={"model_group": "gpt-3.5-turbo"},
request_deadline=deadline,
fallbacks=None,
context_window_fallbacks=None,
@ -820,6 +823,141 @@ async def test_retry_caps_sleep_to_remaining_budget() -> None:
assert all(s <= 2.5 for s in sleep_calls), f"sleep not capped to budget: {sleep_calls}"
@pytest.mark.asyncio
async def test_no_provider_call_after_first_backoff_exhausts_deadline() -> None:
"""Sleep that runs to the deadline must not be followed by a provider call."""
fake_now: Final[list[float]] = [0.0]
deadline: Final = 5.0
call_count: Final[list[int]] = [0]
rate_error: Final = litellm.RateLimitError(
message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
)
rate_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
async def original(*args: object, **kwargs: object) -> object:
call_count[0] += 1
raise rate_error
async def sleep_to_deadline(_host: object, _duration: float) -> None:
fake_now[0] = deadline
router: Final = Router(
model_list=[
{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
],
num_retries=3,
)
with (
patch.object(time, "monotonic", lambda: fake_now[0]),
patch.object(PythonHost, "sleep", new=sleep_to_deadline),
pytest.raises(litellm.RateLimitError),
):
await router.async_function_with_retries(
original_function=original,
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
metadata={"model_group": "gpt-3.5-turbo"},
request_deadline=deadline,
fallbacks=None,
context_window_fallbacks=None,
content_policy_fallbacks=None,
)
assert call_count[0] == 1, f"expected 1 call (initial only), got {call_count[0]}"
@pytest.mark.asyncio
async def test_no_provider_call_after_in_loop_sleep_exhausts_deadline() -> None:
"""After an in-loop backoff, if the remaining budget < sleep duration, no further call starts."""
fake_now: Final[list[float]] = [0.0]
deadline: Final = 5.0
call_count: Final[list[int]] = [0]
rate_error: Final = litellm.RateLimitError(
message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
)
rate_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
async def original(*args: object, **kwargs: object) -> object:
call_count[0] += 1
raise rate_error
async def fake_sleep(_host: object, _duration: float) -> None:
fake_now[0] = deadline - 0.001
router: Final = Router(
model_list=[
{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
],
num_retries=3,
)
with (
patch.object(time, "monotonic", lambda: fake_now[0]),
patch.object(PythonHost, "sleep", new=fake_sleep),
pytest.raises(litellm.RateLimitError),
):
await router.async_function_with_retries(
original_function=original,
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
metadata={"model_group": "gpt-3.5-turbo"},
request_deadline=deadline,
fallbacks=None,
context_window_fallbacks=None,
content_policy_fallbacks=None,
)
assert call_count[0] == 2, f"expected 2 calls (initial + one retry), got {call_count[0]}"
@pytest.mark.asyncio
async def test_retry_proceeds_when_sufficient_budget_remains() -> None:
"""A retry attempt is made when remaining budget outlasts the backoff sleep."""
fake_now: Final[list[float]] = [0.0]
deadline: Final = 5.0
call_count: Final[list[int]] = [0]
rate_error: Final = litellm.RateLimitError(
message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
)
rate_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
async def original(*args: object, **kwargs: object) -> object:
call_count[0] += 1
raise rate_error
async def fake_sleep(_host: object, _duration: float) -> None:
fake_now[0] += 0.1
router: Final = Router(
model_list=[
{"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
],
num_retries=1,
)
with (
patch.object(time, "monotonic", lambda: fake_now[0]),
patch.object(PythonHost, "sleep", new=fake_sleep),
pytest.raises(litellm.RateLimitError),
):
await router.async_function_with_retries(
original_function=original,
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "hi"}],
metadata={"model_group": "gpt-3.5-turbo"},
request_deadline=deadline,
fallbacks=None,
context_window_fallbacks=None,
content_policy_fallbacks=None,
)
assert call_count[0] == 2, f"expected 2 calls (initial + one retry), got {call_count[0]}"
def test_litellm_timeout_takes_priority_over_stainless_timeout() -> None:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup