mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
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:
parent
50d38a7cd8
commit
6586e455c3
6 changed files with 157 additions and 104 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue