| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Prometheus metrics for SimpleCPUOffloadConnector. |
| 4 | |
| 5 | The stats payload uses the same self-describing wire format as the |
| 6 | offloading connector, so it survives IPC serialization without the |
| 7 | metric classes on the receiving side. |
| 8 | """ |
| 9 | |
| 10 | from dataclasses import dataclass |
| 11 | from typing import Any |
| 12 | |
| 13 | from vllm.config import VllmConfig |
| 14 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( |
| 15 | KVConnectorPromMetrics, |
| 16 | KVConnectorStats, |
| 17 | PromMetric, |
| 18 | PromMetricT, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | class MetricName: |
| 23 | """Metric names for SimpleCPUOffloadConnector.""" |
| 24 | |
| 25 | SAVE_OUTCOMES = "vllm:simple_kv_offload_save_outcomes_total" |
| 26 | LOAD_BLOCKS = "vllm:simple_kv_offload_load_blocks_total" |
| 27 | USED_BLOCKS = "vllm:simple_kv_offload_used_blocks" |
| 28 | PENDING_STORE_BLOCKS = "vllm:simple_kv_offload_pending_store_blocks" |
| 29 | INFO = "vllm:simple_kv_offload_info" |
| 30 | |
| 31 | |
| 32 | # outcome label value -> BoundaryStoreStats field name. |
| 33 | # ``published`` is the sum of all outcomes and is not labeled. |
| 34 | OUTCOME_TO_FIELD = { |
| 35 | "stored": "stored", |
| 36 | "dropped_cpu_full": "dropped_cpu_full", |
| 37 | "dropped_request_gone": "dropped_request_gone", |
| 38 | "dropped_null_block": "dropped_null_block", |
| 39 | "dropped_not_hashed": "dropped_not_hashed", |
| 40 | "skipped_already_cached": "skipped_already_cached", |
| 41 | "skipped_in_flight": "skipped_in_flight", |
| 42 | } |
| 43 | |
| 44 | INFO_LABELS = ["backend", "page_cache", "lazy_offload", "capacity_blocks"] |
| 45 | |
| 46 | |
| 47 | class _MetricType: |
| 48 | """Type tags embedded in the serialized stats payload.""" |
| 49 | |
| 50 | COUNTER = "counter" |
| 51 | GAUGE = "gauge" |
| 52 | |
| 53 | |
| 54 | class _StatsKey: |
| 55 | """Top-level keys in the serialized stats dict.""" |
| 56 | |
| 57 | # Maps metric name -> _MetricType value |
| 58 | TYPES = "types" |
| 59 | # Maps metric name -> {label values tuple -> observed value} |
| 60 | DATA = "data" |
| 61 | |
| 62 | |
| 63 | @dataclass |
| 64 | class SimpleCPUOffloadStats(KVConnectorStats): |
| 65 | """Stats payload for SimpleCPUOffloadConnector. |
| 66 | |
| 67 | The ``data`` dict is structured as:: |
| 68 | |
| 69 | { |
| 70 | _StatsKey.TYPES: {name: _MetricType.*, ...}, |
| 71 | _StatsKey.DATA: {name: {labelvalues: value, ...}, ...}, |
| 72 | } |
| 73 | |
| 74 | Counter values are summed when aggregating, gauge values keep the |
| 75 | latest snapshot, and unlabeled metrics use ``()`` as their |
| 76 | labelvalues tuple. |
| 77 | """ |
| 78 | |
| 79 | def __post_init__(self): |
| 80 | if _StatsKey.DATA not in self.data: |
| 81 | self.reset() |
| 82 | |
| 83 | def reset(self): |
| 84 | self.data: dict[str, Any] = { |
| 85 | _StatsKey.TYPES: {}, |
| 86 | _StatsKey.DATA: {}, |
| 87 | } |
| 88 | |
| 89 | @property |
| 90 | def _types(self) -> dict[str, str]: |
| 91 | return self.data[_StatsKey.TYPES] |
| 92 | |
| 93 | @property |
| 94 | def _values(self) -> dict[str, Any]: |
| 95 | return self.data[_StatsKey.DATA] |
| 96 | |
| 97 | def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats": |
| 98 | if other.is_empty(): |
| 99 | return self |
| 100 | assert isinstance(other, SimpleCPUOffloadStats) |
| 101 | for key, other_label_values in other._values.items(): |
| 102 | type_str = other._types.get(key) |
| 103 | if type_str is None: |
| 104 | raise AssertionError(f"Unknown simple KV offload stats key: {key}") |
| 105 | self._types.setdefault(key, type_str) |
| 106 | current_label_values = self._values.setdefault(key, {}) |
| 107 | for labelvalues, value in other_label_values.items(): |
| 108 | if type_str == _MetricType.COUNTER: |
| 109 | assert isinstance(value, int | float) |
| 110 | current_label_values[labelvalues] = ( |
| 111 | current_label_values.get(labelvalues, 0) + value |
| 112 | ) |
| 113 | elif type_str == _MetricType.GAUGE: |
| 114 | assert isinstance(value, int | float) |
| 115 | current_label_values[labelvalues] = value |
| 116 | else: |
| 117 | raise AssertionError( |
| 118 | f"Unknown metric type '{type_str}' for key: {key}" |
| 119 | ) |
| 120 | return self |
| 121 | |
| 122 | def reduce(self) -> dict[str, int | float]: |
| 123 | """Reduce the interval observations to a flat summary for logging.""" |
| 124 | return_dict: dict[str, int | float] = {} |
| 125 | for key, label_value_map in self._values.items(): |
| 126 | if key == MetricName.INFO: |
| 127 | # Constant config fingerprint; exclude from the log summary. |
| 128 | continue |
| 129 | type_str = self._types.get(key) |
| 130 | if type_str is None: |
| 131 | raise AssertionError(f"Unknown simple KV offload stats key: {key}") |
| 132 | for labelvalues, value in label_value_map.items(): |
| 133 | key_with_labels = f"{key}:{labelvalues}" if labelvalues else key |
| 134 | if type_str in (_MetricType.COUNTER, _MetricType.GAUGE): |
| 135 | assert isinstance(value, int | float) |
| 136 | return_dict[key_with_labels] = value |
| 137 | else: |
| 138 | raise AssertionError( |
| 139 | f"Unknown metric type '{type_str}' for key: {key}" |
| 140 | ) |
| 141 | return return_dict |
| 142 | |
| 143 | def is_empty(self) -> bool: |
| 144 | return not self.data.get(_StatsKey.DATA) |
| 145 | |
| 146 | def increase_counter( |
| 147 | self, |
| 148 | counter_name: str, |
| 149 | counter_increase_value: int | float = 1, |
| 150 | labelvalues: tuple[str, ...] = (), |
| 151 | ) -> None: |
| 152 | """Increase a counter on the stats payload.""" |
| 153 | self._types.setdefault(counter_name, _MetricType.COUNTER) |
| 154 | counter_values = self._values.setdefault(counter_name, {}) |
| 155 | counter_values[labelvalues] = ( |
| 156 | counter_values.get(labelvalues, 0) + counter_increase_value |
| 157 | ) |
| 158 | |
| 159 | def set_gauge( |
| 160 | self, |
| 161 | gauge_name: str, |
| 162 | gauge_value: int | float, |
| 163 | labelvalues: tuple[str, ...] = (), |
| 164 | ) -> None: |
| 165 | """Set a gauge snapshot on the stats payload.""" |
| 166 | self._types.setdefault(gauge_name, _MetricType.GAUGE) |
| 167 | gauge_values = self._values.setdefault(gauge_name, {}) |
| 168 | gauge_values[labelvalues] = gauge_value |
| 169 | |
| 170 | |
| 171 | class SimpleCPUOffloadPromMetrics(KVConnectorPromMetrics): |
| 172 | """Prometheus registration and recording for SimpleCPUOffloadConnector.""" |
| 173 | |
| 174 | def __init__( |
| 175 | self, |
| 176 | vllm_config: VllmConfig, |
| 177 | metric_types: dict[type[PromMetric], type[PromMetricT]], |
| 178 | labelnames: list[str], |
| 179 | per_engine_labelvalues: dict[int, list[object]], |
| 180 | ): |
| 181 | super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) |
| 182 | # (engine_idx, metric_name, labelvalues) -> bound child metric |
| 183 | self._metrics: dict[tuple[int, str, tuple[str, ...]], PromMetricT] = {} |
| 184 | |
| 185 | self._counter_save_outcomes = self._counter_cls( |
| 186 | name="vllm:simple_kv_offload_save_outcomes_total", |
| 187 | documentation=( |
| 188 | "Store admission decisions in SimpleCPUOffloadConnector, by " |
| 189 | "outcome. Outcomes classify eager-mode boundary hand-off " |
| 190 | "stores; lazy-mode stores are not classified." |
| 191 | ), |
| 192 | labelnames=self._labelnames + ["outcome"], |
| 193 | ) |
| 194 | self._counter_load_blocks = self._counter_cls( |
| 195 | name="vllm:simple_kv_offload_load_blocks_total", |
| 196 | documentation=( |
| 197 | "KV blocks restored to GPU from the offload pool, counted " |
| 198 | "when a load finishes. This is prefill work avoided " |
| 199 | "through offload cache hits." |
| 200 | ), |
| 201 | labelnames=self._labelnames, |
| 202 | ) |
| 203 | self._gauge_used_blocks = self._gauge_cls( |
| 204 | name="vllm:simple_kv_offload_used_blocks", |
| 205 | documentation=( |
| 206 | "Offload-pool blocks currently pinned by in-flight " |
| 207 | "transfers or cache hits (capacity minus free). Evictable " |
| 208 | "cached blocks are not counted; capacity_blocks is on " |
| 209 | "vllm:simple_kv_offload_info." |
| 210 | ), |
| 211 | labelnames=self._labelnames, |
| 212 | multiprocess_mode="mostrecent", |
| 213 | ) |
| 214 | self._gauge_pending_store_blocks = self._gauge_cls( |
| 215 | name="vllm:simple_kv_offload_pending_store_blocks", |
| 216 | documentation=( |
| 217 | "Offload blocks in pending stores: queued for dispatch, " |
| 218 | "issued to workers, or awaiting release after a cache " |
| 219 | "reset. A persistently growing value indicates a stuck " |
| 220 | "transfer." |
| 221 | ), |
| 222 | labelnames=self._labelnames, |
| 223 | multiprocess_mode="mostrecent", |
| 224 | ) |
| 225 | self._gauge_info = self._gauge_cls( |
| 226 | name="vllm:simple_kv_offload_info", |
| 227 | documentation=( |
| 228 | "SimpleCPUOffloadConnector deployment facts. Value is " |
| 229 | "always 1. backend is cpu or disk; page_cache and " |
| 230 | "lazy_offload are true/false; capacity_blocks is the " |
| 231 | "offload pool size." |
| 232 | ), |
| 233 | labelnames=self._labelnames + INFO_LABELS, |
| 234 | multiprocess_mode="mostrecent", |
| 235 | ) |
| 236 | |
| 237 | self._defs: dict[str, PromMetricT] = { |
| 238 | MetricName.SAVE_OUTCOMES: self._counter_save_outcomes, |
| 239 | MetricName.LOAD_BLOCKS: self._counter_load_blocks, |
| 240 | MetricName.USED_BLOCKS: self._gauge_used_blocks, |
| 241 | MetricName.PENDING_STORE_BLOCKS: self._gauge_pending_store_blocks, |
| 242 | MetricName.INFO: self._gauge_info, |
| 243 | } |
| 244 | self._num_label_values: dict[str, int] = { |
| 245 | MetricName.SAVE_OUTCOMES: 1, |
| 246 | MetricName.LOAD_BLOCKS: 0, |
| 247 | MetricName.USED_BLOCKS: 0, |
| 248 | MetricName.PENDING_STORE_BLOCKS: 0, |
| 249 | MetricName.INFO: len(INFO_LABELS), |
| 250 | } |
| 251 | |
| 252 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 253 | """Observe transfer statistics.""" |
| 254 | metric_types = transfer_stats_data.get(_StatsKey.TYPES, {}) |
| 255 | metric_data = transfer_stats_data.get(_StatsKey.DATA, {}) |
| 256 | for key, label_value_map in metric_data.items(): |
| 257 | prom_metric = self._defs.get(key) |
| 258 | if prom_metric is None: |
| 259 | raise AssertionError(f"Unknown simple KV offload stats key: {key}") |
| 260 | type_str = metric_types.get(key) |
| 261 | for labelvalues, value in label_value_map.items(): |
| 262 | if len(labelvalues) != self._num_label_values[key]: |
| 263 | raise AssertionError( |
| 264 | f"Metric {key} expects " |
| 265 | f"{self._num_label_values[key]} labels, " |
| 266 | f"got {len(labelvalues)}" |
| 267 | ) |
| 268 | child = self._metrics.get((engine_idx, key, labelvalues)) |
| 269 | if child is None: |
| 270 | engine_labelvalues = self.per_engine_labelvalues[engine_idx] |
| 271 | child = prom_metric.labels( |
| 272 | *(engine_labelvalues + list(labelvalues)) |
| 273 | ) |
| 274 | self._metrics[(engine_idx, key, labelvalues)] = child |
| 275 | if type_str == _MetricType.COUNTER: |
| 276 | assert isinstance(value, int | float) |
| 277 | child.inc(value) |
| 278 | elif type_str == _MetricType.GAUGE: |
| 279 | assert isinstance(value, int | float) |
| 280 | child.set(value) |
| 281 | else: |
| 282 | raise AssertionError( |
| 283 | f"Unknown metric type '{type_str}' for key: {key}" |
| 284 | ) |