mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(proxy): isolate poison spend-log rows so one bad record can't drop the whole batch (#31705)
update_spend_logs flushes the queue with a single create_many per batch, so one row carrying bytes Postgres refuses (a residual NUL byte is the canonical case) fails the entire insert and drops every good spend log alongside it. PR #29515 strips NUL bytes from the JSON columns, but the scalar string columns (end_user, model, session_id, ...) still flow through unsanitized, so a poisoned row can still reach the write and take a batch of up to 1000 good rows down with it. On a genuine data-layer rejection the batch is now bisected so the good rows still persist and only the offending row is dropped and logged with its request_id. The classification lives in PrismaDBExceptionHandler.is_prisma_data_error (matched by exact type so systemic subclasses like a missing table are not mistaken for a single poison row), which keeps prisma an in-function import and litellm.proxy.utils importable without the proxy extra. Transport failures, including the "can't reach database server" outage that prisma mislabels as a DataError, are re-raised unchanged so the existing connection-retry path still runs and a transient outage never turns into silent per-row data loss. The bisection carries a per-batch isolation budget so an authenticated caller flooding poisoned rows cannot amplify one failed bulk insert into ~2N failed inserts and N log lines; once the budget is spent the still-failing remainder is dropped wholesale under a single log line. Resolves LIT-4103
This commit is contained in:
parent
6d828e5759
commit
52dc15adfe
4 changed files with 225 additions and 2 deletions
|
|
@ -66,6 +66,29 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def is_prisma_data_error(e: Exception) -> bool:
|
||||
"""True iff ``e`` is a base prisma ``DataError``: the database processed
|
||||
the statement and refused the data itself (e.g. ``invalid byte sequence
|
||||
for encoding "UTF8": 0x00``), as opposed to a connectivity failure.
|
||||
|
||||
Matched by exact type, not ``isinstance``: the specific data-layer
|
||||
subclasses (``UniqueViolationError``, ``TableNotFoundError``,
|
||||
``MissingRequiredValueError`` ...) all derive from ``DataError`` but
|
||||
carry their own semantics, and a systemic one like a missing table must
|
||||
not be mistaken for a single poison row and bisected away. A raw
|
||||
Postgres execution error with no prisma P-code surfaces as the base
|
||||
``DataError``.
|
||||
|
||||
prisma also wraps the P1001 "can't reach database server" outage as a
|
||||
base ``DataError``, so a caller that must not treat an outage as a
|
||||
per-row data rejection has to additionally consult
|
||||
``is_database_service_unavailable_error`` before acting on a True here.
|
||||
"""
|
||||
import prisma
|
||||
|
||||
return type(e) is prisma.errors.DataError
|
||||
|
||||
@staticmethod
|
||||
def is_database_transport_error(e: Exception) -> bool:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -23,7 +23,9 @@ from typing import (
|
|||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Mapping,
|
||||
Optional,
|
||||
Sequence,
|
||||
Tuple,
|
||||
Union,
|
||||
cast,
|
||||
|
|
@ -5194,8 +5196,10 @@ class ProxyUpdateSpend:
|
|||
for j in range(0, len(logs_to_process), BATCH_SIZE):
|
||||
batch = logs_to_process[j : j + BATCH_SIZE]
|
||||
batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in batch]
|
||||
await SpendLogsRepository(prisma_client).table.create_many(
|
||||
data=batch_with_dates, skip_duplicates=True
|
||||
await _create_spend_logs_with_poison_isolation(
|
||||
SpendLogsRepository(prisma_client),
|
||||
batch_with_dates,
|
||||
MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH,
|
||||
)
|
||||
verbose_proxy_logger.debug(f"Flushed {len(batch)} logs to the DB.")
|
||||
# Explicitly clear batch memory
|
||||
|
|
@ -5462,6 +5466,65 @@ async def _monitor_spend_logs_queue(
|
|||
await asyncio.sleep(current_interval)
|
||||
|
||||
|
||||
MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH = 256
|
||||
|
||||
|
||||
async def _create_spend_logs_with_poison_isolation(
|
||||
repo: SpendLogsRepository,
|
||||
rows: Sequence[Mapping[str, object]],
|
||||
attempts_left: int,
|
||||
) -> int:
|
||||
"""Write spend-log rows, isolating any row Postgres rejects on its data.
|
||||
|
||||
``create_many`` writes the whole batch in a single statement, so one row
|
||||
carrying bytes Postgres refuses (a residual NUL byte is the canonical case)
|
||||
fails the entire insert and drops every good row alongside it. On a genuine
|
||||
data-layer rejection the batch is bisected so the good rows still persist
|
||||
and only the offending row is dropped and logged. Transport failures,
|
||||
including the "can't reach database server" outage that prisma mislabels as
|
||||
a ``DataError``, are re-raised unchanged so the caller's connection-retry
|
||||
path still runs.
|
||||
|
||||
``attempts_left`` is a hard ceiling on the number of ``create_many`` calls
|
||||
the isolation may issue for this batch, so an authenticated caller flooding
|
||||
poisoned rows cannot amplify one failed bulk insert into unbounded failed
|
||||
inserts and log lines. It is checked before any insert (so an exhausted
|
||||
budget never even attempts a write), decremented once per ``create_many``
|
||||
call, and threaded through the recursion so the whole bisection shares one
|
||||
allowance; total inserts are therefore bounded by the initial value
|
||||
regardless of how many rows are poisoned. When it runs out the still-failing
|
||||
remainder is dropped wholesale (the pre-existing drop-the-batch behavior)
|
||||
under one log line. Returns the budget left after this subtree.
|
||||
"""
|
||||
if attempts_left <= 0:
|
||||
spend_log_error(
|
||||
"Spend tracking - dropping %d spend log rows without per-row isolation; "
|
||||
"isolation attempt budget exhausted for this flush",
|
||||
len(rows),
|
||||
)
|
||||
return 0
|
||||
try:
|
||||
await repo.table.create_many(data=rows, skip_duplicates=True)
|
||||
return attempts_left - 1
|
||||
except Exception as e:
|
||||
if not PrismaDBExceptionHandler.is_prisma_data_error(e):
|
||||
raise
|
||||
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
|
||||
raise
|
||||
if len(rows) == 1:
|
||||
request_id = rows[0].get("request_id")
|
||||
spend_log_error(
|
||||
"Spend tracking - dropping spend log row Postgres rejected. request_id=%s error=%s",
|
||||
request_id,
|
||||
str(e),
|
||||
exc=e,
|
||||
)
|
||||
return attempts_left - 1
|
||||
mid = len(rows) // 2
|
||||
remaining = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], attempts_left - 1)
|
||||
return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining)
|
||||
|
||||
|
||||
def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_logging_obj: ProxyLogging):
|
||||
"""
|
||||
Raise an exception for failed update spend logs
|
||||
|
|
|
|||
|
|
@ -148,6 +148,37 @@ def test_is_database_service_unavailable_error_prisma_p1001_masquerades_as_datae
|
|||
)
|
||||
|
||||
|
||||
def test_is_prisma_data_error_only_true_for_dataerror():
|
||||
"""The spend-log poison-row isolation gates on this: only a prisma
|
||||
``DataError`` (the DB refused the data, e.g. a NUL byte) may be bisected
|
||||
into a per-row drop. A connectivity failure or any non-prisma exception
|
||||
must not be treated as a data rejection, so the whole batch surfaces."""
|
||||
import httpx
|
||||
|
||||
data_error = DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}})
|
||||
assert PrismaDBExceptionHandler.is_prisma_data_error(data_error) is True
|
||||
|
||||
for non_data in (
|
||||
httpx.ConnectError("conn refused"),
|
||||
PrismaError("can't reach database server"),
|
||||
UniqueViolationError(data={"user_facing_error": {"meta": {"table": "t"}}}),
|
||||
RuntimeError("boom"),
|
||||
):
|
||||
assert PrismaDBExceptionHandler.is_prisma_data_error(non_data) is False
|
||||
|
||||
|
||||
def test_is_prisma_data_error_true_for_connection_masquerade_dataerror():
|
||||
"""The P1001 outage prisma mislabels as a ``DataError`` is still a
|
||||
``DataError`` by type, so this returns True; the spend-log helper relies on
|
||||
``is_database_service_unavailable_error`` (not this check) to keep that
|
||||
outage on the retry path instead of dropping rows."""
|
||||
p1001_as_dataerror = DataError(
|
||||
data={"user_facing_error": {"message": "Can't reach database server at `127.0.0.1`:`5499`"}}
|
||||
)
|
||||
assert PrismaDBExceptionHandler.is_prisma_data_error(p1001_as_dataerror) is True
|
||||
assert PrismaDBExceptionHandler.is_database_service_unavailable_error(p1001_as_dataerror) is True
|
||||
|
||||
|
||||
def test_is_database_service_unavailable_error_cached_plan_escapes_as_503():
|
||||
"""Composes with the cached-plan retry: when that recovery fails and the
|
||||
Postgres "cached plan must not change result type" error escapes (raised by
|
||||
|
|
|
|||
|
|
@ -231,6 +231,112 @@ async def test_update_spend_logs_failure_raises_after_retries(
|
|||
)
|
||||
|
||||
|
||||
def _data_error(message: str) -> Any:
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError({"user_facing_error": {"message": message}})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_isolates_poison_row_and_persists_good_rows(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""One row Postgres rejects (22P05) must not drop the whole batch.
|
||||
|
||||
The good rows still persist and only the offending row is dropped, with no
|
||||
exception bubbling up. On the unfixed single-shot ``create_many`` the first
|
||||
write raises and the entire batch is lost.
|
||||
"""
|
||||
poison_id = "r1"
|
||||
written: List[str] = []
|
||||
|
||||
async def _create_many(*, data: Any, skip_duplicates: bool) -> None:
|
||||
ids = [row["request_id"] for row in data]
|
||||
if poison_id in ids:
|
||||
raise _data_error(
|
||||
"Inconsistent column data: 22P05 invalid byte sequence for encoding UTF8: 0x00"
|
||||
)
|
||||
written.extend(ids)
|
||||
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
|
||||
logs = [make_spend_log_row(request_id=f"r{i}") for i in range(4)]
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=0,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
assert sorted(written) == ["r0", "r2", "r3"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_reraises_connection_masquerade_dataerror(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""A P1001 "can't reach database server" outage that prisma mislabels as a
|
||||
``DataError`` is transient, not a poison row: it must propagate so the batch
|
||||
is surfaced/retried rather than bisected into silent per-row drops.
|
||||
"""
|
||||
err = _data_error("Can't reach database server at db-host:5432")
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
|
||||
with pytest.raises(type(err)):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=0,
|
||||
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"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""A flood of poisoned rows must not amplify one failed bulk insert into
|
||||
unbounded failed inserts. The per-batch attempt budget hard-caps the number
|
||||
of ``create_many`` calls regardless of how many rows are poisoned, so the DB
|
||||
work stays bounded and well below the input row count, and the helper still
|
||||
completes without raising.
|
||||
"""
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
attempt_cap = utils_mod.MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH
|
||||
# single create_many batch (< BATCH_SIZE) whose row count exceeds the attempt
|
||||
# cap, so the bound bites and attempts stay below the input row count
|
||||
n_rows = attempt_cap * 3
|
||||
|
||||
async def _always_poison(*, data: Any, skip_duplicates: bool) -> None:
|
||||
raise _data_error("invalid byte sequence for encoding UTF8: 0x00")
|
||||
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_always_poison)
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
logs = [make_spend_log_row(request_id=f"r{i}") for i in range(n_rows)]
|
||||
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=0,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
|
||||
attempts = mock_prisma_client.db.litellm_spendlogs.create_many.await_count
|
||||
assert attempts <= attempt_cap
|
||||
assert attempts < n_rows
|
||||
|
||||
|
||||
def test_disable_spend_updates_reflects_general_settings(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue