mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
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:
parent
ae51fb266d
commit
c7af908e89
8 changed files with 192 additions and 82 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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 ()
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
103
tests/integration/spend/test_daily_tag_spend_retention.py
Normal file
103
tests/integration/spend/test_daily_tag_spend_retention.py
Normal 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)
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue