mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(proxy): restore spend-log batch to queue when DB write fails
This commit is contained in:
parent
366ec6f487
commit
0136895d28
2 changed files with 92 additions and 8 deletions
|
|
@ -5378,13 +5378,18 @@ async def update_spend_logs_job(
|
|||
logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL]
|
||||
prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[len(logs_to_process) :]
|
||||
|
||||
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,
|
||||
)
|
||||
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,
|
||||
)
|
||||
except Exception:
|
||||
async with prisma_client._spend_log_transactions_lock:
|
||||
prisma_client.spend_log_transactions = logs_to_process + prisma_client.spend_log_transactions
|
||||
raise
|
||||
|
||||
# Guardrail/policy usage tracking (same batch, outside spend-logs update)
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.utils import ProxyUpdateSpend
|
||||
from litellm.proxy.utils import ProxyUpdateSpend, update_spend_logs_job
|
||||
|
||||
|
||||
class _AsyncCM:
|
||||
|
|
@ -337,6 +337,85 @@ async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood(
|
|||
assert attempts < n_rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_job_restores_batch_when_write_fails(
|
||||
mock_prisma_client: Any,
|
||||
make_spend_log_row: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A batch popped by ``update_spend_logs_job`` must return to the queue when
|
||||
the DB write fails for every attempt, so it is retried rather than dropped.
|
||||
|
||||
Reproduces issue #33873: the job removes logs from
|
||||
``spend_log_transactions`` before writing and, on the unfixed code, never
|
||||
restores them once ``update_spend_logs`` raises, permanently losing the rows.
|
||||
"""
|
||||
import httpx
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
|
||||
rows = [make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")]
|
||||
mock_prisma_client.spend_log_transactions = list(rows)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
|
||||
side_effect=httpx.ReadError("db down")
|
||||
)
|
||||
mock_prisma_client.proxy_logging_obj.failure_handler = AsyncMock()
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=mock_prisma_client.proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_job_restores_batch_ahead_of_new_logs(
|
||||
mock_prisma_client: Any,
|
||||
make_spend_log_row: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The restored batch is prepended so its rows keep their place ahead of
|
||||
logs enqueued while the failed write was in flight.
|
||||
"""
|
||||
import httpx
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
async def _fake_sleep(_: float) -> None:
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
|
||||
|
||||
popped = [make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")]
|
||||
newly_enqueued = make_spend_log_row(request_id="c")
|
||||
mock_prisma_client.spend_log_transactions = list(popped)
|
||||
|
||||
async def _fail_then_capture_new_log(*, data: Any, skip_duplicates: bool) -> None:
|
||||
if newly_enqueued not in mock_prisma_client.spend_log_transactions:
|
||||
mock_prisma_client.spend_log_transactions.append(newly_enqueued)
|
||||
raise httpx.ReadError("db down")
|
||||
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
|
||||
side_effect=_fail_then_capture_new_log
|
||||
)
|
||||
mock_prisma_client.proxy_logging_obj.failure_handler = AsyncMock()
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=mock_prisma_client.proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == popped + [newly_enqueued]
|
||||
|
||||
|
||||
def test_disable_spend_updates_reflects_general_settings(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue