| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass, field |
| 4 | from typing import Any, TypeAlias, TypeVar |
| 5 | |
| 6 | from prometheus_client import Counter, Gauge, Histogram |
| 7 | |
| 8 | from vllm.config import KVTransferConfig, VllmConfig |
| 9 | from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory |
| 10 | from vllm.logger import init_logger |
| 11 | |
| 12 | PromMetric: TypeAlias = Gauge | Counter | Histogram |
| 13 | PromMetricT = TypeVar("PromMetricT", bound=PromMetric) |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | @dataclass |
| 19 | class KVConnectorStats: |
| 20 | """Base class for KV Connector Stats, a container for transfer performance |
| 21 | metrics or otherwise important telemetry from the connector. |
| 22 | All sub-classes need to be serializable as stats are sent from worker to |
| 23 | logger process. |
| 24 | """ |
| 25 | |
| 26 | data: dict[str, Any] = field(default_factory=dict) |
| 27 | |
| 28 | def to_dict(self) -> dict[str, Any]: |
| 29 | """Return the serializable connector stats payload.""" |
| 30 | return self.data |
| 31 | |
| 32 | def reset(self): |
| 33 | """Reset the stats, clear the state.""" |
| 34 | raise NotImplementedError |
| 35 | |
| 36 | def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats": |
| 37 | """Aggregate stats with another `KVConnectorStats` object.""" |
| 38 | raise NotImplementedError |
| 39 | |
| 40 | def reduce(self) -> dict[str, int | float]: |
| 41 | """Reduce the observations collected during a time interval to one or |
| 42 | more representative values (eg avg/median/sum of the series). |
| 43 | This is meant to be called by the logger to produce a summary of the |
| 44 | stats for the last time interval. |
| 45 | """ |
| 46 | raise NotImplementedError |
| 47 | |
| 48 | def is_empty(self) -> bool: |
| 49 | """Return True if the stats are empty.""" |
| 50 | raise NotImplementedError |
| 51 | |
| 52 | |
| 53 | class KVConnectorLogging: |
| 54 | def __init__(self, kv_transfer_config: KVTransferConfig | None): |
| 55 | # Instantiate the connector's stats class. |
| 56 | if kv_transfer_config and kv_transfer_config.kv_connector: |
| 57 | self.connector_cls = KVConnectorFactory.get_connector_class( |
| 58 | kv_transfer_config |
| 59 | ) |
| 60 | self.reset() |
| 61 | |
| 62 | def reset(self): |
| 63 | self.transfer_stats_accumulator: KVConnectorStats | None = None |
| 64 | |
| 65 | def observe(self, transfer_stats_data: dict[str, Any]): |
| 66 | # Should not be called when a KVConnector is not configured. |
| 67 | assert self.connector_cls is not None |
| 68 | # Called periodically when connector syncs with the scheduler. |
| 69 | # Note that this is not the same as the logging interval. |
| 70 | # We expect transfer_stats_data to be aggregated across all workers and |
| 71 | # consist of observations from a single connector or a MultiConnector. |
| 72 | transfer_stats = self.connector_cls.build_kv_connector_stats( |
| 73 | transfer_stats_data |
| 74 | ) |
| 75 | if transfer_stats is None: |
| 76 | logger.warning_once( |
| 77 | "The connector %s is collecting stats but " |
| 78 | "does not implement the " |
| 79 | "`build_kv_connector_stats` method. " |
| 80 | "Stats will not be logged.", |
| 81 | self.connector_cls, |
| 82 | ) |
| 83 | return |
| 84 | |
| 85 | if self.transfer_stats_accumulator is None: |
| 86 | self.transfer_stats_accumulator = transfer_stats |
| 87 | else: |
| 88 | # Accumulate last interval stats. |
| 89 | self.transfer_stats_accumulator = self.transfer_stats_accumulator.aggregate( |
| 90 | transfer_stats |
| 91 | ) |
| 92 | |
| 93 | def log(self, log_fn=logger.info): |
| 94 | """Log transfer metrics periodically, similar to throughput logging.""" |
| 95 | if ( |
| 96 | self.transfer_stats_accumulator |
| 97 | and not self.transfer_stats_accumulator.is_empty() |
| 98 | ): |
| 99 | # Produce a single cumulative stats object for the last time |
| 100 | # interval from the recorded observations. |
| 101 | xfer_metrics = self.transfer_stats_accumulator.reduce() |
| 102 | xfer_metrics_str = ", ".join(f"{k}={v}" for k, v in xfer_metrics.items()) |
| 103 | log_fn("KV Transfer metrics: %s", xfer_metrics_str) |
| 104 | |
| 105 | # Reset metrics for next interval |
| 106 | self.reset() |
| 107 | |
| 108 | |
| 109 | class KVConnectorPromMetrics: |
| 110 | """A base class for per-connector Prometheus metric registration |
| 111 | and recording. |
| 112 | """ |
| 113 | |
| 114 | def __init__( |
| 115 | self, |
| 116 | vllm_config: VllmConfig, |
| 117 | metric_types: dict[type[PromMetric], type[PromMetricT]], |
| 118 | labelnames: list[str], |
| 119 | per_engine_labelvalues: dict[int, list[object]], |
| 120 | ): |
| 121 | self._kv_transfer_config = vllm_config.kv_transfer_config |
| 122 | self._gauge_cls = metric_types[Gauge] |
| 123 | self._counter_cls = metric_types[Counter] |
| 124 | self._histogram_cls = metric_types[Histogram] |
| 125 | self._labelnames = labelnames |
| 126 | self.per_engine_labelvalues = per_engine_labelvalues |
| 127 | |
| 128 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 129 | """Record the supplied transfer statistics to Prometheus metrics. These |
| 130 | statistics are engine-specific, and should be recorded to a metric |
| 131 | with the appropriate 'engine' label. These metric instances can be |
| 132 | created using the create_metric_per_engine() helper method. |
| 133 | """ |
| 134 | raise NotImplementedError |
| 135 | |
| 136 | |
| 137 | class KVConnectorProm: |
| 138 | """Support for registering per-connector Prometheus metrics, and |
| 139 | recording transfer statistics to those metrics. Uses |
| 140 | KVConnectorBase.build_prom_metrics(). |
| 141 | """ |
| 142 | |
| 143 | _gauge_cls = Gauge |
| 144 | _counter_cls = Counter |
| 145 | _histogram_cls = Histogram |
| 146 | |
| 147 | def __init__( |
| 148 | self, |
| 149 | vllm_config: VllmConfig, |
| 150 | labelnames: list[str], |
| 151 | per_engine_labelvalues: dict[int, list[object]], |
| 152 | ): |
| 153 | self.prom_metrics: KVConnectorPromMetrics | None = None |
| 154 | kv_transfer_config = vllm_config.kv_transfer_config |
| 155 | if kv_transfer_config and kv_transfer_config.kv_connector: |
| 156 | connector_cls = KVConnectorFactory.get_connector_class(kv_transfer_config) |
| 157 | metric_types = { |
| 158 | Gauge: self._gauge_cls, |
| 159 | Counter: self._counter_cls, |
| 160 | Histogram: self._histogram_cls, |
| 161 | } |
| 162 | self.prom_metrics = connector_cls.build_prom_metrics( |
| 163 | vllm_config, |
| 164 | metric_types, |
| 165 | labelnames, |
| 166 | per_engine_labelvalues, |
| 167 | ) |
| 168 | |
| 169 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 170 | if self.prom_metrics is None: |
| 171 | return |
| 172 | self.prom_metrics.observe(transfer_stats_data, engine_idx) |