From 88873a70d41744ea57e031fb62cc571ea5fdee37 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Fri, 13 Nov 2026 13:36:12 -0800 Subject: [PATCH] fix(proxy): retain unwritten rollups after database disconnects (#51706) * fix(proxy): retain unwritten rollups after database disconnects * fix(proxy): initialize auto-router turns per client --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/proxy/db/autorouter_session_rollup.py | 104 ++++++++++++------ litellm/proxy/db/exception_handler.py | 13 ++- litellm/proxy/utils.py | 48 +++++--- .../db/test_autorouter_session_rollup.py | 48 +++++++- tests/unit/proxy/db/test_exception_handler.py | 6 +- .../prisma_and_spend/test_spend_functions.py | 83 +++++++++++++- 6 files changed, 242 insertions(+), 60 deletions(-) diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 4ddcf7e600b..cc1806e7f77 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -21,7 +21,7 @@ import random from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone -from itertools import groupby +from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, NamedTuple @@ -30,6 +30,7 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES from litellm.proxy.db.create_views import SupportsExecuteRaw from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.exception_handler import db_write_never_sent if TYPE_CHECKING: from litellm.proxy._types import SpendLogsPayload @@ -536,77 +537,110 @@ def _session_partition(transaction: AutoRouterTurnTransaction) -> tuple[str, str return (*identity, transaction.session_id, transaction.router_name) +class UnwrittenTurns(NamedTuple): + """Turns the database provably never received, and the error that proved it.""" + + turns: tuple[AutoRouterTurnTransaction, ...] + error: Exception + + +class _StoppedPartition(NamedTuple): + turns: tuple[AutoRouterTurnTransaction, ...] + error: Exception + + async def _drain_session_partition( prisma_client: PrismaClient, transactions: tuple[AutoRouterTurnTransaction, ...], n_retry_times: int, statement: str, -) -> tuple[AutoRouterTurnTransaction, ...]: +) -> _StoppedPartition | None: for position, transaction in enumerate(transactions): try: await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times, statement) - except Exception as flush_err: # noqa: BLE001 # stop dependent turns without retrying an ambiguous write - verbose_proxy_logger.error( - "Spend tracking - auto-router session rollup flush failed for router %s; " - "%s of %s turn writes stopped in this partition: %s", - transaction.router_name, - len(transactions) - position, - len(transactions), - flush_err, - ) - return transactions[position:] - return () + except Exception as flush_err: # noqa: BLE001 # the caller keeps or drops the stopped turns + return _StoppedPartition(transactions[position:], flush_err) + return None + + +def _log_dropped_turns(stopped: _StoppedPartition, partition_size: int) -> None: + verbose_proxy_logger.error( + "Spend tracking - auto-router session rollup flush failed for router %s; " + "%s of %s turn writes stopped in this partition: %s", + stopped.turns[0].router_name, + len(stopped.turns), + partition_size, + stopped.error, + ) async def _flush_session_partition( prisma_client: PrismaClient, transactions: tuple[AutoRouterTurnTransaction, ...], n_retry_times: int, -) -> None: - failed_suffix: Final = await _drain_session_partition( +) -> UnwrittenTurns | None: + stopped: Final = await _drain_session_partition( prisma_client, transactions, n_retry_times, UPSERT_AUTOROUTER_SESSION_SQL ) - if not failed_suffix or not failed_suffix[0].api_key: - return - failed_user: Final = failed_suffix[0].user_id + if stopped is None: + return None + if db_write_never_sent(stopped.error): + return UnwrittenTurns(stopped.turns, stopped.error) + _log_dropped_turns(stopped, len(transactions)) + if not stopped.turns[0].api_key: + return None + failed_user: Final = stopped.turns[0].user_id other_users: Final = sorted( ( transaction - for transaction in failed_suffix[1:] + for transaction in stopped.turns[1:] if transaction.user_id and transaction.user_id != failed_user ), key=lambda transaction: transaction.user_id, ) - for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id): - await _drain_session_partition( - prisma_client, tuple(user_turns), n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL - ) + user_partitions: Final = tuple( + tuple(user_turns) for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id) + ) + for user_partition in user_partitions: + if ( + projection_stopped := await _drain_session_partition( + prisma_client, user_partition, n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL + ) + ) is not None: + _log_dropped_turns(projection_stopped, len(user_partition)) + return None async def flush_autorouter_turn_transactions( prisma_client: PrismaClient, transactions: Sequence[AutoRouterTurnTransaction], n_retry_times: int = 3, -) -> None: +) -> UnwrittenTurns | None: """Drain a queue batch into the rollup, one upsert per turn. Statements run sequentially in per-session event order: a turn's classification depends on the turns before it, and Postgres rejects one multi-row INSERT touching the same key twice. Only ConnectError is retried, per statement, because it proves - that statement never reached the database. A failed write stops its key and user - histories for this batch. Other users sharing that key can still advance their - independent user histories, with the key projection disabled and the real key - identity preserved. The failed turn is never replayed. Callers must not add their - own retry around this function. + that statement never reached the database. When a failure proves the database could + not take the write at all, the flush stops and returns that turn, the rest of its + session and every session not yet attempted, so the caller can queue them again; + turns already written are never returned. Any other failed write stops its key and + user histories for this batch and is never replayed. Other users sharing that key can + still advance their independent user histories, with the key projection disabled and + the real key identity preserved. Callers must not add their own retry around this + function. """ if not transactions: - return + return None ordered: Final = sorted( transactions, key=lambda transaction: (*_session_partition(transaction), transaction.turn_at), ) - for _, session_group in groupby( - ordered, - key=_session_partition, - ): - await _flush_session_partition(prisma_client, tuple(session_group), n_retry_times) + partitions: Final = tuple(tuple(group) for _, group in groupby(ordered, key=_session_partition)) + for position, partition in enumerate(partitions): + if (unwritten := await _flush_session_partition(prisma_client, partition, n_retry_times)) is not None: + return UnwrittenTurns( + (*unwritten.turns, *chain.from_iterable(partitions[position + 1 :])), + unwritten.error, + ) + return None diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 33d973cccfb..d9d0ece4468 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -776,8 +776,17 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool: def db_write_never_sent(e: Exception) -> bool: - """Whether the database could not take the write at all, because it was unreachable (P1001), - refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again.""" + """Whether the database could not take the write at all, because the client had no query engine to send it + to (ClientNotConnectedError, NotConnectedError, ConnectError), it was unreachable (P1001), refusing connections + (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again.""" + import prisma + import prisma.engine.errors + + unsent_types: Final = _exception_types( + prisma.errors.ClientNotConnectedError, prisma.engine.errors.NotConnectedError + ) + if isinstance(e, (*DB_RETRY_SAFE_ERROR_TYPES, *unsent_types)): + return True if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e): return True if PrismaDBExceptionHandler.is_database_capacity_error(e): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 262cf88f696..76a7046beed 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -337,7 +337,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.models.team import LiteLLM_TeamTableCachedObj - from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction + from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction, UnwrittenTurns from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord from litellm.proxy.db.model_usage_rollup import ModelUsageTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -5705,7 +5705,6 @@ class PrismaClient: _tool_usage_transactions_lock = asyncio.Lock() model_usage_transactions: ClassVar[list["ModelUsageTransaction"]] = [] _model_usage_transactions_lock = asyncio.Lock() - autorouter_turn_transactions: ClassVar[list["AutoRouterTurnTransaction"]] = [] _autorouter_turn_transactions_lock = asyncio.Lock() # How long a health probe failure waits for an in-flight planned engine @@ -5724,6 +5723,7 @@ class PrismaClient: ): ## init logging object self.baseline_accounting_transactions: list[BaselineAccountingRecord] = [] + self.autorouter_turn_transactions: list[AutoRouterTurnTransaction] = [] self.baseline_accounting_lock: Final = asyncio.Lock() self.proxy_logging_obj = proxy_logging_obj self.token_auth: DatabaseTokenAuth | None = resolve_database_token_auth() @@ -9220,25 +9220,27 @@ async def _run_spend_logs_job( await flush_baseline_accounting(prisma_client) - async with prisma_client._autorouter_turn_transactions_lock: + autorouter_queue_lock: Final = prisma_client._autorouter_turn_transactions_lock + async with autorouter_queue_lock: autorouter_turns_to_process: Final = prisma_client.autorouter_turn_transactions[:MAX_LOGS_PER_INTERVAL] remaining_autorouter_turns: Final = prisma_client.autorouter_turn_transactions[ len(autorouter_turns_to_process) : ] prisma_client.autorouter_turn_transactions = remaining_autorouter_turns # rebind-ok: drain under lock - try: - from litellm.proxy.db.autorouter_session_rollup import flush_autorouter_turn_transactions - - await flush_autorouter_turn_transactions( - prisma_client=prisma_client, - transactions=autorouter_turns_to_process, - ) - except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job - verbose_proxy_logger.error( - "Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s", - len(autorouter_turns_to_process), - autorouter_tracking_err, + unwritten_autorouter_turns: Final = await _flush_autorouter_turns(prisma_client, autorouter_turns_to_process) + if unwritten_autorouter_turns is not None: + async with autorouter_queue_lock: + prisma_client.autorouter_turn_transactions = [ # rebind-ok: requeue under lock + *unwritten_autorouter_turns.turns, + *prisma_client.autorouter_turn_transactions, + ] + verbose_proxy_logger.warning( + "Spend tracking - database unavailable during auto-router session rollup flush, " + "requeued %s unwritten turn transactions for the next flush: %s", + len(unwritten_autorouter_turns.turns), + unwritten_autorouter_turns.error, ) + raise unwritten_autorouter_turns.error try: from litellm.proxy.db.shadow_eval_funnel import flush_shadow_eval_funnel @@ -9248,6 +9250,22 @@ async def _run_spend_logs_job( verbose_proxy_logger.error("Spend tracking - shadow eval funnel drain failed: %s", funnel_err) +async def _flush_autorouter_turns( + prisma_client: PrismaClient, turns: "Sequence[AutoRouterTurnTransaction]" +) -> "UnwrittenTurns | None": + from litellm.proxy.db.autorouter_session_rollup import flush_autorouter_turn_transactions + + try: + return await flush_autorouter_turn_transactions(prisma_client=prisma_client, transactions=turns) + except Exception as autorouter_tracking_err: # noqa: BLE001 # a drain bug must not abort the spend job + verbose_proxy_logger.error( + "Spend tracking - auto-router session rollup drain failed; %s turn transactions dropped: %s", + len(turns), + autorouter_tracking_err, + ) + return None + + MAX_SPEND_LOG_DRAIN_ITERATIONS: Final = 20 diff --git a/tests/unit/proxy/db/test_autorouter_session_rollup.py b/tests/unit/proxy/db/test_autorouter_session_rollup.py index 4217feeb83b..6ac49738710 100644 --- a/tests/unit/proxy/db/test_autorouter_session_rollup.py +++ b/tests/unit/proxy/db/test_autorouter_session_rollup.py @@ -26,6 +26,13 @@ from litellm.proxy.db.autorouter_session_rollup import ( ) from tests.unit.proxy.db.fake_prisma_engine import engine_call + +def _client_not_connected() -> Exception: + from prisma.errors import ClientNotConnectedError + + return ClientNotConnectedError() + + ROUTING_DECISION = {"router_model_name": "live-auto", "router_type": "complexity", "routed_model": "haiku"} @@ -272,7 +279,7 @@ class TestBuildTransaction: class _FakeDB: def __init__( self, - failures: "list[Exception] | None" = None, + failures: "list[Exception | None] | None" = None, poison_session: str | None = None, poison_user: str | None = None, commit_then_error_users: frozenset[str] = frozenset(), @@ -290,8 +297,8 @@ class _FakeDB: raise RuntimeError("index row size exceeds btree maximum") if self._poison_user is not None and params[19] == self._poison_user: raise RuntimeError("index row size exceeds btree maximum") - if self._failures: - raise self._failures.pop(0) + if self._failures and (failure := self._failures.pop(0)) is not None: + raise failure self.calls.append((sql, params)) if params[19] in self._commit_then_error_users: raise RuntimeError("commit succeeded but acknowledgement was lost") @@ -301,7 +308,7 @@ class _FakeDB: class _FakeClient: def __init__( self, - failures: "list[Exception] | None" = None, + failures: "list[Exception | None] | None" = None, poison_session: str | None = None, poison_user: str | None = None, commit_then_error_users: frozenset[str] = frozenset(), @@ -411,7 +418,7 @@ class TestFlush: ("z-other", 10, "k2"), ) ) - asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns)))) + assert asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns)))) is None assert client.db.attempts == client.db.calls assert [ @@ -444,9 +451,38 @@ class TestFlush: _transaction(session_id="s1", at=datetime(2026, 8, 1, 12, 0, 10)), _transaction(session_id="s2", at=datetime(2026, 8, 1, 12, 0, 5)), ] - asyncio.run(flush_autorouter_turn_transactions(client, transactions)) + assert asyncio.run(flush_autorouter_turn_transactions(client, transactions)) is None assert [params[1] for _, params in client.db.calls] == ["s2"] + @pytest.mark.parametrize( + "error", + [_client_not_connected(), httpx.ConnectError("connection refused")], + ids=["client-not-connected", "connect-error-retries-exhausted"], + ) + def test_an_unsent_turn_is_returned_with_every_turn_after_it_and_never_double_counts( + self, error: Exception + ) -> None: + client: Final = _FakeClient(failures=[None, error]) + written: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 0)) + unsent: Final = _transaction(user_id="b-user", at=datetime(2026, 8, 1, 12, 0, 1)) + same_key_other_user: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 2)) + other_key: Final = _transaction(api_key="k2", user_id="c-user", at=datetime(2026, 8, 1, 12, 0, 3)) + + unwritten: Final = asyncio.run( + flush_autorouter_turn_transactions( + client, [other_key, same_key_other_user, unsent, written], n_retry_times=0 + ) + ) + + assert unwritten is not None + assert (unwritten.turns, unwritten.error) == ((unsent, same_key_other_user, other_key), error) + assert [params[5] for _, params in client.db.calls] == ["2026-08-01T12:00:00"] + + assert asyncio.run(flush_autorouter_turn_transactions(client, unwritten.turns)) is None + assert [(sql, params[0], params[5]) for sql, params in client.db.calls] == [ + (UPSERT_AUTOROUTER_SESSION_SQL, "k1", f"2026-08-01T12:00:0{second}") for second in range(3) + ] + [(UPSERT_AUTOROUTER_SESSION_SQL, "k2", "2026-08-01T12:00:03")] + def test_an_empty_batch_writes_nothing(self): client = _FakeClient() asyncio.run(flush_autorouter_turn_transactions(client, [])) diff --git a/tests/unit/proxy/db/test_exception_handler.py b/tests/unit/proxy/db/test_exception_handler.py index abc0850194d..2f61864fd35 100644 --- a/tests/unit/proxy/db/test_exception_handler.py +++ b/tests/unit/proxy/db/test_exception_handler.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest from prisma import errors as prisma_errors -from prisma.engine.errors import BinaryNotFoundError, EngineConnectionError, EngineRequestError +from prisma.engine.errors import BinaryNotFoundError, EngineConnectionError, EngineRequestError, NotConnectedError from prisma.errors import ( ClientNotConnectedError, DataError, @@ -1100,7 +1100,11 @@ def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Excep ), True, ), + (ClientNotConnectedError(), True), + (NotConnectedError("Not connected to the query engine"), True), + (httpx.ConnectError("connection refused"), True), (httpx.ReadTimeout("read timed out"), False), + (httpx.ReadError("connection reset"), False), (_localized_error("57014"), False), ], ) diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py index 293f2d49f43..271fb71d06a 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -21,6 +21,11 @@ from unittest.mock import AsyncMock, MagicMock import pytest from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY +from litellm.proxy.db.autorouter_session_rollup import ( + UPSERT_AUTOROUTER_SESSION_SQL, + UPSERT_AUTOROUTER_USER_SESSION_SQL, + AutoRouterTurnTransaction, +) from litellm.proxy.utils import ( MAX_SPEND_LOG_DRAIN_ITERATIONS, drain_spend_logs_queue, @@ -1001,6 +1006,12 @@ def _postgres_restarting() -> Exception: ) +def _client_not_connected() -> Exception: + from prisma.errors import ClientNotConnectedError + + return ClientNotConnectedError() + + def _postgres_unreachable() -> Exception: from prisma.errors import DataError @@ -1257,7 +1268,11 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out @pytest.mark.asyncio -@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"]) +@pytest.mark.parametrize( + "error", + [_postgres_restarting(), _postgres_unreachable(), _client_not_connected()], + ids=["57P03", "P1001", "client-not-connected"], +) async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes( error: Exception, mock_prisma_client: MagicMock, @@ -1284,6 +1299,72 @@ async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_write } == {"queue_after": [first, second], "autorouter_queue_after": 1} +def _autorouter_turn(second: int) -> AutoRouterTurnTransaction: + return AutoRouterTurnTransaction( + api_key="k1", + session_id="s1", + router_name="live-auto", + router_type="complexity", + model="bedrock/haiku", + turn_at=datetime(2026, 10, 2, 12, 0, second), + total_tokens=10, + spend=0.01, + saved_spend=0.0, + classifier_cost=0.0, + covered=False, + cache_hit=False, + cache_ttl_seconds=None, + cache_touched=True, + ) + + +@pytest.mark.asyncio +async def test_unsent_autorouter_turns_go_back_ahead_of_newer_turns_and_stop_the_job( + mock_prisma_client: MagicMock, +) -> None: + """A turn the client could not send is kept with the session turns after it, ahead of a turn + queued while the flush ran, and the turn already written is not queued again.""" + written, unsent, dependent, newer = (_autorouter_turn(second) for second in range(4)) + error: Final = _client_not_connected() + turn_statements: Final = frozenset({UPSERT_AUTOROUTER_SESSION_SQL, UPSERT_AUTOROUTER_USER_SESSION_SQL}) + + async def execute_raw(sql: str, *params: object) -> int: + if sql not in turn_statements: + return 0 + if params[5] == unsent.turn_at.isoformat(): + raise error + mock_prisma_client.autorouter_turn_transactions = [*mock_prisma_client.autorouter_turn_transactions, newer] + return 1 + + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=execute_raw) + mock_prisma_client._autorouter_turn_transactions_lock = asyncio.Lock() + mock_prisma_client.autorouter_turn_transactions = [dependent, unsent, written] + + with pytest.raises(type(error)) as raised: + await update_spend_logs_job( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + + assert { + "raised": raised.value, + "turn_writes": [ + (awaited.args[0], awaited.args[6]) + for awaited in mock_prisma_client.db.execute_raw.await_args_list + if awaited.args[0] in turn_statements + ], + "queue_after": mock_prisma_client.autorouter_turn_transactions, + } == { + "raised": error, + "turn_writes": [ + (UPSERT_AUTOROUTER_SESSION_SQL, written.turn_at.isoformat()), + (UPSERT_AUTOROUTER_SESSION_SQL, unsent.turn_at.isoformat()), + ], + "queue_after": [unsent, dependent, newer], + } + + @pytest.mark.asyncio async def test_rollups_drop_a_batch_whose_session_the_server_terminated( mock_prisma_client: MagicMock,