From 5249565fae996c121f2ec403341a8ba0e42e2df0 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Thu, 16 Apr 2026 22:27:55 +0200 Subject: [PATCH] fix(router): keep timedelta out of redis-backed latency cache --- litellm/caching/redis_cache.py | 18 ++- litellm/router_strategy/lowest_latency.py | 16 ++- .../caching/test_redis_cache_timedelta.py | 89 ++++++++++++ ..._lowest_latency_timedelta_serialization.py | 133 ++++++++++++++++++ 4 files changed, 251 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/caching/test_redis_cache_timedelta.py create mode 100644 tests/test_litellm/test_lowest_latency_timedelta_serialization.py diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 84a2887f527..6ff986478c9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -55,6 +55,20 @@ else: Span = Any +def _json_default(obj: Any) -> Any: + """`json.dumps(default=)` callback used by Redis writers. + + Maps `timedelta -> total_seconds()` (not `str(timedelta)`) so latency values + written by `LowestLatencyLoggingHandler` round-trip as `float` and survive + `_get_available_deployments`' `isinstance(_, float)` filter. Falls back to + `str(obj)` for any other non-JSON-native type, matching the `default=str` + pattern used elsewhere in litellm. + """ + if isinstance(obj, timedelta): + return obj.total_seconds() + return str(obj) + + def _get_call_stack_info(num_frames: int = 2) -> str: """ Get the function names from the previous 1-2 functions in the call stack. @@ -585,7 +599,7 @@ class RedisCache(BaseCache): raise Exception("Redis client cannot set cache. Attribute not found.") result = await _redis_client.set( name=key, - value=json.dumps(value), + value=json.dumps(value, default=_json_default), nx=nx, ex=ttl, ) @@ -643,7 +657,7 @@ class RedisCache(BaseCache): print_verbose( f"Set ASYNC Redis Cache PIPELINE: key: {cache_key}\nValue {cache_value}\nttl={ttl}" ) - json_cache_value = json.dumps(cache_value) + json_cache_value = json.dumps(cache_value, default=_json_default) # Set the value with a TTL if it's provided. _td: Optional[timedelta] = None if ttl is not None: diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 870b3f29d48..3f2c19da62c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -77,6 +77,11 @@ class LowestLatencyLoggingHandler(CustomLogger): precise_minute = f"{current_date}-{current_hour}-{current_minute}" response_ms = end_time - start_time + response_seconds: float = ( + response_ms.total_seconds() + if isinstance(response_ms, timedelta) + else float(response_ms) + ) time_to_first_token_response_time = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: @@ -85,7 +90,7 @@ class LowestLatencyLoggingHandler(CustomLogger): kwargs.get("completion_start_time", end_time) - start_time ) - final_value: Union[float, timedelta] = response_ms + final_value: float = response_seconds time_to_first_token: Optional[float] = None total_tokens = 0 @@ -302,6 +307,11 @@ class LowestLatencyLoggingHandler(CustomLogger): precise_minute = f"{current_date}-{current_hour}-{current_minute}" response_ms = end_time - start_time + response_seconds: float = ( + response_ms.total_seconds() + if isinstance(response_ms, timedelta) + else float(response_ms) + ) time_to_first_token_response_time = None if kwargs.get("stream", None) is not None and kwargs["stream"] is True: # only log ttft for streaming request @@ -309,7 +319,7 @@ class LowestLatencyLoggingHandler(CustomLogger): kwargs.get("completion_start_time", end_time) - start_time ) - final_value: Union[float, timedelta] = response_ms + final_value: float = response_seconds total_tokens = 0 time_to_first_token: Optional[float] = None @@ -331,7 +341,7 @@ class LowestLatencyLoggingHandler(CustomLogger): if final_value is not None: final_value = float(final_value) else: - final_value = response_ms + final_value = response_seconds if time_to_first_token_response_time is not None: if isinstance(time_to_first_token_response_time, timedelta): diff --git a/tests/test_litellm/caching/test_redis_cache_timedelta.py b/tests/test_litellm/caching/test_redis_cache_timedelta.py new file mode 100644 index 00000000000..be13768954f --- /dev/null +++ b/tests/test_litellm/caching/test_redis_cache_timedelta.py @@ -0,0 +1,89 @@ +""" +Tests that RedisCache writers tolerate `datetime.timedelta` values in the payload, +serializing them as `float` seconds (not stringified). Defence-in-depth against +the latency-cache regression class — keeps the cache numeric for the +`_get_available_deployments` read path. +""" + +import json +import os +import sys +from datetime import timedelta +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.caching.redis_cache import RedisCache + + +@pytest.fixture +def redis_no_ping(): + """Patch RedisCache initialization to prevent async ping tasks from being created.""" + with patch("asyncio.get_running_loop") as mock_get_loop: + mock_get_loop.side_effect = RuntimeError("No running event loop") + yield + + +@pytest.mark.asyncio +async def test_async_set_cache_serializes_timedelta_as_float( + monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + + payload = {"latency": [timedelta(seconds=1.5), 2.0]} + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_set_cache(key="embed-group_map", value=payload) + + mock_redis_instance.set.assert_called_once() + call_kwargs = mock_redis_instance.set.call_args.kwargs + assert call_kwargs["name"] == "embed-group_map" + decoded = json.loads(call_kwargs["value"]) + assert decoded == { + "latency": [1.5, 2.0] + }, f"timedelta did not round-trip as float: {decoded!r}" + + +@pytest.mark.asyncio +async def test_async_set_cache_pipeline_serializes_timedelta_as_float( + monkeypatch, redis_no_ping +): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + + mock_pipe = MagicMock() + mock_pipe.set = MagicMock() + mock_pipe.execute = AsyncMock(return_value=[True]) + + pipe_ctx = MagicMock() + pipe_ctx.__aenter__ = AsyncMock(return_value=mock_pipe) + pipe_ctx.__aexit__ = AsyncMock(return_value=None) + mock_redis_instance.pipeline = MagicMock(return_value=pipe_ctx) + + payload = {"latency": [timedelta(seconds=0.25)]} + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_set_cache_pipeline( + cache_list=[("embed-group_map", payload)] + ) + + mock_pipe.set.assert_called_once() + call_kwargs = mock_pipe.set.call_args.kwargs + decoded = json.loads(call_kwargs["value"]) + assert decoded == { + "latency": [0.25] + }, f"pipeline timedelta did not round-trip as float: {decoded!r}" diff --git a/tests/test_litellm/test_lowest_latency_timedelta_serialization.py b/tests/test_litellm/test_lowest_latency_timedelta_serialization.py new file mode 100644 index 00000000000..788631c13da --- /dev/null +++ b/tests/test_litellm/test_lowest_latency_timedelta_serialization.py @@ -0,0 +1,133 @@ +""" +Tests that LowestLatencyLoggingHandler stores latency values as floats +(not raw datetime.timedelta) in the cache, across all success-callback branches. + +Regression test for: timedelta values in the {model_group}_map cache breaking +RedisCache.async_set_cache JSON serialization. Earlier partial fix in PR #14040 +covered only the sync ModelResponse-with-usage branch. +""" + +import os +import sys +from datetime import datetime, timedelta + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +import litellm +from litellm.caching.caching import DualCache +from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler +from litellm.types.utils import EmbeddingResponse + + +def _make_kwargs(deployment_id: str = "1234", model_group: str = "embed-group") -> dict: + return { + "litellm_params": { + "metadata": { + "model_group": model_group, + "deployment": "openai/text-embedding-3-small", + }, + "model_info": {"id": deployment_id}, + } + } + + +def _latency_list(cache: DualCache, model_group: str, deployment_id: str) -> list: + cached = cache.get_cache(key=f"{model_group}_map") + assert cached is not None + assert deployment_id in cached + return cached[deployment_id].get("latency", []) + + +def test_log_success_event_embedding_response_stores_float(): + """EmbeddingResponse hits the non-ModelResponse branch — final_value must be float.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + kwargs = _make_kwargs() + response_obj = EmbeddingResponse( + model="text-embedding-3-small", + usage=litellm.Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), + ) + start_time = datetime.now() + end_time = start_time + timedelta(seconds=1.5) + + handler.log_success_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + latencies = _latency_list(cache, "embed-group", "1234") + assert latencies, "expected one latency entry" + assert all(isinstance(v, float) for v in latencies), f"non-float in {latencies!r}" + + +def test_log_success_event_model_response_no_usage_stores_float(): + """ModelResponse with usage=None hits a sub-branch that skips conversion — must still produce float.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + kwargs = _make_kwargs(model_group="chat-group") + response_obj = litellm.ModelResponse() + response_obj.usage = None # type: ignore[attr-defined] + start_time = datetime.now() + end_time = start_time + timedelta(seconds=0.75) + + handler.log_success_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + latencies = _latency_list(cache, "chat-group", "1234") + assert latencies and all(isinstance(v, float) for v in latencies) + + +@pytest.mark.asyncio +async def test_async_log_success_event_embedding_response_stores_float(): + """Async mirror of the embedding test.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + kwargs = _make_kwargs() + response_obj = EmbeddingResponse( + model="text-embedding-3-small", + usage=litellm.Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), + ) + start_time = datetime.now() + end_time = start_time + timedelta(seconds=2.0) + + await handler.async_log_success_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + latencies = _latency_list(cache, "embed-group", "1234") + assert latencies and all(isinstance(v, float) for v in latencies) + + +@pytest.mark.asyncio +async def test_async_log_success_event_zero_completion_tokens_stores_float(): + """Async ModelResponse with completion_tokens=0 takes the safe_divide_seconds → None + fallback at lowest_latency.py:344, which PR #14040 only fixed in the sync mirror.""" + cache = DualCache() + handler = LowestLatencyLoggingHandler(router_cache=cache) + kwargs = _make_kwargs(model_group="chat-group") + response_obj = litellm.ModelResponse( + usage=litellm.Usage(prompt_tokens=100, completion_tokens=0, total_tokens=100), + ) + start_time = datetime.now() + end_time = start_time + timedelta(seconds=0.5) + + await handler.async_log_success_event( + kwargs=kwargs, + response_obj=response_obj, + start_time=start_time, + end_time=end_time, + ) + + latencies = _latency_list(cache, "chat-group", "1234") + assert latencies and all(isinstance(v, float) for v in latencies)