diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index b4b2b1a334c..c810278f566 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1999,6 +1999,51 @@ class RedisCache(BaseCache): log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e) raise e + @_redis_circuit_breaker_guard + async def async_rpush_and_trim( + self, + key: str, + values: Sequence[str | bytes | int | float], + max_len: int, + ) -> int: + """Append values and keep only the newest ``max_len`` entries in one MULTI/EXEC. + + Returns the list length right after the push, so callers can tell how many + of the oldest entries the trim dropped. + """ + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + start_time: Final = time.time() + try: + async with _redis_client.pipeline(transaction=True) as pipe: + pipe.rpush(namespaced_key, *values) + pipe.ltrim(namespaced_key, -max_len, -1) + results: Final = await pipe.execute() + for r in results: + if isinstance(r, Exception): + raise r + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=time.time() - start_time, + call_type=f"async_rpush_and_trim <- {_get_call_stack_info()}", + ) + ) + return int(results[0]) + except Exception as e: + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=time.time() - start_time, + error=e, + call_type=f"async_rpush_and_trim <- {_get_call_stack_info()}", + ) + ) + log_redis_failure( + verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH+LTRIM: - Got exception from REDIS", e + ) + raise e + async def _pipeline_rpush_helper( self, pipe: pipeline, diff --git a/litellm/constants.py b/litellm/constants.py index d62cad74a36..4c8c51ee860 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -370,6 +370,9 @@ 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 = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) +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 # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) 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 cead63795a2..9044bbb3d3b 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -7,6 +7,7 @@ This is to prevent deadlocks and improve reliability import asyncio import json from collections.abc import Mapping, Sequence +from datetime import datetime from functools import reduce from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast @@ -22,6 +23,8 @@ from litellm.constants import ( 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_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) @@ -48,6 +51,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendUpdateQueue, to_wire_payload, ) +from litellm.proxy.db.spend_log_batching import SpendLogRow from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( RedisPipelineLpopOperation, @@ -93,6 +97,19 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( _ValueT = TypeVar("_ValueT") +def _spend_log_json_default(value: object) -> str: + return value.isoformat() if isinstance(value, datetime) else str(value) + + +def _encode_spend_log_row(row: SpendLogRow) -> str: + return json.dumps(row, default=_spend_log_json_default) + + +def _decode_spend_log_row(encoded: str) -> dict[str, object] | None: + decoded: Final = json.loads(encoded) + return decoded if isinstance(decoded, dict) else None + + def _accumulated_spend(totals: Mapping[str, float], entities: Mapping[str, float]) -> dict[str, float]: return {**totals, **{entity_id: totals.get(entity_id, 0) + amount for entity_id, amount in entities.items()}} @@ -526,6 +543,49 @@ class RedisUpdateBuffer: str(e), ) + async def store_spend_logs_in_redis( + self, + rows: Sequence[SpendLogRow], + max_rows: int = REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, + ) -> bool: + """Park spend-log rows in Redis so they outlive this pod, dropping the oldest past ``max_rows``.""" + if self.redis_cache is None or len(rows) == 0 or not self._should_commit_spend_updates_to_redis(): + return False + try: + buffer_size: Final = await self.redis_cache.async_rpush_and_trim( + key=REDIS_SPEND_LOGS_BUFFER_KEY, + values=[_encode_spend_log_row(row) for row in rows], + max_len=max_rows, + ) + overflow: Final = buffer_size - max_rows + if overflow > 0: + verbose_proxy_logger.error( + "Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs", + max_rows, + overflow, + ) + except Exception as e: # noqa: BLE001 # the caller falls back to the in-memory queue on any Redis fault + verbose_proxy_logger.error( + "Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), str(e) + ) + return False + verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows)) + return True + + 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.""" + 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: + return () + encoded_rows: Final = popped if isinstance(popped, list) else [popped] + decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in encoded_rows) + return tuple(row for row in decoded_rows if row is not None) + @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b078a65759e..434a6179d14 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -51,6 +51,7 @@ from litellm.constants import ( DEFAULT_MODEL_CREATED_AT_TIME, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, MAX_TEAM_LIST_LIMIT, + REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT, SPEND_LOG_QUEUE_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_BYTES, SPEND_LOG_WRITE_BATCH_MAX_ROWS, @@ -4167,6 +4168,7 @@ class PrismaClient: spend_log_flush_requested: "asyncio.Event | None" = None spend_log_queue_bytes: ClassVar[int] = 0 spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None + spend_log_write_lock = asyncio.Lock() tool_usage_transactions: list["ToolUsageTransaction"] = [] _tool_usage_transactions_lock = asyncio.Lock() autorouter_turn_transactions: ClassVar[ @@ -7062,7 +7064,7 @@ class ProxyUpdateSpend: except Exception as e: if not _is_transient_spend_log_write_error(e): if PrismaDBExceptionHandler.is_prisma_error(e): - await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process) verbose_proxy_logger.warning( "Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s", len(logs_to_process), @@ -7077,7 +7079,7 @@ class ProxyUpdateSpend: str(e), ) if i >= n_retry_times: - await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process) raise await asyncio.sleep(2**i) except Exception as e: @@ -7127,6 +7129,7 @@ async def update_spend( ) ### UPDATE SPEND LOGS ### + await recover_parked_spend_logs(prisma_client, proxy_logging_obj) # Check queue size with lock protection queue_size: Final = await _total_queued_spend_transactions(prisma_client) verbose_proxy_logger.debug("Spend Logs transactions: %s", queue_size) @@ -7144,6 +7147,51 @@ async def update_spend( ) +async def _park_spend_logs_in_redis(proxy_logging_obj: ProxyLogging, rows: Sequence[Mapping[str, object]]) -> bool: + try: + return await proxy_logging_obj.db_spend_update_writer.redis_update_buffer.store_spend_logs_in_redis(rows) + except Exception as e: # noqa: BLE001 # a Redis fault falls back to the in-memory queue, never loses the rows + verbose_proxy_logger.warning( + "Spend tracking - could not park spend logs in Redis, keeping them in memory: %s", e + ) + return False + + +async def requeue_spend_logs( + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + rows: Sequence[Mapping[str, object]], +) -> None: + """Park rows from a failed or cancelled write in Redis, falling back to the head of the in-memory queue.""" + if await _park_spend_logs_in_redis(proxy_logging_obj, rows): + return + await enqueue_spend_logs(prisma_client, rows, at_head=True) + + +async def recover_parked_spend_logs( + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + limit: int = REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT, +) -> int: + """Move spend-log rows parked in Redis back to the head of the in-memory queue for the next write.""" + try: + rows: Final = ( + await proxy_logging_obj.db_spend_update_writer.redis_update_buffer.get_spend_logs_from_redis_buffer(limit) + ) + except Exception as e: # noqa: BLE001 # Redis being down must not stop the regular in-memory flush + verbose_proxy_logger.warning("Spend tracking - could not read parked spend logs from Redis: %s", e) + return 0 + if len(rows) == 0: + return 0 + try: + await enqueue_spend_logs(prisma_client, rows, at_head=True) + except BaseException: + await _park_spend_logs_in_redis(proxy_logging_obj, rows) + raise + verbose_proxy_logger.info("Spend tracking - recovered %d parked spend log rows from Redis", len(rows)) + return len(rows) + + async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int: """Pending entries across every request-time spend queue, sized under each queue's lock. Every drain trigger reads this one owner, so a queue added later joins the @@ -7215,14 +7263,19 @@ async def update_spend_logs_job( This job is triggered based on queue size rather than time. Pops the batch once, writes spend logs, then runs guardrail usage tracking. """ - n_retry_times: Final = 3 - MAX_LOGS_PER_INTERVAL: Final = 10000 - - # Atomically pop batch from queue. The tool usage queue counts toward the - # emptiness check: a spend-log write failure aborts a run before the tool - # drain below, and those entries must not strand once the spend queue drains. if await _total_queued_spend_transactions(prisma_client) == 0: return + async with prisma_client.spend_log_write_lock: + await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj) + + +async def _run_spend_logs_job( + prisma_client: PrismaClient, + db_writer_client: AsyncHTTPHandler | None, + proxy_logging_obj: ProxyLogging, +) -> None: + n_retry_times: Final = 3 + MAX_LOGS_PER_INTERVAL: Final = 10000 logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL) @@ -7235,7 +7288,7 @@ async def update_spend_logs_job( logs_to_process=logs_to_process, ) except asyncio.CancelledError: - await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process) verbose_proxy_logger.warning( "Spend tracking - spend log write cancelled, requeued %d rows for the next flush", len(logs_to_process), @@ -7321,14 +7374,22 @@ async def drain_spend_logs_queue( await monitor_task prisma_client.spend_logs_queue_monitor_task = None # rebind-ok: the client owns its monitor handle + async with prisma_client.spend_log_write_lock: + try: + await _drain_spend_logs_queue_to_db(prisma_client, db_writer_client, proxy_logging_obj) + finally: + await _park_remaining_spend_logs(prisma_client, proxy_logging_obj) + + +async def _drain_spend_logs_queue_to_db( + prisma_client: PrismaClient, + db_writer_client: "AsyncHTTPHandler | None", + proxy_logging_obj: ProxyLogging, +) -> None: for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS): if await _total_queued_spend_transactions(prisma_client) == 0: return - await update_spend_logs_job( - prisma_client=prisma_client, - db_writer_client=db_writer_client, - proxy_logging_obj=proxy_logging_obj, - ) + await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj) remaining: Final = await _total_queued_spend_transactions(prisma_client) if remaining > 0: @@ -7339,6 +7400,17 @@ async def drain_spend_logs_queue( ) +async def _park_remaining_spend_logs(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: + rows: Final = await dequeue_spend_logs(prisma_client, sys.maxsize) + if len(rows) == 0 or await _park_spend_logs_in_redis(proxy_logging_obj, rows): + return + await enqueue_spend_logs(prisma_client, rows, at_head=True) + spend_log_error( + "Spend tracking - %d spend log rows could not be written or parked in Redis and will be lost on exit", + len(rows), + ) + + async def _monitor_spend_logs_queue( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, @@ -7372,6 +7444,7 @@ async def _monitor_spend_logs_queue( while True: try: + await recover_parked_spend_logs(prisma_client, proxy_logging_obj) # Check queue sizes with lock protection; the tool usage queue keeps # the monitor firing when a prior failed run left it nonempty. queue_size = await _total_queued_spend_transactions(prisma_client) diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index a28a78cc4a1..8d78b4b61c1 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -42,6 +42,7 @@ class MockPrismaClient: import asyncio self._spend_log_transactions_lock = asyncio.Lock() + self.spend_log_write_lock = asyncio.Lock() self._tool_usage_transactions_lock = asyncio.Lock() self._autorouter_turn_transactions_lock = asyncio.Lock() diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 19638c60b4b..44227b8ed33 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1502,3 +1502,51 @@ async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypat ("ns:u1", '{"user_id": "u1"}', timedelta(seconds=7)), ("ns:org_id:o1", '{"a": 1}', timedelta(seconds=300)), ] + + +class _ListPipeline: + def __init__(self, rows: list[str]) -> None: + self.rows = rows + self.queued: list[tuple[str, ...]] = [] + + async def __aenter__(self) -> "_ListPipeline": + return self + + async def __aexit__(self, *exc: object) -> None: + return None + + def rpush(self, key: str, *values: str) -> None: + self.queued.append(("rpush", key, *values)) + + def ltrim(self, key: str, start: int, end: int) -> None: + self.queued.append(("ltrim", key, str(start), str(end))) + + async def execute(self) -> list[object]: + results: list[object] = [] + for op in self.queued: + if op[0] == "rpush": + self.rows.extend(op[2:]) + results.append(len(self.rows)) + else: + start, end = int(op[2]), int(op[3]) + del self.rows[: max(len(self.rows) + start, 0) if start < 0 else start] + results.append(True) + return results + + +@pytest.mark.asyncio +async def test_async_rpush_and_trim_runs_push_and_trim_in_one_transaction(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + rows = ["a", "b"] + pipe = _ListPipeline(rows) + client = MagicMock() + client.pipeline = MagicMock(return_value=pipe) + + with patch.object(redis_cache, "init_async_client", return_value=client): + pushed_len = await redis_cache.async_rpush_and_trim(key="buf", values=["c", "d"], max_len=3) + + client.pipeline.assert_called_once_with(transaction=True) + assert pushed_len == 4 + assert rows == ["b", "c", "d"] + assert [op[:2] for op in pipe.queued] == [("rpush", "ns:buf"), ("ltrim", "ns:buf")] diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index cc8b10150bd..e04e2402e1b 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -651,3 +651,53 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu restored = await window_queue.flush_and_get_aggregated_window_spend_transactions() assert [payload["spend"] for payload in restored] == [4.0] assert [payload["entity_id"] for payload in restored] == ["team-1"] + + +class _ListRedis: + def __init__(self) -> None: + self.rows: list[str] = [] + + async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: + self.rows.extend(values) + pushed_len = len(self.rows) + del self.rows[:-max_len] + return pushed_len + + async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> list[str] | None: + if not self.rows: + return None + popped = self.rows[:count] + del self.rows[:count] + return popped + + +@pytest.mark.asyncio +async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap(): + redis = _ListRedis() + buffer = RedisUpdateBuffer(redis_cache=redis) + buffer._should_commit_spend_updates_to_redis = MagicMock(return_value=True) + + assert await buffer.store_spend_logs_in_redis([{"request_id": "old"}, {"request_id": "mid"}], max_rows=2) is True + assert await buffer.store_spend_logs_in_redis([{"request_id": "new"}], max_rows=2) is True + + parked = await buffer.get_spend_logs_from_redis_buffer(limit=10) + assert [row["request_id"] for row in parked] == ["mid", "new"] + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () + + +@pytest.mark.asyncio +async def test_store_spend_logs_in_redis_reports_failure_without_redis(): + buffer = RedisUpdateBuffer(redis_cache=None) + + assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False + assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == () + + +@pytest.mark.asyncio +async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_enabled(): + redis = _ListRedis() + buffer = RedisUpdateBuffer(redis_cache=redis) + buffer._should_commit_spend_updates_to_redis = MagicMock(return_value=False) + + assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False + assert redis.rows == [] diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py index fce51c9296c..c502fe4800e 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/conftest.py @@ -130,6 +130,7 @@ def mock_prisma_client() -> MagicMock: client.spend_log_transactions = [] client._spend_log_transactions_lock = asyncio.Lock() client.spend_logs_queue_monitor_task = None + client.spend_log_write_lock = asyncio.Lock() client.tool_usage_transactions = [] client._tool_usage_transactions_lock = asyncio.Lock() client.jsonify_object = lambda data: dict(data) @@ -313,6 +314,54 @@ def make_spend_log_row() -> Callable[..., Dict[str, Any]]: return _make +class FakeRedisList: + def __init__(self) -> None: + self.items: dict[str, list[str]] = {} + self.down = False + + def _check_up(self) -> None: + if self.down: + raise ConnectionError("redis unreachable") + + async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int: + self._check_up() + stored = self.items.setdefault(key, []) + stored.extend(str(v) for v in values) + pushed_len = len(stored) + del stored[:-max_len] + return pushed_len + + async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> str | list[str] | None: + self._check_up() + stored = self.items.get(key, []) + if not stored: + return None + if count is None: + return stored.pop(0) + popped = stored[:count] + del stored[:count] + return popped + + +@pytest.fixture +def fake_redis() -> FakeRedisList: + return FakeRedisList() + + +@pytest.fixture +def proxy_logging_with_redis(fake_redis: FakeRedisList) -> MagicMock: + from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + proxy_logging.db_spend_update_writer = MagicMock() + proxy_logging.db_spend_update_writer.db_update_spend_transaction_handler = AsyncMock() + buffer = RedisUpdateBuffer(redis_cache=fake_redis) + buffer._should_commit_spend_updates_to_redis = MagicMock(return_value=True) + proxy_logging.db_spend_update_writer.redis_update_buffer = buffer + return proxy_logging + + @dataclass class _SentMessage: from_addr: Optional[str] diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index d671a4ffc1f..7099101db1c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -883,3 +883,37 @@ def test_disable_spend_updates_error_when_general_settings_unavailable( monkeypatch.delattr(proxy_server_mod, "general_settings", raising=False) with pytest.raises(ImportError): ProxyUpdateSpend.disable_spend_updates() + + +@pytest.mark.asyncio +async def test_update_spend_logs_parks_failed_batch_in_redis_with_wire_safe_datetimes( + mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any +) -> None: + """Regression: a batch the DB rejected used to go back to process memory only. With Redis + wired in it must be parked there, and datetimes must come back as ISO strings the DB write + accepts, since the row is replayed by a process that never saw the original objects. + """ + from datetime import datetime, timezone + + from prisma.errors import TableNotFoundError + + started = datetime(2026, 9, 19, 20, 0, 5, 123000, tzinfo=timezone.utc) + err = TableNotFoundError( + {"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}} + ) + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err) + mock_prisma_client.spend_log_transactions = [] + + with pytest.raises(TableNotFoundError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + logs_to_process=[make_spend_log_row(request_id="a", startTime=started)], + ) + + buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer + parked = await buffer.get_spend_logs_from_redis_buffer(limit=10) + assert mock_prisma_client.spend_log_transactions == [] + assert [(row["request_id"], row["startTime"]) for row in parked] == [("a", started.isoformat())] diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index c8b87bd671e..d6f41ba55db 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -11,17 +11,20 @@ Symbols pinned here: from __future__ import annotations import asyncio +import json from contextlib import suppress from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, _monitor_spend_logs_queue, _raise_failed_update_spend_exception, drain_spend_logs_queue, + recover_parked_spend_logs, update_daily_tag_spend, update_spend, update_spend_logs_job, @@ -719,3 +722,222 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: with pytest.raises(ValueError, match="specific"): asyncio.run(_runner()) + + +def _table_gone_error() -> Exception: + from prisma.errors import TableNotFoundError + + return TableNotFoundError( + {"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}} + ) + + +def _parked_request_ids(fake_redis: Any) -> list[str]: + return [json.loads(row)["request_id"] for row in fake_redis.items.get(REDIS_SPEND_LOGS_BUFFER_KEY, [])] + + +@pytest.mark.asyncio +async def test_drain_spend_logs_queue_parks_unwritable_rows_in_redis_on_shutdown( + mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any +) -> None: + from prisma.errors import TableNotFoundError + + mock_prisma_client.spend_log_transactions = [ + make_spend_log_row(request_id="r1"), + make_spend_log_row(request_id="r2"), + ] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_table_gone_error()) + + with pytest.raises(TableNotFoundError): + await drain_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + + assert mock_prisma_client.spend_log_transactions == [] + assert sorted(_parked_request_ids(fake_redis)) == ["r1", "r2"] + + +@pytest.mark.asyncio +async def test_drain_spend_logs_queue_waits_for_an_in_flight_write_before_parking( + mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any +) -> None: + db_outage_seen: Final = asyncio.Event() + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="in-flight")] + + async def _fail_once_shutdown_starts(*args: Any, **kwargs: Any) -> None: + await db_outage_seen.wait() + raise _table_gone_error() + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_fail_once_shutdown_starts) + scheduler_write: Final = asyncio.ensure_future( + update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + ) + await asyncio.sleep(0) + assert mock_prisma_client.spend_log_transactions == [] + + async def _release_after_shutdown_started() -> None: + await asyncio.sleep(0.05) + db_outage_seen.set() + + release: Final = asyncio.ensure_future(_release_after_shutdown_started()) + await drain_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + + assert _parked_request_ids(fake_redis) == ["in-flight"] + assert mock_prisma_client.spend_log_transactions == [] + await release + with suppress(Exception): + await scheduler_write + + +@pytest.mark.asyncio +async def test_drain_spend_logs_queue_parks_rows_left_after_max_passes( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, + proxy_logging_with_redis: MagicMock, + fake_redis: Any, +) -> None: + import litellm.proxy.db.spend_log_tool_index as tool_mod + import litellm.proxy.guardrails.usage_tracking as guard_mod + + monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False) + monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False) + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r0")] + + async def _write_and_refill(*args: Any, **kwargs: Any) -> None: + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="late")) + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_write_and_refill) + + await drain_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + + assert mock_prisma_client.spend_log_transactions == [] + assert _parked_request_ids(fake_redis) == ["late"] + + +@pytest.mark.asyncio +async def test_drain_spend_logs_queue_keeps_rows_in_memory_when_redis_is_down( + mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any +) -> None: + from prisma.errors import TableNotFoundError + + fake_redis.down = True + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_table_gone_error()) + + with pytest.raises(TableNotFoundError): + await drain_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + + assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r1"] + assert fake_redis.items == {} + + +@pytest.mark.asyncio +async def test_update_spend_writes_rows_parked_in_redis_by_a_previous_pod( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, + proxy_logging_with_redis: MagicMock, + fake_redis: Any, +) -> None: + import litellm.proxy.db.spend_log_tool_index as tool_mod + import litellm.proxy.guardrails.usage_tracking as guard_mod + + monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False) + monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False) + buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer + assert await buffer.store_spend_logs_in_redis([make_spend_log_row(request_id="parked")]) is True + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock() + + await update_spend( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + + written = mock_prisma_client.db.litellm_spendlogs.create_many.await_args.kwargs["data"] + assert [row["request_id"] for row in written] == ["parked"] + assert _parked_request_ids(fake_redis) == [] + assert mock_prisma_client.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_recover_parked_spend_logs_re_parks_rows_when_the_enqueue_is_cancelled( + mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any +) -> None: + buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer + assert await buffer.store_spend_logs_in_redis([make_spend_log_row(request_id="parked")]) is True + mock_prisma_client.spend_log_transactions = [] + await mock_prisma_client._spend_log_transactions_lock.acquire() + recovery: Final = asyncio.ensure_future( + recover_parked_spend_logs(prisma_client=mock_prisma_client, proxy_logging_obj=proxy_logging_with_redis) + ) + await asyncio.sleep(0.01) + assert _parked_request_ids(fake_redis) == [] + + recovery.cancel() + with pytest.raises(asyncio.CancelledError): + await recovery + mock_prisma_client._spend_log_transactions_lock.release() + + assert _parked_request_ids(fake_redis) == ["parked"] + assert mock_prisma_client.spend_log_transactions == [] + + +@pytest.mark.asyncio +async def test_monitor_spend_logs_queue_pulls_parked_rows_before_each_flush( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, + proxy_logging_with_redis: MagicMock, +) -> None: + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 0.0, raising=False) + buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer + assert await buffer.store_spend_logs_in_redis([make_spend_log_row(request_id="parked")]) is True + mock_prisma_client.spend_log_transactions = [] + seen: list[list[str]] = [] + polls = {"n": 0} + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + seen.append([row["request_id"] for row in mock_prisma_client.spend_log_transactions]) + raise asyncio.CancelledError() + + async def _poll(*args: Any, **kwargs: Any) -> bool: + polls["n"] += 1 + if polls["n"] >= 3: + raise asyncio.CancelledError() + return False + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + monkeypatch.setattr(utils_mod, "_wait_for_spend_log_flush_request", _poll) + + with pytest.raises(asyncio.CancelledError): + await _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging_with_redis, + ) + + assert seen == [["parked"]]