diff --git a/litellm/integrations/cloudzero/cloudzero.py b/litellm/integrations/cloudzero/cloudzero.py index f41bd885442..85f00ce92d3 100644 --- a/litellm/integrations/cloudzero/cloudzero.py +++ b/litellm/integrations/cloudzero/cloudzero.py @@ -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, ) diff --git a/litellm/integrations/focus/focus_logger.py b/litellm/integrations/focus/focus_logger.py index 74ef6f70a65..4d4764e6d98 100644 --- a/litellm/integrations/focus/focus_logger.py +++ b/litellm/integrations/focus/focus_logger.py @@ -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, ) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 965d2792231..c7c9d758c76 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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 diff --git a/litellm/integrations/vantage/vantage_logger.py b/litellm/integrations/vantage/vantage_logger.py index c219ba392ab..2b77baacece 100644 --- a/litellm/integrations/vantage/vantage_logger.py +++ b/litellm/integrations/vantage/vantage_logger.py @@ -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, ) diff --git a/litellm/proxy/common_utils/scheduled_job_metrics.py b/litellm/proxy/common_utils/scheduled_job_metrics.py new file mode 100644 index 00000000000..81081738c05 --- /dev/null +++ b/litellm/proxy/common_utils/scheduled_job_metrics.py @@ -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) diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 4be1331e955..a9cedccaa63 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5adee476f33..233456032b0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 41187af2bd8..16bf0a48c16 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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: diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 3c21f0b4332..fe58af24662 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -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, ...] = () diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_metrics.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_metrics.py new file mode 100644 index 00000000000..b6b6122254a --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_metrics.py @@ -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) diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index 7a1ab60c547..a14d52d7735 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -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 diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index 54d59e690f9..0cf4d0faa91 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -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()