diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5a39b8c610a..00290dfd128 100644 --- a/litellm/proxy/proxy_server.py +++ b/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. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1f453d3b1ba..2f55eb4e7f6 100644 --- a/litellm/proxy/utils.py +++ b/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) diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index a06e7142122..83d3d271c52 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,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() 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..81c892ae742 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, @@ -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,