scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge 09f2926857 into 4bf75e5f51

Browse source
This commit is contained in:
jesus-berri 2026-11-05 16:47:59 -08:00 • committed by GitHub
commit f594d0f5eb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 1202 additions and 515 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_scheduled_job_runs/migration.sql Normal file
View file

@ -0,0 +1,24 @@
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),
"log" TEXT,
"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");

26
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1146,6 +1146,32 @@ 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)
log String?
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

5
litellm/constants.py
View file

@ -1919,6 +1919,11 @@ 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_LOG_MAX_CHARS: Final = get_env_int("SCHEDULED_JOB_LOG_MAX_CHARS", 8000)
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

157
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
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,16 @@ 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,
JobFailed,
JobResult,
JobSucceeded,
Schedule,
ScheduledJob,
cron_schedule,
interval_schedule,
)
from litellm.proxy.utils import PrismaClient
from litellm.types.llms.base import LiteLLMBaseModel
@ -238,6 +250,12 @@ def _cutoff_text(cutoff: Cutoff) -> str:
return cutoff.isoformat() if isinstance(cutoff, datetime) else cutoff
@dataclass(frozen=True, slots=True)
class CleanupRunReport:
outcome: RunOutcome
summary: str
@dataclass(frozen=True, slots=True)
class TableCleanupResult:
"""Outcome of pruning one table, so the caller can report why a run ended."""
@ -395,7 +413,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 +438,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 +449,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,
@ -1409,42 +1423,40 @@ class SpendLogCleanup:
results: tuple[TableCleanupResult, ...],
elapsed_seconds: float,
size_report: SizePassReport | None = None,
) -> None:
) -> str:
per_table: Final = ", ".join(
f"{result.table_name}: deleted={result.rows_deleted} stop_reason={result.stop_reason}" for result in results
)
message: Final = "Spend log cleanup run finished: outcome=%s elapsed=%.1fs [%s]%s"
storage: Final = cls._storage_summary(size_report)
message: Final = (
f"Spend log cleanup run finished: outcome={outcome} elapsed={elapsed_seconds:.1f}s [{per_table}]{storage}"
)
if outcome == "completed" and (size_report is None or size_report.target_reached):
verbose_proxy_logger.info(message, outcome, elapsed_seconds, per_table, storage)
return
verbose_proxy_logger.warning(message, outcome, elapsed_seconds, per_table, storage)
verbose_proxy_logger.info(message)
return message
verbose_proxy_logger.warning(message)
return message
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) -> CleanupRunReport:
"""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) -> CleanupRunReport:
"""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) -> CleanupRunReport:
"""
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) -> CleanupRunReport:
run_started_at: Final = self._clock()
progress: Final = _RunProgress()
progress_token: Final = _run_progress.set(progress)
@ -1471,29 +1483,12 @@ class SpendLogCleanup:
and active_size_rule is None
):
SpendLogCleanupMetrics.record_run("skipped_disabled")
return
return CleanupRunReport(outcome="skipped_disabled", summary="no retention or size rule is configured")
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 CleanupRunReport(outcome="skipped_disabled", summary="the retention period could not be read")
deadline: Final = self._clock() + self.run_budget_seconds
size_group: Final = int(active_size_rule is not None)
@ -1563,7 +1558,10 @@ class SpendLogCleanup:
)
outcome: Final = self._run_outcome(results)
SpendLogCleanupMetrics.record_run(outcome)
self._log_run_summary(outcome, results, self._clock() - run_started_at, size_report)
return CleanupRunReport(
outcome=outcome,
summary=self._log_run_summary(outcome, results, self._clock() - run_started_at, size_report),
)
except asyncio.CancelledError:
verbose_proxy_logger.error(
@ -1583,10 +1581,61 @@ class SpendLogCleanup:
e,
)
SpendLogCleanupMetrics.record_run("aborted")
return # Return after error handling
return CleanupRunReport(outcome="aborted", summary=f"Spend log cleanup failed: {type(e).__name__}: {e}")
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(report: CleanupRunReport) -> JobResult:
return JobFailed(report.summary) if report.outcome == "aborted" else JobSucceeded(report.summary)
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,
)

479
litellm/proxy/db/scheduled_job_runs.py Normal file
View file

@ -0,0 +1,479 @@
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, 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_LOG_MAX_CHARS,
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 prisma import Prisma
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}")
return CronSchedule(expression=value, trigger=trigger)
@dataclass(frozen=True, slots=True)
class JobSucceeded:
log: str
@dataclass(frozen=True, slots=True)
class JobFailed:
log: str
JobResult: TypeAlias = JobSucceeded | JobFailed
@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(),
"log" = concat_ws(
E'\\n', "log", 'lease expired at ' || "lease_expires_at" || ' while ' || "claimed_by" || ' held it'
)
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" <> 'CANCELLED') 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(), "log" = 'superseded by a schedule change or a duplicate'
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(), "log" = $4
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}"))
if not await _try_group_lock(tx, self.job.lock_group):
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, JobFailed("cancelled before it finished, usually a pod shutdown"))
)
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 as e: # 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 JobFailed(f"{type(e).__name__}: {e}")
async def _finish(self, prisma_client: "PrismaClient", run_id: str, result: JobResult) -> None:
status: Final = "SUCCEEDED" if isinstance(result, JobSucceeded) else "FAILED"
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, status, result.log[:SCHEDULED_JOB_LOG_MAX_CHARS]
)
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,
status,
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,
status,
)
return
verbose_proxy_logger.info("Scheduled job %s run %s %s", self.job.name, run_id, status)
async def _try_group_lock(tx: "Prisma", lock_group: str) -> bool:
rows: Final = _LOCKED.validate_python(
await tx.query_raw(
trusted_sql('SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"'),
f"litellm_scheduled_job:{lock_group}",
)
)
return bool(rows) and rows[0].locked
class _Locked(BaseModel):
model_config = ConfigDict(frozen=True)
locked: bool
_LOCKED: Final = TypeAdapter(list[_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)

189
litellm/proxy/proxy_server.py
View file

@ -269,8 +269,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]'`")
@ -365,8 +363,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,
)
@ -601,7 +599,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
@ -656,7 +653,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,
)
@ -682,6 +682,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,
@ -5860,6 +5861,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(
@ -8410,97 +8433,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
@ -11909,70 +11845,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 ###

26
litellm/proxy/schema.prisma
View file

@ -1146,6 +1146,32 @@ 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)
log String?
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

26
schema.prisma
View file

@ -1146,6 +1146,32 @@ 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)
log String?
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

1
tests/integration/database/test_cleanup_transactions.py
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})

282
tests/integration/database/test_scheduled_job_runs.py Normal file
View file

@ -0,0 +1,282 @@
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,
JobSucceeded,
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 JobSucceeded(f"{self.name} done")
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 [row["log"] for row in finished] == ["job done"]
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 "RuntimeError: disk on fire" in str(failed[0]["log"])
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"
assert "dead-pod" in str(reaped["log"])
(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_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"
assert "cancelled" in str(cancelled["log"])

169
tests/unit/proxy/db/test_scheduled_job_runs.py Normal file
View file

@ -0,0 +1,169 @@
from datetime import datetime, timedelta, timezone
from typing import Final
import pytest
from litellm.proxy.db.scheduled_job_runs import (
InvalidSchedule,
JobState,
NewRun,
PendingRun,
Schedule,
TickPlan,
cron_schedule,
interval_schedule,
next_run_at,
plan_tick,
)
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 * * * *", 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)

188
tests/unit/proxy/proxy_server/test_proxy_config.py
View file

@ -18,7 +18,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
@ -4647,30 +4647,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",
@ -4819,7 +4798,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)
@ -4854,33 +4833,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
@ -4899,29 +4851,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,
@ -4957,34 +4886,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",
@ -5024,89 +4925,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
# ---------------------------------------------------------------------------

145
tests/unit/proxy/test_spend_log_cleanup.py
View file

@ -28,10 +28,12 @@ from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
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
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@ -276,17 +278,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 +318,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 +355,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 +390,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 +422,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 +446,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 +671,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 +689,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):
"""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 +698,13 @@ 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
report = 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 report.outcome == "aborted"
assert "LiteLLM_SpendLogs: deleted=0 stop_reason=aborted" in report.summary
def test_cleanup_batch_size_env_var(monkeypatch):
@ -801,7 +741,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 +754,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 +768,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 +783,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 +795,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 +810,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 +837,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 +914,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 +946,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 +1076,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 +1093,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 +1132,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 +1424,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 +1433,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 +1452,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 +1465,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 +1490,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 +1539,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 +1567,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 +1591,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 +1669,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 +2159,22 @@ 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
):
schedule = retention_schedule(settings)
assert not isinstance(schedule, InvalidSchedule)
assert schedule.spec == expected_spec