304 lines
14 KiB
Python
304 lines
14 KiB
Python
# 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
|