diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py index f9a4b37795..6f36a6c814 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_metrics.py @@ -22,7 +22,7 @@ from vllm.v1.kv_offload.base import ( OffloadingGaugeMetadata, OffloadingHistogramMetadata, ) -from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec +from vllm.v1.kv_offload.factory import OffloadingSpecFactory LOAD_BYTES = _TransferMetricName.LOAD_BYTES LOAD_TIME = _TransferMetricName.LOAD_TIME @@ -33,6 +33,8 @@ STORE_SIZE = _TransferMetricName.STORE_SIZE STORES_SKIPPED = "vllm:kv_offload_stores_skipped" PENDING_STORES = "vllm:kv_offload_pending_stores" LOOKUP_LATENCY = "vllm:kv_offload_lookup_latency_seconds" +MY_COUNTER = "my_counter" +MY_LABEL = "my_label" class _FakeMetric: @@ -67,6 +69,20 @@ class _FakeVllmConfig: ) +def _spec_cls_with_metric_definitions( + metric_definitions: dict[str, Any], +) -> type: + """Build a fake offloading spec class reporting the given metric + definitions, so tests don't need to patch the real CPU spec.""" + + class _FakeOffloadingSpec: + @staticmethod + def build_metric_definitions(extra_config): + return metric_definitions + + return _FakeOffloadingSpec + + def _metric_metadata(): return { LOAD_BYTES: OffloadingCounterMetadata( @@ -96,9 +112,17 @@ def _metric_metadata(): LOOKUP_LATENCY: OffloadingHistogramMetadata( documentation="lookup latency", ), + MY_COUNTER: OffloadingCounterMetadata( + documentation="counter with a label", + labelnames=(MY_LABEL,), + ), } +def _unlabeled(values: dict[str, Any], metric_name: str) -> Any: + return values[metric_name][()] + + def test_build_kv_connector_stats_with_none(): """Test that build_kv_connector_stats returns empty stats when given None.""" stats = OffloadingConnector.build_kv_connector_stats(data=None) @@ -131,13 +155,13 @@ def test_build_kv_connector_stats_reconstructs_offload_stats(): STORES_SKIPPED: _MetricType.COUNTER, }, _StatsKey.DATA: { - LOAD_BYTES: 24, - LOAD_TIME: 1.5, - LOAD_SIZE: [16, 8], - STORE_BYTES: 3, - STORE_TIME: 0.3, - STORE_SIZE: [1, 2], - STORES_SKIPPED: 5, + LOAD_BYTES: {(): 24}, + LOAD_TIME: {(): 1.5}, + LOAD_SIZE: {(): [16, 8]}, + STORE_BYTES: {(): 3}, + STORE_TIME: {(): 0.3}, + STORE_SIZE: {(): [1, 2]}, + STORES_SKIPPED: {(): 5}, }, } @@ -145,22 +169,28 @@ def test_build_kv_connector_stats_reconstructs_offload_stats(): assert isinstance(stats, OffloadingConnectorStats) values = stats.data[_StatsKey.DATA] - assert values[LOAD_BYTES] == 24 - assert values[LOAD_TIME] == 1.5 - assert values[LOAD_SIZE] == [16, 8] - assert values[STORE_BYTES] == 3 - assert values[STORE_TIME] == 0.3 - assert values[STORE_SIZE] == [1, 2] - assert values[STORES_SKIPPED] == 5 + assert _unlabeled(values, LOAD_BYTES) == 24 + assert _unlabeled(values, LOAD_TIME) == 1.5 + assert _unlabeled(values, LOAD_SIZE) == [16, 8] + assert _unlabeled(values, STORE_BYTES) == 3 + assert _unlabeled(values, STORE_TIME) == 0.3 + assert _unlabeled(values, STORE_SIZE) == [1, 2] + assert _unlabeled(values, STORES_SKIPPED) == 5 def _make_stats_data( metric_data: dict[str, Any], metric_metadata: dict[str, Any], ) -> dict[str, Any]: - """Build a structured data dict from flat metric data and metadata.""" + """Build a structured data dict from flat metric data and metadata. + + Values for unlabeled metrics may be passed flat (wrapped here under the + empty label tuple); values for labeled metrics must already be passed as + a ``{labelvalues: value}`` map. + """ metric_types = {} - for key in metric_data: + data = {} + for key, value in metric_data.items(): md = metric_metadata[key] if isinstance(md, OffloadingCounterMetadata): metric_types[key] = _MetricType.COUNTER @@ -168,9 +198,10 @@ def _make_stats_data( metric_types[key] = _MetricType.GAUGE elif isinstance(md, OffloadingHistogramMetadata): metric_types[key] = _MetricType.HISTOGRAM + data[key] = value if md.labelnames else {(): value} return { _StatsKey.TYPES: metric_types, - _StatsKey.DATA: metric_data, + _StatsKey.DATA: data, } @@ -215,34 +246,106 @@ def test_aggregate_same_connector(): assert result is stats1 # Should return self values = result.data[_StatsKey.DATA] - assert values[LOAD_BYTES] == 34 - assert values[LOAD_TIME] == 2.6 - assert values[LOAD_SIZE] == [16, 8, 3, 7] - assert values[STORE_BYTES] == 19 - assert values[STORE_TIME] == 2.3 - assert values[STORE_SIZE] == [1, 2, 16] - assert values[STORES_SKIPPED] == 4 - assert values[PENDING_STORES] == 1 - assert values[LOOKUP_LATENCY] == [0.1, 0.2, 0.3] + assert _unlabeled(values, LOAD_BYTES) == 34 + assert _unlabeled(values, LOAD_TIME) == 2.6 + assert _unlabeled(values, LOAD_SIZE) == [16, 8, 3, 7] + assert _unlabeled(values, STORE_BYTES) == 19 + assert _unlabeled(values, STORE_TIME) == 2.3 + assert _unlabeled(values, STORE_SIZE) == [1, 2, 16] + assert _unlabeled(values, STORES_SKIPPED) == 4 + assert _unlabeled(values, PENDING_STORES) == 1 + assert _unlabeled(values, LOOKUP_LATENCY) == [0.1, 0.2, 0.3] + + +def test_aggregate_labeled_metrics(): + metadata = _metric_metadata() + stats1 = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 10, + ("b",): 3, + }, + }, + metadata, + ), + ) + stats2 = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 7, + ("c",): 5, + }, + }, + metadata, + ), + ) + + stats1.aggregate(stats2) + + values = stats1.data[_StatsKey.DATA][MY_COUNTER] + assert values[("a",)] == 17 + assert values[("b",)] == 3 + assert values[("c",)] == 5 + + +def test_aggregate_labeled_metric_missing_from_self(): + """Aggregating a labeled metric that self doesn't have at all yet.""" + metadata = _metric_metadata() + stats1 = OffloadingConnectorStats() + stats2 = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 7, + ("b",): 5, + }, + }, + metadata, + ), + ) + + stats1.aggregate(stats2) + + values = stats1.data[_StatsKey.DATA][MY_COUNTER] + assert values[("a",)] == 7 + assert values[("b",)] == 5 + assert stats1.data[_StatsKey.TYPES][MY_COUNTER] == _MetricType.COUNTER + + +def test_helper_methods_accept_labeled_metrics(): + stats = OffloadingConnectorStats() + + stats.increase_counter(MY_COUNTER, 3, ("a",)) + stats.increase_counter(MY_COUNTER, 4, ("a",)) + stats.set_gauge(PENDING_STORES, 2, ("b",)) + stats.observe_histogram(LOOKUP_LATENCY, 0.1, ("b",)) + stats.observe_histogram(LOOKUP_LATENCY, 0.2, ("b",)) + + values = stats.data[_StatsKey.DATA] + assert values[MY_COUNTER][("a",)] == 7 + assert values[PENDING_STORES][("b",)] == 2 + assert values[LOOKUP_LATENCY][("b",)] == [0.1, 0.2] def test_aggregate_merges_types(): stats1 = OffloadingConnectorStats( data={ _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER}, - _StatsKey.DATA: {LOAD_BYTES: 1}, + _StatsKey.DATA: {LOAD_BYTES: {(): 1}}, }, ) stats2 = OffloadingConnectorStats( data={ _StatsKey.TYPES: {PENDING_STORES: _MetricType.GAUGE}, - _StatsKey.DATA: {PENDING_STORES: 2}, + _StatsKey.DATA: {PENDING_STORES: {(): 2}}, }, ) result = stats1.aggregate(stats2) - assert result.data[_StatsKey.DATA][PENDING_STORES] == 2 + assert _unlabeled(result.data[_StatsKey.DATA], PENDING_STORES) == 2 assert result.data[_StatsKey.TYPES][PENDING_STORES] == _MetricType.GAUGE @@ -283,6 +386,26 @@ def test_reduce(): assert reduced[f"{LOOKUP_LATENCY}_sum"] == sum([0.1, 0.2, 0.3]) +def test_reduce_labeled_metrics(): + metadata = _metric_metadata() + stats = OffloadingConnectorStats( + data=_make_stats_data( + { + MY_COUNTER: { + ("a",): 17, + ("b",): 3, + }, + }, + metadata, + ), + ) + + reduced = stats.reduce() + + assert reduced[f"{MY_COUNTER}:{('a',)}"] == 17 + assert reduced[f"{MY_COUNTER}:{('b',)}"] == 3 + + def test_reset(): """Test that reset() resets all connector stats.""" metadata = _metric_metadata() @@ -326,11 +449,11 @@ def test_prom_metrics_observes_manager_counter(): prom_metrics.observe( { _StatsKey.TYPES: {STORES_SKIPPED: _MetricType.COUNTER}, - _StatsKey.DATA: {STORES_SKIPPED: 7}, + _StatsKey.DATA: {STORES_SKIPPED: {(): 7}}, } ) - counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED)] + counter = prom_metrics.offloading_metrics[(0, STORES_SKIPPED, ())] assert counter.increments == [7] counter_def = prom_metrics._offloading_metric_defs[STORES_SKIPPED] 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(): STORE_SIZE: _MetricType.HISTOGRAM, }, _StatsKey.DATA: { - LOAD_BYTES: 24, - LOAD_TIME: 1.5, - LOAD_SIZE: [16, 8], - STORE_BYTES: 3, - STORE_TIME: 0.3, - STORE_SIZE: [1, 2], + LOAD_BYTES: {(): 24}, + LOAD_TIME: {(): 1.5}, + LOAD_SIZE: {(): [16, 8]}, + STORE_BYTES: {(): 3}, + STORE_TIME: {(): 0.3}, + STORE_SIZE: {(): [1, 2]}, }, } ) - assert prom_metrics.offloading_metrics[(0, LOAD_BYTES)].increments == [24] - assert prom_metrics.offloading_metrics[(0, LOAD_TIME)].increments == [1.5] - assert prom_metrics.offloading_metrics[(0, LOAD_SIZE)].observed == [16, 8] - assert prom_metrics.offloading_metrics[(0, STORE_BYTES)].increments == [3] - assert prom_metrics.offloading_metrics[(0, STORE_TIME)].increments == [0.3] - assert prom_metrics.offloading_metrics[(0, STORE_SIZE)].observed == [1, 2] + assert prom_metrics.offloading_metrics[(0, LOAD_BYTES, ())].increments == [24] + assert prom_metrics.offloading_metrics[(0, LOAD_TIME, ())].increments == [1.5] + assert prom_metrics.offloading_metrics[(0, LOAD_SIZE, ())].observed == [16, 8] + assert prom_metrics.offloading_metrics[(0, STORE_BYTES, ())].increments == [3] + assert prom_metrics.offloading_metrics[(0, STORE_TIME, ())].increments == [0.3] + assert prom_metrics.offloading_metrics[(0, STORE_SIZE, ())].observed == [1, 2] assert prom_metrics.counter_kv_bytes[(0, "CPU_to_GPU")].increments == [24] 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(): ), } with patch.object( - CPUOffloadingSpec, "build_metric_definitions", return_value=metric_definitions + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), ): prom_metrics = OffloadPromMetrics( vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] @@ -416,20 +541,91 @@ def test_prom_metrics_observes_manager_gauge_and_histogram(): LOOKUP_LATENCY: _MetricType.HISTOGRAM, }, _StatsKey.DATA: { - PENDING_STORES: 5, - LOOKUP_LATENCY: [0.2, 0.4], + PENDING_STORES: {(): 5}, + LOOKUP_LATENCY: {(): [0.2, 0.4]}, }, } ) - gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES)] - histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY)] + gauge = prom_metrics.offloading_metrics[(0, PENDING_STORES, ())] + histogram = prom_metrics.offloading_metrics[(0, LOOKUP_LATENCY, ())] assert gauge.set_values == [5] assert histogram.observed == [0.2, 0.4] histogram_def = prom_metrics._offloading_metric_defs[LOOKUP_LATENCY] assert histogram_def.kwargs["buckets"] == (0.1, 1.0) +def test_prom_metrics_lazily_observes_labeled_metric(): + metric_definitions = { + MY_COUNTER: OffloadingCounterMetadata( + documentation="counter with a label", + labelnames=(MY_LABEL,), + ), + } + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + assert (0, MY_COUNTER, ("a",)) not in prom_metrics.offloading_metrics + + prom_metrics.observe( + { + _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER}, + _StatsKey.DATA: {MY_COUNTER: {("a",): 7}}, + } + ) + + counter = prom_metrics.offloading_metrics[(0, MY_COUNTER, ("a",))] + assert counter.increments == [7] + assert counter.labelvalues == ("model", "0", "a") + counter_def = prom_metrics._offloading_metric_defs[MY_COUNTER] + assert counter_def.kwargs["labelnames"] == ["model_name", "engine", MY_LABEL] + + +def test_prom_metrics_rejects_wrong_label_count(): + metric_definitions = { + MY_COUNTER: OffloadingCounterMetadata( + documentation="counter with a label", + labelnames=(MY_LABEL,), + ), + } + with patch.object( + OffloadingSpecFactory, + "get_spec_cls", + return_value=_spec_cls_with_metric_definitions(metric_definitions), + ): + prom_metrics = OffloadPromMetrics( + vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] + metric_types={ + Gauge: _FakeMetric, + Counter: _FakeMetric, + Histogram: _FakeMetric, + }, + labelnames=["model_name", "engine"], + per_engine_labelvalues={0: ["model", "0"]}, + ) + + with pytest.raises(AssertionError, match="expects 1 labels"): + prom_metrics.observe( + { + _StatsKey.TYPES: {MY_COUNTER: _MetricType.COUNTER}, + _StatsKey.DATA: {MY_COUNTER: {("a", "extra"): 7}}, + } + ) + + def test_prom_metrics_uses_configured_manager_metrics(): prom_metrics = OffloadPromMetrics( vllm_config=_FakeVllmConfig(store_threshold=0), # type: ignore[arg-type] @@ -458,9 +654,9 @@ def test_aggregate_into_empty_stats(): PENDING_STORES: _MetricType.GAUGE, }, _StatsKey.DATA: { - LOAD_BYTES: 42, - LOAD_SIZE: [10, 20], - PENDING_STORES: 3, + LOAD_BYTES: {(): 42}, + LOAD_SIZE: {(): [10, 20]}, + PENDING_STORES: {(): 3}, }, }, ) @@ -469,9 +665,9 @@ def test_aggregate_into_empty_stats(): assert result is empty values = result.data[_StatsKey.DATA] - assert values[LOAD_BYTES] == 42 - assert values[LOAD_SIZE] == [10, 20] - assert values[PENDING_STORES] == 3 + assert _unlabeled(values, LOAD_BYTES) == 42 + assert _unlabeled(values, LOAD_SIZE) == [10, 20] + assert _unlabeled(values, PENDING_STORES) == 3 def test_prom_metrics_multi_engine_routing(): @@ -490,14 +686,13 @@ def test_prom_metrics_multi_engine_routing(): prom_metrics.observe( { _StatsKey.TYPES: {LOAD_BYTES: _MetricType.COUNTER}, - _StatsKey.DATA: {LOAD_BYTES: 100}, + _StatsKey.DATA: {LOAD_BYTES: {(): 100}}, }, engine_idx=1, ) - engine0 = prom_metrics.offloading_metrics[(0, LOAD_BYTES)] - engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES)] - assert engine0.increments == [] + assert (0, LOAD_BYTES, ()) not in prom_metrics.offloading_metrics + engine1 = prom_metrics.offloading_metrics[(1, LOAD_BYTES, ())] assert engine1.increments == [100] @@ -518,6 +713,6 @@ def test_prom_metrics_rejects_undeclared_metric(): prom_metrics.observe( { _StatsKey.TYPES: {"unknown:metric": _MetricType.COUNTER}, - _StatsKey.DATA: {"unknown:metric": 1}, + _StatsKey.DATA: {"unknown:metric": {(): 1}}, } ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index 3e4463924b..a90250d285 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -112,7 +112,7 @@ class _StatsKey: # Maps metric name -> _MetricType value TYPES = "types" - # Maps metric name -> observed value (number or list) + # Maps metric name -> {label values tuple -> observed value (number or list)} DATA = "data" @@ -125,15 +125,17 @@ class OffloadingConnectorStats(KVConnectorStats): { _StatsKey.TYPES: {name: _MetricType.*, ...}, - _StatsKey.DATA: {name: value, ...}, + _StatsKey.DATA: {name: {labelvalues: value, ...}, ...}, } This structure is self-describing: it survives IPC serialization without needing the full ``OffloadingMetricMetadata`` objects on the receiving side. - Counter values are aggregated by summing, gauge values use the latest - snapshot, and histogram values are lists of observed samples. + Counter values are aggregated by summing per-label-tuple, gauge values + use the latest snapshot per-label-tuple, and histogram values are lists of + observed samples per-label-tuple. Unlabeled metrics use ``()`` as their + labelvalues tuple. """ def __post_init__(self): @@ -160,26 +162,32 @@ class OffloadingConnectorStats(KVConnectorStats): assert isinstance(other, OffloadingConnectorStats) other_types = other._types other_values = other._values - for key, value in other_values.items(): + for key, other_label_values in other_values.items(): type_str = other_types.get(key) if type_str is None: raise AssertionError(f"Unknown offloading stats key: {key}") self._types.setdefault(key, type_str) - if type_str == _MetricType.HISTOGRAM: - assert isinstance(value, list) - if key not in self._values: - self._values[key] = value + current_label_values = self._values.setdefault(key, {}) + for labelvalues, value in other_label_values.items(): + if type_str == _MetricType.HISTOGRAM: + assert isinstance(value, list) + if labelvalues not in current_label_values: + current_label_values[labelvalues] = list(value) + else: + assert isinstance(current_label_values[labelvalues], list) + current_label_values[labelvalues].extend(value) + elif type_str == _MetricType.COUNTER: + assert isinstance(value, int | float) + current_label_values[labelvalues] = ( + current_label_values.get(labelvalues, 0) + value + ) + elif type_str == _MetricType.GAUGE: + assert isinstance(value, int | float) + current_label_values[labelvalues] = value else: - assert isinstance(self._values[key], list) - self._values[key].extend(value) - elif type_str == _MetricType.COUNTER: - assert isinstance(value, int | float) - self._values[key] = self._values.get(key, 0) + value - elif type_str == _MetricType.GAUGE: - assert isinstance(value, int | float) - self._values[key] = value - else: - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}") + raise AssertionError( + f"Unknown metric type '{type_str}' for key: {key}" + ) return self def reduce(self) -> dict[str, int | float]: @@ -190,44 +198,62 @@ class OffloadingConnectorStats(KVConnectorStats): stats for the last time interval. """ return_dict: dict[str, int | float] = {} - for key, value in self._values.items(): + for key, label_value_map in self._values.items(): type_str = self._types.get(key) if type_str is None: raise AssertionError(f"Unknown offloading stats key: {key}") - if type_str == _MetricType.HISTOGRAM: - assert isinstance(value, list) - return_dict[f"{key}_count"] = len(value) - return_dict[f"{key}_sum"] = sum(value) - elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE): - assert isinstance(value, int | float) - return_dict[key] = value - else: - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}") + for labelvalues, value in label_value_map.items(): + key_with_labels = f"{key}:{labelvalues}" if labelvalues else key + if type_str == _MetricType.HISTOGRAM: + assert isinstance(value, list) + return_dict[f"{key_with_labels}_count"] = len(value) + return_dict[f"{key_with_labels}_sum"] = sum(value) + elif type_str in (_MetricType.COUNTER, _MetricType.GAUGE): + assert isinstance(value, int | float) + return_dict[key_with_labels] = value + else: + raise AssertionError( + f"Unknown metric type '{type_str}' for key: {key}" + ) return return_dict def is_empty(self) -> bool: return not self.data.get(_StatsKey.DATA) def increase_counter( - self, counter_name: str, counter_increase_value: int | float + self, + counter_name: str, + counter_increase_value: int | float, + labelvalues: tuple[str, ...] = (), ) -> None: """Increase a counter on the stats payload.""" self._types.setdefault(counter_name, _MetricType.COUNTER) - self._values[counter_name] = ( - self._values.get(counter_name, 0) + counter_increase_value + counter_values = self._values.setdefault(counter_name, {}) + counter_values[labelvalues] = ( + counter_values.get(labelvalues, 0) + counter_increase_value ) - def set_gauge(self, gauge_name: str, gauge_value: int | float) -> None: + def set_gauge( + self, + gauge_name: str, + gauge_value: int | float, + labelvalues: tuple[str, ...] = (), + ) -> None: """Set a gauge snapshot on the stats payload.""" self._types.setdefault(gauge_name, _MetricType.GAUGE) - self._values[gauge_name] = gauge_value + gauge_values = self._values.setdefault(gauge_name, {}) + gauge_values[labelvalues] = gauge_value def observe_histogram( - self, histogram_name: str, histogram_value: int | float + self, + histogram_name: str, + histogram_value: int | float, + labelvalues: tuple[str, ...] = (), ) -> None: """Record a histogram observation on the stats payload.""" self._types.setdefault(histogram_name, _MetricType.HISTOGRAM) - self._values.setdefault(histogram_name, []).append(histogram_value) + histogram_values = self._values.setdefault(histogram_name, {}) + histogram_values.setdefault(labelvalues, []).append(histogram_value) class OffloadPromMetrics(KVConnectorPromMetrics): @@ -255,7 +281,10 @@ class OffloadPromMetrics(KVConnectorPromMetrics): self._observe_deprecated_metrics = issubclass(spec_cls, CPUOffloadingSpec) self._offloading_metric_defs: dict[str, PromMetricT] = {} - self.offloading_metrics: dict[tuple[int, str], PromMetricT] = {} + # (engine_idx, metric_name, labelvalues) -> metric with bound labels + self.offloading_metrics: dict[ + tuple[int, str, tuple[str, ...]], PromMetricT + ] = {} self._counter_kv_bytes = self._counter_cls( name=_DEPRECATED_TOTAL_BYTES, @@ -301,10 +330,6 @@ class OffloadPromMetrics(KVConnectorPromMetrics): self._offloading_metric_defs[metric_name] = self._create_metric( metric_name, metadata ) - for engine_idx, labelvalues in per_engine_labelvalues.items(): - self.offloading_metrics[(engine_idx, metric_name)] = ( - self._offloading_metric_defs[metric_name].labels(*labelvalues) - ) def _create_metric( self, metric_name: str, metadata: OffloadingMetricMetadata @@ -312,7 +337,7 @@ class OffloadPromMetrics(KVConnectorPromMetrics): kwargs: dict[str, Any] = { "name": metric_name, "documentation": metadata.documentation, - "labelnames": self._labelnames, + "labelnames": self._labelnames + list(metadata.labelnames), } if isinstance(metadata, OffloadingCounterMetadata): metric_cls = self._counter_cls @@ -326,11 +351,37 @@ class OffloadPromMetrics(KVConnectorPromMetrics): raise AssertionError(f"Unknown offloading metric metadata: {metadata}") return metric_cls(**kwargs) + def _get_prometheus_metric( + self, + metric_name: str, + labelvalues: tuple[str, ...], + engine_idx: int, + ) -> PromMetric: + metadata = self._offloading_metric_metadata[metric_name] + if len(labelvalues) != len(metadata.labelnames): + raise AssertionError( + f"Metric {metric_name} expects {len(metadata.labelnames)} labels, " + f"got {len(labelvalues)}" + ) + key = (engine_idx, metric_name, labelvalues) + prom_metric = self.offloading_metrics.get(key) + if prom_metric is None: + engine_labelvalues = self.per_engine_labelvalues[engine_idx] + prom_metric = self._offloading_metric_defs[metric_name].labels( + *(engine_labelvalues + list(labelvalues)) + ) + self.offloading_metrics[key] = prom_metric + return prom_metric + def _increase_counter( - self, metric_name: str, value: int | float, engine_idx: int + self, + metric_name: str, + value: int | float, + labelvalues: tuple[str, ...], + engine_idx: int, ) -> None: - self.offloading_metrics[(engine_idx, metric_name)].inc(value) - if not self._observe_deprecated_metrics: + self._get_prometheus_metric(metric_name, labelvalues, engine_idx).inc(value) + if labelvalues or not self._observe_deprecated_metrics: return # Keep deprecated CPU offload transfer metrics updated during the # transition to flat metric names. @@ -343,15 +394,26 @@ class OffloadPromMetrics(KVConnectorPromMetrics): elif metric_name == _TransferMetricName.STORE_TIME: self.counter_kv_transfer_time[(engine_idx, _TransferType.STORE)].inc(value) - def _set_gauge(self, metric_name: str, value: int | float, engine_idx: int) -> None: - self.offloading_metrics[(engine_idx, metric_name)].set(value) + def _set_gauge( + self, + metric_name: str, + value: int | float, + labelvalues: tuple[str, ...], + engine_idx: int, + ) -> None: + self._get_prometheus_metric(metric_name, labelvalues, engine_idx).set(value) def _observe_histogram( - self, metric_name: str, value: list[int | float], engine_idx: int + self, + metric_name: str, + value: list[int | float], + labelvalues: tuple[str, ...], + engine_idx: int, ) -> None: + prom_metric = self._get_prometheus_metric(metric_name, labelvalues, engine_idx) for observation in value: - self.offloading_metrics[(engine_idx, metric_name)].observe(observation) - if not self._observe_deprecated_metrics: + prom_metric.observe(observation) + if labelvalues or not self._observe_deprecated_metrics: continue # Keep deprecated CPU offload transfer metrics updated during the # transition to flat metric names. @@ -368,20 +430,23 @@ class OffloadPromMetrics(KVConnectorPromMetrics): """Observe transfer statistics.""" metric_types = transfer_stats_data.get(_StatsKey.TYPES, {}) metric_data = transfer_stats_data.get(_StatsKey.DATA, {}) - for key, value in metric_data.items(): + for key, label_value_map in metric_data.items(): type_str = metric_types.get(key) if type_str is None: raise AssertionError(f"Unknown offloading stats key: {key}") assert key in self._offloading_metric_defs - if type_str == _MetricType.COUNTER: - assert isinstance(value, int | float) - self._increase_counter(key, value, engine_idx) - elif type_str == _MetricType.GAUGE: - assert isinstance(value, int | float) - self._set_gauge(key, value, engine_idx) - elif type_str == _MetricType.HISTOGRAM: - assert isinstance(value, list) - assert all(isinstance(v, int | float) for v in value) - self._observe_histogram(key, value, engine_idx) - else: - raise AssertionError(f"Unknown metric type '{type_str}' for key: {key}") + for labelvalues, value in label_value_map.items(): + if type_str == _MetricType.COUNTER: + assert isinstance(value, int | float) + self._increase_counter(key, value, labelvalues, engine_idx) + elif type_str == _MetricType.GAUGE: + assert isinstance(value, int | float) + self._set_gauge(key, value, labelvalues, engine_idx) + elif type_str == _MetricType.HISTOGRAM: + assert isinstance(value, list) + assert all(isinstance(v, int | float) for v in value) + self._observe_histogram(key, value, labelvalues, engine_idx) + else: + raise AssertionError( + f"Unknown metric type '{type_str}' for key: {key}" + ) diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index 2d27c14fe8..904003bdcb 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -129,6 +129,7 @@ The class provides the following primitives: @dataclass(frozen=True) class OffloadingMetricMetadata: documentation: str + labelnames: tuple[str, ...] = () @dataclass(frozen=True)