diff --git a/litellm/constants.py b/litellm/constants.py index ec632435e45..7eccfc58221 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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)) diff --git a/litellm/proxy/db/spend_log_batching.py b/litellm/proxy/db/spend_log_batching.py index a8fced5485d..daba63c54ad 100644 --- a/litellm/proxy/db/spend_log_batching.py +++ b/litellm/proxy/db/spend_log_batching.py @@ -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] diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 81a86ebe34d..ce8fe68d76c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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), diff --git a/tests/proxy_unit_tests/test_update_spend.py b/tests/proxy_unit_tests/test_update_spend.py index 6b8973fbad2..2df381c8190 100644 --- a/tests/proxy_unit_tests/test_update_spend.py +++ b/tests/proxy_unit_tests/test_update_spend.py @@ -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 = [] diff --git a/tests/test_litellm/proxy/db/test_spend_log_batching.py b/tests/test_litellm/proxy/db/test_spend_log_batching.py index 2069490e7a0..bc26e17dac4 100644 --- a/tests/test_litellm/proxy/db/test_spend_log_batching.py +++ b/tests/test_litellm/proxy/db/test_spend_log_batching.py @@ -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, + ] diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index 93c99c7fd04..048fddb10d6 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -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,