Merge f4e6824a08 into f82e4c80ac

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

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

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
@@ -2885,9 +2886,34 @@ 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":
-                    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"]))
@@ -2936,9 +2962,34 @@ 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":
-                    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

tests/unit/test_router/core/test_retries_cooldowns.py

@@ -6,7 +6,9 @@ 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
@@ -28,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"}]
 
@@ -830,3 +833,143 @@ 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])
+
+
+@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'))}"
+
+

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"