| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copy |
| 4 | import time |
| 5 | |
| 6 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm |
| 7 | from vllm.v1.metrics.loggers import PrometheusStatLogger |
| 8 | from vllm.v1.metrics.perf import PerfMetricsProm |
| 9 | from vllm.v1.spec_decode.metrics import SpecDecodingProm |
| 10 | |
| 11 | try: |
| 12 | from ray import serve as ray_serve |
| 13 | from ray.util import metrics as ray_metrics |
| 14 | from ray.util.metrics import Metric |
| 15 | except ImportError: |
| 16 | ray_metrics = None |
| 17 | ray_serve = None |
| 18 | import regex as re |
| 19 | |
| 20 | |
| 21 | def _get_replica_id() -> str | None: |
| 22 | """Get the current Ray Serve replica ID, or None if not in a Serve context.""" |
| 23 | if ray_serve is None: |
| 24 | return None |
| 25 | try: |
| 26 | return ray_serve.get_replica_context().replica_id.unique_id |
| 27 | except ray_serve.exceptions.RayServeException: |
| 28 | return None |
| 29 | |
| 30 | |
| 31 | class RayPrometheusMetric: |
| 32 | _is_labeled: bool = False |
| 33 | |
| 34 | def __init__(self): |
| 35 | if ray_metrics is None: |
| 36 | raise ImportError("RayPrometheusMetric requires Ray to be installed.") |
| 37 | self.metric: Metric = None |
| 38 | self._tags: dict[str, str] = {"ReplicaId": _get_replica_id() or ""} |
| 39 | |
| 40 | @staticmethod |
| 41 | def _get_tag_keys(labelnames: list[str] | None) -> tuple[str, ...]: |
| 42 | labels = list(labelnames) if labelnames else [] |
| 43 | labels.append("ReplicaId") |
| 44 | return tuple(labels) |
| 45 | |
| 46 | def _build_tags(self, *labels, **labelskwargs) -> dict[str, str]: |
| 47 | if labels: |
| 48 | # -1 because ReplicaId was added automatically |
| 49 | expected = len(self.metric._tag_keys) - 1 |
| 50 | if len(labels) != expected: |
| 51 | raise ValueError( |
| 52 | "Number of labels must match the number of tag keys. " |
| 53 | f"Expected {expected}, got {len(labels)}" |
| 54 | ) |
| 55 | labelskwargs.update(zip(self.metric._tag_keys, labels)) |
| 56 | |
| 57 | labelskwargs["ReplicaId"] = _get_replica_id() or "" |
| 58 | |
| 59 | return {k: v if isinstance(v, str) else str(v) for k, v in labelskwargs.items()} |
| 60 | |
| 61 | def labels(self, *labels, **labelskwargs) -> "RayPrometheusMetric": |
| 62 | if self._is_labeled: |
| 63 | raise ValueError("labels() cannot be called on an already-labeled metric.") |
| 64 | clone = copy.copy(self) |
| 65 | clone._tags = self._build_tags(*labels, **labelskwargs) |
| 66 | clone._is_labeled = True |
| 67 | return clone |
| 68 | |
| 69 | @staticmethod |
| 70 | def _get_sanitized_opentelemetry_name(name: str) -> str: |
| 71 | """For compatibility with Ray + OpenTelemetry, the metric name must be |
| 72 | sanitized. In particular, this replaces disallowed character (e.g., ':') |
| 73 | with '_' in the metric name. |
| 74 | Allowed characters: a-z, A-Z, 0-9, _ |
| 75 | |
| 76 | # ruff: noqa: E501 |
| 77 | Ref: https://github.com/open-telemetry/opentelemetry-cpp/blob/main/sdk/src/metrics/instrument_metadata_validator.cc#L22-L23 |
| 78 | Ref: https://github.com/ray-project/ray/blob/master/src/ray/stats/metric.cc#L107 |
| 79 | """ |
| 80 | return re.sub(r"[^a-zA-Z0-9_]", "_", name) |
| 81 | |
| 82 | |
| 83 | class RayGaugeWrapper(RayPrometheusMetric): |
| 84 | """Wraps around ray.util.metrics.Gauge to provide same API as |
| 85 | prometheus_client.Gauge""" |
| 86 | |
| 87 | def __init__( |
| 88 | self, |
| 89 | name: str, |
| 90 | documentation: str | None = "", |
| 91 | labelnames: list[str] | None = None, |
| 92 | multiprocess_mode: str | None = "", |
| 93 | ): |
| 94 | # All Ray metrics are keyed by WorkerId, so multiprocess modes like |
| 95 | # "mostrecent", "all", "sum" do not apply. This logic can be manually |
| 96 | # implemented at the observability layer (Prometheus/Grafana). |
| 97 | del multiprocess_mode |
| 98 | |
| 99 | super().__init__() |
| 100 | tag_keys = self._get_tag_keys(labelnames) |
| 101 | name = self._get_sanitized_opentelemetry_name(name) |
| 102 | |
| 103 | self.metric = ray_metrics.Gauge( |
| 104 | name=name, |
| 105 | description=documentation, |
| 106 | tag_keys=tag_keys, |
| 107 | ) |
| 108 | |
| 109 | def set(self, value: int | float): |
| 110 | return self.metric.set(value, tags=self._tags) |
| 111 | |
| 112 | def set_to_current_time(self): |
| 113 | # ray metrics doesn't have set_to_current time, https://docs.ray.io/en/latest/_modules/ray/util/metrics.html |
| 114 | return self.set(time.time()) |
| 115 | |
| 116 | |
| 117 | class RayCounterWrapper(RayPrometheusMetric): |
| 118 | """Wraps around ray.util.metrics.Counter to provide same API as |
| 119 | prometheus_client.Counter""" |
| 120 | |
| 121 | def __init__( |
| 122 | self, |
| 123 | name: str, |
| 124 | documentation: str | None = "", |
| 125 | labelnames: list[str] | None = None, |
| 126 | ): |
| 127 | super().__init__() |
| 128 | tag_keys = self._get_tag_keys(labelnames) |
| 129 | name = self._get_sanitized_opentelemetry_name(name) |
| 130 | self.metric = ray_metrics.Counter( |
| 131 | name=name, |
| 132 | description=documentation, |
| 133 | tag_keys=tag_keys, |
| 134 | ) |
| 135 | |
| 136 | def inc(self, value: int | float = 1.0): |
| 137 | if value == 0: |
| 138 | return |
| 139 | return self.metric.inc(value, tags=self._tags) |
| 140 | |
| 141 | |
| 142 | class RayHistogramWrapper(RayPrometheusMetric): |
| 143 | """Wraps around ray.util.metrics.Histogram to provide same API as |
| 144 | prometheus_client.Histogram""" |
| 145 | |
| 146 | def __init__( |
| 147 | self, |
| 148 | name: str, |
| 149 | documentation: str | None = "", |
| 150 | labelnames: list[str] | None = None, |
| 151 | buckets: list[float] | None = None, |
| 152 | ): |
| 153 | super().__init__() |
| 154 | tag_keys = self._get_tag_keys(labelnames) |
| 155 | name = self._get_sanitized_opentelemetry_name(name) |
| 156 | |
| 157 | boundaries = buckets if buckets else [] |
| 158 | self.metric = ray_metrics.Histogram( |
| 159 | name=name, |
| 160 | description=documentation, |
| 161 | tag_keys=tag_keys, |
| 162 | boundaries=boundaries, |
| 163 | ) |
| 164 | |
| 165 | def observe(self, value: int | float): |
| 166 | return self.metric.observe(value, tags=self._tags) |
| 167 | |
| 168 | |
| 169 | class RaySpecDecodingProm(SpecDecodingProm): |
| 170 | """RaySpecDecodingProm is used by RayMetrics to log to Ray metrics. |
| 171 | Provides the same metrics as SpecDecodingProm but uses Ray's |
| 172 | util.metrics library. |
| 173 | """ |
| 174 | |
| 175 | _counter_cls = RayCounterWrapper |
| 176 | |
| 177 | |
| 178 | class RayKVConnectorProm(KVConnectorProm): |
| 179 | """RayKVConnectorProm is used by RayMetrics to log Ray |
| 180 | metrics. Provides the same metrics as KV connectors but |
| 181 | uses Ray's util.metrics library. |
| 182 | """ |
| 183 | |
| 184 | _gauge_cls = RayGaugeWrapper |
| 185 | _counter_cls = RayCounterWrapper |
| 186 | _histogram_cls = RayHistogramWrapper |
| 187 | |
| 188 | |
| 189 | class RayPerfMetricsProm(PerfMetricsProm): |
| 190 | """RayPerfMetricsProm is used by RayMetrics to log Ray |
| 191 | metrics. Provides the same MFU metrics as PerfMetricsProm |
| 192 | uses Ray's util.metrics library. |
| 193 | """ |
| 194 | |
| 195 | _counter_cls = RayCounterWrapper |
| 196 | |
| 197 | |
| 198 | class RayPrometheusStatLogger(PrometheusStatLogger): |
| 199 | """RayPrometheusStatLogger uses Ray metrics instead.""" |
| 200 | |
| 201 | _gauge_cls = RayGaugeWrapper |
| 202 | _counter_cls = RayCounterWrapper |
| 203 | _histogram_cls = RayHistogramWrapper |
| 204 | _spec_decoding_cls = RaySpecDecodingProm |
| 205 | _kv_connector_cls = RayKVConnectorProm |
| 206 | _perf_metrics_cls = RayPerfMetricsProm |
| 207 | |
| 208 | @staticmethod |
| 209 | def _unregister_vllm_metrics(): |
| 210 | # No-op on purpose |
| 211 | pass |