mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): requeue spend logs when a Postgres restart tears down the prisma engine mid-write (#47571)
Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Backport of #47571 to stable/1.106.x
Cherry-picked from ccd0d4ea49 (main)
This commit is contained in:
parent
6a6c52d240
commit
e8c33f3568
2 changed files with 55 additions and 1 deletions
|
|
@ -8239,7 +8239,11 @@ MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
|
|||
|
||||
|
||||
def _is_transient_spend_log_write_error(e: Exception) -> bool:
|
||||
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
return (
|
||||
PrismaDBExceptionHandler.is_database_transport_error(e)
|
||||
or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e)
|
||||
or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
)
|
||||
|
||||
|
||||
async def _create_spend_logs_with_poison_isolation(
|
||||
|
|
|
|||
|
|
@ -952,3 +952,53 @@ async def test_update_spend_logs_requeues_batch_when_postgres_is_out_of_connecti
|
|||
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1
|
||||
sleep.assert_not_awaited()
|
||||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r0", "r1", "r2", "r3", "e"]
|
||||
|
||||
|
||||
async def _create_many_while_postgres_is_restarting(**_: object) -> int:
|
||||
from prisma.engine.utils import handle_response_errors
|
||||
|
||||
handle_response_errors(
|
||||
MagicMock(),
|
||||
[
|
||||
{
|
||||
"error": "Error in connector: Can't reach database server at `127.0.0.1:5432`",
|
||||
"user_facing_error": {"error_code": "P1001", "message": "Can't reach database server", "meta": None},
|
||||
}
|
||||
],
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_retries_and_requeues_batch_when_postgres_restart_tears_down_prisma_engine(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""While Postgres restarts the prisma engine answers the write with a P1001
|
||||
payload whose ``meta`` is ``None``, which prisma's own response parser trips
|
||||
over before it can build a typed error. That outage must be retried like
|
||||
any other, and once retries exhaust the dequeued batch goes back to the head
|
||||
of the queue so the rows land when Postgres is back.
|
||||
"""
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
|
||||
side_effect=_create_many_while_postgres_is_restarting
|
||||
)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
|
||||
|
||||
with pytest.raises(AttributeError):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=2,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
|
||||
)
|
||||
|
||||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
|
||||
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue