feat(proxy): run spend log cleanup from a DB-backed scheduled job table (#50131)

* feat(proxy): run spend log cleanup from a DB-backed scheduled job table

Spend log retention ran from an in-process APScheduler timer that reset on
every restart and needed Redis for leader election, so pods that restart
more often than the interval never cleaned up at all.

Add LiteLLM_ScheduledJobRun, a generic table with one row per job run.
Every pod polls it, plans the next run idempotently and claims a due run
under a transaction-scoped advisory lock, so exactly one pod runs each
occurrence. Runs carry a lease, a status and a log. Intervals count from
the last finished run and survive restarts. Schedule edits cancel and
replace the pending run. The retention and size-check jobs share a lock
group so they never overlap, and the Redis pod lock is gone.

* fix(proxy): plan an earlier cron edit from now, and cover the job runner

Moving the cleanup cron earlier queued tomorrow's run because the next
occurrence was planned after the pending run being cancelled. Plan from
finished runs only.

Register LiteLLM_ScheduledJobRun as a span table, keep the advisory lock
query inside the tick's span, move the non-string cron startup test to
the schedule parser, and add unit tests for the runner and job wiring.

* test(proxy): annotate new spend log cleanup test locals as Final

* test(proxy): type the parameters of the new spend log cleanup tests

* test(proxy): leave an hour of margin when moving a cron earlier so an hour boundary cannot flake it

* fix(proxy): reject a cron that never fires instead of failing every tick

A crontab like 0 0 31 2 * parses but has no next fire time, so next_run_at raised on every tick and the job was never scheduled again. cron_schedule now returns InvalidSchedule for it, which logs and keeps the pending run like any other bad config

* refactor(proxy): drop the informational log column from LiteLLM_ScheduledJobRun

Nothing read the column back. Run outcomes stay in status, and the reasons
already go to the proxy logs, so JobResult becomes a SUCCEEDED or FAILED
literal and the cleanup returns its RunOutcome instead of a summary report
This commit is contained in:
jesus-berri 2026-11-11 18:22:02 -08:00 • committed by GitHub
parent e7fd0d6cba
commit d40dedf23b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 1417 additions and 570 deletions

View file

@ -0,0 +1,23 @@
SET LOCAL lock_timeout = '10s';
CREATE TYPE "ScheduledJobRunStatus" AS ENUM ('PENDING', 'RUNNING', 'SUCCEEDED', 'FAILED', 'CANCELLED');
CREATE TABLE IF NOT EXISTS "LiteLLM_ScheduledJobRun" (
"id" TEXT NOT NULL,
"job_name" TEXT NOT NULL,
"lock_group" TEXT NOT NULL,
"schedule" TEXT NOT NULL,
"scheduled_for" TIMESTAMPTZ(3) NOT NULL,
"status" "ScheduledJobRunStatus" NOT NULL DEFAULT 'PENDING',
"claimed_by" TEXT,
"claimed_at" TIMESTAMPTZ(3),
"lease_expires_at" TIMESTAMPTZ(3),
"finished_at" TIMESTAMPTZ(3),
"created_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_ScheduledJobRun_pkey" PRIMARY KEY ("id")
);
CREATE INDEX IF NOT EXISTS "LiteLLM_ScheduledJobRun_job_name_status_idx" ON "LiteLLM_ScheduledJobRun"("job_name", "status");
CREATE INDEX IF NOT EXISTS "LiteLLM_ScheduledJobRun_lock_group_status_idx" ON "LiteLLM_ScheduledJobRun"("lock_group", "status");

View file

@ -1155,6 +1155,31 @@ enum JobStatus {
INACTIVE
}
model LiteLLM_ScheduledJobRun {
id String @id @default(uuid())
job_name String
lock_group String
schedule String
scheduled_for DateTime @client.Timestamptz(3)
status ScheduledJobRunStatus @default(PENDING)
claimed_by String?
claimed_at DateTime? @client.Timestamptz(3)
lease_expires_at DateTime? @client.Timestamptz(3)
finished_at DateTime? @client.Timestamptz(3)
created_at DateTime @default(now()) @client.Timestamptz(3)
@@index([job_name, status])
@@index([lock_group, status])
}
enum ScheduledJobRunStatus {
PENDING
RUNNING
SUCCEEDED
FAILED
CANCELLED
}
model LiteLLM_ManagedFileTable {
id String @id @default(uuid())
unified_file_id String @unique // The base64 encoded unified file ID

View file

@ -1931,6 +1931,10 @@ SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = get_env_float("SPEND_LOG_CLEANU
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = get_env_int("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", 100000)
SPEND_LOG_SIZE_CHECK_JOB_NAME: Final = "spend_log_size_check"
SPEND_LOG_SIZE_CHECK_INTERVAL: Final = os.getenv("SPEND_LOG_SIZE_CHECK_INTERVAL", "1d")
SCHEDULED_JOB_POLL_INTERVAL_SECONDS: Final = get_env_int("SCHEDULED_JOB_POLL_INTERVAL_SECONDS", 30)
SCHEDULED_JOB_TICK_TIMEOUT_SECONDS: Final = get_env_float("SCHEDULED_JOB_TICK_TIMEOUT_SECONDS", 10.0)
SCHEDULED_JOB_LEASE_GRACE_SECONDS: Final = get_env_float("SCHEDULED_JOB_LEASE_GRACE_SECONDS", 300.0)
SCHEDULED_JOB_RUN_HISTORY_LIMIT: Final = get_env_int("SCHEDULED_JOB_RUN_HISTORY_LIMIT", 1000)
SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES: Final = get_env_int("SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES", 10)
SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = get_env_float(
"SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", 5.0

View file

@ -328,6 +328,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_PromptTable",
"LiteLLM_HealthCheckTable",
"LiteLLM_SearchToolsTable",
"LiteLLM_ScheduledJobRun",
"LiteLLM_SSOConfig",
"LiteLLM_ManagedVectorStoreIndexTable",
"LiteLLM_CacheConfig",

View file

@ -2,7 +2,7 @@ import asyncio
import math
import re
import time
from collections.abc import Awaitable, Callable
from collections.abc import Awaitable, Callable, Mapping
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
@ -11,8 +11,8 @@ from typing import Final, Literal, TypeAlias
from pydantic import ConfigDict, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
SCHEDULED_JOB_LEASE_GRACE_SECONDS,
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
SPEND_LOG_CLEANUP_BATCH_SIZE,
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS,
@ -21,6 +21,8 @@ from litellm.constants import (
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
SPEND_LOG_RUN_LOOPS,
SPEND_LOG_SIZE_CHECK_INTERVAL,
SPEND_LOG_SIZE_CHECK_JOB_NAME,
SPEND_LOG_SIZE_CLEANUP_RECHECK_BATCHES,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
@ -35,6 +37,14 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
bounded_tx,
trusted_sql,
)
from litellm.proxy.db.scheduled_job_runs import (
InvalidSchedule,
JobResult,
Schedule,
ScheduledJob,
cron_schedule,
interval_schedule,
)
from litellm.proxy.utils import PrismaClient
from litellm.types.llms.base import LiteLLMBaseModel
@ -395,7 +405,8 @@ class SpendLogCleanup:
When LiteLLM_SpendLogs is range-partitioned, expired data is reclaimed by
dropping whole partitions (instant, frees disk immediately). Otherwise it
falls back to deleting logs in batches.
Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
One pod runs it at a time, through the LiteLLM_ScheduledJobRun claim of SPEND_LOG_RETENTION_JOB and
SPEND_LOG_SIZE_CHECK_JOB, which share a lock group so the two never overlap either.
Every run is bounded so it can never monopolise the database: a wall-clock
budget shared across all tables, a per-table batch cap, and a Postgres
@ -419,7 +430,6 @@ class SpendLogCleanup:
def __init__(
self,
general_settings=None,
redis_cache: RedisCache | None = None,
partition_manager: SpendLogsPartitionManager | None = None,
*,
clock: Callable[[], float] = time.monotonic,
@ -431,10 +441,6 @@ class SpendLogCleanup:
self.general_settings = general_settings or default_settings
self._refresh_bounds()
from litellm.proxy.proxy_server import proxy_logging_obj
pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
self.pod_lock_manager = pod_lock_manager
verbose_proxy_logger.info(
"SpendLogCleanup initialized: batch_size=%s max_batches=%s run_budget=%ss batch_timeout=%ss",
self.batch_size,
@ -1420,31 +1426,26 @@ class SpendLogCleanup:
return
verbose_proxy_logger.warning(message, outcome, elapsed_seconds, per_table, storage)
async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None:
"""
The scheduled retention run: every age rule, plus the size rule when one cron drives both.
If pod_lock_manager is available, ensures only one pod runs cleanup.
If no pod_lock_manager, runs cleanup without distributed locking.
"""
await self._run(prisma_client, age_rules=True, size_rule=self._size_pass_runs_with_age_pass())
async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> RunOutcome:
"""The scheduled retention run: every age rule, plus the size rule when one cron drives both."""
return await self._run(prisma_client, age_rules=True, size_rule=self._size_pass_runs_with_age_pass())
async def cleanup_spend_logs_by_size(self, prisma_client: PrismaClient) -> None:
async def cleanup_spend_logs_by_size(self, prisma_client: PrismaClient) -> RunOutcome:
"""The size-check run scheduled on maximum_spend_logs_size_check_interval: the size rule only."""
await self._run(prisma_client, age_rules=False, size_rule=True)
return await self._run(prisma_client, age_rules=False, size_rule=True)
async def _run(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None:
async def _run(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> RunOutcome:
"""
One pass at a time per pod. The age job and the size job fire from separate schedules, and the pod
lock only exists with Redis, so this is what keeps them apart in a pod without it. The later one waits
rather than skips, because the other job applies a different rule and would not do its work for it.
One pass at a time per pod, so a pass started outside the job claim (a manual call) cannot overlap a
scheduled one in the same pod. The later one waits rather than skips, because the other job applies a
different rule and would not do its work for it.
"""
if _PASS_LOCK.locked():
verbose_proxy_logger.info("Another cleanup pass is running in this pod; waiting for it to finish")
async with _PASS_LOCK:
await self._run_passes(prisma_client, age_rules, size_rule)
return await self._run_passes(prisma_client, age_rules, size_rule)
async def _run_passes(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> None:
lock_acquired = False
async def _run_passes(self, prisma_client: PrismaClient, age_rules: bool, size_rule: bool) -> RunOutcome:
run_started_at: Final = self._clock()
progress: Final = _RunProgress()
progress_token: Final = _run_progress.set(progress)
@ -1471,29 +1472,12 @@ class SpendLogCleanup:
and active_size_rule is None
):
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
return "skipped_disabled"
if delete_spend_logs and self.retention_seconds is None:
verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup")
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
# If we have a pod lock manager, try to acquire the lock
if self.pod_lock_manager and self.pod_lock_manager.redis_cache:
lock_acquired = (
await self.pod_lock_manager.acquire_lock(
cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME,
)
or False
)
verbose_proxy_logger.info(
"Lock acquisition attempt: %s at %s", "successful" if lock_acquired else "failed", datetime.now()
)
if not lock_acquired:
verbose_proxy_logger.info("Another pod is already running cleanup")
SpendLogCleanupMetrics.record_run("skipped_locked")
return
return "skipped_disabled"
deadline: Final = self._clock() + self.run_budget_seconds
size_group: Final = int(active_size_rule is not None)
@ -1583,10 +1567,63 @@ class SpendLogCleanup:
e,
)
SpendLogCleanupMetrics.record_run("aborted")
return # Return after error handling
return "aborted"
else:
return outcome
finally:
_run_progress.reset(progress_token)
# Only release the lock if it was actually acquired
if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache:
await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME)
verbose_proxy_logger.info("Released cleanup lock")
def retention_schedule(settings: Mapping[str, object]) -> Schedule | InvalidSchedule:
cron: Final = settings.get("maximum_spend_logs_cleanup_cron")
if cron:
return cron_schedule("maximum_spend_logs_cleanup_cron", cron)
return interval_schedule(
"maximum_spend_logs_retention_interval", settings.get("maximum_spend_logs_retention_interval", "1d")
)
def size_check_schedule(settings: Mapping[str, object]) -> Schedule | InvalidSchedule:
return interval_schedule(
"maximum_spend_logs_size_check_interval",
settings.get("maximum_spend_logs_size_check_interval", SPEND_LOG_SIZE_CHECK_INTERVAL),
)
def _live_settings() -> Mapping[str, object]:
from litellm.proxy.proxy_server import general_settings
return general_settings # pyright: ignore[reportUnknownVariableType] # general_settings is an untyped module dict
def _cleanup_lease() -> timedelta:
run_budget: Final = SpendLogCleanup(general_settings=_live_settings()).run_budget_seconds
return timedelta(seconds=run_budget + SCHEDULED_JOB_LEASE_GRACE_SECONDS)
def _job_result(outcome: RunOutcome) -> JobResult:
return "FAILED" if outcome == "aborted" else "SUCCEEDED"
async def _run_retention(prisma_client: PrismaClient) -> JobResult:
return _job_result(await SpendLogCleanup().cleanup_old_spend_logs(prisma_client))
async def _run_size_check(prisma_client: PrismaClient) -> JobResult:
return _job_result(await SpendLogCleanup().cleanup_spend_logs_by_size(prisma_client))
SPEND_LOG_RETENTION_JOB: Final = ScheduledJob(
name=SPEND_LOG_CLEANUP_JOB_NAME,
lock_group=SPEND_LOG_CLEANUP_JOB_NAME,
schedule=lambda: retention_schedule(_live_settings()),
lease=_cleanup_lease,
run=_run_retention,
)
SPEND_LOG_SIZE_CHECK_JOB: Final = ScheduledJob(
name=SPEND_LOG_SIZE_CHECK_JOB_NAME,
lock_group=SPEND_LOG_CLEANUP_JOB_NAME,
schedule=lambda: size_check_schedule(_live_settings()),
lease=_cleanup_lease,
run=_run_size_check,
)

View file

@ -0,0 +1,454 @@
import asyncio
import os
import socket
import uuid
from collections.abc import Awaitable, Callable, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from functools import lru_cache
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
from apscheduler.triggers import cron # pyright: ignore[reportMissingTypeStubs, reportUnknownVariableType] # untyped
from pydantic import AwareDatetime, BaseModel, ConfigDict, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
SCHEDULED_JOB_RUN_HISTORY_LIMIT,
SCHEDULED_JOB_TICK_TIMEOUT_SECONDS,
)
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import bounded_tx, trusted_sql
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_TABLE: Final = "LiteLLM_ScheduledJobRun"
_FIRE_TIME: Final = TypeAdapter[datetime](AwareDatetime)
@dataclass(frozen=True, slots=True)
class IntervalSchedule:
seconds: float
@property
def spec(self) -> str:
return f"interval:{self.seconds:g}s"
class FireTimes(Protocol):
def get_next_fire_time(self, previous_fire_time: datetime | None, now: datetime) -> datetime | None: ...
class JobScheduler(Protocol):
"""The slice of APScheduler's AsyncIOScheduler that registers a poll, which ships no type information."""
def get_job(self, job_id: str) -> object | None: ...
def remove_job(self, job_id: str) -> None: ...
def add_job(
self,
func: Callable[["PrismaClient | None"], Awaitable[None]],
trigger: str,
*,
args: Sequence[object],
id: str,
replace_existing: bool,
max_instances: int,
misfire_grace_time: int,
seconds: float,
) -> object: ...
@dataclass(frozen=True, slots=True)
class CronSchedule:
expression: str
trigger: FireTimes
@property
def spec(self) -> str:
return f"cron:{self.expression}"
@dataclass(frozen=True, slots=True)
class InvalidSchedule:
reason: str
Schedule: TypeAlias = IntervalSchedule | CronSchedule
def interval_schedule(setting_name: str, value: object) -> Schedule | InvalidSchedule:
if not isinstance(value, str):
return InvalidSchedule(f"{setting_name} must be a duration string like '1d', got {value!r}")
try:
seconds: Final = duration_in_seconds(value)
timedelta(seconds=seconds)
except (ValueError, OverflowError):
return InvalidSchedule(f"{setting_name} is not a valid duration: {value!r}")
if seconds <= 0:
return InvalidSchedule(f"{setting_name} must be positive, got {value!r}")
return IntervalSchedule(seconds=seconds)
def _cron_trigger(expression: str) -> FireTimes:
return cron.CronTrigger.from_crontab(expression, timezone=timezone.utc) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped
def cron_schedule(setting_name: str, value: object) -> Schedule | InvalidSchedule:
if not isinstance(value, str):
return InvalidSchedule(f"{setting_name} must be a crontab string, got {value!r}")
try:
trigger: Final = _cron_trigger(value)
except (ValueError, TypeError):
return InvalidSchedule(f"{setting_name} is not a valid crontab expression: {value!r}")
if trigger.get_next_fire_time(None, datetime.now(timezone.utc)) is None:
return InvalidSchedule(f"{setting_name} never fires: {value!r}")
return CronSchedule(expression=value, trigger=trigger)
JobResult: TypeAlias = Literal["SUCCEEDED", "FAILED"]
@dataclass(frozen=True, slots=True)
class ScheduledJob:
"""
A job that runs on exactly one pod per occurrence.
Jobs sharing a lock_group never run at the same time, on any pod. schedule and lease are read on every tick,
so a settings change applies without re-registering the job.
"""
name: str
lock_group: str
schedule: Callable[[], Schedule | InvalidSchedule]
lease: Callable[[], timedelta]
run: Callable[["PrismaClient"], Awaitable[JobResult]]
class PendingRun(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
schedule: str
scheduled_for: AwareDatetime
class JobState(BaseModel):
model_config = ConfigDict(frozen=True)
now: AwareDatetime
group_running: bool
job_running: bool
last_finished_at: AwareDatetime | None
last_scheduled_for: AwareDatetime | None
@dataclass(frozen=True, slots=True)
class NewRun:
id: str
schedule: str
scheduled_for: datetime
@dataclass(frozen=True, slots=True)
class TickPlan:
cancel: tuple[str, ...]
insert: NewRun | None
claim: str | None
def next_run_at(schedule: Schedule, state: JobState) -> datetime:
"""
When the next occurrence is due.
An interval counts from the last finished run, so a restart never resets it and a missed run happens once, right
away. A cron fires at its next time after both now and the last occurrence, so an outage collapses into the one
occurrence that was already pending.
"""
match schedule:
case IntervalSchedule(seconds=seconds):
if state.last_finished_at is None:
return state.now
return state.last_finished_at + timedelta(seconds=seconds)
case CronSchedule(trigger=trigger):
anchor: Final = max(state.now, state.last_scheduled_for or state.now)
return _FIRE_TIME.validate_python(trigger.get_next_fire_time(None, anchor + timedelta(microseconds=1)))
def plan_tick(
schedule: Schedule | InvalidSchedule, state: JobState, pending: tuple[PendingRun, ...], new_id: str
) -> TickPlan:
"""
Decide one tick from the job's rows. Pure, so every pod reaches the same decision from the same rows.
A pending row planned from a different schedule is cancelled and replaced, so a settings change applies at once.
With an invalid schedule nothing is replanned, and a pending row from the last valid one still runs.
"""
if isinstance(schedule, InvalidSchedule):
return TickPlan(cancel=(), insert=None, claim=_claimable(pending[:1], state))
current: Final = tuple(run for run in pending if run.schedule == schedule.spec)
keep: Final = current[:1]
cancel: Final = tuple(run.id for run in pending if run not in keep)
if keep or state.job_running:
return TickPlan(cancel=cancel, insert=None, claim=_claimable(keep, state))
new_run: Final = NewRun(id=new_id, schedule=schedule.spec, scheduled_for=next_run_at(schedule, state))
claim: Final = new_run.id if new_run.scheduled_for <= state.now and not state.group_running else None
return TickPlan(cancel=cancel, insert=new_run, claim=claim)
def _claimable(pending: tuple[PendingRun, ...], state: JobState) -> str | None:
if not pending or state.group_running or pending[0].scheduled_for > state.now:
return None
return pending[0].id
def default_worker_id() -> str:
return f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4().hex[:8]}"
class _Claimed(BaseModel):
model_config = ConfigDict(frozen=True)
id: str
job_name: str
_STATE: Final = TypeAdapter(list[JobState])
_PENDING: Final = TypeAdapter(list[PendingRun])
_CLAIMED: Final = TypeAdapter(list[_Claimed])
_REAP_SQL: Final = f"""
UPDATE "{_TABLE}"
SET "status" = 'FAILED', "finished_at" = now()
WHERE "lock_group" = $1 AND "status" = 'RUNNING' AND "lease_expires_at" < now()
RETURNING "id", "job_name"
"""
_STATE_SQL: Final = f"""
SELECT now() AS "now",
EXISTS (SELECT 1 FROM "{_TABLE}" WHERE "lock_group" = $2 AND "status" = 'RUNNING') AS "group_running",
EXISTS (SELECT 1 FROM "{_TABLE}" WHERE "job_name" = $1 AND "status" = 'RUNNING') AS "job_running",
(SELECT max("finished_at") FROM "{_TABLE}"
WHERE "job_name" = $1 AND "status" IN ('SUCCEEDED', 'FAILED')) AS "last_finished_at",
(SELECT max("scheduled_for") FROM "{_TABLE}"
WHERE "job_name" = $1 AND "status" NOT IN ('CANCELLED', 'PENDING')) AS "last_scheduled_for"
"""
_PENDING_SQL: Final = f"""
SELECT "id", "schedule", "scheduled_for" FROM "{_TABLE}"
WHERE "job_name" = $1 AND "status" = 'PENDING'
ORDER BY "scheduled_for", "created_at"
"""
_CANCEL_SQL: Final = f"""
UPDATE "{_TABLE}" SET "status" = 'CANCELLED', "finished_at" = now()
WHERE "id" = $1 AND "status" = 'PENDING'
"""
_INSERT_SQL: Final = f"""
INSERT INTO "{_TABLE}" ("id", "job_name", "lock_group", "schedule", "scheduled_for")
VALUES ($1, $2, $3, $4, $5::timestamptz)
"""
_CLAIM_SQL: Final = f"""
UPDATE "{_TABLE}"
SET "status" = 'RUNNING', "claimed_by" = $2, "claimed_at" = now(),
"lease_expires_at" = now() + $3::double precision * interval '1 second'
WHERE "id" = $1 AND "status" = 'PENDING'
RETURNING "id", "job_name"
"""
_FINISH_SQL: Final = f"""
UPDATE "{_TABLE}" SET "status" = $3::"ScheduledJobRunStatus", "finished_at" = now()
WHERE "id" = $1 AND "claimed_by" = $2 AND "status" = 'RUNNING'
"""
_PRUNE_SQL: Final = f"""
DELETE FROM "{_TABLE}" WHERE "id" IN (
SELECT "id" FROM "{_TABLE}" WHERE "job_name" = $1 AND "finished_at" IS NOT NULL
ORDER BY "finished_at" DESC OFFSET $2
)
"""
class ScheduledJobRunner:
"""
Polls one ScheduledJob's rows in LiteLLM_ScheduledJobRun. Every pod runs the same tick.
Each tick runs in one transaction under a transaction-scoped advisory lock on the job's lock_group, so only one
pod plans and claims for a group at a time and a pod that dies mid-tick releases it. The claimed run executes
outside any transaction under a lease. A pod that dies mid-run leaves its row RUNNING until the lease expires,
when the next tick on any pod marks it FAILED and plans the next occurrence.
"""
def __init__(
self,
job: ScheduledJob,
*,
worker_id: str | None = None,
new_id: Callable[[], str] = lambda: str(uuid.uuid4()),
) -> None:
self.job = job
self.worker_id = worker_id or default_worker_id()
self._new_id = new_id
self._running = False
async def tick(self, prisma_client: "PrismaClient | None") -> None:
"""
Plan, claim and run whatever is due. Awaits the run, so graceful shutdown sees it as an in-flight job.
Register it with max_instances=2: the extra instance lets the polls that land during a long run return at
once instead of APScheduler warning that it skipped them.
"""
if prisma_client is None or self._running:
return
self._running = True
try:
claimed: Final = await self._safe_plan_and_claim(prisma_client)
if claimed is not None:
await self._execute(prisma_client, claimed)
finally:
self._running = False
async def _safe_plan_and_claim(self, prisma_client: "PrismaClient") -> str | None:
try:
return await self._plan_and_claim(prisma_client)
except Exception as e: # noqa: BLE001 - a failed tick is retried by the next one
verbose_proxy_logger.warning("Scheduled job %s tick failed: %s: %s", self.job.name, type(e).__name__, e)
return None
async def _plan_and_claim(self, prisma_client: "PrismaClient") -> str | None:
schedule: Final = self.job.schedule()
if isinstance(schedule, InvalidSchedule):
_log_invalid_schedule(self.job.name, schedule.reason)
timeout_ms: Final = int(SCHEDULED_JOB_TICK_TIMEOUT_SECONDS * 1000)
async with db_span("scheduled_job_tick", _TABLE), bounded_tx(prisma_client, timeout_ms) as tx:
await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
await tx.execute_raw(trusted_sql(f"SET LOCAL lock_timeout = {timeout_ms}"))
locked: Final = _LOCKED.validate_python(
await tx.query_raw(trusted_sql(_GROUP_LOCK_SQL), f"litellm_scheduled_job:{self.job.lock_group}")
)
if not locked or not locked[0].locked:
return None
for reaped in _CLAIMED.validate_python(await tx.query_raw(trusted_sql(_REAP_SQL), self.job.lock_group)):
verbose_proxy_logger.warning(
"Scheduled job %s run %s lost its lease; marked FAILED", reaped.job_name, reaped.id
)
state: Final = _STATE.validate_python(
await tx.query_raw(trusted_sql(_STATE_SQL), self.job.name, self.job.lock_group)
)[0]
pending: Final = tuple(
_PENDING.validate_python(await tx.query_raw(trusted_sql(_PENDING_SQL), self.job.name))
)
plan: Final = plan_tick(schedule, state, pending, self._new_id())
for run_id in plan.cancel:
await tx.execute_raw(trusted_sql(_CANCEL_SQL), run_id)
if plan.insert is not None:
await tx.execute_raw(
trusted_sql(_INSERT_SQL),
plan.insert.id,
self.job.name,
self.job.lock_group,
plan.insert.schedule,
plan.insert.scheduled_for.isoformat(),
)
if plan.claim is None:
return None
claimed: Final = _CLAIMED.validate_python(
await tx.query_raw(
trusted_sql(_CLAIM_SQL), plan.claim, self.worker_id, self.job.lease().total_seconds()
)
)
return claimed[0].id if claimed else None
async def _execute(self, prisma_client: "PrismaClient", run_id: str) -> None:
verbose_proxy_logger.info("Scheduled job %s run %s claimed by %s", self.job.name, run_id, self.worker_id)
try:
result: Final = await self._run_job(prisma_client)
except asyncio.CancelledError:
await asyncio.shield(self._finish(prisma_client, run_id, "FAILED"))
raise
await self._finish(prisma_client, run_id, result)
next_run_id: Final = await self._safe_plan_and_claim(prisma_client)
if next_run_id is not None:
await self._execute(prisma_client, next_run_id)
async def _run_job(self, prisma_client: "PrismaClient") -> JobResult:
try:
return await self.job.run(prisma_client)
except Exception: # noqa: BLE001 - the failure is recorded on the run row instead of escaping the task
verbose_proxy_logger.exception("Scheduled job %s raised", self.job.name)
return "FAILED"
async def _finish(self, prisma_client: "PrismaClient", run_id: str, result: JobResult) -> None:
timeout_ms: Final = int(SCHEDULED_JOB_TICK_TIMEOUT_SECONDS * 1000)
try:
async with db_span("scheduled_job_finish", _TABLE), bounded_tx(prisma_client, timeout_ms) as tx:
await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
updated: Final = await tx.execute_raw(trusted_sql(_FINISH_SQL), run_id, self.worker_id, result)
await tx.execute_raw(trusted_sql(_PRUNE_SQL), self.job.name, SCHEDULED_JOB_RUN_HISTORY_LIMIT)
except Exception as e: # noqa: BLE001 - the lease expiry marks the row when this write cannot
verbose_proxy_logger.warning(
"Scheduled job %s run %s could not record %s: %s: %s",
self.job.name,
run_id,
result,
type(e).__name__,
e,
)
return
if updated == 0:
verbose_proxy_logger.warning(
"Scheduled job %s run %s finished after its lease expired; its %s result was not recorded",
self.job.name,
run_id,
result,
)
return
verbose_proxy_logger.info("Scheduled job %s run %s %s", self.job.name, run_id, result)
class _Locked(BaseModel):
model_config = ConfigDict(frozen=True)
locked: bool
_LOCKED: Final = TypeAdapter(list[_Locked])
_GROUP_LOCK_SQL: Final = 'SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"'
@lru_cache(maxsize=64)
def _log_invalid_schedule(job_name: str, reason: str) -> None:
verbose_proxy_logger.error(
"Scheduled job %s has an invalid schedule and will not be replanned: %s", job_name, reason
)
def register_poll(
scheduler: JobScheduler,
prisma_client: "PrismaClient | None",
job_id: str,
job: ScheduledJob,
*,
wanted: bool,
poll_seconds: float,
misfire_grace_time: int,
) -> None:
"""
Make every pod poll job, or stop polling it. A poll that already exists is kept: its runner reads the schedule
on every tick, so a schedule edit needs no re-registration.
"""
registered: Final = scheduler.get_job(job_id) is not None
if registered and not wanted:
scheduler.remove_job(job_id)
verbose_proxy_logger.info("Stopped polling scheduled job %s", job.name)
return
if registered or not wanted:
return
scheduler.add_job(
ScheduledJobRunner(job).tick,
"interval",
args=(prisma_client,),
id=job_id,
replace_existing=True,
max_instances=2,
misfire_grace_time=misfire_grace_time,
seconds=poll_seconds,
)
verbose_proxy_logger.info("Polling scheduled job %s every %ss", job.name, poll_seconds)

View file

@ -277,8 +277,6 @@ try:
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]'`")
@ -370,8 +368,8 @@ from litellm.constants import (
REALTIME_SESSION_FAILURE_LOGGED_KEY,
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG,
SCHEDULED_JOB_POLL_INTERVAL_SECONDS,
SPEND_CAPTURE_RATE_CHECK_JOB_ID,
SPEND_LOG_SIZE_CHECK_INTERVAL,
USER_SPEND_ALERTS_JOB_ID,
WEEKLY_SPEND_REPORT_JOB_ID,
)
@ -617,7 +615,6 @@ from litellm.proxy.common_utils.scheduled_job_stagger import (
apply_scheduled_job_stagger,
attach_job_timing_logger,
parse_stagger_settings,
stagger_trigger,
)
from litellm.proxy.common_utils.slow_gc_logger import log_slow_gc_passes
from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES
@ -671,7 +668,10 @@ from litellm.proxy.container_endpoints.endpoints import router as container_rout
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
from litellm.proxy.db.create_views import SupportsRawQueries
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SPEND_LOG_RETENTION_JOB,
SPEND_LOG_SIZE_CHECK_JOB,
)
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
@ -697,6 +697,7 @@ from litellm.proxy.db.request_error_tracking import (
flush_request_errors,
request_error_accumulator,
)
from litellm.proxy.db.scheduled_job_runs import JobScheduler, register_poll
from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed
from litellm.proxy.discovery_endpoints import (
agent_skills_discovery_router,
@ -5911,6 +5912,28 @@ def _wants_cleanup_jobs(settings: Mapping[str, object]) -> bool:
return _wants_retention_job(settings) or _wants_size_check_job(settings)
_SPEND_LOG_CLEANUP_POLLS: Final = (
(_SPEND_LOG_CLEANUP_JOB_ID, _wants_retention_job, SPEND_LOG_RETENTION_JOB),
(_SPEND_LOG_SIZE_CHECK_JOB_ID, _wants_size_check_job, SPEND_LOG_SIZE_CHECK_JOB),
)
def _register_spend_log_cleanup_jobs(
job_scheduler: JobScheduler, db: PrismaClient | None, settings: Mapping[str, object]
) -> None:
"""Poll LiteLLM_ScheduledJobRun for each cleanup job the settings ask for. Every pod polls, one claims each run."""
for job_id, wanted, job in _SPEND_LOG_CLEANUP_POLLS:
register_poll(
job_scheduler,
db,
job_id,
job,
wanted=wanted(settings),
poll_seconds=SCHEDULED_JOB_POLL_INTERVAL_SECONDS,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
@lru_cache(maxsize=4096)
def _log_ignored_cost_map_copy(model_id: str, fields: tuple[str, ...]) -> None:
verbose_proxy_logger.warning(
@ -8488,97 +8511,10 @@ class ProxyConfig:
)
async def _reschedule_spend_log_cleanup_job(self):
"""
Reschedule the spend log cleanup job based on current general_settings.
This is called when maximum_spend_logs_retention_period is updated dynamically.
If the retention period is None, the job will be removed.
"""
global scheduler, general_settings, prisma_client
"""Add or remove the cleanup polls when the settings that ask for them change."""
if scheduler is None:
return
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SpendLogCleanup,
)
if not _wants_retention_job(general_settings):
if scheduler.get_job(_SPEND_LOG_CLEANUP_JOB_ID) is not None:
scheduler.remove_job(_SPEND_LOG_CLEANUP_JOB_ID)
verbose_proxy_logger.info("Removed existing spend log cleanup job")
else:
trigger: Final = self._spend_log_cleanup_trigger()
if trigger is not None:
scheduler.add_job(
SpendLogCleanup().cleanup_old_spend_logs,
trigger,
args=[prisma_client],
id=_SPEND_LOG_CLEANUP_JOB_ID,
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log cleanup rescheduled with trigger: %s", trigger)
if not _wants_size_check_job(general_settings): # pyright: ignore[reportUnknownArgumentType] # general_settings is an untyped module dict
if scheduler.get_job(_SPEND_LOG_SIZE_CHECK_JOB_ID) is not None: # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped
scheduler.remove_job(_SPEND_LOG_SIZE_CHECK_JOB_ID) # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped
verbose_proxy_logger.info("Removed existing spend log size check job")
return
size_trigger: Final = self._spend_log_size_check_trigger()
if size_trigger is None:
return
scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped
SpendLogCleanup().cleanup_spend_logs_by_size,
size_trigger,
args=[prisma_client],
id=_SPEND_LOG_SIZE_CHECK_JOB_ID,
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log size check rescheduled with trigger: %s", size_trigger)
def _spend_log_size_check_trigger(self) -> object | None:
return self._staggered_interval_trigger(
_SPEND_LOG_SIZE_CHECK_JOB_ID,
"maximum_spend_logs_size_check_interval",
general_settings.get("maximum_spend_logs_size_check_interval", SPEND_LOG_SIZE_CHECK_INTERVAL), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # general_settings is an untyped module dict
)
@staticmethod
def _staggered_interval_trigger(job_id: str, setting_name: str, interval: object) -> object | None:
if not isinstance(interval, str):
verbose_proxy_logger.error("Invalid %s value: %r", setting_name, 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 settings
try:
interval_seconds: Final = duration_in_seconds(interval)
return stagger_trigger(
job_id=job_id,
trigger=IntervalTrigger(seconds=interval_seconds),
period_seconds=interval_seconds,
settings=parse_stagger_settings(general_settings), # pyright: ignore[reportUnknownArgumentType] # general_settings is an untyped module dict
)
except (ValueError, OverflowError):
verbose_proxy_logger.error("Invalid %s value: %r", setting_name, interval)
return None
def _spend_log_cleanup_trigger(self) -> object | None:
cleanup_cron: Final[object] = general_settings.get("maximum_spend_logs_cleanup_cron")
if cleanup_cron:
from apscheduler.triggers.cron import CronTrigger
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
return self._staggered_interval_trigger(
_SPEND_LOG_CLEANUP_JOB_ID,
"maximum_spend_logs_retention_interval",
general_settings.get("maximum_spend_logs_retention_interval", "1d"),
)
_register_spend_log_cleanup_jobs(scheduler, prisma_client, general_settings) # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # apscheduler is untyped
async def _update_general_settings(self, db_general_settings: Mapping[str, SettingsJsonValue] | str | None) -> None:
global general_settings
@ -12026,70 +11962,7 @@ class ProxyStartupEvent:
### SPEND LOG CLEANUP ###
cleanup_settings: Final = _current_general_settings()
if _wants_size_check_job(cleanup_settings):
size_check_interval: Final = cleanup_settings.get(
"maximum_spend_logs_size_check_interval", SPEND_LOG_SIZE_CHECK_INTERVAL
)
try:
if not isinstance(size_check_interval, str):
raise TypeError(size_check_interval)
size_check_seconds: Final = duration_in_seconds(size_check_interval)
scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped
SpendLogCleanup().cleanup_spend_logs_by_size,
"interval",
seconds=size_check_seconds + random.randint(0, 60),
args=[prisma_client],
id=_SPEND_LOG_SIZE_CHECK_JOB_ID,
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log size check scheduled every %s", size_check_interval)
except (TypeError, ValueError, OverflowError):
verbose_proxy_logger.error(
"Invalid maximum_spend_logs_size_check_interval value: %r", size_check_interval
)
if _wants_retention_job(cleanup_settings):
spend_log_cleanup: Final = SpendLogCleanup()
cleanup_cron: Final = cleanup_settings.get("maximum_spend_logs_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:
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_ID,
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info("Spend log cleanup scheduled 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)
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,
"interval",
seconds=interval_seconds + random.randint(0, 60),
args=[prisma_client],
id=_SPEND_LOG_CLEANUP_JOB_ID,
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
except (ValueError, OverflowError):
verbose_proxy_logger.error(
"Invalid maximum_spend_logs_retention_interval value: %r", retention_interval
)
_register_spend_log_cleanup_jobs(scheduler, prisma_client, cleanup_settings) # pyright: ignore[reportUnknownArgumentType, reportArgumentType] # apscheduler is untyped
if _wants_cleanup_jobs(cleanup_settings):
proxy_config.record_cleanup_schedule_attempt(cleanup_settings)
### CHECK BATCH COST ###

View file

@ -1155,6 +1155,31 @@ enum JobStatus {
INACTIVE
}
model LiteLLM_ScheduledJobRun {
id String @id @default(uuid())
job_name String
lock_group String
schedule String
scheduled_for DateTime @client.Timestamptz(3)
status ScheduledJobRunStatus @default(PENDING)
claimed_by String?
claimed_at DateTime? @client.Timestamptz(3)
lease_expires_at DateTime? @client.Timestamptz(3)
finished_at DateTime? @client.Timestamptz(3)
created_at DateTime @default(now()) @client.Timestamptz(3)
@@index([job_name, status])
@@index([lock_group, status])
}
enum ScheduledJobRunStatus {
PENDING
RUNNING
SUCCEEDED
FAILED
CANCELLED
}
model LiteLLM_ManagedFileTable {
id String @id @default(uuid())
unified_file_id String @unique // The base64 encoded unified file ID

View file

@ -1155,6 +1155,31 @@ enum JobStatus {
INACTIVE
}
model LiteLLM_ScheduledJobRun {
id String @id @default(uuid())
job_name String
lock_group String
schedule String
scheduled_for DateTime @client.Timestamptz(3)
status ScheduledJobRunStatus @default(PENDING)
claimed_by String?
claimed_at DateTime? @client.Timestamptz(3)
lease_expires_at DateTime? @client.Timestamptz(3)
finished_at DateTime? @client.Timestamptz(3)
created_at DateTime @default(now()) @client.Timestamptz(3)
@@index([job_name, status])
@@index([lock_group, status])
}
enum ScheduledJobRunStatus {
PENDING
RUNNING
SUCCEEDED
FAILED
CANCELLED
}
model LiteLLM_ManagedFileTable {
id String @id @default(uuid())
unified_file_id String @unique // The base64 encoded unified file ID

View file

@ -202,7 +202,6 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does
"maximum_spend_logs_cleanup_batch_size": 100,
}
)
cleanup.pod_lock_manager = None
async def size_check() -> None:
database: Final = Prisma(datasource={"url": scoped_url})

View file

@ -0,0 +1,294 @@
import asyncio
import time
from collections.abc import AsyncIterator, Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Final, cast
import psycopg
import pytest
from integration._support.database import scratch_database
from prisma import Prisma
from psycopg.rows import dict_row
from litellm.proxy.db.scheduled_job_runs import (
JobResult,
Schedule,
ScheduledJob,
ScheduledJobRunner,
cron_schedule,
interval_schedule,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_MIGRATION: Final = next(
Path(__file__)
.resolve()
.parents[3]
.joinpath("litellm-proxy-extras/litellm_proxy_extras/migrations")
.glob("*_scheduled_job_runs/migration.sql")
)
_PODS: Final = 4
@dataclass(frozen=True)
class _Connection:
db: Prisma
@dataclass(frozen=True)
class _Database:
url: str
pods: tuple["PrismaClient", ...]
def rows(self, where: str = "TRUE") -> list[dict[str, object]]:
with psycopg.connect(self.url, row_factory=dict_row) as connection:
return connection.execute(
f'SELECT * FROM "LiteLLM_ScheduledJobRun" WHERE {where} ORDER BY "created_at", "scheduled_for"'
).fetchall()
def execute(self, statement: str, parameters: tuple[object, ...] = ()) -> None:
with psycopg.connect(self.url, autocommit=True) as connection:
connection.execute(statement, parameters)
class _Job:
"""A job whose runs are counted and can be held open until the test releases them."""
def __init__(self, name: str = "job", lock_group: str = "group", schedule: str = "1h") -> None:
self.name = name
self.lock_group = lock_group
self.schedule_value = schedule
self.started = 0
self.release = asyncio.Event()
self.release.set()
self.outcome: Callable[[], Awaitable[JobResult]] = self._succeed
async def _succeed(self) -> JobResult:
return "SUCCEEDED"
async def _run(self, _prisma_client: object) -> JobResult:
self.started += 1
await self.release.wait()
return await self.outcome()
def _schedule(self) -> Schedule:
parsed: Final = (
cron_schedule("cron", self.schedule_value)
if " " in self.schedule_value
else interval_schedule("interval", self.schedule_value)
)
assert isinstance(parsed, Schedule)
return parsed
def definition(self) -> ScheduledJob:
return ScheduledJob(
name=self.name,
lock_group=self.lock_group,
schedule=self._schedule,
lease=lambda: timedelta(minutes=5),
run=self._run,
)
@asynccontextmanager
async def _scratch() -> AsyncIterator[_Database]:
with scratch_database("scheduled_jobs") as url:
with psycopg.connect(url, autocommit=True) as setup:
setup.execute(_MIGRATION.read_text())
connections: Final = tuple(_Connection(Prisma(datasource={"url": url})) for _ in range(_PODS))
for connection in connections:
await connection.db.connect()
try:
yield _Database(
url=url,
pods=tuple(
cast("PrismaClient", connection) # cast-ok: the runner only reaches the client through .db
for connection in connections
),
)
finally:
for connection in connections:
await connection.db.disconnect()
def _runners(job: _Job) -> tuple[ScheduledJobRunner, ...]:
return tuple(ScheduledJobRunner(job.definition(), worker_id=f"pod-{index}") for index in range(_PODS))
async def _tick_all(runners: tuple[ScheduledJobRunner, ...], database: _Database) -> None:
await asyncio.gather(*(runner.tick(pod) for runner, pod in zip(runners, database.pods)))
async def _until(condition: Callable[[], bool], seconds: float = 10) -> None:
deadline: Final = time.monotonic() + seconds
while not condition():
assert time.monotonic() < deadline, "condition never held"
await asyncio.sleep(0.02)
async def test_every_pod_ticks_but_exactly_one_runs_and_the_next_run_is_planned_one_interval_after_it_finished() -> (
None
):
async with _scratch() as database:
job: Final = _Job()
job.release.clear()
runners: Final = _runners(job)
ticks: Final = asyncio.gather(*(runner.tick(pod) for runner, pod in zip(runners, database.pods)))
await _until(lambda: job.started == 1)
await asyncio.sleep(0.5)
running: Final = database.rows("\"status\" = 'RUNNING'")
assert job.started == 1
assert len(running) == 1
job.release.set()
await ticks
finished: Final = database.rows("\"status\" = 'SUCCEEDED'")
pending: Final = database.rows("\"status\" = 'PENDING'")
assert len(finished) == 1
assert finished[0]["claimed_by"] == running[0]["claimed_by"]
assert len(pending) == 1
assert pending[0]["scheduled_for"] == finished[0]["finished_at"] + timedelta(hours=1)
await _tick_all(runners, database)
assert job.started == 1, "a run was started before the next occurrence was due"
assert len(database.rows("\"status\" = 'PENDING'")) == 1
async def test_jobs_in_one_lock_group_never_run_at_the_same_time_on_different_pods() -> None:
async with _scratch() as database:
first: Final = _Job(name="first")
second: Final = _Job(name="second")
first.release.clear()
first_runner: Final = ScheduledJobRunner(first.definition(), worker_id="pod-a")
second_runner: Final = ScheduledJobRunner(second.definition(), worker_id="pod-b")
first_tick: Final = asyncio.create_task(first_runner.tick(database.pods[0]))
await _until(lambda: first.started == 1)
await second_runner.tick(database.pods[1])
assert second.started == 0
assert [row["status"] for row in database.rows("\"job_name\" = 'second'")] == ["PENDING"]
first.release.set()
await first_tick
await second_runner.tick(database.pods[1])
assert second.started == 1
async def test_a_run_that_raises_is_recorded_as_failed_and_the_next_run_is_still_planned() -> None:
async with _scratch() as database:
job: Final = _Job()
async def explode() -> JobResult:
raise RuntimeError("disk on fire")
job.outcome = explode
await _tick_all(_runners(job), database)
failed: Final = database.rows("\"status\" = 'FAILED'")
assert len(failed) == 1
assert len(database.rows("\"status\" = 'PENDING'")) == 1
async def test_a_run_held_by_a_dead_pod_is_failed_once_its_lease_expires_and_the_schedule_carries_on() -> None:
async with _scratch() as database:
database.execute(
'INSERT INTO "LiteLLM_ScheduledJobRun" ("id", "job_name", "lock_group", "schedule", "scheduled_for", '
'"status", "claimed_by", "claimed_at", "lease_expires_at") '
"VALUES ('dead', 'job', 'group', 'interval:3600s', now() - interval '2 hours', 'RUNNING', 'dead-pod', "
"now() - interval '2 hours', now() - interval '1 second')"
)
job: Final = _Job()
await _tick_all(_runners(job), database)
(reaped,) = database.rows("\"id\" = 'dead'")
assert reaped["status"] == "FAILED"
(pending,) = database.rows("\"status\" = 'PENDING'")
assert pending["scheduled_for"] == reaped["finished_at"] + timedelta(hours=1)
assert job.started == 0
async def test_a_run_held_by_a_live_pod_is_left_alone_and_nothing_else_starts() -> None:
async with _scratch() as database:
job: Final = _Job()
job.release.clear()
runners: Final = _runners(job)
holder: Final = asyncio.create_task(runners[0].tick(database.pods[0]))
await _until(lambda: job.started == 1)
await asyncio.gather(*(runner.tick(pod) for runner, pod in zip(runners[1:], database.pods[1:])))
assert [row["status"] for row in database.rows()] == ["RUNNING"]
job.release.set()
await holder
assert job.started == 1
async def test_an_interval_already_overdue_when_the_pods_restart_runs_on_the_first_tick() -> None:
async with _scratch() as database:
database.execute(
'INSERT INTO "LiteLLM_ScheduledJobRun" ("id", "job_name", "lock_group", "schedule", "scheduled_for", '
'"status", "finished_at") '
"VALUES ('yesterday', 'job', 'group', 'interval:86400s', now() - interval '2 days', 'SUCCEEDED', "
"now() - interval '2 days')"
)
job: Final = _Job(schedule="1d")
await _tick_all(_runners(job), database)
assert job.started == 1
assert [row["status"] for row in database.rows("\"id\" <> 'yesterday'")] == ["SUCCEEDED", "PENDING"]
async def test_a_schedule_change_replaces_the_pending_run_planned_from_the_old_schedule() -> None:
async with _scratch() as database:
job: Final = _Job(schedule="0 4 * * *")
runners: Final = _runners(job)
await _tick_all(runners, database)
(old,) = database.rows("\"status\" = 'PENDING'")
assert old["scheduled_for"].astimezone(timezone.utc).hour == 4
job.schedule_value = "0 5 * * *"
await _tick_all(runners, database)
assert database.rows(f"\"id\" = '{old['id']}'")[0]["status"] == "CANCELLED"
(replacement,) = database.rows("\"status\" = 'PENDING'")
assert replacement["schedule"] == "cron:0 5 * * *"
assert replacement["scheduled_for"].astimezone(timezone.utc).hour == 5
assert replacement["scheduled_for"] > datetime.now(timezone.utc)
assert job.started == 0
async def test_moving_a_cron_earlier_plans_the_earlier_occurrence_instead_of_skipping_a_day() -> None:
async with _scratch() as database:
now: Final = datetime.now(timezone.utc)
job: Final = _Job(schedule=f"0 {(now + timedelta(hours=3)).hour} * * *")
runners: Final = _runners(job)
await _tick_all(runners, database)
(old,) = database.rows("\"status\" = 'PENDING'")
job.schedule_value = f"0 {(now + timedelta(hours=2)).hour} * * *"
await _tick_all(runners, database)
(replacement,) = database.rows("\"status\" = 'PENDING'")
assert replacement["scheduled_for"] < old["scheduled_for"]
assert replacement["scheduled_for"] - datetime.now(timezone.utc) < timedelta(hours=2)
async def test_a_run_cancelled_by_shutdown_is_recorded_as_failed_instead_of_holding_its_lease() -> None:
async with _scratch() as database:
job: Final = _Job()
job.release.clear()
runner: Final = ScheduledJobRunner(job.definition(), worker_id="pod-shutting-down")
tick: Final = asyncio.create_task(runner.tick(database.pods[0]))
await _until(lambda: job.started == 1)
tick.cancel()
with pytest.raises(asyncio.CancelledError):
await tick
(cancelled,) = database.rows()
assert cancelled["status"] == "FAILED"

View file

@ -0,0 +1,355 @@
import asyncio
import logging
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Final, cast
import pytest
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.db.scheduled_job_runs import (
InvalidSchedule,
JobResult,
JobScheduler,
JobState,
NewRun,
PendingRun,
Schedule,
ScheduledJob,
ScheduledJobRunner,
TickPlan,
cron_schedule,
interval_schedule,
next_run_at,
plan_tick,
register_poll,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
NOW: Final = datetime(2026, 11, 4, 12, 0, tzinfo=timezone.utc)
def _schedule(value: str) -> Schedule:
parsed: Final = cron_schedule("cron", value) if " " in value else interval_schedule("interval", value)
assert not isinstance(parsed, InvalidSchedule)
return parsed
def _state(
*,
group_running: bool = False,
job_running: bool = False,
last_finished_at: datetime | None = None,
last_scheduled_for: datetime | None = None,
) -> JobState:
return JobState(
now=NOW,
group_running=group_running,
job_running=job_running,
last_finished_at=last_finished_at,
last_scheduled_for=last_scheduled_for,
)
def _pending(run_id: str, schedule: str, scheduled_for: datetime) -> PendingRun:
return PendingRun(id=run_id, schedule=schedule, scheduled_for=scheduled_for)
@pytest.mark.parametrize(
("schedule", "state", "expected"),
[
pytest.param("1h", _state(), NOW, id="interval-never-run-is-due-now"),
pytest.param(
"1h",
_state(last_finished_at=NOW - timedelta(hours=3)),
NOW - timedelta(hours=2),
id="interval-counts-from-last-finish-so-a-restart-cannot-postpone-it",
),
pytest.param(
"1h",
_state(last_finished_at=NOW - timedelta(minutes=10)),
NOW + timedelta(minutes=50),
id="interval-not-yet-elapsed",
),
pytest.param("0 4 * * *", _state(), datetime(2026, 11, 5, 4, tzinfo=timezone.utc), id="cron-next-fire"),
pytest.param(
"0 12 * * *",
_state(last_scheduled_for=NOW),
datetime(2026, 11, 5, 12, tzinfo=timezone.utc),
id="cron-never-replans-the-occurrence-it-just-had",
),
pytest.param(
"0 4 * * *",
_state(last_scheduled_for=NOW - timedelta(days=3)),
datetime(2026, 11, 5, 4, tzinfo=timezone.utc),
id="cron-outage-collapses-to-one-next-fire",
),
],
)
def test_next_run_at(schedule: str, state: JobState, expected: datetime) -> None:
assert next_run_at(_schedule(schedule), state) == expected
DUE: Final = NOW - timedelta(minutes=1)
LATER: Final = NOW + timedelta(hours=1)
@pytest.mark.parametrize(
("schedule", "state", "pending", "expected"),
[
pytest.param(
"1h",
_state(),
(),
TickPlan(cancel=(), insert=NewRun(id="new", schedule="interval:3600s", scheduled_for=NOW), claim="new"),
id="first-ever-tick-plans-and-claims",
),
pytest.param(
"1h",
_state(),
(_pending("a", "interval:3600s", DUE),),
TickPlan(cancel=(), insert=None, claim="a"),
id="due-pending-is-claimed",
),
pytest.param(
"1h",
_state(),
(_pending("a", "interval:3600s", LATER),),
TickPlan(cancel=(), insert=None, claim=None),
id="future-pending-waits",
),
pytest.param(
"1h",
_state(group_running=True),
(_pending("a", "interval:3600s", DUE),),
TickPlan(cancel=(), insert=None, claim=None),
id="due-pending-waits-while-its-lock-group-runs",
),
pytest.param(
"1h",
_state(group_running=True),
(),
TickPlan(cancel=(), insert=NewRun(id="new", schedule="interval:3600s", scheduled_for=NOW), claim=None),
id="a-sibling-run-still-lets-the-next-run-be-planned",
),
pytest.param(
"1h",
_state(group_running=True, job_running=True),
(),
TickPlan(cancel=(), insert=None, claim=None),
id="nothing-is-planned-while-the-job-itself-runs",
),
pytest.param(
"2h",
_state(last_finished_at=NOW - timedelta(minutes=30)),
(_pending("old", "interval:3600s", NOW + timedelta(minutes=30)),),
TickPlan(
cancel=("old",),
insert=NewRun(id="new", schedule="interval:7200s", scheduled_for=NOW + timedelta(minutes=90)),
claim=None,
),
id="schedule-change-replaces-the-pending-run",
),
pytest.param(
"1h",
_state(),
(_pending("a", "interval:3600s", DUE), _pending("b", "interval:3600s", LATER)),
TickPlan(cancel=("b",), insert=None, claim="a"),
id="a-duplicate-pending-run-is-cancelled",
),
],
)
def test_plan_tick(schedule: str, state: JobState, pending: tuple[PendingRun, ...], expected: TickPlan) -> None:
assert plan_tick(_schedule(schedule), state, pending, "new") == expected
def test_an_invalid_schedule_replans_nothing_but_still_runs_what_the_last_valid_one_planned() -> None:
pending: Final = (_pending("a", "interval:3600s", DUE),)
plan: Final = plan_tick(InvalidSchedule("bad cron"), _state(), pending, "new")
assert plan == TickPlan(cancel=(), insert=None, claim="a")
@pytest.mark.parametrize("value", ["not a cron", "0 0 * * * *", "0 0 31 2 *", "0 0 30 2 *", 5, None])
def test_cron_schedule_rejects_what_is_not_a_crontab(value: object) -> None:
assert isinstance(cron_schedule("maximum_spend_logs_cleanup_cron", value), InvalidSchedule)
@pytest.mark.parametrize("value", ["soon", "99999999999d", "0s", 86400, None])
def test_interval_schedule_rejects_what_is_not_a_positive_duration(value: object) -> None:
assert isinstance(interval_schedule("maximum_spend_logs_retention_interval", value), InvalidSchedule)
class _ScriptedDatabase:
"""Answers the runner's statements the way Postgres would for one due run, and records every write."""
def __init__(self, *, locked: bool = True, finish_updates: int = 1, broken: bool = False) -> None:
self.locked = locked
self.finish_updates = finish_updates
self.broken = broken
self.writes: tuple[tuple[str, tuple[object, ...]], ...] = ()
self.db = self
@asynccontextmanager
async def tx(self, timeout: timedelta) -> AsyncIterator["_ScriptedDatabase"]:
if self.broken:
raise ConnectionError("database unreachable")
yield self
async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]:
if "pg_try_advisory_xact_lock" in sql:
return [{"locked": self.locked}]
if "SELECT now()" in sql:
return [
{
"now": NOW,
"group_running": False,
"job_running": False,
"last_finished_at": None,
"last_scheduled_for": None,
}
]
if "'RUNNING', \"claimed_by\"" in sql:
self.writes = (*self.writes, ("claim", args))
return [{"id": args[0], "job_name": "job"}]
return []
async def execute_raw(self, sql: str, *args: object) -> int:
if sql.lstrip().startswith("SET LOCAL"):
return 0
if 'SET "status" = $3' in sql:
self.writes = (*self.writes, ("finish", args))
self.locked = False
return self.finish_updates
if sql.lstrip().startswith("INSERT"):
self.writes = (*self.writes, ("insert", args))
return 1
def client(self) -> "PrismaClient":
return cast("PrismaClient", self) # cast-ok: the runner only reaches the client through .db
def _finishes(database: _ScriptedDatabase) -> tuple[tuple[object, ...], ...]:
return tuple(args for kind, args in database.writes if kind == "finish")
class _Job:
def __init__(self, result: JobResult | BaseException = "SUCCEEDED") -> None:
self.result = result
self.runs = 0
self.release = asyncio.Event()
self.release.set()
async def run(self, _prisma_client: object) -> JobResult:
self.runs += 1
await self.release.wait()
if isinstance(self.result, BaseException):
raise self.result
return self.result
def runner(self) -> ScheduledJobRunner:
definition: Final = ScheduledJob(
name="job",
lock_group="group",
schedule=lambda: _schedule("1h"),
lease=lambda: timedelta(minutes=5),
run=self.run,
)
return ScheduledJobRunner(definition, worker_id="pod-a", new_id=lambda: "run-1")
async def test_a_due_run_is_claimed_by_this_pod_run_and_recorded_as_succeeded() -> None:
database: Final = _ScriptedDatabase()
job: Final = _Job("SUCCEEDED")
await job.runner().tick(database.client())
assert job.runs == 1
assert ("claim", ("run-1", "pod-a", 300.0)) in database.writes
assert _finishes(database) == (("run-1", "pod-a", "SUCCEEDED"),)
@pytest.mark.parametrize(
"result",
[pytest.param("FAILED", id="reported"), pytest.param(RuntimeError("disk on fire"), id="raised")],
)
async def test_a_failed_run_is_recorded_as_failed(result: JobResult | BaseException) -> None:
database: Final = _ScriptedDatabase()
await _Job(result).runner().tick(database.client())
assert _finishes(database) == (("run-1", "pod-a", "FAILED"),)
async def test_a_pod_that_does_not_get_the_group_lock_neither_plans_nor_runs() -> None:
database: Final = _ScriptedDatabase(locked=False)
job: Final = _Job()
await job.runner().tick(database.client())
assert job.runs == 0
assert database.writes == ()
async def test_an_unreachable_database_skips_the_tick_instead_of_raising_into_the_scheduler() -> None:
job: Final = _Job()
await job.runner().tick(_ScriptedDatabase(broken=True).client())
assert job.runs == 0
async def test_a_poll_landing_while_this_pod_still_runs_the_job_does_not_touch_the_database() -> None:
database: Final = _ScriptedDatabase()
job: Final = _Job()
job.release.clear()
runner: Final = job.runner()
first: Final = asyncio.create_task(runner.tick(database.client()))
while job.runs == 0:
await asyncio.sleep(0)
writes_during_run: Final = database.writes
await runner.tick(database.client())
assert database.writes == writes_during_run
job.release.set()
await first
assert job.runs == 1
async def test_a_run_cancelled_mid_flight_records_failed_and_still_cancels() -> None:
database: Final = _ScriptedDatabase()
job: Final = _Job()
job.release.clear()
tick: Final = asyncio.create_task(job.runner().tick(database.client()))
while job.runs == 0:
await asyncio.sleep(0)
tick.cancel()
with pytest.raises(asyncio.CancelledError):
await tick
assert _finishes(database) == (("run-1", "pod-a", "FAILED"),)
async def test_a_result_that_arrives_after_the_lease_expired_is_reported_not_recorded(
caplog: pytest.LogCaptureFixture,
) -> None:
database: Final = _ScriptedDatabase(finish_updates=0)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await _Job().runner().tick(database.client())
assert any("after its lease expired" in record.getMessage() for record in caplog.records)
async def test_register_poll_adds_a_wanted_poll_once_and_removes_it_when_no_longer_wanted() -> None:
scheduler: Final = AsyncIOScheduler()
scheduler.start(paused=True)
job_scheduler: Final = cast(JobScheduler, scheduler) # cast-ok: apscheduler ships no type information
definition: Final = _Job().runner().job
def register(wanted: bool) -> None:
register_poll(job_scheduler, None, "poll-id", definition, wanted=wanted, poll_seconds=30, misfire_grace_time=60)
try:
register(True)
first: Final = scheduler.get_job("poll-id")
assert first is not None
register(True)
assert scheduler.get_job("poll-id") is first, "re-registering replaced the live poll and its runner"
register(False)
assert scheduler.get_job("poll-id") is None
finally:
scheduler.shutdown(wait=False)

View file

@ -19,7 +19,7 @@ import sys
from collections.abc import Awaitable, Callable, Mapping
from contextlib import nullcontext
from dataclasses import dataclass
from datetime import datetime, timedelta
from datetime import datetime
from pathlib import Path
from types import MappingProxyType, SimpleNamespace
from typing import Any, Dict, Final
@ -4717,30 +4717,9 @@ async def test_ProxyConfig__reschedule_spend_log_cleanup_job_no_scheduler(monkey
assert snapshot == {"raised": False, "called": True, "scheduler_was": "none"}
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_invalid_cron(monkeypatch):
fake_scheduler = MagicMock()
fake_scheduler.remove_job = MagicMock()
fake_scheduler.add_job = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", fake_scheduler)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{
"maximum_spend_logs_retention_period": "1d",
"maximum_spend_logs_cleanup_cron": "INVALID CRON STRING",
},
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
# Invalid cron is caught and logged — does not raise outward.
await pc._reschedule_spend_log_cleanup_job()
# But add_job should not have been called for the invalid cron path.
assert fake_scheduler.add_job.call_count == 0
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_health_check_retention(monkeypatch):
fake_scheduler = MagicMock()
fake_scheduler = MagicMock(**{"get_job.return_value": None})
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", fake_scheduler)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
@ -4889,7 +4868,7 @@ async def test_ProxyConfig__reschedule_spend_log_cleanup_job_daily_tag_spend_ret
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"
assert job.func.__name__ == "tick"
finally:
real_scheduler.shutdown(wait=False)
@ -4924,33 +4903,6 @@ async def test_ProxyConfig__update_general_settings_schedules_cleanup_when_db_ro
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
@ -4969,29 +4921,6 @@ async def test_ProxyConfig__update_general_settings_retries_a_schedule_that_rais
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,
@ -5027,34 +4956,6 @@ async def test_ProxyConfig__update_general_settings_leaves_first_registration_to
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 apscheduler.triggers.interval import IntervalTrigger
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()
pc.settings.load_yaml({"scheduled_job_stagger": {"offsets": {"spend_log_cleanup_job": 120}}})
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"]
trigger: Final = jobs[0].trigger
assert isinstance(trigger, _OffsetTrigger), repr(trigger)
assert trigger.offset == timedelta(seconds=120)
assert isinstance(trigger.base, IntervalTrigger)
assert trigger.base.interval == timedelta(days=1)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"bad_schedule",
@ -5094,89 +4995,6 @@ async def test_ProxyConfig__update_general_settings_keeps_the_live_cleanup_job_w
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
# ---------------------------------------------------------------------------

View file

@ -6,8 +6,6 @@ import json
import logging
import os
import re
import socket
import subprocess
import sys
import time
import types
@ -19,9 +17,6 @@ from typing import Final
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
import click
import fastapi.routing
import httpx
import pytest
@ -34,7 +29,6 @@ from starlette.datastructures import URL
import litellm
import litellm.proxy.proxy_server as proxy_server_module
from litellm._internal_context import current_service_target
from litellm._logging import turn_on_json, verbose_proxy_logger
from litellm.caching.caching import RedisCache
from litellm.caching.dual_cache import DualCache
@ -56,16 +50,17 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.login_throttle import LoginThrottle
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.banner import LITELLM_BANNER
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendUpdateQueue
from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash
from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler
from litellm.litellm_core_utils.model_cost_table import publish
from litellm.proxy.common_utils.banner import LITELLM_BANNER
from tests._master_key import MASTER_KEY
from litellm.redis.stores.connection import connect
from litellm.redis.stores.rate_limits import RateLimitScope, RateLimitStore
from litellm.redis.stores.spend_counters import SpendCounterStore
from litellm.rust_bridge._native import NativeRedisStoreError
from tests._master_key import MASTER_KEY
from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234"
@ -1057,43 +1052,6 @@ async def test_initialize_scheduled_jobs_registers_cleanup_when_retention_lives_
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):
"""
@ -1756,7 +1714,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path):
with open(config_path, "w") as f:
yaml.dump(test_config, f)
print(f"SET ENV VARIABLE - CONFIG_FILE_PATH, str(config_path): {str(config_path)}")
print(f"SET ENV VARIABLE - CONFIG_FILE_PATH, str(config_path): {config_path!s}")
# Second setting of CONFIG_FILE_PATH to a different value
monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path))
print(f"config_path: {config_path}")
@ -1775,7 +1733,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path):
yaml.dump(empty_config, f)
monkeypatch.setenv("LITELLM_MASTER_KEY", test_env_master_key)
print("test_env_master_key: {}".format(test_env_master_key))
print(f"test_env_master_key: {test_env_master_key}")
async with proxy_startup_event(app):
from litellm.proxy.proxy_server import master_key
@ -2098,7 +2056,7 @@ def test_team_info_masking():
all_teams_config=[team1_info],
)
print("Got exception: {}".format(exc_info.value))
print(f"Got exception: {exc_info.value}")
assert "secret-test-key" not in str(exc_info.value)
assert "public-test-key" not in str(exc_info.value)
@ -2110,7 +2068,6 @@ async def test_get_all_team_models():
"""
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy._types import LiteLLM_TeamTable
from litellm.proxy.proxy_server import get_all_team_models
# Mock team data
@ -4092,7 +4049,7 @@ async def test_load_config_warns_per_worker_login_counters_without_general_setti
Redis must hear that its counters are per worker even when the config has no general_settings."""
import logging
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy import proxy_server
from litellm.proxy.auth.login_throttle import warn_login_counters_are_per_worker
from litellm.proxy.proxy_server import ProxyConfig
@ -4252,7 +4209,6 @@ async def test_load_environment_variables_blocks_dangerous_keys():
Test that _load_environment_variables rejects dangerous env var keys
like PATH, LD_PRELOAD, PYTHONPATH, etc.
"""
import logging
from litellm.proxy.proxy_server import ProxyConfig
@ -4334,7 +4290,7 @@ async def test_write_config_to_file(monkeypatch):
"""
Do not write config to file if store_model_in_db is True
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, patch
from litellm.proxy.proxy_server import ProxyConfig
@ -4382,7 +4338,7 @@ async def test_write_config_to_file_when_store_model_in_db_false(monkeypatch):
"""
Test that config IS written to file when store_model_in_db is False
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import patch
from litellm.proxy.proxy_server import ProxyConfig
@ -4479,7 +4435,7 @@ async def test_async_data_generator_midstream_error():
try:
async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data):
yielded_data.append(data)
except Exception as e:
except Exception:
# If there's an exception, that's also part of what we want to test
pass
@ -5581,7 +5537,7 @@ async def test_add_router_settings_from_db_config_merge_logic():
This tests how router settings from config file and database are combined,
including scenarios where nested dictionaries should be properly merged.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.proxy_server import ProxyConfig
@ -6753,7 +6709,7 @@ async def test_init_sso_settings_in_db_error_handling():
"""
Test that _init_sso_settings_in_db handles database errors gracefully.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, MagicMock
from litellm.proxy.proxy_server import ProxyConfig
@ -6883,7 +6839,6 @@ async def test_init_hashicorp_vault_config_override_retries_on_transport_error()
invocations.append(None)
if len(invocations) == 1:
raise prisma.errors.ClientNotConnectedError()
return None # No config in DB → function returns early after retry.
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock(side_effect=_flaky_find_unique)
@ -7031,9 +6986,7 @@ def test_root_redirect_when_docs_url_not_root_and_redirect_url_set(monkeypatch):
routes_to_remove = []
for route in app.routes:
if hasattr(route, "path") and route.path == "/":
if hasattr(route, "methods") and "GET" in route.methods:
routes_to_remove.append(route)
elif not hasattr(route, "methods"): # Catch-all routes
if hasattr(route, "methods") and "GET" in route.methods or not hasattr(route, "methods"):
routes_to_remove.append(route)
for route in routes_to_remove:
@ -9869,7 +9822,7 @@ async def test_reseed_locks_dict_is_bounded():
long-lived deployments with high counter-key churn. Inserting more
than the cap evicts the oldest entries.
"""
import litellm.constants as constants
from litellm import constants
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
orig_locks = SpendCounterReseed._locks.copy()

View file

@ -22,16 +22,20 @@ from litellm.constants import (
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
SPEND_LOG_CLEANUP_BOUND_SETTINGS,
SPEND_LOG_RETENTION_JOB,
SPEND_LOG_SIZE_CHECK_JOB,
SizeCleanupState,
SpendLogCleanup,
SpendLogStorage,
StorageProbe,
TableCleanupResult,
advance_size_state,
retention_schedule,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
SpendLogCleanupMetrics,
)
from litellm.proxy.db.scheduled_job_runs import InvalidSchedule, ScheduledJob
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@ -276,17 +280,8 @@ async def test_cleanup_old_spend_logs_batch_deletion():
# Wire up mocks
mock_prisma_client.db = mock_db
# Mock Redis cache and pod_lock_manager
mock_redis_cache = MagicMock()
mock_pod_lock_manager = MagicMock()
mock_pod_lock_manager.redis_cache = mock_redis_cache
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
mock_pod_lock_manager.release_lock = AsyncMock()
# Run cleanup with mocked pod_lock_manager
test_settings = {"maximum_spend_logs_retention_period": "7d"}
cleaner = SpendLogCleanup(general_settings=test_settings)
cleaner.pod_lock_manager = mock_pod_lock_manager
assert cleaner._should_delete_spend_logs() is True
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -325,17 +320,8 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff():
mock_db.execute_raw = AsyncMock(return_value=0)
mock_prisma_client.db = mock_db
# Mock Redis cache and pod_lock_manager
mock_redis_cache = MagicMock()
mock_pod_lock_manager = MagicMock()
mock_pod_lock_manager.redis_cache = mock_redis_cache
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
mock_pod_lock_manager.release_lock = AsyncMock()
# Run cleanup with mocked pod_lock_manager
test_settings = {"maximum_spend_logs_retention_period": "24h"}
cleaner = SpendLogCleanup(general_settings=test_settings)
cleaner.pod_lock_manager = mock_pod_lock_manager
assert cleaner._should_delete_spend_logs() is True
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -371,8 +357,6 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned():
},
partition_manager=partition_manager,
)
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -408,8 +392,6 @@ async def test_cleanup_uses_delete_when_partitioning_not_enabled():
general_settings={"maximum_spend_logs_retention_period": "7d"},
partition_manager=partition_manager,
)
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -442,8 +424,6 @@ async def test_cleanup_uses_delete_when_not_partitioned():
},
partition_manager=partition_manager,
)
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -468,32 +448,6 @@ async def test_cleanup_old_spend_logs_no_retention_period():
mock_prisma_client.db.execute_raw.assert_not_called()
@pytest.mark.asyncio
async def test_lock_not_released_when_not_acquired():
"""
Lock release should be skipped when _should_delete_spend_logs returns False
before the lock is ever acquired.
"""
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock()
mock_redis_cache = MagicMock()
mock_pod_lock_manager = MagicMock()
mock_pod_lock_manager.redis_cache = mock_redis_cache
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
mock_pod_lock_manager.release_lock = AsyncMock()
# No retention setting → _should_delete_spend_logs() returns False before lock is acquired
cleaner = SpendLogCleanup(general_settings={})
cleaner.pod_lock_manager = mock_pod_lock_manager
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
mock_pod_lock_manager.acquire_lock.assert_not_called()
mock_pod_lock_manager.release_lock.assert_not_called()
@pytest.mark.asyncio
async def test_integer_retention_treated_as_days():
"""
@ -719,7 +673,6 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch):
_wire_tx(mock_prisma_client.db)
# Force the outer try/except to fire by making _should_delete_spend_logs raise.
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = None
def boom():
raise RuntimeError("simulated prisma timeout")
@ -738,9 +691,8 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch):
@pytest.mark.asyncio
async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch):
"""Even when batch deletion aborts due to consecutive failures, the pod lock
must still be released so the next scheduled run isn't permanently blocked."""
async def test_a_run_whose_batches_keep_failing_reports_aborted(monkeypatch: pytest.MonkeyPatch) -> None:
"""An aborted run is what the scheduled job records as FAILED, so the abort must reach the report."""
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2)
@ -748,23 +700,12 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_db = MagicMock()
_wire_tx(mock_db)
mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
mock_prisma_client.db = mock_db
mock_pod_lock_manager = MagicMock()
mock_pod_lock_manager.redis_cache = MagicMock()
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
mock_pod_lock_manager.release_lock = AsyncMock()
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = mock_pod_lock_manager
outcome: Final = await cleaner.cleanup_old_spend_logs(mock_prisma_client)
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
# Cleanup didn't crash; the abort-after-failures path returned cleanly.
mock_pod_lock_manager.release_lock.assert_awaited_once()
assert outcome == "aborted"
def test_cleanup_batch_size_env_var(monkeypatch):
@ -801,7 +742,6 @@ def _mock_prisma_for_retention(side_effect: list) -> "MagicMock":
async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup():
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 any('"LiteLLM_SpendLogs"' in sql for sql in tables)
@ -815,7 +755,6 @@ async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup():
async def test_session_retention_alone_cleans_both_session_rollups_and_the_daily_rollup():
client = _mock_prisma_for_retention([0, 0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_autorouter_session_retention_period": "365d"})
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(client)
calls = client.db.execute_raw.call_args_list
tables = [call[0][0] for call in calls]
@ -830,7 +769,6 @@ async def test_session_retention_alone_cleans_both_session_rollups_and_the_daily
async def test_health_check_retention_alone_cleans_only_the_health_check_table():
client = _mock_prisma_for_retention([0])
cleaner = SpendLogCleanup(general_settings={"maximum_health_check_retention_period": "30d"})
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
@ -846,7 +784,6 @@ async def test_health_check_retention_alone_cleans_only_the_health_check_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
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
@ -859,7 +796,6 @@ async def test_daily_tag_spend_retention_alone_prunes_only_that_table_by_calenda
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)
@ -875,7 +811,6 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
"maximum_health_check_retention_period": "30d",
}
)
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(client)
cutoffs = {
(
@ -903,7 +838,6 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon():
async def test_no_retention_keys_means_no_cleanup_at_all():
client = _mock_prisma_for_retention([])
cleaner = SpendLogCleanup(general_settings={})
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(client)
assert client.db.execute_raw.await_count == 0
@ -981,7 +915,6 @@ async def test_run_budget_is_shared_across_tables_not_granted_per_table():
},
clock=clock,
)
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -1014,7 +947,6 @@ async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete()
},
clock=clock,
)
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -1145,7 +1077,6 @@ async def test_a_bound_changed_after_construction_reaches_the_next_run(
every dashboard change inert until the process restarts."""
settings = {"maximum_spend_logs_retention_period": "7d"}
cleaner = SpendLogCleanup(general_settings=settings)
cleaner.pod_lock_manager = None
assert getattr(cleaner, attribute) != expected
settings[setting_name] = setting_value
@ -1163,7 +1094,6 @@ async def test_a_bound_cleared_after_construction_falls_back_to_its_default(clea
spells the clear as an explicit None or as an absent key."""
settings = {"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_batch_size": 137}
cleaner = SpendLogCleanup(general_settings=settings)
cleaner.pod_lock_manager = None
assert cleaner.batch_size == 137
if cleared_to_none:
@ -1203,32 +1133,6 @@ async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table():
assert mock_db.query_raw.call_args[0][2] == SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP
@pytest.mark.asyncio
async def test_a_run_skipped_because_another_pod_holds_the_lock_is_reported():
"""Operators need to tell "nothing to do" apart from "someone else is doing
it", so a lock-skipped run is recorded under its own outcome."""
recorded: list[str] = []
original_record_run = SpendLogCleanupMetrics.record_run
mock_prisma_client = MagicMock()
_wire_tx(mock_prisma_client.db)
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = MagicMock()
cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
cleaner.pod_lock_manager.release_lock = AsyncMock()
SpendLogCleanupMetrics.record_run = classmethod(lambda cls, outcome: recorded.append(outcome))
try:
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
finally:
SpendLogCleanupMetrics.record_run = original_record_run
assert recorded == ["skipped_locked"]
cleaner.pod_lock_manager.release_lock.assert_not_awaited()
@pytest.mark.asyncio
async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
"""
@ -1521,10 +1425,6 @@ async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_r
mock_prisma_client.db.execute_raw = _execute_raw
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = MagicMock()
cleaner.pod_lock_manager.redis_cache = MagicMock()
cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
cleaner.pod_lock_manager.release_lock = AsyncMock()
run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client))
await asyncio.wait_for(third_batch_reached.wait(), timeout=5)
@ -1534,7 +1434,6 @@ async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_r
assert _runs_recorded("aborted") == aborted_runs_before + 1
assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before
cleaner.pod_lock_manager.release_lock.assert_awaited_once()
mock_logger.exception.assert_not_called()
(error_call,) = mock_logger.error.call_args_list
rendered = error_call[0][0] % error_call[0][1:]
@ -1554,7 +1453,6 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0])
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()])
@ -1568,7 +1466,7 @@ async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatc
@pytest.mark.asyncio
async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_pod_waits_for_it(monkeypatch):
"""Without Redis there is no pod lock, and the two jobs are scheduled on two SpendLogCleanup instances"""
"""Two passes started in one pod, on two SpendLogCleanup instances, must not overlap"""
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
mock_logger = MagicMock()
@ -1593,7 +1491,6 @@ async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_p
size_client.db = size_db
age_cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
age_cleaner.pod_lock_manager = None
size_cleaner = _size_cleaner({})
slow_run = asyncio.ensure_future(age_cleaner.cleanup_old_spend_logs(slow_client))
@ -1643,7 +1540,6 @@ async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup():
},
clock=clock,
)
cleaner.pod_lock_manager = None
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -1672,7 +1568,6 @@ async def test_run_that_leaves_backlog_logs_a_warning_summary_naming_each_table(
"maximum_spend_logs_cleanup_max_batches": 2,
}
)
cleaner.pod_lock_manager = None
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -1697,7 +1592,6 @@ async def test_run_that_drains_every_table_logs_the_summary_at_info_not_warning(
mock_prisma_client.db = mock_db
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
cleaner.pod_lock_manager = None
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
await cleaner.cleanup_old_spend_logs(mock_prisma_client)
@ -1776,7 +1670,6 @@ def _wire_size_probes(
def _size_cleaner(settings: dict, **kwargs) -> SpendLogCleanup:
cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_size": "1GB", **settings}, **kwargs)
cleaner.pod_lock_manager = None
return cleaner
@ -2267,3 +2160,71 @@ async def test_size_state_round_trips_through_its_litellm_config_row_and_a_lost_
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await cleaner._save_size_state(prisma_client, _far_deadline(), _RUN_STATE)
assert any("Could not save the spend-log size cleanup state" in record.getMessage() for record in caplog.records)
@pytest.mark.parametrize(
("settings", "expected_spec"),
[
({}, "interval:86400s"),
({"maximum_spend_logs_retention_interval": "6h"}, "interval:21600s"),
(
{"maximum_spend_logs_retention_interval": "6h", "maximum_spend_logs_cleanup_cron": "0 4 * * *"},
"cron:0 4 * * *",
),
],
)
def test_the_retention_job_follows_the_cleanup_cron_when_set_and_the_retention_interval_otherwise(
settings: dict[str, str], expected_spec: str
) -> None:
schedule: Final = retention_schedule(settings)
assert not isinstance(schedule, InvalidSchedule)
assert schedule.spec == expected_spec
@pytest.mark.parametrize("cron", [5, ["0 4 * * *"], "not a cron"])
def test_an_invalid_cleanup_cron_is_rejected_instead_of_falling_back_to_the_retention_interval(cron: object) -> None:
schedule: Final = retention_schedule(
{"maximum_spend_logs_cleanup_cron": cron, "maximum_spend_logs_retention_interval": "1h"}
)
assert isinstance(schedule, InvalidSchedule)
async def test_the_retention_job_records_an_aborted_cleanup_as_failed(monkeypatch: pytest.MonkeyPatch) -> None:
import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 1)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_retention_period": "7d"})
mock_prisma_client: Final = MagicMock()
_wire_tx(mock_prisma_client.db)
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
result: Final = await SPEND_LOG_RETENTION_JOB.run(mock_prisma_client)
assert result == "FAILED"
@pytest.mark.parametrize("job", [SPEND_LOG_RETENTION_JOB, SPEND_LOG_SIZE_CHECK_JOB], ids=["retention", "size"])
async def test_a_cleanup_job_with_no_rule_configured_succeeds_without_touching_the_database(
monkeypatch: pytest.MonkeyPatch, job: ScheduledJob
) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.execute_raw = AsyncMock()
result: Final = await job.run(mock_prisma_client)
assert result == "SUCCEEDED"
mock_prisma_client.db.execute_raw.assert_not_awaited()
def test_a_cleanup_lease_outlasts_the_configured_run_budget(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "10m"})
for job in (SPEND_LOG_RETENTION_JOB, SPEND_LOG_SIZE_CHECK_JOB):
assert job.lease().total_seconds() > 600
def test_the_size_check_job_follows_the_size_check_interval(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_size_check_interval": "6h"})
schedule: Final = SPEND_LOG_SIZE_CHECK_JOB.schedule()
assert not isinstance(schedule, InvalidSchedule)
assert schedule.spec == "interval:21600s"