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 e1e1d3690ff..7f5efcdaec3 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -121,10 +121,12 @@ def advance_size_state(state: SizeCleanupState | None, probe: StorageProbe) -> S """ Roll the reclaimable count forward to what the catalog shows now. - Without a prior state, or after the cumulative counters went backwards (a statistics reset), the count - starts from the dead rows Postgres still knows about, at the files' average bytes per row. A probe with no - row counters at all (`stats_are_stale`) cannot seed anything: the prior count is carried forward, and with - no prior count the files are taken as fully reclaimable so nothing is deleted on its say-so. + Without a prior state the count starts from the dead rows Postgres still knows about, at the files' average + bytes per row. After the cumulative counters went backwards (a statistics reset, e.g. a crash) the prior + count is carried forward less any shrink of the files, or the dead-row seed if that is larger, because the + reset changed the counters and not the files. A probe with no row counters at all (`stats_are_stale`) + cannot seed anything: the prior count is carried forward, and with no prior count the files are taken as + fully reclaimable so nothing is deleted on its say-so. """ if probe.stats_are_stale: return SizeCleanupState( @@ -138,10 +140,26 @@ def advance_size_state(state: SizeCleanupState | None, probe: StorageProbe) -> S ) if state is None or probe.inserted_total < state.inserted_total or probe.deleted_total < state.deleted_total: seed_bytes_per_row: Final = probe.allocated_bytes / max(probe.live_rows + probe.dead_rows, 1) + seeded: Final = min(probe.allocated_bytes, round(probe.dead_rows * seed_bytes_per_row)) + if state is None: + return SizeCleanupState( + allocated_bytes=probe.allocated_bytes, + reclaimable_bytes=seeded, + bytes_per_row=max(seed_bytes_per_row, 1.0), + inserted_total=probe.inserted_total, + deleted_total=probe.deleted_total, + ) + carried: Final = max( + seeded, + min(probe.allocated_bytes, state.reclaimable_bytes - max(0, state.allocated_bytes - probe.allocated_bytes)), + ) + carried_live_bytes: Final = probe.allocated_bytes - carried return SizeCleanupState( allocated_bytes=probe.allocated_bytes, - reclaimable_bytes=min(probe.allocated_bytes, round(probe.dead_rows * seed_bytes_per_row)), - bytes_per_row=max(seed_bytes_per_row, 1.0), + reclaimable_bytes=carried, + bytes_per_row=max(carried_live_bytes / probe.live_rows, 1.0) + if probe.live_rows > 0 + else state.bytes_per_row, inserted_total=probe.inserted_total, deleted_total=probe.deleted_total, ) @@ -357,6 +375,11 @@ _STORAGE_SQL: Final = """ # The age job and the size job are scheduled on separate SpendLogCleanup instances, so the lock is per process. _PASS_LOCK: Final = asyncio.Lock() +_ANALYZE_SQL: Final = 'ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"' +# pg_stat_user_tables publishes what ANALYZE counted up to a second later (the stats collector on Postgres 14). +_STATS_SETTLE_SECONDS: Final = 1.0 +_STATS_SETTLE_ATTEMPTS: Final = 5 + SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = ( "maximum_spend_logs_cleanup_batch_size", "maximum_spend_logs_cleanup_max_batches", @@ -586,6 +609,48 @@ class SpendLogCleanup: deleted_total=row.deleted_total, ) + async def _reprobe_after_analyze( + self, prisma_client: PrismaClient, deadline: float, stale: StorageProbe + ) -> StorageProbe | None: + """ + Row counters that are all zero under a file with data were reset, e.g. by a crash. ANALYZE rebuilds + n_live_tup and n_dead_tup from a 30,000-row sample in seconds and blocks nothing, so the pass runs it and + reads the catalog again rather than trusting an estimate of zero or waiting a day for autoanalyze. The + new counts show up in pg_stat_user_tables a moment after ANALYZE returns, hence the settled re-reads. + """ + verbose_proxy_logger.warning( + "Spend-log row statistics are stale (allocated=%s bytes, live_rows=0, dead_rows=0); running %s before " + "the size pass", + stale.allocated_bytes, + _ANALYZE_SQL, + ) + if time.monotonic() >= deadline: + return None + timeout_ms: Final = self._timeout_ms(deadline) + try: + async with db_span("analyze", "LiteLLM_SpendLogs"), bounded_tx(prisma_client, timeout_ms) as tx: + await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}")) + await tx.execute_raw(trusted_sql(_ANALYZE_SQL)) + except Exception as e: # noqa: BLE001 - a failed ANALYZE skips the size pass, it must not fail the run + verbose_proxy_logger.warning("Could not ANALYZE the spend-log tables: %s: %s", type(e).__name__, e) + return None + refreshed: Final = await self._probe_once_settled(prisma_client, deadline, _STATS_SETTLE_ATTEMPTS) + if refreshed is None or refreshed.stats_are_stale: + verbose_proxy_logger.warning( + "Spend-log row statistics are still stale after ANALYZE; skipping the size pass" + ) + return None + return refreshed + + async def _probe_once_settled( + self, prisma_client: PrismaClient, deadline: float, attempts_left: int + ) -> StorageProbe | None: + await asyncio.sleep(_STATS_SETTLE_SECONDS) + refreshed: Final = await self._probe_storage(prisma_client, deadline) + if refreshed is None or not refreshed.stats_are_stale or attempts_left <= 1: + return refreshed + return await self._probe_once_settled(prisma_client, deadline, attempts_left - 1) + async def _load_size_state(self, prisma_client: PrismaClient, deadline: float) -> SizeCleanupState | None: """The reclaimable count saved by the last check, or None on the first check or when it cannot be read.""" if time.monotonic() >= deadline: @@ -1042,15 +1107,15 @@ class SpendLogCleanup: second. The run's batch cap and budget apply to the whole pass, so a pass that stops short reports how far above the target it still is and the next check continues from there. """ - probe: Final = await self._probe_storage(prisma_client, deadline) - if probe is None: + measured: Final = await self._probe_storage(prisma_client, deadline) + if measured is None: return None - if probe.stats_are_stale: - verbose_proxy_logger.warning( - "Spend-log row statistics are stale (allocated=%s bytes, live_rows=0, dead_rows=0); skipping the " - 'size pass until pg_stat_user_tables is repopulated, ANALYZE "LiteLLM_SpendLogs" refreshes it', - probe.allocated_bytes, - ) + probe: Final = ( + await self._reprobe_after_analyze(prisma_client, deadline, measured) + if measured.stats_are_stale + else measured + ) + if probe is None: return None state: Final = advance_size_state(await self._load_size_state(prisma_client, deadline), probe) before: Final = SpendLogStorage( diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index a0f222ed3e0..876ffe96584 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -1877,19 +1877,41 @@ async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_th @pytest.mark.asyncio -async def test_size_pass_skips_on_stale_row_statistics_instead_of_trusting_an_estimate_of_zero(caplog): +async def test_size_pass_analyzes_and_measures_again_when_the_row_statistics_were_reset(caplog): + """A crash zeroes pg_stat_user_tables under a full file; the pass rebuilds the counters instead of skipping. + The first read after ANALYZE still shows zeros, as the stats collector publishes them a moment later""" db = MagicMock() _wire_tx(db) - db.execute_raw = AsyncMock(return_value=1000) - _wire_size_probes(db, [_storage(allocated=5 * 1024**3, live=0, dead=0)], cutoff=datetime.now(timezone.utc)) + statements = AsyncMock(return_value=0) + db.execute_raw = statements + saves = _wire_size_probes( + db, + [ + _storage(allocated=10 * 1024**3, live=0, dead=0, inserted=0, deleted=0), + _storage(allocated=10 * 1024**3, live=0, dead=0, inserted=0, deleted=0), + _storage(10 * 1024**3, 1000, 0, 0, 0), + ], + cutoff=datetime.now(timezone.utc), + saved_state=SizeCleanupState( + allocated_bytes=10 * 1024**3, + reclaimable_bytes=9 * 1024**3 + 512 * 1024**2, + bytes_per_row=1024.0, + inserted_total=50_000, + deleted_total=49_000, + ), + ) client = MagicMock() client.db = db with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): await _size_cleaner({}).cleanup_spend_logs_by_size(client) - assert db.deletes.await_count == 0 + issued = [call.args[0] for call in statements.await_args_list] + assert issued == ['ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"'], ( + "the only statement is the ANALYZE: no DELETE, the carried count puts the table under 1GB" + ) assert any("row statistics are stale" in record.getMessage() for record in caplog.records) + assert saves[0].reclaimable_bytes == 9 * 1024**3 + 512 * 1024**2 @pytest.mark.asyncio @@ -1957,9 +1979,19 @@ _RUN_STATE = SizeCleanupState( ), ("the retention job deleted 2000 more rows", _storage(32_030_720, 3000, 2000, 20_000, 17_000), 27_217_000), ( - "statistics reset: counters restart from the dead rows", + "statistics reset (a crash, then ANALYZE): the count is carried forward, not restarted from the dead rows", _storage(32_030_720, 5000, 1000, 5000, 1000), - 5_338_453, + 24_015_000, + ), + ( + "statistics reset on a file VACUUM FULL shrank meanwhile: the carried count loses the shrink", + _storage(8_015_720, 5000, 0, 5000, 0), + 0, + ), + ( + "statistics reset with more dead rows than the carried count explains: the dead rows win", + _storage(32_030_720, 1000, 19_000, 1000, 19_000), + 30_429_184, ), ], )