Merge bea74c144e into a6b7384094

litellm/proxy/proxy_server.py

@@ -934,7 +934,15 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N
     if worker_heartbeat is not None and prisma_client:
         await worker_heartbeat.deregister()
     if prisma_client:
-        # Drain the SGR fold first: it lives in memory, so an un-drained interval
+        # Request-time spend queues live in this process. A worker recycle
+        # (`--max_requests_before_restart`) or deploy otherwise drops whatever
+        # the periodic flush has not written. Drain while Prisma is still
+        # connected; a write after disconnect is ClientNotConnectedError and
+        # the rows never reach LiteLLM_SpendLogs. The same emptiness owner
+        # (`_total_queued_spend_transactions`) also covers
+        # tool_usage_transactions and autorouter_turn_transactions.
+        await _flush_spend_logs_queue_on_shutdown()
+        # Drain the SGR fold next: it lives in memory, so an un-drained interval
         # is lost, and a write attempted after disconnect raises
         # ClientNotConnectedError rather than persisting anything. Ordering this
         # inside the same guard is what keeps the two from drifting apart.

litellm/proxy/utils.py

@@ -6527,6 +6527,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.
@@ -6546,12 +6547,15 @@ 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(
-            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,
+        await asyncio.wait_for(
+            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,
+            ),
+            timeout=timeout,
         )
     except asyncio.CancelledError:
         await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
@@ -6560,6 +6564,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:
@@ -6626,6 +6638,7 @@ async def update_spend_logs_job(
 
 
 MAX_SPEND_LOG_DRAIN_ITERATIONS: Final = 20
+MAX_SPEND_LOG_DRAIN_SECONDS: Final = 15.0
 
 
 async def drain_spend_logs_queue(
@@ -6640,13 +6653,18 @@ 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
 
+    deadline: Final = time.monotonic() + MAX_SPEND_LOG_DRAIN_SECONDS
     for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS):
         if await _total_queued_spend_transactions(prisma_client) == 0:
             return
+        remaining_seconds = 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/proxy_server/test_lifecycle.py

@@ -98,6 +98,7 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch):
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock()
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", AsyncMock(), raising=False)
     monkeypatch.setattr(ps, "master_key", "sk-x", raising=False)
 
     fake_jwt = MagicMock()
@@ -142,6 +143,7 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", AsyncMock(), raising=False)
 
     async def _record_flush(client, accumulator):
         calls.append("flush")
@@ -164,6 +166,50 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke
     assert calls == ["flush", "disconnect"]
 
 
+@pytest.mark.asyncio
+async def test_proxy_shutdown_drains_spend_logs_before_disconnecting(monkeypatch):
+    """
+    SpendLogs, tool-usage, and auto-router turn queues live in memory.
+
+    ``proxy_shutdown_event`` is the function that disconnects Prisma, so it
+    must drain those queues first. The lifespan helper already flushes once,
+    but anything queued after that — or a caller that hits this function
+    directly — is otherwise discarded on worker recycle with no log line.
+    Ordering is the behavior, so assert drain then disconnect.
+    """
+    calls: list[str] = []  # mutable-ok: records call order, which is the assertion
+
+    fake_prisma = MagicMock()
+    fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect"))
+    monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+
+    async def _record_drain() -> None:
+        calls.append("drain_spend")
+
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", _record_drain, raising=False)
+
+    async def _record_flush(client: object, accumulator: object) -> None:
+        del accumulator
+        calls.append("flush_gateway")
+        assert client is fake_prisma
+
+    monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False)
+
+    fake_jwt = MagicMock()
+    fake_jwt.close = AsyncMock()
+    monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False)
+    monkeypatch.setattr(ps, "db_writer_client", None, raising=False)
+
+    import litellm
+
+    monkeypatch.setattr(litellm, "cache", None, raising=False)
+    monkeypatch.setattr(litellm, "success_callback", [], raising=False)
+
+    await proxy_shutdown_event()
+
+    assert calls == ["drain_spend", "flush_gateway", "disconnect"]
+
+
 @pytest.mark.asyncio
 async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch):
     """No prisma client means nothing to drain to, and no attempt is made."""
@@ -191,6 +237,7 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch):
     fake_prisma = MagicMock()
     fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone"))
     monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
+    monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", AsyncMock(), raising=False)
 
     fake_jwt = MagicMock()
     fake_jwt.close = AsyncMock()

tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py

@@ -19,6 +19,7 @@ import pytest
 
 from litellm.proxy.utils import (
     MAX_SPEND_LOG_DRAIN_ITERATIONS,
+    MAX_SPEND_LOG_DRAIN_SECONDS,
     _monitor_spend_logs_queue,
     _raise_failed_update_spend_exception,
     drain_spend_logs_queue,
@@ -350,8 +351,9 @@ async def test_drain_spend_logs_queue_flushes_rows_queued_while_draining(
 
     written: list[str] = []
 
-    async def _write(*args: Any, **kwargs: Any) -> None:
-        written.extend(row["request_id"] for row in kwargs["data"])
+    async def _write(*, data: list[dict[str, object]], skip_duplicates: bool) -> None:
+        del skip_duplicates
+        written.extend(row["request_id"] for row in data)
         if len(written) == 1:
             mock_prisma_client.spend_log_transactions.append(
                 make_spend_log_row(request_id="r2")
@@ -391,12 +393,13 @@ async def test_drain_spend_logs_queue_stops_monitor_and_keeps_its_popped_rows(
     written: list[str] = []
     write_calls = {"n": 0}
 
-    async def _write(*args: Any, **kwargs: Any) -> None:
+    async def _write(*, data: list[dict[str, object]], skip_duplicates: bool) -> None:
+        del skip_duplicates
         write_calls["n"] += 1
         if write_calls["n"] == 1:
             write_started.set()
             await asyncio.Event().wait()
-        written.extend(row["request_id"] for row in kwargs["data"])
+        written.extend(row["request_id"] for row in data)
 
     mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_write)
 
@@ -439,7 +442,10 @@ async def test_drain_spend_logs_queue_gives_up_after_max_passes(
     proxy_logging.failure_handler = AsyncMock()
     mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")]
 
-    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
         mock_prisma_client.spend_log_transactions.append(make_spend_log_row())
 
     mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
@@ -458,6 +464,75 @@ async def test_drain_spend_logs_queue_gives_up_after_max_passes(
     )
 
 
+@pytest.mark.asyncio
+async def test_drain_spend_logs_queue_gives_up_when_time_budget_expires(
+    mock_prisma_client: Any, make_spend_log_row: 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
+    import litellm.proxy.utils as utils_mod
+
+    monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False)
+    monkeypatch.setattr(tool_mod, "process_spend_logs_tool_usage", AsyncMock(), raising=False)
+
+    proxy_logging = MagicMock()
+    proxy_logging.failure_handler = AsyncMock()
+    mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")]
+
+    monotonic = {"t": 0.0}
+
+    def _now() -> float:
+        return monotonic["t"]
+
+    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())
+
+    monkeypatch.setattr(utils_mod.time, "monotonic", _now)
+    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,
+    )
+
+    # First pass starts under the deadline; the write then consumes the whole
+    # budget, so the next loop exits instead of spinning to max iterations.
+    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,