scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(spend): bound each spend-log write statement by payload bytes (#34956)

Browse source 
The Prisma query engine is a separate Rust process whose resident memory is
a high-water mark: it grows with the payload of the largest single statement
it executes and glibc never returns that memory to the OS, so a pod's memory
floor ratchets up to its worst-ever write and stays there for the life of
the worker. Memory-based autoscaling then reads a number that reflects the
largest write the pod has ever done rather than what it is doing now.

The spend-log flush handed Prisma a fixed 1000 rows per create_many. With
store_prompts_in_spend_logs enabled a single row carries the full prompt and
response, so one statement can be tens of megabytes and permanently costs
hundreds of megabytes of RSS. Row counts cannot express that budget: the
same 1000 rows range from well under a megabyte to tens of megabytes.

Split each flush into statements bounded by encoded payload size
(SPEND_LOG_WRITE_BATCH_MAX_BYTES, default 2MB) on top of the existing
1000-row cap. What is measured is the encoded statement, so the budget
counts what actually goes on the wire: the JSON escaping of quotes and
newlines, multibyte characters at their encoded width, the field names and
separators a 25-column row carries, and the brackets and row separators the
rows carry as one collection. Deployments that do not store prompts keep one
statement per 1000 rows and are unaffected; prompt-carrying flushes get
several small statements instead of one huge one. A row larger than the
budget is still written on its own rather than dropped, and a row the
serializer refuses counts as zero rather than raising out of the flush and
dropping every row queued behind it.

Splitting a flush must not multiply what a poison-row flood costs, so the
poison-isolation allowance is threaded through every statement of a 1000-row
group instead of being handed out fresh per statement. That is only safe
because the allowance now counts failed inserts rather than every insert:
the one insert a statement needs when nothing is poisoned is not charged, so
a healthy flush never runs the allowance down however many statements it
splits into, and a statement reached after the allowance is spent is still
attempted so clean rows behind a flood still persist. Failed inserts for a
group are bounded by the allowance plus one baseline insert per statement,
which restores the constant-per-group ceiling the single-statement path had.

Resolves LIT-4765
This commit is contained in:
Yassin Kortam 2026-08-01 13:53:36 -07:00 • committed by GitHub
parent a8cc6a921a
commit 4178a857ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 416 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -1461,6 +1461,7 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS = float(
TOOL_SPEND_TOP_TOOLS = 100
SPEND_LOG_PARTITION_INTERVAL = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
SPEND_LOG_PARTITION_PRECREATE_AHEAD = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
SPEND_LOG_WRITE_BATCH_MAX_BYTES = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_BYTES", 2_000_000)))
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))

86
litellm/proxy/db/spend_log_batching.py Normal file
View file

@ -0,0 +1,86 @@
"""Split a spend-log flush into statements the Prisma query engine can afford.
The query engine is a separate Rust process whose resident memory is a
high-water mark: it grows with the payload of the largest single statement it
is asked to execute and glibc never returns that memory to the OS, so a pod's
memory floor ratchets up to its worst-ever write and stays there for the life
of the worker. Under ``store_prompts_in_spend_logs`` a single spend-log row
carries the full prompt and response, so a fixed 1000-row ``create_many``
hands the engine tens of megabytes in one statement and permanently costs
hundreds of megabytes of RSS, which is what makes memory-based autoscaling
read the wrong number.
Bounding each statement by payload size instead caps that floor. Row-count
batching alone cannot: the same 1000 rows range from well under a megabyte
(spend counters only) to tens of megabytes (prompts stored), and only the
byte budget tracks what the engine actually allocates.
"""
import json
from collections.abc import Iterator, Mapping, Sequence
SpendLogRow = Mapping[str, object]
_STATEMENT_FRAMING_BYTES = len(json.dumps([]))
_ROW_SEPARATOR_BYTES = len(json.dumps([0, 0])) - len(json.dumps([0])) - len(json.dumps(0))
def _row_payload_bytes(row: SpendLogRow) -> int:
"""Bytes this row contributes to the encoded write statement.
The whole row is serialized rather than its values summed, so the count
includes the field names, separators and braces the row carries on the
wire and not only its payload. Those are what make the difference between
a measurement and an estimate for a row of many small columns, where the
keys outweigh the values.
Serializing is also what makes the count a byte count. Character counts
under-measure a prompt in a non-Latin script by its bytes-per-character
factor, and even an all-ASCII prompt grows when JSON escapes its quotes,
backslashes and newlines (about 18% for a realistic stored prompt, and up
to double for escape-dense content). ``json.dumps`` escapes non-ASCII to
``\\uXXXX`` and defaults to ASCII output, so its length never under-states
the wire size. ``default=str`` covers the datetimes and other scalars a
row carries.
A row the serializer refuses (a self-reference is the reachable case)
counts as zero rather than raising: measuring a row must never be what
loses spend data, since raising here would propagate out of the flush and
drop every row queued behind it. Such a row is still written, it just does
not contribute to the budget.
"""
try:
return len(json.dumps(row, default=str))
except (TypeError, ValueError):
return 0
def spend_log_write_batches(
rows: Sequence[SpendLogRow],
max_bytes: int,
) -> Iterator[Sequence[SpendLogRow]]:
"""Yield consecutive slices of ``rows`` whose payload fits ``max_bytes``.
What is measured is the encoded slice, not the sum of its rows: rows become
one collection on the wire, so the brackets around them and the separator
between each pair count too. Summing rows alone under-states a slice by one
separator per row, which is negligible for prompt-carrying rows and is not
for a slice of many small ones, where the budget would be exceeded by the
row count. The two framing constants are derived from the serializer rather
than written down so they cannot drift from it.
Slices preserve input order and together cover every row exactly once. A
row larger than ``max_bytes`` on its own is yielded alone rather than
dropped: the budget is a memory guardrail, not an admission filter, and
losing spend data to protect RSS would be the worse failure.
"""
sizes = tuple(_row_payload_bytes(row) for row in rows)
start = 0
while start < len(rows):
end = start + 1
used = _STATEMENT_FRAMING_BYTES + sizes[start]
while end < len(rows) and used + _ROW_SEPARATOR_BYTES + sizes[end] <= max_bytes:
used += _ROW_SEPARATOR_BYTES + sizes[end]
end += 1
yield rows[start:end]
start = end

68
litellm/proxy/utils.py
View file

@ -38,6 +38,7 @@ from litellm.constants import (
DEFAULT_MODEL_CREATED_AT_TIME,
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
MAX_TEAM_LIST_LIMIT,
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
)
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
@ -135,6 +136,7 @@ from litellm.proxy.db.prisma_client import (
parse_iam_endpoint_from_url,
)
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
from litellm.proxy.db.spend_log_batching import spend_log_write_batches
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -5477,11 +5479,15 @@ 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 _create_spend_logs_with_poison_isolation(
SpendLogsRepository(prisma_client),
batch_with_dates,
MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH,
)
isolation_budget = MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH
for statement_rows in spend_log_write_batches(
batch_with_dates, SPEND_LOG_WRITE_BATCH_MAX_BYTES
):
isolation_budget = await _create_spend_logs_with_poison_isolation(
SpendLogsRepository(prisma_client),
statement_rows,
isolation_budget,
)
verbose_proxy_logger.debug(f"Flushed {len(batch)} logs to the DB.")
# Explicitly clear batch memory
del batch, batch_with_dates
@ -5764,13 +5770,13 @@ async def _monitor_spend_logs_queue(
await asyncio.sleep(current_interval)
MAX_SPEND_LOG_ISOLATION_ATTEMPTS_PER_BATCH = 256
MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH = 256
async def _create_spend_logs_with_poison_isolation(
repo: SpendLogsRepository,
rows: Sequence[Mapping[str, object]],
attempts_left: int,
failure_budget: int,
) -> int:
"""Write spend-log rows, isolating any row Postgres rejects on its data.
@ -5783,32 +5789,26 @@ async def _create_spend_logs_with_poison_isolation(
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.
``failure_budget`` caps the *failed* inserts the isolation may issue, which
is the work an authenticated caller flooding poisoned rows can amplify. The
one insert a statement needs when nothing is poisoned is not charged, so a
caller can thread a single budget through every statement of a flush and
bound the whole flush's failed inserts and log lines by the initial value,
without a large healthy flush ever running out and losing rows. When the
budget is spent the still-failing remainder is dropped wholesale (the
pre-existing drop-the-batch behavior) under one log line, and a statement
reached afterwards is still attempted, so clean rows behind a poison flood
persist. 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
return failure_budget
except Exception as e:
if not PrismaDBExceptionHandler.is_prisma_data_error(e):
raise
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
raise
budget_left = max(failure_budget - 1, 0)
if len(rows) == 1:
request_id = rows[0].get("request_id")
spend_log_error(
@ -5817,9 +5817,23 @@ async def _create_spend_logs_with_poison_isolation(
str(e),
exc=e,
)
return attempts_left - 1
return budget_left
if budget_left <= 0:
spend_log_error(
"Spend tracking - dropping %d spend log rows without per-row isolation; "
"isolation failure budget exhausted for this flush",
len(rows),
)
return 0
mid = len(rows) // 2
remaining = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], attempts_left - 1)
remaining = await _create_spend_logs_with_poison_isolation(repo, rows[:mid], budget_left)
if remaining <= 0:
spend_log_error(
"Spend tracking - dropping %d spend log rows without per-row isolation; "
"isolation failure budget exhausted for this flush",
len(rows) - mid,
)
return 0
return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining)

150
tests/test_litellm/proxy/db/test_spend_log_batching.py Normal file
View file

@ -0,0 +1,150 @@
"""Payload-size batching for spend-log writes.
The Prisma query engine's resident memory is a high-water mark set by the
largest single statement it executes, so these tests pin that no batch
exceeds the byte budget while every row is still written exactly once.
Symbols pinned here:
- ``spend_log_write_batches``
- ``_row_payload_bytes``
"""
import json
from typing import Any, Dict, List
from litellm.proxy.db.spend_log_batching import (
_row_payload_bytes,
spend_log_write_batches,
)
def _row(request_id: str, blob_bytes: int = 0) -> Dict[str, Any]:
return {
"request_id": request_id,
"spend": 0.01,
"total_tokens": 10,
"messages": json.dumps({"content": "x" * blob_bytes}) if blob_bytes else "{}",
}
def test_rows_are_split_when_the_payload_exceeds_the_budget() -> None:
rows = [_row(f"r{i}", blob_bytes=1000) for i in range(10)]
# The encoded size of a three-row statement, so three rows fit and four do not.
budget = len(json.dumps(rows[:3], default=str))
batches = list(spend_log_write_batches(rows, max_bytes=budget))
assert [len(batch) for batch in batches] == [3, 3, 3, 1]
assert all(len(json.dumps(list(batch), default=str)) <= budget for batch in batches)
def test_every_row_is_written_exactly_once_and_in_order() -> None:
rows = [_row(f"r{i}", blob_bytes=500) for i in range(37)]
flattened: List[Any] = [row for batch in spend_log_write_batches(rows, max_bytes=1700) for row in batch]
assert [row["request_id"] for row in flattened] == [row["request_id"] for row in rows]
def test_small_rows_stay_in_one_statement() -> None:
rows = [_row(f"r{i}") for i in range(1000)]
batches = list(spend_log_write_batches(rows, max_bytes=2_000_000))
assert [len(batch) for batch in batches] == [1000]
def test_a_row_larger_than_the_budget_is_written_alone_not_dropped() -> None:
rows = [_row("small"), _row("huge", blob_bytes=50_000), _row("small2")]
batches = list(spend_log_write_batches(rows, max_bytes=1000))
assert [[row["request_id"] for row in batch] for batch in batches] == [
["small"],
["huge"],
["small2"],
]
def test_field_names_and_separators_are_counted() -> None:
"""A spend-log row carries ~25 columns, so for rows of many small values
the field names and separators outweigh the values themselves. Counting
only the values would let such a statement run well past the budget."""
row = {f"column_with_a_long_name_{i}": "v" for i in range(25)}
value_bytes_only = sum(len(json.dumps(value)) for value in row.values())
assert _row_payload_bytes(row) > 4 * value_bytes_only
# Every row fits the budget counting values alone, and only three fit once
# the keys are counted, so the split is what proves they are counted.
budget = len(json.dumps([row] * 3, default=str))
assert [len(batch) for batch in spend_log_write_batches([row] * 6, max_bytes=budget)] == [3, 3]
def test_an_unserializable_value_does_not_break_the_flush() -> None:
"""Measuring a row must never be what loses spend data. A value the
serializer refuses (a self-referencing list is the reachable case) counts
as zero and the row is still written, rather than raising out of the
flush and dropping every row queued behind it."""
circular: List[Any] = []
circular.append(circular)
row = {"request_id": "r", "messages": circular}
assert _row_payload_bytes(row) == 0
assert [[r["request_id"] for r in batch] for batch in spend_log_write_batches([row], max_bytes=10)] == [["r"]]
def test_every_statement_fits_the_budget_when_encoded_whole() -> None:
"""The budget bounds the statement, not the sum of its rows. Rows become
one collection on the wire, so a slice also carries the brackets around it
and a separator between each pair; summing rows alone runs a slice of many
small rows over the budget by roughly its row count."""
rows = [_row(f"r{i}", blob_bytes=20 * (i % 7)) for i in range(400)]
# Exactly the encoded size of the first 40 rows, so a batcher that ignored
# the framing would fit 40 of them and overshoot by the 39 separators.
budget = len(json.dumps(rows[:40], default=str))
batches = [list(batch) for batch in spend_log_write_batches(rows, max_bytes=budget)]
encoded = [len(json.dumps(batch, default=str)) for batch in batches]
assert len(batches) > 1
assert max(len(batch) for batch in batches) > 1
assert max(encoded) <= budget
def test_empty_input_yields_no_statements() -> None:
assert list(spend_log_write_batches([], max_bytes=1000)) == []
def test_non_ascii_payloads_are_measured_in_bytes_not_characters() -> None:
"""A prompt in a non-Latin script encodes to several bytes per character,
so counting characters would let a statement carry a multiple of the
budget, which is the whole thing the budget exists to prevent."""
row = {"request_id": "r", "messages": "你好" * 2000}
characters = len(row["messages"])
assert _row_payload_bytes(row) >= len(row["messages"].encode("utf-8"))
budget = characters + 1000 # comfortably over the character count, under the encoded size
assert [len(batch) for batch in spend_log_write_batches([row, row], max_bytes=budget)] == [1, 1]
def test_json_escaping_growth_is_counted() -> None:
"""An all-ASCII prompt still grows when JSON escapes its quotes and
newlines, so a raw character count would let a statement exceed the
budget by that expansion factor."""
row = {"request_id": "r", "messages": '"quoted"\n' * 1000}
characters = len(row["messages"])
assert _row_payload_bytes(row) > characters
# Both rows fit the budget when counted as raw characters, and do not once
# the escaping is counted, so the split is what proves the escaping is measured.
budget = 2 * characters + 200
assert [len(batch) for batch in spend_log_write_batches([row, row], max_bytes=budget)] == [1, 1]
def test_unserialized_list_payloads_are_measured_not_ignored() -> None:
"""``jsonify_object`` only stringifies dicts, so a list-valued ``messages``
reaches the batcher raw; counting it as zero would let the largest rows
bypass the budget entirely."""
row = {"request_id": "r", "messages": [{"content": "x" * 5000}]}
assert _row_payload_bytes(row) > 5000
assert [len(batch) for batch in spend_log_write_batches([row, row], max_bytes=5100)] == [1, 1]

143
tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py
View file

@ -9,11 +9,13 @@ Symbols pinned here:
from __future__ import annotations
import asyncio
import json
from typing import Any, Dict, List
from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm.proxy.utils as utils_mod
from litellm.proxy.utils import ProxyUpdateSpend
@ -168,6 +170,40 @@ async def test_update_spend_logs_writes_batches_via_create_many(
}
@pytest.mark.asyncio
async def test_update_spend_logs_bounds_each_statement_by_payload_bytes(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A flush of prompt-carrying rows must reach Prisma as several small
statements rather than one huge one: the query engine's resident memory is
a high-water mark set by the largest statement it ever executes, and it
never returns that memory to the OS."""
monkeypatch.setattr(utils_mod, "SPEND_LOG_WRITE_BATCH_MAX_BYTES", 50_000)
blob = json.dumps({"content": "x" * 10_000})
logs = [make_spend_log_row(request_id=f"r{i}", messages=blob, response=blob) for i in range(50)]
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock()
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
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,
)
calls = mock_prisma_client.db.litellm_spendlogs.create_many.await_args_list
written = [row["request_id"] for call in calls for row in call.kwargs["data"]]
# Each statement is encoded whole rather than summed row by row, so the
# assertion covers the collection framing the rows carry on the wire and
# not only the payload the batcher adds up.
largest_statement_bytes = max(len(json.dumps(list(call.kwargs["data"]), default=str)) for call in calls)
assert written == [f"r{i}" for i in range(50)]
assert [len(call.kwargs["data"]) for call in calls] == [2] * 25
assert largest_statement_bytes <= 50_000
@pytest.mark.asyncio
async def test_update_spend_logs_uses_spend_logs_url_when_set(
mock_prisma_client: Any,
@ -327,14 +363,14 @@ 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.
unbounded failed inserts. The per-batch failure budget hard-caps the number
of failed ``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
attempt_cap = utils_mod.MAX_SPEND_LOG_ISOLATION_FAILURES_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
@ -360,6 +396,103 @@ async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood(
assert attempts < n_rows
async def _flush_and_count_create_many(
mock_prisma_client: Any,
logs: List[Any],
max_bytes: int,
poison: bool,
monkeypatch: pytest.MonkeyPatch,
) -> int:
"""Run one flush and return how many ``create_many`` calls it issued."""
async def _create_many(*, data: Any, skip_duplicates: bool) -> None:
if poison:
raise _data_error("invalid byte sequence for encoding UTF8: 0x00")
monkeypatch.setattr(utils_mod, "SPEND_LOG_WRITE_BATCH_MAX_BYTES", max_bytes)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many)
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
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,
)
return int(mock_prisma_client.db.litellm_spendlogs.create_many.await_count)
@pytest.mark.asyncio
async def test_poison_flood_cost_does_not_grow_with_the_number_of_statements(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Splitting a flush by payload bytes must not multiply what a poison flood
costs. The same poisoned rows are flushed as one statement and as several;
the split run may only pay the one unavoidable insert per extra statement,
not a fresh isolation budget each time.
"""
blob = json.dumps({"content": "x" * 2_000})
logs = [make_spend_log_row(request_id=f"r{i}", messages=blob, response=blob) for i in range(300)]
split_bytes = 250_000
statements = await _flush_and_count_create_many(
mock_prisma_client, logs, max_bytes=split_bytes, poison=False, monkeypatch=monkeypatch
)
split_attempts = await _flush_and_count_create_many(
mock_prisma_client, logs, max_bytes=split_bytes, poison=True, monkeypatch=monkeypatch
)
single_attempts = await _flush_and_count_create_many(
mock_prisma_client, logs, max_bytes=1_000_000_000, poison=True, monkeypatch=monkeypatch
)
# A clean flush issues exactly one call per statement, so this is the split
# count; asserting it is >1 is what proves the split was really exercised.
assert statements > 1
assert single_attempts == utils_mod.MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH
assert split_attempts <= single_attempts + statements
@pytest.mark.asyncio
async def test_clean_statement_is_still_written_after_a_poison_flood(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Sharing the isolation budget must not let a poison flood in one
statement silently drop the clean statements behind it. The budget bounds
failed inserts only, so a later statement is still attempted and its rows
persist.
"""
blob = json.dumps({"content": "x" * 2_000})
poisoned = [make_spend_log_row(request_id=f"bad{i}", messages=blob, response=blob) for i in range(300)]
# Larger than the budget, so it is always yielded as its own statement and
# the assertion cannot pass by riding along with a poisoned one.
lone_blob = json.dumps({"content": "x" * 300_000})
clean = [make_spend_log_row(request_id="good", messages=lone_blob, response=lone_blob)]
written: List[str] = []
async def _create_many(*, data: Any, skip_duplicates: bool) -> None:
ids = [row["request_id"] for row in data]
if any(request_id.startswith("bad") for request_id in ids):
raise _data_error("invalid byte sequence for encoding UTF8: 0x00")
written.extend(ids)
monkeypatch.setattr(utils_mod, "SPEND_LOG_WRITE_BATCH_MAX_BYTES", 250_000)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many)
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
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=poisoned + clean,
)
assert written == ["good"]
def test_disable_spend_updates_reflects_general_settings(
monkeypatch: pytest.MonkeyPatch,
) -> None: