diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 924189fed4b..6adeee1b806 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5538,6 +5538,33 @@ async def update_daily_tag_spend( verbose_proxy_logger.error(f"Error updating daily tag spend: {e}") +def _consume_task_exception(task: "asyncio.Future[None]") -> None: + if not task.cancelled(): + task.exception() + + +async def _requeue_spend_log_transactions(prisma_client: PrismaClient, logs: Sequence[Mapping[str, Any]]) -> None: + """Put a popped batch back at the head of the spend-log queue. + + Re-writing a row that already landed is harmless: the bulk insert runs with + ``skip_duplicates=True``, so a requeued batch that partially made it to the + DB writes only the rows that are still missing. + """ + if not logs: + return + async with prisma_client._spend_log_transactions_lock: + prisma_client.spend_log_transactions = [*logs, *prisma_client.spend_log_transactions] + + +async def _requeue_tool_usage_transactions( + prisma_client: PrismaClient, transactions: Sequence["ToolUsageTransaction"] +) -> None: + if not transactions: + return + async with prisma_client._tool_usage_transactions_lock: + prisma_client.tool_usage_transactions = [*transactions, *prisma_client.tool_usage_transactions] + + async def update_spend_logs_job( prisma_client: PrismaClient, db_writer_client: Optional[AsyncHTTPHandler], @@ -5548,6 +5575,12 @@ 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. + + The batch leaves the in-memory queue before the DB write is awaited, so a + cancellation landing in that window (task cancelled, worker shutting down, + ``wait_for`` timeout) would otherwise drop those rows for good. The write is + shielded so it still runs to completion, and the batch is requeued so a + later flush retries it if the write never finished. """ n_retry_times = 3 MAX_LOGS_PER_INTERVAL = 10000 @@ -5566,13 +5599,21 @@ async def update_spend_logs_job( logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[len(logs_to_process) :] - await 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, + write_task = asyncio.ensure_future( + 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, + ) ) + try: + await asyncio.shield(write_task) + except asyncio.CancelledError: + write_task.add_done_callback(_consume_task_exception) + await _requeue_spend_log_transactions(prisma_client, logs_to_process) + raise # Guardrail/policy usage tracking (same batch, outside spend-logs update) try: @@ -5599,10 +5640,18 @@ async def update_spend_logs_job( try: from litellm.proxy.db.spend_log_tool_index import flush_tool_usage_transactions - await flush_tool_usage_transactions( - prisma_client=prisma_client, - transactions=tool_usage_to_process, + tool_flush_task = asyncio.ensure_future( + flush_tool_usage_transactions( + prisma_client=prisma_client, + transactions=tool_usage_to_process, + ) ) + try: + await asyncio.shield(tool_flush_task) + except asyncio.CancelledError: + tool_flush_task.add_done_callback(_consume_task_exception) + await _requeue_tool_usage_transactions(prisma_client, tool_usage_to_process) + raise except Exception as tool_tracking_err: verbose_proxy_logger.error( "Spend tracking - tool usage flush failed; %s tool usage transactions dropped: %s", 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 d9eeb168611..3ab4cfd778c 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 @@ -263,6 +263,119 @@ async def test_update_spend_logs_job_processes_and_clears_queue( } +async def _wait_for(predicate: Any, iterations: int = 200) -> None: + for _ in range(iterations): + if predicate(): + return + await asyncio.sleep(0) + raise AssertionError("condition never became true") + + +@pytest.mark.asyncio +async def test_update_spend_logs_job_requeues_batch_when_cancelled_mid_write( + mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """Regression: the batch is popped from the in-memory queue before the DB + write is awaited, so a cancellation arriving in that window used to lose + the rows permanently. The shielded write must still complete and the batch + must go back to the head of the queue, behind nothing that was enqueued + while the flush was in flight. + """ + 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) + + write_started = asyncio.Event() + release_write = asyncio.Event() + written: List[List[Dict[str, Any]]] = [] + + async def _slow_create_many(data: Any, skip_duplicates: bool = False) -> None: + write_started.set() + await release_write.wait() + written.append(list(data)) + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_slow_create_many) + mock_prisma_client.spend_log_transactions = [ + make_spend_log_row(request_id="r1"), + make_spend_log_row(request_id="r2"), + ] + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + job = asyncio.create_task( + update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + ) + await write_started.wait() + assert mock_prisma_client.spend_log_transactions == [] + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r3")) + + job.cancel() + with pytest.raises(asyncio.CancelledError): + await job + + requeued = [row["request_id"] for row in mock_prisma_client.spend_log_transactions] + + release_write.set() + await _wait_for(lambda: bool(written)) + + assert requeued == ["r1", "r2", "r3"] + assert [row["request_id"] for row in written[0]] == ["r1", "r2"] + + +@pytest.mark.asyncio +async def test_update_spend_logs_job_requeues_tool_usage_when_cancelled_mid_flush( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> 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) + + flush_started = asyncio.Event() + release_flush = asyncio.Event() + flushed: List[Any] = [] + + async def _slow_flush(prisma_client: Any, transactions: Any) -> None: + flush_started.set() + await release_flush.wait() + flushed.append(list(transactions)) + + monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", _slow_flush, raising=False) + + tool_transaction = MagicMock() + mock_prisma_client.spend_log_transactions = [] + mock_prisma_client.tool_usage_transactions = [tool_transaction] + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + job = asyncio.create_task( + update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + ) + ) + await flush_started.wait() + assert mock_prisma_client.tool_usage_transactions == [] + + job.cancel() + with pytest.raises(asyncio.CancelledError): + await job + + requeued = list(mock_prisma_client.tool_usage_transactions) + + release_flush.set() + await _wait_for(lambda: bool(flushed)) + + assert requeued == [tool_transaction] + + @pytest.mark.asyncio async def test_monitor_spend_logs_queue_invokes_job_when_queue_nonempty( mock_prisma_client: Any,