# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """OmniPrometheusStatLogger — wrap upstream PrometheusStatLogger. Rewrites the upstream ``engine`` single-label scheme into a ``stage`` + ``replica`` two-label scheme so the ~37 ``vllm:*`` metric families gain per-(stage, replica) visibility for multi-replica deployments. Contents: - ``_ENGINE_INDEX_MAP``: process-wide engine_idx → (stage_name, replica_id) lookup, populated by ``OmniPrometheusStatLogger.__init__``. - ``_RelabelMixin``: rewrites ``labelnames`` at family creation and translates ``.labels()`` calls; applied via ``_RelabelGauge`` / ``_RelabelCounter`` / ``_RelabelHistogram``. - ``_OmniPerfMetricsProm`` / ``_OmniSpecDecodingProm`` / ``_OmniKVConnectorProm``: helper-class wraps so the upstream sub-collectors construct their internal families through the relabel mixin too. - ``OmniPrometheusStatLogger``: the subclass that wires everything together and rewrites ``per_engine_labelvalues`` from 2-tuple to 3-tuple at setter time. """ from __future__ import annotations from prometheus_client import Counter, Gauge, Histogram from prometheus_client.metrics import MetricWrapperBase from vllm.config import VllmConfig from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm, KVConnectorPromMetrics from vllm.v1.metrics.loggers import PrometheusStatLogger from vllm.v1.metrics.perf import PerfMetricsProm from vllm.v1.spec_decode.metrics import SpecDecodingProm # Process-wide translation table written by OmniPrometheusStatLogger at init. # Keys are flat engine_idx values (as upstream PrometheusStatLogger sees them); # values are the (stage_name, replica_id_str) tuple we expose as labels. # # Module-level rather than per-instance because the wrapper metric classes are # constructed by upstream's __init__ and never get a back-reference to the # StatLogger that owns them. vLLM runs a single Orchestrator/StatLogger per # process, so a module global is safe; tests isolate by .clear()ing first. _ENGINE_INDEX_MAP: dict[int, tuple[str, str]] = {} def _rewrite_labelnames(labelnames): """Replace `engine` in ``labelnames`` with (`stage`, `replica`) in place. Preserves ordering (so ``["model_name", "engine", "reason"]`` becomes ``["model_name", "stage", "replica", "reason"]``) and the original container type (list vs tuple). """ if labelnames is None: return labelnames seq = list(labelnames) if "engine" not in seq: return labelnames out: list[str] = [] for name in seq: if name == "engine": out.extend(("stage", "replica")) else: out.append(name) return type(labelnames)(out) if not isinstance(labelnames, list) else out def _engine_to_stage_replica(engine_value) -> tuple[str, str]: """Look up (stage, replica) for an engine_idx, accepting int or str input. Upstream emits engine values in two flavors: - int form, e.g. ``gauge_engine_sleep_state.labels(engine=idx, ...)`` (loggers.py:510) - str form, e.g. ``info_gauge.labels(**metrics_info)`` where ``metrics_info["engine"] = str(idx)`` (loggers.py:1055) Raises ``KeyError`` when the value is missing from the map — fail-fast is preferable to silently emitting series under a wrong (stage, replica). """ key = int(engine_value) if isinstance(engine_value, str) else engine_value return _ENGINE_INDEX_MAP[key] class _RelabelMixin(MetricWrapperBase): """Mixin: rewrite ``labelnames`` at family creation and ``.labels()`` calls. Handles all four upstream forms encountered in ``vllm.v1.metrics.loggers.PrometheusStatLogger``: 1. ``.labels(engine=idx, ...)`` kwarg with int engine (loggers.py:510) 2. ``.labels(model_name, str(idx), source)`` positional with str engine (loggers.py:646, 679) 3. ``.labels(**metrics_info)`` kwarg with str engine (loggers.py:1056) 4. Families without an ``engine`` label — passthrough (e.g. lora_info) Drops into upstream's ``_gauge_cls`` / ``_counter_cls`` / ``_histogram_cls`` class slots. """ _omni_parent: MetricWrapperBase def __init__(self, *args, **kwargs): # Remember where `engine` sat in the original labelnames so positional # `.labels()` calls can splice (stage, replica) at the right offset. labelnames = kwargs.get("labelnames") if labelnames is not None: original = list(labelnames) self._engine_label_index = original.index("engine") if "engine" in original else -1 kwargs["labelnames"] = _rewrite_labelnames(labelnames) else: self._engine_label_index = -1 super().__init__(*args, **kwargs) def labels(self, *args, **kwargs): if self._engine_label_index >= 0: if args: # Positional form. There are TWO upstream patterns: # # (a) Pre-rewritten path: create_metric_per_engine fans # `per_engine_labelvalues` (already a 3-tuple # [model_name, stage, replica] thanks to the property- # descriptor setter on OmniPrometheusStatLogger) into # `metric.labels(*values)`. len(args) matches the # rewritten label set already, so just pass through. # # (b) Legacy 2-tuple path: upstream sites like # `counter_request_success.labels(model_name, str(idx), # str(reason))` pass values shaped to the *original* # labelnames (engine still present at idx). Here # len(args) is short by 1 — splice (stage, replica) # in place of the engine value at engine_label_index. if len(args) == len(self._labelnames): child = super().labels(*args, **kwargs) child._omni_parent = self return child idx = self._engine_label_index if idx < len(args): stage, replica = _engine_to_stage_replica(args[idx]) args = (*args[:idx], stage, replica, *args[idx + 1 :]) elif "engine" in kwargs: stage, replica = _engine_to_stage_replica(kwargs.pop("engine")) kwargs["stage"] = stage kwargs["replica"] = replica child = super().labels(*args, **kwargs) child._omni_parent = self return child def _extend_metric_maps(value: object, engine_idx: int, stage: str, replica: str) -> None: """Bind new children in upstream's per-engine metric tables. Nested tables keep extra labels such as finish reason or token source. Reuse their existing family rather than rebuilding collectors, which would unregister metrics and reset counters for the running replicas. """ if isinstance(value, dict): template = next(iter(value.values()), None) if isinstance(template, _RelabelMixin) and "stage" in template._labelnames: labels = dict(zip(template._labelnames, template._labelvalues)) labels.update(stage=stage, replica=replica) value[engine_idx] = template._omni_parent.labels(**labels) elif isinstance(template, list) and template and all(isinstance(item, _RelabelMixin) for item in template): value[engine_idx] = [ item._omni_parent.labels( **(dict(zip(item._labelnames, item._labelvalues)) | {"stage": stage, "replica": replica}) ) for item in template ] else: for child in value.values(): _extend_metric_maps(child, engine_idx, stage, replica) elif isinstance(value, (list, tuple)): for child in value: _extend_metric_maps(child, engine_idx, stage, replica) elif isinstance(value, (PerfMetricsProm, SpecDecodingProm, KVConnectorProm, KVConnectorPromMetrics)): _extend_metric_maps(vars(value), engine_idx, stage, replica) class _RelabelGauge(_RelabelMixin, Gauge): pass class _RelabelCounter(_RelabelMixin, Counter): pass class _RelabelHistogram(_RelabelMixin, Histogram): pass # ---------------------------------------------------------------------------- # Helper-class wraps for the three sub-metric collectors that upstream # PrometheusStatLogger constructs in its __init__ (loggers.py:438-446): # # self.spec_decoding_prom = self._spec_decoding_cls(...) # self.kv_connector_prom = self._kv_connector_cls(...) # self.perf_metrics_prom = self._perf_metrics_cls(...) # # Each helper receives raw `labelnames` as a constructor argument and uses # its own class-level `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots # to build internal Counter/Gauge/Histogram families. The slot overrides on # OmniPrometheusStatLogger only reach families created via *its* slots, so # the helpers would otherwise still construct 2-label families and then hit # `Incorrect label count` when create_metric_per_engine feeds the rewritten # 3-element per_engine_labelvalues. Subclassing each helper and overriding # its slots routes the relabel mixin through to the helper-internal families # too. The helper kept seeing the OLD 2-element labelnames param, but that # is fine because the wrapper rewrites it at family-creation time. # ---------------------------------------------------------------------------- class _OmniPerfMetricsProm(PerfMetricsProm): _counter_cls = _RelabelCounter class _OmniSpecDecodingProm(SpecDecodingProm): _counter_cls = _RelabelCounter class _OmniKVConnectorProm(KVConnectorProm): _gauge_cls = _RelabelGauge _counter_cls = _RelabelCounter _histogram_cls = _RelabelHistogram class OmniPrometheusStatLogger(PrometheusStatLogger): """Wrap upstream PrometheusStatLogger to expose per-(stage, replica) labels. Replaces the upstream single ``engine`` label with two labels ``stage`` and ``replica`` so the ~37 ``vllm:*`` metric families gain per-replica visibility for multi-replica deployments. The orchestrator builds ``stage_replica_map`` from the static stage_pools config; flat engine_idx values map 1:1 to (stage_name, replica_id) tuples. New replicas are registered before their first statistics are recorded. Indices and counters are retained across detach/rejoin of the same replica. """ # Inject our wrapper metric classes into upstream's class-level slots so # every ~37 family is created with `engine` rewritten to `stage`+`replica`. _gauge_cls = _RelabelGauge _counter_cls = _RelabelCounter _histogram_cls = _RelabelHistogram # Inject helper-class wraps too so the perf / spec-decoding / kv-connector # sub-collectors get the same labelname rewrite and don't crash with # `Incorrect label count` when create_metric_per_engine fans out the # rewritten 3-element per_engine_labelvalues over their internal families. _perf_metrics_cls = _OmniPerfMetricsProm _spec_decoding_cls = _OmniSpecDecodingProm _kv_connector_cls = _OmniKVConnectorProm def __init__( self, vllm_config: VllmConfig, stage_replica_map: dict[int, tuple[str, str]], ) -> None: self._stage_replica_map = stage_replica_map # Populate the process-level translation table that wrapper metric # classes consult on every `.labels()` call. Cleared first so a # second OmniPrometheusStatLogger in the same process (e.g. tests, # orchestrator restart) starts from a clean slate. _ENGINE_INDEX_MAP.clear() _ENGINE_INDEX_MAP.update(stage_replica_map) self.vllm_config = vllm_config self.engine_indexes = list(stage_replica_map) self._omni_per_engine_labelvalues: dict[int, list[object]] = {} # A head can start without local replicas. Initialize collectors on # its first join so every per-engine table has a binding template. if stage_replica_map: super().__init__(vllm_config=vllm_config, engine_indexes=self.engine_indexes) def register_replica(self, engine_idx: int, stage: str, replica: str) -> None: """Register a replica without resetting existing metric series.""" pair = (stage, replica) if engine_idx in self._stage_replica_map: if self._stage_replica_map[engine_idx] != pair: raise ValueError(f"Engine index {engine_idx} already belongs to another replica") return self._stage_replica_map[engine_idx] = pair _ENGINE_INDEX_MAP[engine_idx] = pair if not self.engine_indexes: super().__init__(vllm_config=self.vllm_config, engine_indexes=[engine_idx]) return _extend_metric_maps(vars(self), engine_idx, stage, replica) self.per_engine_labelvalues[engine_idx] = [self.vllm_config.model_config.served_model_name, stage, replica] self.engine_indexes.append(engine_idx) self.gauge_engine_sleep_state["awake"][engine_idx].set(1) @property def stage_replica_map(self) -> dict[int, tuple[str, str]]: return self._stage_replica_map @property def per_engine_labelvalues(self) -> dict[int, list[object]]: return self._omni_per_engine_labelvalues @per_engine_labelvalues.setter def per_engine_labelvalues(self, value: dict[int, list[object]]) -> None: # Upstream sets {idx: [model_name, str(idx)]} (loggers.py:433); we drop # the engine str and append (stage, replica) so labelvalues match the # 3-element labelnames our wrapper classes produce. rewritten: dict[int, list[object]] = {} for idx, vals in value.items(): model_name = vals[0] stage, replica = self._stage_replica_map[idx] rewritten[idx] = [model_name, stage, replica] self._omni_per_engine_labelvalues = rewritten