diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_scheduled_job_runs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_scheduled_job_runs/migration.sql new file mode 100644 index 00000000000..4e9b168921c --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_scheduled_job_runs/migration.sql @@ -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"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 5a5d56cc6e2..4314a23e011 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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 diff --git a/litellm/constants.py b/litellm/constants.py index 966a5b71b8c..066b682362d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 8f9e2c8c6db..ad275f64096 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -328,6 +328,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_PromptTable", "LiteLLM_HealthCheckTable", "LiteLLM_SearchToolsTable", + "LiteLLM_ScheduledJobRun", "LiteLLM_SSOConfig", "LiteLLM_ManagedVectorStoreIndexTable", "LiteLLM_CacheConfig", diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index 29f727ea41b..7ca3075732f 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -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, +) diff --git a/litellm/proxy/db/scheduled_job_runs.py b/litellm/proxy/db/scheduled_job_runs.py new file mode 100644 index 00000000000..f38ab375345 --- /dev/null +++ b/litellm/proxy/db/scheduled_job_runs.py @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8320667a3a..0c34f77890f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 ### diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 5a5d56cc6e2..4314a23e011 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 diff --git a/schema.prisma b/schema.prisma index 5a5d56cc6e2..4314a23e011 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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 diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py index 202657d180e..be2487d5498 100644 --- a/tests/integration/database/test_cleanup_transactions.py +++ b/tests/integration/database/test_cleanup_transactions.py @@ -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}) diff --git a/tests/integration/database/test_scheduled_job_runs.py b/tests/integration/database/test_scheduled_job_runs.py new file mode 100644 index 00000000000..913111600bf --- /dev/null +++ b/tests/integration/database/test_scheduled_job_runs.py @@ -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" diff --git a/tests/unit/proxy/db/test_scheduled_job_runs.py b/tests/unit/proxy/db/test_scheduled_job_runs.py new file mode 100644 index 00000000000..9555a5ea61c --- /dev/null +++ b/tests/unit/proxy/db/test_scheduled_job_runs.py @@ -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) diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 91ff818da4c..9302b16f4f5 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -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 # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 4f96af7f02c..3d4d8f74c10 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -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() diff --git a/tests/unit/proxy/test_spend_log_cleanup.py b/tests/unit/proxy/test_spend_log_cleanup.py index 114c970e163..01289c5c37e 100644 --- a/tests/unit/proxy/test_spend_log_cleanup.py +++ b/tests/unit/proxy/test_spend_log_cleanup.py @@ -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"