feat(proxy): rebase daily tag spend retention onto the run-budgeted cleanup job

Reworks the cleanup on top of the refactored SpendLogCleanup: the daily tag spend table is pruned through the shared batched delete with a text cutoff on the indexed ISO date column, the setting is picked up by /config/update and the scheduler registration, and an integration test proves rows older than the period are pruned while the cutoff day and unset retention are left alone

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-23 07:00:59 +00:00
parent ae51fb266d
commit c7af908e89
8 changed files with 192 additions and 82 deletions

View file

@ -2961,10 +2961,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
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."
"Maximum retention period for per-day tag spend aggregate rows (e.g., '90d'). Rows whose day is older "
"than this are deleted by the spend log cleanup job, on that job's schedule. Unset means rows are never "
"deleted. Only historical tag usage analytics are affected; tag budgets read the lifetime counter."
),
)
use_spend_logs_partitioning: bool | None = Field(

View file

@ -32,6 +32,17 @@ from litellm.proxy.utils import PrismaClient
StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reached", "aborted"]
Cutoff: TypeAlias = datetime | str
"""Rows strictly older than this are expired: a timestamp, or an ISO calendar day for tables keyed by day"""
def _cutoff_cast(cutoff: Cutoff) -> str:
return "timestamptz" if isinstance(cutoff, datetime) else "text"
def _cutoff_text(cutoff: Cutoff) -> str:
return cutoff.isoformat() if isinstance(cutoff, datetime) else cutoff
@dataclass(frozen=True, slots=True)
class TableCleanupResult:
@ -277,7 +288,7 @@ class SpendLogCleanup:
return remaining
async def _execute_delete_batch(
self, prisma_client: PrismaClient, delete_sql: str, cutoff: datetime | str, deadline: float
self, prisma_client: PrismaClient, delete_sql: str, cutoff_date: Cutoff, deadline: float
) -> int | None:
"""
Run one delete batch under a Postgres statement and lock timeout.
@ -296,17 +307,11 @@ 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, self.batch_size)
deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size)
return deleted_result if isinstance(deleted_result, int) else None
async def _count_remaining(
self,
prisma_client: PrismaClient,
cutoff: datetime | str,
table_name: str,
time_column: str,
time_cast: str,
deadline: float,
self, prisma_client: PrismaClient, cutoff_date: Cutoff, table_name: str, time_column: str, deadline: float
) -> int | None:
"""
Count expired rows still outstanding, stopping at a cap.
@ -319,7 +324,7 @@ class SpendLogCleanup:
count_sql: Final = f"""
SELECT count(*)::int AS remaining FROM (
SELECT 1 FROM "{table_name}"
WHERE "{time_column}" < $1::{time_cast}
WHERE "{time_column}" < $1::{_cutoff_cast(cutoff_date)}
LIMIT $2
) capped
"""
@ -327,7 +332,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, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP)
await tx.query_raw(count_sql, cutoff_date, 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)
@ -337,12 +342,11 @@ class SpendLogCleanup:
async def _delete_old_rows_batched(
self,
prisma_client: PrismaClient,
cutoff: datetime | str,
cutoff_date: Cutoff,
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.
@ -356,7 +360,7 @@ class SpendLogCleanup:
DELETE FROM "{table_name}"
WHERE ({key_list}) IN (
SELECT {key_list} FROM "{table_name}"
WHERE "{time_column}" < $1::{time_cast}
WHERE "{time_column}" < $1::{_cutoff_cast(cutoff_date)}
LIMIT $2
)
"""
@ -371,33 +375,19 @@ class SpendLogCleanup:
total_deleted,
)
return await self._finish_table(
prisma_client,
cutoff,
table_name,
time_column,
time_cast,
total_deleted,
"budget_exhausted",
deadline,
prisma_client, cutoff_date, table_name, time_column, 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,
table_name,
time_column,
time_cast,
total_deleted,
"batch_cap_reached",
deadline,
prisma_client, cutoff_date, table_name, time_column, 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, deadline)
batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff_date, deadline)
except Exception as batch_exc:
if time.monotonic() >= deadline:
# The statement timeout was clamped to the budget that was
@ -412,14 +402,7 @@ class SpendLogCleanup:
total_deleted,
)
return await self._finish_table(
prisma_client,
cutoff,
table_name,
time_column,
time_cast,
total_deleted,
"budget_exhausted",
deadline,
prisma_client, cutoff_date, table_name, time_column, 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.
@ -433,7 +416,7 @@ class SpendLogCleanup:
run_count,
consecutive_failures,
self.batch_size,
cutoff.isoformat() if isinstance(cutoff, datetime) else cutoff,
_cutoff_text(cutoff_date),
total_deleted,
type(batch_exc).__name__,
batch_exc,
@ -446,7 +429,7 @@ class SpendLogCleanup:
total_deleted,
)
return await self._finish_table(
prisma_client, cutoff, table_name, time_column, time_cast, total_deleted, "aborted", deadline
prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
)
await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
continue
@ -457,7 +440,7 @@ class SpendLogCleanup:
table_name,
)
return await self._finish_table(
prisma_client, cutoff, table_name, time_column, time_cast, total_deleted, "aborted", deadline
prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
)
consecutive_failures = 0
@ -468,7 +451,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, table_name, time_column, time_cast, total_deleted, "exhausted", deadline
prisma_client, cutoff_date, table_name, time_column, total_deleted, "exhausted", deadline
)
total_deleted += deleted_count
@ -481,10 +464,9 @@ class SpendLogCleanup:
async def _finish_table(
self,
prisma_client: PrismaClient,
cutoff: datetime | str,
cutoff_date: Cutoff,
table_name: str,
time_column: str,
time_cast: str,
rows_deleted: int,
stop_reason: StopReason,
deadline: float,
@ -502,9 +484,7 @@ 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, table_name, time_column, time_cast, deadline
)
remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline)
if remaining is not None:
SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining)
return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
@ -574,8 +554,6 @@ class SpendLogCleanup:
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,
@ -583,7 +561,6 @@ class SpendLogCleanup:
key_columns=("id",),
time_column="date",
deadline=deadline,
time_cast="text",
)
async def _clean_spend_log_tables(
@ -670,12 +647,14 @@ class SpendLogCleanup:
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,)
"""
Prune per-day tag spend rows whose ISO day sorts before the horizon day; the horizon day itself is kept.
"""
horizon: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))
cutoff_day: Final = horizon.date().isoformat()
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", result.rows_deleted)
return (result,)
@staticmethod
def _run_outcome(results: tuple[TableCleanupResult, ...]) -> RunOutcome:
@ -768,6 +747,9 @@ class SpendLogCleanup:
+ int(health_check_retention_seconds is not None)
+ int(daily_tag_spend_retention_seconds is not None)
)
remaining_groups_after_sessions: Final = 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(
prisma_client,
@ -777,9 +759,6 @@ 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,
@ -790,11 +769,7 @@ class SpendLogCleanup:
else ()
)
daily_tag_spend_results: Final = (
await self._clean_daily_tag_spend(
prisma_client,
daily_tag_spend_retention_seconds,
deadline,
)
await self._clean_daily_tag_spend(prisma_client, daily_tag_spend_retention_seconds, deadline)
if daily_tag_spend_retention_seconds is not None
else ()
)

View file

@ -7526,6 +7526,7 @@ class ProxyConfig:
"maximum_spend_logs_retention_period",
"maximum_autorouter_session_retention_period",
"maximum_health_check_retention_period",
"maximum_daily_tag_spend_retention_period",
)
)
@ -17744,6 +17745,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
"store_prompts_in_spend_logs": "Boolean",
"maximum_spend_logs_retention_period": "String",
"maximum_health_check_retention_period": "String",
"maximum_daily_tag_spend_retention_period": "String",
"maximum_spend_logs_cleanup_batch_size": "Integer",
"maximum_spend_logs_cleanup_max_batches": "Integer",
"maximum_spend_logs_cleanup_run_budget": "String",

View file

@ -0,0 +1,103 @@
import os
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final
import psycopg
import pytest
import yaml
from pydantic import JsonValue, TypeAdapter
from tests.integration._support.client import Gateway, eventually
from tests.integration._support.database import read_rows
from tests.integration._support.process import owned_proxy
CLEANUP_EVERY_MINUTE: Final = "* * * * *"
_CONFIG: Final = TypeAdapter(dict[str, dict[str, JsonValue]])
def _day(days_ago: int) -> str:
return (datetime.now(timezone.utc) - timedelta(days=days_ago)).strftime("%Y-%m-%d")
def _seed_daily_tag_spend(tag: str, days: tuple[str, ...]) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
for day in days:
connection.execute(
'INSERT INTO "LiteLLM_DailyTagSpend" (id, tag, date, api_key, model, spend, updated_at) '
"VALUES (%s, %s, %s, %s, %s, 1.0, now())",
(uuid.uuid4().hex, tag, day, f"integration-{tag}", "gpt-4o-mini"),
)
def _seed_old_spend_log(request_id: str, days_ago: int) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute(
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, api_key, spend, "startTime", "endTime") '
"VALUES (%s, 'acompletion', %s, 0, now() - make_interval(days => %s), now() - make_interval(days => %s))",
(request_id, f"integration-{request_id}", str(days_ago), str(days_ago)),
)
def _delete_daily_tag_spend(tag: str) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
connection.execute('DELETE FROM "LiteLLM_DailyTagSpend" WHERE tag = %s', (tag,))
def _remaining_days(tag: str) -> tuple[str, ...]:
rows: Final = read_rows('SELECT date FROM "LiteLLM_DailyTagSpend" WHERE tag = %s ORDER BY date', (tag,))
return tuple(str(row["date"]) for row in rows)
def _spend_log_present(request_id: str) -> bool:
return bool(read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)))
def _cleanup_config(tmp_path: Path, retention: dict[str, JsonValue]) -> Path:
base: Final = _CONFIG.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {
**base,
"general_settings": {
**base["general_settings"],
**retention,
"maximum_spend_logs_cleanup_cron": CLEANUP_EVERY_MINUTE,
"scheduled_job_stagger": {"enabled": False},
},
}
path: Final = tmp_path / "retention.yaml"
path.write_text(yaml.safe_dump(config))
return path
@pytest.mark.covers("spend.daily_tag_spend.retention_prunes_rows_older_than_the_period_and_keeps_the_rest")
def test_daily_tag_spend_retention_prunes_only_rows_older_than_the_period(gateway: Gateway, tmp_path: Path) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
expired, on_the_cutoff, today = _day(200), _day(30), _day(0)
_seed_daily_tag_spend(tag, (expired, on_the_cutoff, today))
try:
config: Final = _cleanup_config(tmp_path, {"maximum_daily_tag_spend_retention_period": "30d"})
with owned_proxy(gateway, tmp_path, {}, config=config):
remaining: Final = eventually(
lambda: _remaining_days(tag),
lambda days: expired not in days,
seconds=150,
)
assert remaining == (on_the_cutoff, today), remaining
finally:
_delete_daily_tag_spend(tag)
@pytest.mark.covers("spend.daily_tag_spend.unset_retention_never_deletes_even_while_spend_logs_are_pruned")
def test_daily_tag_spend_is_kept_forever_when_its_retention_is_unset(gateway: Gateway, tmp_path: Path) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
request_id: Final = f"integration-retention-{uuid.uuid4().hex}"
_seed_daily_tag_spend(tag, (_day(200),))
_seed_old_spend_log(request_id, days_ago=200)
try:
config: Final = _cleanup_config(tmp_path, {"maximum_spend_logs_retention_period": "30d"})
with owned_proxy(gateway, tmp_path, {}, config=config):
eventually(lambda: _spend_log_present(request_id), lambda present: not present, seconds=150)
assert _remaining_days(tag) == (_day(200),)
finally:
_delete_daily_tag_spend(tag)

View file

@ -79,6 +79,7 @@ _PREVIOUSLY_DB_WINS: Final[tuple[str, ...]] = (
"maximum_spend_logs_retention_period",
"maximum_autorouter_session_retention_period",
"maximum_health_check_retention_period",
"maximum_daily_tag_spend_retention_period",
"maximum_spend_logs_cleanup_batch_size",
"maximum_spend_logs_cleanup_max_batches",
"maximum_spend_logs_cleanup_run_budget",

View file

@ -3872,6 +3872,35 @@ async def test_ProxyConfig__update_general_settings_updates_health_check_retenti
reschedule.assert_awaited_once()
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_daily_tag_spend_retention(monkeypatch):
fake_scheduler = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", fake_scheduler)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"maximum_daily_tag_spend_retention_period": "90d"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
await pc._reschedule_spend_log_cleanup_job()
assert fake_scheduler.add_job.call_count == 1
assert fake_scheduler.add_job.call_args.kwargs["id"] == "spend_log_cleanup_job"
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_updates_daily_tag_spend_retention(monkeypatch):
settings = {}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", settings)
pc = ProxyConfig()
reschedule = AsyncMock()
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
from litellm.proxy import proxy_server
assert proxy_server.general_settings["maximum_daily_tag_spend_retention_period"] == "90d"
reschedule.assert_awaited_once()
# ---------------------------------------------------------------------------
# ProxyConfig._update_general_settings
# ---------------------------------------------------------------------------

View file

@ -827,7 +827,7 @@ async def test_health_check_retention_alone_cleans_only_the_health_check_table()
@pytest.mark.asyncio
async def test_daily_tag_spend_retention_alone_cleans_only_the_daily_tag_spend_table():
async def test_daily_tag_spend_retention_alone_prunes_only_that_table_by_calendar_day():
client = _mock_prisma_for_retention([0])
cleaner = SpendLogCleanup(general_settings={"maximum_daily_tag_spend_retention_period": "90d"})
cleaner.pod_lock_manager = None
@ -835,12 +835,19 @@ async def test_daily_tag_spend_retention_alone_cleans_only_the_daily_tag_spend_t
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]
assert '"date" < $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
assert cutoff_day == (datetime.now(timezone.utc) - timedelta(days=90)).date().isoformat()
@pytest.mark.asyncio
async def test_spend_logs_retention_alone_keeps_daily_tag_spend_forever():
client = _mock_prisma_for_retention([0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
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 not any('"LiteLLM_DailyTagSpend"' in sql for sql in tables)
@pytest.mark.asyncio
@ -851,7 +858,6 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
"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
@ -864,8 +870,6 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
if '"LiteLLM_AutoRouterUserSession"' 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
@ -875,7 +879,6 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365
assert cutoffs["LiteLLM_AutoRouterUserSession"] == cutoffs["LiteLLM_AutoRouterSession"]
assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30
assert cutoffs["LiteLLM_DailyTagSpend"] == (now - timedelta(days=90)).strftime("%Y-%m-%d")
@pytest.mark.asyncio
@ -1240,7 +1243,6 @@ async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
datetime.now(timezone.utc) - timedelta(days=7),
"LiteLLM_SpendLogs",
"startTime",
"timestamptz",
_far_deadline(),
)
@ -1307,7 +1309,6 @@ 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,

View file

@ -28116,7 +28116,7 @@ export interface components {
maximum_autorouter_session_retention_period?: string | null;
/**
* Maximum Daily Tag Spend Retention Period
* @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.
* @description Maximum retention period for per-day tag spend aggregate rows (e.g., '90d'). Rows whose day is older than this are deleted by the spend log cleanup job, on that job's schedule. Unset means rows are never deleted. Only historical tag usage analytics are affected; tag budgets read the lifetime counter.
*/
maximum_daily_tag_spend_retention_period?: string | null;
/**