mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(router): guard every attempt transition against expired deadline
The sleep guard alone was insufficient: after a capped sleep that finishes past the deadline due to scheduler jitter, the following case "attempt" would start a new provider call with no check. Add a monotonic() >= deadline guard at the case "attempt" transition in both async _run_steps and sync _run_sync, so every route through the retry loop is protected regardless of whether a sleep step preceded it. Add four controlled-clock tests that patch litellm.router_rust.core.monotonic to isolate timing from wall-clock reality: - Sleep advances past deadline -> next attempt blocked, provider called once (async) - Same scenario for the sync _run_sync path - Budget outlasts backoff -> second attempt allowed (positive case) - Fallback without sleep step -> attempt guard fires at attempt transition
This commit is contained in:
parent
7fbc1c7f48
commit
f4e6824a08
2 changed files with 111 additions and 0 deletions
|
|
@ -2886,6 +2886,16 @@ class RustCore(
|
|||
await await_all(driver.outcomes.take_pending())
|
||||
match step["kind"]:
|
||||
case "attempt":
|
||||
_raw_dl_at: Final = driver._request.get("request_deadline") # pyright: ignore[reportPrivateUsage] # same-module internal
|
||||
_dl_at: Final[float | None] = float(_raw_dl_at) if isinstance(_raw_dl_at, (int, float)) else None
|
||||
if _dl_at is not None and monotonic() >= _dl_at:
|
||||
from litellm.exceptions import Timeout
|
||||
|
||||
raise Timeout(
|
||||
message="request deadline exceeded before retry attempt",
|
||||
model=str(driver._request.get("model", "")), # pyright: ignore[reportPrivateUsage] # same-module internal
|
||||
llm_provider="",
|
||||
)
|
||||
outcome = await driver.attempt(step)
|
||||
case "sleep":
|
||||
_secs: Final = _sleep_seconds(step)
|
||||
|
|
@ -2952,6 +2962,16 @@ class RustCore(
|
|||
driver.log_outcomes_sync()
|
||||
match step["kind"]:
|
||||
case "attempt":
|
||||
_raw_dl_ats: Final = driver._request.get("request_deadline") # pyright: ignore[reportPrivateUsage] # same-module internal
|
||||
_dl_ats: Final[float | None] = float(_raw_dl_ats) if isinstance(_raw_dl_ats, (int, float)) else None
|
||||
if _dl_ats is not None and monotonic() >= _dl_ats:
|
||||
from litellm.exceptions import Timeout
|
||||
|
||||
raise Timeout(
|
||||
message="request deadline exceeded before retry attempt",
|
||||
model=str(driver._request.get("model", "")), # pyright: ignore[reportPrivateUsage] # same-module internal
|
||||
llm_provider="",
|
||||
)
|
||||
outcome = driver.attempt_sync(step)
|
||||
case "sleep":
|
||||
_secs_s: Final = _sleep_seconds(step)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ inside ``[base, base + JITTER)`` rather than exactly.
|
|||
from collections.abc import Callable, Mapping
|
||||
from time import monotonic
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -29,6 +30,7 @@ from tests.unit.test_router.core.scripted_host import (
|
|||
server_error,
|
||||
timeout,
|
||||
)
|
||||
from tests.unit.test_router.core.retries_health_helpers import attempt_sync
|
||||
|
||||
MESSAGES: Final = [{"role": "user", "content": "hi"}]
|
||||
|
||||
|
|
@ -882,3 +884,92 @@ async def test_no_deadline_leaves_backoff_unchanged(make_core: MakeCore, host: S
|
|||
assert len(host.sleeps) >= 2
|
||||
assert all(3 <= slept < 3 + JITTER for slept in host.sleeps[:2])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attempt_guarded_after_sleep_advances_clock_past_deadline_async(
|
||||
make_core: MakeCore, host: ScriptedHost
|
||||
) -> None:
|
||||
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
|
||||
host.answer("openai/a", lambda **_: rate_limited(retry_after="0.001"))
|
||||
|
||||
t0: Final = 100.0
|
||||
deadline: Final = t0 + 0.050
|
||||
# First two monotonic() calls are the first-attempt check and the sleep remaining check.
|
||||
# All subsequent calls simulate scheduler overrun, returning a time past the deadline.
|
||||
call_n: Final[list[int]] = [0] # mutable-ok: monotonic() call counter
|
||||
|
||||
def fake_monotonic() -> float:
|
||||
call_n[0] += 1 # mutable-ok: incrementing counter
|
||||
return t0 if call_n[0] <= 2 else t0 + 0.100
|
||||
|
||||
with patch("litellm.router_rust.core.monotonic", fake_monotonic):
|
||||
error: Final = await attempt(core, request_deadline=deadline)
|
||||
|
||||
assert isinstance(error, litellm.Timeout), f"expected Timeout, got {error!r}"
|
||||
assert len(host.calls_to("acompletion")) == 1, f"expected 1 provider call, got {len(host.calls_to('acompletion'))}"
|
||||
|
||||
|
||||
def test_attempt_guarded_after_sleep_advances_clock_past_deadline_sync(
|
||||
make_core: MakeCore, host: ScriptedHost
|
||||
) -> None:
|
||||
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3)
|
||||
host.answer("openai/a", lambda **_: rate_limited(retry_after="0.001"))
|
||||
|
||||
t0: Final = 100.0
|
||||
deadline: Final = t0 + 0.050
|
||||
call_n: Final[list[int]] = [0] # mutable-ok: monotonic() call counter
|
||||
|
||||
def fake_monotonic() -> float:
|
||||
call_n[0] += 1 # mutable-ok: incrementing counter
|
||||
return t0 if call_n[0] <= 2 else t0 + 0.100
|
||||
|
||||
with patch("litellm.router_rust.core.monotonic", fake_monotonic):
|
||||
error: Final = attempt_sync(core, request_deadline=deadline)
|
||||
|
||||
assert isinstance(error, litellm.Timeout), f"expected Timeout, got {error!r}"
|
||||
assert len(host.calls_to("completion")) == 1, f"expected 1 provider call, got {len(host.calls_to('completion'))}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attempt_allowed_when_budget_outlasts_backoff_fake_clock(
|
||||
make_core: MakeCore, host: ScriptedHost
|
||||
) -> None:
|
||||
core: Final = make_core(model_list=[deployment("g", "a")], num_retries=1)
|
||||
host.answer("openai/a", lambda **_: rate_limited(retry_after="0.001"), None)
|
||||
|
||||
t0: Final = 100.0
|
||||
deadline: Final = t0 + 0.050
|
||||
call_n: Final[list[int]] = [0] # mutable-ok: monotonic() call counter
|
||||
|
||||
def fake_monotonic() -> float:
|
||||
call_n[0] += 1 # mutable-ok: incrementing counter
|
||||
return t0 + 0.001 * call_n[0] # advances 1ms per call, stays well within 50ms budget
|
||||
|
||||
with patch("litellm.router_rust.core.monotonic", fake_monotonic):
|
||||
response: Final = await attempt(core, request_deadline=deadline)
|
||||
|
||||
assert not isinstance(response, Exception), f"expected success, got {response!r}"
|
||||
assert len(host.calls_to("acompletion")) == 2, "expected initial attempt + one retry"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_attempt_guarded_on_fallback_without_sleep(make_core: MakeCore, host: ScriptedHost) -> None:
|
||||
core: Final = make_core(model_list=two_deployments())
|
||||
host.answer("openai/a", lambda **_: auth_error())
|
||||
|
||||
t0: Final = 100.0
|
||||
deadline: Final = t0 + 0.050
|
||||
# Clock before deadline on first attempt check, past deadline on any subsequent check.
|
||||
call_n: Final[list[int]] = [0] # mutable-ok: monotonic() call counter
|
||||
|
||||
def fake_monotonic() -> float:
|
||||
call_n[0] += 1 # mutable-ok: incrementing counter
|
||||
return t0 if call_n[0] <= 1 else t0 + 0.100
|
||||
|
||||
with patch("litellm.router_rust.core.monotonic", fake_monotonic):
|
||||
error: Final = await attempt(core, request_deadline=deadline)
|
||||
|
||||
assert isinstance(error, litellm.Timeout), f"expected Timeout, got {error!r}"
|
||||
assert len(host.calls_to("acompletion")) == 1, f"expected 1 provider call, got {len(host.calls_to('acompletion'))}"
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue