fix(spend): bound each spend-log write statement by row count as well as bytes (#37758)

The Prisma query engine is a separate process whose resident memory grows with
what it is asked to hold and glibc never returns it, so a pod's memory floor
ratchets up to its worst statement and stays there for the life of the worker.
#34956 bounded a spend-log flush by payload bytes, which caps that floor when
prompts are stored and does nothing when they are not: rows carrying only
attribution metadata run about 1.2 KB, so a 1000-row statement is roughly
1.2 MB, the 2 MB byte budget never binds, and every statement stays at 1000
rows forever.

The engine charges per row as well as per byte. Measured on a container running
the same engine build (5.4.2) against real Postgres, with rows shaped like a
store_prompts_in_spend_logs=false deployment, writing the same 200,000 rows:

  rows/statement   engine RSS still resident after the flush
  1000             179 MB
  500               91 MB
  250               41 MB
  100               19 MB

None of those statements came near the byte budget, so the whole difference is
row count. The floor is a plateau rather than a leak: 1,000,000 rows written at
1000 per statement settles around 229 MB and stops climbing.

Adds SPEND_LOG_WRITE_BATCH_MAX_ROWS, default 100, applied alongside the
existing byte budget so whichever binds first splits the statement. Both are
needed, since bytes are what track a prompt-carrying row and rows are what
track the engine's per-row bookkeeping.

One consequence worth naming: a flush now issues more statements, and a
statement that fails under a poison flood costs one insert before any
isolation runs, so the irreducible floor rises by the statement count. The
isolation budget still caps the amplification on top of that, and the tests
assert the bound derived from the configured row cap rather than a constant.
This commit is contained in:
Yassin Kortam 2026-08-21 09:49:51 -07:00 committed by GitHub
parent 7da34e8aed
commit 40b8300ac2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 169 additions and 59 deletions

View file

@ -1534,6 +1534,7 @@ TOOL_SPEND_TOP_TOOLS: Final = 100
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
SPEND_LOG_WRITE_BATCH_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_BYTES", 2_000_000)))
SPEND_LOG_WRITE_BATCH_MAX_ROWS: Final = max(1, int(os.getenv("SPEND_LOG_WRITE_BATCH_MAX_ROWS", "100")))
SPEND_LOG_QUEUE_SIZE_THRESHOLD: Final = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
SPEND_LOG_QUEUE_MAX_BYTES: Final = max(1, int(os.getenv("SPEND_LOG_QUEUE_MAX_BYTES", "64000000")))
SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))

View file

@ -10,10 +10,15 @@ 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.
Bounding each statement caps that floor, and it takes two budgets because the
engine charges for both terms. A byte budget is what tracks a prompt-carrying
row, whose size swings by orders of magnitude, and a row budget is what tracks
the engine's per-row bookkeeping, which a byte budget cannot see: rows holding
attribution metadata only stay far under any useful byte budget, so it never
binds and every statement runs at the caller's row cap. Measured on such a
flush, the same 100,000 rows cost 151 MB of permanently resident engine RSS at
1000 rows per statement against 25 MB at 100, with no statement anywhere near
a 2 MB byte budget.
"""
import json
@ -99,16 +104,28 @@ def spend_log_queue_within_budget(
def spend_log_write_batches(
rows: Sequence[SpendLogRow],
max_bytes: int,
max_rows: int,
) -> Iterator[Sequence[SpendLogRow]]:
"""Yield consecutive slices of ``rows`` whose payload fits ``max_bytes``.
"""Yield consecutive slices of ``rows`` within both ``max_bytes`` and ``max_rows``.
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.
What is measured for the byte budget 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.
Both budgets are needed because the engine's cost has two terms. Payload
bytes dominate when prompts are stored, and per-row bookkeeping dominates
when they are not: a slice of narrow rows costs the engine far more than
its bytes suggest, so a byte budget alone never binds on a deployment whose
rows carry no prompts and every statement stays at the caller's row cap.
Measured on a spend-log flush of rows carrying attribution metadata only,
writing the same 100,000 rows at 1000 rows per statement left 151 MB of
engine RSS resident against 25 MB at 100, with neither reaching a 2 MB byte
budget.
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
@ -120,7 +137,7 @@ def spend_log_write_batches(
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:
while end < len(rows) and end - start < max_rows and used + _ROW_SEPARATOR_BYTES + sizes[end] <= max_bytes:
used += _ROW_SEPARATOR_BYTES + sizes[end]
end += 1
yield rows[start:end]

View file

@ -28,6 +28,7 @@ from litellm.constants import (
MAX_TEAM_LIST_LIMIT,
SPEND_LOG_QUEUE_MAX_BYTES,
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
)
from litellm.proxy._types import (
CommonProxyErrors,
@ -6048,7 +6049,9 @@ class ProxyUpdateSpend:
batch_with_dates = [prisma_client.jsonify_object({**entry}) for entry in 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
batch_with_dates,
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
):
isolation_budget = await _create_spend_logs_with_poison_isolation(
SpendLogsRepository(prisma_client),

View file

@ -15,8 +15,20 @@ from unittest.mock import MagicMock, patch, AsyncMock
import httpx
import math
from litellm.constants import SPEND_LOG_WRITE_BATCH_MAX_ROWS
from litellm.proxy.utils import update_spend
# The flush chunks the queue by BATCH_SIZE and then splits each chunk by the row
# budget, so statement counts below are derived from both rather than hardcoded.
_OUTER_BATCH_SIZE = 1000
def _statements_for(rows: int) -> int:
full, remainder = divmod(rows, _OUTER_BATCH_SIZE)
chunks = [_OUTER_BATCH_SIZE] * full + ([remainder] if remainder else [])
return sum(math.ceil(chunk / SPEND_LOG_WRITE_BATCH_MAX_ROWS) for chunk in chunks)
class MockPrismaClient:
def __init__(self):
@ -242,25 +254,16 @@ async def test_update_spend_logs_multiple_batches_success():
await update_spend(prisma_client, None, proxy_logging_obj)
# Verify
assert create_many_mock.call_count == 2 # Should have made 2 batch calls
assert create_many_mock.call_count == _statements_for(1500)
# Get the actual data from each batch call
first_batch = create_many_mock.call_args_list[0][1]["data"]
second_batch = create_many_mock.call_args_list[1][1]["data"]
# No statement may exceed the row budget, which is what bounds the query
# engine's resident memory.
batches = [call[1]["data"] for call in create_many_mock.call_args_list]
assert all(len(batch) <= SPEND_LOG_WRITE_BATCH_MAX_ROWS for batch in batches)
# Verify batch sizes
assert len(first_batch) == 1000
assert len(second_batch) == 500
# Verify exact IDs in each batch
expected_first_batch_ids = {str(i) for i in range(1000)}
expected_second_batch_ids = {str(i) for i in range(1000, 1500)}
actual_first_batch_ids = {item["id"] for item in first_batch}
actual_second_batch_ids = {item["id"] for item in second_batch}
assert actual_first_batch_ids == expected_first_batch_ids
assert actual_second_batch_ids == expected_second_batch_ids
# Every row is written exactly once and in order, whatever the split.
written_ids = [item["id"] for batch in batches for item in batch]
assert written_ids == [str(i) for i in range(1500)]
# Verify all logs were processed
assert len(prisma_client.spend_log_transactions) == 0
@ -298,8 +301,9 @@ async def test_update_spend_logs_multiple_batches_with_failure():
# Execute
await update_spend(prisma_client, None, proxy_logging_obj)
# Verify
assert create_many_mock.call_count == 6 # 4 batches + 2 retries for failed batch
# The first attempt aborts on its second statement, then the whole flush
# replays, so the total is those two calls plus one complete pass.
assert create_many_mock.call_count == 2 + _statements_for(4000)
# Verify all batches were processed
all_processed_logs = []

View file

@ -20,6 +20,9 @@ from litellm.proxy.db.spend_log_batching import (
)
_ROWS_UNBOUNDED = 10_000
def _row(request_id: str, blob_bytes: int = 0) -> Dict[str, Any]:
return {
"request_id": request_id,
@ -33,7 +36,7 @@ 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))
batches = list(spend_log_write_batches(rows, max_bytes=budget, max_rows=_ROWS_UNBOUNDED))
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)
@ -41,21 +44,63 @@ def test_rows_are_split_when_the_payload_exceeds_the_budget() -> None:
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]
flattened: List[Any] = [
row for batch in spend_log_write_batches(rows, max_bytes=1700, max_rows=_ROWS_UNBOUNDED) 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:
def test_small_rows_are_not_split_by_the_byte_budget() -> None:
rows = [_row(f"r{i}") for i in range(1000)]
batches = list(spend_log_write_batches(rows, max_bytes=2_000_000))
batches = list(spend_log_write_batches(rows, max_bytes=2_000_000, max_rows=_ROWS_UNBOUNDED))
assert [len(batch) for batch in batches] == [1000]
def test_the_row_budget_splits_a_statement_the_byte_budget_never_would() -> None:
"""Rows carrying no prompts stay far under any useful byte budget, so the
byte budget never binds and every statement would otherwise run at the
caller's row cap."""
rows = [_row(f"r{i}") for i in range(1000)]
generous_bytes = 100 * len(json.dumps(rows, default=str))
batches = list(spend_log_write_batches(rows, max_bytes=generous_bytes, max_rows=100))
assert [len(batch) for batch in batches] == [100] * 10
# Without this, a batcher bounded only by bytes would still pass the line above.
assert max(len(json.dumps(list(batch), default=str)) for batch in batches) < generous_bytes / 10
def test_whichever_budget_binds_first_is_the_one_that_splits() -> None:
"""Fat rows are bounded by bytes and narrow rows by count, so neither
budget can be dropped in favour of the other."""
fat = [_row(f"f{i}", blob_bytes=1000) for i in range(10)]
narrow = [_row(f"n{i}") for i in range(10)]
two_fat_rows = len(json.dumps(fat[:2], default=str))
assert [len(b) for b in spend_log_write_batches(fat, max_bytes=two_fat_rows, max_rows=5)] == [2] * 5
assert [len(b) for b in spend_log_write_batches(narrow, max_bytes=two_fat_rows, max_rows=5)] == [5, 5]
def test_the_row_budget_still_writes_every_row_exactly_once_and_in_order() -> None:
rows = [_row(f"r{i}") for i in range(37)]
flattened: List[Any] = [
row for batch in spend_log_write_batches(rows, max_bytes=2_000_000, max_rows=10) for row in batch
]
assert [row["request_id"] for row in flattened] == [row["request_id"] for row in rows]
def test_a_row_budget_of_one_yields_one_statement_per_row() -> None:
rows = [_row(f"r{i}") for i in range(4)]
assert [len(b) for b in spend_log_write_batches(rows, max_bytes=2_000_000, max_rows=1)] == [1, 1, 1, 1]
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))
batches = list(spend_log_write_batches(rows, max_bytes=1000, max_rows=_ROWS_UNBOUNDED))
assert [[row["request_id"] for row in batch] for batch in batches] == [
["small"],
@ -76,7 +121,10 @@ def test_field_names_and_separators_are_counted() -> None:
# 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]
assert [len(batch) for batch in spend_log_write_batches([row] * 6, max_bytes=budget, max_rows=_ROWS_UNBOUNDED)] == [
3,
3,
]
def test_an_unserializable_value_does_not_break_the_flush() -> None:
@ -89,7 +137,10 @@ def test_an_unserializable_value_does_not_break_the_flush() -> None:
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"]]
assert [
[r["request_id"] for r in batch]
for batch in spend_log_write_batches([row], max_bytes=10, max_rows=_ROWS_UNBOUNDED)
] == [["r"]]
def test_every_statement_fits_the_budget_when_encoded_whole() -> None:
@ -102,7 +153,7 @@ def test_every_statement_fits_the_budget_when_encoded_whole() -> None:
# 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)]
batches = [list(batch) for batch in spend_log_write_batches(rows, max_bytes=budget, max_rows=_ROWS_UNBOUNDED)]
encoded = [len(json.dumps(batch, default=str)) for batch in batches]
assert len(batches) > 1
@ -111,7 +162,7 @@ def test_every_statement_fits_the_budget_when_encoded_whole() -> None:
def test_empty_input_yields_no_statements() -> None:
assert list(spend_log_write_batches([], max_bytes=1000)) == []
assert list(spend_log_write_batches([], max_bytes=1000, max_rows=_ROWS_UNBOUNDED)) == []
def test_non_ascii_payloads_are_measured_in_bytes_not_characters() -> None:
@ -124,7 +175,9 @@ def test_non_ascii_payloads_are_measured_in_bytes_not_characters() -> None:
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]
assert [
len(batch) for batch in spend_log_write_batches([row, row], max_bytes=budget, max_rows=_ROWS_UNBOUNDED)
] == [1, 1]
def test_json_escaping_growth_is_counted() -> None:
@ -139,7 +192,9 @@ def test_json_escaping_growth_is_counted() -> None:
# 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]
assert [
len(batch) for batch in spend_log_write_batches([row, row], max_bytes=budget, max_rows=_ROWS_UNBOUNDED)
] == [1, 1]
def test_queue_within_budget_drops_the_oldest_rows_and_reports_what_is_left() -> None:
@ -174,4 +229,7 @@ def test_unserialized_list_payloads_are_measured_not_ignored() -> None:
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]
assert [len(batch) for batch in spend_log_write_batches([row, row], max_bytes=5100, max_rows=_ROWS_UNBOUNDED)] == [
1,
1,
]

View file

@ -358,9 +358,7 @@ async def test_update_spend_logs_failure_raises_after_retries(
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(
side_effect=httpx.ReadError("network blip")
)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=httpx.ReadError("network blip"))
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
with pytest.raises(httpx.ReadError):
@ -395,9 +393,7 @@ async def test_update_spend_logs_isolates_poison_row_and_persists_good_rows(
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"
)
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)
@ -576,7 +572,7 @@ async def test_update_spend_logs_does_not_requeue_non_transport_failures(
@pytest.mark.asyncio
async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood(
mock_prisma_client: Any, make_spend_log_row: Any
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A flood of poisoned rows must not amplify one failed bulk insert into
unbounded failed inserts. The per-batch failure budget hard-caps the number
@ -590,6 +586,9 @@ async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood(
# 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
# One statement, so this measures the isolation cap alone. The per-statement
# floor the row budget adds is pinned separately below.
monkeypatch.setattr(utils_mod, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", n_rows)
async def _always_poison(*, data: Any, skip_duplicates: bool) -> None:
raise _data_error("invalid byte sequence for encoding UTF8: 0x00")
@ -612,20 +611,52 @@ async def test_update_spend_logs_caps_isolation_attempts_under_poison_flood(
assert attempts < n_rows
@pytest.mark.asyncio
async def test_row_budget_costs_at_most_one_extra_attempt_per_statement(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Splitting a flush into more statements must not buy the poison flood a
fresh isolation budget each time. Every statement costs the one insert it
takes to discover it is poisoned, and the shared budget caps everything
above that, so the whole flush stays within the cap plus the statement
count however finely it is split.
"""
attempt_cap = utils_mod.MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH
n_rows = attempt_cap * 3
logs = [make_spend_log_row(request_id=f"r{i}") for i in range(n_rows)]
split = {"max_bytes": 2_000_000, "max_rows": 100, "monkeypatch": monkeypatch}
# A clean flush issues exactly one call per statement, so this is the observed
# split count rather than an arithmetic one; asserting it is >1 is what proves
# the row budget really divided the flush.
statements = await _flush_and_count_create_many(mock_prisma_client, logs, poison=False, **split)
attempts = await _flush_and_count_create_many(mock_prisma_client, logs, poison=True, **split)
assert statements > 1
assert attempts <= attempt_cap + statements
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,
max_rows: int = 10_000,
) -> int:
"""Run one flush and return how many ``create_many`` calls it issued."""
"""Run one flush and return how many ``create_many`` calls it issued.
``max_rows`` defaults high enough not to bind so a caller varying
``max_bytes`` measures the byte budget alone.
"""
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)
monkeypatch.setattr(utils_mod, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", max_rows)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many)
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
@ -717,15 +748,11 @@ def test_disable_spend_updates_reflects_general_settings(
"""
import litellm.proxy.proxy_server as proxy_server_mod
monkeypatch.setattr(
proxy_server_mod, "general_settings", {"disable_spend_updates": True}
)
monkeypatch.setattr(proxy_server_mod, "general_settings", {"disable_spend_updates": True})
pinned = {
"with_flag_true": ProxyUpdateSpend.disable_spend_updates(),
"type_is_bool": isinstance(ProxyUpdateSpend.disable_spend_updates(), bool),
"method_is_static": isinstance(
ProxyUpdateSpend.__dict__["disable_spend_updates"], staticmethod
),
"method_is_static": isinstance(ProxyUpdateSpend.__dict__["disable_spend_updates"], staticmethod),
}
assert pinned == {
"with_flag_true": True,