diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 4d932287d24..d69ca08eba4 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -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) diff --git a/tests/unit/test_router/core/test_retries_cooldowns.py b/tests/unit/test_router/core/test_retries_cooldowns.py index c0fd549379a..3d0575d4d85 100644 --- a/tests/unit/test_router/core/test_retries_cooldowns.py +++ b/tests/unit/test_router/core/test_retries_cooldowns.py @@ -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'))}" + +