diff --git a/litellm/proxy/common_utils/scheduled_job_metrics.py b/litellm/proxy/common_utils/scheduled_job_metrics.py index 81081738c05..147d5c52f6d 100644 --- a/litellm/proxy/common_utils/scheduled_job_metrics.py +++ b/litellm/proxy/common_utils/scheduled_job_metrics.py @@ -88,8 +88,11 @@ 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. + report it. Runs are keyed by job id AND scheduled run time, because + ``max_instances`` and ``coalesce`` are both env-overridable: raising the + first puts several runs of one job in flight at once, and disabling the + second makes one submission produce a completion per missed run time. Keying + on the job id alone would let those runs consume each other's start times. """ def __init__(self, *, monotonic: Final = time.monotonic) -> None: @@ -109,10 +112,16 @@ class ScheduledJobMetricsListener: except Exception as e: # noqa: BLE001 # telemetry must not disturb the scheduler verbose_proxy_logger.debug("scheduled job metrics listener failed: %s", e) + @staticmethod + def _key(job_id: str, scheduled_run_time: object) -> str: + return f"{job_id}@{scheduled_run_time}" + 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() + now: Final = self._monotonic() + for scheduled in getattr(event, "scheduled_run_times", ()) or (None,): + self._started_at[self._key(event.job_id, scheduled)] = now return None # Neither of these follows a submission of its own. MAX_INSTANCES in @@ -124,7 +133,9 @@ class ScheduledJobMetricsListener: 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) + started_at: Final = self._started_at.pop( + self._key(event.job_id, getattr(event, "scheduled_run_time", None)), None + ) duration: Final = None if started_at is None else self._monotonic() - started_at if event.code == EVENT_JOB_ERROR: 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 a9cedccaa63..0627fa66c16 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -15,14 +15,13 @@ else: ProxyLogging = Any -def _record_lock_attempt(cronjob_id: str, acquired: bool | None) -> None: +def _record_lock_attempt(cronjob_id: str, result: str) -> 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. + Each result is a distinct operational state: no Redis means no pod can ever + be elected, an error means the attempt itself failed, and not_acquired means + another pod simply won. """ - result: Final = "no_redis" if acquired is None else ("acquired" if acquired else "not_acquired") try: from litellm.integrations.prometheus import PrometheusLogger @@ -65,11 +64,21 @@ end ) -> 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. + Owns the no-Redis and error exits so the metric can tell a failed attempt + apart from losing the election, which the inner method reports as the + same False. """ - acquired: Final = await self._attempt_acquire_lock(cronjob_id, ttl=ttl, allow_reentrant=allow_reentrant) - _record_lock_attempt(cronjob_id, acquired) + if self.redis_cache is None: + verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") + _record_lock_attempt(cronjob_id, "no_redis") + return None + try: + acquired: Final = await self._attempt_acquire_lock(cronjob_id, ttl=ttl, allow_reentrant=allow_reentrant) + except Exception as e: + verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e) + _record_lock_attempt(cronjob_id, "error") + return False + _record_lock_attempt(cronjob_id, "acquired" if acquired else "not_acquired") return acquired async def _attempt_acquire_lock( @@ -93,57 +102,52 @@ end may redo it before the TTL expires. """ if self.redis_cache is None: - verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") return None - try: - lock_ttl: Final = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS - verbose_proxy_logger.debug( - "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", + lock_ttl: Final = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + verbose_proxy_logger.debug( + "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", + self.pod_id, + cronjob_id, + lock_ttl, + ) + # Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX) + # and with an expiration (EX) to avoid deadlocks. + lock_key: Final = PodLockManager.get_redis_lock_key(cronjob_id) + acquired: Final = await self.redis_cache.async_set_cache( + lock_key, + self.pod_id, + nx=True, + ttl=lock_ttl, + ) + if acquired: + verbose_proxy_logger.info( + "Pod %s successfully acquired Redis lock for cronjob_id=%s", self.pod_id, cronjob_id, - lock_ttl, ) - # Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX) - # and with an expiration (EX) to avoid deadlocks. - lock_key: Final = PodLockManager.get_redis_lock_key(cronjob_id) - acquired: Final = await self.redis_cache.async_set_cache( - lock_key, - self.pod_id, - nx=True, - ttl=lock_ttl, - ) - if acquired: - verbose_proxy_logger.info( - "Pod %s successfully acquired Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - return True - else: - # Check if the current pod already holds the lock - current_value = await self.redis_cache.async_get_cache(lock_key) - if current_value is not None: - if isinstance(current_value, bytes): - current_value = current_value.decode("utf-8") - if current_value == self.pod_id and allow_reentrant: - verbose_proxy_logger.info( - "Pod %s already holds the Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - self._emit_acquired_lock_event(cronjob_id, self.pod_id) - return True + return True + else: + # Check if the current pod already holds the lock + current_value = await self.redis_cache.async_get_cache(lock_key) + if current_value is not None: + if isinstance(current_value, bytes): + current_value = current_value.decode("utf-8") + if current_value == self.pod_id and allow_reentrant: verbose_proxy_logger.info( - "Pod %s could not acquire lock for cronjob_id=%s, held by pod %s.", + "Pod %s already holds the Redis lock for cronjob_id=%s", self.pod_id, cronjob_id, - current_value, ) - return False - except Exception as e: - verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e) - return False + self._emit_acquired_lock_event(cronjob_id, self.pod_id) + return True + verbose_proxy_logger.info( + "Pod %s could not acquire lock for cronjob_id=%s, held by pod %s.", + self.pod_id, + cronjob_id, + current_value, + ) + return False async def release_lock( self, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 16bf0a48c16..41187af2bd8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6068,11 +6068,9 @@ async def update_spend( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, -) -> int: +): """ - 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. + Batch write updates to db. Triggered every minute. @@ -6102,19 +6100,12 @@ 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: - 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) + 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, + ) async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int: 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 index b6b6122254a..68eaec050b9 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_metrics.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_metrics.py @@ -254,3 +254,55 @@ def test_every_job_litellm_registers_pins_an_explicit_id(): 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) + + +def test_concurrent_runs_of_one_job_keep_their_own_durations(): + """APSCHEDULER_MAX_INSTANCES is env-overridable, so two runs of the same job + can be in flight at once. Keying start times by job id alone would let the + first completion consume the second run's start.""" + import datetime + + from apscheduler.events import ( + EVENT_JOB_EXECUTED, + EVENT_JOB_SUBMITTED, + JobExecutionEvent, + JobSubmissionEvent, + ) + + first = datetime.datetime(2026, 1, 1, 0, 0, 0) + second = datetime.datetime(2026, 1, 1, 0, 0, 30) + ticks = iter([100.0, 110.0, 105.0, 140.0]) + listener = ScheduledJobMetricsListener(monotonic=lambda: next(ticks)) + + listener._to_run(JobSubmissionEvent(EVENT_JOB_SUBMITTED, "slow_job", "default", [first])) + listener._to_run(JobSubmissionEvent(EVENT_JOB_SUBMITTED, "slow_job", "default", [second])) + run_a = listener._to_run(JobExecutionEvent(EVENT_JOB_EXECUTED, "slow_job", "default", first, retval=None)) + run_b = listener._to_run(JobExecutionEvent(EVENT_JOB_EXECUTED, "slow_job", "default", second, retval=None)) + + assert run_a is not None and run_a.duration_seconds == pytest.approx(5.0) + assert run_b is not None and run_b.duration_seconds == pytest.approx(30.0) + + +def test_one_submission_of_several_run_times_pairs_each_completion(): + """With coalesce=False a single submission carries several run times and + produces a completion for each.""" + import datetime + + from apscheduler.events import ( + EVENT_JOB_EXECUTED, + EVENT_JOB_SUBMITTED, + JobExecutionEvent, + JobSubmissionEvent, + ) + + times = [datetime.datetime(2026, 1, 1, 0, 0, s) for s in (0, 1, 2)] + ticks = iter([100.0, 101.0, 102.0, 103.0]) + listener = ScheduledJobMetricsListener(monotonic=lambda: next(ticks)) + + listener._to_run(JobSubmissionEvent(EVENT_JOB_SUBMITTED, "job", "default", times)) + durations = [ + listener._to_run(JobExecutionEvent(EVENT_JOB_EXECUTED, "job", "default", t, retval=None)).duration_seconds + for t in times + ] + + assert all(d is not None for d in durations), f"every run time must pair, got {durations}" 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 a14d52d7735..031ab51688a 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 @@ -504,3 +504,22 @@ async def test_recording_the_lock_outcome_never_blocks_the_job(): side_effect=RuntimeError("metrics down"), ): assert await manager.acquire_lock(cronjob_id="db_spend_update_job") is True + + +@pytest.mark.asyncio +async def test_a_redis_failure_is_not_reported_as_losing_the_election(): + """not_acquired means another pod won. An attempt that errored is a + different operational state and must not be read as healthy contention.""" + 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(side_effect=ConnectionError("redis down")) + manager = PodLockManager(redis_cache=cache) + logger = MagicMock() + + with patch("litellm.integrations.prometheus.PrometheusLogger.get_instance", return_value=logger): + assert await manager.acquire_lock(cronjob_id="db_spend_update_job") is False + + logger.record_cronjob_lock_attempt.assert_called_once_with("db_spend_update_job", "error") 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 0cf4d0faa91..28dbe0a44a3 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 @@ -584,50 +584,26 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: @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.""" +async def test_update_spend_reports_no_item_count(): + """The spend-log queue is drained by _monitor_spend_logs_queue too, so any + count derived from queue depth here would attribute that task's work to this + job, or hide work this job did when the queue refilled. Reporting nothing is + the honest option until a job can report a count it owns.""" 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._total_queued_spend_transactions", AsyncMock(return_value=10)), 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( + result = await update_spend( prisma_client=MagicMock(), db_writer_client=None, proxy_logging_obj=proxy_logging_obj, ) - assert drained == 0 - logs_job.assert_not_awaited() + assert result is None, "a count that cannot be attributed to this job must not be published"