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"