diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f2b2cf066fe..0599da86dcc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6285,6 +6285,7 @@ async def update_spend_logs_job( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, + timeout: float | None = None, ): """ Job to process spend_log_transactions queue. @@ -6304,13 +6305,17 @@ async def update_spend_logs_job( logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL) try: - await ProxyUpdateSpend.update_spend_logs( + update_spend_logs = ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, db_writer_client=db_writer_client, logs_to_process=logs_to_process, ) + if timeout is None: + await update_spend_logs + else: + await asyncio.wait_for(update_spend_logs, timeout=timeout) except asyncio.CancelledError: await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) verbose_proxy_logger.warning( @@ -6318,6 +6323,14 @@ async def update_spend_logs_job( len(logs_to_process), ) raise + except asyncio.TimeoutError: + await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True) + verbose_proxy_logger.warning( + "Spend tracking - spend log write timed out after %.2f seconds; requeued %d rows", + timeout or 0.0, + len(logs_to_process), + ) + return # Guardrail/policy usage tracking (same batch, outside spend-logs update) try: @@ -6396,12 +6409,14 @@ async def drain_spend_logs_queue( for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS): if await _total_queued_spend_transactions(prisma_client) == 0: return - if time.monotonic() >= deadline: + remaining_seconds: Final = deadline - time.monotonic() + if remaining_seconds <= 0: break await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, proxy_logging_obj=proxy_logging_obj, + timeout=remaining_seconds, ) remaining: Final = await _total_queued_spend_transactions(prisma_client) 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 b4c7802d757..20f651c2e35 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 @@ -479,7 +479,10 @@ async def test_drain_spend_logs_queue_gives_up_when_time_budget_expires( def _now() -> float: return monotonic["t"] - async def _write_and_refill(*args: Any, **kwargs: Any) -> None: + async def _write_and_refill( + *, data: list[dict[str, object]], skip_duplicates: bool + ) -> None: + del data, skip_duplicates monotonic["t"] += MAX_SPEND_LOG_DRAIN_SECONDS mock_prisma_client.spend_log_transactions.append(make_spend_log_row()) @@ -497,6 +500,34 @@ async def test_drain_spend_logs_queue_gives_up_when_time_budget_expires( assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1 +@pytest.mark.asyncio +async def test_drain_spend_logs_queue_requeues_when_write_exceeds_time_budget( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + import litellm.proxy.utils as utils_mod + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + row = make_spend_log_row(request_id="r1") + mock_prisma_client.spend_log_transactions = [row] + monkeypatch.setattr(utils_mod, "MAX_SPEND_LOG_DRAIN_SECONDS", 0.01) + + async def _hang(**_: object) -> None: + await asyncio.Event().wait() + + write_mock = AsyncMock(side_effect=_hang) + monkeypatch.setattr(utils_mod.ProxyUpdateSpend, "update_spend_logs", write_mock) + + await drain_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + + assert write_mock.await_count == 1 + assert mock_prisma_client.spend_log_transactions == [row] + + @pytest.mark.asyncio async def test_monitor_spend_logs_queue_invokes_job_when_queue_nonempty( mock_prisma_client: Any,