From 22177263ec2d201823b5754562253eed0b48c2e9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 23:28:00 +0000 Subject: [PATCH] feat(proxy): add maximum_daily_tag_spend_retention_period cleanup setting Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 9 ++ .../db_transaction_queue/spend_log_cleanup.py | 117 ++++++++++++++---- litellm/proxy/proxy_server.py | 9 +- .../proxy/test_spend_log_cleanup.py | 25 +++- 4 files changed, 137 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e0a2097919b..bc3116b906a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2640,6 +2640,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "Set this well above health_check_interval because /health and the UI read the latest row per model." ), ) + maximum_daily_tag_spend_retention_period: str | None = Field( + None, + description=( + "Maximum retention period for LiteLLM_DailyTagSpend rows (e.g., '90d'). Rows whose date is older than " + "this are deleted by the spend log cleanup job, on that job's schedule. The table only feeds usage " + "analytics (tag usage dashboards, /spend/tags), so deleting old rows truncates historical tag usage " + "charts but does not affect budget enforcement. Unset means rows are never deleted." + ), + ) use_spend_logs_partitioning: bool | None = Field( None, description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.", 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 e97e9f6e683..1a11c27e79e 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -254,7 +254,7 @@ class SpendLogCleanup: return remaining async def _execute_delete_batch( - self, prisma_client: PrismaClient, delete_sql: str, cutoff_date: datetime, deadline: float + self, prisma_client: PrismaClient, delete_sql: str, cutoff: datetime | str, deadline: float ) -> int | None: """ Run one delete batch under a Postgres statement and lock timeout. @@ -273,11 +273,17 @@ class SpendLogCleanup: async with prisma_client.db.tx() as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}") await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}") - deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size) + deleted_result: Final = await tx.execute_raw(delete_sql, cutoff, self.batch_size) return deleted_result if isinstance(deleted_result, int) else None async def _count_remaining( - self, prisma_client: PrismaClient, cutoff_date: datetime, table_name: str, time_column: str, deadline: float + self, + prisma_client: PrismaClient, + cutoff: datetime | str, + table_name: str, + time_column: str, + time_cast: str, + deadline: float, ) -> int | None: """ Count expired rows still outstanding, stopping at a cap. @@ -290,7 +296,7 @@ class SpendLogCleanup: count_sql: Final = f""" SELECT count(*)::int AS remaining FROM ( SELECT 1 FROM "{table_name}" - WHERE "{time_column}" < $1::timestamptz + WHERE "{time_column}" < $1::{time_cast} LIMIT $2 ) capped """ @@ -298,7 +304,7 @@ class SpendLogCleanup: async with prisma_client.db.tx() as tx: await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}") rows: Final = _REMAINING_ROWS.validate_python( - await tx.query_raw(count_sql, cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP) + await tx.query_raw(count_sql, cutoff, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP) ) except Exception as e: # noqa: BLE001 - an observability probe must never fail the cleanup run verbose_proxy_logger.warning("Could not count remaining %s rows: %s", table_name, e) @@ -308,11 +314,12 @@ class SpendLogCleanup: async def _delete_old_rows_batched( self, prisma_client: PrismaClient, - cutoff_date: datetime, + cutoff: datetime | str, table_name: str, key_columns: tuple[str, ...], time_column: str, deadline: float, + time_cast: str = "timestamptz", ) -> TableCleanupResult: """ Delete a table's rows older than the cutoff in batches. @@ -326,7 +333,7 @@ class SpendLogCleanup: DELETE FROM "{table_name}" WHERE ({key_list}) IN ( SELECT {key_list} FROM "{table_name}" - WHERE "{time_column}" < $1::timestamptz + WHERE "{time_column}" < $1::{time_cast} LIMIT $2 ) """ @@ -341,19 +348,33 @@ class SpendLogCleanup: total_deleted, ) return await self._finish_table( - prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline + prisma_client, + cutoff, + table_name, + time_column, + time_cast, + total_deleted, + "budget_exhausted", + deadline, ) if run_count >= self.max_batches: verbose_proxy_logger.info( "Max batches reached for %s cleanup, remaining rows will be deleted in next run", table_name ) return await self._finish_table( - prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline + prisma_client, + cutoff, + table_name, + time_column, + time_cast, + total_deleted, + "batch_cap_reached", + deadline, ) # Find rows and delete them in one go without fetching to application batch_started_at = time.monotonic() try: - batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff_date, deadline) + batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff, deadline) except Exception as batch_exc: if time.monotonic() >= deadline: # The statement timeout was clamped to the budget that was @@ -368,7 +389,14 @@ class SpendLogCleanup: total_deleted, ) return await self._finish_table( - prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline + prisma_client, + cutoff, + table_name, + time_column, + time_cast, + total_deleted, + "budget_exhausted", + deadline, ) # A single batch failure (e.g. Prisma/DB timeout) must not abort # the whole run — subsequent batches may still succeed. @@ -382,7 +410,7 @@ class SpendLogCleanup: run_count, consecutive_failures, self.batch_size, - cutoff_date.isoformat(), + cutoff.isoformat() if isinstance(cutoff, datetime) else cutoff, total_deleted, type(batch_exc).__name__, batch_exc, @@ -395,7 +423,7 @@ class SpendLogCleanup: total_deleted, ) return await self._finish_table( - prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline + prisma_client, cutoff, table_name, time_column, time_cast, total_deleted, "aborted", deadline ) await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS) continue @@ -406,7 +434,7 @@ class SpendLogCleanup: table_name, ) return await self._finish_table( - prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline + prisma_client, cutoff, table_name, time_column, time_cast, total_deleted, "aborted", deadline ) consecutive_failures = 0 @@ -417,7 +445,7 @@ class SpendLogCleanup: if deleted_count == 0: verbose_proxy_logger.info("No more %s rows to delete. Total deleted: %s", table_name, total_deleted) return await self._finish_table( - prisma_client, cutoff_date, table_name, time_column, total_deleted, "exhausted", deadline + prisma_client, cutoff, table_name, time_column, time_cast, total_deleted, "exhausted", deadline ) total_deleted += deleted_count @@ -429,9 +457,10 @@ class SpendLogCleanup: async def _finish_table( self, prisma_client: PrismaClient, - cutoff_date: datetime, + cutoff: datetime | str, table_name: str, time_column: str, + time_cast: str, rows_deleted: int, stop_reason: StopReason, deadline: float, @@ -449,7 +478,9 @@ class SpendLogCleanup: """ if time.monotonic() >= deadline: return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason) - remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline) + remaining: Final = await self._count_remaining( + prisma_client, cutoff, table_name, time_column, time_cast, deadline + ) if remaining is not None: SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining) return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason) @@ -504,6 +535,21 @@ class SpendLogCleanup: deadline=deadline, ) + async def _delete_old_daily_tag_spend_rows( + self, prisma_client: PrismaClient, cutoff_day: str, deadline: float + ) -> TableCleanupResult: + # "date" is a YYYY-MM-DD string, so lexicographic order matches chronological + # order and the comparison stays sargable on the existing date index. + return await self._delete_old_rows_batched( + prisma_client, + cutoff_day, + table_name="LiteLLM_DailyTagSpend", + key_columns=("id",), + time_column="date", + deadline=deadline, + time_cast="text", + ) + async def _clean_spend_log_tables( self, prisma_client: PrismaClient, deadline: float ) -> tuple[TableCleanupResult, ...]: @@ -566,6 +612,16 @@ class SpendLogCleanup: ) return (health_checks_result,) + async def _clean_daily_tag_spend( + self, prisma_client: PrismaClient, retention_seconds: int, deadline: float + ) -> tuple[TableCleanupResult, ...]: + cutoff_day: Final = (datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))).strftime( + "%Y-%m-%d" + ) + tag_spend_result: Final = await self._delete_old_daily_tag_spend_rows(prisma_client, cutoff_day, deadline) + verbose_proxy_logger.info("Deleted %s expired daily tag spend rows", tag_spend_result.rows_deleted) + return (tag_spend_result,) + @staticmethod def _run_outcome(results: tuple[TableCleanupResult, ...]) -> RunOutcome: """ @@ -599,10 +655,14 @@ class SpendLogCleanup: "maximum_autorouter_session_retention_period" ) health_check_retention_seconds: Final = self._retention_seconds_for("maximum_health_check_retention_period") + daily_tag_spend_retention_seconds: Final = self._retention_seconds_for( + "maximum_daily_tag_spend_retention_period" + ) if ( not delete_spend_logs and autorouter_retention_seconds is None and health_check_retention_seconds is None + and daily_tag_spend_retention_seconds is None ): SpendLogCleanupMetrics.record_run("skipped_disabled") return @@ -634,6 +694,7 @@ class SpendLogCleanup: int(delete_spend_logs and self.retention_seconds is not None) + int(autorouter_retention_seconds is not None) + int(health_check_retention_seconds is not None) + + int(daily_tag_spend_retention_seconds is not None) ) spend_log_results: Final = ( @@ -644,8 +705,10 @@ class SpendLogCleanup: if delete_spend_logs and self.retention_seconds is not None else () ) - remaining_groups_after_spend_logs: Final = int(autorouter_retention_seconds is not None) + int( - health_check_retention_seconds is not None + remaining_groups_after_spend_logs: Final = ( + int(autorouter_retention_seconds is not None) + + int(health_check_retention_seconds is not None) + + int(daily_tag_spend_retention_seconds is not None) ) session_results: Final = ( await self._clean_session_rollup( @@ -656,18 +719,30 @@ class SpendLogCleanup: if autorouter_retention_seconds is not None else () ) + remaining_groups_after_sessions: Final = int(health_check_retention_seconds is not None) + int( + daily_tag_spend_retention_seconds is not None + ) health_check_results: Final = ( await self._clean_health_checks( prisma_client, health_check_retention_seconds, - deadline, + self._group_deadline(deadline, remaining_groups_after_sessions), ) if health_check_retention_seconds is not None else () ) + daily_tag_spend_results: Final = ( + await self._clean_daily_tag_spend( + prisma_client, + daily_tag_spend_retention_seconds, + deadline, + ) + if daily_tag_spend_retention_seconds is not None + else () + ) SpendLogCleanupMetrics.record_run( - self._run_outcome(spend_log_results + session_results + health_check_results) + self._run_outcome(spend_log_results + session_results + health_check_results + daily_tag_spend_results) ) except Exception as e: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77a80ea0052..695d8d6c88c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6572,7 +6572,13 @@ class ProxyConfig: retention_period: Final = general_settings.get("maximum_spend_logs_retention_period") autorouter_retention: Final = general_settings.get("maximum_autorouter_session_retention_period") health_check_retention: Final = general_settings.get("maximum_health_check_retention_period") - if retention_period is not None or autorouter_retention is not None or health_check_retention is not None: + daily_tag_spend_retention: Final = general_settings.get("maximum_daily_tag_spend_retention_period") + if ( + retention_period is not None + or autorouter_retention is not None + or health_check_retention is not None + or daily_tag_spend_retention is not None + ): from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( SpendLogCleanup, ) @@ -9508,6 +9514,7 @@ class ProxyStartupEvent: general_settings.get("maximum_spend_logs_retention_period") is not None or general_settings.get("maximum_autorouter_session_retention_period") is not None or general_settings.get("maximum_health_check_retention_period") is not None + or general_settings.get("maximum_daily_tag_spend_retention_period") is not None ): spend_log_cleanup: Final = SpendLogCleanup() cleanup_cron: Final = general_settings.get("maximum_spend_logs_cleanup_cron") diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index bf1538183ab..04651f8baef 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -823,14 +823,32 @@ async def test_health_check_retention_alone_cleans_only_the_health_check_table() assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1 +@pytest.mark.asyncio +async def test_daily_tag_spend_retention_alone_cleans_only_the_daily_tag_spend_table(): + client = _mock_prisma_for_retention([0]) + cleaner = SpendLogCleanup(general_settings={"maximum_daily_tag_spend_retention_period": "90d"}) + cleaner.pod_lock_manager = None + await cleaner.cleanup_old_spend_logs(client) + tables = [call[0][0] for call in client.db.execute_raw.call_args_list] + assert len(tables) == 1 + assert '"LiteLLM_DailyTagSpend"' in tables[0] + assert '"id"' in tables[0] + assert '"date"' in tables[0] + assert "$1::text" in tables[0] + cutoff_day = client.db.execute_raw.call_args[0][1] + expected_cutoff_day = (datetime.now(timezone.utc) - timedelta(days=90)).strftime("%Y-%m-%d") + assert cutoff_day == expected_cutoff_day + + @pytest.mark.asyncio async def test_each_retention_key_cuts_off_at_its_own_horizon(): - client = _mock_prisma_for_retention([0, 0, 0, 0]) + client = _mock_prisma_for_retention([0, 0, 0, 0, 0]) cleaner = SpendLogCleanup( general_settings={ "maximum_spend_logs_retention_period": "7d", "maximum_autorouter_session_retention_period": "365d", "maximum_health_check_retention_period": "30d", + "maximum_daily_tag_spend_retention_period": "90d", } ) cleaner.pod_lock_manager = None @@ -841,6 +859,8 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon(): if '"LiteLLM_AutoRouterSession"' in call[0][0] else "LiteLLM_HealthCheckTable" if '"LiteLLM_HealthCheckTable"' in call[0][0] + else "LiteLLM_DailyTagSpend" + if '"LiteLLM_DailyTagSpend"' in call[0][0] else "logs" ): call[0][1] for call in client.db.execute_raw.call_args_list @@ -849,6 +869,7 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon(): assert (now - cutoffs["logs"]).days == 7 assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365 assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30 + assert cutoffs["LiteLLM_DailyTagSpend"] == (now - timedelta(days=90)).strftime("%Y-%m-%d") @pytest.mark.asyncio @@ -1213,6 +1234,7 @@ async def test_the_outstanding_rows_probe_carries_a_statement_timeout(): datetime.now(timezone.utc) - timedelta(days=7), "LiteLLM_SpendLogs", "startTime", + "timestamptz", _far_deadline(), ) @@ -1279,6 +1301,7 @@ async def test_no_statement_is_issued_once_the_budget_is_spent(): datetime.now(timezone.utc) - timedelta(days=7), "LiteLLM_SpendLogs", "startTime", + "timestamptz", 123, "budget_exhausted", time.monotonic() - 1,