feat(proxy): add maximum_daily_tag_spend_retention_period cleanup setting (#39221)

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

* chore(ui): regenerate schema.d.ts for new retention setting

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

* fix(proxy): schedule the cleanup job when a retention db row lands before the side effects run

A config reload applies the db row to the SettingsStore before _update_general_settings snapshots the previous retention values, so the before/after compare saw no change and a retention period first set through /config/update never scheduled the cleanup job. Also reschedule when the job is missing but a retention period is set

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): accept list-valued top-level keys in the base integration proxy config

The shared tests/integration/proxy_config.yaml now carries list-valued top-level keys, so the retention config helper validates only the mapping it merges into. Also drops a SQL-shape assertion from the unit test in favor of the behavioral cutoff-day check

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): cover runtime update, invalid value, independent horizons and worker loss for daily tag spend retention

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): restore the shared retention setting, capture seeded days once and kill a listening worker

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): retry a failed cleanup schedule only when its settings change

_apply_retention_settings rescheduled whenever retention was set and no job existed, so an unparseable cleanup cron was retried on every config reload. Remember the last attempted retention, cron and interval tuple and retry only when it differs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): give daily tag spend retention tests a 240s timeout

Each node boots a proxy and waits for a whole-minute cleanup cron tick, so the global 90s pytest-timeout can expire during teardown on a slow runner, as integration-accounting did on pipeline 90302

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): reschedule cleanup when only the cron or interval changes at runtime

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): reschedule cleanup when the first db sync changes only the cron or interval

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(ui): add text input for String general settings so retention periods can be set from the Admin UI

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): record a cleanup schedule attempt only after it did not raise

Records _last_cleanup_schedule_attempt after _reschedule_spend_log_cleanup_job returns, so a transient add_job error is retried on the next config sync while an invalid cron, which is caught and logged inside the reschedule, is still attempted once per settings value

Also adds --num_workers 2 to the dev proxy command in AGENTS.md as requested on the PR

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* docs: revert unrelated AGENTS.md dev command change

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* Revert "docs: revert unrelated AGENTS.md dev command change"

This reverts commit 047706a623.

* Revert "fix(proxy): record a cleanup schedule attempt only after it did not raise"

This reverts commit 678f7c72b4.

* Revert "feat(ui): add text input for String general settings so retention periods can be set from the Admin UI"

This reverts commit 24e49d71d7.

* fix(proxy): record a cleanup schedule attempt only after it did not raise

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): validate the cleanup schedule before swapping the job and leave startup registration to the startup block

_reschedule_spend_log_cleanup_job builds the new trigger first and only touches the live job once it parsed, so an invalid cron or interval (including a non string value) keeps the previous schedule running instead of removing it. An error raised while rescheduling is logged and retried on the next sync, so it no longer stops the rest of the general settings sync. _apply_retention_settings skips the job-missing path while the scheduler is still stopped, so the startup block is the only registration before start and the cross-replica stagger it applies to pending jobs survives

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): skip cleanup rescheduling while the scheduler is stopped and retry a failed replacement

The stopped-scheduler guard only covered the missing-job path, so the first DB sync (which runs before the startup block) still registered the cleanup job whenever the DB schedule differed from yaml, and startup then replaced it. Every runtime path now defers to the startup block while the scheduler is stopped.

A raised add_job that was replacing a live job was never retried because the live job kept wants_job == has_job; the sync now remembers the failure and retries on the next sync until the schedule is applied.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(proxy): drop redundant docstring on _spend_log_cleanup_trigger

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): schedule DB-only retention at boot and log overflowing cleanup intervals once

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(proxy): drop explanatory comment from startup cleanup block

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): reject a non-string cleanup cron at startup and drop legacy covers markers

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): reschedule spend log cleanup when the reload path already applied a DB cron or interval edit

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): assert cleanup scheduling on a real paused scheduler instead of mock call counts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-09-26 15:07:56 -07:00 • committed by GitHub
parent e47b1f2a3f
commit 1474ea53e6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 884 additions and 93 deletions

View file

@ -2983,6 +2983,14 @@ 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 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(
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.",

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:
@ -278,7 +289,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_date: Cutoff, deadline: float
) -> int | None:
"""
Run one delete batch under a Postgres statement and lock timeout.
@ -301,7 +312,7 @@ class SpendLogCleanup:
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_date: Cutoff, table_name: str, time_column: str, deadline: float
) -> int | None:
"""
Count expired rows still outstanding, stopping at a cap.
@ -314,7 +325,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::{_cutoff_cast(cutoff_date)}
LIMIT $2
) capped
"""
@ -332,7 +343,7 @@ class SpendLogCleanup:
async def _delete_old_rows_batched(
self,
prisma_client: PrismaClient,
cutoff_date: datetime,
cutoff_date: Cutoff,
table_name: str,
key_columns: tuple[str, ...],
time_column: str,
@ -350,7 +361,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::{_cutoff_cast(cutoff_date)}
LIMIT $2
)
"""
@ -406,7 +417,7 @@ class SpendLogCleanup:
run_count,
consecutive_failures,
self.batch_size,
cutoff_date.isoformat(),
_cutoff_text(cutoff_date),
total_deleted,
type(batch_exc).__name__,
batch_exc,
@ -454,7 +465,7 @@ class SpendLogCleanup:
async def _finish_table(
self,
prisma_client: PrismaClient,
cutoff_date: datetime,
cutoff_date: Cutoff,
table_name: str,
time_column: str,
rows_deleted: int,
@ -541,6 +552,18 @@ class SpendLogCleanup:
deadline=deadline,
)
async def _delete_old_daily_tag_spend_rows(
self, prisma_client: PrismaClient, cutoff_day: str, deadline: float
) -> TableCleanupResult:
return await self._delete_old_rows_batched(
prisma_client,
cutoff_day,
table_name="LiteLLM_DailyTagSpend",
key_columns=("id",),
time_column="date",
deadline=deadline,
)
async def _clean_spend_log_tables(
self, prisma_client: PrismaClient, deadline: float
) -> tuple[TableCleanupResult, ...]:
@ -624,6 +647,18 @@ class SpendLogCleanup:
)
return (health_checks_result,)
async def _clean_daily_tag_spend(
self, prisma_client: PrismaClient, retention_seconds: int, deadline: float
) -> tuple[TableCleanupResult, ...]:
"""
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:
"""
@ -671,10 +706,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
@ -706,6 +745,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 = (
@ -716,8 +756,13 @@ 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)
)
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(
@ -732,13 +777,18 @@ class SpendLogCleanup:
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 ()
)
results: Final = spend_log_results + session_results + health_check_results
results: Final = spend_log_results + session_results + health_check_results + daily_tag_spend_results
outcome: Final = self._run_outcome(results)
SpendLogCleanupMetrics.record_run(outcome)
self._log_run_summary(outcome, results, time.monotonic() - run_started_at)

View file

@ -213,6 +213,8 @@ try:
import orjson
import yaml
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.schedulers.base import STATE_STOPPED
from apscheduler.triggers.base import BaseTrigger
from apscheduler.triggers.interval import IntervalTrigger
except ImportError as e:
raise ImportError(f"Missing dependency {e}. Run `pip install 'litellm[proxy]'`")
@ -5078,6 +5080,20 @@ def _current_general_settings() -> Mapping[str, object]:
return general_settings
_CLEANUP_SCHEDULE_KEYS: Final = (
"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_cron",
"maximum_spend_logs_retention_interval",
)
def _cleanup_schedule_of(settings: Mapping[str, object]) -> tuple[object, ...]:
return tuple(settings.get(key) for key in _CLEANUP_SCHEDULE_KEYS)
@lru_cache(maxsize=4096)
def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None:
verbose_proxy_logger.warning(
@ -5100,6 +5116,8 @@ class ProxyConfig:
self._last_websearch_interception_config: dict[str, object] | None = None
self._last_hashicorp_vault_config: dict[str, object] | None = None
self._last_cyberark_config: dict[str, object] | None = None # mutable-ok: change-detection cache
self._last_cleanup_schedule_attempt: tuple[object, ...] | None = None
self._cleanup_reschedule_failed: bool = False
self._cyberark_boot_env: dict[str, str | None] | None = None # mutable-ok: deployment env snapshot, set once
self.worker_registry: list[WorkerRegistryEntry] = []
self.config_sync_subscriber: ConfigSyncSubscriber | None = None
@ -7455,69 +7473,67 @@ class ProxyConfig:
if scheduler is None:
return
# Remove existing job if it exists
try:
scheduler.remove_job("spend_log_cleanup_job")
verbose_proxy_logger.info("Removed existing spend log cleanup job")
except Exception:
pass # Job might not exist, which is fine
# Schedule new job if retention period is set (not None)
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:
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SpendLogCleanup,
wants_job: Final = any(
general_settings.get(key) is not None
for key in (
"maximum_spend_logs_retention_period",
"maximum_autorouter_session_retention_period",
"maximum_health_check_retention_period",
"maximum_daily_tag_spend_retention_period",
)
)
if not wants_job:
if scheduler.get_job("spend_log_cleanup_job") is not None:
scheduler.remove_job("spend_log_cleanup_job")
verbose_proxy_logger.info("Removed existing spend log cleanup job")
return
spend_log_cleanup: Final = SpendLogCleanup()
cleanup_cron: Final = general_settings.get("maximum_spend_logs_cleanup_cron")
trigger: Final = self._spend_log_cleanup_trigger()
if trigger is None:
return
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SpendLogCleanup,
)
if cleanup_cron:
from apscheduler.triggers.cron import CronTrigger
scheduler.add_job(
SpendLogCleanup().cleanup_old_spend_logs,
trigger,
args=[prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log cleanup rescheduled with trigger: %s", trigger)
try:
cron_trigger: Final = CronTrigger.from_crontab(cleanup_cron)
scheduler.add_job(
spend_log_cleanup.cleanup_old_spend_logs,
cron_trigger,
args=[prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log cleanup rescheduled with cron: %s", cleanup_cron)
except ValueError:
verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron)
else:
# Interval-based scheduling (existing behavior)
from litellm.litellm_core_utils.duration_parser import (
duration_in_seconds,
)
def _spend_log_cleanup_trigger(self) -> BaseTrigger | None:
cleanup_cron: Final[object] = general_settings.get("maximum_spend_logs_cleanup_cron")
if cleanup_cron:
from apscheduler.triggers.cron import CronTrigger
retention_interval: Final = general_settings.get("maximum_spend_logs_retention_interval", "1d")
try:
interval_seconds: Final = duration_in_seconds(retention_interval)
# this runs against a started scheduler, which the startup stagger sweep
# cannot reach, so the offset is applied here or the job reconverges across
# replicas the first time an admin edits the retention settings
scheduler.add_job(
spend_log_cleanup.cleanup_old_spend_logs,
stagger_trigger(
job_id="spend_log_cleanup_job",
trigger=IntervalTrigger(seconds=interval_seconds),
period_seconds=interval_seconds,
settings=parse_stagger_settings(general_settings),
),
args=[prisma_client],
id="spend_log_cleanup_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log cleanup rescheduled with interval: %s", retention_interval)
except ValueError:
verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value")
try:
cron_trigger: Final[BaseTrigger] = CronTrigger.from_crontab(cleanup_cron)
except (ValueError, TypeError, AttributeError):
verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron)
return None
return cron_trigger
retention_interval: Final[object] = general_settings.get("maximum_spend_logs_retention_interval", "1d")
if not isinstance(retention_interval, str):
verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value: %r", retention_interval)
return None
# this runs against a started scheduler, which the startup stagger sweep
# cannot reach, so the offset is applied here or the job reconverges across
# replicas the first time an admin edits the retention settings
try:
interval_seconds: Final = duration_in_seconds(retention_interval)
return stagger_trigger(
job_id="spend_log_cleanup_job",
trigger=IntervalTrigger(seconds=interval_seconds),
period_seconds=interval_seconds,
settings=parse_stagger_settings(general_settings),
)
except (ValueError, OverflowError):
verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value: %r", retention_interval)
return None
async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | None) -> None:
global general_settings
@ -7526,32 +7542,28 @@ class ProxyConfig:
if not isinstance(general_settings, SettingsStore):
self.settings.load_yaml(_as_settings_mapping(general_settings))
cache_size_was_db: Final = self.settings.source("user_api_key_cache_max_size") == "db"
previous_retention_values: Final = self._resolved_retention_values()
previous_cleanup_schedule: Final = self._resolved_cleanup_schedule()
previous_pass_through_endpoints: Final = self.settings.get("pass_through_endpoints")
self.settings.apply_db_row("general_settings", db_general_settings)
_bind_general_settings_store(self.settings)
await self._apply_general_settings_side_effects(
db_general_settings,
cache_size_was_db,
previous_retention_values,
previous_cleanup_schedule,
previous_pass_through_endpoints,
)
def _resolved_retention_values(self) -> tuple[SettingsJsonValue | None, ...]:
return tuple(
self.settings.get(key)
for key in (
"maximum_spend_logs_retention_period",
"maximum_autorouter_session_retention_period",
"maximum_health_check_retention_period",
)
)
def _resolved_cleanup_schedule(self) -> tuple[object, ...]:
return _cleanup_schedule_of(self.settings)
def record_cleanup_schedule_attempt(self, settings: Mapping[str, object]) -> None:
self._last_cleanup_schedule_attempt = _cleanup_schedule_of(settings)
async def _apply_general_settings_side_effects(
self,
db_values: Mapping[str, SettingsJsonValue],
cache_size_was_db: bool,
previous_retention_values: tuple[SettingsJsonValue | None, ...],
previous_cleanup_schedule: tuple[object, ...],
previous_pass_through_endpoints: SettingsJsonValue | None,
) -> None:
effects: Final = (
@ -7560,7 +7572,7 @@ class ProxyConfig:
self._apply_boolean_settings,
partial(self._apply_cache_size_setting, cache_size_was_db=cache_size_was_db),
self._apply_store_model_in_db_setting,
partial(self._apply_retention_settings, previous_retention_values=previous_retention_values),
partial(self._apply_retention_settings, previous_cleanup_schedule=previous_cleanup_schedule),
self._apply_ssrf_settings,
)
for effect in effects:
@ -7655,10 +7667,36 @@ class ProxyConfig:
async def _apply_retention_settings(
self,
db_values: Mapping[str, SettingsJsonValue],
previous_retention_values: tuple[SettingsJsonValue | None, ...],
previous_cleanup_schedule: tuple[object, ...],
) -> None:
if previous_retention_values != self._resolved_retention_values():
# while the scheduler is still stopped the startup block owns the first registration
if scheduler is not None and scheduler.state == STATE_STOPPED:
return
schedule: Final = self._resolved_cleanup_schedule()
wants_job: Final = any(value is not None for value in schedule[:4])
has_job: Final = scheduler is not None and scheduler.get_job("spend_log_cleanup_job") is not None
baseline: Final = (
self._last_cleanup_schedule_attempt
if has_job and self._last_cleanup_schedule_attempt is not None
else previous_cleanup_schedule
)
retry_due: Final = (
wants_job
and (not has_job or self._cleanup_reschedule_failed)
and schedule != self._last_cleanup_schedule_attempt
)
if not (baseline != schedule or retry_due or (has_job and not wants_job)):
return
try:
await self._reschedule_spend_log_cleanup_job()
except Exception as exc:
self._cleanup_reschedule_failed = True
verbose_proxy_logger.exception(
"Spend log cleanup could not be rescheduled, will retry on next sync: %s", exc
)
return
self._cleanup_reschedule_failed = False
self._last_cleanup_schedule_attempt = schedule
async def _apply_ssrf_settings(self, db_values: Mapping[str, SettingsJsonValue]) -> None:
_apply_ssrf_general_settings(db_values)
@ -10535,15 +10573,19 @@ class ProxyStartupEvent:
)
### SPEND LOG CLEANUP ###
cleanup_settings: Final = _current_general_settings()
if (
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
cleanup_settings.get("maximum_spend_logs_retention_period") is not None
or cleanup_settings.get("maximum_autorouter_session_retention_period") is not None
or cleanup_settings.get("maximum_health_check_retention_period") is not None
or cleanup_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")
cleanup_cron: Final = cleanup_settings.get("maximum_spend_logs_cleanup_cron")
if cleanup_cron:
if cleanup_cron and not isinstance(cleanup_cron, str):
verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %r", cleanup_cron)
elif isinstance(cleanup_cron, str) and cleanup_cron:
from apscheduler.triggers.cron import CronTrigger
try:
@ -10561,8 +10603,10 @@ class ProxyStartupEvent:
verbose_proxy_logger.error("Invalid maximum_spend_logs_cleanup_cron value: %s", cleanup_cron)
else:
# Interval-based scheduling (existing behavior)
retention_interval: Final = general_settings.get("maximum_spend_logs_retention_interval", "1d")
retention_interval: Final = cleanup_settings.get("maximum_spend_logs_retention_interval", "1d")
try:
if not isinstance(retention_interval, str):
raise ValueError(retention_interval)
interval_seconds: Final = duration_in_seconds(retention_interval)
scheduler.add_job(
spend_log_cleanup.cleanup_old_spend_logs,
@ -10573,8 +10617,11 @@ class ProxyStartupEvent:
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
except ValueError:
verbose_proxy_logger.error("Invalid maximum_spend_logs_retention_interval value")
except (ValueError, OverflowError):
verbose_proxy_logger.error(
"Invalid maximum_spend_logs_retention_interval value: %r", retention_interval
)
proxy_config.record_cleanup_schedule_attempt(cleanup_settings)
### CHECK BATCH COST ###
if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
try:
@ -17909,6 +17956,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,247 @@
import json
import os
import signal
import uuid
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Final
import psutil
import psycopg
import pytest
import yaml
from pydantic import JsonValue, TypeAdapter
from tests.integration._support.client import Gateway, eventually, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process
CLEANUP_EVERY_MINUTE: Final = "* * * * *"
RETENTION_SETTING: Final = "maximum_daily_tag_spend_retention_period"
_MAPPING: Final = TypeAdapter(dict[str, JsonValue])
_SETTINGS: Final = TypeAdapter(list[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 _stored_retention_setting() -> JsonValue:
rows: Final = read_rows(
'SELECT param_value -> %s AS value FROM "LiteLLM_Config" WHERE param_name = %s',
(RETENTION_SETTING, "general_settings"),
)
return rows[0]["value"] if rows else None
def _store_retention_setting(value: JsonValue) -> None:
with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection:
if value is None:
connection.execute(
'UPDATE "LiteLLM_Config" SET param_value = param_value - %s WHERE param_name = %s',
(RETENTION_SETTING, "general_settings"),
)
return
connection.execute(
'UPDATE "LiteLLM_Config" SET param_value = jsonb_set(param_value, ARRAY[%s], %s::jsonb) '
"WHERE param_name = %s",
(RETENTION_SETTING, json.dumps(value), "general_settings"),
)
def _listening_workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]:
port: Final = owned.gateway.client.base_url.port
return tuple(
child
for child in psutil.Process(owned.process.pid).children(recursive=True)
if any(conn.status == psutil.CONN_LISTEN and conn.laddr.port == port for conn in child.net_connections("inet"))
)
def _listed_retention_value(gateway: Gateway) -> JsonValue:
listed: Final = _SETTINGS.validate_json(
gateway.request("GET", "/config/list", params={"config_type": "general_settings"}).content
)
matching: Final = tuple(entry for entry in listed if entry["field_name"] == RETENTION_SETTING)
return matching[0]["field_value"] if matching else "not listed"
def _completion_id(gateway: Gateway, model: str) -> str:
return string_value(gateway.chat(model, text=f"retention audit {uuid.uuid4().hex}")["id"])
def _cleanup_config(tmp_path: Path, retention: dict[str, JsonValue]) -> Path:
base: Final = _MAPPING.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
config: Final = {
**base,
"general_settings": {
**_MAPPING.validate_python(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.timeout(240)
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.timeout(240)
def test_config_update_turns_on_daily_tag_spend_cleanup_without_a_restart(gateway: Gateway, tmp_path: Path) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
expired, yesterday_of_cutoff, on_the_cutoff, today = _day(200), _day(31), _day(30), _day(0)
_seed_daily_tag_spend(tag, (expired, yesterday_of_cutoff, on_the_cutoff, today))
previously_stored: Final = _stored_retention_setting()
_store_retention_setting(None)
try:
config: Final = _cleanup_config(tmp_path, {})
with owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as owned, owned.scenario() as scenario:
model: Final = scenario.model()
assert _listed_retention_value(owned) is None
owned.post("/config/update", {"general_settings": {RETENTION_SETTING: "30d"}})
assert _listed_retention_value(owned) == "30d"
remaining: Final = eventually(
lambda: _remaining_days(tag),
lambda days: yesterday_of_cutoff not in days,
seconds=150,
)
assert remaining == (on_the_cutoff, today), remaining
assert _completion_id(owned, model).startswith("chatcmpl-")
finally:
_store_retention_setting(previously_stored)
_delete_daily_tag_spend(tag)
@pytest.mark.timeout(240)
def test_unparseable_daily_tag_spend_retention_deletes_nothing_and_keeps_serving(
gateway: Gateway, tmp_path: Path
) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
request_id: Final = f"integration-retention-{uuid.uuid4().hex}"
expired: Final = _day(200)
_seed_daily_tag_spend(tag, (expired,))
_seed_old_spend_log(request_id, days_ago=200)
try:
config: Final = _cleanup_config(
tmp_path, {RETENTION_SETTING: "soon", "maximum_spend_logs_retention_period": "30d"}
)
with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario:
model: Final = scenario.model()
eventually(lambda: _spend_log_present(request_id), lambda present: not present, seconds=150)
assert _remaining_days(tag) == (expired,)
assert _completion_id(owned, model).startswith("chatcmpl-")
finally:
_delete_daily_tag_spend(tag)
@pytest.mark.timeout(240)
def test_daily_tag_spend_keeps_days_the_shorter_spend_log_horizon_already_pruned(
gateway: Gateway, tmp_path: Path
) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
request_id: Final = f"integration-retention-{uuid.uuid4().hex}"
expired, inside_tag_horizon = _day(200), _day(60)
_seed_daily_tag_spend(tag, (expired, inside_tag_horizon))
_seed_old_spend_log(request_id, days_ago=60)
try:
config: Final = _cleanup_config(
tmp_path, {RETENTION_SETTING: "90d", "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)
remaining: Final = eventually(lambda: _remaining_days(tag), lambda days: expired not in days, seconds=150)
assert remaining == (inside_tag_horizon,), remaining
finally:
_delete_daily_tag_spend(tag)
@pytest.mark.timeout(240)
def test_daily_tag_spend_cleanup_completes_after_one_of_two_workers_is_killed(gateway: Gateway, tmp_path: Path) -> None:
tag: Final = f"integration-retention-{uuid.uuid4().hex}"
expired, today = _day(200), _day(0)
_seed_daily_tag_spend(tag, (expired, today))
try:
config: Final = _cleanup_config(tmp_path, {RETENTION_SETTING: "30d"})
with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned:
with owned.gateway.scenario() as scenario:
model: Final = scenario.model()
workers: Final = eventually(
lambda: _listening_workers(owned), lambda found: len(found) == 2, seconds=30
)
workers[0].send_signal(signal.SIGKILL)
eventually(lambda: workers[0].is_running(), lambda alive: not alive, seconds=10)
ids: Final = tuple(_completion_id(owned.gateway, model) for _ in range(6))
assert len(set(ids)) == 6 and all(identity.startswith("chatcmpl-") for identity in ids), ids
remaining: Final = eventually(
lambda: _remaining_days(tag), lambda days: expired not in days, seconds=150
)
assert remaining == (today,), remaining
finally:
_delete_daily_tag_spend(tag)
@pytest.mark.timeout(240)
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}"
expired: Final = _day(200)
_seed_daily_tag_spend(tag, (expired,))
_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) == (expired,)
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

@ -3862,6 +3862,7 @@ async def test_ProxyConfig__reschedule_spend_log_cleanup_job_health_check_retent
async def test_ProxyConfig__update_general_settings_updates_health_check_retention(monkeypatch):
settings = {}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", settings)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", MagicMock(**{"get_job.return_value": None}))
pc = ProxyConfig()
reschedule = AsyncMock()
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
@ -3872,6 +3873,329 @@ async def test_ProxyConfig__update_general_settings_updates_health_check_retenti
reschedule.assert_awaited_once()
def _paused_scheduler(monkeypatch):
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
return real_scheduler
def _scheduler_whose_first_add_job_raises(monkeypatch):
from apscheduler.schedulers.asyncio import AsyncIOScheduler
class FirstAddJobRaises(AsyncIOScheduler):
raised = False
def add_job(self, *args, **kwargs):
if not self.raised:
self.raised = True
raise RuntimeError("scheduler busy")
return super().add_job(*args, **kwargs)
real_scheduler = FirstAddJobRaises()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
return real_scheduler
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_daily_tag_spend_retention(monkeypatch):
real_scheduler = _paused_scheduler(monkeypatch)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"maximum_daily_tag_spend_retention_period": "90d"},
)
pc = ProxyConfig()
try:
await pc._reschedule_spend_log_cleanup_job()
job = real_scheduler.get_job("spend_log_cleanup_job")
assert job is not None, "daily tag spend retention alone did not schedule the cleanup job"
assert job.func.__name__ == "cleanup_old_spend_logs"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_updates_daily_tag_spend_retention(monkeypatch):
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
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"
assert real_scheduler.get_job("spend_log_cleanup_job") is not None, "runtime retention did not schedule cleanup"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_schedules_cleanup_when_db_row_was_already_applied(monkeypatch):
"""A config reload applies the db row to the store before the side effects run, so the
before/after snapshot is equal; the job must still be scheduled when none is running."""
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
pc.settings.apply_db_row("general_settings", {"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
assert real_scheduler.get_job("spend_log_cleanup_job") is not None, "DB-only retention never scheduled cleanup"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_retries_a_failed_schedule_once_per_settings_value(
monkeypatch, caplog
):
"""An unparseable cron leaves no job behind; reloads must not retry it every tick, only when the
cron or a retention value changes."""
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
bad_cron = {"maximum_daily_tag_spend_retention_period": "90d", "maximum_spend_logs_cleanup_cron": "not a cron"}
pc.settings.apply_db_row("general_settings", bad_cron)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
for _ in range(3):
await pc._update_general_settings(bad_cron)
assert real_scheduler.get_job("spend_log_cleanup_job") is None
cron_errors = [r for r in caplog.records if "maximum_spend_logs_cleanup_cron" in r.getMessage()]
assert len(cron_errors) == 1, f"invalid cron was retried on every reload: {len(cron_errors)} error lines"
await pc._update_general_settings({**bad_cron, "maximum_spend_logs_cleanup_cron": "* * * * *"})
job = real_scheduler.get_job("spend_log_cleanup_job")
assert job is not None, "a corrected cron did not schedule cleanup"
assert "minute='*'" in str(job.trigger)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_retries_a_schedule_that_raised(monkeypatch):
"""A transient add_job failure must not be remembered as a completed attempt; the next
reload with the same settings tries again."""
real_scheduler = _scheduler_whose_first_add_job_raises(monkeypatch)
pc = ProxyConfig()
retention = {"maximum_daily_tag_spend_retention_period": "90d"}
pc.settings.apply_db_row("general_settings", retention)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings(retention)
assert real_scheduler.get_job("spend_log_cleanup_job") is None
await pc._update_general_settings(retention)
assert real_scheduler.get_job("spend_log_cleanup_job") is not None, "raised add_job was not retried"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_retries_a_failed_replacement_of_the_live_job(monkeypatch):
"""A cron change whose add_job raised keeps the old job running, so the next reload with the
same settings must try the replacement again instead of leaving the new cron unapplied."""
real_scheduler = _scheduler_whose_first_add_job_raises(monkeypatch)
pc = ProxyConfig()
pc.settings.load_yaml({"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
real_scheduler.raised = True
await pc._reschedule_spend_log_cleanup_job()
real_scheduler.raised = False
try:
new_cron = {"maximum_spend_logs_cleanup_cron": "0 3 * * *"}
await pc._update_general_settings(new_cron)
assert "hour='3'" not in str(real_scheduler.get_job("spend_log_cleanup_job").trigger), "old job was lost"
await pc._update_general_settings(new_cron)
assert "hour='3'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger), (
"failed replacement was not retried on the next sync"
)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_leaves_a_changed_db_schedule_to_startup_while_scheduler_is_stopped(
monkeypatch,
):
"""The first DB sync runs before the scheduler starts and usually differs from the yaml; it
must still leave registration to the startup block instead of adding a job it will replace."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
assert real_scheduler.get_jobs() == [], "DB sync registered the cleanup job before the scheduler started"
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_leaves_first_registration_to_startup_while_scheduler_is_stopped(
monkeypatch,
):
"""The DB sync that runs before the scheduler starts must not register the cleanup job; the
startup block does, once, so the cross-replica stagger it applies to pending jobs survives."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
pc.settings.load_yaml({"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
await pc._update_general_settings({"unrelated_key": "value"})
assert real_scheduler.get_jobs() == [], "DB sync registered the cleanup job before the scheduler started"
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_runtime_interval_job_carries_the_stagger_offset(monkeypatch):
"""Once the scheduler is running the sync owns registration and the job it adds is staggered."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.common_utils.scheduled_job_stagger import _OffsetTrigger
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
jobs = real_scheduler.get_jobs()
assert [job.id for job in jobs] == ["spend_log_cleanup_job"]
assert isinstance(jobs[0].trigger, _OffsetTrigger), repr(jobs[0].trigger)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"bad_schedule",
[
{"maximum_spend_logs_cleanup_cron": "not a cron"},
{"maximum_spend_logs_cleanup_cron": "0 0 * * * *"},
{"maximum_spend_logs_retention_interval": "soon"},
{"maximum_spend_logs_retention_interval": 86400},
],
)
async def test_ProxyConfig__update_general_settings_keeps_the_live_cleanup_job_when_the_new_schedule_is_invalid(
monkeypatch, bad_schedule
):
"""A schedule edit that does not parse must leave the old cleanup job running and must not
stop the rest of the general settings sync."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
ssrf_sync = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server._apply_ssrf_general_settings", ssrf_sync)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
old_trigger = real_scheduler.get_job("spend_log_cleanup_job").trigger
ssrf_sync.reset_mock()
for _ in range(2):
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d", **bad_schedule})
live_job = real_scheduler.get_job("spend_log_cleanup_job")
assert live_job is not None, "invalid schedule removed the cleanup job"
assert live_job.trigger is old_trigger
assert ssrf_sync.call_count == 2, "schedule error blocked the rest of the settings sync"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_logs_an_overflowing_interval_once(monkeypatch, caplog):
"""An interval that parses but overflows the trigger must keep the live job and log one
error, not a traceback on every sync."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
old_trigger = real_scheduler.get_job("spend_log_cleanup_job").trigger
overflowing = {
"maximum_daily_tag_spend_retention_period": "90d",
"maximum_spend_logs_retention_interval": "99999999999d",
}
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
for _ in range(5):
await pc._update_general_settings(overflowing)
errors = [record for record in caplog.records if record.levelno >= logging.ERROR]
assert len(errors) == 1, [record.getMessage() for record in errors]
assert real_scheduler.get_job("spend_log_cleanup_job").trigger is old_trigger
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_reschedules_when_only_the_cron_changes(monkeypatch):
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
pc.settings.load_yaml({"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
await pc._reschedule_spend_log_cleanup_job()
try:
interval_job = real_scheduler.get_job("spend_log_cleanup_job")
assert interval_job is not None and "hour='3'" not in str(interval_job.trigger)
await pc._update_general_settings({"maximum_spend_logs_cleanup_cron": "0 3 * * *"})
cron_job = real_scheduler.get_job("spend_log_cleanup_job")
assert "hour='3'" in str(cron_job.trigger), "cron-only change did not reschedule"
await pc._update_general_settings({"maximum_spend_logs_cleanup_cron": "0 3 * * *"})
assert real_scheduler.get_job("spend_log_cleanup_job") is cron_job, "unchanged cron replaced the job"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_reschedules_a_cron_edit_the_reload_path_already_applied(
monkeypatch,
):
"""The periodic reload applies the DB row through _update_config_from_db before
_update_general_settings snapshots the previous schedule, so a cron edited in the DB must
still replace the live job's trigger."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
first_row = {"maximum_daily_tag_spend_retention_period": "90d", "maximum_spend_logs_cleanup_cron": "0 3 * * *"}
pc.settings.apply_db_row("general_settings", first_row)
await pc._update_general_settings(first_row)
assert "hour='3'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger)
edited_row = {**first_row, "maximum_spend_logs_cleanup_cron": "0 5 * * *"}
pc.settings.apply_db_row("general_settings", edited_row)
await pc._update_general_settings(edited_row)
assert "hour='5'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger), "DB cron edit was ignored"
pc.settings.apply_db_row("general_settings", edited_row)
await pc._update_general_settings(edited_row)
assert "hour='5'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger)
finally:
real_scheduler.shutdown(wait=False)
# ---------------------------------------------------------------------------
# ProxyConfig._update_general_settings
# ---------------------------------------------------------------------------
@ -4003,6 +4327,7 @@ async def test_ProxyConfig__update_general_settings_skips_redundant_retention_re
pc = ProxyConfig()
reschedule: Final = AsyncMock()
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "scheduler", MagicMock())
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"maximum_health_check_retention_period": "30d"})
@ -4021,6 +4346,7 @@ async def test_ProxyConfig__update_general_settings_reschedules_after_retention_
pc = ProxyConfig()
reschedule: Final = AsyncMock()
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "scheduler", MagicMock(**{"get_job.return_value": None}))
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"maximum_health_check_retention_period": "30d"})
@ -4052,7 +4378,7 @@ async def test_ProxyConfig__update_general_settings_dispatches_every_side_effect
if name == "_apply_cache_size_setting":
handler.assert_awaited_once_with({}, cache_size_was_db=False)
elif name == "_apply_retention_settings":
handler.assert_awaited_once_with({}, previous_retention_values=())
handler.assert_awaited_once_with({}, previous_cleanup_schedule=())
elif name == "_apply_pass_through_settings":
handler.assert_awaited_once_with({}, previous_endpoints=None)
else:

View file

@ -935,6 +935,89 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat
scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_registers_cleanup_when_retention_lives_only_in_the_db(monkeypatch):
"""With no config file, the startup DB sync rebinds general_settings to a store holding the
retention period; the cleanup job must be registered from that live value, not the stale
empty dict the caller passed in."""
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
mock_proxy_logging = MagicMock(spec=ProxyLogging)
mock_proxy_logging.slack_alerting_instance = MagicMock()
mock_proxy_logging.db_spend_update_writer = MagicMock()
mock_proxy_config = _mock_scheduled_proxy_config()
db_settings = proxy_server_module.ProxyConfig().settings
db_settings.apply_db_row("general_settings", {"maximum_daily_tag_spend_retention_period": "30d"})
async def sync_from_db(*args: object, **kwargs: object) -> None:
proxy_server_module._bind_general_settings_store(db_settings)
mock_proxy_config.add_deployment.side_effect = sync_from_db
scheduler = AsyncIOScheduler()
try:
with (
patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings={},
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
assert scheduler.get_job("spend_log_cleanup_job") is not None, "DB-only retention was not scheduled at boot"
finally:
scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_does_not_fall_back_to_the_interval_for_a_non_string_cron(monkeypatch):
"""A truthy non-string cron is invalid, so startup must log it and register no cleanup job
rather than silently pruning on the default interval the admin never configured."""
monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False)
monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False)
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.proxy_server import ProxyStartupEvent
from litellm.proxy.utils import ProxyLogging
mock_prisma_client = MagicMock()
mock_proxy_logging = MagicMock(spec=ProxyLogging)
mock_proxy_logging.slack_alerting_instance = MagicMock()
mock_proxy_logging.db_spend_update_writer = MagicMock()
settings = {"maximum_daily_tag_spend_retention_period": "30d", "maximum_spend_logs_cleanup_cron": 5}
scheduler = AsyncIOScheduler()
try:
with (
patch("litellm.proxy.proxy_server.proxy_config", _mock_scheduled_proxy_config()),
patch("litellm.proxy.proxy_server.store_model_in_db", False),
patch("litellm.proxy.proxy_server.general_settings", settings),
patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler),
):
await ProxyStartupEvent.initialize_scheduled_background_jobs(
general_settings=settings,
prisma_client=mock_prisma_client,
proxy_budget_rescheduler_min_time=1,
proxy_budget_rescheduler_max_time=2,
proxy_batch_write_at=5,
proxy_logging_obj=mock_proxy_logging,
)
assert scheduler.get_job("spend_log_cleanup_job") is None, "invalid cron fell back to the interval"
finally:
scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval(monkeypatch):
"""

View file

@ -827,6 +827,29 @@ 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_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
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]
cutoff_day = client.db.execute_raw.call_args[0][1]
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
async def test_each_retention_key_cuts_off_at_its_own_horizon():
client = _mock_prisma_for_retention([0, 0, 0, 0, 0])

View file

@ -28325,6 +28325,11 @@ export interface components {
* @description Maximum retention period for auto-router benchmark session rollup rows (e.g., '365d'). Rows whose last turn is older than this are deleted by the spend log cleanup job, on that job's schedule. Unset means rollup rows are never deleted.
*/
maximum_autorouter_session_retention_period?: string | null;
/**
* Maximum Daily Tag Spend Retention Period
* @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;
/**
* Maximum Health Check Retention Period
* @description Maximum retention period for health-check rows (e.g., '30d'). Rows whose checked_at is older than this are deleted by the spend log cleanup job, on that job's schedule. Unset means rows are never deleted. Set this well above health_check_interval because /health and the UI read the latest row per model.