mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
feat(observability): expose scheduled background job and cron lock telemetry
Nothing recorded which background job ran on a pod, when, for how long, whether it succeeded, or how much work it moved, so during an incident job activity could only be inferred from database load. A job that overran its interval and started being skipped left no trace at all. One APScheduler listener instruments every registered job at once rather than each job growing its own instrumentation, and reports max_instances skips as a first-class result so a job falling behind its schedule is visible. The single-owner cron lock outcome becomes a metric that separates winning the lock from losing it from having no Redis to elect with. Three integrations registered their export job without an explicit id, leaving APScheduler to generate a uuid that would have grown the job_name label without bound; they now pin the id they already had a constant for. Refs LIT-5435
This commit is contained in:
parent
2aa65cdd81
commit
50d38a7cd8
12 changed files with 644 additions and 11 deletions
|
|
@ -4,7 +4,10 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import CLOUDZERO_EXPORT_INTERVAL_MINUTES
|
||||
from litellm.constants import (
|
||||
CLOUDZERO_EXPORT_INTERVAL_MINUTES,
|
||||
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -360,5 +363,7 @@ class CloudZeroLogger(CustomLogger):
|
|||
scheduler.add_job(
|
||||
cloudzero_logger.initialize_cloudzero_export_job,
|
||||
"interval",
|
||||
id=CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME,
|
||||
replace_existing=True,
|
||||
minutes=CLOUDZERO_EXPORT_INTERVAL_MINUTES,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -150,6 +150,8 @@ class FocusLogger(CustomLogger):
|
|||
trigger_kwargs: Final = focus_logger._build_scheduler_trigger()
|
||||
scheduler.add_job(
|
||||
focus_logger.initialize_focus_export_job,
|
||||
id=FOCUS_USAGE_DATA_JOB_NAME,
|
||||
replace_existing=True,
|
||||
**trigger_kwargs,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import asyncio
|
|||
import math
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast
|
||||
|
|
@ -59,6 +60,7 @@ if TYPE_CHECKING:
|
|||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from prometheus_client.metrics import MetricWrapperBase
|
||||
|
||||
from litellm.proxy.common_utils.scheduled_job_metrics import JobRun
|
||||
from litellm.proxy.db.db_pool_metrics import DBPoolMetricsUpdate
|
||||
else:
|
||||
AsyncIOScheduler = Any
|
||||
|
|
@ -77,6 +79,8 @@ _NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
|
|||
"purpose",
|
||||
"file_type",
|
||||
"result",
|
||||
"job_name",
|
||||
"cronjob_id",
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -713,6 +717,51 @@ class PrometheusLogger(CustomLogger):
|
|||
labelnames=[],
|
||||
)
|
||||
|
||||
########################################
|
||||
# Scheduled background jobs
|
||||
########################################
|
||||
self.litellm_scheduled_job_runs_total = self._counter_factory(
|
||||
name="litellm_scheduled_job_runs_total",
|
||||
documentation=(
|
||||
"Scheduled background job runs by outcome. result is one of success, error, missed, "
|
||||
"max_instances; max_instances means the previous run had not finished"
|
||||
),
|
||||
labelnames=("job_name", "result"),
|
||||
)
|
||||
|
||||
self.litellm_scheduled_job_duration_seconds = self._histogram_factory(
|
||||
"litellm_scheduled_job_duration_seconds",
|
||||
"Wall-clock duration of a scheduled background job run",
|
||||
labelnames=("job_name",),
|
||||
buckets=self.latency_buckets,
|
||||
)
|
||||
|
||||
self.litellm_scheduled_job_last_run_timestamp = self._gauge_factory(
|
||||
"litellm_scheduled_job_last_run_timestamp",
|
||||
"Unix timestamp of the last completed run of a scheduled background job",
|
||||
labelnames=("job_name",),
|
||||
multiprocess_mode="livemax",
|
||||
)
|
||||
|
||||
self.litellm_scheduled_job_items_processed_total = self._counter_factory(
|
||||
name="litellm_scheduled_job_items_processed_total",
|
||||
documentation=(
|
||||
"Items scheduled background jobs reported processing. A counter, not a last-value gauge: "
|
||||
"queues drain between runs, so the most recent cycle is usually zero and a gauge would hide "
|
||||
"the bursts"
|
||||
),
|
||||
labelnames=("job_name",),
|
||||
)
|
||||
|
||||
self.litellm_cronjob_lock_acquisitions_total = self._counter_factory(
|
||||
name="litellm_cronjob_lock_acquisitions_total",
|
||||
documentation=(
|
||||
"Attempts to take the single-owner lock for a cron job. result is one of acquired, "
|
||||
"not_acquired, no_redis; no_redis means no Redis is configured, so no pod can be elected"
|
||||
),
|
||||
labelnames=("cronjob_id", "result"),
|
||||
)
|
||||
|
||||
########################################
|
||||
# Database connection pool saturation
|
||||
########################################
|
||||
|
|
@ -3140,8 +3189,6 @@ class PrometheusLogger(CustomLogger):
|
|||
jobs_polled: Number of unprocessed batches found
|
||||
processed_models: List of (model, api_provider) tuples for processed jobs
|
||||
"""
|
||||
import time
|
||||
|
||||
try:
|
||||
self.litellm_check_batch_cost_last_run_timestamp.set(time.time())
|
||||
self.litellm_check_batch_cost_jobs_polled.set(jobs_polled)
|
||||
|
|
@ -3155,6 +3202,19 @@ class PrometheusLogger(CustomLogger):
|
|||
except Exception as e:
|
||||
verbose_logger.warning("Error recording check batch cost metrics: %s", e)
|
||||
|
||||
def record_scheduled_job_run(self, run: JobRun) -> None:
|
||||
"""Publish one completed run of a scheduled background job."""
|
||||
self.litellm_scheduled_job_runs_total.labels(job_name=run.job_name, result=run.result.value).inc()
|
||||
self.litellm_scheduled_job_last_run_timestamp.labels(job_name=run.job_name).set(time.time())
|
||||
if run.duration_seconds is not None:
|
||||
self.litellm_scheduled_job_duration_seconds.labels(job_name=run.job_name).observe(run.duration_seconds)
|
||||
if run.items_processed:
|
||||
self.litellm_scheduled_job_items_processed_total.labels(job_name=run.job_name).inc(run.items_processed)
|
||||
|
||||
def record_cronjob_lock_attempt(self, cronjob_id: str, result: str) -> None:
|
||||
"""Publish the outcome of one single-owner lock attempt."""
|
||||
self.litellm_cronjob_lock_acquisitions_total.labels(cronjob_id=cronjob_id, result=result).inc()
|
||||
|
||||
def record_db_pool_sample(self, update: DBPoolMetricsUpdate) -> None:
|
||||
"""Publish one reading of this worker's Prisma connection pool."""
|
||||
sample: Final = update.sample
|
||||
|
|
|
|||
|
|
@ -125,6 +125,8 @@ class VantageLogger(FocusLogger):
|
|||
trigger_kwargs: Final = vantage_logger._build_scheduler_trigger()
|
||||
scheduler.add_job(
|
||||
vantage_logger.initialize_focus_export_job,
|
||||
id=VANTAGE_USAGE_DATA_JOB_NAME,
|
||||
replace_existing=True,
|
||||
**trigger_kwargs,
|
||||
)
|
||||
|
||||
|
|
|
|||
148
litellm/proxy/common_utils/scheduled_job_metrics.py
Normal file
148
litellm/proxy/common_utils/scheduled_job_metrics.py
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
"""Turns APScheduler's own lifecycle events into scheduled-job telemetry.
|
||||
|
||||
Background jobs were previously invisible: nothing recorded which job ran on a
|
||||
pod, when, for how long, whether it succeeded, or how much work it moved. During
|
||||
an incident that left operators inferring job activity from database load.
|
||||
|
||||
A scheduler listener covers every registered job at once, including ones added
|
||||
later, rather than each job growing its own instrumentation.
|
||||
|
||||
Job ids are the label, and every job litellm registers pins one. A caller that
|
||||
omits ``id=`` gets a fresh uuid from APScheduler instead, which as a label would
|
||||
grow without bound across pods and restarts, so those collapse into a single
|
||||
bucket rather than being trusted.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from apscheduler.events import (
|
||||
EVENT_JOB_ERROR,
|
||||
EVENT_JOB_EXECUTED,
|
||||
EVENT_JOB_MAX_INSTANCES,
|
||||
EVENT_JOB_MISSED,
|
||||
EVENT_JOB_SUBMITTED,
|
||||
)
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from apscheduler.events import JobEvent
|
||||
from apscheduler.schedulers.base import BaseScheduler
|
||||
|
||||
# APScheduler assigns `uuid4().hex` when a caller omits `id=`.
|
||||
_GENERATED_JOB_ID: Final = re.compile(r"\A[0-9a-f]{32}\Z")
|
||||
_UNNAMED_JOB: Final = "unnamed_job"
|
||||
|
||||
_LISTENER_EVENT_MASK: Final = (
|
||||
EVENT_JOB_SUBMITTED | EVENT_JOB_EXECUTED | EVENT_JOB_ERROR | EVENT_JOB_MISSED | EVENT_JOB_MAX_INSTANCES
|
||||
)
|
||||
|
||||
|
||||
class JobResult(str, Enum):
|
||||
"""Closed set of outcomes, so the ``result`` label stays bounded."""
|
||||
|
||||
SUCCESS = "success"
|
||||
ERROR = "error"
|
||||
# The trigger fired but the run was skipped entirely.
|
||||
MISSED = "missed"
|
||||
# A previous run of the same job was still going. With max_instances=1 this
|
||||
# is how a job that overruns its interval reports falling behind.
|
||||
MAX_INSTANCES = "max_instances"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class JobRun:
|
||||
"""One completed run of a scheduled job."""
|
||||
|
||||
job_name: str
|
||||
result: JobResult
|
||||
duration_seconds: float | None
|
||||
items_processed: int | None
|
||||
|
||||
|
||||
def _label_for(job_id: str) -> str:
|
||||
"""The job id, unless APScheduler generated it."""
|
||||
return _UNNAMED_JOB if _GENERATED_JOB_ID.match(job_id) else job_id
|
||||
|
||||
|
||||
def _items_processed(retval: object) -> int | None:
|
||||
"""The item count a job reported, if it reported one.
|
||||
|
||||
Jobs signal how much work they moved by returning a count. A job that
|
||||
returns anything else, which is most of them, simply has no count to
|
||||
publish. ``bool`` is excluded because it is an ``int`` subclass and a job
|
||||
returning ``True`` means success, not one item.
|
||||
"""
|
||||
if isinstance(retval, bool) or not isinstance(retval, int):
|
||||
return None
|
||||
return retval
|
||||
|
||||
|
||||
class ScheduledJobMetricsListener:
|
||||
"""Pairs APScheduler's submit and completion events into job runs.
|
||||
|
||||
Duration is measured across those two events because APScheduler does not
|
||||
report it. Start times are held per job id; ``max_instances=1`` means a job
|
||||
id has at most one run in flight, so a plain mapping is sufficient.
|
||||
"""
|
||||
|
||||
def __init__(self, *, monotonic: Final = time.monotonic) -> None:
|
||||
self._monotonic: Final = monotonic
|
||||
self._started_at: dict[str, float] = {} # mutable-ok: start times arrive one scheduler event at a time
|
||||
|
||||
def register(self, scheduler: BaseScheduler) -> None:
|
||||
scheduler.add_listener(self.handle, _LISTENER_EVENT_MASK)
|
||||
|
||||
def handle(self, event: JobEvent) -> None:
|
||||
"""Never raises. A listener that throws is swallowed by APScheduler and
|
||||
would leave the scheduler running with no telemetry and no explanation."""
|
||||
try:
|
||||
run: Final = self._to_run(event)
|
||||
if run is not None:
|
||||
self._publish(run)
|
||||
except Exception as e: # noqa: BLE001 # telemetry must not disturb the scheduler
|
||||
verbose_proxy_logger.debug("scheduled job metrics listener failed: %s", e)
|
||||
|
||||
def _to_run(self, event: JobEvent) -> JobRun | None:
|
||||
job_name: Final = _label_for(event.job_id)
|
||||
if event.code == EVENT_JOB_SUBMITTED:
|
||||
self._started_at[event.job_id] = self._monotonic()
|
||||
return None
|
||||
|
||||
# Neither of these follows a submission of its own. MAX_INSTANCES in
|
||||
# particular arrives while the previous run is still going, so popping
|
||||
# before this point would steal that run's start time and drop the
|
||||
# duration of exactly the overruns worth measuring.
|
||||
if event.code == EVENT_JOB_MISSED:
|
||||
return JobRun(job_name, JobResult.MISSED, None, None)
|
||||
if event.code == EVENT_JOB_MAX_INSTANCES:
|
||||
return JobRun(job_name, JobResult.MAX_INSTANCES, None, None)
|
||||
|
||||
started_at: Final = self._started_at.pop(event.job_id, None)
|
||||
duration: Final = None if started_at is None else self._monotonic() - started_at
|
||||
|
||||
if event.code == EVENT_JOB_ERROR:
|
||||
return JobRun(job_name, JobResult.ERROR, duration, None)
|
||||
return JobRun(job_name, JobResult.SUCCESS, duration, _items_processed(getattr(event, "retval", None)))
|
||||
|
||||
@staticmethod
|
||||
def _publish(run: JobRun) -> None:
|
||||
verbose_proxy_logger.info(
|
||||
"scheduled_job_completed job=%s result=%s duration_seconds=%s items_processed=%s",
|
||||
run.job_name,
|
||||
run.result.value,
|
||||
"unknown" if run.duration_seconds is None else f"{run.duration_seconds:.3f}",
|
||||
"unknown" if run.items_processed is None else run.items_processed,
|
||||
)
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
logger: Final = PrometheusLogger.get_instance()
|
||||
if logger is not None:
|
||||
logger.record_scheduled_job_run(run)
|
||||
|
|
@ -15,6 +15,24 @@ else:
|
|||
ProxyLogging = Any
|
||||
|
||||
|
||||
def _record_lock_attempt(cronjob_id: str, acquired: bool | None) -> None:
|
||||
"""Publish the outcome of one single-owner lock attempt.
|
||||
|
||||
``None`` means no Redis is configured, which is a different operational
|
||||
state from losing the race, so it gets its own result rather than being
|
||||
folded into a failure.
|
||||
"""
|
||||
result: Final = "no_redis" if acquired is None else ("acquired" if acquired else "not_acquired")
|
||||
try:
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
logger = PrometheusLogger.get_instance()
|
||||
if logger is not None:
|
||||
logger.record_cronjob_lock_attempt(cronjob_id, result)
|
||||
except Exception as e: # noqa: BLE001 # telemetry must not stop a job from running
|
||||
verbose_proxy_logger.debug("cronjob lock metric failed: %s", e)
|
||||
|
||||
|
||||
class PodLockManager:
|
||||
"""
|
||||
Manager for acquiring and releasing locks for cron jobs using Redis.
|
||||
|
|
@ -44,6 +62,21 @@ end
|
|||
cronjob_id: str,
|
||||
ttl: int | None = None,
|
||||
allow_reentrant: bool = True,
|
||||
) -> bool | None:
|
||||
"""Attempt the lock and record the outcome, then hand back the raw result.
|
||||
|
||||
Wraps the attempt rather than instrumenting each of its exits, so the
|
||||
three-state contract below reaches callers untouched.
|
||||
"""
|
||||
acquired: Final = await self._attempt_acquire_lock(cronjob_id, ttl=ttl, allow_reentrant=allow_reentrant)
|
||||
_record_lock_attempt(cronjob_id, acquired)
|
||||
return acquired
|
||||
|
||||
async def _attempt_acquire_lock(
|
||||
self,
|
||||
cronjob_id: str,
|
||||
ttl: int | None = None,
|
||||
allow_reentrant: bool = True,
|
||||
) -> bool | None:
|
||||
"""
|
||||
Attempt to acquire the lock for a specific cron job using Redis.
|
||||
|
|
|
|||
|
|
@ -348,6 +348,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import (
|
|||
)
|
||||
from litellm.proxy.common_utils.proxy_state import ProxyState
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.proxy.common_utils.scheduled_job_metrics import ScheduledJobMetricsListener
|
||||
from litellm.proxy.common_utils.scheduled_job_stagger import (
|
||||
apply_scheduled_job_stagger,
|
||||
attach_job_timing_logger,
|
||||
|
|
@ -9134,6 +9135,8 @@ class ProxyStartupEvent:
|
|||
scheduler=scheduler,
|
||||
settings=parse_stagger_settings(general_settings),
|
||||
)
|
||||
# Registered before start so the first run of every job is observed.
|
||||
ScheduledJobMetricsListener().register(scheduler)
|
||||
|
||||
# Start the scheduler immediately without processing backlogs
|
||||
scheduler.start(paused=False)
|
||||
|
|
|
|||
|
|
@ -6068,9 +6068,11 @@ async def update_spend(
|
|||
prisma_client: PrismaClient,
|
||||
db_writer_client: AsyncHTTPHandler | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
) -> int:
|
||||
"""
|
||||
Batch write updates to db.
|
||||
Batch write updates to db. Returns how many queued spend transactions this
|
||||
cycle actually drained, which the scheduled-job listener publishes as the
|
||||
run's item count.
|
||||
|
||||
Triggered every minute.
|
||||
|
||||
|
|
@ -6100,12 +6102,19 @@ async def update_spend(
|
|||
# See update_spend_logs_job and _monitor_spend_logs_queue for the new behavior.
|
||||
# Safe to keep: under high concurrency this can take up to ~30s to run,
|
||||
# so it's unlikely to overlap with monitor_spend_logs_queue.
|
||||
if queue_size > 0:
|
||||
await update_spend_logs_job(
|
||||
prisma_client=prisma_client,
|
||||
db_writer_client=db_writer_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if queue_size == 0:
|
||||
return 0
|
||||
|
||||
await update_spend_logs_job(
|
||||
prisma_client=prisma_client,
|
||||
db_writer_client=db_writer_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# What actually drained, not what was pending on entry: a partial failure or
|
||||
# a queue that refilled mid-run would otherwise be reported as processed.
|
||||
remaining: Final = await _total_queued_spend_transactions(prisma_client)
|
||||
return max(0, queue_size - remaining)
|
||||
|
||||
|
||||
async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int:
|
||||
|
|
|
|||
|
|
@ -273,6 +273,12 @@ DEFINED_PROMETHEUS_METRICS = Literal[
|
|||
# MCP tool call metrics
|
||||
"litellm_mcp_tool_calls_total",
|
||||
"litellm_mcp_tool_call_spend_metric",
|
||||
# Scheduled background jobs
|
||||
"litellm_scheduled_job_runs_total",
|
||||
"litellm_scheduled_job_duration_seconds",
|
||||
"litellm_scheduled_job_last_run_timestamp",
|
||||
"litellm_scheduled_job_items_processed_total",
|
||||
"litellm_cronjob_lock_acquisitions_total",
|
||||
# Database connection pool saturation
|
||||
"litellm_db_pool_connections_max",
|
||||
"litellm_db_pool_connections_busy",
|
||||
|
|
@ -776,6 +782,17 @@ class PrometheusMetricLabels:
|
|||
|
||||
litellm_check_batch_cost_last_run_timestamp: list[str] = []
|
||||
|
||||
# Scheduled background jobs. Labels are closed sets fixed at startup: job ids
|
||||
# come from the scheduler registration, cronjob ids from the lock call sites,
|
||||
# and results from an enum. No pod label: pod identity is unbounded, and the
|
||||
# lock result already distinguishes the pod that owns a job from the ones
|
||||
# that skipped it.
|
||||
litellm_scheduled_job_runs_total: tuple[str, ...] = ()
|
||||
litellm_scheduled_job_duration_seconds: tuple[str, ...] = ()
|
||||
litellm_scheduled_job_last_run_timestamp: tuple[str, ...] = ()
|
||||
litellm_scheduled_job_items_processed_total: tuple[str, ...] = ()
|
||||
litellm_cronjob_lock_acquisitions_total: tuple[str, ...] = ()
|
||||
|
||||
# Unlabelled: the pool is per-worker, so key/team/user labels would only add
|
||||
# cardinality.
|
||||
litellm_db_pool_connections_max: tuple[str, ...] = ()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,256 @@
|
|||
"""Scheduled-job telemetry, exercised against a real APScheduler.
|
||||
|
||||
The listener's whole job is to interpret APScheduler's event stream, so the
|
||||
events come from a running scheduler rather than from hand-built objects. A
|
||||
fabricated event proves only that the code reads the fields it was written to
|
||||
read.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
from litellm.proxy.common_utils.scheduled_job_metrics import (
|
||||
JobResult,
|
||||
ScheduledJobMetricsListener,
|
||||
)
|
||||
|
||||
|
||||
async def _drain(recorded, *, expected: int, timeout: float = 5.0):
|
||||
"""Wait for the scheduler to deliver `expected` runs."""
|
||||
deadline = asyncio.get_running_loop().time() + timeout
|
||||
while len(recorded) < expected and asyncio.get_running_loop().time() < deadline:
|
||||
await asyncio.sleep(0.02)
|
||||
return recorded
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def recorded():
|
||||
runs = []
|
||||
logger = MagicMock()
|
||||
logger.record_scheduled_job_run = runs.append
|
||||
with patch(
|
||||
"litellm.integrations.prometheus.PrometheusLogger.get_instance",
|
||||
return_value=logger,
|
||||
):
|
||||
yield runs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_successful_job_reports_its_name_duration_and_result(recorded):
|
||||
scheduler = AsyncIOScheduler()
|
||||
ScheduledJobMetricsListener().register(scheduler)
|
||||
|
||||
async def quick_job():
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
scheduler.add_job(quick_job, "interval", seconds=60, id="quick_job", next_run_time=None)
|
||||
scheduler.start(paused=False)
|
||||
try:
|
||||
scheduler.get_job("quick_job").modify(next_run_time=__import__("datetime").datetime.now())
|
||||
await _drain(recorded, expected=1)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
assert len(recorded) == 1
|
||||
run = recorded[0]
|
||||
assert run.job_name == "quick_job"
|
||||
assert run.result is JobResult.SUCCESS
|
||||
assert run.duration_seconds is not None and run.duration_seconds >= 0.05, (
|
||||
f"duration must span the real execution, got {run.duration_seconds}"
|
||||
)
|
||||
assert run.items_processed is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failing_job_is_recorded_as_an_error_not_a_success(recorded):
|
||||
scheduler = AsyncIOScheduler()
|
||||
ScheduledJobMetricsListener().register(scheduler)
|
||||
|
||||
async def broken_job():
|
||||
raise RuntimeError("job blew up")
|
||||
|
||||
scheduler.add_job(broken_job, "interval", seconds=60, id="broken_job", next_run_time=None)
|
||||
scheduler.start(paused=False)
|
||||
try:
|
||||
scheduler.get_job("broken_job").modify(next_run_time=__import__("datetime").datetime.now())
|
||||
await _drain(recorded, expected=1)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
assert len(recorded) == 1
|
||||
assert recorded[0].result is JobResult.ERROR
|
||||
assert recorded[0].job_name == "broken_job"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_that_returns_a_count_publishes_it(recorded):
|
||||
"""Item counts ride on the return value so a job opts in with one line
|
||||
instead of reaching for the metrics layer itself."""
|
||||
scheduler = AsyncIOScheduler()
|
||||
ScheduledJobMetricsListener().register(scheduler)
|
||||
|
||||
async def counting_job():
|
||||
return 42
|
||||
|
||||
scheduler.add_job(counting_job, "interval", seconds=60, id="counting_job", next_run_time=None)
|
||||
scheduler.start(paused=False)
|
||||
try:
|
||||
scheduler.get_job("counting_job").modify(next_run_time=__import__("datetime").datetime.now())
|
||||
await _drain(recorded, expected=1)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
assert recorded[0].items_processed == 42
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_returning_true_is_not_read_as_one_item(recorded):
|
||||
"""bool is an int subclass, so a job returning True would otherwise publish
|
||||
an item count of 1 that it never meant."""
|
||||
scheduler = AsyncIOScheduler()
|
||||
ScheduledJobMetricsListener().register(scheduler)
|
||||
|
||||
async def boolean_job():
|
||||
return True
|
||||
|
||||
scheduler.add_job(boolean_job, "interval", seconds=60, id="boolean_job", next_run_time=None)
|
||||
scheduler.start(paused=False)
|
||||
try:
|
||||
scheduler.get_job("boolean_job").modify(next_run_time=__import__("datetime").datetime.now())
|
||||
await _drain(recorded, expected=1)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
assert recorded[0].items_processed is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_job_that_overruns_its_interval_reports_max_instances(recorded):
|
||||
"""With max_instances=1 this is how a job falling behind its schedule
|
||||
surfaces, which is exactly the signal an operator wants during an incident."""
|
||||
import datetime
|
||||
|
||||
scheduler = AsyncIOScheduler(job_defaults={"max_instances": 1, "coalesce": False})
|
||||
ScheduledJobMetricsListener().register(scheduler)
|
||||
|
||||
async def slow_job():
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
scheduler.add_job(slow_job, "interval", seconds=1, id="slow_job", next_run_time=datetime.datetime.now())
|
||||
scheduler.start(paused=False)
|
||||
try:
|
||||
await _drain(recorded, expected=1, timeout=6.0)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
results = {run.result for run in recorded}
|
||||
assert JobResult.MAX_INSTANCES in results, f"expected a max_instances skip, saw {results}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_listener_failure_never_disturbs_the_scheduler(recorded):
|
||||
"""APScheduler swallows listener exceptions, so a throwing listener would
|
||||
leave the scheduler running with no telemetry and nothing to explain it."""
|
||||
scheduler = AsyncIOScheduler()
|
||||
listener = ScheduledJobMetricsListener()
|
||||
listener.register(scheduler)
|
||||
|
||||
ran = asyncio.Event()
|
||||
|
||||
async def job():
|
||||
ran.set()
|
||||
|
||||
with patch.object(ScheduledJobMetricsListener, "_publish", side_effect=RuntimeError("metrics down")):
|
||||
scheduler.add_job(job, "interval", seconds=60, id="job", next_run_time=__import__("datetime").datetime.now())
|
||||
scheduler.start(paused=False)
|
||||
try:
|
||||
await asyncio.wait_for(ran.wait(), timeout=5.0)
|
||||
finally:
|
||||
scheduler.shutdown(wait=False)
|
||||
|
||||
|
||||
def test_duration_is_unknown_when_the_submit_event_was_missed():
|
||||
"""A job already in flight when the listener registers has no start time.
|
||||
Reporting zero would put a false value on the duration histogram."""
|
||||
from apscheduler.events import EVENT_JOB_EXECUTED, JobExecutionEvent
|
||||
|
||||
listener = ScheduledJobMetricsListener()
|
||||
event = JobExecutionEvent(EVENT_JOB_EXECUTED, "orphan_job", "default", None, retval=None)
|
||||
|
||||
run = listener._to_run(event)
|
||||
|
||||
assert run is not None
|
||||
assert run.duration_seconds is None
|
||||
assert run.result is JobResult.SUCCESS
|
||||
|
||||
|
||||
def test_max_instances_does_not_steal_the_running_job_start_time():
|
||||
"""APScheduler emits MAX_INSTANCES *instead of* SUBMITTED, while the previous
|
||||
run is still going, so popping the start time on that event would drop the
|
||||
duration of exactly the overrunning runs this metric exists to surface."""
|
||||
from apscheduler.events import (
|
||||
EVENT_JOB_EXECUTED,
|
||||
EVENT_JOB_MAX_INSTANCES,
|
||||
EVENT_JOB_SUBMITTED,
|
||||
JobExecutionEvent,
|
||||
JobSubmissionEvent,
|
||||
)
|
||||
|
||||
# Exactly two reads: the submit, and the completion. A MAX_INSTANCES skip in
|
||||
# between must consume neither a tick nor the stored start time.
|
||||
ticks = iter([100.0, 103.5])
|
||||
listener = ScheduledJobMetricsListener(monotonic=lambda: next(ticks))
|
||||
|
||||
listener._to_run(JobSubmissionEvent(EVENT_JOB_SUBMITTED, "slow_job", "default", [None]))
|
||||
skipped = listener._to_run(JobExecutionEvent(EVENT_JOB_MAX_INSTANCES, "slow_job", "default", None))
|
||||
finished = listener._to_run(JobExecutionEvent(EVENT_JOB_EXECUTED, "slow_job", "default", None, retval=None))
|
||||
|
||||
assert skipped is not None and skipped.result is JobResult.MAX_INSTANCES
|
||||
assert finished is not None and finished.result is JobResult.SUCCESS
|
||||
assert finished.duration_seconds == pytest.approx(3.5), (
|
||||
"the overrunning run must keep its duration through a MAX_INSTANCES skip"
|
||||
)
|
||||
|
||||
|
||||
def test_an_apscheduler_generated_job_id_does_not_become_an_unbounded_label():
|
||||
"""A caller that omits id= gets uuid4().hex, which as a Prometheus label would
|
||||
grow without bound across pods and restarts."""
|
||||
from apscheduler.events import EVENT_JOB_EXECUTED, JobExecutionEvent
|
||||
|
||||
listener = ScheduledJobMetricsListener()
|
||||
|
||||
generated = listener._to_run(
|
||||
JobExecutionEvent(EVENT_JOB_EXECUTED, "e849e76a882d45ad9dc9965cd1c8a335", "default", None, retval=None)
|
||||
)
|
||||
pinned = listener._to_run(
|
||||
JobExecutionEvent(EVENT_JOB_EXECUTED, "update_spend_job", "default", None, retval=None)
|
||||
)
|
||||
|
||||
assert generated is not None and generated.job_name == "unnamed_job"
|
||||
assert pinned is not None and pinned.job_name == "update_spend_job"
|
||||
|
||||
|
||||
def test_every_job_litellm_registers_pins_an_explicit_id():
|
||||
"""The label is only bounded because each add_job passes id=. This fails if a
|
||||
new registration forgets one."""
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import litellm
|
||||
|
||||
root = Path(litellm.__file__).parent
|
||||
offenders = []
|
||||
for path in root.rglob("*.py"):
|
||||
source = path.read_text(encoding="utf-8", errors="ignore")
|
||||
for match in re.finditer(r"scheduler\.add_job\((.*?)\n\s*\)", source, re.DOTALL):
|
||||
if "id=" not in match.group(1):
|
||||
offenders.append(f"{path.relative_to(root)}: {match.group(1).strip().splitlines()[0]}")
|
||||
|
||||
assert not offenders, "scheduler.add_job without an explicit id=: " + "; ".join(offenders)
|
||||
|
|
@ -456,3 +456,51 @@ async def test_acquire_lock_own_lock_not_reentrant(pod_lock_manager, mock_redis)
|
|||
|
||||
assert await pod_lock_manager.acquire_lock(cronjob_id="test_job", allow_reentrant=False) is False
|
||||
assert await pod_lock_manager.acquire_lock(cronjob_id="test_job") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"redis_cache, granted, expected_result",
|
||||
[
|
||||
(None, None, "no_redis"),
|
||||
("present", True, "acquired"),
|
||||
("present", False, "not_acquired"),
|
||||
],
|
||||
)
|
||||
async def test_every_lock_outcome_is_recorded_under_its_own_result(redis_cache, granted, expected_result):
|
||||
"""no_redis is a different operational state from losing the race: one means
|
||||
no pod can ever be elected, the other means another pod won."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
|
||||
cache = None
|
||||
if redis_cache == "present":
|
||||
cache = MagicMock()
|
||||
cache.async_set_cache = AsyncMock(return_value=granted)
|
||||
cache.async_get_cache = AsyncMock(return_value="some-other-pod")
|
||||
|
||||
manager = PodLockManager(redis_cache=cache)
|
||||
logger = MagicMock()
|
||||
|
||||
with patch("litellm.integrations.prometheus.PrometheusLogger.get_instance", return_value=logger):
|
||||
await manager.acquire_lock(cronjob_id="db_spend_update_job")
|
||||
|
||||
logger.record_cronjob_lock_attempt.assert_called_once_with("db_spend_update_job", expected_result)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recording_the_lock_outcome_never_blocks_the_job():
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
|
||||
cache = MagicMock()
|
||||
cache.async_set_cache = AsyncMock(return_value=True)
|
||||
manager = PodLockManager(redis_cache=cache)
|
||||
|
||||
with patch(
|
||||
"litellm.integrations.prometheus.PrometheusLogger.get_instance",
|
||||
side_effect=RuntimeError("metrics down"),
|
||||
):
|
||||
assert await manager.acquire_lock(cronjob_id="db_spend_update_job") is True
|
||||
|
|
|
|||
|
|
@ -581,3 +581,53 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None:
|
|||
|
||||
with pytest.raises(ValueError, match="specific"):
|
||||
asyncio.run(_runner())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_reports_what_it_drained_not_what_was_queued():
|
||||
"""The scheduled-job listener publishes this as an items-processed count, so
|
||||
a queue that only partially drains must not be reported as fully processed."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy.utils import update_spend
|
||||
|
||||
prisma_client = MagicMock()
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.db_update_spend_transaction_handler = AsyncMock()
|
||||
|
||||
# 10 queued on entry, 4 still queued after the drain
|
||||
with (
|
||||
patch("litellm.proxy.utils._total_queued_spend_transactions", AsyncMock(side_effect=[10, 4])),
|
||||
patch("litellm.proxy.utils.update_spend_logs_job", AsyncMock()),
|
||||
):
|
||||
drained = await update_spend(
|
||||
prisma_client=prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert drained == 6, f"expected the drained count, got {drained}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_reports_zero_when_nothing_was_queued():
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy.utils import update_spend
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.db_spend_update_writer.db_update_spend_transaction_handler = AsyncMock()
|
||||
logs_job = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.utils._total_queued_spend_transactions", AsyncMock(return_value=0)),
|
||||
patch("litellm.proxy.utils.update_spend_logs_job", logs_job),
|
||||
):
|
||||
drained = await update_spend(
|
||||
prisma_client=MagicMock(),
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert drained == 0
|
||||
logs_job.assert_not_awaited()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue