@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c441ad1c07

[KV Offloading] Add labeled metrics support (#45957)

Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>

Srinivas Krovvidi committed Jun 21, 2026 at 13:04 UTC c441ad1c07cbfe0240a5699e4386fd0b5cc8aa82
3 files changed +384 -123
tests/v1/kv_connector/unit/offloading_connector/test_metrics.py
+255 -60
@@ -22,7 +22,7 @@ from vllm.v1.kv_offload.base import (
22 OffloadingGaugeMetadata,
23 OffloadingHistogramMetadata,
24 )
25 -from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec
25 +from vllm.v1.kv_offload.factory import OffloadingSpecFactory
26
27 LOAD_BYTES = _TransferMetricName.LOAD_BYTES
28 LOAD_TIME = _TransferMetricName.LOAD_TIME
@@ -33,6 +33,8 @@ STORE_SIZE = _TransferMetricName.STORE_SIZE
33 STORES_SKIPPED = "vllm:kv_offload_stores_skipped"
34 PENDING_STORES = "vllm:kv_offload_pending_stores"
35 LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds"
36 +MY_COUNTER = "my_counter"
37 +MY_LABEL = "my_label"
38
39
40 class _FakeMetric:
@@ -67,6 +69,20 @@ class _FakeVllmConfig:
69 )
70
71
72 +def _spec_cls_with_metric_definitions(
73 + metric_definitions: dict[str, Any],
74 +) -> type:
75 + """Build a fake offloading spec class reporting the given metric
76 + definitions, so tests don't need to patch the real CPU spec."""
77 +
78 + class _FakeOffloadingSpec:
79 + @staticmethod
80 + def build_metric_definitions(extra_config):
81 + return metric_definitions
82 +
83 + return _FakeOffloadingSpec
84 +
85 +
86 def _metric_metadata():
87 return {
88 LOAD_BYTES: OffloadingCounterMetadata(
@@ -96,9 +112,17 @@ def _metric_metadata():
112 LOOKUP_LATENCY: OffloadingHistogramMetadata(
113 documentation="lookup latency",
114 ),
115 + MY_COUNTER: OffloadingCounterMetadata(
116 + documentation="counter with a label",
117 + labelnames=(MY_LABEL,),
118 + ),
119 }
120
121
122 +def _unlabeled(values: dict[str, Any], metric_name: str) -> Any:
123 + return values[metric_name][()]
124 +
125 +
126 def test_build_kv_connector_stats_with_none():
127 """Test that build_kv_connector_stats returns empty stats when given None."""
128 stats = OffloadingConnector.build_kv_connector_stats(data=None)
@@ -131,13 +155,13 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
155 STORES_SKIPPED: _MetricType.COUNTER,
156 },
157 _StatsKey.DATA: {
134 - LOAD_BYTES: 24,
135 - LOAD_TIME: 1.5,
136 - LOAD_SIZE: [16, 8],
137 - STORE_BYTES: 3,
138 - STORE_TIME: 0.3,
139 - STORE_SIZE: [1, 2],
140 - STORES_SKIPPED: 5,
158 + LOAD_BYTES: {(): 24},
159 + LOAD_TIME: {(): 1.5},
160 + LOAD_SIZE: {(): [16, 8]},
161 + STORE_BYTES: {(): 3},
162 + STORE_TIME: {(): 0.3},
163 + STORE_SIZE: {(): [1, 2]},
164 + STORES_SKIPPED: {(): 5},
165 },
166 }
167
@@ -145,22 +169,28 @@ def test_build_kv_connector_stats_reconstructs_offload_stats():
169
170 assert isinstance(stats, OffloadingConnectorStats)
171 values = stats.data[_StatsKey.DATA]
148 - assert values[LOAD_BYTES] == 24
149 - assert values[LOAD_TIME] == 1.5
150 - assert values[LOAD_SIZE] == [16, 8]
151 - assert values[STORE_BYTES] == 3
152 - assert values[STORE_TIME] == 0.3
153 - assert values[STORE_SIZE] == [1, 2]
154 - assert values[STORES_SKIPPED] == 5
172 + assert _unlabeled(values, LOAD_BYTES) == 24
173 + assert _unlabeled(values, LOAD_TIME) == 1.5
174 + assert _unlabeled(values, LOAD_SIZE) == [16, 8]
175 + assert _unlabeled(values, STORE_BYTES) == 3
176 + assert _unlabeled(values, STORE_TIME) == 0.3
177 + assert _unlabeled(values, STORE_SIZE) == [1, 2]
178 + assert _unlabeled(values, STORES_SKIPPED) == 5
179
180
181 def _make_stats_data(
182 metric_data: dict[str, Any],
183 metric_metadata: dict[str, Any],
184 ) -> dict[str, Any]:
161 - """Build a structured data dict from flat metric data and metadata."""
185 + """Build a structured data dict from flat metric data and metadata.
186 +
187 + Values for unlabeled metrics may be passed flat (wrapped here under the
188 + empty label tuple); values for labeled metrics must already be passed as
189 + a ``{labelvalues: value}`` map.
190 + """
191 metric_types = {}
163 - for key in metric_data:
192 + data = {}
193 + for key, value in metric_data.items():
194 md = metric_metadata[key]
195 if isinstance(md, OffloadingCounterMetadata):
196 metric_types[key] = _MetricType.COUNTER
@@ -168,9 +198,10 @@ def _make_stats_data(
198 metric_types[key] = _MetricType.GAUGE
199 elif isinstance(md, OffloadingHistogramMetadata):
200 metric_types[key] = _MetricType.HISTOGRAM
201 + data[key] = value if md.labelnames else {(): value}
202 return {
203 _StatsKey.TYPES: metric_types,
173 - _StatsKey.DATA: metric_data,
204 + _StatsKey.DATA: data,
205 }
206
207
@@ -215,34 +246,106 @@ def test_aggregate_same_connector():
246
247 assert result is stats1 # Should return self
248 values = result.data[_StatsKey.DATA]
218 - assert values[LOAD_BYTES] == 34
219 - assert values[LOAD_TIME] == 2.6
220 - assert values[LOAD_SIZE] == [16, 8, 3, 7]
221 - assert values[STORE_BYTES] == 19
222 - assert values[STORE_TIME] == 2.3
223 - assert values[STORE_SIZE] == [1, 2, 16]
224 - assert values[STORES_SKIPPED] == 4
225 - assert values[PENDING_STORES] == 1
226 - assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3]
249 + assert _unlabeled(values, LOAD_BYTES) == 34
250 + assert _unlabeled(values, LOAD_TIME) == 2.6
251 + assert _unlabeled(values, LOAD_SIZE) == [16, 8, 3, 7]
252 + assert _unlabeled(values, STORE_BYTES) == 19
253 + assert _unlabeled(values, STORE_TIME) == 2.3
254 + assert _unlabeled(values, STORE_SIZE) == [1, 2, 16]
255 + assert _unlabeled(values, STORES_SKIPPED) == 4
256 + assert _unlabeled(values, PENDING_STORES) == 1
257 + assert _unlabeled(values, LOOKUP_LATENCY) == [0.1, 0.2, 0.3]
258 +
259 +
260 +def test_aggregate_labeled_metrics():
261 + metadata = _metric_metadata()
262 + stats1 = OffloadingConnectorStats(
263 + data=_make_stats_data(
264 + {
265 + MY_COUNTER: {
266 + ("a",): 10,
267 + ("b",): 3,
268 + },
269 + },
270 + metadata,
271 + ),
272 + )
273 + stats2 = OffloadingConnectorStats(
274 + data=_make_stats_data(
275 + {
276 + MY_COUNTER: {
277 + ("a",): 7,
278 + ("c",): 5,
279 + },
280 + },
281 + metadata,
282 + ),
283 + )
284 +
285 + stats1.aggregate(stats2)
286 +
287 + values = stats1.data[_StatsKey.DATA][MY_COUNTER]
288 + assert values[("a",)] == 17
289 + assert values[("b",)] == 3
290 + assert values[("c",)] == 5
291 +
292 +
293 +def test_aggregate_labeled_metric_missing_from_self():
294 + """Aggregating a labeled metric that self doesn't have at all yet."""
295 + metadata = _metric_metadata()
296 + stats1 = OffloadingConnectorStats()
297 + stats2 = OffloadingConnectorStats(
298 + data=_make_stats_data(
299 + {
300 + MY_COUNTER: {
301 + ("a",): 7,
302 + ("b",): 5,
303 + },
304 + },
305 + metadata,
306 + ),
307 + )
308 +
309 + stats1.aggregate(stats2)
310 +
311 + values = stats1.data[_StatsKey.DATA][MY_COUNTER]
312 + assert values[("a",)] == 7
313 + assert values[("b",)] == 5
314 + assert stats1.data[_StatsKey.TYPES][MY_COUNTER] == _MetricType.COUNTER
315 +
316 +
317 +def test_helper_methods_accept_labeled_metrics():
318 + stats = OffloadingConnectorStats()
319 +
320 + stats.increase_counter(MY_COUNTER, 3, ("a",))
321 + stats.increase_counter(MY_COUNTER, 4, ("a",))
322 + stats.set_gauge(PENDING_STORES, 2, ("b",))
323 + stats.observe_histogram(LOOKUP_LATENCY, 0.1, ("b",))
324 + stats.observe_histogram(LOOKUP_LATENCY, 0.2, ("b",))
325 +
326 + values = stats.data[_StatsKey.DATA]
327 + assert values[MY_COUNTER][("a",)] == 7
328 + assert values[PENDING_STORES][("b",)] == 2
329 + assert values[LOOKUP_LATENCY][("b",)] == [0.1, 0.2]
330
331
332 def test_aggregate_merges_types():
333 stats1 = OffloadingConnectorStats(
334 data={
335 _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
233 - _StatsKey.DATA: {LOAD_BYTES: 1},
336 + _StatsKey.DATA: {LOAD_BYTES: {(): 1}},
337 },
338 )
339 stats2 = OffloadingConnectorStats(
340 data={
341 _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE},
239 - _StatsKey.DATA: {PENDING_STORES: 2},
342 + _StatsKey.DATA: {PENDING_STORES: {(): 2}},
343 },
344 )
345
346 result = stats1.aggregate(stats2)
347
245 - assert result.data[_StatsKey.DATA][PENDING_STORES] == 2
348 + assert _unlabeled(result.data[_StatsKey.DATA], PENDING_STORES) == 2
349 assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE
350
351
@@ -283,6 +386,26 @@ def test_reduce():
386 assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3])
387
388
389 +def test_reduce_labeled_metrics():
390 + metadata = _metric_metadata()
391 + stats = OffloadingConnectorStats(
392 + data=_make_stats_data(
393 + {
394 + MY_COUNTER: {
395 + ("a",): 17,
396 + ("b",): 3,
397 + },
398 + },
399 + metadata,
400 + ),
401 + )
402 +
403 + reduced = stats.reduce()
404 +
405 + assert reduced[f"{MY_COUNTER}:{('a',)}"] == 17
406 + assert reduced[f"{MY_COUNTER}:{('b',)}"] == 3
407 +
408 +
409 def test_reset():
410 """Test that reset() resets all connector stats."""
411 metadata = _metric_metadata()
@@ -326,11 +449,11 @@ def test_prom_metrics_observes_manager_counter():
449 prom_metrics.observe(
450 {
451 _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER},
329 - _StatsKey.DATA: {STORES_SKIPPED: 7},
452 + _StatsKey.DATA: {STORES_SKIPPED: {(): 7}},
453 }
454 )
455
333 - counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)]
456 + counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED, ())]
457 assert counter.increments == [7]
458 counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED]
459 assert counter_def.kwargs["name"] == "vllm:kv_offload_stores_skipped"
@@ -360,22 +483,22 @@ def test_prom_metrics_observes_flat_transfer_metrics_and_legacy_metrics():
483 STORE_SIZE: _MetricType.HISTOGRAM,
484 },
485 _StatsKey.DATA: {
363 - LOAD_BYTES: 24,
364 - LOAD_TIME: 1.5,
365 - LOAD_SIZE: [16, 8],
366 - STORE_BYTES: 3,
367 - STORE_TIME: 0.3,
368 - STORE_SIZE: [1, 2],
486 + LOAD_BYTES: {(): 24},
487 + LOAD_TIME: {(): 1.5},
488 + LOAD_SIZE: {(): [16, 8]},
489 + STORE_BYTES: {(): 3},
490 + STORE_TIME: {(): 0.3},
491 + STORE_SIZE: {(): [1, 2]},
492 },
493 }
494 )
495
373 - assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24]
374 - assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5]
375 - assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8]
376 - assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3]
377 - assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3]
378 - assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2]
496 + assert prom_metrics.offloading_metrics[(0, LOAD_BYTES, ())].increments == [24]
497 + assert prom_metrics.offloading_metrics[(0, LOAD_TIME, ())].increments == [1.5]
498 + assert prom_metrics.offloading_metrics[(0, LOAD_SIZE, ())].observed == [16, 8]
499 + assert prom_metrics.offloading_metrics[(0, STORE_BYTES, ())].increments == [3]
500 + assert prom_metrics.offloading_metrics[(0, STORE_TIME, ())].increments == [0.3]
501 + assert prom_metrics.offloading_metrics[(0, STORE_SIZE, ())].observed == [1, 2]
502
503 assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24]
504 assert prom_metrics.counter_kv_transfer_time[(0, "CPU_to_GPU")].increments == [1.5]
@@ -396,7 +519,9 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
519 ),
520 }
521 with patch.object(
399 - CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions
522 + OffloadingSpecFactory,
523 + "get_spec_cls",
524 + return_value=_spec_cls_with_metric_definitions(metric_definitions),
525 ):
526 prom_metrics = OffloadPromMetrics(
527 vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
@@ -416,20 +541,91 @@ def test_prom_metrics_observes_manager_gauge_and_histogram():
541 LOOKUP_LATENCY: _MetricType.HISTOGRAM,
542 },
543 _StatsKey.DATA: {
419 - PENDING_STORES: 5,
420 - LOOKUP_LATENCY: [0.2, 0.4],
544 + PENDING_STORES: {(): 5},
545 + LOOKUP_LATENCY: {(): [0.2, 0.4]},
546 },
547 }
548 )
549
425 - gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)]
426 - histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)]
550 + gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES, ())]
551 + histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY, ())]
552 assert gauge.set_values == [5]
553 assert histogram.observed == [0.2, 0.4]
554 histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY]
555 assert histogram_def.kwargs["buckets"] == (0.1, 1.0)
556
557
558 +def test_prom_metrics_lazily_observes_labeled_metric():
559 + metric_definitions = {
560 + MY_COUNTER: OffloadingCounterMetadata(
561 + documentation="counter with a label",
562 + labelnames=(MY_LABEL,),
563 + ),
564 + }
565 + with patch.object(
566 + OffloadingSpecFactory,
567 + "get_spec_cls",
568 + return_value=_spec_cls_with_metric_definitions(metric_definitions),
569 + ):
570 + prom_metrics = OffloadPromMetrics(
571 + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
572 + metric_types={
573 + Gauge: _FakeMetric,
574 + Counter: _FakeMetric,
575 + Histogram: _FakeMetric,
576 + },
577 + labelnames=["model_name", "engine"],
578 + per_engine_labelvalues={0: ["model", "0"]},
579 + )
580 +
581 + assert (0, MY_COUNTER, ("a",)) not in prom_metrics.offloading_metrics
582 +
583 + prom_metrics.observe(
584 + {
585 + _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
586 + _StatsKey.DATA: {MY_COUNTER: {("a",): 7}},
587 + }
588 + )
589 +
590 + counter = prom_metrics.offloading_metrics[(0, MY_COUNTER, ("a",))]
591 + assert counter.increments == [7]
592 + assert counter.labelvalues == ("model", "0", "a")
593 + counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER]
594 + assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL]
595 +
596 +
597 +def test_prom_metrics_rejects_wrong_label_count():
598 + metric_definitions = {
599 + MY_COUNTER: OffloadingCounterMetadata(
600 + documentation="counter with a label",
601 + labelnames=(MY_LABEL,),
602 + ),
603 + }
604 + with patch.object(
605 + OffloadingSpecFactory,
606 + "get_spec_cls",
607 + return_value=_spec_cls_with_metric_definitions(metric_definitions),
608 + ):
609 + prom_metrics = OffloadPromMetrics(
610 + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
611 + metric_types={
612 + Gauge: _FakeMetric,
613 + Counter: _FakeMetric,
614 + Histogram: _FakeMetric,
615 + },
616 + labelnames=["model_name", "engine"],
617 + per_engine_labelvalues={0: ["model", "0"]},
618 + )
619 +
620 + with pytest.raises(AssertionError, match="expects 1 labels"):
621 + prom_metrics.observe(
622 + {
623 + _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER},
624 + _StatsKey.DATA: {MY_COUNTER: {("a", "extra"): 7}},
625 + }
626 + )
627 +
628 +
629 def test_prom_metrics_uses_configured_manager_metrics():
630 prom_metrics = OffloadPromMetrics(
631 vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type]
@@ -458,9 +654,9 @@ def test_aggregate_into_empty_stats():
654 PENDING_STORES: _MetricType.GAUGE,
655 },
656 _StatsKey.DATA: {
461 - LOAD_BYTES: 42,
462 - LOAD_SIZE: [10, 20],
463 - PENDING_STORES: 3,
657 + LOAD_BYTES: {(): 42},
658 + LOAD_SIZE: {(): [10, 20]},
659 + PENDING_STORES: {(): 3},
660 },
661 },
662 )
@@ -469,9 +665,9 @@ def test_aggregate_into_empty_stats():
665
666 assert result is empty
667 values = result.data[_StatsKey.DATA]
472 - assert values[LOAD_BYTES] == 42
473 - assert values[LOAD_SIZE] == [10, 20]
474 - assert values[PENDING_STORES] == 3
668 + assert _unlabeled(values, LOAD_BYTES) == 42
669 + assert _unlabeled(values, LOAD_SIZE) == [10, 20]
670 + assert _unlabeled(values, PENDING_STORES) == 3
671
672
673 def test_prom_metrics_multi_engine_routing():
@@ -490,14 +686,13 @@ def test_prom_metrics_multi_engine_routing():
686 prom_metrics.observe(
687 {
688 _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER},
493 - _StatsKey.DATA: {LOAD_BYTES: 100},
689 + _StatsKey.DATA: {LOAD_BYTES: {(): 100}},
690 },
691 engine_idx=1,
692 )
693
498 - engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)]
499 - engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)]
500 - assert engine0.increments == []
694 + assert (0, LOAD_BYTES, ()) not in prom_metrics.offloading_metrics
695 + engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES, ())]
696 assert engine1.increments == [100]
697
698
@@ -518,6 +713,6 @@ def test_prom_metrics_rejects_undeclared_metric():
713 prom_metrics.observe(
714 {
715 _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER},
521 - _StatsKey.DATA: {"unknown:metric": 1},
716 + _StatsKey.DATA: {"unknown:metric": {(): 1}},
717 }
718 )
vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py
+128 -63
@@ -112,7 +112,7 @@ class _StatsKey:
112
113 # Maps metric name -> _MetricType value
114 TYPES = "types"
115 - # Maps metric name -> observed value (number or list)
115 + # Maps metric name -> {label values tuple -> observed value (number or list)}
116 DATA = "data"
117
118
@@ -125,15 +125,17 @@ class OffloadingConnectorStats(KVConnectorStats):
125
126 {
127 _StatsKey.TYPES: {name: _MetricType.*, ...},
128 - _StatsKey.DATA: {name: value, ...},
128 + _StatsKey.DATA: {name: {labelvalues: value, ...}, ...},
129 }
130
131 This structure is self-describing: it survives IPC serialization
132 without needing the full ``OffloadingMetricMetadata`` objects on the
133 receiving side.
134
135 - Counter values are aggregated by summing, gauge values use the latest
136 - snapshot, and histogram values are lists of observed samples.
135 + Counter values are aggregated by summing per-label-tuple, gauge values
136 + use the latest snapshot per-label-tuple, and histogram values are lists of
137 + observed samples per-label-tuple. Unlabeled metrics use ``()`` as their
138 + labelvalues tuple.
139 """
140
141 def __post_init__(self):
@@ -160,26 +162,32 @@ class OffloadingConnectorStats(KVConnectorStats):
162 assert isinstance(other, OffloadingConnectorStats)
163 other_types = other._types
164 other_values = other._values
163 - for key, value in other_values.items():
165 + for key, other_label_values in other_values.items():
166 type_str = other_types.get(key)
167 if type_str is None:
168 raise AssertionError(f"Unknown offloading stats key: {key}")
169 self._types.setdefault(key, type_str)
168 - if type_str == _MetricType.HISTOGRAM:
169 - assert isinstance(value, list)
170 - if key not in self._values:
171 - self._values[key] = value
170 + current_label_values = self._values.setdefault(key, {})
171 + for labelvalues, value in other_label_values.items():
172 + if type_str == _MetricType.HISTOGRAM:
173 + assert isinstance(value, list)
174 + if labelvalues not in current_label_values:
175 + current_label_values[labelvalues] = list(value)
176 + else:
177 + assert isinstance(current_label_values[labelvalues], list)
178 + current_label_values[labelvalues].extend(value)
179 + elif type_str == _MetricType.COUNTER:
180 + assert isinstance(value, int | float)
181 + current_label_values[labelvalues] = (
182 + current_label_values.get(labelvalues, 0) + value
183 + )
184 + elif type_str == _MetricType.GAUGE:
185 + assert isinstance(value, int | float)
186 + current_label_values[labelvalues] = value
187 else:
173 - assert isinstance(self._values[key], list)
174 - self._values[key].extend(value)
175 - elif type_str == _MetricType.COUNTER:
176 - assert isinstance(value, int | float)
177 - self._values[key] = self._values.get(key, 0) + value
178 - elif type_str == _MetricType.GAUGE:
179 - assert isinstance(value, int | float)
180 - self._values[key] = value
181 - else:
182 - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
188 + raise AssertionError(
189 + f"Unknown metric type '{type_str}' for key: {key}"
190 + )
191 return self
192
193 def reduce(self) -> dict[str, int | float]:
@@ -190,44 +198,62 @@ class OffloadingConnectorStats(KVConnectorStats):
198 stats for the last time interval.
199 """
200 return_dict: dict[str, int | float] = {}
193 - for key, value in self._values.items():
201 + for key, label_value_map in self._values.items():
202 type_str = self._types.get(key)
203 if type_str is None:
204 raise AssertionError(f"Unknown offloading stats key: {key}")
197 - if type_str == _MetricType.HISTOGRAM:
198 - assert isinstance(value, list)
199 - return_dict[f"{key}_count"] = len(value)
200 - return_dict[f"{key}_sum"] = sum(value)
201 - elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
202 - assert isinstance(value, int | float)
203 - return_dict[key] = value
204 - else:
205 - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
205 + for labelvalues, value in label_value_map.items():
206 + key_with_labels = f"{key}:{labelvalues}" if labelvalues else key
207 + if type_str == _MetricType.HISTOGRAM:
208 + assert isinstance(value, list)
209 + return_dict[f"{key_with_labels}_count"] = len(value)
210 + return_dict[f"{key_with_labels}_sum"] = sum(value)
211 + elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
212 + assert isinstance(value, int | float)
213 + return_dict[key_with_labels] = value
214 + else:
215 + raise AssertionError(
216 + f"Unknown metric type '{type_str}' for key: {key}"
217 + )
218 return return_dict
219
220 def is_empty(self) -> bool:
221 return not self.data.get(_StatsKey.DATA)
222
223 def increase_counter(
212 - self, counter_name: str, counter_increase_value: int | float
224 + self,
225 + counter_name: str,
226 + counter_increase_value: int | float,
227 + labelvalues: tuple[str, ...] = (),
228 ) -> None:
229 """Increase a counter on the stats payload."""
230 self._types.setdefault(counter_name, _MetricType.COUNTER)
216 - self._values[counter_name] = (
217 - self._values.get(counter_name, 0) + counter_increase_value
231 + counter_values = self._values.setdefault(counter_name, {})
232 + counter_values[labelvalues] = (
233 + counter_values.get(labelvalues, 0) + counter_increase_value
234 )
235
220 - def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None:
236 + def set_gauge(
237 + self,
238 + gauge_name: str,
239 + gauge_value: int | float,
240 + labelvalues: tuple[str, ...] = (),
241 + ) -> None:
242 """Set a gauge snapshot on the stats payload."""
243 self._types.setdefault(gauge_name, _MetricType.GAUGE)
223 - self._values[gauge_name] = gauge_value
244 + gauge_values = self._values.setdefault(gauge_name, {})
245 + gauge_values[labelvalues] = gauge_value
246
247 def observe_histogram(
226 - self, histogram_name: str, histogram_value: int | float
248 + self,
249 + histogram_name: str,
250 + histogram_value: int | float,
251 + labelvalues: tuple[str, ...] = (),
252 ) -> None:
253 """Record a histogram observation on the stats payload."""
254 self._types.setdefault(histogram_name, _MetricType.HISTOGRAM)
230 - self._values.setdefault(histogram_name, []).append(histogram_value)
255 + histogram_values = self._values.setdefault(histogram_name, {})
256 + histogram_values.setdefault(labelvalues, []).append(histogram_value)
257
258
259 class OffloadPromMetrics(KVConnectorPromMetrics):
@@ -255,7 +281,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
281
282 self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec)
283 self._offloading_metric_defs: dict[str, PromMetricT] = {}
258 - self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {}
284 + # (engine_idx, metric_name, labelvalues) -> metric with bound labels
285 + self.offloading_metrics: dict[
286 + tuple[int, str, tuple[str, ...]], PromMetricT
287 + ] = {}
288
289 self._counter_kv_bytes = self._counter_cls(
290 name=_DEPRECATED_TOTAL_BYTES,
@@ -301,10 +330,6 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
330 self._offloading_metric_defs[metric_name] = self._create_metric(
331 metric_name, metadata
332 )
304 - for engine_idx, labelvalues in per_engine_labelvalues.items():
305 - self.offloading_metrics[(engine_idx, metric_name)] = (
306 - self._offloading_metric_defs[metric_name].labels(*labelvalues)
307 - )
333
334 def _create_metric(
335 self, metric_name: str, metadata: OffloadingMetricMetadata
@@ -312,7 +337,7 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
337 kwargs: dict[str, Any] = {
338 "name": metric_name,
339 "documentation": metadata.documentation,
315 - "labelnames": self._labelnames,
340 + "labelnames": self._labelnames + list(metadata.labelnames),
341 }
342 if isinstance(metadata, OffloadingCounterMetadata):
343 metric_cls = self._counter_cls
@@ -326,11 +351,37 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
351 raise AssertionError(f"Unknown offloading metric metadata: {metadata}")
352 return metric_cls(**kwargs)
353
354 + def _get_prometheus_metric(
355 + self,
356 + metric_name: str,
357 + labelvalues: tuple[str, ...],
358 + engine_idx: int,
359 + ) -> PromMetric:
360 + metadata = self._offloading_metric_metadata[metric_name]
361 + if len(labelvalues) != len(metadata.labelnames):
362 + raise AssertionError(
363 + f"Metric {metric_name} expects {len(metadata.labelnames)} labels, "
364 + f"got {len(labelvalues)}"
365 + )
366 + key = (engine_idx, metric_name, labelvalues)
367 + prom_metric = self.offloading_metrics.get(key)
368 + if prom_metric is None:
369 + engine_labelvalues = self.per_engine_labelvalues[engine_idx]
370 + prom_metric = self._offloading_metric_defs[metric_name].labels(
371 + *(engine_labelvalues + list(labelvalues))
372 + )
373 + self.offloading_metrics[key] = prom_metric
374 + return prom_metric
375 +
376 def _increase_counter(
330 - self, metric_name: str, value: int | float, engine_idx: int
377 + self,
378 + metric_name: str,
379 + value: int | float,
380 + labelvalues: tuple[str, ...],
381 + engine_idx: int,
382 ) -> None:
332 - self.offloading_metrics[(engine_idx, metric_name)].inc(value)
333 - if not self._observe_deprecated_metrics:
383 + self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value)
384 + if labelvalues or not self._observe_deprecated_metrics:
385 return
386 # Keep deprecated CPU offload transfer metrics updated during the
387 # transition to flat metric names.
@@ -343,15 +394,26 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
394 elif metric_name == _TransferMetricName.STORE_TIME:
395 self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value)
396
346 - def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None:
347 - self.offloading_metrics[(engine_idx, metric_name)].set(value)
397 + def _set_gauge(
398 + self,
399 + metric_name: str,
400 + value: int | float,
401 + labelvalues: tuple[str, ...],
402 + engine_idx: int,
403 + ) -> None:
404 + self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value)
405
406 def _observe_histogram(
350 - self, metric_name: str, value: list[int | float], engine_idx: int
407 + self,
408 + metric_name: str,
409 + value: list[int | float],
410 + labelvalues: tuple[str, ...],
411 + engine_idx: int,
412 ) -> None:
413 + prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx)
414 for observation in value:
353 - self.offloading_metrics[(engine_idx, metric_name)].observe(observation)
354 - if not self._observe_deprecated_metrics:
415 + prom_metric.observe(observation)
416 + if labelvalues or not self._observe_deprecated_metrics:
417 continue
418 # Keep deprecated CPU offload transfer metrics updated during the
419 # transition to flat metric names.
@@ -368,20 +430,23 @@ class OffloadPromMetrics(KVConnectorPromMetrics):
430 """Observe transfer statistics."""
431 metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
432 metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
371 - for key, value in metric_data.items():
433 + for key, label_value_map in metric_data.items():
434 type_str = metric_types.get(key)
435 if type_str is None:
436 raise AssertionError(f"Unknown offloading stats key: {key}")
437 assert key in self._offloading_metric_defs
376 - if type_str == _MetricType.COUNTER:
377 - assert isinstance(value, int | float)
378 - self._increase_counter(key, value, engine_idx)
379 - elif type_str == _MetricType.GAUGE:
380 - assert isinstance(value, int | float)
381 - self._set_gauge(key, value, engine_idx)
382 - elif type_str == _MetricType.HISTOGRAM:
383 - assert isinstance(value, list)
384 - assert all(isinstance(v, int | float) for v in value)
385 - self._observe_histogram(key, value, engine_idx)
386 - else:
387 - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}")
438 + for labelvalues, value in label_value_map.items():
439 + if type_str == _MetricType.COUNTER:
440 + assert isinstance(value, int | float)
441 + self._increase_counter(key, value, labelvalues, engine_idx)
442 + elif type_str == _MetricType.GAUGE:
443 + assert isinstance(value, int | float)
444 + self._set_gauge(key, value, labelvalues, engine_idx)
445 + elif type_str == _MetricType.HISTOGRAM:
446 + assert isinstance(value, list)
447 + assert all(isinstance(v, int | float) for v in value)
448 + self._observe_histogram(key, value, labelvalues, engine_idx)
449 + else:
450 + raise AssertionError(
451 + f"Unknown metric type '{type_str}' for key: {key}"
452 + )
vllm/v1/kv_offload/base.py
+1
@@ -129,6 +129,7 @@ The class provides the following primitives:
129 @dataclass(frozen=True)
130 class OffloadingMetricMetadata:
131 documentation: str
132 + labelnames: tuple[str, ...] = ()
133
134
135 @dataclass(frozen=True)