fix(observability): correct job duration pairing, lock error reporting and the item count

Three defects surfaced in review.

Durations were keyed by job id alone, but APSCHEDULER_MAX_INSTANCES and
coalesce are both env-overridable, so several runs of one job can be in flight
or one submission can carry several run times. Runs are now keyed by job id and
scheduled run time, so concurrent runs stop consuming each other's start times.

A failed lock attempt was reported as ordinary contention, because the inner
method returns False both when another pod wins and when Redis errors. The
wrapper now owns the no-Redis and error exits and emits a distinct result, with
the bool | None contract to callers unchanged.

update_spend no longer reports an item count. The spend-log queue is drained by
_monitor_spend_logs_queue as well, so a count derived from queue depth here
could credit that task's work or hide work this job did. The listener still
publishes counts for any job that can report one it owns.
This commit is contained in:
Yucheng Zhu 2026-08-12 10:18:53 -07:00
parent 50d38a7cd8
commit 6586e455c3
6 changed files with 157 additions and 104 deletions

View file

@ -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:

View file

@ -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,

View file

@ -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:

View file

@ -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}"

View file

@ -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")

View file

@ -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"