diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e55b254ab8e..549d439d4d0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -896,7 +896,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. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8cbf5b685fd..095cea9ea28 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6405,6 +6405,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( @@ -6419,9 +6420,12 @@ 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 + if time.monotonic() >= deadline: + break await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..3ff99fb81b5 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/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,49 @@ 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 = [] # 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(): + calls.append("drain_spend") + + monkeypatch.setattr(ps, "_flush_spend_logs_queue_on_shutdown", _record_drain, raising=False) + + async def _record_flush(client, 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 +236,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() 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 a1eb88a7834..b4c7802d757 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 @@ -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, @@ -458,6 +459,44 @@ 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(*args: Any, **kwargs: Any) -> None: + 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_monitor_spend_logs_queue_invokes_job_when_queue_nonempty( mock_prisma_client: Any,