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 459f95972ce..45a477be89e 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -20,17 +20,10 @@ from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( MAX_REDIS_BUFFER_DEQUEUE_COUNT, - REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, - REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -61,6 +54,13 @@ 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_update_buffer import ( + SpendBuffer, + SpendUpdateBufferStore, + TransportSpendUpdateBufferStore, + spend_buffer_key, +) +from litellm.redis.transport.contract import RedisFailure, raise_redis_failure from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( RedisPipelineLpopOperation, @@ -107,6 +107,19 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( _ValueT = TypeVar("_ValueT") _DECODED_JSON: Final = TypeAdapter(object) +_DRAINED_BUFFERS: Final[tuple[SpendBuffer, ...]] = ( + "spend", + "daily_user", + "daily_team", + "daily_org", + "daily_end_user", + "daily_agent", + "window", +) + + +def _spend_buffers(redis_cache: RedisCache) -> SpendUpdateBufferStore: + return TransportSpendUpdateBufferStore(redis_cache.transport()) def _zero_for_non_finite_constant(constant: str) -> float: @@ -191,7 +204,7 @@ class RedisUpdateBuffer: async def _store_transactions_in_redis( self, transactions: dict[str, BaseDailySpendTransaction] | None, # mutable-ok: add_update requires dict - redis_key: str, + buffer: SpendBuffer, service_type: ServiceTypes, *, source_queue: DailySpendUpdateQueue, @@ -201,39 +214,34 @@ class RedisUpdateBuffer: Args: transactions: The transactions to store - redis_key: The Redis key to store under + buffer: The Redis spend buffer to store under service_type: The service type for event emission source_queue: The in-memory queue to restore transactions to if Redis push fails """ if transactions is None or len(transactions) == 0: return - list_of_transactions: Final = [safe_dumps(transactions)] if self.redis_cache is None: return - try: - current_redis_buffer_size: Final = await self.redis_cache.async_rpush( - key=redis_key, - values=list_of_transactions, - ) - except Exception as e: + current_redis_buffer_size: Final = await _spend_buffers(self.redis_cache).push(buffer, transactions) + if isinstance(current_redis_buffer_size, RedisFailure): verbose_proxy_logger.error( "Spend tracking - failed to push spend updates to Redis (redis_key=%s). " "Restoring them to the in-memory queue for retry on next tick. Error: %s", - redis_key, - str(e), + spend_buffer_key(buffer), + current_redis_buffer_size.message, ) await source_queue.add_update(transactions) - else: - verbose_proxy_logger.debug( - "Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s", - redis_key, - current_redis_buffer_size, - ) - await self._emit_new_item_added_to_redis_buffer_event( - queue_size=current_redis_buffer_size, - service=service_type, - ) + return + verbose_proxy_logger.debug( + "Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s", + spend_buffer_key(buffer), + current_redis_buffer_size, + ) + await self._emit_new_item_added_to_redis_buffer_event( + queue_size=current_redis_buffer_size, + service=service_type, + ) @with_service_target(SPEND_QUEUE_TARGET) async def store_in_memory_spend_updates_in_redis( @@ -323,65 +331,54 @@ class RedisUpdateBuffer: verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) # Build a list of rpush operations, skipping empty/None transaction sets - _queue_configs: Final[list[tuple[BufferedSpendTransactions | None, str, ServiceTypes]]] = [ + _queue_configs: Final[tuple[tuple[object, SpendBuffer, ServiceTypes], ...]] = ( ( db_spend_update_transactions, - REDIS_UPDATE_BUFFER_KEY, + "spend", ServiceTypes.REDIS_SPEND_UPDATE_QUEUE, ), ( daily_spend_update_transactions, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + "daily_user", ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, ), ( daily_team_spend_update_transactions, - REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + "daily_team", ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE, ), ( daily_org_spend_update_transactions, - REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + "daily_org", ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE, ), ( daily_end_user_spend_update_transactions, - REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + "daily_end_user", ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, ), ( daily_agent_spend_update_transactions, - REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + "daily_agent", ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, ), ( tuple(map(to_wire_payload, window_spend_update_transactions)), - REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + "window", ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE, ), - ] + ) - rpush_list: Final[list[RedisPipelineRpushOperation]] = [] - service_types: Final[list[ServiceTypes]] = [] - for transactions, redis_key, service_type in _queue_configs: - if transactions is None or len(transactions) == 0: - continue - rpush_list.append( - RedisPipelineRpushOperation( - key=redis_key, - values=[safe_dumps(transactions)], - ) - ) - service_types.append(service_type) - - if len(rpush_list) == 0: + pending: Final[tuple[tuple[SpendBuffer, object, ServiceTypes], ...]] = tuple( + (config[1], config[0], config[2]) for config in _queue_configs if config[0] + ) + if len(pending) == 0: return - try: - result_lengths: Final = await self.redis_cache.async_rpush_pipeline( - rpush_list=rpush_list, - ) - except Exception as e: + result_lengths: Final = await _spend_buffers(self.redis_cache).push_all( + tuple((buffer, transactions) for buffer, transactions, _ in pending) + ) + if isinstance(result_lengths, RedisFailure): # The in-memory queues were already drained above. If we let the # exception propagate without restoring, the aggregated spend is # permanently lost. Re-enqueue so the next scheduler tick retries. @@ -389,8 +386,8 @@ class RedisUpdateBuffer: "Spend tracking - failed to push aggregated spend updates to Redis. " "Restoring %d transaction sets to in-memory queues for retry on next tick. " "Error: %s", - len(rpush_list), - str(e), + len(pending), + result_lengths.message, ) await self._restore_spend_updates_to_in_memory_queues( db_spend_update_transactions=db_spend_update_transactions, @@ -411,12 +408,8 @@ class RedisUpdateBuffer: return # Emit gauge events for each queue - for i, queue_size in enumerate(result_lengths): - if i < len(service_types): - await self._emit_new_item_added_to_redis_buffer_event( - queue_size=queue_size, - service=service_types[i], - ) + for queue_size, (_, _, service_type) in zip(result_lengths, pending, strict=False): + await self._emit_new_item_added_to_redis_buffer_event(queue_size=queue_size, service=service_type) @staticmethod async def _restore_spend_updates_to_in_memory_queues( @@ -542,42 +535,42 @@ class RedisUpdateBuffer: if self.redis_cache is None: return - restore_configs: Final = ( - (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY), - (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY), - (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY), - (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY), - (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), - (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), - (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + restore_configs: Final[tuple[tuple[object, SpendBuffer], ...]] = ( + (db_spend_update_transactions, "spend"), + (daily_spend_update_transactions, "daily_user"), + (daily_team_spend_update_transactions, "daily_team"), + (daily_org_spend_update_transactions, "daily_org"), + (daily_end_user_spend_update_transactions, "daily_end_user"), + (daily_agent_spend_update_transactions, "daily_agent"), + (daily_tag_spend_update_transactions, "daily_tag"), ( None if window_spend_update_transactions is None else tuple(map(to_wire_payload, window_spend_update_transactions)), - REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + "window", ), ) - rpush_list: Final = tuple( - RedisPipelineRpushOperation(key=redis_key, values=(safe_dumps(transactions),)) - for transactions, redis_key in restore_configs - if transactions + batches: Final[tuple[tuple[SpendBuffer, object], ...]] = tuple( + (buffer, transactions) for transactions, buffer in restore_configs if transactions ) - if len(rpush_list) == 0: + if len(batches) == 0: return - try: - await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + restored: Final = await _spend_buffers(self.redis_cache).push_all(batches) + if not isinstance(restored, RedisFailure): verbose_proxy_logger.info( "Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.", - len(rpush_list), - ) - except RedisError as e: - verbose_proxy_logger.error( - "Spend tracking - failed to restore uncommitted transactions to Redis. " - "These spend updates are lost. Error: %s", - str(e), + len(batches), ) + return + if not isinstance(restored.exception, RedisError): + raise_redis_failure(restored) + verbose_proxy_logger.error( + "Spend tracking - failed to restore uncommitted transactions to Redis. " + "These spend updates are lost. Error: %s", + restored.message, + ) @with_service_target(SPEND_QUEUE_TARGET) async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None: @@ -734,22 +727,23 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final[str | list[str] | None] = await self.redis_cache.async_lpop( - key=REDIS_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_transactions: Final = await _spend_buffers(self.redis_cache).pop( + "spend", count=MAX_REDIS_BUFFER_DEQUEUE_COUNT ) + if isinstance(list_of_transactions, RedisFailure): + raise_redis_failure(list_of_transactions) if list_of_transactions is None: return None verbose_proxy_logger.info( "Spend tracking - popped %d spend update batches from Redis buffer (key=%s). " "These items are now removed from Redis and must be committed to DB.", - len(list_of_transactions) if isinstance(list_of_transactions, list) else 1, + len(list_of_transactions), REDIS_UPDATE_BUFFER_KEY, ) # Parse the list of transactions from JSON strings - parsed_transactions: Final = self._parse_list_of_transactions(list_of_transactions) + parsed_transactions: Final = self._parse_list_of_transactions(list(list_of_transactions)) # If there are no transactions, return None if len(parsed_transactions) == 0: @@ -787,44 +781,16 @@ class RedisUpdateBuffer: if self.redis_cache is None: return None, None, None, None, None, None, None - lpop_list: Final[list[RedisPipelineLpopOperation]] = [ - RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), - RedisPipelineLpopOperation( - key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), - RedisPipelineLpopOperation( - key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), - RedisPipelineLpopOperation( - key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), - RedisPipelineLpopOperation( - key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), - RedisPipelineLpopOperation( - key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), - RedisPipelineLpopOperation( - key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, - ), - ] - - raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) - - # Pad with None if pipeline returned fewer results than expected - while len(raw_results) < 7: - raw_results.append(None) + raw_results: Final = await _spend_buffers(self.redis_cache).pop_all( + _DRAINED_BUFFERS, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT + ) + if isinstance(raw_results, RedisFailure): + raise_redis_failure(raw_results) # Slot 0: DBSpendUpdateTransactions db_spend: DBSpendUpdateTransactions | None = None if raw_results[0] is not None: - parsed: Final = self._parse_list_of_transactions(raw_results[0]) + parsed: Final = self._parse_list_of_transactions(list(raw_results[0])) if len(parsed) > 0: db_spend = self._combine_list_of_transactions(parsed) @@ -869,7 +835,7 @@ class RedisUpdateBuffer: ) await self._store_transactions_in_redis( transactions=daily_tag_spend_update_transactions, - redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + buffer="daily_tag", service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, source_queue=daily_tag_spend_update_queue, ) @@ -877,17 +843,18 @@ class RedisUpdateBuffer: @with_service_target(SPEND_QUEUE_TARGET) async def _lpop_daily_spend_transactions( self, - redis_key: str, + buffer: SpendBuffer, ) -> list[dict[str, BaseDailySpendTransaction]] | None: """ Drains a daily spend buffer key and parses each popped item as JSON. """ if self.redis_cache is None: return None - list_of_transactions: Final[list[str] | None] = await self.redis_cache.async_lpop( - key=redis_key, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_transactions: Final = await _spend_buffers(self.redis_cache).pop( + buffer, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT ) + if isinstance(list_of_transactions, RedisFailure): + raise_redis_failure(list_of_transactions) if list_of_transactions is None: return None return [_load_buffered_transaction(transaction) for transaction in list_of_transactions] @@ -900,9 +867,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY - ) + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_user") if list_of_daily_spend_update_transactions is None: return None return cast( @@ -920,9 +885,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( - REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY - ) + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_team") if list_of_daily_spend_update_transactions is None: return None return cast( @@ -940,9 +903,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( - REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY - ) + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_org") if list_of_daily_spend_update_transactions is None: return None return cast( @@ -960,9 +921,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( - REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY - ) + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_end_user") if list_of_daily_spend_update_transactions is None: return None return cast( @@ -980,9 +939,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( - REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY - ) + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_agent") if list_of_daily_spend_update_transactions is None: return None return cast( @@ -1000,9 +957,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( - REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY - ) + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_tag") if list_of_daily_spend_update_transactions is None: return None return cast( diff --git a/litellm/redis/stores/spend_update_buffer.py b/litellm/redis/stores/spend_update_buffer.py new file mode 100644 index 00000000000..44e5b7adfa8 --- /dev/null +++ b/litellm/redis/stores/spend_update_buffer.py @@ -0,0 +1,89 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, Protocol, TypeAlias + +from litellm.constants import ( + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, +) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.redis.transport.contract import Popped, RedisFailure, RedisTransport + +SpendBuffer: TypeAlias = Literal[ + "spend", "daily_user", "daily_team", "daily_org", "daily_end_user", "daily_agent", "daily_tag", "window" +] + +_KEYS: Final = MappingProxyType[SpendBuffer, str]( + { + "spend": REDIS_UPDATE_BUFFER_KEY, + "daily_user": REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + "daily_team": REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + "daily_org": REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + "daily_end_user": REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + "daily_agent": REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + "daily_tag": REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + "window": REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + } +) + + +def spend_buffer_key(buffer: SpendBuffer) -> str: + return _KEYS[buffer] + + +class SpendUpdateBufferStore(Protocol): + async def push(self, buffer: SpendBuffer, transactions: object) -> int | RedisFailure: ... + + async def push_all(self, batches: Sequence[tuple[SpendBuffer, object]]) -> tuple[int, ...] | RedisFailure: ... + + async def pop(self, buffer: SpendBuffer, *, count: int) -> tuple[str, ...] | None | RedisFailure: ... + + async def pop_all( + self, buffers: Sequence[SpendBuffer], *, count: int + ) -> tuple[tuple[str, ...] | None, ...] | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportSpendUpdateBufferStore: + transport: RedisTransport + + async def push(self, buffer: SpendBuffer, transactions: object) -> int | RedisFailure: + return await self.transport.rpush(spend_buffer_key(buffer), safe_dumps(transactions)) + + async def push_all(self, batches: Sequence[tuple[SpendBuffer, object]]) -> tuple[int, ...] | RedisFailure: + batch: Final = self.transport.batch(name="spend_update_buffer") + pushed: Final = tuple( + batch.rpush(spend_buffer_key(buffer), safe_dumps(transactions)) for buffer, transactions in batches + ) + lengths: Final = tuple([await reply for reply in pushed]) + failure: Final = next((length for length in lengths if isinstance(length, RedisFailure)), None) + return failure if failure is not None else tuple(length for length in lengths if isinstance(length, int)) + + async def pop(self, buffer: SpendBuffer, *, count: int) -> tuple[str, ...] | None | RedisFailure: + return _decode(await self.transport.lpop(spend_buffer_key(buffer), count)) + + async def pop_all( + self, buffers: Sequence[SpendBuffer], *, count: int + ) -> tuple[tuple[str, ...] | None, ...] | RedisFailure: + batch: Final = self.transport.batch(name="spend_update_buffer") + popped: Final = tuple(batch.lpop(spend_buffer_key(buffer), count) for buffer in buffers) + decoded: Final = tuple([_decode(await reply) for reply in popped]) + failure: Final = next((entries for entries in decoded if isinstance(entries, RedisFailure)), None) + return failure if failure is not None else tuple(e for e in decoded if not isinstance(e, RedisFailure)) + + +def _decode(popped: Popped | RedisFailure) -> tuple[str, ...] | None | RedisFailure: + 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 spend update buffer entry: {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 65da5630f33..925ecf2b5d6 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 @@ -4,15 +4,28 @@ from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import pytest +import pytest_asyncio from litellm._logging import verbose_proxy_logger -from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY +from litellm.caching.redis_cache import RedisCache +from litellm.constants import ( + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.proxy_server import ProxyStartupEvent +from tests._support.fake_redis import fake_redis_cache def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction]: @@ -45,9 +58,9 @@ def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction @pytest.mark.asyncio async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure( - redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, fake_server ): - mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away")) + fake_server.connected = False queue: Final = DailySpendUpdateQueue() original_transactions: Final = _daily_tag_spend_update_transactions() await queue.add_update(original_transactions) @@ -59,31 +72,36 @@ async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure( @pytest.mark.asyncio async def test_daily_tag_spend_restored_after_rpush_failure_is_pushed_once_on_next_tick( - redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, fake_server, raw_redis ): - mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away")) + fake_server.connected = False queue: Final = DailySpendUpdateQueue() original_transactions: Final = _daily_tag_spend_update_transactions() await queue.add_update(original_transactions) await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) - mock_redis_cache.async_rpush = AsyncMock(return_value=1) + fake_server.connected = True await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue) - mock_redis_cache.async_rpush.assert_called_once() - redis_key: Final = mock_redis_cache.async_rpush.call_args.kwargs["key"] - values: Final = mock_redis_cache.async_rpush.call_args.kwargs["values"] - assert redis_key == REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY - assert json.loads(values[0]) == json.loads(safe_dumps(original_transactions)) + values: Final = await raw_redis.lrange(REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, 0, -1) + assert values == [safe_dumps(original_transactions).encode()] assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == {} @pytest.fixture -def mock_redis_cache(): - """Create a mock RedisCache instance""" - mock = AsyncMock() - return mock +def fake_server() -> fakeredis.FakeServer: + return fakeredis.FakeServer() + + +@pytest.fixture +def raw_redis(fake_server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=fake_server) + + +@pytest_asyncio.fixture(loop_scope="function") +async def mock_redis_cache(fake_server: fakeredis.FakeServer) -> RedisCache: + return fake_redis_cache(fake_server) @pytest.fixture @@ -93,13 +111,11 @@ def redis_update_buffer(mock_redis_cache): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): +async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache, raw_redis): """ Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once with the correct operations and skips empty queues. """ - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2]) - # Create mock queues - only 3 of 6 have data spend_update_queue = AsyncMock() spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( @@ -135,17 +151,23 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, daily_agent_spend_update_queue=daily_agent_queue, ) - # Should be called exactly once (pipeline) - mock_redis_cache.async_rpush_pipeline.assert_called_once() - - # Verify only 3 operations were included (empty ones skipped) - call_args = mock_redis_cache.async_rpush_pipeline.call_args - rpush_list = call_args.kwargs["rpush_list"] - assert len(rpush_list) == 3 + assert await raw_redis.lrange(REDIS_UPDATE_BUFFER_KEY, 0, -1) == [b'{"key_list_transactions": {"key1": 1.0}}'] + assert await raw_redis.lrange(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 0, -1) == [b'{"user_key1": {"spend": 1.0}}'] + assert await raw_redis.lrange(REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, 0, -1) == [b'{"team_key1": {"spend": 2.0}}'] + assert ( + await raw_redis.exists( + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + ) + == 0 + ) @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache): +async def test_store_in_memory_spend_updates_restores_on_rpush_failure( + redis_update_buffer, mock_redis_cache, fake_server +): """ If async_rpush_pipeline raises, the already-drained transactions must be put back into the in-memory queues so the next scheduler tick retries. @@ -159,7 +181,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_upd SpendUpdateQueue, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) + fake_server.connected = False spend_queue = SpendUpdateQueue() daily_user_queue = DailySpendUpdateQueue() @@ -245,7 +267,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat @pytest.mark.asyncio -async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache): +async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache, raw_redis): """ Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. @@ -278,17 +300,10 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff ] ) - mock_redis_cache.async_lpop_pipeline = AsyncMock( - return_value=[ - [db_spend_json], # slot 0: db spend updates - [daily_user_json], # slot 1: daily user - [daily_team_json], # slot 2: daily team - None, # slot 3: daily org (empty) - None, # slot 4: daily end-user (empty) - None, # slot 5: daily agent (empty) - [window_spend_json, window_spend_json], # slot 6: budget window spend - ] - ) + await raw_redis.rpush(REDIS_UPDATE_BUFFER_KEY, db_spend_json) + await raw_redis.rpush(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, daily_user_json) + await raw_redis.rpush(REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, daily_team_json) + await raw_redis.rpush(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, window_spend_json, window_spend_json) result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() @@ -328,17 +343,12 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff assert daily_end_user is None assert daily_agent is None - # Verify pipeline was called once with correct keys - mock_redis_cache.async_lpop_pipeline.assert_called_once() - from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY - - popped_keys = [op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]] - assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + assert await raw_redis.exists(REDIS_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY) == 0 @pytest.mark.asyncio async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_failure( - redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, fake_server, raw_redis ): from litellm.proxy._types import Litellm_EntityType from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -350,16 +360,14 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail member_key: Final = "organization_id::org-1::user_id::user-1" pod_json: Final = json.dumps({"org_member_list_transactions": {member_key: 0.25}}) - mock_redis_cache.async_lpop_pipeline = AsyncMock( - return_value=[[pod_json, pod_json], None, None, None, None, None, None] - ) + await raw_redis.rpush(REDIS_UPDATE_BUFFER_KEY, pod_json, pod_json) (db_spend, *_rest) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() assert db_spend is not None assert db_spend["org_member_list_transactions"] == {member_key: 0.5} - mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) + fake_server.connected = False spend_queue: Final = SpendUpdateQueue() await spend_queue.add_update( { @@ -383,10 +391,11 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail @pytest.mark.asyncio async def test_a_single_parked_batch_with_a_nan_spend_drains_as_zero( - redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, raw_redis ): - mock_redis_cache.async_lpop = AsyncMock( - return_value=json.dumps({"key_list_transactions": {"nan-key": float("nan"), "ok-key": 0.5}}, allow_nan=True) + await raw_redis.rpush( + REDIS_UPDATE_BUFFER_KEY, + json.dumps({"key_list_transactions": {"nan-key": float("nan"), "ok-key": 0.5}}, allow_nan=True), ) drained: Final = await redis_update_buffer.get_all_update_transactions_from_redis_buffer() @@ -416,7 +425,7 @@ def _parked_daily_batch(spend: float) -> str: @pytest.mark.asyncio async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_as_zero( - redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, raw_redis ): entity_batch: Final = json.dumps({"key_list_transactions": {"hashed-key": float("nan")}}, allow_nan=True) window_batches: Final = tuple( @@ -435,17 +444,11 @@ async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_a ) for spend in (float("nan"), 3.0) ) - mock_redis_cache.async_lpop_pipeline = AsyncMock( - return_value=[ - [entity_batch], - [_parked_daily_batch(float("nan")), _parked_daily_batch(0.0042)], - None, - None, - None, - None, - list(window_batches), - ] + await raw_redis.rpush(REDIS_UPDATE_BUFFER_KEY, entity_batch) + await raw_redis.rpush( + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, _parked_daily_batch(float("nan")), _parked_daily_batch(0.0042) ) + await raw_redis.rpush(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, *window_batches) ( db_spend, @@ -465,10 +468,10 @@ async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_a @pytest.mark.asyncio async def test_a_nan_spend_parked_by_an_older_version_drains_from_a_daily_key_as_zero( - redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, caplog: pytest.LogCaptureFixture + redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, caplog: pytest.LogCaptureFixture, raw_redis ): - mock_redis_cache.async_lpop = AsyncMock( - return_value=[_parked_daily_batch(float("inf")), _parked_daily_batch(0.0042)] + await raw_redis.rpush( + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, _parked_daily_batch(float("inf")), _parked_daily_batch(0.0042) ) with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): @@ -488,18 +491,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): @pytest.mark.asyncio -async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache): +async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache, raw_redis): """ restore_transactions_to_redis re-pushes only the transaction sets it was given, to their matching buffer keys, so uncommitted spend can be retried. """ - from litellm.constants import ( - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - REDIS_UPDATE_BUFFER_KEY, - ) - - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1, 1]) - db_spend = {"key_list_transactions": {"key1": 1.0}} daily_user = {"user_key1": {"spend": 1.0}} @@ -508,24 +504,19 @@ async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_b daily_spend_update_transactions=daily_user, ) - mock_redis_cache.async_rpush_pipeline.assert_called_once() - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - pushed_keys = {op["key"] for op in rpush_list} - assert pushed_keys == { - REDIS_UPDATE_BUFFER_KEY, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - } - # Payloads round-trip through the same JSON encoding used on the store path - payloads = {op["key"]: json.loads(op["values"][0]) for op in rpush_list} - assert payloads[REDIS_UPDATE_BUFFER_KEY] == db_spend - assert payloads[REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY] == daily_user + assert sorted(await raw_redis.keys("*")) == sorted( + (REDIS_UPDATE_BUFFER_KEY.encode(), REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY.encode()) + ) + assert await raw_redis.lrange(REDIS_UPDATE_BUFFER_KEY, 0, -1) == [safe_dumps(db_spend).encode()] + assert await raw_redis.lrange(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 0, -1) == [safe_dumps(daily_user).encode()] @pytest.mark.asyncio -async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): +async def test_restored_window_spend_transactions_drain_back_unchanged( + redis_update_buffer, mock_redis_cache, raw_redis +): """A window commit that fails after the destructive lpop must be re-pushed in the store path's encoding, so the next drain returns the same increments.""" - from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) @@ -540,16 +531,9 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), ), ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) - await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY] - - mock_redis_cache.async_lpop_pipeline = AsyncMock( - return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])] - ) + assert await raw_redis.keys("*") == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY.encode()] drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() assert drained[6] == window_transactions @@ -564,17 +548,18 @@ async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer @pytest.mark.asyncio -async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): +async def test_restore_transactions_to_redis_swallows_redis_error( + redis_update_buffer, mock_redis_cache, fake_server, raw_redis +): """A Redis failure during restore must not propagate to the caller's finally block.""" - from redis.exceptions import RedisError - - mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) + fake_server.connected = False await redis_update_buffer.restore_transactions_to_redis( db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, ) - mock_redis_cache.async_rpush_pipeline.assert_called_once() + fake_server.connected = True + assert await raw_redis.llen(REDIS_UPDATE_BUFFER_KEY) == 0 def test_validate_redis_transaction_buffer_raises_without_redis(): @@ -684,19 +669,18 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_update_buffer, mock_redis_cache): +async def test_store_in_memory_spend_updates_pushes_budget_window_spend( + redis_update_buffer, mock_redis_cache, raw_redis +): """The budget window queue has to ride the same rpush as the daily queues, otherwise multi-pod deployments never persist per-window spend.""" from datetime import datetime, timezone - from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendUpdateQueue, build_window_spend_transaction, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) - empty_queue = AsyncMock() empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) empty_daily_queue = AsyncMock() @@ -724,10 +708,9 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up window_spend_update_queue=window_queue, ) - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - assert len(rpush_list) == 1 - assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY - pushed = json.loads(rpush_list[0]["values"][0]) + assert await raw_redis.keys("*") == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY.encode()] + (pushed_bytes,) = await raw_redis.lrange(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 0, -1) + pushed = json.loads(pushed_bytes) assert pushed == [ { "entity_type": "key", @@ -742,7 +725,9 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up @pytest.mark.asyncio -async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_update_buffer, mock_redis_cache): +async def test_budget_window_payloads_keep_request_ids_for_older_workers( + redis_update_buffer, mock_redis_cache, raw_redis +): """A leader from before the field was dropped indexes request_ids while merging what it popped, and the pop is destructive, so a payload without the key would cost a rolling deploy those increments.""" @@ -753,7 +738,6 @@ async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_u build_window_spend_transaction, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) window_queue = WindowSpendUpdateQueue() await window_queue.add_update( build_window_spend_transaction( @@ -769,14 +753,14 @@ async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_u window_spend_update_transactions=await window_queue.flush_and_get_aggregated_window_spend_transactions(), ) - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - restored = json.loads(rpush_list[0]["values"][0]) + (restored_bytes,) = await raw_redis.lrange(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 0, -1) + restored = json.loads(restored_bytes) assert [payload["request_ids"] for payload in restored] == [[]] @pytest.mark.asyncio async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure( - redis_update_buffer, mock_redis_cache + redis_update_buffer, mock_redis_cache, fake_server ): """The window queue is drained before the rpush, so a Redis hiccup would silently drop per-window spend without the restore.""" @@ -787,7 +771,7 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu build_window_spend_transaction, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) + fake_server.connected = False empty_queue = AsyncMock() empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) @@ -859,7 +843,7 @@ async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap(): @pytest.mark.asyncio @pytest.mark.parametrize("spend", [float("nan"), float("inf")]) async def test_a_non_finite_daily_spend_buffered_in_redis_drains_as_zero(spend: float): - buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis()) + buffer: Final = RedisUpdateBuffer(redis_cache=fake_redis_cache(fakeredis.FakeServer())) queue: Final = DailySpendUpdateQueue() await queue.add_update(json.loads(_parked_daily_batch(spend))) await queue.add_update(json.loads(_parked_daily_batch(0.0042))) diff --git a/tests/unit/redis/stores/test_spend_update_buffer.py b/tests/unit/redis/stores/test_spend_update_buffer.py new file mode 100644 index 00000000000..af3f9dd6489 --- /dev/null +++ b/tests/unit/redis/stores/test_spend_update_buffer.py @@ -0,0 +1,63 @@ +from typing import Final + +import fakeredis +import pytest + +from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY +from litellm.redis.stores.spend_update_buffer import TransportSpendUpdateBufferStore +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache + + +def _store(server: fakeredis.FakeServer) -> TransportSpendUpdateBufferStore: + return TransportSpendUpdateBufferStore(fake_redis_cache(server, namespace="ns").transport()) + + +@pytest.mark.asyncio +async def test_batches_land_on_the_legacy_lists_as_safe_dumps_json(server: fakeredis.FakeServer) -> None: + store: Final = _store(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + + assert await store.push("daily_tag", {"k": {"spend": 1.5}}) == 1 + assert await store.push_all((("spend", {"key_list_transactions": {"a": 0.5}}), ("daily_tag", {"k": {}}))) == (1, 2) + + assert await raw.lrange(f"ns:{REDIS_UPDATE_BUFFER_KEY}", 0, -1) == [b'{"key_list_transactions": {"a": 0.5}}'] + assert await raw.lrange(f"ns:{REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY}", 0, -1) == [ + b'{"k": {"spend": 1.5}}', + b'{"k": {}}', + ] + assert await store.pop("daily_tag", count=1) == ('{"k": {"spend": 1.5}}',) + assert await store.pop_all(("spend", "daily_tag", "window"), count=10) == ( + ('{"key_list_transactions": {"a": 0.5}}',), + ('{"k": {}}',), + None, + ) + assert await raw.exists(f"ns:{REDIS_UPDATE_BUFFER_KEY}", f"ns:{REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY}") == 0 + + +@pytest.mark.asyncio +async def test_a_rejected_push_fails_the_batch_but_the_other_lists_keep_their_push( + server: fakeredis.FakeServer, +) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(f"ns:{REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY}", "not a list") + store: Final = _store(server) + + pushed: Final = await store.push_all((("spend", {"a": 1}), ("daily_tag", {"b": 2}))) + + assert isinstance(pushed, RedisFailure) and pushed.kind == "server" + assert await raw.lrange(f"ns:{REDIS_UPDATE_BUFFER_KEY}", 0, -1) == [b'{"a": 1}'] + + +@pytest.mark.asyncio +async def test_undecodable_entries_and_outages_are_failures(server: fakeredis.FakeServer) -> None: + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.rpush(f"ns:{REDIS_UPDATE_BUFFER_KEY}", b"\xff\xfe") + store: Final = _store(server) + + undecodable: Final = await store.pop("spend", count=10) + server.connected = False + outage: Final = await store.pop_all(("spend",), count=10) + + assert isinstance(undecodable, RedisFailure) and undecodable.kind == "decode" + assert isinstance(outage, RedisFailure) and outage.kind == "transport"