diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 45a477be89e..2ba570efacb 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -7,7 +7,6 @@ This is to prevent deadlocks and improve reliability import asyncio import json from collections.abc import Mapping, Sequence -from datetime import datetime from functools import partial, reduce from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast @@ -20,7 +19,6 @@ from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( MAX_REDIS_BUFFER_DEQUEUE_COUNT, - REDIS_SPEND_LOGS_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, @@ -54,6 +52,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( to_wire_payload, ) from litellm.proxy.db.spend_log_batching import SpendLogRow +from litellm.redis.stores.spend_logs_parking import SpendLogsParkingStore, TransportSpendLogsParkingStore from litellm.redis.stores.spend_update_buffer import ( SpendBuffer, SpendUpdateBufferStore, @@ -122,6 +121,10 @@ def _spend_buffers(redis_cache: RedisCache) -> SpendUpdateBufferStore: return TransportSpendUpdateBufferStore(redis_cache.transport()) +def _parked_spend_logs(redis_cache: RedisCache) -> SpendLogsParkingStore: + return TransportSpendLogsParkingStore(redis_cache.transport()) + + def _zero_for_non_finite_constant(constant: str) -> float: verbose_proxy_logger.warning( "Redis spend buffer - reading a parked %s spend as 0, the spend it held is lost", constant @@ -132,14 +135,6 @@ def _zero_for_non_finite_constant(constant: str) -> float: _load_buffered_transaction: Final = partial(json.loads, parse_constant=_zero_for_non_finite_constant) -def _spend_log_json_default(value: object) -> str: - return value.isoformat() if isinstance(value, datetime) else str(value) - - -def _encode_spend_log_row(row: SpendLogRow) -> str: - return json.dumps(row, default=_spend_log_json_default) - - def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: decoded: Final = _DECODED_JSON.validate_python(_load_buffered_transaction(encoded)) return decoded if isinstance(decoded, dict) else None @@ -631,24 +626,19 @@ class RedisUpdateBuffer: """Park spend-log rows in Redis so they outlive this pod, dropping the oldest past ``max_rows``.""" if self.redis_cache is None or len(rows) == 0 or not self.should_commit_spend_updates_to_redis(): return False - try: - buffer_size: Final = await self.redis_cache.async_rpush_and_trim( - key=REDIS_SPEND_LOGS_BUFFER_KEY, - values=tuple(_encode_spend_log_row(row) for row in rows), - max_len=max_rows, - ) - overflow: Final = buffer_size - max_rows - if overflow > 0: - verbose_proxy_logger.error( - "Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs", - max_rows, - overflow, - ) - except Exception as e: # noqa: BLE001 # the caller falls back to the in-memory queue on any Redis fault + buffer_size: Final = await _parked_spend_logs(self.redis_cache).park(rows, max_rows=max_rows) + if isinstance(buffer_size, RedisFailure): verbose_proxy_logger.error( - "Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), str(e) + "Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), buffer_size.message ) return False + overflow: Final = buffer_size - max_rows + if overflow > 0: + verbose_proxy_logger.error( + "Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs", + max_rows, + overflow, + ) verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows)) return True @@ -657,14 +647,12 @@ class RedisUpdateBuffer: """Atomically take up to ``limit`` parked spend-log rows out of Redis.""" if self.redis_cache is None or not self.should_commit_spend_updates_to_redis(): return () - popped: Final[str | list[str] | None] = await self.redis_cache.async_lpop( - key=REDIS_SPEND_LOGS_BUFFER_KEY, - count=limit, - ) + popped: Final = await _parked_spend_logs(self.redis_cache).take(limit) + if isinstance(popped, RedisFailure): + raise_redis_failure(popped) if popped is None: return () - encoded_rows: Final = tuple(popped) if isinstance(popped, list) else (popped,) - decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in encoded_rows) + decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in popped) return tuple(row for row in decoded_rows if row is not None) @staticmethod diff --git a/litellm/redis/stores/spend_logs_parking.py b/litellm/redis/stores/spend_logs_parking.py new file mode 100644 index 00000000000..d62f3321964 --- /dev/null +++ b/litellm/redis/stores/spend_logs_parking.py @@ -0,0 +1,45 @@ +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Protocol + +from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +def _json_default(value: object) -> str: + return value.isoformat() if isinstance(value, datetime) else str(value) + + +def encode_spend_log_row(row: Mapping[str, object]) -> str: + return json.dumps(row, default=_json_default) + + +class SpendLogsParkingStore(Protocol): + async def park(self, rows: Sequence[Mapping[str, object]], *, max_rows: int) -> int | RedisFailure: ... + + async def take(self, limit: int) -> tuple[str, ...] | None | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportSpendLogsParkingStore: + transport: RedisTransport + + async def park(self, rows: Sequence[Mapping[str, object]], *, max_rows: int) -> int | RedisFailure: + batch: Final = self.transport.batch(transaction=True, name="spend_logs_parking") + pushed: Final = batch.rpush(REDIS_SPEND_LOGS_BUFFER_KEY, *(encode_spend_log_row(row) for row in rows)) + trimmed: Final = batch.ltrim(REDIS_SPEND_LOGS_BUFFER_KEY, -max_rows, -1) + length: Final = await pushed + trim: Final = await trimmed + return trim if isinstance(trim, RedisFailure) and not isinstance(length, RedisFailure) else length + + async def take(self, limit: int) -> tuple[str, ...] | None | RedisFailure: + popped: Final = await self.transport.lpop(REDIS_SPEND_LOGS_BUFFER_KEY, limit) + if popped is None or isinstance(popped, RedisFailure): + return popped + entries: Final = (popped,) if isinstance(popped, bytes) else popped + try: + return tuple(entry.decode("utf-8") for entry in entries) + except UnicodeDecodeError as error: + return RedisFailure(kind="decode", message=f"Invalid parked spend log row: {error.reason}") diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 925ecf2b5d6..cd1289e7189 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -804,32 +804,9 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu assert [payload["entity_id"] for payload in restored] == ["team-1"] -class _ListRedis: - def __init__(self) -> None: - self.rows: list[str] = [] - - async def async_rpush(self, key: str, values: list[str]) -> int: - self.rows.extend(values) - return len(self.rows) - - async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: - self.rows.extend(values) - pushed_len = len(self.rows) - del self.rows[:-max_len] - return pushed_len - - async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> list[str] | None: - if not self.rows: - return None - popped = self.rows[:count] - del self.rows[:count] - return popped - - @pytest.mark.asyncio async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap(): - redis = _ListRedis() - buffer = RedisUpdateBuffer(redis_cache=redis) + buffer = RedisUpdateBuffer(redis_cache=fake_redis_cache(fakeredis.FakeServer())) buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) assert await buffer.store_spend_logs_in_redis([{"request_id": "old"}, {"request_id": "mid"}], max_rows=2) is True @@ -858,7 +835,7 @@ async def test_a_non_finite_daily_spend_buffered_in_redis_drains_as_zero(spend: @pytest.mark.asyncio async def test_a_spend_log_row_parked_with_a_nan_spend_is_read_back_with_zero_spend(): - buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis()) + buffer: Final = RedisUpdateBuffer(redis_cache=fake_redis_cache(fakeredis.FakeServer())) buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) assert await buffer.store_spend_logs_in_redis([{"request_id": "nan-row", "spend": float("nan")}]) is True @@ -876,12 +853,12 @@ async def test_store_spend_logs_in_redis_reports_failure_without_redis(): @pytest.mark.asyncio async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_enabled(): - redis = _ListRedis() - buffer = RedisUpdateBuffer(redis_cache=redis) + server: Final = fakeredis.FakeServer() + buffer = RedisUpdateBuffer(redis_cache=fake_redis_cache(server)) buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=False) assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False - assert redis.rows == [] + assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == [] @pytest.mark.asyncio diff --git a/tests/unit/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py index aefb95516af..16f262b5363 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -13,15 +13,23 @@ needing a generated Prisma client or a real database. from __future__ import annotations import asyncio +import json import sys import threading +from collections.abc import Callable, Iterator from dataclasses import dataclass, field from email.message import EmailMessage from pathlib import Path -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any from unittest.mock import AsyncMock, MagicMock +import fakeredis import pytest +import pytest_asyncio + +from litellm.caching.redis_cache import RedisCache +from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY +from tests._support.fake_redis import fake_redis_cache sys.path.insert(0, str(Path(__file__).resolve().parents[5])) @@ -51,16 +59,13 @@ VOLATILE_KEYS = frozenset( def normalize(data: Any, volatile: frozenset = VOLATILE_KEYS) -> Any: """Recursively replace values for volatile keys with ''.""" if isinstance(data, dict): - return { - k: ("" if k in volatile else normalize(v, volatile)) - for k, v in data.items() - } + return {k: ("" if k in volatile else normalize(v, volatile)) for k, v in data.items()} if isinstance(data, list): return [normalize(v, volatile) for v in data] return data -_PRISMA_TABLES: List[str] = [ +_PRISMA_TABLES: list[str] = [ "litellm_verificationtoken", "litellm_teamtable", "litellm_usertable", @@ -156,7 +161,7 @@ def mock_dual_cache() -> MagicMock: ``cache._store`` directly in a test for assertion convenience. """ cache = MagicMock(name="MockDualCache") - cache._store: Dict[str, Any] = {} + cache._store: dict[str, Any] = {} def _sync_get(key: str, **_: Any) -> Any: return cache._store.get(key) @@ -191,6 +196,7 @@ def patched_prisma_import(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock directly and restore in teardown. """ import prisma as _prisma_pkg + import litellm.proxy.utils as _utils_mod fake_prisma = MagicMock(name="FakePrisma") @@ -261,7 +267,7 @@ class FakeClock: """ now: float = 0.0 - sleep_calls: List[float] = field(default_factory=list) + sleep_calls: list[float] = field(default_factory=list) def advance(self, seconds: float) -> None: self.now += seconds @@ -290,7 +296,7 @@ def fake_clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock: @pytest.fixture -def make_spend_log_row() -> Callable[..., Dict[str, Any]]: +def make_spend_log_row() -> Callable[..., dict[str, Any]]: """Factory for fake LiteLLM_SpendLogs rows.""" def _make( @@ -298,7 +304,7 @@ def make_spend_log_row() -> Callable[..., Dict[str, Any]]: spend: float = 0.01, model: str = "gpt-4o-mini", **overrides: Any, - ) -> Dict[str, Any]: + ) -> dict[str, Any]: row = { "request_id": request_id, "spend": spend, @@ -316,38 +322,21 @@ def make_spend_log_row() -> Callable[..., Dict[str, Any]]: return _make +@dataclass(frozen=True) class FakeRedisList: - def __init__(self) -> None: - self.items: dict[str, list[str]] = {} - self.down = False + server: fakeredis.FakeServer + cache: RedisCache + raw: fakeredis.FakeAsyncRedis - def _check_up(self) -> None: - if self.down: - raise ConnectionError("redis unreachable") - - async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: - self._check_up() - stored = self.items.setdefault(key, []) - stored.extend(str(v) for v in values) - pushed_len = len(stored) - del stored[:-max_len] - return pushed_len - - async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> str | list[str] | None: - self._check_up() - stored = self.items.get(key, []) - if not stored: - return None - if count is None: - return stored.pop(0) - popped = stored[:count] - del stored[:count] - return popped + async def parked_request_ids(self) -> list[str]: + rows = await self.raw.lrange(REDIS_SPEND_LOGS_BUFFER_KEY, 0, -1) + return [json.loads(row)["request_id"] for row in rows] -@pytest.fixture -def fake_redis() -> FakeRedisList: - return FakeRedisList() +@pytest_asyncio.fixture(loop_scope="function") +async def fake_redis() -> FakeRedisList: + server = fakeredis.FakeServer() + return FakeRedisList(server=server, cache=fake_redis_cache(server), raw=fakeredis.FakeAsyncRedis(server=server)) @pytest.fixture @@ -358,7 +347,7 @@ def proxy_logging_with_redis(fake_redis: FakeRedisList) -> MagicMock: proxy_logging.failure_handler = AsyncMock() proxy_logging.db_spend_update_writer = MagicMock() proxy_logging.db_spend_update_writer.db_update_spend_transaction_handler = AsyncMock() - buffer = RedisUpdateBuffer(redis_cache=fake_redis) + buffer = RedisUpdateBuffer(redis_cache=fake_redis.cache) buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) proxy_logging.db_spend_update_writer.redis_update_buffer = buffer return proxy_logging @@ -366,12 +355,12 @@ def proxy_logging_with_redis(fake_redis: FakeRedisList) -> MagicMock: @dataclass class _SentMessage: - from_addr: Optional[str] + from_addr: str | None to_addrs: Any - subject: Optional[str] - body: Optional[str] + subject: str | None + body: str | None starttls_called: bool - login_args: Optional[tuple] + login_args: tuple | None thread_ident: int @@ -379,9 +368,9 @@ class _SentMessage: class InMemorySMTP: """Captures outbound SMTP traffic for ``send_email`` tests.""" - sent: List[_SentMessage] = field(default_factory=list) - raise_on_send: Optional[Exception] = None - connection_kwargs: List[Dict[str, Any]] = field(default_factory=list) + sent: list[_SentMessage] = field(default_factory=list) + raise_on_send: Exception | None = None + connection_kwargs: list[dict[str, Any]] = field(default_factory=list) def server_factory(self) -> Callable[..., Any]: outer = self @@ -389,9 +378,9 @@ class InMemorySMTP: class _Conn: def __init__(self) -> None: self._starttls_called = False - self._login_args: Optional[tuple] = None + self._login_args: tuple | None = None - def __enter__(self) -> "_Conn": + def __enter__(self) -> _Conn: return self def __exit__(self, *exc: Any) -> None: @@ -406,7 +395,7 @@ class InMemorySMTP: def send_message( self, msg: EmailMessage, - from_addr: Optional[str] = None, + from_addr: str | None = None, to_addrs: Any = None, ) -> None: if outer.raise_on_send is not None: diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py index fc7ae531c34..ce65370d49d 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -11,7 +11,6 @@ Symbols pinned here: from __future__ import annotations import asyncio -import json from collections.abc import Callable from contextlib import suppress from datetime import datetime, timezone @@ -20,12 +19,11 @@ from unittest.mock import AsyncMock, MagicMock import pytest -from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, + drain_spend_logs_queue, monitor_spend_logs_queue, raise_failed_update_spend_exception, - drain_spend_logs_queue, recover_parked_spend_logs, update_daily_tag_spend, update_spend, @@ -719,10 +717,6 @@ def _table_gone_error() -> Exception: ) -def _parked_request_ids(fake_redis: Any) -> list[str]: - return [json.loads(row)["request_id"] for row in fake_redis.items.get(REDIS_SPEND_LOGS_BUFFER_KEY, [])] - - @pytest.mark.asyncio async def test_drain_spend_logs_queue_parks_unwritable_rows_in_redis_on_shutdown( mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any @@ -743,7 +737,7 @@ async def test_drain_spend_logs_queue_parks_unwritable_rows_in_redis_on_shutdown ) assert mock_prisma_client.spend_log_transactions == [] - assert sorted(_parked_request_ids(fake_redis)) == ["r1", "r2"] + assert sorted(await fake_redis.parked_request_ids()) == ["r1", "r2"] @pytest.mark.asyncio @@ -779,7 +773,7 @@ async def test_drain_spend_logs_queue_waits_for_an_in_flight_write_before_parkin proxy_logging_obj=proxy_logging_with_redis, ) - assert _parked_request_ids(fake_redis) == ["in-flight"] + assert await fake_redis.parked_request_ids() == ["in-flight"] assert mock_prisma_client.spend_log_transactions == [] await release with suppress(Exception): @@ -813,7 +807,7 @@ async def test_drain_spend_logs_queue_parks_rows_left_after_max_passes( ) assert mock_prisma_client.spend_log_transactions == [] - assert _parked_request_ids(fake_redis) == ["late"] + assert await fake_redis.parked_request_ids() == ["late"] @pytest.mark.asyncio @@ -822,7 +816,7 @@ async def test_drain_spend_logs_queue_keeps_rows_in_memory_when_redis_is_down( ) -> None: from prisma.errors import TableNotFoundError - fake_redis.down = True + fake_redis.server.connected = False mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")] mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_table_gone_error()) @@ -834,7 +828,8 @@ async def test_drain_spend_logs_queue_keeps_rows_in_memory_when_redis_is_down( ) assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r1"] - assert fake_redis.items == {} + fake_redis.server.connected = True + assert await fake_redis.raw.keys("*") == [] @pytest.mark.asyncio @@ -863,7 +858,7 @@ async def test_update_spend_writes_rows_parked_in_redis_by_a_previous_pod( written = mock_prisma_client.db.litellm_spendlogs.create_many.await_args.kwargs["data"] assert [row["request_id"] for row in written] == ["parked"] - assert _parked_request_ids(fake_redis) == [] + assert await fake_redis.parked_request_ids() == [] assert mock_prisma_client.spend_log_transactions == [] @@ -879,14 +874,14 @@ async def test_recover_parked_spend_logs_re_parks_rows_when_the_enqueue_is_cance recover_parked_spend_logs(prisma_client=mock_prisma_client, proxy_logging_obj=proxy_logging_with_redis) ) await asyncio.sleep(0.01) - assert _parked_request_ids(fake_redis) == [] + assert await fake_redis.parked_request_ids() == [] recovery.cancel() with pytest.raises(asyncio.CancelledError): await recovery mock_prisma_client._spend_log_transactions_lock.release() - assert _parked_request_ids(fake_redis) == ["parked"] + assert await fake_redis.parked_request_ids() == ["parked"] assert mock_prisma_client.spend_log_transactions == [] diff --git a/tests/unit/redis/stores/test_spend_logs_parking.py b/tests/unit/redis/stores/test_spend_logs_parking.py new file mode 100644 index 00000000000..92a36e0b003 --- /dev/null +++ b/tests/unit/redis/stores/test_spend_logs_parking.py @@ -0,0 +1,53 @@ +from datetime import datetime, timezone +from typing import Final + +import fakeredis +import pytest + +from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY +from litellm.redis.stores.spend_logs_parking import TransportSpendLogsParkingStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + +KEY: Final = f"ns:{REDIS_SPEND_LOGS_BUFFER_KEY}" + + +def _store(server: fakeredis.FakeServer) -> TransportSpendLogsParkingStore: + return TransportSpendLogsParkingStore(fake_redis_cache(server, namespace="ns").transport()) + + +@pytest.mark.asyncio +async def test_rows_are_parked_as_json_with_iso_datetimes_and_capped_to_the_newest( + server: fakeredis.FakeServer, +) -> None: + store: Final = _store(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + started: Final = datetime(2026, 9, 19, 20, 0, 5, tzinfo=timezone.utc) + + assert await store.park(({"request_id": "old"}, {"request_id": "mid", "startTime": started}), max_rows=2) == 2 + assert await store.park(({"request_id": "new"},), max_rows=2) == 3 + + assert await raw.lrange(KEY, 0, -1) == [ + b'{"request_id": "mid", "startTime": "2026-09-19T20:00:05+00:00"}', + b'{"request_id": "new"}', + ] + assert await store.take(1) == ('{"request_id": "mid", "startTime": "2026-09-19T20:00:05+00:00"}',) + assert await store.take(10) == ('{"request_id": "new"}',) + assert await store.take(10) is None + + +@pytest.mark.asyncio +async def test_rejections_undecodable_rows_and_outages_are_failures(server: fakeredis.FakeServer) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + store: Final = _store(server) + await raw.set(KEY, "not a list") + rejected: Final = await store.park(({"request_id": "a"},), max_rows=10) + await raw.delete(KEY) + await raw.rpush(KEY, b"\xff") + undecodable: Final = await store.take(10) + server.connected = False + outage: Final = await store.take(10) + + assert isinstance(rejected, RedisFailure) and rejected.kind == "server" + assert isinstance(undecodable, RedisFailure) and undecodable.kind == "decode" + assert isinstance(outage, RedisFailure) and outage.kind == "transport"