Merge 95b3bf4d6b into 429f814a04

litellm/proxy/_types.py

@@ -4534,6 +4534,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_python_core.py

@@ -5609,29 +5609,43 @@ class PythonCore(DeprecatedRouterMembers):
         self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams
     ) -> object:
         """Helper to get stream timeout from kwargs or deployment params"""
-        return (
-            kwargs.get("stream_timeout", None)  # the params dynamically set by user
+        configured: Final = (
+            kwargs.get("stream_timeout", None)
             or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier")
-            or data.get("stream_timeout", None)  # timeout set on litellm_params for this deployment
-            or self.stream_timeout  # timeout set on router
-            or self.request_timeout  # litellm_settings.request_timeout (per-attempt)
+            or data.get("stream_timeout", None)
+            or self.stream_timeout
+            or self.request_timeout
             or self._core_default_litellm_params.get("stream_timeout", None)
         )
+        _deadline: Final = kwargs.get("request_deadline")
+        if not isinstance(_deadline, float):
+            return configured
+        _remaining: Final = _deadline - time.monotonic()
+        if not isinstance(configured, (int, float)):
+            return _remaining
+        return min(float(configured), _remaining)
 
     def _get_non_stream_timeout(
         self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams
     ) -> object:
         """Helper to get non-stream timeout from kwargs or deployment params"""
-        return (
-            kwargs.get("timeout", None)  # the params dynamically set by user
-            or kwargs.get("request_timeout", None)  # the params dynamically set by user
+        configured: Final = (
+            kwargs.get("timeout", None)
+            or kwargs.get("request_timeout", None)
             or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier")
-            or data.get("timeout", None)  # timeout set on litellm_params for this deployment
-            or data.get("request_timeout", None)  # timeout set on litellm_params for this deployment
-            or self.request_timeout  # litellm_settings.request_timeout (per-attempt)
-            or self.timeout  # timeout set on router (router_settings.timeout)
+            or data.get("timeout", None)
+            or data.get("request_timeout", None)
+            or self.request_timeout
+            or self.timeout
             or self._core_default_litellm_params.get("timeout", None)
         )
+        _deadline: Final = kwargs.get("request_deadline")
+        if not isinstance(_deadline, float):
+            return configured
+        _remaining: Final = _deadline - time.monotonic()
+        if not isinstance(configured, (int, float)):
+            return _remaining
+        return min(float(configured), _remaining)
 
     def _get_timeout(self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams) -> object:
         """Helper to get timeout from kwargs or deployment params"""
@@ -9676,6 +9690,9 @@ class PythonCore(DeprecatedRouterMembers):
             if model_list is not None:
                 _metadata.update({"model_group_size": len(model_list)})
 
+        _deadline_raw: Final = kwargs.get("request_deadline")
+        request_deadline: Final[float | None] = _deadline_raw if isinstance(_deadline_raw, float) else None
+
         verbose_router_logger.debug(
             "async function w/ retries: original_function - %s, num_retries - %s", original_function, num_retries
         )
@@ -9777,9 +9794,17 @@ class PythonCore(DeprecatedRouterMembers):
                 all_deployments=_all_deployments,
             )
 
-            await self._host.sleep(retry_after)
+            if request_deadline is not None:
+                _remaining: Final = request_deadline - time.monotonic()
+                if _remaining <= 0 or retry_after >= _remaining:
+                    raise original_exception
+                await self._host.sleep(min(retry_after, _remaining))
+            else:
+                await self._host.sleep(retry_after)
 
             for current_attempt in range(num_retries):
+                if request_deadline is not None and request_deadline - time.monotonic() <= 0:
+                    raise original_exception
                 try:
                     # Update retry tracking metadata before each retry attempt
                     _metadata["attempted_retries"] = current_attempt + 1
@@ -9850,14 +9875,20 @@ class PythonCore(DeprecatedRouterMembers):
                         _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), skipped_ids
                     ):
                         raise
-                    _timeout = self._time_to_sleep_before_retry(
+                    _sleep_duration: Final = self._time_to_sleep_before_retry(
                         e=e,
                         remaining_retries=remaining_retries,
                         num_retries=num_retries,
                         healthy_deployments=_healthy_deployments,
                         all_deployments=_all_deployments,
                     )
-                    await self._host.sleep(_timeout)
+                    if request_deadline is not None:
+                        _loop_remaining: Final = request_deadline - time.monotonic()
+                        if _loop_remaining <= 0 or _sleep_duration >= _loop_remaining:
+                            raise
+                        await self._host.sleep(min(_sleep_duration, _loop_remaining))
+                    else:
+                        await self._host.sleep(_sleep_duration)
 
             if type(original_exception) in self._host.dynamic().exception_types:
                 setattr(original_exception, "max_retries", num_retries)

tests/unit/test_router/test_router_timeout.py

@@ -15,6 +15,7 @@ import respx
 
 import litellm
 from litellm import Router
+from litellm.router_python_host import PythonHost
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 from litellm.router_utils.cooldown_handlers import is_caller_timeout_408
@@ -690,4 +691,346 @@ def test_azure_router_timeout_raises_after_httpx_read_timeout(
             messages=[{"role": "user", "content": "timeout contract"}],
         )
 
-    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_request_deadline_body_timeout_header_wins() -> 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={"x-litellm-timeout": "5"},
+            request_data={"timeout": 30},
+            user_api_key_dict=UserAPIKeyAuth(),
+        )
+
+    assert result.get("timeout") == 5.0, "header timeout must win in data[timeout]"
+    assert result.get("request_deadline") == 5.0, "deadline must use header timeout, not body"
+
+
+@pytest.mark.asyncio
+async def test_retry_skips_sleep_when_deadline_already_expired() -> None:
+    sleep_calls: Final[list[float]] = []
+
+    async def fake_sleep(_host: object, duration: float) -> None:
+        sleep_calls.append(duration)
+
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"},
+            }
+        ],
+        num_retries=3,
+    )
+
+    error: Final = litellm.RateLimitError(
+        message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
+    )
+    error.litellm_response_headers = httpx.Headers({"retry-after": "60"})
+
+    with (
+        patch.object(PythonHost, "sleep", new=fake_sleep),
+        pytest.raises(litellm.RateLimitError),
+    ):
+        await router.async_function_with_retries(
+            original_function=AsyncMock(side_effect=error),
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+            metadata={"model_group": "gpt-3.5-turbo"},
+            request_deadline=time.monotonic() - 1.0,
+            fallbacks=None,
+            context_window_fallbacks=None,
+            content_policy_fallbacks=None,
+        )
+
+    assert sleep_calls == [], f"expected no sleeps when deadline already past, got {sleep_calls}"
+
+
+@pytest.mark.asyncio
+async def test_retry_caps_sleep_to_remaining_budget() -> None:
+    sleep_calls: Final[list[float]] = []
+
+    async def fake_sleep(_host: object, duration: float) -> None:
+        sleep_calls.append(duration)
+
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"},
+            }
+        ],
+        num_retries=1,
+    )
+
+    error: Final = litellm.RateLimitError(
+        message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
+    )
+    error.litellm_response_headers = httpx.Headers({"retry-after": "60"})
+
+    deadline: Final = time.monotonic() + 2.0
+
+    with (
+        patch.object(PythonHost, "sleep", new=fake_sleep),
+        pytest.raises(litellm.RateLimitError),
+    ):
+        await router.async_function_with_retries(
+            original_function=AsyncMock(side_effect=error),
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+            metadata={"model_group": "gpt-3.5-turbo"},
+            request_deadline=deadline,
+            fallbacks=None,
+            context_window_fallbacks=None,
+            content_policy_fallbacks=None,
+        )
+
+    assert all(s <= 2.5 for s in sleep_calls), f"sleep not capped to budget: {sleep_calls}"
+
+
+@pytest.mark.asyncio
+async def test_no_provider_call_after_first_backoff_exhausts_deadline() -> None:
+    """Sleep that runs to the deadline must not be followed by a provider call."""
+    fake_now: Final[list[float]] = [0.0]
+    deadline: Final = 5.0
+    call_count: Final[list[int]] = [0]
+
+    rate_error: Final = litellm.RateLimitError(
+        message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
+    )
+    rate_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
+
+    async def original(*args: object, **kwargs: object) -> object:
+        call_count[0] += 1
+        raise rate_error
+
+    async def sleep_to_deadline(_host: object, _duration: float) -> None:
+        fake_now[0] = deadline
+
+    router: Final = Router(
+        model_list=[
+            {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
+        ],
+        num_retries=3,
+    )
+
+    with (
+        patch.object(time, "monotonic", lambda: fake_now[0]),
+        patch.object(PythonHost, "sleep", new=sleep_to_deadline),
+        pytest.raises(litellm.RateLimitError),
+    ):
+        await router.async_function_with_retries(
+            original_function=original,
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+            metadata={"model_group": "gpt-3.5-turbo"},
+            request_deadline=deadline,
+            fallbacks=None,
+            context_window_fallbacks=None,
+            content_policy_fallbacks=None,
+        )
+
+    assert call_count[0] == 1, f"expected 1 call (initial only), got {call_count[0]}"
+
+
+@pytest.mark.asyncio
+async def test_no_provider_call_after_in_loop_sleep_exhausts_deadline() -> None:
+    """After an in-loop backoff, if the remaining budget < sleep duration, no further call starts."""
+    fake_now: Final[list[float]] = [0.0]
+    deadline: Final = 5.0
+    call_count: Final[list[int]] = [0]
+
+    rate_error: Final = litellm.RateLimitError(
+        message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
+    )
+    rate_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
+
+    async def original(*args: object, **kwargs: object) -> object:
+        call_count[0] += 1
+        raise rate_error
+
+    async def fake_sleep(_host: object, _duration: float) -> None:
+        fake_now[0] = deadline - 0.001
+
+    router: Final = Router(
+        model_list=[
+            {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
+        ],
+        num_retries=3,
+    )
+
+    with (
+        patch.object(time, "monotonic", lambda: fake_now[0]),
+        patch.object(PythonHost, "sleep", new=fake_sleep),
+        pytest.raises(litellm.RateLimitError),
+    ):
+        await router.async_function_with_retries(
+            original_function=original,
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+            metadata={"model_group": "gpt-3.5-turbo"},
+            request_deadline=deadline,
+            fallbacks=None,
+            context_window_fallbacks=None,
+            content_policy_fallbacks=None,
+        )
+
+    assert call_count[0] == 2, f"expected 2 calls (initial + one retry), got {call_count[0]}"
+
+
+@pytest.mark.asyncio
+async def test_retry_proceeds_when_sufficient_budget_remains() -> None:
+    """A retry attempt is made when remaining budget outlasts the backoff sleep."""
+    fake_now: Final[list[float]] = [0.0]
+    deadline: Final = 5.0
+    call_count: Final[list[int]] = [0]
+
+    rate_error: Final = litellm.RateLimitError(
+        message="rate limited", model="gpt-3.5-turbo", llm_provider="openai"
+    )
+    rate_error.litellm_response_headers = httpx.Headers({"retry-after": "1"})
+
+    async def original(*args: object, **kwargs: object) -> object:
+        call_count[0] += 1
+        raise rate_error
+
+    async def fake_sleep(_host: object, _duration: float) -> None:
+        fake_now[0] += 0.1
+
+    router: Final = Router(
+        model_list=[
+            {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
+        ],
+        num_retries=1,
+    )
+
+    with (
+        patch.object(time, "monotonic", lambda: fake_now[0]),
+        patch.object(PythonHost, "sleep", new=fake_sleep),
+        pytest.raises(litellm.RateLimitError),
+    ):
+        await router.async_function_with_retries(
+            original_function=original,
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "hi"}],
+            metadata={"model_group": "gpt-3.5-turbo"},
+            request_deadline=deadline,
+            fallbacks=None,
+            context_window_fallbacks=None,
+            content_policy_fallbacks=None,
+        )
+
+    assert call_count[0] == 2, f"expected 2 calls (initial + one retry), got {call_count[0]}"
+
+
+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"
+
+
+def test_per_attempt_timeout_clamped_to_deadline_remaining() -> None:
+    """When configured timeout > remaining budget, _get_non_stream_timeout must clamp to remaining."""
+    fake_now: Final[list[float]] = [0.0]
+    router: Final = Router(
+        model_list=[
+            {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
+        ]
+    )
+    with patch.object(time, "monotonic", lambda: fake_now[0]):
+        result: Final = router._get_non_stream_timeout(
+            kwargs={"timeout": 30.0, "request_deadline": 2.0},
+            data={},
+        )
+    assert result == 2.0, f"expected timeout clamped to 2.0 (remaining budget), got {result}"
+
+
+def test_per_attempt_timeout_not_extended_beyond_configured() -> None:
+    """When configured timeout < remaining budget, _get_non_stream_timeout must return configured."""
+    fake_now: Final[list[float]] = [0.0]
+    router: Final = Router(
+        model_list=[
+            {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"}}
+        ]
+    )
+    with patch.object(time, "monotonic", lambda: fake_now[0]):
+        result: Final = router._get_non_stream_timeout(
+            kwargs={"timeout": 1.0, "request_deadline": 10.0},
+            data={},
+        )
+    assert result == 1.0, f"expected configured timeout 1.0 to win over deadline remaining 10.0, got {result}"