| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Stats and Prometheus metrics for the HiSparse connector.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | from typing import Any |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( |
| 10 | KVConnectorPromMetrics, |
| 11 | KVConnectorStats, |
| 12 | PromMetric, |
| 13 | PromMetricT, |
| 14 | ) |
| 15 | from vllm.v1.metrics.utils import create_metric_per_engine |
| 16 | |
| 17 | _HISPARSE_COUNTERS: tuple[tuple[str, str], ...] = ( |
| 18 | ("cache_hits", "Number of HiSparse device hot-buffer hits."), |
| 19 | ("cache_misses", "Number of HiSparse device hot-buffer misses."), |
| 20 | ( |
| 21 | "host_to_device_bytes", |
| 22 | "Bytes transferred from host KV storage to HiSparse hot buffers.", |
| 23 | ), |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | @dataclass |
| 28 | class HiSparseKVConnectorStats(KVConnectorStats): |
| 29 | """Container for HiSparse hot-buffer residency metrics. |
| 30 | |
| 31 | Each list entry is the delta recorded over one device counter snapshot |
| 32 | interval, so a list can hold multiple snapshots per logging interval. |
| 33 | """ |
| 34 | |
| 35 | def __post_init__(self): |
| 36 | if not self.data: |
| 37 | # Empty container init, no data is passed in. |
| 38 | self.reset() |
| 39 | |
| 40 | def reset(self): |
| 41 | # Must be serializable |
| 42 | self.data: dict[str, list[int]] = { |
| 43 | "cache_hits": [], |
| 44 | "cache_misses": [], |
| 45 | "host_to_device_bytes": [], |
| 46 | } |
| 47 | |
| 48 | def record_snapshot(self, hits: int, misses: int, host_to_device_bytes: int): |
| 49 | self.data["cache_hits"].append(hits) |
| 50 | self.data["cache_misses"].append(misses) |
| 51 | self.data["host_to_device_bytes"].append(host_to_device_bytes) |
| 52 | |
| 53 | def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: |
| 54 | if not other.is_empty(): |
| 55 | for k, v in other.data.items(): |
| 56 | accumulator = self.data[k] |
| 57 | assert isinstance(accumulator, list) |
| 58 | accumulator.extend(v) |
| 59 | return self |
| 60 | |
| 61 | def reduce(self) -> dict[str, int | float]: |
| 62 | # Compute compact representative stats suitable for CLI logging. |
| 63 | return { |
| 64 | "HiSparse hot-buffer hits": sum(self.data["cache_hits"]), |
| 65 | "HiSparse hot-buffer misses": sum(self.data["cache_misses"]), |
| 66 | "HiSparse host-to-device bytes": sum(self.data["host_to_device_bytes"]), |
| 67 | } |
| 68 | |
| 69 | def is_empty(self) -> bool: |
| 70 | return all(len(values) == 0 for values in self.data.values()) |
| 71 | |
| 72 | |
| 73 | class HiSparsePromMetrics(KVConnectorPromMetrics): |
| 74 | def __init__( |
| 75 | self, |
| 76 | vllm_config: VllmConfig, |
| 77 | metric_types: dict[type[PromMetric], type[PromMetricT]], |
| 78 | labelnames: list[str], |
| 79 | per_engine_labelvalues: dict[int, list[object]], |
| 80 | ): |
| 81 | super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) |
| 82 | |
| 83 | self.hisparse_counters: dict[str, Any] = {} |
| 84 | for name, documentation in _HISPARSE_COUNTERS: |
| 85 | counter = self._counter_cls( |
| 86 | name=f"vllm:hisparse_{name}", |
| 87 | documentation=documentation, |
| 88 | labelnames=labelnames, |
| 89 | ) |
| 90 | self.hisparse_counters[name] = create_metric_per_engine( |
| 91 | counter, self.per_engine_labelvalues |
| 92 | ) |
| 93 | |
| 94 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 95 | for name, counter in self.hisparse_counters.items(): |
| 96 | for value in transfer_stats_data.get(name, []): |
| 97 | counter[engine_idx].inc(value) |