mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(proxy): run spend log cleanup from a DB-backed scheduled job table (#50131)
* feat(proxy): run spend log cleanup from a DB-backed scheduled job table Spend log retention ran from an in-process APScheduler timer that reset on every restart and needed Redis for leader election, so pods that restart more often than the interval never cleaned up at all. Add LiteLLM_ScheduledJobRun, a generic table with one row per job run. Every pod polls it, plans the next run idempotently and claims a due run under a transaction-scoped advisory lock, so exactly one pod runs each occurrence. Runs carry a lease, a status and a log. Intervals count from the last finished run and survive restarts. Schedule edits cancel and replace the pending run. The retention and size-check jobs share a lock group so they never overlap, and the Redis pod lock is gone. * fix(proxy): plan an earlier cron edit from now, and cover the job runner Moving the cleanup cron earlier queued tomorrow's run because the next occurrence was planned after the pending run being cancelled. Plan from finished runs only. Register LiteLLM_ScheduledJobRun as a span table, keep the advisory lock query inside the tick's span, move the non-string cron startup test to the schedule parser, and add unit tests for the runner and job wiring. * test(proxy): annotate new spend log cleanup test locals as Final * test(proxy): type the parameters of the new spend log cleanup tests * test(proxy): leave an hour of margin when moving a cron earlier so an hour boundary cannot flake it * fix(proxy): reject a cron that never fires instead of failing every tick A crontab like 0 0 31 2 * parses but has no next fire time, so next_run_at raised on every tick and the job was never scheduled again. cron_schedule now returns InvalidSchedule for it, which logs and keeps the pending run like any other bad config * refactor(proxy): drop the informational log column from LiteLLM_ScheduledJobRun Nothing read the column back. Run outcomes stay in status, and the reasons already go to the proxy logs, so JobResult becomes a SUCCEEDED or FAILED literal and the cleanup returns its RunOutcome instead of a summary report
This commit is contained in:
parent
e7fd0d6cba
commit
d40dedf23b
15 changed files with 1417 additions and 570 deletions
|
|
@ -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");
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -328,6 +328,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
|||
"LiteLLM_PromptTable",
|
||||
"LiteLLM_HealthCheckTable",
|
||||
"LiteLLM_SearchToolsTable",
|
||||
"LiteLLM_ScheduledJobRun",
|
||||
"LiteLLM_SSOConfig",
|
||||
"LiteLLM_ManagedVectorStoreIndexTable",
|
||||
"LiteLLM_CacheConfig",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
454
litellm/proxy/db/scheduled_job_runs.py
Normal file
454
litellm/proxy/db/scheduled_job_runs.py
Normal file
|
|
@ -0,0 +1,454 @@
|
|||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
import uuid
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias
|
||||
|
||||
from apscheduler.triggers import cron # pyright: ignore[reportMissingTypeStubs, reportUnknownVariableType] # untyped
|
||||
from pydantic import AwareDatetime, BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
SCHEDULED_JOB_RUN_HISTORY_LIMIT,
|
||||
SCHEDULED_JOB_TICK_TIMEOUT_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy.db.db_span import db_span
|
||||
from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import bounded_tx, trusted_sql
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_TABLE: Final = "LiteLLM_ScheduledJobRun"
|
||||
_FIRE_TIME: Final = TypeAdapter[datetime](AwareDatetime)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class IntervalSchedule:
|
||||
seconds: float
|
||||
|
||||
@property
|
||||
def spec(self) -> str:
|
||||
return f"interval:{self.seconds:g}s"
|
||||
|
||||
|
||||
class FireTimes(Protocol):
|
||||
def get_next_fire_time(self, previous_fire_time: datetime | None, now: datetime) -> datetime | None: ...
|
||||
|
||||
|
||||
class JobScheduler(Protocol):
|
||||
"""The slice of APScheduler's AsyncIOScheduler that registers a poll, which ships no type information."""
|
||||
|
||||
def get_job(self, job_id: str) -> object | None: ...
|
||||
|
||||
def remove_job(self, job_id: str) -> None: ...
|
||||
|
||||
def add_job(
|
||||
self,
|
||||
func: Callable[["PrismaClient | None"], Awaitable[None]],
|
||||
trigger: str,
|
||||
*,
|
||||
args: Sequence[object],
|
||||
id: str,
|
||||
replace_existing: bool,
|
||||
max_instances: int,
|
||||
misfire_grace_time: int,
|
||||
seconds: float,
|
||||
) -> object: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CronSchedule:
|
||||
expression: str
|
||||
trigger: FireTimes
|
||||
|
||||
@property
|
||||
def spec(self) -> str:
|
||||
return f"cron:{self.expression}"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InvalidSchedule:
|
||||
reason: str
|
||||
|
||||
|
||||
Schedule: TypeAlias = IntervalSchedule | CronSchedule
|
||||
|
||||
|
||||
def interval_schedule(setting_name: str, value: object) -> Schedule | InvalidSchedule:
|
||||
if not isinstance(value, str):
|
||||
return InvalidSchedule(f"{setting_name} must be a duration string like '1d', got {value!r}")
|
||||
try:
|
||||
seconds: Final = duration_in_seconds(value)
|
||||
timedelta(seconds=seconds)
|
||||
except (ValueError, OverflowError):
|
||||
return InvalidSchedule(f"{setting_name} is not a valid duration: {value!r}")
|
||||
if seconds <= 0:
|
||||
return InvalidSchedule(f"{setting_name} must be positive, got {value!r}")
|
||||
return IntervalSchedule(seconds=seconds)
|
||||
|
||||
|
||||
def _cron_trigger(expression: str) -> FireTimes:
|
||||
return cron.CronTrigger.from_crontab(expression, timezone=timezone.utc) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # untyped
|
||||
|
||||
|
||||
def cron_schedule(setting_name: str, value: object) -> Schedule | InvalidSchedule:
|
||||
if not isinstance(value, str):
|
||||
return InvalidSchedule(f"{setting_name} must be a crontab string, got {value!r}")
|
||||
try:
|
||||
trigger: Final = _cron_trigger(value)
|
||||
except (ValueError, TypeError):
|
||||
return InvalidSchedule(f"{setting_name} is not a valid crontab expression: {value!r}")
|
||||
if trigger.get_next_fire_time(None, datetime.now(timezone.utc)) is None:
|
||||
return InvalidSchedule(f"{setting_name} never fires: {value!r}")
|
||||
return CronSchedule(expression=value, trigger=trigger)
|
||||
|
||||
|
||||
JobResult: TypeAlias = Literal["SUCCEEDED", "FAILED"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ScheduledJob:
|
||||
"""
|
||||
A job that runs on exactly one pod per occurrence.
|
||||
|
||||
Jobs sharing a lock_group never run at the same time, on any pod. schedule and lease are read on every tick,
|
||||
so a settings change applies without re-registering the job.
|
||||
"""
|
||||
|
||||
name: str
|
||||
lock_group: str
|
||||
schedule: Callable[[], Schedule | InvalidSchedule]
|
||||
lease: Callable[[], timedelta]
|
||||
run: Callable[["PrismaClient"], Awaitable[JobResult]]
|
||||
|
||||
|
||||
class PendingRun(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
schedule: str
|
||||
scheduled_for: AwareDatetime
|
||||
|
||||
|
||||
class JobState(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
now: AwareDatetime
|
||||
group_running: bool
|
||||
job_running: bool
|
||||
last_finished_at: AwareDatetime | None
|
||||
last_scheduled_for: AwareDatetime | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NewRun:
|
||||
id: str
|
||||
schedule: str
|
||||
scheduled_for: datetime
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TickPlan:
|
||||
cancel: tuple[str, ...]
|
||||
insert: NewRun | None
|
||||
claim: str | None
|
||||
|
||||
|
||||
def next_run_at(schedule: Schedule, state: JobState) -> datetime:
|
||||
"""
|
||||
When the next occurrence is due.
|
||||
|
||||
An interval counts from the last finished run, so a restart never resets it and a missed run happens once, right
|
||||
away. A cron fires at its next time after both now and the last occurrence, so an outage collapses into the one
|
||||
occurrence that was already pending.
|
||||
"""
|
||||
match schedule:
|
||||
case IntervalSchedule(seconds=seconds):
|
||||
if state.last_finished_at is None:
|
||||
return state.now
|
||||
return state.last_finished_at + timedelta(seconds=seconds)
|
||||
case CronSchedule(trigger=trigger):
|
||||
anchor: Final = max(state.now, state.last_scheduled_for or state.now)
|
||||
return _FIRE_TIME.validate_python(trigger.get_next_fire_time(None, anchor + timedelta(microseconds=1)))
|
||||
|
||||
|
||||
def plan_tick(
|
||||
schedule: Schedule | InvalidSchedule, state: JobState, pending: tuple[PendingRun, ...], new_id: str
|
||||
) -> TickPlan:
|
||||
"""
|
||||
Decide one tick from the job's rows. Pure, so every pod reaches the same decision from the same rows.
|
||||
|
||||
A pending row planned from a different schedule is cancelled and replaced, so a settings change applies at once.
|
||||
With an invalid schedule nothing is replanned, and a pending row from the last valid one still runs.
|
||||
"""
|
||||
if isinstance(schedule, InvalidSchedule):
|
||||
return TickPlan(cancel=(), insert=None, claim=_claimable(pending[:1], state))
|
||||
current: Final = tuple(run for run in pending if run.schedule == schedule.spec)
|
||||
keep: Final = current[:1]
|
||||
cancel: Final = tuple(run.id for run in pending if run not in keep)
|
||||
if keep or state.job_running:
|
||||
return TickPlan(cancel=cancel, insert=None, claim=_claimable(keep, state))
|
||||
new_run: Final = NewRun(id=new_id, schedule=schedule.spec, scheduled_for=next_run_at(schedule, state))
|
||||
claim: Final = new_run.id if new_run.scheduled_for <= state.now and not state.group_running else None
|
||||
return TickPlan(cancel=cancel, insert=new_run, claim=claim)
|
||||
|
||||
|
||||
def _claimable(pending: tuple[PendingRun, ...], state: JobState) -> str | None:
|
||||
if not pending or state.group_running or pending[0].scheduled_for > state.now:
|
||||
return None
|
||||
return pending[0].id
|
||||
|
||||
|
||||
def default_worker_id() -> str:
|
||||
return f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4().hex[:8]}"
|
||||
|
||||
|
||||
class _Claimed(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
id: str
|
||||
job_name: str
|
||||
|
||||
|
||||
_STATE: Final = TypeAdapter(list[JobState])
|
||||
_PENDING: Final = TypeAdapter(list[PendingRun])
|
||||
_CLAIMED: Final = TypeAdapter(list[_Claimed])
|
||||
|
||||
_REAP_SQL: Final = f"""
|
||||
UPDATE "{_TABLE}"
|
||||
SET "status" = 'FAILED', "finished_at" = now()
|
||||
WHERE "lock_group" = $1 AND "status" = 'RUNNING' AND "lease_expires_at" < now()
|
||||
RETURNING "id", "job_name"
|
||||
"""
|
||||
_STATE_SQL: Final = f"""
|
||||
SELECT now() AS "now",
|
||||
EXISTS (SELECT 1 FROM "{_TABLE}" WHERE "lock_group" = $2 AND "status" = 'RUNNING') AS "group_running",
|
||||
EXISTS (SELECT 1 FROM "{_TABLE}" WHERE "job_name" = $1 AND "status" = 'RUNNING') AS "job_running",
|
||||
(SELECT max("finished_at") FROM "{_TABLE}"
|
||||
WHERE "job_name" = $1 AND "status" IN ('SUCCEEDED', 'FAILED')) AS "last_finished_at",
|
||||
(SELECT max("scheduled_for") FROM "{_TABLE}"
|
||||
WHERE "job_name" = $1 AND "status" NOT IN ('CANCELLED', 'PENDING')) AS "last_scheduled_for"
|
||||
"""
|
||||
_PENDING_SQL: Final = f"""
|
||||
SELECT "id", "schedule", "scheduled_for" FROM "{_TABLE}"
|
||||
WHERE "job_name" = $1 AND "status" = 'PENDING'
|
||||
ORDER BY "scheduled_for", "created_at"
|
||||
"""
|
||||
_CANCEL_SQL: Final = f"""
|
||||
UPDATE "{_TABLE}" SET "status" = 'CANCELLED', "finished_at" = now()
|
||||
WHERE "id" = $1 AND "status" = 'PENDING'
|
||||
"""
|
||||
_INSERT_SQL: Final = f"""
|
||||
INSERT INTO "{_TABLE}" ("id", "job_name", "lock_group", "schedule", "scheduled_for")
|
||||
VALUES ($1, $2, $3, $4, $5::timestamptz)
|
||||
"""
|
||||
_CLAIM_SQL: Final = f"""
|
||||
UPDATE "{_TABLE}"
|
||||
SET "status" = 'RUNNING', "claimed_by" = $2, "claimed_at" = now(),
|
||||
"lease_expires_at" = now() + $3::double precision * interval '1 second'
|
||||
WHERE "id" = $1 AND "status" = 'PENDING'
|
||||
RETURNING "id", "job_name"
|
||||
"""
|
||||
_FINISH_SQL: Final = f"""
|
||||
UPDATE "{_TABLE}" SET "status" = $3::"ScheduledJobRunStatus", "finished_at" = now()
|
||||
WHERE "id" = $1 AND "claimed_by" = $2 AND "status" = 'RUNNING'
|
||||
"""
|
||||
_PRUNE_SQL: Final = f"""
|
||||
DELETE FROM "{_TABLE}" WHERE "id" IN (
|
||||
SELECT "id" FROM "{_TABLE}" WHERE "job_name" = $1 AND "finished_at" IS NOT NULL
|
||||
ORDER BY "finished_at" DESC OFFSET $2
|
||||
)
|
||||
"""
|
||||
|
||||
|
||||
class ScheduledJobRunner:
|
||||
"""
|
||||
Polls one ScheduledJob's rows in LiteLLM_ScheduledJobRun. Every pod runs the same tick.
|
||||
|
||||
Each tick runs in one transaction under a transaction-scoped advisory lock on the job's lock_group, so only one
|
||||
pod plans and claims for a group at a time and a pod that dies mid-tick releases it. The claimed run executes
|
||||
outside any transaction under a lease. A pod that dies mid-run leaves its row RUNNING until the lease expires,
|
||||
when the next tick on any pod marks it FAILED and plans the next occurrence.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
job: ScheduledJob,
|
||||
*,
|
||||
worker_id: str | None = None,
|
||||
new_id: Callable[[], str] = lambda: str(uuid.uuid4()),
|
||||
) -> None:
|
||||
self.job = job
|
||||
self.worker_id = worker_id or default_worker_id()
|
||||
self._new_id = new_id
|
||||
self._running = False
|
||||
|
||||
async def tick(self, prisma_client: "PrismaClient | None") -> None:
|
||||
"""
|
||||
Plan, claim and run whatever is due. Awaits the run, so graceful shutdown sees it as an in-flight job.
|
||||
|
||||
Register it with max_instances=2: the extra instance lets the polls that land during a long run return at
|
||||
once instead of APScheduler warning that it skipped them.
|
||||
"""
|
||||
if prisma_client is None or self._running:
|
||||
return
|
||||
self._running = True
|
||||
try:
|
||||
claimed: Final = await self._safe_plan_and_claim(prisma_client)
|
||||
if claimed is not None:
|
||||
await self._execute(prisma_client, claimed)
|
||||
finally:
|
||||
self._running = False
|
||||
|
||||
async def _safe_plan_and_claim(self, prisma_client: "PrismaClient") -> str | None:
|
||||
try:
|
||||
return await self._plan_and_claim(prisma_client)
|
||||
except Exception as e: # noqa: BLE001 - a failed tick is retried by the next one
|
||||
verbose_proxy_logger.warning("Scheduled job %s tick failed: %s: %s", self.job.name, type(e).__name__, e)
|
||||
return None
|
||||
|
||||
async def _plan_and_claim(self, prisma_client: "PrismaClient") -> str | None:
|
||||
schedule: Final = self.job.schedule()
|
||||
if isinstance(schedule, InvalidSchedule):
|
||||
_log_invalid_schedule(self.job.name, schedule.reason)
|
||||
timeout_ms: Final = int(SCHEDULED_JOB_TICK_TIMEOUT_SECONDS * 1000)
|
||||
async with db_span("scheduled_job_tick", _TABLE), bounded_tx(prisma_client, timeout_ms) as tx:
|
||||
await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
|
||||
await tx.execute_raw(trusted_sql(f"SET LOCAL lock_timeout = {timeout_ms}"))
|
||||
locked: Final = _LOCKED.validate_python(
|
||||
await tx.query_raw(trusted_sql(_GROUP_LOCK_SQL), f"litellm_scheduled_job:{self.job.lock_group}")
|
||||
)
|
||||
if not locked or not locked[0].locked:
|
||||
return None
|
||||
for reaped in _CLAIMED.validate_python(await tx.query_raw(trusted_sql(_REAP_SQL), self.job.lock_group)):
|
||||
verbose_proxy_logger.warning(
|
||||
"Scheduled job %s run %s lost its lease; marked FAILED", reaped.job_name, reaped.id
|
||||
)
|
||||
state: Final = _STATE.validate_python(
|
||||
await tx.query_raw(trusted_sql(_STATE_SQL), self.job.name, self.job.lock_group)
|
||||
)[0]
|
||||
pending: Final = tuple(
|
||||
_PENDING.validate_python(await tx.query_raw(trusted_sql(_PENDING_SQL), self.job.name))
|
||||
)
|
||||
plan: Final = plan_tick(schedule, state, pending, self._new_id())
|
||||
for run_id in plan.cancel:
|
||||
await tx.execute_raw(trusted_sql(_CANCEL_SQL), run_id)
|
||||
if plan.insert is not None:
|
||||
await tx.execute_raw(
|
||||
trusted_sql(_INSERT_SQL),
|
||||
plan.insert.id,
|
||||
self.job.name,
|
||||
self.job.lock_group,
|
||||
plan.insert.schedule,
|
||||
plan.insert.scheduled_for.isoformat(),
|
||||
)
|
||||
if plan.claim is None:
|
||||
return None
|
||||
claimed: Final = _CLAIMED.validate_python(
|
||||
await tx.query_raw(
|
||||
trusted_sql(_CLAIM_SQL), plan.claim, self.worker_id, self.job.lease().total_seconds()
|
||||
)
|
||||
)
|
||||
return claimed[0].id if claimed else None
|
||||
|
||||
async def _execute(self, prisma_client: "PrismaClient", run_id: str) -> None:
|
||||
verbose_proxy_logger.info("Scheduled job %s run %s claimed by %s", self.job.name, run_id, self.worker_id)
|
||||
try:
|
||||
result: Final = await self._run_job(prisma_client)
|
||||
except asyncio.CancelledError:
|
||||
await asyncio.shield(self._finish(prisma_client, run_id, "FAILED"))
|
||||
raise
|
||||
await self._finish(prisma_client, run_id, result)
|
||||
next_run_id: Final = await self._safe_plan_and_claim(prisma_client)
|
||||
if next_run_id is not None:
|
||||
await self._execute(prisma_client, next_run_id)
|
||||
|
||||
async def _run_job(self, prisma_client: "PrismaClient") -> JobResult:
|
||||
try:
|
||||
return await self.job.run(prisma_client)
|
||||
except Exception: # noqa: BLE001 - the failure is recorded on the run row instead of escaping the task
|
||||
verbose_proxy_logger.exception("Scheduled job %s raised", self.job.name)
|
||||
return "FAILED"
|
||||
|
||||
async def _finish(self, prisma_client: "PrismaClient", run_id: str, result: JobResult) -> None:
|
||||
timeout_ms: Final = int(SCHEDULED_JOB_TICK_TIMEOUT_SECONDS * 1000)
|
||||
try:
|
||||
async with db_span("scheduled_job_finish", _TABLE), bounded_tx(prisma_client, timeout_ms) as tx:
|
||||
await tx.execute_raw(trusted_sql(f"SET LOCAL statement_timeout = {timeout_ms}"))
|
||||
updated: Final = await tx.execute_raw(trusted_sql(_FINISH_SQL), run_id, self.worker_id, result)
|
||||
await tx.execute_raw(trusted_sql(_PRUNE_SQL), self.job.name, SCHEDULED_JOB_RUN_HISTORY_LIMIT)
|
||||
except Exception as e: # noqa: BLE001 - the lease expiry marks the row when this write cannot
|
||||
verbose_proxy_logger.warning(
|
||||
"Scheduled job %s run %s could not record %s: %s: %s",
|
||||
self.job.name,
|
||||
run_id,
|
||||
result,
|
||||
type(e).__name__,
|
||||
e,
|
||||
)
|
||||
return
|
||||
if updated == 0:
|
||||
verbose_proxy_logger.warning(
|
||||
"Scheduled job %s run %s finished after its lease expired; its %s result was not recorded",
|
||||
self.job.name,
|
||||
run_id,
|
||||
result,
|
||||
)
|
||||
return
|
||||
verbose_proxy_logger.info("Scheduled job %s run %s %s", self.job.name, run_id, result)
|
||||
|
||||
|
||||
class _Locked(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
locked: bool
|
||||
|
||||
|
||||
_LOCKED: Final = TypeAdapter(list[_Locked])
|
||||
_GROUP_LOCK_SQL: Final = 'SELECT pg_try_advisory_xact_lock(hashtext($1)) AS "locked"'
|
||||
|
||||
|
||||
@lru_cache(maxsize=64)
|
||||
def _log_invalid_schedule(job_name: str, reason: str) -> None:
|
||||
verbose_proxy_logger.error(
|
||||
"Scheduled job %s has an invalid schedule and will not be replanned: %s", job_name, reason
|
||||
)
|
||||
|
||||
|
||||
def register_poll(
|
||||
scheduler: JobScheduler,
|
||||
prisma_client: "PrismaClient | None",
|
||||
job_id: str,
|
||||
job: ScheduledJob,
|
||||
*,
|
||||
wanted: bool,
|
||||
poll_seconds: float,
|
||||
misfire_grace_time: int,
|
||||
) -> None:
|
||||
"""
|
||||
Make every pod poll job, or stop polling it. A poll that already exists is kept: its runner reads the schedule
|
||||
on every tick, so a schedule edit needs no re-registration.
|
||||
"""
|
||||
registered: Final = scheduler.get_job(job_id) is not None
|
||||
if registered and not wanted:
|
||||
scheduler.remove_job(job_id)
|
||||
verbose_proxy_logger.info("Stopped polling scheduled job %s", job.name)
|
||||
return
|
||||
if registered or not wanted:
|
||||
return
|
||||
scheduler.add_job(
|
||||
ScheduledJobRunner(job).tick,
|
||||
"interval",
|
||||
args=(prisma_client,),
|
||||
id=job_id,
|
||||
replace_existing=True,
|
||||
max_instances=2,
|
||||
misfire_grace_time=misfire_grace_time,
|
||||
seconds=poll_seconds,
|
||||
)
|
||||
verbose_proxy_logger.info("Polling scheduled job %s every %ss", job.name, poll_seconds)
|
||||
|
|
@ -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 ###
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
|
|||
294
tests/integration/database/test_scheduled_job_runs.py
Normal file
294
tests/integration/database/test_scheduled_job_runs.py
Normal file
|
|
@ -0,0 +1,294 @@
|
|||
import asyncio
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
from integration._support.database import scratch_database
|
||||
from prisma import Prisma
|
||||
from psycopg.rows import dict_row
|
||||
|
||||
from litellm.proxy.db.scheduled_job_runs import (
|
||||
JobResult,
|
||||
Schedule,
|
||||
ScheduledJob,
|
||||
ScheduledJobRunner,
|
||||
cron_schedule,
|
||||
interval_schedule,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_MIGRATION: Final = next(
|
||||
Path(__file__)
|
||||
.resolve()
|
||||
.parents[3]
|
||||
.joinpath("litellm-proxy-extras/litellm_proxy_extras/migrations")
|
||||
.glob("*_scheduled_job_runs/migration.sql")
|
||||
)
|
||||
_PODS: Final = 4
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Connection:
|
||||
db: Prisma
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Database:
|
||||
url: str
|
||||
pods: tuple["PrismaClient", ...]
|
||||
|
||||
def rows(self, where: str = "TRUE") -> list[dict[str, object]]:
|
||||
with psycopg.connect(self.url, row_factory=dict_row) as connection:
|
||||
return connection.execute(
|
||||
f'SELECT * FROM "LiteLLM_ScheduledJobRun" WHERE {where} ORDER BY "created_at", "scheduled_for"'
|
||||
).fetchall()
|
||||
|
||||
def execute(self, statement: str, parameters: tuple[object, ...] = ()) -> None:
|
||||
with psycopg.connect(self.url, autocommit=True) as connection:
|
||||
connection.execute(statement, parameters)
|
||||
|
||||
|
||||
class _Job:
|
||||
"""A job whose runs are counted and can be held open until the test releases them."""
|
||||
|
||||
def __init__(self, name: str = "job", lock_group: str = "group", schedule: str = "1h") -> None:
|
||||
self.name = name
|
||||
self.lock_group = lock_group
|
||||
self.schedule_value = schedule
|
||||
self.started = 0
|
||||
self.release = asyncio.Event()
|
||||
self.release.set()
|
||||
self.outcome: Callable[[], Awaitable[JobResult]] = self._succeed
|
||||
|
||||
async def _succeed(self) -> JobResult:
|
||||
return "SUCCEEDED"
|
||||
|
||||
async def _run(self, _prisma_client: object) -> JobResult:
|
||||
self.started += 1
|
||||
await self.release.wait()
|
||||
return await self.outcome()
|
||||
|
||||
def _schedule(self) -> Schedule:
|
||||
parsed: Final = (
|
||||
cron_schedule("cron", self.schedule_value)
|
||||
if " " in self.schedule_value
|
||||
else interval_schedule("interval", self.schedule_value)
|
||||
)
|
||||
assert isinstance(parsed, Schedule)
|
||||
return parsed
|
||||
|
||||
def definition(self) -> ScheduledJob:
|
||||
return ScheduledJob(
|
||||
name=self.name,
|
||||
lock_group=self.lock_group,
|
||||
schedule=self._schedule,
|
||||
lease=lambda: timedelta(minutes=5),
|
||||
run=self._run,
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _scratch() -> AsyncIterator[_Database]:
|
||||
with scratch_database("scheduled_jobs") as url:
|
||||
with psycopg.connect(url, autocommit=True) as setup:
|
||||
setup.execute(_MIGRATION.read_text())
|
||||
connections: Final = tuple(_Connection(Prisma(datasource={"url": url})) for _ in range(_PODS))
|
||||
for connection in connections:
|
||||
await connection.db.connect()
|
||||
try:
|
||||
yield _Database(
|
||||
url=url,
|
||||
pods=tuple(
|
||||
cast("PrismaClient", connection) # cast-ok: the runner only reaches the client through .db
|
||||
for connection in connections
|
||||
),
|
||||
)
|
||||
finally:
|
||||
for connection in connections:
|
||||
await connection.db.disconnect()
|
||||
|
||||
|
||||
def _runners(job: _Job) -> tuple[ScheduledJobRunner, ...]:
|
||||
return tuple(ScheduledJobRunner(job.definition(), worker_id=f"pod-{index}") for index in range(_PODS))
|
||||
|
||||
|
||||
async def _tick_all(runners: tuple[ScheduledJobRunner, ...], database: _Database) -> None:
|
||||
await asyncio.gather(*(runner.tick(pod) for runner, pod in zip(runners, database.pods)))
|
||||
|
||||
|
||||
async def _until(condition: Callable[[], bool], seconds: float = 10) -> None:
|
||||
deadline: Final = time.monotonic() + seconds
|
||||
while not condition():
|
||||
assert time.monotonic() < deadline, "condition never held"
|
||||
await asyncio.sleep(0.02)
|
||||
|
||||
|
||||
async def test_every_pod_ticks_but_exactly_one_runs_and_the_next_run_is_planned_one_interval_after_it_finished() -> (
|
||||
None
|
||||
):
|
||||
async with _scratch() as database:
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
runners: Final = _runners(job)
|
||||
ticks: Final = asyncio.gather(*(runner.tick(pod) for runner, pod in zip(runners, database.pods)))
|
||||
await _until(lambda: job.started == 1)
|
||||
await asyncio.sleep(0.5)
|
||||
running: Final = database.rows("\"status\" = 'RUNNING'")
|
||||
assert job.started == 1
|
||||
assert len(running) == 1
|
||||
job.release.set()
|
||||
await ticks
|
||||
|
||||
finished: Final = database.rows("\"status\" = 'SUCCEEDED'")
|
||||
pending: Final = database.rows("\"status\" = 'PENDING'")
|
||||
assert len(finished) == 1
|
||||
assert finished[0]["claimed_by"] == running[0]["claimed_by"]
|
||||
assert len(pending) == 1
|
||||
assert pending[0]["scheduled_for"] == finished[0]["finished_at"] + timedelta(hours=1)
|
||||
|
||||
await _tick_all(runners, database)
|
||||
assert job.started == 1, "a run was started before the next occurrence was due"
|
||||
assert len(database.rows("\"status\" = 'PENDING'")) == 1
|
||||
|
||||
|
||||
async def test_jobs_in_one_lock_group_never_run_at_the_same_time_on_different_pods() -> None:
|
||||
async with _scratch() as database:
|
||||
first: Final = _Job(name="first")
|
||||
second: Final = _Job(name="second")
|
||||
first.release.clear()
|
||||
first_runner: Final = ScheduledJobRunner(first.definition(), worker_id="pod-a")
|
||||
second_runner: Final = ScheduledJobRunner(second.definition(), worker_id="pod-b")
|
||||
first_tick: Final = asyncio.create_task(first_runner.tick(database.pods[0]))
|
||||
await _until(lambda: first.started == 1)
|
||||
|
||||
await second_runner.tick(database.pods[1])
|
||||
assert second.started == 0
|
||||
assert [row["status"] for row in database.rows("\"job_name\" = 'second'")] == ["PENDING"]
|
||||
|
||||
first.release.set()
|
||||
await first_tick
|
||||
await second_runner.tick(database.pods[1])
|
||||
assert second.started == 1
|
||||
|
||||
|
||||
async def test_a_run_that_raises_is_recorded_as_failed_and_the_next_run_is_still_planned() -> None:
|
||||
async with _scratch() as database:
|
||||
job: Final = _Job()
|
||||
|
||||
async def explode() -> JobResult:
|
||||
raise RuntimeError("disk on fire")
|
||||
|
||||
job.outcome = explode
|
||||
await _tick_all(_runners(job), database)
|
||||
|
||||
failed: Final = database.rows("\"status\" = 'FAILED'")
|
||||
assert len(failed) == 1
|
||||
assert len(database.rows("\"status\" = 'PENDING'")) == 1
|
||||
|
||||
|
||||
async def test_a_run_held_by_a_dead_pod_is_failed_once_its_lease_expires_and_the_schedule_carries_on() -> None:
|
||||
async with _scratch() as database:
|
||||
database.execute(
|
||||
'INSERT INTO "LiteLLM_ScheduledJobRun" ("id", "job_name", "lock_group", "schedule", "scheduled_for", '
|
||||
'"status", "claimed_by", "claimed_at", "lease_expires_at") '
|
||||
"VALUES ('dead', 'job', 'group', 'interval:3600s', now() - interval '2 hours', 'RUNNING', 'dead-pod', "
|
||||
"now() - interval '2 hours', now() - interval '1 second')"
|
||||
)
|
||||
job: Final = _Job()
|
||||
await _tick_all(_runners(job), database)
|
||||
|
||||
(reaped,) = database.rows("\"id\" = 'dead'")
|
||||
assert reaped["status"] == "FAILED"
|
||||
(pending,) = database.rows("\"status\" = 'PENDING'")
|
||||
assert pending["scheduled_for"] == reaped["finished_at"] + timedelta(hours=1)
|
||||
assert job.started == 0
|
||||
|
||||
|
||||
async def test_a_run_held_by_a_live_pod_is_left_alone_and_nothing_else_starts() -> None:
|
||||
async with _scratch() as database:
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
runners: Final = _runners(job)
|
||||
holder: Final = asyncio.create_task(runners[0].tick(database.pods[0]))
|
||||
await _until(lambda: job.started == 1)
|
||||
|
||||
await asyncio.gather(*(runner.tick(pod) for runner, pod in zip(runners[1:], database.pods[1:])))
|
||||
assert [row["status"] for row in database.rows()] == ["RUNNING"]
|
||||
|
||||
job.release.set()
|
||||
await holder
|
||||
assert job.started == 1
|
||||
|
||||
|
||||
async def test_an_interval_already_overdue_when_the_pods_restart_runs_on_the_first_tick() -> None:
|
||||
async with _scratch() as database:
|
||||
database.execute(
|
||||
'INSERT INTO "LiteLLM_ScheduledJobRun" ("id", "job_name", "lock_group", "schedule", "scheduled_for", '
|
||||
'"status", "finished_at") '
|
||||
"VALUES ('yesterday', 'job', 'group', 'interval:86400s', now() - interval '2 days', 'SUCCEEDED', "
|
||||
"now() - interval '2 days')"
|
||||
)
|
||||
job: Final = _Job(schedule="1d")
|
||||
await _tick_all(_runners(job), database)
|
||||
|
||||
assert job.started == 1
|
||||
assert [row["status"] for row in database.rows("\"id\" <> 'yesterday'")] == ["SUCCEEDED", "PENDING"]
|
||||
|
||||
|
||||
async def test_a_schedule_change_replaces_the_pending_run_planned_from_the_old_schedule() -> None:
|
||||
async with _scratch() as database:
|
||||
job: Final = _Job(schedule="0 4 * * *")
|
||||
runners: Final = _runners(job)
|
||||
await _tick_all(runners, database)
|
||||
(old,) = database.rows("\"status\" = 'PENDING'")
|
||||
assert old["scheduled_for"].astimezone(timezone.utc).hour == 4
|
||||
|
||||
job.schedule_value = "0 5 * * *"
|
||||
await _tick_all(runners, database)
|
||||
|
||||
assert database.rows(f"\"id\" = '{old['id']}'")[0]["status"] == "CANCELLED"
|
||||
(replacement,) = database.rows("\"status\" = 'PENDING'")
|
||||
assert replacement["schedule"] == "cron:0 5 * * *"
|
||||
assert replacement["scheduled_for"].astimezone(timezone.utc).hour == 5
|
||||
assert replacement["scheduled_for"] > datetime.now(timezone.utc)
|
||||
assert job.started == 0
|
||||
|
||||
|
||||
async def test_moving_a_cron_earlier_plans_the_earlier_occurrence_instead_of_skipping_a_day() -> None:
|
||||
async with _scratch() as database:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
job: Final = _Job(schedule=f"0 {(now + timedelta(hours=3)).hour} * * *")
|
||||
runners: Final = _runners(job)
|
||||
await _tick_all(runners, database)
|
||||
(old,) = database.rows("\"status\" = 'PENDING'")
|
||||
|
||||
job.schedule_value = f"0 {(now + timedelta(hours=2)).hour} * * *"
|
||||
await _tick_all(runners, database)
|
||||
|
||||
(replacement,) = database.rows("\"status\" = 'PENDING'")
|
||||
assert replacement["scheduled_for"] < old["scheduled_for"]
|
||||
assert replacement["scheduled_for"] - datetime.now(timezone.utc) < timedelta(hours=2)
|
||||
|
||||
|
||||
async def test_a_run_cancelled_by_shutdown_is_recorded_as_failed_instead_of_holding_its_lease() -> None:
|
||||
async with _scratch() as database:
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
runner: Final = ScheduledJobRunner(job.definition(), worker_id="pod-shutting-down")
|
||||
tick: Final = asyncio.create_task(runner.tick(database.pods[0]))
|
||||
await _until(lambda: job.started == 1)
|
||||
|
||||
tick.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await tick
|
||||
|
||||
(cancelled,) = database.rows()
|
||||
assert cancelled["status"] == "FAILED"
|
||||
355
tests/unit/proxy/db/test_scheduled_job_runs.py
Normal file
355
tests/unit/proxy/db/test_scheduled_job_runs.py
Normal file
|
|
@ -0,0 +1,355 @@
|
|||
import asyncio
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Final, cast
|
||||
|
||||
import pytest
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
from litellm.proxy.db.scheduled_job_runs import (
|
||||
InvalidSchedule,
|
||||
JobResult,
|
||||
JobScheduler,
|
||||
JobState,
|
||||
NewRun,
|
||||
PendingRun,
|
||||
Schedule,
|
||||
ScheduledJob,
|
||||
ScheduledJobRunner,
|
||||
TickPlan,
|
||||
cron_schedule,
|
||||
interval_schedule,
|
||||
next_run_at,
|
||||
plan_tick,
|
||||
register_poll,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
NOW: Final = datetime(2026, 11, 4, 12, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _schedule(value: str) -> Schedule:
|
||||
parsed: Final = cron_schedule("cron", value) if " " in value else interval_schedule("interval", value)
|
||||
assert not isinstance(parsed, InvalidSchedule)
|
||||
return parsed
|
||||
|
||||
|
||||
def _state(
|
||||
*,
|
||||
group_running: bool = False,
|
||||
job_running: bool = False,
|
||||
last_finished_at: datetime | None = None,
|
||||
last_scheduled_for: datetime | None = None,
|
||||
) -> JobState:
|
||||
return JobState(
|
||||
now=NOW,
|
||||
group_running=group_running,
|
||||
job_running=job_running,
|
||||
last_finished_at=last_finished_at,
|
||||
last_scheduled_for=last_scheduled_for,
|
||||
)
|
||||
|
||||
|
||||
def _pending(run_id: str, schedule: str, scheduled_for: datetime) -> PendingRun:
|
||||
return PendingRun(id=run_id, schedule=schedule, scheduled_for=scheduled_for)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("schedule", "state", "expected"),
|
||||
[
|
||||
pytest.param("1h", _state(), NOW, id="interval-never-run-is-due-now"),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(last_finished_at=NOW - timedelta(hours=3)),
|
||||
NOW - timedelta(hours=2),
|
||||
id="interval-counts-from-last-finish-so-a-restart-cannot-postpone-it",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(last_finished_at=NOW - timedelta(minutes=10)),
|
||||
NOW + timedelta(minutes=50),
|
||||
id="interval-not-yet-elapsed",
|
||||
),
|
||||
pytest.param("0 4 * * *", _state(), datetime(2026, 11, 5, 4, tzinfo=timezone.utc), id="cron-next-fire"),
|
||||
pytest.param(
|
||||
"0 12 * * *",
|
||||
_state(last_scheduled_for=NOW),
|
||||
datetime(2026, 11, 5, 12, tzinfo=timezone.utc),
|
||||
id="cron-never-replans-the-occurrence-it-just-had",
|
||||
),
|
||||
pytest.param(
|
||||
"0 4 * * *",
|
||||
_state(last_scheduled_for=NOW - timedelta(days=3)),
|
||||
datetime(2026, 11, 5, 4, tzinfo=timezone.utc),
|
||||
id="cron-outage-collapses-to-one-next-fire",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_next_run_at(schedule: str, state: JobState, expected: datetime) -> None:
|
||||
assert next_run_at(_schedule(schedule), state) == expected
|
||||
|
||||
|
||||
DUE: Final = NOW - timedelta(minutes=1)
|
||||
LATER: Final = NOW + timedelta(hours=1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("schedule", "state", "pending", "expected"),
|
||||
[
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(),
|
||||
(),
|
||||
TickPlan(cancel=(), insert=NewRun(id="new", schedule="interval:3600s", scheduled_for=NOW), claim="new"),
|
||||
id="first-ever-tick-plans-and-claims",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(),
|
||||
(_pending("a", "interval:3600s", DUE),),
|
||||
TickPlan(cancel=(), insert=None, claim="a"),
|
||||
id="due-pending-is-claimed",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(),
|
||||
(_pending("a", "interval:3600s", LATER),),
|
||||
TickPlan(cancel=(), insert=None, claim=None),
|
||||
id="future-pending-waits",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(group_running=True),
|
||||
(_pending("a", "interval:3600s", DUE),),
|
||||
TickPlan(cancel=(), insert=None, claim=None),
|
||||
id="due-pending-waits-while-its-lock-group-runs",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(group_running=True),
|
||||
(),
|
||||
TickPlan(cancel=(), insert=NewRun(id="new", schedule="interval:3600s", scheduled_for=NOW), claim=None),
|
||||
id="a-sibling-run-still-lets-the-next-run-be-planned",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(group_running=True, job_running=True),
|
||||
(),
|
||||
TickPlan(cancel=(), insert=None, claim=None),
|
||||
id="nothing-is-planned-while-the-job-itself-runs",
|
||||
),
|
||||
pytest.param(
|
||||
"2h",
|
||||
_state(last_finished_at=NOW - timedelta(minutes=30)),
|
||||
(_pending("old", "interval:3600s", NOW + timedelta(minutes=30)),),
|
||||
TickPlan(
|
||||
cancel=("old",),
|
||||
insert=NewRun(id="new", schedule="interval:7200s", scheduled_for=NOW + timedelta(minutes=90)),
|
||||
claim=None,
|
||||
),
|
||||
id="schedule-change-replaces-the-pending-run",
|
||||
),
|
||||
pytest.param(
|
||||
"1h",
|
||||
_state(),
|
||||
(_pending("a", "interval:3600s", DUE), _pending("b", "interval:3600s", LATER)),
|
||||
TickPlan(cancel=("b",), insert=None, claim="a"),
|
||||
id="a-duplicate-pending-run-is-cancelled",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_plan_tick(schedule: str, state: JobState, pending: tuple[PendingRun, ...], expected: TickPlan) -> None:
|
||||
assert plan_tick(_schedule(schedule), state, pending, "new") == expected
|
||||
|
||||
|
||||
def test_an_invalid_schedule_replans_nothing_but_still_runs_what_the_last_valid_one_planned() -> None:
|
||||
pending: Final = (_pending("a", "interval:3600s", DUE),)
|
||||
plan: Final = plan_tick(InvalidSchedule("bad cron"), _state(), pending, "new")
|
||||
assert plan == TickPlan(cancel=(), insert=None, claim="a")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["not a cron", "0 0 * * * *", "0 0 31 2 *", "0 0 30 2 *", 5, None])
|
||||
def test_cron_schedule_rejects_what_is_not_a_crontab(value: object) -> None:
|
||||
assert isinstance(cron_schedule("maximum_spend_logs_cleanup_cron", value), InvalidSchedule)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["soon", "99999999999d", "0s", 86400, None])
|
||||
def test_interval_schedule_rejects_what_is_not_a_positive_duration(value: object) -> None:
|
||||
assert isinstance(interval_schedule("maximum_spend_logs_retention_interval", value), InvalidSchedule)
|
||||
|
||||
|
||||
class _ScriptedDatabase:
|
||||
"""Answers the runner's statements the way Postgres would for one due run, and records every write."""
|
||||
|
||||
def __init__(self, *, locked: bool = True, finish_updates: int = 1, broken: bool = False) -> None:
|
||||
self.locked = locked
|
||||
self.finish_updates = finish_updates
|
||||
self.broken = broken
|
||||
self.writes: tuple[tuple[str, tuple[object, ...]], ...] = ()
|
||||
self.db = self
|
||||
|
||||
@asynccontextmanager
|
||||
async def tx(self, timeout: timedelta) -> AsyncIterator["_ScriptedDatabase"]:
|
||||
if self.broken:
|
||||
raise ConnectionError("database unreachable")
|
||||
yield self
|
||||
|
||||
async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]:
|
||||
if "pg_try_advisory_xact_lock" in sql:
|
||||
return [{"locked": self.locked}]
|
||||
if "SELECT now()" in sql:
|
||||
return [
|
||||
{
|
||||
"now": NOW,
|
||||
"group_running": False,
|
||||
"job_running": False,
|
||||
"last_finished_at": None,
|
||||
"last_scheduled_for": None,
|
||||
}
|
||||
]
|
||||
if "'RUNNING', \"claimed_by\"" in sql:
|
||||
self.writes = (*self.writes, ("claim", args))
|
||||
return [{"id": args[0], "job_name": "job"}]
|
||||
return []
|
||||
|
||||
async def execute_raw(self, sql: str, *args: object) -> int:
|
||||
if sql.lstrip().startswith("SET LOCAL"):
|
||||
return 0
|
||||
if 'SET "status" = $3' in sql:
|
||||
self.writes = (*self.writes, ("finish", args))
|
||||
self.locked = False
|
||||
return self.finish_updates
|
||||
if sql.lstrip().startswith("INSERT"):
|
||||
self.writes = (*self.writes, ("insert", args))
|
||||
return 1
|
||||
|
||||
def client(self) -> "PrismaClient":
|
||||
return cast("PrismaClient", self) # cast-ok: the runner only reaches the client through .db
|
||||
|
||||
|
||||
def _finishes(database: _ScriptedDatabase) -> tuple[tuple[object, ...], ...]:
|
||||
return tuple(args for kind, args in database.writes if kind == "finish")
|
||||
|
||||
|
||||
class _Job:
|
||||
def __init__(self, result: JobResult | BaseException = "SUCCEEDED") -> None:
|
||||
self.result = result
|
||||
self.runs = 0
|
||||
self.release = asyncio.Event()
|
||||
self.release.set()
|
||||
|
||||
async def run(self, _prisma_client: object) -> JobResult:
|
||||
self.runs += 1
|
||||
await self.release.wait()
|
||||
if isinstance(self.result, BaseException):
|
||||
raise self.result
|
||||
return self.result
|
||||
|
||||
def runner(self) -> ScheduledJobRunner:
|
||||
definition: Final = ScheduledJob(
|
||||
name="job",
|
||||
lock_group="group",
|
||||
schedule=lambda: _schedule("1h"),
|
||||
lease=lambda: timedelta(minutes=5),
|
||||
run=self.run,
|
||||
)
|
||||
return ScheduledJobRunner(definition, worker_id="pod-a", new_id=lambda: "run-1")
|
||||
|
||||
|
||||
async def test_a_due_run_is_claimed_by_this_pod_run_and_recorded_as_succeeded() -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
job: Final = _Job("SUCCEEDED")
|
||||
await job.runner().tick(database.client())
|
||||
|
||||
assert job.runs == 1
|
||||
assert ("claim", ("run-1", "pod-a", 300.0)) in database.writes
|
||||
assert _finishes(database) == (("run-1", "pod-a", "SUCCEEDED"),)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result",
|
||||
[pytest.param("FAILED", id="reported"), pytest.param(RuntimeError("disk on fire"), id="raised")],
|
||||
)
|
||||
async def test_a_failed_run_is_recorded_as_failed(result: JobResult | BaseException) -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
await _Job(result).runner().tick(database.client())
|
||||
assert _finishes(database) == (("run-1", "pod-a", "FAILED"),)
|
||||
|
||||
|
||||
async def test_a_pod_that_does_not_get_the_group_lock_neither_plans_nor_runs() -> None:
|
||||
database: Final = _ScriptedDatabase(locked=False)
|
||||
job: Final = _Job()
|
||||
await job.runner().tick(database.client())
|
||||
assert job.runs == 0
|
||||
assert database.writes == ()
|
||||
|
||||
|
||||
async def test_an_unreachable_database_skips_the_tick_instead_of_raising_into_the_scheduler() -> None:
|
||||
job: Final = _Job()
|
||||
await job.runner().tick(_ScriptedDatabase(broken=True).client())
|
||||
assert job.runs == 0
|
||||
|
||||
|
||||
async def test_a_poll_landing_while_this_pod_still_runs_the_job_does_not_touch_the_database() -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
runner: Final = job.runner()
|
||||
first: Final = asyncio.create_task(runner.tick(database.client()))
|
||||
while job.runs == 0:
|
||||
await asyncio.sleep(0)
|
||||
writes_during_run: Final = database.writes
|
||||
|
||||
await runner.tick(database.client())
|
||||
assert database.writes == writes_during_run
|
||||
|
||||
job.release.set()
|
||||
await first
|
||||
assert job.runs == 1
|
||||
|
||||
|
||||
async def test_a_run_cancelled_mid_flight_records_failed_and_still_cancels() -> None:
|
||||
database: Final = _ScriptedDatabase()
|
||||
job: Final = _Job()
|
||||
job.release.clear()
|
||||
tick: Final = asyncio.create_task(job.runner().tick(database.client()))
|
||||
while job.runs == 0:
|
||||
await asyncio.sleep(0)
|
||||
tick.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await tick
|
||||
assert _finishes(database) == (("run-1", "pod-a", "FAILED"),)
|
||||
|
||||
|
||||
async def test_a_result_that_arrives_after_the_lease_expired_is_reported_not_recorded(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
database: Final = _ScriptedDatabase(finish_updates=0)
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
await _Job().runner().tick(database.client())
|
||||
assert any("after its lease expired" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
async def test_register_poll_adds_a_wanted_poll_once_and_removes_it_when_no_longer_wanted() -> None:
|
||||
scheduler: Final = AsyncIOScheduler()
|
||||
scheduler.start(paused=True)
|
||||
job_scheduler: Final = cast(JobScheduler, scheduler) # cast-ok: apscheduler ships no type information
|
||||
definition: Final = _Job().runner().job
|
||||
|
||||
def register(wanted: bool) -> None:
|
||||
register_poll(job_scheduler, None, "poll-id", definition, wanted=wanted, poll_seconds=30, misfire_grace_time=60)
|
||||
|
||||
try:
|
||||
register(True)
|
||||
first: Final = scheduler.get_job("poll-id")
|
||||
assert first is not None
|
||||
register(True)
|
||||
assert scheduler.get_job("poll-id") is first, "re-registering replaced the live poll and its runner"
|
||||
register(False)
|
||||
assert scheduler.get_job("poll-id") is None
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue