mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #41891 from BerriAI/litellm_overhead_window_from_proxy_receive
fix(timing): anchor response duration and overhead at proxy receive time
This commit is contained in:
kerry-berri
GitHub
commit
a1f3124e18
7 changed files with 353 additions and 20 deletions
|
|
@ -1,11 +1,12 @@
|
|||
import datetime
|
||||
from collections.abc import Mapping
|
||||
from functools import reduce
|
||||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.constants import LITELLM_DETAILED_TIMING
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, process_response_headers
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
from litellm.litellm_core_utils.logging_utils import LiteLLMLoggingObject
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -16,19 +17,59 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
def _timing_window_start(
|
||||
start_time: datetime.datetime, logging_obj: LiteLLMLoggingObject
|
||||
) -> tuple[datetime.datetime, bool]:
|
||||
received_at: Final = get_litellm_metadata_from_kwargs(logging_obj.model_call_details).get("litellm_received_at")
|
||||
if isinstance(received_at, datetime.datetime):
|
||||
return received_at, True
|
||||
return start_time, False
|
||||
|
||||
|
||||
def _union_duration_ms(windows: object, lower: float, upper: float) -> float | None:
|
||||
if not isinstance(windows, (list, tuple)):
|
||||
return None
|
||||
clipped: Final[tuple[tuple[float, float], ...]] = tuple(
|
||||
(max(lower, float(window[0])), min(upper, float(window[1])))
|
||||
for window in windows
|
||||
if isinstance(window, (list, tuple))
|
||||
and len(window) == 2
|
||||
and isinstance(window[0], (int, float))
|
||||
and isinstance(window[1], (int, float))
|
||||
and max(lower, float(window[0])) < min(upper, float(window[1]))
|
||||
)
|
||||
if not clipped:
|
||||
return None
|
||||
|
||||
ordered: Final[tuple[tuple[float, float], ...]] = tuple(sorted(clipped))
|
||||
|
||||
def merge_window(
|
||||
merged: tuple[tuple[float, float], ...], current: tuple[float, float]
|
||||
) -> tuple[tuple[float, float], ...]:
|
||||
if not merged or current[0] > merged[-1][1]:
|
||||
return (*merged, current)
|
||||
return (*merged[:-1], (merged[-1][0], max(merged[-1][1], current[1])))
|
||||
|
||||
merged: Final[tuple[tuple[float, float], ...]] = reduce(merge_window, ordered, ())
|
||||
return sum(end - start for start, end in merged) * 1000
|
||||
|
||||
|
||||
def response_timing_metrics(
|
||||
start_time: datetime.datetime,
|
||||
end_time: datetime.datetime,
|
||||
logging_obj: LiteLLMLoggingObject,
|
||||
include_overhead: bool = True,
|
||||
) -> Mapping[str, float]:
|
||||
"""``_response_ms`` for the whole call, plus ``litellm_overhead_time_ms`` when it can be derived.
|
||||
"""``_response_ms`` for the window starting at proxy receive time when stamped, else ``start_time``.
|
||||
|
||||
On a cache hit the overhead is the total minus the cache read; otherwise it is the total minus
|
||||
the provider call (``llm_api_duration_ms``). It is omitted when neither duration was recorded,
|
||||
and when ``include_overhead`` is False because the two durations cover different windows.
|
||||
"""
|
||||
total_response_time_ms: Final = (end_time - start_time).total_seconds() * 1000
|
||||
timing_window: Final = _timing_window_start(start_time, logging_obj)
|
||||
window_start: Final = timing_window[0]
|
||||
receive_anchored: Final = timing_window[1]
|
||||
total_response_time_ms: Final = (end_time.timestamp() - window_start.timestamp()) * 1000
|
||||
if not include_overhead:
|
||||
return {"_response_ms": total_response_time_ms} # mutable-ok: read-only timing result
|
||||
caching_details: Final = logging_obj.caching_details
|
||||
|
|
@ -37,11 +78,22 @@ def response_timing_metrics(
|
|||
if caching_details is not None and caching_details.get("cache_hit") is True
|
||||
else None
|
||||
)
|
||||
metadata: Final[Mapping[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details)
|
||||
llm_api_duration_ms: Final = logging_obj.model_call_details.get("llm_api_duration_ms")
|
||||
if cache_duration_ms is not None:
|
||||
overhead_ms: float | None = total_response_time_ms - cache_duration_ms
|
||||
elif llm_api_duration_ms is not None:
|
||||
overhead_ms = round(total_response_time_ms - llm_api_duration_ms, 4)
|
||||
provider_duration_ms: Final[float | None] = (
|
||||
_union_duration_ms(
|
||||
metadata.get("llm_api_timing_windows"),
|
||||
window_start.timestamp(),
|
||||
end_time.timestamp(),
|
||||
)
|
||||
if receive_anchored
|
||||
else None
|
||||
)
|
||||
effective: Final = provider_duration_ms if provider_duration_ms is not None else llm_api_duration_ms
|
||||
overhead_ms = round(total_response_time_ms - effective, 4) if isinstance(effective, (int, float)) else None
|
||||
else:
|
||||
overhead_ms = None
|
||||
if overhead_ms is None:
|
||||
|
|
@ -152,7 +204,8 @@ class ResponseMetadata:
|
|||
# pre-processing = time from request start to LLM API call start
|
||||
api_call_start: Final[datetime.datetime | None] = logging_obj.model_call_details.get("api_call_start_time")
|
||||
if api_call_start is not None and start_time is not None:
|
||||
pre_ms: Final = (api_call_start - start_time).total_seconds() * 1000
|
||||
anchor: Final = _timing_window_start(start_time, logging_obj)[0]
|
||||
pre_ms: Final = (api_call_start.timestamp() - anchor.timestamp()) * 1000
|
||||
detailed["timing_pre_processing_ms"] = round(pre_ms, 4)
|
||||
|
||||
# post-processing = total - pre - llm_api
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from litellm.constants import (
|
|||
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS,
|
||||
MAX_BASE64_LENGTH_FOR_LOGGING,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
|
|
@ -286,6 +287,20 @@ def _set_duration_in_model_call_details(
|
|||
duration_ms: Final = (end_time - start_time).total_seconds() * 1000
|
||||
if logging_obj and hasattr(logging_obj, "model_call_details"):
|
||||
logging_obj.model_call_details["llm_api_duration_ms"] = duration_ms
|
||||
metadata: Final[dict[str, object]] = get_litellm_metadata_from_kwargs(logging_obj.model_call_details)
|
||||
recorded: Final = metadata.get("llm_api_timing_windows")
|
||||
earlier: Final[tuple[tuple[float, float], ...]] = tuple(
|
||||
(float(window[0]), float(window[1]))
|
||||
for window in (recorded if isinstance(recorded, (list, tuple)) else ())
|
||||
if isinstance(window, (list, tuple))
|
||||
and len(window) == 2
|
||||
and isinstance(window[0], (int, float))
|
||||
and isinstance(window[1], (int, float))
|
||||
)
|
||||
metadata["llm_api_timing_windows"] = (
|
||||
*earlier,
|
||||
(start_time.timestamp(), end_time.timestamp()),
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug("`logging_obj` not found - unable to track `llm_api_duration_ms")
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -2367,6 +2367,7 @@ async def add_litellm_data_to_request(
|
|||
# OTel layer can compute pre-request latency, including on the failure
|
||||
# path after the logging object is popped.
|
||||
data[_metadata_variable_name]["litellm_received_at"] = getattr(request.state, "litellm_received_at", None)
|
||||
data[_metadata_variable_name]["llm_api_timing_windows"] = ()
|
||||
|
||||
# OTEL Controls / Tracing
|
||||
# Add the OTEL Parent Trace before sending it LiteLLM
|
||||
|
|
|
|||
|
|
@ -9,11 +9,14 @@ import asyncio
|
|||
import datetime
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm.litellm_core_utils.llm_response_utils.response_metadata as response_metadata_mod
|
||||
import litellm.proxy.common_request_processing as common_request_processing_mod
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
ResponseMetadata,
|
||||
_union_duration_ms,
|
||||
response_timing_metrics,
|
||||
update_response_metadata,
|
||||
)
|
||||
|
|
@ -70,9 +73,7 @@ class TestCallbackDurationMs:
|
|||
def test_update_response_metadata_includes_callback_duration(self):
|
||||
"""End-to-end: update_response_metadata should propagate callback_duration_ms."""
|
||||
result = ModelResponse()
|
||||
logging_obj = self._make_logging_obj(
|
||||
callback_duration_ms=5.5, llm_api_duration_ms=800.0
|
||||
)
|
||||
logging_obj = self._make_logging_obj(callback_duration_ms=5.5, llm_api_duration_ms=800.0)
|
||||
logging_obj._response_cost_calculator = MagicMock(return_value=0.001)
|
||||
logging_obj.litellm_call_id = "test-call-id"
|
||||
|
||||
|
|
@ -231,11 +232,24 @@ class TestResponseTimingMetrics:
|
|||
START = datetime.datetime(2025, 1, 1, 0, 0, 0)
|
||||
END = datetime.datetime(2025, 1, 1, 0, 0, 1)
|
||||
|
||||
def _make_logging_obj(self, llm_api_duration_ms=None, caching_details=None):
|
||||
def _make_logging_obj(
|
||||
self,
|
||||
llm_api_duration_ms: float | None = None,
|
||||
llm_api_timing_windows: object = None,
|
||||
caching_details: dict[str, object] | None = None,
|
||||
received_at: datetime.datetime | str | None = None,
|
||||
) -> MagicMock:
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {}
|
||||
if llm_api_duration_ms is not None:
|
||||
logging_obj.model_call_details["llm_api_duration_ms"] = llm_api_duration_ms
|
||||
if received_at is not None or llm_api_timing_windows is not None:
|
||||
metadata = {}
|
||||
if received_at is not None:
|
||||
metadata["litellm_received_at"] = received_at
|
||||
if llm_api_timing_windows is not None:
|
||||
metadata["llm_api_timing_windows"] = llm_api_timing_windows
|
||||
logging_obj.model_call_details["litellm_params"] = {"metadata": metadata}
|
||||
logging_obj.caching_details = caching_details
|
||||
return logging_obj
|
||||
|
||||
|
|
@ -246,6 +260,104 @@ class TestResponseTimingMetrics:
|
|||
"litellm_overhead_time_ms": 100.0,
|
||||
}
|
||||
|
||||
def test_window_starts_at_proxy_receive_when_stamped(self):
|
||||
received_at = self.START.astimezone(datetime.timezone.utc) - datetime.timedelta(seconds=3)
|
||||
logging_obj = self._make_logging_obj(llm_api_duration_ms=900.0, received_at=received_at)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["_response_ms"] == pytest.approx(4000.0)
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(3100.0)
|
||||
|
||||
def test_receive_anchored_window_subtracts_all_provider_attempts(self):
|
||||
logging_obj = self._make_logging_obj(
|
||||
llm_api_duration_ms=300.0,
|
||||
llm_api_timing_windows=(
|
||||
(self.START.timestamp(), self.START.timestamp() + 0.3),
|
||||
(self.START.timestamp() + 0.4, self.START.timestamp() + 0.8),
|
||||
),
|
||||
received_at=self.START,
|
||||
)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["_response_ms"] == pytest.approx(1000.0)
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(300.0)
|
||||
|
||||
def test_sdk_window_subtracts_current_provider_attempt(self):
|
||||
logging_obj = self._make_logging_obj(
|
||||
llm_api_duration_ms=300.0,
|
||||
llm_api_timing_windows=((self.START.timestamp(), self.START.timestamp() + 0.3),),
|
||||
)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["_response_ms"] == pytest.approx(1000.0)
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(700.0)
|
||||
|
||||
def test_receive_anchored_window_unions_nested_and_retry_windows(self):
|
||||
windows = (
|
||||
(self.START.timestamp(), self.START.timestamp() + 0.3),
|
||||
(self.START.timestamp(), self.START.timestamp() + 0.3),
|
||||
(self.START.timestamp() + 0.4, self.START.timestamp() + 0.7),
|
||||
)
|
||||
logging_obj = self._make_logging_obj(
|
||||
llm_api_duration_ms=300.0,
|
||||
llm_api_timing_windows=windows,
|
||||
received_at=self.START,
|
||||
)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(400.0)
|
||||
assert _union_duration_ms(windows, self.START.timestamp(), self.END.timestamp()) == pytest.approx(600.0)
|
||||
|
||||
def test_receive_anchored_window_ignores_seeded_windows_outside_window(self):
|
||||
windows = (
|
||||
(self.START.timestamp() - 10.0, self.START.timestamp() - 1.0),
|
||||
(self.END.timestamp() + 1.0, self.END.timestamp() + 2.0),
|
||||
)
|
||||
logging_obj = self._make_logging_obj(
|
||||
llm_api_duration_ms=300.0,
|
||||
llm_api_timing_windows=windows,
|
||||
received_at=self.START,
|
||||
)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(700.0)
|
||||
|
||||
def test_receive_anchored_window_falls_back_to_current_provider_attempt(self):
|
||||
logging_obj = self._make_logging_obj(
|
||||
llm_api_duration_ms=300.0,
|
||||
received_at=self.START,
|
||||
)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["_response_ms"] == pytest.approx(1000.0)
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(700.0)
|
||||
|
||||
def test_cache_hit_window_starts_at_proxy_receive_when_stamped(self):
|
||||
received_at = self.START.astimezone(datetime.timezone.utc) - datetime.timedelta(seconds=3)
|
||||
logging_obj = self._make_logging_obj(
|
||||
caching_details={"cache_hit": True, "cache_duration_ms": 250.0},
|
||||
received_at=received_at,
|
||||
)
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["_response_ms"] == pytest.approx(4000.0)
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(3750.0)
|
||||
|
||||
def test_non_datetime_proxy_receive_falls_back_to_start_time(self):
|
||||
logging_obj = self._make_logging_obj(llm_api_duration_ms=900.0, received_at="bad")
|
||||
|
||||
result = response_timing_metrics(self.START, self.END, logging_obj)
|
||||
|
||||
assert result["_response_ms"] == pytest.approx(1000.0)
|
||||
assert result["litellm_overhead_time_ms"] == pytest.approx(100.0)
|
||||
|
||||
def test_overhead_omitted_when_no_provider_or_cache_duration_recorded(self):
|
||||
logging_obj = self._make_logging_obj()
|
||||
assert response_timing_metrics(self.START, self.END, logging_obj) == {"_response_ms": 1000.0}
|
||||
|
|
@ -358,6 +470,28 @@ class TestDetailedTiming:
|
|||
assert hidden.get("timing_pre_processing_ms") == 20.0
|
||||
assert hidden.get("timing_post_processing_ms") == 10.0 # 530 - 20 - 500
|
||||
|
||||
def test_detailed_timing_pre_processing_uses_receive_anchor(self, monkeypatch):
|
||||
monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", True)
|
||||
|
||||
result = ModelResponse()
|
||||
received_at = datetime.datetime.now(datetime.timezone.utc)
|
||||
start = received_at + datetime.timedelta(milliseconds=200)
|
||||
api_call_start = start.replace(tzinfo=None)
|
||||
end = start + datetime.timedelta(milliseconds=530)
|
||||
logging_obj = self._make_logging_obj(
|
||||
llm_api_duration_ms=500.0,
|
||||
api_call_start_time=api_call_start,
|
||||
)
|
||||
logging_obj.model_call_details["litellm_params"] = {"metadata": {"litellm_received_at": received_at}}
|
||||
|
||||
metadata = ResponseMetadata(result)
|
||||
metadata.set_timing_metrics(start, end, logging_obj)
|
||||
metadata.apply()
|
||||
|
||||
hidden = result._hidden_params
|
||||
assert hidden.get("timing_pre_processing_ms") == pytest.approx(200.0)
|
||||
assert hidden.get("timing_post_processing_ms") == pytest.approx(30.0)
|
||||
|
||||
def test_detailed_timing_absent_when_disabled(self, monkeypatch):
|
||||
"""When LITELLM_DETAILED_TIMING is false, no detailed timing keys."""
|
||||
monkeypatch.setattr(response_metadata_mod, "LITELLM_DETAILED_TIMING", False)
|
||||
|
|
@ -377,9 +511,7 @@ class TestDetailedTiming:
|
|||
|
||||
def test_detailed_timing_headers_in_custom_headers(self, monkeypatch):
|
||||
"""When LITELLM_DETAILED_TIMING is true, headers flow to get_custom_headers."""
|
||||
monkeypatch.setattr(
|
||||
common_request_processing_mod, "LITELLM_DETAILED_TIMING", True
|
||||
)
|
||||
monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", True)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
|
||||
hidden_params = {
|
||||
|
|
@ -402,9 +534,7 @@ class TestDetailedTiming:
|
|||
|
||||
def test_detailed_timing_headers_absent_when_disabled(self, monkeypatch):
|
||||
"""When LITELLM_DETAILED_TIMING is false, no timing headers emitted."""
|
||||
monkeypatch.setattr(
|
||||
common_request_processing_mod, "LITELLM_DETAILED_TIMING", False
|
||||
)
|
||||
monkeypatch.setattr(common_request_processing_mod, "LITELLM_DETAILED_TIMING", False)
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="sk-test")
|
||||
hidden_params = {
|
||||
|
|
|
|||
|
|
@ -2,18 +2,42 @@
|
|||
Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import threading
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.litellm_core_utils import logging_utils
|
||||
from litellm.litellm_core_utils.logging_utils import (
|
||||
format_base64_size,
|
||||
_set_duration_in_model_call_details,
|
||||
_truncate_base64_in_string,
|
||||
format_base64_size,
|
||||
truncate_base64_in_messages,
|
||||
truncate_base64_in_messages_async,
|
||||
)
|
||||
|
||||
|
||||
class TestSetDurationInModelCallDetails:
|
||||
def test_records_provider_attempt_windows_in_shared_metadata(self):
|
||||
metadata = {"request_id": "test"}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model_call_details = {"litellm_params": {"metadata": metadata}}
|
||||
first_start = datetime.datetime(2025, 1, 1, 0, 0, 0)
|
||||
first_end = first_start + datetime.timedelta(milliseconds=300)
|
||||
second_start = datetime.datetime(2025, 1, 1, 0, 0, 1)
|
||||
second_end = second_start + datetime.timedelta(milliseconds=700)
|
||||
|
||||
_set_duration_in_model_call_details(logging_obj, first_start, first_end)
|
||||
_set_duration_in_model_call_details(logging_obj, second_start, second_end)
|
||||
|
||||
assert metadata["llm_api_timing_windows"] == (
|
||||
(first_start.timestamp(), first_end.timestamp()),
|
||||
(second_start.timestamp(), second_end.timestamp()),
|
||||
)
|
||||
assert logging_obj.model_call_details["llm_api_duration_ms"] == pytest.approx(700.0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# format_base64_size
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -157,10 +181,7 @@ class TestTruncateBase64InMessages:
|
|||
}
|
||||
]
|
||||
result = truncate_base64_in_messages(messages)
|
||||
assert (
|
||||
result[0]["content"][0]["image_url"]["url"]
|
||||
== f"data:image/png;base64,{short}"
|
||||
)
|
||||
assert result[0]["content"][0]["image_url"]["url"] == f"data:image/png;base64,{short}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -333,6 +333,36 @@ async def test_arrival_time_prefers_litellm_received_at_over_time_time():
|
|||
assert updated_data["proxy_server_request"]["arrival_time"] == received_at.timestamp()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_clears_client_supplied_timing_windows():
|
||||
request_mock = MagicMock(spec=Request)
|
||||
request_mock.url = MagicMock()
|
||||
request_mock.url.path = "/v1/chat/completions"
|
||||
request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions"
|
||||
request_mock.method = "POST"
|
||||
request_mock.query_params = {}
|
||||
request_mock.headers = {"Content-Type": "application/json"}
|
||||
request_mock.client = MagicMock()
|
||||
request_mock.client.host = "127.0.0.1"
|
||||
request_mock.state = SimpleNamespace(litellm_received_at=datetime.now(timezone.utc))
|
||||
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata={})
|
||||
|
||||
updated_data = await add_litellm_data_to_request(
|
||||
data={
|
||||
"model": "gpt-3.5-turbo",
|
||||
"metadata": {"llm_api_timing_windows": ((0.0, 1.0),)},
|
||||
},
|
||||
request=request_mock,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
proxy_config=MagicMock(),
|
||||
general_settings={},
|
||||
version="test-version",
|
||||
)
|
||||
|
||||
assert updated_data["metadata"]["llm_api_timing_windows"] == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_arrival_time_falls_back_to_time_time_without_litellm_received_at():
|
||||
"""Callers that never went through user_api_key_auth (no stamp on request.state)
|
||||
|
|
|
|||
|
|
@ -10,12 +10,24 @@ Verifies that:
|
|||
Regression tests for https://github.com/BerriAI/litellm/issues/21343
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import (
|
||||
_union_duration_ms,
|
||||
response_timing_metrics,
|
||||
)
|
||||
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
|
||||
from litellm.litellm_core_utils.rules import Rules
|
||||
from litellm.utils import function_setup
|
||||
|
||||
|
||||
def _make_rate_limit_error(message="Rate limited"):
|
||||
|
|
@ -274,3 +286,74 @@ async def test_not_found_error_in_retry_loop_raises_immediately():
|
|||
|
||||
# Only 2 calls: initial + first retry that hits non-retryable
|
||||
assert call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retry_attempts_accumulate_timing_in_shared_request_metadata():
|
||||
received_at: Final = datetime.datetime.now()
|
||||
metadata: dict[str, object] = {
|
||||
"model_group": "test-model",
|
||||
"litellm_received_at": received_at,
|
||||
}
|
||||
logging_obj_raw, _ = function_setup(
|
||||
"acompletion",
|
||||
Rules(),
|
||||
datetime.datetime.now(),
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
metadata=metadata,
|
||||
litellm_call_id="retry-timing-test",
|
||||
is_async_call=True,
|
||||
)
|
||||
assert isinstance(logging_obj_raw, Logging)
|
||||
logging_obj: Final[Logging] = logging_obj_raw
|
||||
attempt_numbers: list[int] = []
|
||||
metadata_ids: list[int] = []
|
||||
|
||||
@track_llm_api_timing()
|
||||
async def timed_attempt(*, logging_obj: Logging, **kwargs: object) -> str:
|
||||
del kwargs
|
||||
attempt_numbers.append(len(attempt_numbers) + 1)
|
||||
metadata_ids.append(id(logging_obj.model_call_details["litellm_params"]["metadata"]))
|
||||
await asyncio.sleep(0.01)
|
||||
if len(attempt_numbers) == 1:
|
||||
raise _make_rate_limit_error()
|
||||
return "success"
|
||||
|
||||
async def invoke(original_function: Callable[..., Awaitable[str]], *args: object, **kwargs: object) -> str:
|
||||
return await original_function(*args, **kwargs)
|
||||
|
||||
router = _create_router(num_retries=1)
|
||||
with (
|
||||
patch.object(router, "make_call", new=AsyncMock(side_effect=invoke)),
|
||||
patch.object(
|
||||
router,
|
||||
"_async_get_healthy_deployments",
|
||||
new=AsyncMock(return_value=(["d1"], ["d1"])),
|
||||
),
|
||||
patch.object(router, "_time_to_sleep_before_retry", return_value=0),
|
||||
):
|
||||
result = await router.async_function_with_retries(
|
||||
original_function=timed_attempt,
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "test"}],
|
||||
metadata=metadata,
|
||||
logging_obj=logging_obj,
|
||||
num_retries=1,
|
||||
)
|
||||
|
||||
request_metadata: Final = logging_obj.model_call_details["litellm_params"]["metadata"]
|
||||
windows: Final = request_metadata["llm_api_timing_windows"]
|
||||
end_time: Final = datetime.datetime.fromtimestamp(max(window[1] for window in windows))
|
||||
timing_metrics: Final = response_timing_metrics(received_at, end_time, logging_obj)
|
||||
assert result == "success"
|
||||
assert attempt_numbers == [1, 2]
|
||||
assert request_metadata is metadata
|
||||
assert metadata_ids == [id(metadata), id(metadata)]
|
||||
assert len(windows) == 2
|
||||
union_duration_ms: Final = _union_duration_ms(windows, received_at.timestamp(), end_time.timestamp())
|
||||
assert union_duration_ms is not None
|
||||
total_response_time_ms: Final = (end_time.timestamp() - received_at.timestamp()) * 1000
|
||||
assert timing_metrics["litellm_overhead_time_ms"] == pytest.approx(
|
||||
round(total_response_time_ms - union_duration_ms, 4)
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue