main
py 502 lines 19.5 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 from typing import Any
5
6 from vllm.config import VllmConfig
7 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
8 KVConnectorPromMetrics,
9 KVConnectorStats,
10 PromMetric,
11 PromMetricT,
12 )
13 from vllm.v1.kv_offload.base import (
14 OffloadingCounterMetadata,
15 OffloadingGaugeMetadata,
16 OffloadingHistogramMetadata,
17 OffloadingMetricMetadata,
18 )
19 from vllm.v1.kv_offload.factory import OffloadingSpecFactory
20
21
22 class _TransferMetricName:
23 """Flat metric names for GPU↔offload-medium transfer operations."""
24
25 LOAD_BYTES = "vllm:kv_offload_load_bytes"
26 LOAD_TIME = "vllm:kv_offload_load_time"
27 LOAD_SIZE = "vllm:kv_offload_load_size"
28 STORE_BYTES = "vllm:kv_offload_store_bytes"
29 STORE_TIME = "vllm:kv_offload_store_time"
30 STORE_SIZE = "vllm:kv_offload_store_size"
31
32
33 class _ConnectorMetricName:
34 """Connector-side metrics emitted by scheduler-side offloading code."""
35
36 LOOKUP_SYNC_DELAY = "vllm:kv_offload_lookup_sync_delay_seconds"
37 LOOKUP_ASYNC_DELAY = "vllm:kv_offload_lookup_async_delay_seconds"
38 ALLOCATION_FAILURE = "vllm:kv_offload_allocation_failure"
39
40
41 class _TransferType:
42 """Transfer direction labels for deprecated CPU offload metrics."""
43
44 LOAD = "CPU_to_GPU"
45 STORE = "GPU_to_CPU"
46 ALL = (LOAD, STORE)
47
48
49 TRANSFER_SIZE_BUCKETS = (
50 1e6,
51 5e6,
52 10e6,
53 20e6,
54 40e6,
55 60e6,
56 80e6,
57 100e6,
58 150e6,
59 200e6,
60 )
61
62
63 def get_connector_metric_definitions() -> dict[str, OffloadingMetricMetadata]:
64 return {
65 _TransferMetricName.LOAD_BYTES: OffloadingCounterMetadata(
66 documentation="Total bytes loaded from offload storage to GPU.",
67 ),
68 _TransferMetricName.LOAD_TIME: OffloadingCounterMetadata(
69 documentation="Total load time from offload storage to GPU, in seconds.",
70 ),
71 _TransferMetricName.LOAD_SIZE: OffloadingHistogramMetadata(
72 documentation="Histogram of KV offload load operation size, in bytes.",
73 buckets=TRANSFER_SIZE_BUCKETS,
74 ),
75 _TransferMetricName.STORE_BYTES: OffloadingCounterMetadata(
76 documentation="Total bytes stored from GPU to offload storage.",
77 ),
78 _TransferMetricName.STORE_TIME: OffloadingCounterMetadata(
79 documentation="Total store time from GPU to offload storage, in seconds.",
80 ),
81 _TransferMetricName.STORE_SIZE: OffloadingHistogramMetadata(
82 documentation="Histogram of KV offload store operation size, in bytes.",
83 buckets=TRANSFER_SIZE_BUCKETS,
84 ),
85 _ConnectorMetricName.LOOKUP_SYNC_DELAY: OffloadingHistogramMetadata(
86 documentation=(
87 "Histogram of the time spent in a single offload lookup call, "
88 "in seconds."
89 ),
90 buckets=(
91 0.00001,
92 0.00005,
93 0.0001,
94 0.0005,
95 0.001,
96 0.005,
97 0.01,
98 0.05,
99 0.1,
100 0.5,
101 1,
102 ),
103 ),
104 _ConnectorMetricName.LOOKUP_ASYNC_DELAY: OffloadingHistogramMetadata(
105 documentation=(
106 "Histogram of time between a request's offload lookup first "
107 "deferring and the following lookup resolving, or request "
108 "finish, in seconds."
109 ),
110 buckets=(
111 0.0001,
112 0.0005,
113 0.001,
114 0.005,
115 0.01,
116 0.05,
117 0.1,
118 0.5,
119 1,
120 5,
121 10,
122 ),
123 ),
124 _ConnectorMetricName.ALLOCATION_FAILURE: OffloadingCounterMetadata(
125 documentation=(
126 "Number of KV offload store allocation attempts that failed."
127 ),
128 ),
129 }
130
131
132 _DEPRECATED_TOTAL_BYTES = "vllm:kv_offload_total_bytes"
133 _DEPRECATED_TOTAL_TIME = "vllm:kv_offload_total_time"
134 _DEPRECATED_SIZE = "vllm:kv_offload_size"
135
136 # Deprecated legacy transfer metrics, kept during the migration to the flat
137 # metric names above. These stay in a separate definition block because they
138 # use a transfer_type label, but are emitted from the same flat stats payload
139 # for compatibility.
140 _DEPRECATED_CONNECTOR_METRIC_DEFINITIONS: dict[str, OffloadingMetricMetadata] = {
141 _DEPRECATED_TOTAL_BYTES: OffloadingCounterMetadata(
142 documentation="Number of bytes offloaded by KV connector",
143 ),
144 _DEPRECATED_TOTAL_TIME: OffloadingCounterMetadata(
145 documentation="Total time measured by all KV offloading operations",
146 ),
147 _DEPRECATED_SIZE: OffloadingHistogramMetadata(
148 documentation="Histogram of KV offload transfer size, in bytes.",
149 buckets=TRANSFER_SIZE_BUCKETS,
150 ),
151 }
152
153
154 class _MetricType:
155 """Type tags embedded in the serialized stats payload."""
156
157 COUNTER = "counter"
158 GAUGE = "gauge"
159 HISTOGRAM = "histogram"
160
161
162 class _StatsKey:
163 """Top-level keys in the serialized stats dict."""
164
165 # Maps metric name -> _MetricType value
166 TYPES = "types"
167 # Maps metric name -> {label values tuple -> observed value (number or list)}
168 DATA = "data"
169
170
171 @dataclass
172 class OffloadingConnectorStats(KVConnectorStats):
173 """Offloading connector stats use flat metric names as keys.
174
175 The ``data`` dict is structured using ``_StatsKey`` / ``_MetricType``::
176
177 {
178 _StatsKey.TYPES: {name: _MetricType.*, ...},
179 _StatsKey.DATA: {name: {labelvalues: value, ...}, ...},
180 }
181
182 This structure is self-describing: it survives IPC serialization
183 without needing the full ``OffloadingMetricMetadata`` objects on the
184 receiving side.
185
186 Counter values are aggregated by summing per-label-tuple, gauge values
187 use the latest snapshot per-label-tuple, and histogram values are lists of
188 observed samples per-label-tuple. Unlabeled metrics use ``()`` as their
189 labelvalues tuple.
190 """
191
192 def __post_init__(self):
193 if _StatsKey.DATA not in self.data:
194 self.reset()
195
196 def reset(self):
197 self.data: dict[str, Any] = {
198 _StatsKey.TYPES: {},
199 _StatsKey.DATA: {},
200 }
201
202 @property
203 def _types(self) -> dict[str, str]:
204 return self.data[_StatsKey.TYPES]
205
206 @property
207 def _values(self) -> dict[str, Any]:
208 return self.data[_StatsKey.DATA]
209
210 def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
211 if other.is_empty():
212 return self
213 assert isinstance(other, OffloadingConnectorStats)
214 other_types = other._types
215 other_values = other._values
216 for key, other_label_values in other_values.items():
217 type_str = other_types.get(key)
218 if type_str is None:
219 raise AssertionError(f"Unknown offloading stats key: {key}")
220 self._types.setdefault(key, type_str)
221 current_label_values = self._values.setdefault(key, {})
222 for labelvalues, value in other_label_values.items():
223 if type_str == _MetricType.HISTOGRAM:
224 assert isinstance(value, list)
225 if labelvalues not in current_label_values:
226 current_label_values[labelvalues] = list(value)
227 else:
228 assert isinstance(current_label_values[labelvalues], list)
229 current_label_values[labelvalues].extend(value)
230 elif type_str == _MetricType.COUNTER:
231 assert isinstance(value, int | float)
232 current_label_values[labelvalues] = (
233 current_label_values.get(labelvalues, 0) + value
234 )
235 elif type_str == _MetricType.GAUGE:
236 assert isinstance(value, int | float)
237 current_label_values[labelvalues] = value
238 else:
239 raise AssertionError(
240 f"Unknown metric type '{type_str}' for key: {key}"
241 )
242 return self
243
244 def reduce(self) -> dict[str, int | float]:
245 """Reduce the observations collected during a time interval to one or
246 more representative values (eg avg/median/sum of the series).
247 This is meant to be called by the logger to produce a summary of the
248 stats for the last time interval.
249 """
250 return_dict: dict[str, int | float] = {}
251 for key, label_value_map in self._values.items():
252 type_str = self._types.get(key)
253 if type_str is None:
254 raise AssertionError(f"Unknown offloading stats key: {key}")
255 for labelvalues, value in label_value_map.items():
256 key_with_labels = f"{key}:{labelvalues}" if labelvalues else key
257 if type_str == _MetricType.HISTOGRAM:
258 assert isinstance(value, list)
259 return_dict[f"{key_with_labels}_count"] = len(value)
260 return_dict[f"{key_with_labels}_sum"] = sum(value)
261 elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
262 assert isinstance(value, int | float)
263 return_dict[key_with_labels] = value
264 else:
265 raise AssertionError(
266 f"Unknown metric type '{type_str}' for key: {key}"
267 )
268 return return_dict
269
270 def is_empty(self) -> bool:
271 return not self.data.get(_StatsKey.DATA)
272
273 def increase_counter(
274 self,
275 counter_name: str,
276 counter_increase_value: int | float = 1,
277 labelvalues: tuple[str, ...] = (),
278 ) -> None:
279 """Increase a counter on the stats payload."""
280 self._types.setdefault(counter_name, _MetricType.COUNTER)
281 counter_values = self._values.setdefault(counter_name, {})
282 counter_values[labelvalues] = (
283 counter_values.get(labelvalues, 0) + counter_increase_value
284 )
285
286 def set_gauge(
287 self,
288 gauge_name: str,
289 gauge_value: int | float,
290 labelvalues: tuple[str, ...] = (),
291 ) -> None:
292 """Set a gauge snapshot on the stats payload."""
293 self._types.setdefault(gauge_name, _MetricType.GAUGE)
294 gauge_values = self._values.setdefault(gauge_name, {})
295 gauge_values[labelvalues] = gauge_value
296
297 def observe_histogram(
298 self,
299 histogram_name: str,
300 histogram_value: int | float,
301 labelvalues: tuple[str, ...] = (),
302 ) -> None:
303 """Record a histogram observation on the stats payload."""
304 self._types.setdefault(histogram_name, _MetricType.HISTOGRAM)
305 histogram_values = self._values.setdefault(histogram_name, {})
306 histogram_values.setdefault(labelvalues, []).append(histogram_value)
307
308
309 class OffloadPromMetrics(KVConnectorPromMetrics):
310 def __init__(
311 self,
312 vllm_config: VllmConfig,
313 metric_types: dict[type[PromMetric], type[PromMetricT]],
314 labelnames: list[str],
315 per_engine_labelvalues: dict[int, list[object]],
316 ):
317 super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
318 # (engine_idx, transfer_type) -> (metric with bounded labels)
319 self.histogram_transfer_size: dict[tuple[int, str], PromMetricT] = {}
320 self.counter_kv_bytes: dict[tuple[int, str], PromMetricT] = {}
321 self.counter_kv_transfer_time: dict[tuple[int, str], PromMetricT] = {}
322 kv_transfer_config = vllm_config.kv_transfer_config
323 assert kv_transfer_config is not None
324 extra_config = kv_transfer_config.kv_connector_extra_config
325 spec_cls = OffloadingSpecFactory.get_spec_cls(extra_config)
326 self._offloading_metric_metadata: dict[str, OffloadingMetricMetadata] = {
327 **spec_cls.build_metric_definitions(extra_config),
328 **get_connector_metric_definitions(),
329 }
330 from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
331
332 self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec)
333 self._offloading_metric_defs: dict[str, PromMetricT] = {}
334 # (engine_idx, metric_name, labelvalues) -> metric with bound labels
335 self.offloading_metrics: dict[
336 tuple[int, str, tuple[str, ...]], PromMetricT
337 ] = {}
338
339 self._counter_kv_bytes = self._counter_cls(
340 name=_DEPRECATED_TOTAL_BYTES,
341 documentation=_DEPRECATED_CONNECTOR_METRIC_DEFINITIONS[
342 _DEPRECATED_TOTAL_BYTES
343 ].documentation,
344 labelnames=labelnames + ["transfer_type"],
345 )
346
347 self._counter_kv_transfer_time = self._counter_cls(
348 name=_DEPRECATED_TOTAL_TIME,
349 documentation=_DEPRECATED_CONNECTOR_METRIC_DEFINITIONS[
350 _DEPRECATED_TOTAL_TIME
351 ].documentation,
352 labelnames=labelnames + ["transfer_type"],
353 )
354
355 deprecated_size_metadata = _DEPRECATED_CONNECTOR_METRIC_DEFINITIONS[
356 _DEPRECATED_SIZE
357 ]
358 assert isinstance(deprecated_size_metadata, OffloadingHistogramMetadata)
359 self._histogram_transfer_size = self._histogram_cls(
360 name=_DEPRECATED_SIZE,
361 documentation=deprecated_size_metadata.documentation,
362 buckets=deprecated_size_metadata.buckets,
363 labelnames=labelnames + ["transfer_type"],
364 )
365
366 for engine_idx, labelvalues in per_engine_labelvalues.items():
367 for transfer_type in _TransferType.ALL:
368 bounded_labelvalues = labelvalues + [transfer_type]
369 self.histogram_transfer_size[(engine_idx, transfer_type)] = (
370 self._histogram_transfer_size.labels(*bounded_labelvalues)
371 )
372 self.counter_kv_bytes[(engine_idx, transfer_type)] = (
373 self._counter_kv_bytes.labels(*bounded_labelvalues)
374 )
375 self.counter_kv_transfer_time[(engine_idx, transfer_type)] = (
376 self._counter_kv_transfer_time.labels(*bounded_labelvalues)
377 )
378
379 for metric_name, metadata in self._offloading_metric_metadata.items():
380 self._offloading_metric_defs[metric_name] = self._create_metric(
381 metric_name, metadata
382 )
383
384 def _create_metric(
385 self, metric_name: str, metadata: OffloadingMetricMetadata
386 ) -> Any:
387 kwargs: dict[str, Any] = {
388 "name": metric_name,
389 "documentation": metadata.documentation,
390 "labelnames": self._labelnames + list(metadata.labelnames),
391 }
392 if isinstance(metadata, OffloadingCounterMetadata):
393 metric_cls = self._counter_cls
394 elif isinstance(metadata, OffloadingGaugeMetadata):
395 metric_cls = self._gauge_cls
396 elif isinstance(metadata, OffloadingHistogramMetadata):
397 metric_cls = self._histogram_cls
398 if metadata.buckets is not None:
399 kwargs["buckets"] = metadata.buckets
400 else:
401 raise AssertionError(f"Unknown offloading metric metadata: {metadata}")
402 return metric_cls(**kwargs)
403
404 def _get_prometheus_metric(
405 self,
406 metric_name: str,
407 labelvalues: tuple[str, ...],
408 engine_idx: int,
409 ) -> PromMetric:
410 metadata = self._offloading_metric_metadata[metric_name]
411 if len(labelvalues) != len(metadata.labelnames):
412 raise AssertionError(
413 f"Metric {metric_name} expects {len(metadata.labelnames)} labels, "
414 f"got {len(labelvalues)}"
415 )
416 key = (engine_idx, metric_name, labelvalues)
417 prom_metric = self.offloading_metrics.get(key)
418 if prom_metric is None:
419 engine_labelvalues = self.per_engine_labelvalues[engine_idx]
420 prom_metric = self._offloading_metric_defs[metric_name].labels(
421 *(engine_labelvalues + list(labelvalues))
422 )
423 self.offloading_metrics[key] = prom_metric
424 return prom_metric
425
426 def _increase_counter(
427 self,
428 metric_name: str,
429 value: int | float,
430 labelvalues: tuple[str, ...],
431 engine_idx: int,
432 ) -> None:
433 self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value)
434 if labelvalues or not self._observe_deprecated_metrics:
435 return
436 # Keep deprecated CPU offload transfer metrics updated during the
437 # transition to flat metric names.
438 if metric_name == _TransferMetricName.LOAD_BYTES:
439 self.counter_kv_bytes[(engine_idx, _TransferType.LOAD)].inc(value)
440 elif metric_name == _TransferMetricName.LOAD_TIME:
441 self.counter_kv_transfer_time[(engine_idx, _TransferType.LOAD)].inc(value)
442 elif metric_name == _TransferMetricName.STORE_BYTES:
443 self.counter_kv_bytes[(engine_idx, _TransferType.STORE)].inc(value)
444 elif metric_name == _TransferMetricName.STORE_TIME:
445 self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value)
446
447 def _set_gauge(
448 self,
449 metric_name: str,
450 value: int | float,
451 labelvalues: tuple[str, ...],
452 engine_idx: int,
453 ) -> None:
454 self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value)
455
456 def _observe_histogram(
457 self,
458 metric_name: str,
459 value: list[int | float],
460 labelvalues: tuple[str, ...],
461 engine_idx: int,
462 ) -> None:
463 prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx)
464 for observation in value:
465 prom_metric.observe(observation)
466 if labelvalues or not self._observe_deprecated_metrics:
467 continue
468 # Keep deprecated CPU offload transfer metrics updated during the
469 # transition to flat metric names.
470 if metric_name == _TransferMetricName.LOAD_SIZE:
471 self.histogram_transfer_size[(engine_idx, _TransferType.LOAD)].observe(
472 observation
473 )
474 elif metric_name == _TransferMetricName.STORE_SIZE:
475 self.histogram_transfer_size[(engine_idx, _TransferType.STORE)].observe(
476 observation
477 )
478
479 def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
480 """Observe transfer statistics."""
481 metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
482 metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
483 for key, label_value_map in metric_data.items():
484 type_str = metric_types.get(key)
485 if type_str is None:
486 raise AssertionError(f"Unknown offloading stats key: {key}")
487 assert key in self._offloading_metric_defs
488 for labelvalues, value in label_value_map.items():
489 if type_str == _MetricType.COUNTER:
490 assert isinstance(value, int | float)
491 self._increase_counter(key, value, labelvalues, engine_idx)
492 elif type_str == _MetricType.GAUGE:
493 assert isinstance(value, int | float)
494 self._set_gauge(key, value, labelvalues, engine_idx)
495 elif type_str == _MetricType.HISTOGRAM:
496 assert isinstance(value, list)
497 assert all(isinstance(v, int | float) for v in value)
498 self._observe_histogram(key, value, labelvalues, engine_idx)
499 else:
500 raise AssertionError(
501 f"Unknown metric type '{type_str}' for key: {key}"
502 )