fix: bound spend log shutdown drain

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)

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,