From 7fbc1c7f48b7e8a32ba8b17aef8752e851b784bc Mon Sep 17 00:00:00 2001 From: KJyothirmayi-2025sl93025 <2025sl93025@wilp.bits-pilani.ac.in> Date: Fri, 6 Nov 2026 09:37:38 +0530 Subject: [PATCH] fix(router): enforce caller deadline across retries and Retry-After sleeps Closes #49751. When a request carries x-litellm-timeout or x-stainless-timeout, the proxy now computes a monotonic request_deadline (now + timeout_secs) and stores it in LitellmDataForBackendLLMCall. The Rust router core reads that deadline before every backoff sleep: if the deadline has already passed, or the requested sleep would exceed it, the loop raises litellm.Timeout immediately rather than sleeping past the caller's budget. When budget remains, the sleep is capped to the remaining seconds. - litellm/proxy/_types.py: add request_deadline: ReadOnly[float] to LitellmDataForBackendLLMCall - litellm/proxy/litellm_pre_call_utils.py: read x-stainless-timeout as fallback when x-litellm-timeout is absent; compute and store request_deadline from both header and body timeout sources - litellm/router_rust/core.py: check deadline in _run_steps (async) and _run_sync before each sleep step - tests/unit/test_router/core/test_retries_cooldowns.py: four regression tests using ScriptedHost + RouterCore - tests/unit/test_router/test_router_timeout.py: five proxy-layer tests for deadline propagation --- litellm/proxy/_types.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 13 ++- litellm/router_rust/core.py | 35 +++++++- .../core/test_retries_cooldowns.py | 52 ++++++++++++ tests/unit/test_router/test_router_timeout.py | 79 +++++++++++++++++++ 5 files changed, 177 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..99584c5653c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4575,6 +4575,7 @@ class LitellmDataForBackendLLMCall(TypedDict, total=False): # cooldown_handlers._trigger_cooldown_for_failed_deployment). client_side_timeout: bool keepalive_seconds: float | None + request_deadline: ReadOnly[float] class LitellmMetadataFromRequestHeaders(TypedDict, total=False): diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..795636d4ac6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1657,7 +1657,7 @@ class LiteLLMProxyRequestSetup: }); ``` """ - timeout_header: Final = headers.get("x-litellm-timeout", None) + timeout_header: Final = headers.get("x-litellm-timeout") or headers.get("x-stainless-timeout") if timeout_header is not None: return float(timeout_header) return None @@ -1895,6 +1895,17 @@ class LiteLLMProxyRequestSetup: if header_timeout is not None: data["timeout"] = header_timeout + _body_timeout: Final = request_data.get("timeout") + _deadline_secs: Final = ( + header_timeout + if header_timeout is not None + else float(_body_timeout) + if isinstance(_body_timeout, (int, float)) + else None + ) + if _deadline_secs is not None: + data["request_deadline"] = time.monotonic() + _deadline_secs + header_stream_timeout: Final = LiteLLMProxyRequestSetup._get_stream_timeout_from_request(headers) if header_stream_timeout is not None: data["stream_timeout"] = header_stream_timeout diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index c6ca4a1f399..4d932287d24 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -25,6 +25,7 @@ from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, from contextvars import ContextVar from datetime import date, datetime, time from functools import partial +from time import monotonic from time import perf_counter from time import time as unix_time from types import MappingProxyType @@ -2887,7 +2888,22 @@ class RustCore( case "attempt": outcome = await driver.attempt(step) case "sleep": - await self._host.sleep(_sleep_seconds(step)) + _secs: Final = _sleep_seconds(step) + _raw_dl: Final = driver._request.get("request_deadline") # pyright: ignore[reportPrivateUsage] # same-module internal + _dl: Final[float | None] = float(_raw_dl) if isinstance(_raw_dl, (int, float)) else None + if _dl is not None: + _remaining: Final = _dl - monotonic() + if _remaining <= 0 or _secs >= _remaining: + from litellm.exceptions import Timeout + + raise Timeout( + message="request deadline exceeded during retry backoff", + model=str(driver._request.get("model", "")), # pyright: ignore[reportPrivateUsage] # same-module internal + llm_provider="", + ) + await self._host.sleep(min(_secs, _remaining)) + else: + await self._host.sleep(_secs) outcome = None case "allow_fallback": outcome = await driver.allow_fallback(str(step["target"])) @@ -2938,7 +2954,22 @@ class RustCore( case "attempt": outcome = driver.attempt_sync(step) case "sleep": - self._host.sleep_sync(_sleep_seconds(step)) + _secs_s: Final = _sleep_seconds(step) + _raw_dl_s: Final = driver._request.get("request_deadline") # pyright: ignore[reportPrivateUsage] # same-module internal + _dl_s: Final[float | None] = float(_raw_dl_s) if isinstance(_raw_dl_s, (int, float)) else None + if _dl_s is not None: + _remaining_s: Final = _dl_s - monotonic() + if _remaining_s <= 0 or _secs_s >= _remaining_s: + from litellm.exceptions import Timeout + + raise Timeout( + message="request deadline exceeded during retry backoff", + model=str(driver._request.get("model", "")), # pyright: ignore[reportPrivateUsage] # same-module internal + llm_provider="", + ) + self._host.sleep_sync(min(_secs_s, _remaining_s)) + else: + self._host.sleep_sync(_secs_s) outcome = None case "allow_fallback": from litellm.litellm_core_utils import asyncify diff --git a/tests/unit/test_router/core/test_retries_cooldowns.py b/tests/unit/test_router/core/test_retries_cooldowns.py index 004b1762240..c0fd549379a 100644 --- a/tests/unit/test_router/core/test_retries_cooldowns.py +++ b/tests/unit/test_router/core/test_retries_cooldowns.py @@ -6,6 +6,7 @@ inside ``[base, base + JITTER)`` rather than exactly. """ from collections.abc import Callable, Mapping +from time import monotonic from typing import Final import httpx @@ -830,3 +831,54 @@ async def test_health_record_failure_cools_down_past_allowed_fails(make_core: Ma assert (first, second) == (False, True) assert set(cooled_down(core, ["a", "b"])) == {"a"} assert core.health().failures_this_minute("a") == 2 + + +# --- request_deadline enforcement --------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_expired_deadline_raises_timeout_without_sleeping(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="60")) + + past_deadline: Final = monotonic() - 1.0 + error: Final = await attempt(core, request_deadline=past_deadline) + + assert isinstance(error, litellm.Timeout), f"expected Timeout, got {error!r}" + assert host.sleeps == [], f"no sleep should occur when deadline already passed: {host.sleeps}" + + +@pytest.mark.asyncio +async def test_retry_after_header_capped_to_remaining_budget(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="60")) + + deadline: Final = monotonic() + 2.0 + error: Final = await attempt(core, request_deadline=deadline) + + assert isinstance(error, litellm.Timeout), f"expected Timeout, got {error!r}" + assert all(s <= 2.5 for s in host.sleeps), f"sleep exceeded 2-second budget: {host.sleeps}" + + +@pytest.mark.asyncio +async def test_retry_proceeds_when_budget_outlasts_backoff(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) + + deadline: Final = monotonic() + 30.0 + 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_no_deadline_leaves_backoff_unchanged(make_core: MakeCore, host: ScriptedHost) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=2) + host.answer("openai/a", lambda **_: rate_limited(retry_after="3")) + + await attempt(core) + + assert len(host.sleeps) >= 2 + assert all(3 <= slept < 3 + JITTER for slept in host.sleeps[:2]) + diff --git a/tests/unit/test_router/test_router_timeout.py b/tests/unit/test_router/test_router_timeout.py index 3fafb9e0f04..5673b1f5abb 100644 --- a/tests/unit/test_router/test_router_timeout.py +++ b/tests/unit/test_router/test_router_timeout.py @@ -650,3 +650,82 @@ def test_azure_router_timeout_raises_after_httpx_read_timeout( ) assert route.call_count == 1 + + +def test_request_deadline_set_when_litellm_timeout_header_present() -> None: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + before: Final = time.monotonic() + result: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={"x-litellm-timeout": "10"}, + request_data={}, + user_api_key_dict=UserAPIKeyAuth(), + ) + after: Final = time.monotonic() + + assert result.get("timeout") == 10.0 + deadline: Final = result.get("request_deadline") + assert isinstance(deadline, float) + assert before + 10.0 <= deadline <= after + 10.0 + + +def test_request_deadline_set_when_stainless_timeout_header_present() -> None: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + before: Final = time.monotonic() + result: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={"x-stainless-timeout": "15"}, + request_data={}, + user_api_key_dict=UserAPIKeyAuth(), + ) + after: Final = time.monotonic() + + assert result.get("timeout") == 15.0 + deadline: Final = result.get("request_deadline") + assert isinstance(deadline, float) + assert before + 15.0 <= deadline <= after + 15.0 + + +def test_request_deadline_not_set_when_no_timeout_header() -> None: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + result: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={}, + request_data={}, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert result.get("request_deadline") is None + + +def test_request_deadline_set_from_body_timeout() -> None: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + fake_now: Final[list[float]] = [0.0] + with patch.object(time, "monotonic", lambda: fake_now[0]): + result: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={}, + request_data={"timeout": 10}, + user_api_key_dict=UserAPIKeyAuth(), + ) + + deadline: Final = result.get("request_deadline") + assert isinstance(deadline, float), "request_deadline must be set when body timeout is present" + assert deadline == 10.0, f"expected deadline 10.0 (now=0 + 10), got {deadline}" + + +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 + + result: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers={"x-litellm-timeout": "5", "x-stainless-timeout": "30"}, + request_data={}, + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert result.get("timeout") == 5.0, "x-litellm-timeout must win when both headers present"