diff --git a/litellm/constants.py b/litellm/constants.py index 49514fc4d0e..752a67f371d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1546,6 +1546,8 @@ AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX: Final = "core.windows.net" PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int( os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5) ) +PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE: Final = "other" +PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX: Final = "litellm_admitted_series_" CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = int(os.getenv("CLOUDZERO_EXPORT_INTERVAL_MINUTES", 60)) MCP_TOOL_NAME_PREFIX: Final = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 14546dda8c4..ac48a9c954e 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -20,7 +20,11 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._internal_context import with_service_target from litellm._logging import print_verbose, verbose_logger -from litellm.constants import PROXY_LLM_PROVIDER_FALLBACK, PROXY_REJECTED_BEFORE_ROUTING_KEY +from litellm.constants import ( + PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE, + PROXY_LLM_PROVIDER_FALLBACK, + PROXY_REJECTED_BEFORE_ROUTING_KEY, +) from litellm.exceptions import ( validate_rate_limit_category, validate_rate_limit_type, @@ -34,6 +38,9 @@ from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker i BoundedPrometheusSeriesTracker, PrometheusSeriesLimits, ) +from litellm.integrations.prometheus_helpers.shared_prometheus_series_admissions import ( + SharedPrometheusSeriesAdmissions, +) from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -102,7 +109,6 @@ _TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel) _DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0 UNRECOGNIZED_REQUESTED_MODEL_LABEL: Final = "other" -OVERFLOW_SERIES_LABEL_VALUE: Final = "other" _NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset( ( @@ -172,11 +178,10 @@ class _LabeledMetric: call so the emitted arguments match the metric's real label set. With ``limits.max_series`` set, only that many label sets get a series of their own: a counter or histogram records every later label set on one series whose labels are all ``other``, so totals stay exact, and a gauge skips it, since one shared gauge - value would mean nothing. ``evicts`` is False in multi-process mode, where the prometheus client cannot - remove a series.""" + value would mean nothing. In multi-process mode the tracker is the one the workers share, and ``remove`` does + nothing there, since the prometheus client cannot remove a series.""" __slots__ = ( - "_evicts", "_excluded_labels", "_limits", "_metric", @@ -192,9 +197,8 @@ class _LabeledMetric: metric_name: str, original_labelnames: tuple[str, ...], excluded_labels: frozenset[str], - tracker: BoundedPrometheusSeriesTracker, + tracker: BoundedPrometheusSeriesTracker | SharedPrometheusSeriesAdmissions, limits: PrometheusSeriesLimits, - evicts: bool, shares_overflow_series: bool, ) -> None: kept_label_count: Final = len(tuple(name for name in original_labelnames if name not in excluded_labels)) @@ -204,9 +208,8 @@ class _LabeledMetric: self._excluded_labels = excluded_labels self._tracker = tracker self._limits = limits - self._evicts = evicts self._overflow_child: Callable[[], MetricWrapperBase | NoOpMetric] = ( - partial(metric.labels, *(OVERFLOW_SERIES_LABEL_VALUE,) * kept_label_count) + partial(metric.labels, *(PROMETHEUS_OVERFLOW_SERIES_LABEL_VALUE,) * kept_label_count) if shares_overflow_series else NoOpMetric ) @@ -216,18 +219,33 @@ class _LabeledMetric: kept_values: Final = self._kept_values(values) if not kept_values: return self._metric - if not self._limits.enabled or self._tracker.admit_series( - metric=self._metric, metric_name=self._metric_name, label_values=kept_values, limits=self._limits - ): + if not self._limits.enabled: return self._metric.labels(*kept_values) + with self._tracker.lock: + if self._admits(kept_values): + return self._metric.labels(*kept_values) return self._overflow_child() def remove(self, *labelvalues: object) -> None: - if not self._evicts: - return - kept_values: Final = self._kept_values(labelvalues) - self._tracker.forget_series(self._metric_name, kept_values) - self._metric.remove(*kept_values) + match self._tracker: + case SharedPrometheusSeriesAdmissions(): + return + case BoundedPrometheusSeriesTracker(): + kept_values: Final = self._kept_values(labelvalues) + with self._tracker.lock: + self._tracker.forget_series(self._metric_name, kept_values) + self._metric.remove(*kept_values) + + def _admits(self, kept_values: tuple[str, ...]) -> bool: + match self._tracker: + case SharedPrometheusSeriesAdmissions(): + return self._limits.max_series is None or self._tracker.admit_series( + metric_name=self._metric_name, label_values=kept_values, max_series=self._limits.max_series + ) + case BoundedPrometheusSeriesTracker(): + return self._tracker.admit_series( + metric=self._metric, metric_name=self._metric_name, label_values=kept_values, limits=self._limits + ) def _kept_values(self, values: tuple[object, ...]) -> tuple[str, ...]: return tuple( @@ -349,9 +367,13 @@ class PrometheusLogger(CustomLogger): _custom_buckets: Final = litellm.prometheus_latency_buckets self.latency_buckets = tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS self._bounded_prometheus_series_tracker = BoundedPrometheusSeriesTracker() - self._multiprocess_mode = "PROMETHEUS_MULTIPROC_DIR" in os.environ - self._series_cap_tracker = BoundedPrometheusSeriesTracker() - self._series_limits = self._configured_series_limits(multiprocess_mode=self._multiprocess_mode) + _multiproc_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR") + self._series_cap_tracker = ( + BoundedPrometheusSeriesTracker() + if _multiproc_dir is None + else SharedPrometheusSeriesAdmissions(directory=_multiproc_dir) + ) + self._series_limits = self._configured_series_limits(multiprocess_mode=_multiproc_dir is not None) # Create metric factory functions self._counter_factory = self._create_metric_factory(Counter) @@ -745,13 +767,13 @@ class PrometheusLogger(CustomLogger): self.litellm_deployment_successful_fallbacks = self._counter_factory( "litellm_deployment_successful_fallbacks", "LLM Deployment Analytics - Number of successful fallback requests from primary model -> fallback model", - self.get_labels_for_metric("litellm_deployment_successful_fallbacks"), + labelnames=self.get_labels_for_metric("litellm_deployment_successful_fallbacks"), ) self.litellm_deployment_failed_fallbacks = self._counter_factory( "litellm_deployment_failed_fallbacks", "LLM Deployment Analytics - Number of failed fallback requests from primary model -> fallback model", - self.get_labels_for_metric("litellm_deployment_failed_fallbacks"), + labelnames=self.get_labels_for_metric("litellm_deployment_failed_fallbacks"), ) # Callback Logging Failure Metrics @@ -1255,7 +1277,6 @@ class PrometheusLogger(CustomLogger): excluded_labels=self.exclude_labels, tracker=self._series_cap_tracker, limits=self._series_limits, - evicts=not self._multiprocess_mode, shares_overflow_series=shares_overflow_series, ) diff --git a/litellm/integrations/prometheus_helpers/shared_prometheus_series_admissions.py b/litellm/integrations/prometheus_helpers/shared_prometheus_series_admissions.py new file mode 100644 index 00000000000..598dd0cef40 --- /dev/null +++ b/litellm/integrations/prometheus_helpers/shared_prometheus_series_admissions.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import os +from threading import RLock +from typing import Final + +from pydantic import TypeAdapter + +from litellm.constants import PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX + +_LABEL_VALUES: Final = TypeAdapter(tuple[str, ...]) + + +class _MetricAdmissions: + __slots__ = ("_label_sets", "_max_series", "_path", "_read_offset") + + def __init__(self, path: str, max_series: int) -> None: + self._path = path + self._max_series = max_series + self._label_sets: set[tuple[str, ...]] = ( # mutable-ok: a frozenset copy per admission is quadratic in the cap + set() + ) + self._read_offset = 0 + + def admit(self, label_values: tuple[str, ...]) -> bool: + if label_values in self._label_sets: + return True + if self._is_full(): + return False + self._read_new_admissions() + if label_values not in self._label_sets and not self._is_full(): + self._append(label_values) + self._read_new_admissions() + return label_values in self._label_sets + + def _is_full(self) -> bool: + return len(self._label_sets) >= self._max_series + + def _append(self, label_values: tuple[str, ...]) -> None: + descriptor: Final = os.open(self._path, os.O_WRONLY | os.O_APPEND | os.O_CREAT, 0o600) + try: + os.write(descriptor, _LABEL_VALUES.dump_json(label_values) + b"\n") + finally: + os.close(descriptor) + + def _read_new_admissions(self) -> None: + try: + with open(self._path, "rb") as admissions_file: + admissions_file.seek(self._read_offset) + unread: Final = admissions_file.read() + except FileNotFoundError: + return + complete_lines, newline, _ = unread.rpartition(b"\n") + if not newline: + return + self._read_offset += len(complete_lines) + len(newline) + for line in complete_lines.split(b"\n"): + if self._is_full(): + return + self._label_sets.add(_LABEL_VALUES.validate_json(line)) + + +class SharedPrometheusSeriesAdmissions: + """Picks which label sets get a series when several worker processes write to one + ``PROMETHEUS_MULTIPROC_DIR``. Each metric has one append-only file there, and its first ``max_series`` + distinct lines are the admitted label sets. Every worker reads the same lines in the same order, so all of + them, including a worker that replaces an exited one, admit the same label sets and a scrape that merges + the workers stays at the cap.""" + + def __init__(self, directory: str) -> None: + self._directory = directory + self._admissions: dict[str, _MetricAdmissions] = {} # mutable-ok: one entry per metric, added on first use + self.lock = RLock() + + def admit_series(self, metric_name: str, label_values: tuple[str, ...], max_series: int) -> bool: + with self.lock: + if metric_name not in self._admissions: + self._admissions[metric_name] = _MetricAdmissions( + path=os.path.join(self._directory, f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}{metric_name}"), + max_series=max_series, + ) + return self._admissions[metric_name].admit(label_values) diff --git a/litellm/proxy/prometheus_cleanup.py b/litellm/proxy/prometheus_cleanup.py index c65aaeedfaa..391fff1c964 100644 --- a/litellm/proxy/prometheus_cleanup.py +++ b/litellm/proxy/prometheus_cleanup.py @@ -1,7 +1,7 @@ """ Prometheus multiprocess directory cleanup utilities. -Wipes all .db files on startup so workers start with a clean slate. +Wipes all .db files and admitted-series files on startup so workers start with a clean slate. """ from __future__ import annotations @@ -12,13 +12,17 @@ import re from typing import Final from litellm._logging import verbose_proxy_logger +from litellm.constants import PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX _LIVE_GAUGE_PID: Final = re.compile(r"gauge_live[a-z]*_(\d+)\.db$") def wipe_directory(directory: str) -> None: - """Delete all .db files in the directory. Called once before workers fork.""" - files: Final = glob.glob(os.path.join(directory, "*.db")) + """Delete all .db files and admitted-series files in the directory. Called once before workers fork.""" + files: Final = ( + *glob.glob(os.path.join(directory, "*.db")), + *glob.glob(os.path.join(directory, f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}*")), + ) deleted = 0 for filepath in files: try: @@ -27,7 +31,7 @@ def wipe_directory(directory: str) -> None: except OSError as e: verbose_proxy_logger.warning("Failed to delete stale prometheus file %s: %s", filepath, e) if deleted: - verbose_proxy_logger.info("Prometheus cleanup: wiped %s stale .db files from %s", deleted, directory) + verbose_proxy_logger.info("Prometheus cleanup: wiped %s stale files from %s", deleted, directory) def mark_worker_exit(worker_pid: int) -> None: diff --git a/tests/unit/integrations/test_prometheus_series_cardinality.py b/tests/unit/integrations/test_prometheus_series_cardinality.py index 7693eb1a81d..9894d760788 100644 --- a/tests/unit/integrations/test_prometheus_series_cardinality.py +++ b/tests/unit/integrations/test_prometheus_series_cardinality.py @@ -1,13 +1,24 @@ import re +from pathlib import Path +from threading import Thread from time import monotonic from typing import Final import pytest -from prometheus_client import REGISTRY, generate_latest +from prometheus_client import REGISTRY, CollectorRegistry, Counter, generate_latest import litellm -from litellm.integrations.prometheus import PrometheusLogger, prometheus_label_factory +from litellm.constants import PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX +from litellm.integrations.prometheus import PrometheusLogger, _LabeledMetric, prometheus_label_factory from litellm.integrations.prometheus_helpers import bounded_prometheus_series_tracker +from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker import ( + BoundedPrometheusSeriesTracker, + PrometheusSeriesLimits, +) +from litellm.integrations.prometheus_helpers.shared_prometheus_series_admissions import ( + SharedPrometheusSeriesAdmissions, +) +from litellm.proxy.prometheus_cleanup import wipe_directory from litellm.types.integrations.prometheus import UserAPIKeyLabelValues SERIES_SETTINGS: Final = ( @@ -32,12 +43,15 @@ def _unregister_everything() -> None: @pytest.fixture(autouse=True) def isolated_registry_and_settings(monkeypatch): + collectors_before: Final = tuple(REGISTRY._collector_to_names) _unregister_everything() monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False) for setting in SERIES_SETTINGS: monkeypatch.setattr(litellm, setting, getattr(litellm, setting)) yield _unregister_everything() + for collector in collectors_before: + REGISTRY.register(collector) @pytest.fixture @@ -47,8 +61,8 @@ def clock(monkeypatch): return now -def _scraped_series(sample_name: str) -> frozenset[str]: - exposition: Final = generate_latest(REGISTRY).decode() +def _scraped_series(sample_name: str, registry: CollectorRegistry = REGISTRY) -> frozenset[str]: + exposition: Final = generate_latest(registry).decode() return frozenset(line for line in exposition.splitlines() if line.startswith(f"{sample_name}{{")) @@ -142,6 +156,114 @@ def test_cap_holds_and_ttl_is_ignored_in_multiprocess_mode(monkeypatch, tmp_path assert _label_values(series, "user_agent") == {"first-agent", "second-agent", "other"} +def test_workers_sharing_a_multiprocess_dir_admit_the_same_label_sets(tmp_path: Path): + first_worker: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + second_worker: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + + assert first_worker.admit_series("litellm_requests_metric", ("user-a",), max_series=2) + assert second_worker.admit_series("litellm_requests_metric", ("user-b",), max_series=2) + + replacement_worker: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + for worker in (first_worker, second_worker, replacement_worker): + assert worker.admit_series("litellm_requests_metric", ("user-a",), max_series=2) + assert worker.admit_series("litellm_requests_metric", ("user-b",), max_series=2) + assert not worker.admit_series("litellm_requests_metric", ("user-c",), max_series=2) + assert first_worker.admit_series("litellm_spend_metric", ("user-c",), max_series=2) + + +def test_workers_agree_when_racing_appends_overfill_the_admissions_file(tmp_path: Path): + racing_workers: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + for user in ("user-a", "user-b", "user-c"): + assert racing_workers.admit_series("litellm_requests_metric", (user,), max_series=3) + + worker: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + + assert not worker.admit_series("litellm_requests_metric", ("user-c",), max_series=2) + assert worker.admit_series("litellm_requests_metric", ("user-a",), max_series=2) + assert worker.admit_series("litellm_requests_metric", ("user-b",), max_series=2) + + +def test_a_line_another_worker_is_still_writing_is_read_once_it_is_complete(tmp_path: Path): + admissions_file: Final = tmp_path / f"{PROMETHEUS_ADMITTED_SERIES_FILE_PREFIX}litellm_requests_metric" + assert SharedPrometheusSeriesAdmissions(directory=str(tmp_path)).admit_series( + "litellm_requests_metric", ("user-a",), max_series=2 + ) + reader: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + + with admissions_file.open("ab") as write_in_progress: + write_in_progress.write(b'["user') + write_in_progress.flush() + assert reader.admit_series("litellm_requests_metric", ("user-a",), max_series=2) + write_in_progress.write(b'-b"]\n') + + assert not reader.admit_series("litellm_requests_metric", ("user-c",), max_series=2) + assert reader.admit_series("litellm_requests_metric", ("user-b",), max_series=2) + + +def test_wiping_the_multiprocess_dir_frees_every_admitted_slot(tmp_path: Path): + before_restart: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + assert before_restart.admit_series("litellm_requests_metric", ("user-a",), max_series=1) + + wipe_directory(str(tmp_path)) + + after_restart: Final = SharedPrometheusSeriesAdmissions(directory=str(tmp_path)) + assert after_restart.admit_series("litellm_requests_metric", ("user-b",), max_series=1) + assert not after_restart.admit_series("litellm_requests_metric", ("user-a",), max_series=1) + + +def test_eviction_racing_a_new_series_cannot_leave_it_untracked(): + registry: Final = CollectorRegistry() + counter: Final = Counter("requests", "requests", labelnames=("user",), registry=registry) + + class _EvictedWhileBeingCreated: + def labels(self, *labelvalues: str): + if labelvalues == ("evicted-user",): + eviction.start() + eviction.join(timeout=0.05) + return counter.labels(*labelvalues) + + def remove(self, *labelvalues: str) -> None: + counter.remove(*labelvalues) + + labeled: Final = _LabeledMetric( + metric=_EvictedWhileBeingCreated(), + metric_name="requests", + original_labelnames=("user",), + excluded_labels=frozenset(), + tracker=BoundedPrometheusSeriesTracker(), + limits=PrometheusSeriesLimits(max_series=1, ttl_seconds=None, cleanup_interval_seconds=None), + shares_overflow_series=True, + ) + eviction: Final = Thread(target=labeled.remove, args=("evicted-user",)) + + labeled.labels("evicted-user").inc() + eviction.join() + labeled.labels("next-user").inc() + + assert _label_values(_scraped_series("requests_total", registry), "user") == {"next-user"} + + +@pytest.mark.parametrize( + "metric_name", ["litellm_deployment_successful_fallbacks", "litellm_deployment_failed_fallbacks"] +) +def test_cap_applies_to_the_fallback_counters(metric_name: str): + litellm.prometheus_metrics_max_series_per_metric = 2 + litellm.prometheus_metrics_ttl_seconds = None + logger: Final = PrometheusLogger() + + for index in range(4): + PrometheusLogger._inc_labeled_counter( + logger, + getattr(logger, metric_name), + metric_name, + UserAPIKeyLabelValues(fallback_model=f"model-{index}"), + ) + + series: Final = _scraped_series(f"{metric_name}_total") + assert _label_values(series, "fallback_model") == {"model-0", "model-1", "other"} + assert _sample_value(series, "fallback_model", "other") == 2 + + def test_end_user_eviction_keeps_the_series_and_its_slot_in_multiprocess_mode(monkeypatch, tmp_path): monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) litellm.enable_end_user_cost_tracking_prometheus_only = True