mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
7da34e8aed
commit
40b8300ac2
6 changed files with 169 additions and 59 deletions
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue