diff --git a/litellm/accounting/metrics.py b/litellm/accounting/metrics.py index ade715d5d14..d4650813549 100644 --- a/litellm/accounting/metrics.py +++ b/litellm/accounting/metrics.py @@ -1,30 +1,55 @@ -from typing import Final - -from prometheus_client import Gauge, Histogram +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final from litellm.repositories.accounting_delivery_repository import DeliveryMetrics -_PENDING: Final = Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax") -_PARKED: Final = Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax") -_OLDEST: Final = Gauge( - "litellm_accounting_oldest_pending_seconds", - "Age of the oldest pending SpendLogs delivery", - multiprocess_mode="livemax", -) -_BATCH_TIME: Final = Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time") -_BATCH_COUNT: Final = Gauge( - "litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax" -) -_TABLE_BYTES: Final = Gauge( - "litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax" -) +if TYPE_CHECKING: + from prometheus_client import Gauge, Histogram + + +@dataclass(frozen=True, slots=True) +class _Instruments: + pending: "Gauge" + parked: "Gauge" + oldest: "Gauge" + batch_time: "Histogram" + batch_count: "Gauge" + table_bytes: "Gauge" + + +def _instruments() -> _Instruments | None: + try: + # prometheus-client ships only in the proxy-runtime extra, so a plain litellm[proxy] install lands here + from prometheus_client import Gauge, Histogram + except ImportError: + return None + return _Instruments( + pending=Gauge("litellm_accounting_pending", "Pending SpendLogs deliveries", multiprocess_mode="livemax"), + parked=Gauge("litellm_accounting_parked", "Parked SpendLogs deliveries", multiprocess_mode="livemax"), + oldest=Gauge( + "litellm_accounting_oldest_pending_seconds", + "Age of the oldest pending SpendLogs delivery", + multiprocess_mode="livemax", + ), + batch_time=Histogram("litellm_accounting_batch_seconds", "SpendLogs delivery worker batch time"), + batch_count=Gauge( + "litellm_accounting_batch_deliveries", "Retained batch deliveries", multiprocess_mode="livemax" + ), + table_bytes=Gauge( + "litellm_accounting_delivery_table_bytes", "Delivery table size in bytes", multiprocess_mode="livemax" + ), + ) + + +_INSTRUMENTS: Final = _instruments() def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None: - _PENDING.set(metrics.pending_count) - _PARKED.set(metrics.parked_count) - _OLDEST.set(max(0.0, metrics.oldest_pending_seconds)) - _BATCH_TIME.observe(seconds) - - _BATCH_COUNT.set(metrics.batch_count) - _TABLE_BYTES.set(metrics.table_bytes) + if _INSTRUMENTS is None: + return + _INSTRUMENTS.pending.set(metrics.pending_count) + _INSTRUMENTS.parked.set(metrics.parked_count) + _INSTRUMENTS.oldest.set(max(0.0, metrics.oldest_pending_seconds)) + _INSTRUMENTS.batch_time.observe(seconds) + _INSTRUMENTS.batch_count.set(metrics.batch_count) + _INSTRUMENTS.table_bytes.set(metrics.table_bytes) diff --git a/tests/unit/accounting/test_metrics.py b/tests/unit/accounting/test_metrics.py index 49b46eb36c2..054ceaba967 100644 --- a/tests/unit/accounting/test_metrics.py +++ b/tests/unit/accounting/test_metrics.py @@ -1,9 +1,17 @@ +import importlib.util +import sys + +import pytest from prometheus_client.metrics import MetricWrapperBase from litellm.accounting import metrics from litellm.accounting.metrics import record_deliveries from litellm.repositories.accounting_delivery_repository import DeliveryMetrics +_DELIVERY_METRICS = DeliveryMetrics( + pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192 +) + def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: family = metric.collect()[0] @@ -11,18 +19,31 @@ def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None: - batches_before = _sample(metrics._BATCH_TIME, "_count") + instruments = metrics._INSTRUMENTS + assert instruments is not None + batches_before = _sample(instruments.batch_time, "_count") - record_deliveries( - DeliveryMetrics(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192), - 0.5, - ) + record_deliveries(_DELIVERY_METRICS, 0.5) assert ( - _sample(metrics._PENDING), - _sample(metrics._PARKED), - _sample(metrics._OLDEST), - _sample(metrics._BATCH_COUNT), - _sample(metrics._TABLE_BYTES), + _sample(instruments.pending), + _sample(instruments.parked), + _sample(instruments.oldest), + _sample(instruments.batch_count), + _sample(instruments.table_bytes), ) == (4.0, 1.0, 0.0, 7.0, 8192.0) - assert _sample(metrics._BATCH_TIME, "_count") == batches_before + 1 + assert _sample(instruments.batch_time, "_count") == batches_before + 1 + + +def test_delivery_metrics_are_skipped_when_prometheus_client_is_not_installed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setitem(sys.modules, "prometheus_client", None) + spec = importlib.util.spec_from_file_location("accounting_metrics_without_prometheus", metrics.__file__) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + + spec.loader.exec_module(module) + + assert module.record_deliveries(_DELIVERY_METRICS, 0.5) is None