From 7f0bb8424c5e075e84b18ca726e022ed2fb5be6c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 11:20:37 -0700 Subject: [PATCH] feat(proxy): size pass drops the oldest spend log partitions first (#47451) * feat(proxy): size pass drops the oldest spend log partitions first On a range-partitioned LiteLLM_SpendLogs the size pass drops fully elapsed partitions oldest first, never the current or DEFAULT one, remeasures after each drop and falls back to oldest-first row deletes only while still above the target * fix(proxy): a partition drop rebases the size state instead of reading as a statistics reset Bugbot on #47449: dropping a partition removes that child's n_tup_ins and n_tup_del from the summed probe, which advance_size_state took for a reset and reseeded from n_dead_tup. The pass now saves the counters of the last post-drop probe as the state's base and scales the reclaimable count to the surviving files * style(proxy): keep the mutable-ok reason on the line it suppresses * fix(proxy): size pass stops when a post-drop re-measure fails and keeps day partitions ensured A failed storage probe after a partition drop no longer reads as under target. The pass logs a warning, skips the row step and saves no count, so neither live rows get deleted on the pre-drop numbers nor the next check reads the drop as a statistics reset. The partitioned size path also calls ensure_partitions before dropping, so a size-only setup keeps getting day partitions instead of writing into DEFAULT --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../db_transaction_queue/spend_log_cleanup.py | 74 ++++++++++-- .../spend_logs_partition_manager.py | 48 +++++++- tests/unit/proxy/test_spend_log_cleanup.py | 108 ++++++++++++++++++ 3 files changed, 219 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 268f0230d7a..726bc2b221c 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -2,7 +2,7 @@ import asyncio import math import re import time -from collections.abc import Callable +from collections.abc import Awaitable, Callable from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -254,6 +254,7 @@ class SizePassReport: rule: SizeRule before: SpendLogStorage after: SpendLogStorage + partitions_dropped: tuple[str, ...] results: tuple[TableCleanupResult, ...] @property @@ -1028,7 +1029,9 @@ class SpendLogCleanup: Bring the spend-log tables under the size target, oldest rows first. Live bytes are the files on disk less the reclaimable count carried in LiteLLM_Config (see - `SizeCleanupState`), rolled forward from the catalog's cumulative insert and delete counters. Rows go in steps: the rows still needed to reach the target follow + `SizeCleanupState`), rolled forward from the catalog's cumulative insert and delete counters. A + partitioned table gives its oldest fully elapsed partitions back first, because a DROP returns disk at + once, re-measuring after each. Then rows go in steps: the rows still needed to reach the target follow from bytes per row, and each step turns that many of the oldest rows (at most SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES batches) into a cutoff and runs the same bounded DELETE the age rule uses, for the spend logs and the tool index rows derived from them. Progress is counted from the @@ -1066,7 +1069,7 @@ class SpendLogCleanup: ) SpendLogCleanupMetrics.set_storage_bytes("above_target", 0) await self._save_size_state(prisma_client, deadline, state) - return SizePassReport(rule=rule, before=before, after=before, results=()) + return SizePassReport(rule=rule, before=before, after=before, partitions_dropped=(), results=()) verbose_proxy_logger.warning( "Spend-log storage is over maximum_spend_logs_size (estimated_live=%s allocated=%s reclaimable=%s " "limit=%s bytes); deleting the oldest spend logs down to %s bytes", @@ -1076,9 +1079,40 @@ class SpendLogCleanup: rule.limit_bytes, rule.target_bytes, ) - rows_to_delete: Final = max( - 0, math.ceil((before.estimated_live_bytes - rule.target_bytes) / before.bytes_per_row) + latest: Final[list[SpendLogStorage]] = [before] # mutable-ok: last measurement, rewritten per partition drop + baseline: Final[list[StorageProbe]] = [probe] # mutable-ok: counters after the last drop + + remeasure_failed: Final[list[bool]] = [False] # mutable-ok: set when a post-drop probe fails + + async def still_over_target() -> bool: + measured: Final = await self._probe_storage(prisma_client, deadline) + if measured is None: + remeasure_failed[0] = True + return False + kept: Final = min(1.0, measured.allocated_bytes / max(latest[0].allocated_bytes, 1)) + latest[0] = SpendLogStorage( + allocated_bytes=measured.allocated_bytes, + reclaimable_bytes=min(measured.allocated_bytes, round(latest[0].reclaimable_bytes * kept)), + bytes_per_row=latest[0].bytes_per_row, + live_rows=measured.live_rows, + ) + baseline[0] = measured + return latest[0].estimated_live_bytes > rule.target_bytes + + partitions_dropped: Final = await self._drop_oldest_partitions_for_size( + prisma_client, deadline, still_over_target ) + if remeasure_failed[0]: + verbose_proxy_logger.warning( + "Spend-log storage could not be re-measured after dropping partitions %s; the row step and the saved " + "count wait for the next check rather than run on the pre-drop numbers", + partitions_dropped, + ) + return SizePassReport( + rule=rule, before=before, after=latest[0], partitions_dropped=partitions_dropped, results=() + ) + base: Final = latest[0] + rows_to_delete: Final = max(0, math.ceil((base.estimated_live_bytes - rule.target_bytes) / base.bytes_per_row)) def rows_still_needed(rows_deleted: int) -> int: return rows_to_delete - rows_deleted @@ -1086,20 +1120,39 @@ class SpendLogCleanup: results: Final = await self._delete_oldest_rows_for_size(prisma_client, deadline, rows_still_needed) rows_deleted: Final = sum(result.rows_deleted for result in results) remeasured: Final = await self._probe_storage(prisma_client, deadline) - after: Final = before.after_deleting( - rows_deleted, remeasured.allocated_bytes if remeasured is not None else before.allocated_bytes + after: Final = base.after_deleting( + rows_deleted, remeasured.allocated_bytes if remeasured is not None else base.allocated_bytes ) await self._save_size_state( prisma_client, deadline, - after.to_state(inserted_total=probe.inserted_total, deleted_total=probe.deleted_total + rows_deleted), + after.to_state( + inserted_total=baseline[0].inserted_total, deleted_total=baseline[0].deleted_total + rows_deleted + ), + ) + report: Final = SizePassReport( + rule=rule, before=before, after=after, partitions_dropped=partitions_dropped, results=results ) - report: Final = SizePassReport(rule=rule, before=before, after=after, results=results) SpendLogCleanupMetrics.set_storage_bytes("allocated", after.allocated_bytes) SpendLogCleanupMetrics.set_storage_bytes("estimated_live", after.estimated_live_bytes) SpendLogCleanupMetrics.set_storage_bytes("above_target", report.above_target_bytes) return report + async def _drop_oldest_partitions_for_size( + self, prisma_client: PrismaClient, deadline: float, still_over_target: Callable[[], Awaitable[bool]] + ) -> tuple[str, ...]: + remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline) + if time.monotonic() >= deadline or not self._setting("use_spend_logs_partitioning"): + return () + if not await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms): + return () + await self.partition_manager.ensure_partitions(prisma_client, remaining_timeout_ms) + dropped: Final = await self.partition_manager.drop_oldest_partitions_while( + prisma_client, datetime.now(timezone.utc), remaining_timeout_ms, still_over_target + ) + verbose_proxy_logger.info("Dropped %d oldest spend-log partitions for the size rule: %s", len(dropped), dropped) + return tuple(dropped) + async def _delete_oldest_rows_for_size( self, prisma_client: PrismaClient, @@ -1275,7 +1328,8 @@ class SpendLogCleanup: f"allocated_after={report.after.allocated_bytes} estimated_live_after={report.after.estimated_live_bytes} " f"reclaimable_after={report.after.reclaimable_bytes} " f"limit={report.rule.limit_bytes} target={report.rule.target_bytes} " - f"target_reached={str(report.target_reached).lower()} above_target={report.above_target_bytes}{reclaim_note}" + f"target_reached={str(report.target_reached).lower()} above_target={report.above_target_bytes} " + f"partitions_dropped={len(report.partitions_dropped)}{reclaim_note}" ) @classmethod diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py index 3f9384d5d15..049a6ecbeee 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py +++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py @@ -14,7 +14,7 @@ keeps the batched-DELETE path, so existing deployments are untouched. """ import re -from collections.abc import Callable +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence from datetime import date, datetime, timedelta, timezone from typing import ( TYPE_CHECKING, @@ -265,6 +265,52 @@ class SpendLogsPartitionManager: ) return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows] + async def drop_oldest_partitions_while( + self, + prisma_client: "PrismaClient", + now: datetime, + remaining_timeout_ms: RemainingTimeoutMs, + still_over_target: Callable[[], Awaitable[bool]], + ) -> tuple[str, ...]: + """ + DROP fully elapsed partitions oldest first until `still_over_target` says the size rule is satisfied. + + Only partitions whose whole range lies before `now` are candidates, so the partition receiving + current writes and the DEFAULT partition are never dropped. The callback runs after every drop + because a DROP returns its disk at once, so a re-measure is what tells the loop to stop. + """ + list_budget_ms: Final = remaining_timeout_ms() + if list_budget_ms is None: + return () + now_naive: Final = now.astimezone(timezone.utc).replace(tzinfo=None) + partitions: Final = await self._list_partitions(prisma_client, list_budget_ms) + elapsed: Final = sorted( + ((upper, name) for name, upper in partitions if upper is not None and upper <= now_naive), + ) + drops: Final = self._drop_each_while(prisma_client, elapsed, remaining_timeout_ms, still_over_target) + return tuple([name async for name in drops]) + + async def _drop_each_while( + self, + prisma_client: "PrismaClient", + elapsed: Sequence[tuple[datetime, str]], + remaining_timeout_ms: RemainingTimeoutMs, + still_over_target: Callable[[], Awaitable[bool]], + ) -> AsyncIterator[str]: + for _, name in elapsed: + budget_ms: int | None = remaining_timeout_ms() + if budget_ms is None: + verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run") + return + try: + await self._execute_bounded_ddl(prisma_client, f'DROP TABLE IF EXISTS "{name}"', budget_ms) + except Exception as e: # noqa: BLE001 # a partition this run cannot drop is left for the next one + verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e) + return + yield name + if not await still_over_target(): + return + async def drop_partitions_older_than( self, prisma_client: "PrismaClient", cutoff: datetime, remaining_timeout_ms: RemainingTimeoutMs ) -> list[str]: diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index e70b93f7955..8c97f829025 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -1998,6 +1998,114 @@ async def test_the_retention_run_carries_the_size_pass_only_when_one_cron_drives assert measured is size_pass_expected +@pytest.mark.asyncio +async def test_partitioned_size_pass_drops_the_oldest_partitions_first_and_skips_row_deletes_once_under_target( + caplog, +): + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=1000) + saves = _wire_size_probes( + db, + [_storage(10 * 1024**3, live=1000, dead=0), _storage(256 * 1024**2, live=25, dead=0)], + cutoff=datetime.now(timezone.utc), + ) + client = MagicMock() + client.db = db + verdicts = [] + + async def _drop(prisma_client, now, remaining_timeout_ms, still_over_target): + verdicts.append(await still_over_target()) + return ("LiteLLM_SpendLogs_p20260601",) + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=True) + partition_manager.ensure_partitions = AsyncMock() + partition_manager.drop_oldest_partitions_while = AsyncMock(side_effect=_drop) + cleaner = _size_cleaner( + {"use_spend_logs_partitioning": True, "maximum_spend_logs_cleanup_target_size": "512MB"}, + partition_manager=partition_manager, + ) + + with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): + await cleaner.cleanup_spend_logs_by_size(client) + + partition_manager.ensure_partitions.assert_awaited_once(), "a size-only setup still needs tomorrow's partition" + partition_manager.drop_oldest_partitions_while.assert_awaited_once() + assert verdicts == [False] + assert db.deletes.await_count == 0 + summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage()) + assert "partitions_dropped=1" in summary + assert "target_reached=true" in summary + assert (saves[-1].inserted_total, saves[-1].deleted_total, saves[-1].allocated_bytes) == (25, 0, 256 * 1024**2), ( + "the saved counters must be the post-drop ones, or the next check reads the drop as a statistics reset" + ) + + +@pytest.mark.asyncio +async def test_a_failed_remeasure_after_a_partition_drop_stops_the_pass_instead_of_deleting_on_old_numbers(caplog): + db = MagicMock() + _wire_tx(db) + db.execute_raw = AsyncMock(return_value=1000) + saves = _wire_size_probes(db, [_storage(10 * 1024**3, live=1000, dead=0)], cutoff=datetime.now(timezone.utc)) + client = MagicMock() + client.db = db + + async def _drop(prisma_client, now, remaining_timeout_ms, still_over_target): + db.query_raw = AsyncMock(side_effect=RuntimeError("connection reset")) + assert await still_over_target() is False + return ("LiteLLM_SpendLogs_p20260601",) + + partition_manager = MagicMock() + partition_manager.is_partitioned = AsyncMock(return_value=True) + partition_manager.ensure_partitions = AsyncMock() + partition_manager.drop_oldest_partitions_while = AsyncMock(side_effect=_drop) + cleaner = _size_cleaner( + {"use_spend_logs_partitioning": True, "maximum_spend_logs_cleanup_target_size": "512MB"}, + partition_manager=partition_manager, + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await cleaner.cleanup_spend_logs_by_size(client) + + assert db.deletes.await_count == 0, "the row step ran on the pre-drop measurement" + assert saves == [], "a count saved from pre-drop counters reads the drop as a statistics reset next check" + assert any("could not be re-measured after dropping partitions" in r.getMessage() for r in caplog.records) + + +def test_drop_oldest_partitions_while_only_drops_fully_elapsed_partitions_oldest_first(): + from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import SpendLogsPartitionManager + + manager = SpendLogsPartitionManager() + now = datetime(2026, 10, 17, 12, tzinfo=timezone.utc) + manager._list_partitions = AsyncMock( + return_value=[ + ("LiteLLM_SpendLogs_p20261017", datetime(2026, 10, 18)), # today's, still receiving writes + ("LiteLLM_SpendLogs_p20261015", datetime(2026, 10, 16)), + ("LiteLLM_SpendLogs_default", None), + ("LiteLLM_SpendLogs_p20261014", datetime(2026, 10, 15)), + ] + ) + dropped_sql: list[str] = [] + + async def _ddl(prisma_client, statement, timeout_ms): + dropped_sql.append(statement) + + manager._execute_bounded_ddl = _ddl + verdicts = iter([True, False]) + + async def still_over_target(): + return next(verdicts) + + dropped = asyncio.run(manager.drop_oldest_partitions_while(MagicMock(), now, lambda: 1000, still_over_target)) + + assert dropped == ("LiteLLM_SpendLogs_p20261014", "LiteLLM_SpendLogs_p20261015") + assert dropped_sql == [ + 'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20261014"', + 'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20261015"', + ] + + @pytest.mark.parametrize( ("settings", "retention_job", "size_job"), [