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:
KJyothirmayi-2025sl93025 2026-11-06 10:57:46 +05:30
parent 7fbc1c7f48
commit f4e6824a08
2 changed files with 111 additions and 0 deletions

View file

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

View file

@ -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'))}"