scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-28 06:23:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 94

feat(proxy): size pass drops the oldest spend log partitions first (#47451)

Browse source 
* 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>
This commit is contained in:
Kerry Lu 2026-10-20 11:20:37 -07:00 • committed by GitHub
parent 2773ef0084
commit 7f0bb8424c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 219 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

74
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -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

48
litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
View file

@ -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]:

108
tests/unit/proxy/test_spend_log_cleanup.py
View file

@ -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"),
[