From 6b55fa8901e3954e9d1262fce14735ef84db2e1e Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 18 Oct 2026 09:45:43 +0000 Subject: [PATCH] fix(proxy): keep Redis spend buffer items recoverable after a drain reply timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 208 ++++++++++- litellm/caching/redis_cluster_cache.py | 4 + litellm/constants.py | 1 + litellm/integrations/otel/model/spans.py | 2 + .../redis_update_buffer.py | 200 +++++++---- tests/unit/caching/test_redis_cache.py | 118 ++++++- .../unit/caching/test_redis_cluster_cache.py | 28 +- .../test_redis_update_buffer.py | 329 ++++++++++++++++-- .../proxy/utils/prisma_and_spend/conftest.py | 25 +- .../prisma_and_spend/test_spend_functions.py | 3 + tests/unit/test_redis.py | 52 +++ 11 files changed, 864 insertions(+), 106 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 461ac5d947f..3e5be6e24f8 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,7 +20,7 @@ import threading import time import weakref from collections import Counter -from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -72,6 +72,43 @@ else: Span = Any +class _AsyncRedisListPipeline(Protocol): + async def __aenter__(self) -> "_AsyncRedisListPipeline": ... + + async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> None: ... + + def lmove(self, first_list: str, second_list: str, src: str, dest: str) -> "_AsyncRedisListPipeline": ... + + def lpop(self, name: str) -> "_AsyncRedisListPipeline": ... + + def expire(self, name: str, time: int) -> "_AsyncRedisListPipeline": ... + + async def execute(self) -> list[object]: ... + + +_LPOP_TO_PROCESSING_SCRIPT: Final = """ +local moved = {} +for i = 1, tonumber(ARGV[1]) do + local v = redis.call('LPOP', KEYS[1]) + if not v then break end + redis.call('RPUSH', KEYS[2], v) + moved[#moved + 1] = v +end +redis.call('EXPIRE', KEYS[2], tonumber(ARGV[2])) +return moved +""" + + +class _RedisServiceLogger(Protocol): + def async_service_success_hook( + self, service: ServiceTypes, call_type: str, duration: float, caller: str + ) -> Coroutine[object, object, None]: ... + + def async_service_failure_hook( + self, service: ServiceTypes, duration: float, error: Exception, call_type: str, caller: str + ) -> Coroutine[object, object, None]: ... + + class _AsyncRedisCommands(Protocol): """Async redis commands this cache issues. @@ -91,6 +128,8 @@ class _AsyncRedisCommands(Protocol): def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... + def lrange(self, name: str, start: int, end: int) -> Awaitable[list[bytes | str]]: ... + def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... @@ -135,6 +174,7 @@ _INCREMENT_WITH_FLOOR_LUA: Final = ( ) _LUA_COUNT: Final = TypeAdapter(int) +_LUA_LIST_RESULT: Final = TypeAdapter(list[bytes | str | None]) _OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) @@ -144,6 +184,14 @@ def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, . ) +def _decode_redis_list(values: Sequence[bytes | str | None]) -> list[str]: + return [ + value.decode("utf-8", errors="replace") if isinstance(value, bytes) else value + for value in values + if value is not None + ] + + def _is_generic_caller_frame(frame: FrameType) -> bool: module: Final = frame.f_globals.get("__name__") return module in _GENERIC_CALLER_MODULES or frame.f_code.co_name in _GENERIC_CALLER_FRAME_NAMES @@ -680,7 +728,7 @@ class RedisCache(BaseCache): self.redis_flush_size: int = 100 else: self.redis_flush_size = redis_flush_size - self.redis_version: str = "Unknown" + self.redis_version: str | int | float = "Unknown" try: if not coroutine_checker.is_async_callable(self.redis_client): self.redis_version = self.redis_client.info()["redis_version"] @@ -2406,6 +2454,162 @@ class RedisCache(BaseCache): decoded_results.append(None) return decoded_results + @functools.cached_property + def _lpop_to_processing_script(self) -> Callable[..., Awaitable[object]]: + return self.async_register_script(_LPOP_TO_PROCESSING_SCRIPT) + + async def _async_move_list_to_processing( + self, + redis_client: _AsyncRedisCommands, + namespaced_key: str, + processing_key: str, + count: int, + ttl_seconds: int, + ) -> list[str]: + major_version: Final = self._parse_redis_major_version() + version_parts: Final = str(self.redis_version).strip().split(".") + minor_part: Final = ( + version_parts[1].split("-", maxsplit=1)[0] if major_version == 6 and len(version_parts) > 1 else "" + ) + minor_version: Final = int(minor_part) if minor_part.isdigit() else 0 + use_lmove: Final = major_version > 6 or (major_version == 6 and minor_version >= 2) + if not use_lmove: + script_results: Final = await self._lpop_to_processing_script( + keys=(namespaced_key, processing_key), + args=(count, ttl_seconds), + ) + return _decode_redis_list(_LUA_LIST_RESULT.validate_python(script_results)) + + list_pipeline: Final = cast( # cast-ok: RedisCluster stubs omit list commands + _AsyncRedisListPipeline, + redis_client.pipeline(transaction=False), + ) + async with list_pipeline as pipe: + for _ in range(count): + pipe.lmove(namespaced_key, processing_key, src="LEFT", dest="RIGHT") + pipe.expire(processing_key, ttl_seconds) + raw_results: Final = await pipe.execute() + + filtered_results: Final = tuple( + result for result in raw_results[:-1] if result is None or isinstance(result, (bytes, str)) + ) + return _decode_redis_list(filtered_results) + + def _processing_list_key(self, namespaced_key: str, worker_id: str) -> str: + return f"{namespaced_key}:processing:{worker_id}" + + @_redis_circuit_breaker_guard + async def async_drain_list_to_processing(self, key: str, worker_id: str, count: int, ttl_seconds: int) -> list[str]: + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key) + processing_key: Final = self._processing_list_key( + namespaced_key=namespaced_key, + worker_id=worker_id, + ) + start_time: Final = time.time() + service_logger: Final = cast( # cast-ok: logger has the required hooks + _RedisServiceLogger, + self.service_logger_obj, + ) + + try: + pending_items: Final = _decode_redis_list(await _redis_client.lrange(processing_key, 0, -1)) + drained_items: Final = ( + pending_items + if pending_items + else await self._async_move_list_to_processing( + redis_client=_redis_client, + namespaced_key=namespaced_key, + processing_key=processing_key, + count=count, + ttl_seconds=ttl_seconds, + ) + ) + end_time: Final = time.time() + _duration: Final = end_time - start_time + asyncio.create_task( + service_logger.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type="async_drain_list_to_processing", + caller=_get_call_stack_info(), + ) + ) + except Exception as e: + failure_end_time: Final = time.time() + failure_duration: Final = failure_end_time - start_time + asyncio.create_task( + service_logger.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=failure_duration, + error=e, + call_type="async_drain_list_to_processing", + caller=_get_call_stack_info(), + ) + ) + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_drain_list_to_processing() - Got exception from REDIS", + e, + ) + raise + else: + return drained_items + + @_redis_circuit_breaker_guard + async def async_ack_processing_list(self, key: str, worker_id: str) -> None: + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key) + processing_key: Final = self._processing_list_key( + namespaced_key=namespaced_key, + worker_id=worker_id, + ) + start_time: Final = time.time() + service_logger: Final = cast( # cast-ok: logger has the required hooks + _RedisServiceLogger, + self.service_logger_obj, + ) + + try: + for attempt in range(3): + try: + await _redis_client.delete(processing_key) + break + except Exception: + if attempt == 2: + raise + + end_time: Final = time.time() + _duration: Final = end_time - start_time + asyncio.create_task( + service_logger.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type="async_ack_processing_list", + caller=_get_call_stack_info(), + ) + ) + except Exception as e: + failure_end_time: Final = time.time() + failure_duration: Final = failure_end_time - start_time + asyncio.create_task( + service_logger.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=failure_duration, + error=e, + call_type="async_ack_processing_list", + caller=_get_call_stack_info(), + ) + ) + log_redis_failure( + verbose_logger, + logging.ERROR, + "LiteLLM Redis Caching: async_ack_processing_list() - Got exception from REDIS", + e, + ) + raise + @_redis_circuit_breaker_guard async def async_lpop_pipeline( self, diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 6fb84df331e..594ff479526 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -126,6 +126,10 @@ class RedisClusterCache(RedisCache): return call() return self._connection_gate.run(call, self._get_async_client_cache_key()) + def _processing_list_key(self, namespaced_key: str, worker_id: str) -> str: + namespace_prefix: Final = f"{self.namespace}:" if self.namespace else "" + return f"{namespace_prefix}{{{namespaced_key}}}:processing:{worker_id}" + def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ Overrides `_run_redis_mget_operation` in redis_cache.py diff --git a/litellm/constants.py b/litellm/constants.py index da95b97036a..91e55c28805 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -453,6 +453,7 @@ REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_up REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100) +REDIS_BUFFER_PROCESSING_LIST_TTL_SECONDS: Final = 86400 REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer" REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000 REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT: Final = 1000 diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 76fc969cfc0..8b9e3e3addf 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -218,6 +218,8 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = { "async_delete_cache": "delete", "async_rpush": "rpush", "async_lpop": "lpop", + "async_drain_list_to_processing": "lmove", + "async_ack_processing_list": "delete", "async_scan_iter": "scan", "async_lpop_pipeline": "lpop", "async_rpush_pipeline": "rpush", 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 898aee493df..09d770de906 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -6,9 +6,12 @@ This is to prevent deadlocks and improve reliability import asyncio import json +import uuid from collections.abc import Mapping, Sequence +from contextlib import AsyncExitStack from datetime import datetime from functools import reduce +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from redis.exceptions import RedisError @@ -18,6 +21,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache from litellm.constants import ( MAX_REDIS_BUFFER_DEQUEUE_COUNT, + REDIS_BUFFER_PROCESSING_LIST_TTL_SECONDS, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, @@ -55,7 +59,6 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( from litellm.proxy.db.spend_log_batching import SpendLogRow from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( - RedisPipelineLpopOperation, RedisPipelineRpushOperation, ) from litellm.types.services import ServiceTypes @@ -145,6 +148,26 @@ class RedisUpdateBuffer: redis_cache: RedisCache | None = None, ): self.redis_cache = redis_cache + self._drain_worker_id: Final = str(uuid.uuid4()) + self._pipeline_drain_keys: Final = ( + REDIS_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ) + self._drain_locks: Final = MappingProxyType( + { + key: asyncio.Lock() + for key in ( + *self._pipeline_drain_keys, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + REDIS_SPEND_LOGS_BUFFER_KEY, + ) + } + ) @staticmethod def should_commit_spend_updates_to_redis() -> bool: @@ -503,11 +526,11 @@ class RedisUpdateBuffer: window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None, ) -> None: """ - Re-push transactions that were popped from Redis but not committed to the DB. + Re-push transactions drained from Redis but not committed to the DB. - The leader drains the buffers with a destructive ``lpop`` before committing to - the database. When a commit fails after its retries are exhausted, the popped - transactions must be pushed back so a later scheduler tick can retry them; + The leader removes drained transactions from Redis before committing to the + database. When a commit fails after its retries are exhausted, the transactions + must be pushed back so a later scheduler tick can retry them; otherwise the aggregated spend is lost permanently. The re-pushed payloads use the same JSON encoding as the store path, so the next drain parses them normally. """ @@ -581,19 +604,40 @@ class RedisUpdateBuffer: verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows)) return True + async def _drain_to_processing( + self, + key: str, + count: int, + ) -> list[str] | None: # mutable-ok: preserve the Redis list drain contract + if self.redis_cache is None: + return None + drained_items: Final = await self.redis_cache.async_drain_list_to_processing( + key=key, + worker_id=self._drain_worker_id, + count=count, + ttl_seconds=REDIS_BUFFER_PROCESSING_LIST_TTL_SECONDS, + ) + return drained_items if drained_items else None + + async def _drain(self, key: str, count: int) -> list[str] | None: + if self.redis_cache is None: + return None + async with self._drain_locks[key]: + drained_items: Final = await self._drain_to_processing(key=key, count=count) + if drained_items is None: + return None + await self.redis_cache.async_ack_processing_list(key=key, worker_id=self._drain_worker_id) + return drained_items + @with_service_target(SPEND_QUEUE_TARGET) async def get_spend_logs_from_redis_buffer(self, limit: int) -> tuple[dict[str, object], ...]: - """Atomically take up to ``limit`` parked spend-log rows out of Redis.""" + """Drain up to ``limit`` parked spend-log rows from 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, - ) - if popped is None: + drained_items: Final = await self._drain(key=REDIS_SPEND_LOGS_BUFFER_KEY, count=limit) + if drained_items 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 drained_items) return tuple(row for row in decoded_rows if row is not None) @staticmethod @@ -656,7 +700,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final[str | list[str] | None] = await self.redis_cache.async_lpop( + list_of_transactions: Final = await self._drain( key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) @@ -666,7 +710,7 @@ class RedisUpdateBuffer: 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, ) @@ -695,7 +739,7 @@ class RedisUpdateBuffer: tuple[WindowSpendTransaction, ...] | None, ]: """ - Drains the main 7 Redis buffer queues in a single pipeline round-trip. + Concurrently drains the main 7 Redis buffer queues. Returns a 7-tuple of parsed results in this order: 0: DBSpendUpdateTransactions @@ -709,63 +753,77 @@ 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) - - # 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]) - if len(parsed) > 0: - db_spend = self._combine_list_of_transactions(parsed) - - # Slots 1-5: daily spend categories - daily_results: Final[list[dict[str, BaseDailySpendTransaction] | None]] = [] - for slot in range(1, 6): - slot_result = raw_results[slot] - if slot_result is None: - daily_results.append(None) - else: - list_of_daily = [json.loads(t) for t in slot_result] - aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily) - daily_results.append(aggregated) + async with AsyncExitStack() as drain_lock_stack: + for key in self._pipeline_drain_keys: + await drain_lock_stack.enter_async_context(self._drain_locks[key]) + drain_results_with_errors: Final = tuple( + await asyncio.gather( + *( + self._drain_to_processing(key=key, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT) + for key in self._pipeline_drain_keys + ), + return_exceptions=True, + ) + ) + drain_errors: Final = tuple( + result for result in drain_results_with_errors if isinstance(result, BaseException) + ) + if drain_errors: + raise drain_errors[0] + raw_results: Final = tuple( + result for result in drain_results_with_errors if not isinstance(result, BaseException) + ) + acknowledgement_indexes: Final = tuple( + index for index, items in enumerate(raw_results) if items is not None + ) + acknowledgement_results: Final = tuple( + await asyncio.gather( + *( + self.redis_cache.async_ack_processing_list( + key=self._pipeline_drain_keys[index], + worker_id=self._drain_worker_id, + ) + for index in acknowledgement_indexes + ), + return_exceptions=True, + ) + ) + failed_acknowledgements: Final = tuple( + (index, result) + for index, result in zip(acknowledgement_indexes, acknowledgement_results) + if isinstance(result, BaseException) + ) + for index, error in failed_acknowledgements: + verbose_proxy_logger.warning( + "Spend tracking - failed to acknowledge Redis buffer drain (key=%s). Items will be retried: %s", + self._pipeline_drain_keys[index], + error, + ) + failed_acknowledgement_indexes: Final = frozenset(index for index, _ in failed_acknowledgements) + results_ready_for_parse: Final = tuple( + None if index in failed_acknowledgement_indexes else items for index, items in enumerate(raw_results) + ) + parsed_transactions: Final = ( + self._parse_list_of_transactions(results_ready_for_parse[0]) + if results_ready_for_parse[0] is not None + else () + ) + db_spend: Final = self._combine_list_of_transactions(parsed_transactions) if parsed_transactions else None + daily_results: Final = tuple( + None + if transactions is None + else DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + [json.loads(transaction) for transaction in transactions] + ) + for transactions in results_ready_for_parse[1:6] + ) + window_transactions: Final = results_ready_for_parse[6] window_spend: Final = ( WindowSpendUpdateQueue.get_aggregated_window_spend_transactions( - tuple(json.loads(transaction) for transaction in raw_results[6]) + tuple(json.loads(transaction) for transaction in window_transactions) ) - if raw_results[6] is not None + if window_transactions is not None else None ) @@ -801,11 +859,11 @@ class RedisUpdateBuffer: redis_key: str, ) -> list[dict[str, BaseDailySpendTransaction]] | None: """ - Drains a daily spend buffer key and parses each popped item as JSON. + Drains a daily spend buffer key and parses each item as JSON. """ if self.redis_cache is None: return None - list_of_transactions: Final[list[str] | None] = await self.redis_cache.async_lpop( + list_of_transactions: Final = await self._drain( key=redis_key, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..ed7e9101d1f 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -8,6 +8,8 @@ from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeServer +from fakeredis.aioredis import FakeRedis import litellm from litellm._service_logger import ServiceLogging @@ -23,6 +25,22 @@ def redis_no_ping(): yield +def _cache_with_fakeredis( + monkeypatch: pytest.MonkeyPatch, namespace: str, redis_version: str | float = "7.0.0" +) -> tuple[RedisCache, FakeRedis]: + redis_client: Final = FakeRedis(server=FakeServer()) + sync_redis_client: Final = MagicMock() + sync_redis_client.info.return_value = {"redis_version": redis_version} + with ( + patch("asyncio.get_running_loop", side_effect=RuntimeError("No running event loop")), + patch("litellm._redis.get_redis_client", return_value=sync_redis_client), + patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()), + ): + redis_cache: Final = RedisCache(namespace=namespace) + monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client) + return redis_cache, redis_client + + @pytest.fixture def sync_batch_redis_cache(redis_no_ping): with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point @@ -59,6 +77,101 @@ def test_check_and_fix_namespace_prefixes_keys_sharing_the_namespace_prefix( assert redis_cache.check_and_fix_namespace(key=key) == expected +@pytest.mark.asyncio +async def test_async_drain_list_to_processing_respects_count_pending_order_and_namespace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch, namespace="tenant") + await redis_client.rpush("tenant:buffer", "first", "second", "third") + + drained: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-1", + count=2, + ttl_seconds=60, + ) + + processing_key: Final = "tenant:buffer:processing:worker-1" + assert drained == ["first", "second"] + assert await redis_client.lrange(processing_key, 0, -1) == [b"first", b"second"] + await redis_client.rpush("tenant:buffer", "newer") + + pending: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-1", + count=2, + ttl_seconds=60, + ) + + assert pending == ["first", "second"] + assert await redis_client.lrange("tenant:buffer", 0, -1) == [b"third", b"newer"] + redis_keys: Final = await redis_client.keys("*") + assert set(redis_keys) == {b"tenant:buffer", processing_key.encode()} + assert all(redis_key.startswith(b"tenant:") for redis_key in redis_keys) + + +@pytest.mark.asyncio +async def test_async_drain_list_to_processing_replaces_invalid_utf8_and_acknowledges( + monkeypatch: pytest.MonkeyPatch, +) -> None: + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch, namespace="tenant") + await redis_client.rpush("tenant:buffer", b"\xff", b"valid") + + drained: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-utf8", + count=2, + ttl_seconds=60, + ) + assert drained == ["\ufffd", "valid"] + + await redis_cache.async_ack_processing_list(key="buffer", worker_id="worker-utf8") + assert ( + await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-utf8", + count=2, + ttl_seconds=60, + ) + == [] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("redis_version", "expected_items", "expected_processing", "expected_buffer"), + [ + ("6.2.0", ("first", "second"), (b"first", b"second"), (b"third",)), + (7.2, ("first", "second"), (b"first", b"second"), (b"third",)), + (6.2, ("first", "second"), (b"first", b"second"), (b"third",)), + ], +) +async def test_async_drain_list_to_processing_preserves_pending_items_for_redis_version( + monkeypatch: pytest.MonkeyPatch, + redis_version: str | float, + expected_items: tuple[str, ...], + expected_processing: tuple[bytes, ...], + expected_buffer: tuple[bytes, ...], +) -> None: + redis_cache, redis_client = _cache_with_fakeredis( + monkeypatch, + namespace="tenant", + redis_version=redis_version, + ) + await redis_client.rpush("tenant:buffer", "first", "second", "third") + + drained: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-1", + count=2, + ttl_seconds=60, + ) + + assert drained == list(expected_items) + assert tuple(await redis_client.lrange("tenant:buffer:processing:worker-1", 0, -1)) == expected_processing + assert tuple(await redis_client.lrange("tenant:buffer", 0, -1)) == expected_buffer + + @pytest.mark.parametrize("namespace", [None, "litellm"]) @pytest.mark.asyncio async def test_async_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping): @@ -1033,8 +1146,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) gc.collect() ServiceLogging() - assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( - redis_breaker_state_counts()["closed"] + assert ( + REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) + == (redis_breaker_state_counts()["closed"]) ) diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 85422819926..d6d46f93e94 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -5,13 +5,14 @@ import ssl import weakref from collections.abc import Awaitable, Callable from importlib import import_module -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from redis.asyncio import Redis, RedisCluster from redis.asyncio.cluster import ClusterNode from redis.asyncio.connection import SSLConnection +from redis.crc import key_slot from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache @@ -44,6 +45,31 @@ def test_redis_cluster_batch_get(mock_init_redis_cluster): assert not mock_redis.mget.called +@pytest.mark.parametrize( + ("namespace", "expected_source_key", "expected_processing_prefix"), + [ + (None, "buffer", "{buffer}:"), + ("tenant", "tenant:buffer", "tenant:{tenant:buffer}:"), + ], +) +def test_redis_cluster_processing_list_key_preserves_namespace_and_hash_slot( + namespace: str | None, + expected_source_key: str, + expected_processing_prefix: str, +) -> None: + with patch("litellm._redis.init_redis_cluster", return_value=MagicMock()): + cache: Final = RedisClusterCache( + startup_nodes=[{"host": "localhost", "port": 6379}], + namespace=namespace, + ) + + source_namespaced_key: Final = cache.check_and_fix_namespace("buffer") + processing_key: Final = cache._processing_list_key(source_namespaced_key, "worker-1") + assert source_namespaced_key == expected_source_key + assert processing_key.startswith(expected_processing_prefix) + assert key_slot(processing_key.encode()) == key_slot(source_namespaced_key.encode()) + + @pytest.mark.asyncio @patch("litellm._redis.init_redis_cluster") async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): 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 5ecc8784384..e1cfd8be035 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 @@ -1,11 +1,17 @@ +import asyncio import json from datetime import datetime, timezone +from itertools import count from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeServer +from fakeredis.aioredis import FakeRedis +from redis.asyncio.client import Pipeline +from redis.exceptions import TimeoutError as RedisTimeoutError - +from litellm.caching.redis_cache import RedisCache from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.proxy_server import ProxyStartupEvent @@ -23,6 +29,65 @@ def redis_update_buffer(mock_redis_cache): return RedisUpdateBuffer(redis_cache=mock_redis_cache) +def _cache_with_fakeredis(monkeypatch: pytest.MonkeyPatch) -> tuple[RedisCache, FakeRedis]: + redis_client: Final = FakeRedis(server=FakeServer()) + sync_redis_client: Final = MagicMock() + sync_redis_client.info.return_value = {"redis_version": "7.0.0"} + with ( + patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")), + patch("litellm._redis.get_redis_client", return_value=sync_redis_client), + patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()), + ): + redis_cache: Final = RedisCache(namespace="litellm") + monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client) + return redis_cache, redis_client + + +def _timeout_first_pipeline_execute(monkeypatch: pytest.MonkeyPatch, redis_client: FakeRedis) -> None: + original_pipeline: Final = redis_client.pipeline + pipeline_calls: Final = count() + + def pipeline_with_timeout(transaction: bool = True) -> Pipeline: + pipeline: Final = original_pipeline(transaction=transaction) + if next(pipeline_calls) == 0: + original_execute: Final = pipeline.execute + + async def execute_then_timeout() -> list[object]: + await original_execute() + raise RedisTimeoutError("simulated response timeout") + + monkeypatch.setattr(pipeline, "execute", execute_then_timeout) + return pipeline + + monkeypatch.setattr(redis_client, "pipeline", pipeline_with_timeout) + + +def _hold_first_pipeline_until_release( + monkeypatch: pytest.MonkeyPatch, + redis_client: FakeRedis, + move_executed: asyncio.Event, + release_move: asyncio.Event, +) -> None: + original_pipeline: Final = redis_client.pipeline + pipeline_calls: Final = count() + + def pipeline_with_timeout(transaction: bool = True) -> Pipeline: + pipeline: Final = original_pipeline(transaction=transaction) + if next(pipeline_calls) == 0: + original_execute: Final = pipeline.execute + + async def execute_after_release() -> list[object]: + results: Final = await original_execute() + move_executed.set() + await release_move.wait() + return results + + monkeypatch.setattr(pipeline, "execute", execute_after_release) + return pipeline + + monkeypatch.setattr(redis_client, "pipeline", pipeline_with_timeout) + + @pytest.mark.asyncio async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): """ @@ -179,7 +244,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache): """ Verify get_all_transactions_from_redis_buffer_pipeline correctly parses - and aggregates results from async_lpop_pipeline. + and aggregates results from each buffer drain. """ # Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories, # slot 6 = budget window spend @@ -209,17 +274,18 @@ 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 - ] + mock_redis_cache.async_drain_list_to_processing = AsyncMock( + side_effect=( + [db_spend_json], + [daily_user_json], + [daily_team_json], + [], + [], + [], + [window_spend_json, window_spend_json], + ) ) + mock_redis_cache.async_ack_processing_list = AsyncMock() result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() @@ -259,12 +325,29 @@ 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 + 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_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + 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 + drained_keys: Final = tuple( + call.kwargs["key"] for call in mock_redis_cache.async_drain_list_to_processing.await_args_list + ) + assert drained_keys == ( + REDIS_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, + ) + assert mock_redis_cache.async_ack_processing_list.await_count == 4 @pytest.mark.asyncio @@ -281,8 +364,8 @@ 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] + mock_redis_cache.async_drain_list_to_processing = AsyncMock( + side_effect=([pod_json, pod_json], [], [], [], [], [], []) ) (db_spend, *_rest) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() @@ -356,7 +439,7 @@ async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_b @pytest.mark.asyncio async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): - """A window commit that fails after the destructive lpop must be re-pushed + """A window commit that fails after a drain 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 ( @@ -380,8 +463,8 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd 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"])] + mock_redis_cache.async_drain_list_to_processing = AsyncMock( + side_effect=([], [], [], [], [], [], list(rpush_list[0]["values"])) ) drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() @@ -410,6 +493,190 @@ async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_b mock_redis_cache.async_rpush_pipeline.assert_called_once() +@pytest.mark.asyncio +async def test_pipeline_drain_recovers_spend_after_move_reply_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy._types import Litellm_EntityType + from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue + from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue + + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch) + buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache) + spend_queue: Final = SpendUpdateQueue() + await spend_queue.add_update({"entity_type": Litellm_EntityType.KEY, "entity_id": "key-abc", "response_cost": 0.25}) + await spend_queue.add_update({"entity_type": Litellm_EntityType.KEY, "entity_id": "key-abc", "response_cost": 0.5}) + await spend_queue.add_update( + {"entity_type": Litellm_EntityType.TEAM, "entity_id": "team-xyz", "response_cost": 1.25} + ) + await buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=spend_queue, + daily_spend_update_queue=DailySpendUpdateQueue(), + daily_team_spend_update_queue=DailySpendUpdateQueue(), + daily_org_spend_update_queue=DailySpendUpdateQueue(), + daily_end_user_spend_update_queue=DailySpendUpdateQueue(), + daily_agent_spend_update_queue=DailySpendUpdateQueue(), + ) + + _timeout_first_pipeline_execute(monkeypatch, redis_client) + with pytest.raises(RedisTimeoutError, match="simulated response timeout"): + await buffer.get_all_transactions_from_redis_buffer_pipeline() + + recovered: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline() + assert recovered[0] is not None + assert recovered[0]["key_list_transactions"] == {"key-abc": 0.75} + assert recovered[0]["team_list_transactions"] == {"team-xyz": 1.25} + assert recovered[1:] == (None, None, None, None, None, None) + assert await buffer.get_all_transactions_from_redis_buffer_pipeline() == (None, None, None, None, None, None, None) + + +@pytest.mark.asyncio +async def test_pipeline_drain_keeps_all_buffers_pending_after_partial_move_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.constants import ( + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + ) + + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch) + buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache) + db_spend: Final = json.dumps({"key_list_transactions": {"key-partial": 1.25}}) + daily_spend: Final = json.dumps({"user-partial": {"spend": 2.5, "api_requests": 1}}) + await redis_cache.async_rpush(key=REDIS_UPDATE_BUFFER_KEY, values=(db_spend,)) + await redis_cache.async_rpush(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, values=(daily_spend,)) + + _timeout_first_pipeline_execute(monkeypatch, redis_client) + with pytest.raises(RedisTimeoutError, match="simulated response timeout"): + await buffer.get_all_transactions_from_redis_buffer_pipeline() + + recovered: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline() + assert recovered[0] is not None + assert recovered[0]["key_list_transactions"] == {"key-partial": 1.25} + assert recovered[1] == {"user-partial": {"spend": 2.5, "api_requests": 1}} + assert recovered[2:] == (None, None, None, None, None) + assert await buffer.get_all_transactions_from_redis_buffer_pipeline() == (None, None, None, None, None, None, None) + + +@pytest.mark.asyncio +async def test_pipeline_drain_returns_only_keys_with_successful_acknowledgements( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.constants import ( + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY, + ) + + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch) + buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache) + db_spend: Final = json.dumps({"key_list_transactions": {"key-ack": 1.5}}) + daily_spend: Final = json.dumps({"user-ack": {"spend": 2.75, "api_requests": 1}}) + await redis_cache.async_rpush(key=REDIS_UPDATE_BUFFER_KEY, values=(db_spend,)) + await redis_cache.async_rpush(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, values=(daily_spend,)) + + original_delete: Final = redis_client.delete + + async def fail_main_buffer_delete(key: str) -> int: + if REDIS_UPDATE_BUFFER_KEY in key: + raise RedisTimeoutError("DEL failed") + return await original_delete(key) + + monkeypatch.setattr(redis_client, "delete", fail_main_buffer_delete) + first_result: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline() + assert first_result[0] is None + assert first_result[1] == {"user-ack": {"spend": 2.75, "api_requests": 1}} + assert first_result[2:] == (None, None, None, None, None) + + monkeypatch.setattr(redis_client, "delete", original_delete) + second_result: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline() + assert second_result[0] is not None + assert second_result[0]["key_list_transactions"] == {"key-ack": 1.5} + assert second_result[1:] == (None, None, None, None, None, None) + assert await buffer.get_all_transactions_from_redis_buffer_pipeline() == (None, None, None, None, None, None, None) + + +@pytest.mark.asyncio +async def test_spend_logs_are_recovered_once_after_move_reply_timeout(monkeypatch: pytest.MonkeyPatch) -> None: + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch) + buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache) + buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) + rows: Final = ({"request_id": "req-a"}, {"request_id": "req-b"}) + assert await buffer.store_spend_logs_in_redis(rows) is True + + _timeout_first_pipeline_execute(monkeypatch, redis_client) + with pytest.raises(RedisTimeoutError, match="simulated response timeout"): + await buffer.get_spend_logs_from_redis_buffer(limit=10) + + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == rows + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () + + +@pytest.mark.asyncio +async def test_overlapping_spend_log_drains_serialize_and_recover_rows_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch) + buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache) + buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True) + rows: Final = ({"request_id": "req-a"}, {"request_id": "req-b"}) + assert await buffer.store_spend_logs_in_redis(rows) is True + + move_executed: Final = asyncio.Event() + release_move: Final = asyncio.Event() + _hold_first_pipeline_until_release( + monkeypatch, + redis_client, + move_executed, + release_move, + ) + + first_call: Final = asyncio.create_task(buffer.get_spend_logs_from_redis_buffer(limit=10)) + await move_executed.wait() + + second_started: Final = asyncio.Event() + + async def start_second_drain() -> tuple[dict[str, object], ...]: + second_started.set() + return await buffer.get_spend_logs_from_redis_buffer(limit=10) + + second_call: Final = asyncio.create_task(start_second_drain()) + await second_started.wait() + + release_move.set() + first_result: Final = await first_call + second_result: Final = await second_call + request_ids: Final = tuple(str(row["request_id"]) for row in first_result + second_result) + assert sorted(request_ids) == ["req-a", "req-b"] + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () + + +@pytest.mark.asyncio +async def test_ack_failure_leaves_items_for_the_next_drain(monkeypatch: pytest.MonkeyPatch) -> None: + redis_cache, redis_client = _cache_with_fakeredis(monkeypatch) + buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache) + payload: Final = json.dumps( + { + "key_list_transactions": {"key-ack": 1.5}, + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + } + ) + await redis_cache.async_rpush(key="litellm_spend_update_buffer", values=(payload,)) + + original_delete: Final = redis_client.delete + failed_delete: Final = AsyncMock(side_effect=RedisTimeoutError("DEL failed")) + monkeypatch.setattr(redis_client, "delete", failed_delete) + with pytest.raises(RedisTimeoutError, match="DEL failed"): + await buffer.get_all_update_transactions_from_redis_buffer() + + monkeypatch.setattr(redis_client, "delete", original_delete) + recovered: Final = await buffer.get_all_update_transactions_from_redis_buffer() + assert recovered is not None + assert recovered["key_list_transactions"] == {"key-ack": 1.5} + + def test_validate_redis_transaction_buffer_raises_without_redis(): """ When use_redis_transaction_buffer=true but no Redis cache is configured, @@ -577,7 +844,7 @@ 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): """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 + merging what it drained, and the drain removes the payload, so one without the key would cost a rolling deploy those increments.""" from datetime import datetime, timezone @@ -656,6 +923,7 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu class _ListRedis: def __init__(self) -> None: self.rows: list[str] = [] + self.processing: list[str] = [] async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: self.rows.extend(values) @@ -663,12 +931,15 @@ class _ListRedis: 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 + async def async_drain_list_to_processing(self, key: str, worker_id: str, count: int, ttl_seconds: int) -> list[str]: + if self.processing: + return self.processing + self.processing = self.rows[:count] + self.rows = self.rows[count:] + return self.processing + + async def async_ack_processing_list(self, key: str, worker_id: str) -> None: + self.processing = [] @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..66aa5377467 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -18,7 +18,7 @@ import threading 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, Callable, Dict, Final, Iterator, List, Optional from unittest.mock import AsyncMock, MagicMock import pytest @@ -344,6 +344,29 @@ class FakeRedisList: del stored[:count] return popped + async def async_drain_list_to_processing( + self, + key: str, + worker_id: str, + count: int, + ttl_seconds: int, + ) -> list[str]: + self._check_up() + processing_key: Final = f"{key}:processing:{worker_id}" + pending: Final = self.items.get(processing_key, []) + if pending: + return list(pending) + source: Final = self.items.get(key, []) + moved: Final = source[:count] + if moved: + self.items[key] = source[len(moved) :] + self.items[processing_key] = moved + return list(moved) + + async def async_ack_processing_list(self, key: str, worker_id: str) -> None: + self._check_up() + self.items.pop(f"{key}:processing:{worker_id}", None) + @pytest.fixture def fake_redis() -> FakeRedisList: 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..80a389b89a5 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 @@ -887,6 +887,9 @@ async def test_recover_parked_spend_logs_re_parks_rows_when_the_enqueue_is_cance mock_prisma_client._spend_log_transactions_lock.release() assert _parked_request_ids(fake_redis) == ["parked"] + recovered_rows: Final = await buffer.get_spend_logs_from_redis_buffer(limit=10) + assert [row["request_id"] for row in recovered_rows] == ["parked"] + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () assert mock_prisma_client.spend_log_transactions == [] diff --git a/tests/unit/test_redis.py b/tests/unit/test_redis.py index bfe3b55497d..60127c9649a 100644 --- a/tests/unit/test_redis.py +++ b/tests/unit/test_redis.py @@ -1,5 +1,7 @@ import inspect import json +import shutil +from pathlib import Path from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -36,6 +38,7 @@ from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL from litellm.proxy._types import CoordinationRedisParams +from tests.integration._support.redis_process import owned_redis class _StubCredentialProvider(CredentialProvider): @@ -1960,3 +1963,52 @@ def test_async_sentinel_keeps_the_credential_provider_off_the_monitors(markers, master_kwargs = mock_sentinel_cls.return_value.master_for.call_args[1] assert isinstance(master_kwargs["credential_provider"], provider_cls) assert "password" not in master_kwargs + + +@pytest.mark.asyncio +@pytest.mark.skipif(shutil.which("redis-server") is None, reason="Lua drain needs a real redis-server") +@pytest.mark.parametrize("redis_version", ["6.0", 6.1]) +async def test_async_drain_list_to_processing_keeps_legacy_redis_rows_pending_in_order( + redis_version: str | float, + tmp_path: Path, +) -> None: + with owned_redis(tmp_path) as redis_server: + redis_cache: Final = RedisCache(host=redis_server.host, port=redis_server.port, namespace="drain-test") + redis_cache.redis_version = redis_version + redis_client: Final = async_redis.Redis(host=redis_server.host, port=redis_server.port) + try: + rows: Final = ("row-1", "row-2", "row-3", "row-4", "row-5") + await redis_client.rpush("drain-test:buffer", *rows) + + first_batch: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-1", + count=2, + ttl_seconds=60, + ) + processing_key: Final = "drain-test:buffer:processing:worker-1" + assert first_batch == ["row-1", "row-2"] + assert await redis_client.llen(processing_key) == 2 + assert await redis_client.lrange("drain-test:buffer", 0, -1) == [b"row-3", b"row-4", b"row-5"] + + pending_batch: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-1", + count=2, + ttl_seconds=60, + ) + assert pending_batch == ["row-1", "row-2"] + assert await redis_client.lrange(processing_key, 0, -1) == [b"row-1", b"row-2"] + + await redis_cache.async_ack_processing_list(key="buffer", worker_id="worker-1") + next_batch: Final = await redis_cache.async_drain_list_to_processing( + key="buffer", + worker_id="worker-1", + count=2, + ttl_seconds=60, + ) + assert next_batch == ["row-3", "row-4"] + assert await redis_client.lrange("drain-test:buffer", 0, -1) == [b"row-5"] + finally: + await redis_client.aclose() + await redis_cache.disconnect()