| 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 ECTransferConfig, VllmConfig |
| 9 | from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory |
| 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 ECConnectorStats: |
| 20 | """Base class for EC 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 reset(self): |
| 29 | """Reset the stats, clear the state.""" |
| 30 | raise NotImplementedError |
| 31 | |
| 32 | def aggregate(self, other: "ECConnectorStats") -> "ECConnectorStats": |
| 33 | """Aggregate stats with another `ECConnectorStats` object.""" |
| 34 | raise NotImplementedError |
| 35 | |
| 36 | def reduce(self) -> dict[str, int | float]: |
| 37 | """Reduce the observations collected during a time interval to one or |
| 38 | more representative values (eg avg/median/sum of the series). |
| 39 | This is meant to be called by the logger to produce a summary of the |
| 40 | stats for the last time interval. |
| 41 | """ |
| 42 | raise NotImplementedError |
| 43 | |
| 44 | def is_empty(self) -> bool: |
| 45 | """Return True if the stats are empty.""" |
| 46 | raise NotImplementedError |
| 47 | |
| 48 | |
| 49 | class ECConnectorLogging: |
| 50 | def __init__(self, ec_transfer_config: ECTransferConfig | None): |
| 51 | # Instantiate the connector's stats class. |
| 52 | if ec_transfer_config and ec_transfer_config.ec_connector: |
| 53 | self.connector_cls = ECConnectorFactory.get_connector_class( |
| 54 | ec_transfer_config |
| 55 | ) |
| 56 | self.reset() |
| 57 | |
| 58 | def reset(self): |
| 59 | self.transfer_stats_accumulator: ECConnectorStats | None = None |
| 60 | |
| 61 | def observe(self, transfer_stats_data: dict[str, Any]): |
| 62 | # Should not be called when an ECConnector is not configured. |
| 63 | assert self.connector_cls is not None |
| 64 | # Called periodically when connector syncs with the scheduler. |
| 65 | # Note that this is not the same as the logging interval. |
| 66 | # We expect transfer_stats_data to be aggregated across all workers and |
| 67 | # consist of observations from a single connector. |
| 68 | transfer_stats = self.connector_cls.build_ec_connector_stats( |
| 69 | transfer_stats_data |
| 70 | ) |
| 71 | if transfer_stats is None: |
| 72 | logger.warning_once( |
| 73 | "The connector %s is collecting stats but " |
| 74 | "does not implement the " |
| 75 | "`build_ec_connector_stats` method. " |
| 76 | "Stats will not be logged.", |
| 77 | self.connector_cls, |
| 78 | ) |
| 79 | return |
| 80 | |
| 81 | if self.transfer_stats_accumulator is None: |
| 82 | self.transfer_stats_accumulator = transfer_stats |
| 83 | else: |
| 84 | # Accumulate last interval stats. |
| 85 | self.transfer_stats_accumulator = self.transfer_stats_accumulator.aggregate( |
| 86 | transfer_stats |
| 87 | ) |
| 88 | |
| 89 | def log(self, log_fn=logger.info): |
| 90 | """Log transfer metrics periodically, similar to throughput logging""" |
| 91 | if ( |
| 92 | self.transfer_stats_accumulator |
| 93 | and not self.transfer_stats_accumulator.is_empty() |
| 94 | ): |
| 95 | # Produce a single cumulative stats object for the last time |
| 96 | # interval from the recorded observations. |
| 97 | xfer_metrics = self.transfer_stats_accumulator.reduce() |
| 98 | xfer_metrics_str = ", ".join(f"{k}={v}" for k, v in xfer_metrics.items()) |
| 99 | log_fn("EC Transfer metrics: %s", xfer_metrics_str) |
| 100 | |
| 101 | # Reset metrics for next interval |
| 102 | self.reset() |
| 103 | |
| 104 | |
| 105 | class ECConnectorPromMetrics: |
| 106 | """A base class for per-connector Prometheus metric registration |
| 107 | and recording. |
| 108 | """ |
| 109 | |
| 110 | def __init__( |
| 111 | self, |
| 112 | vllm_config: VllmConfig, |
| 113 | metric_types: dict[type[PromMetric], type[PromMetricT]], |
| 114 | labelnames: list[str], |
| 115 | per_engine_labelvalues: dict[int, list[object]], |
| 116 | ): |
| 117 | self._ec_transfer_config = vllm_config.ec_transfer_config |
| 118 | self._gauge_cls = metric_types[Gauge] |
| 119 | self._counter_cls = metric_types[Counter] |
| 120 | self._histogram_cls = metric_types[Histogram] |
| 121 | self._labelnames = labelnames |
| 122 | self.per_engine_labelvalues = per_engine_labelvalues |
| 123 | |
| 124 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 125 | """Record the supplied transfer statistics to Prometheus metrics. These |
| 126 | statistics are engine-specific, and should be recorded to a metric |
| 127 | with the appropriate 'engine' label. These metric instances can be |
| 128 | created using the create_metric_per_engine() helper method. |
| 129 | """ |
| 130 | raise NotImplementedError |
| 131 | |
| 132 | |
| 133 | class ECConnectorProm: |
| 134 | """Support for registering per-connector Prometheus metrics, and |
| 135 | recording transfer statistics to those metrics. Uses |
| 136 | ECConnectorBase.build_prom_metrics(). |
| 137 | """ |
| 138 | |
| 139 | _gauge_cls = Gauge |
| 140 | _counter_cls = Counter |
| 141 | _histogram_cls = Histogram |
| 142 | |
| 143 | def __init__( |
| 144 | self, |
| 145 | vllm_config: VllmConfig, |
| 146 | labelnames: list[str], |
| 147 | per_engine_labelvalues: dict[int, list[object]], |
| 148 | ): |
| 149 | self.prom_metrics: ECConnectorPromMetrics | None = None |
| 150 | ec_transfer_config = vllm_config.ec_transfer_config |
| 151 | if ec_transfer_config and ec_transfer_config.ec_connector: |
| 152 | connector_cls = ECConnectorFactory.get_connector_class(ec_transfer_config) |
| 153 | metric_types = { |
| 154 | Gauge: self._gauge_cls, |
| 155 | Counter: self._counter_cls, |
| 156 | Histogram: self._histogram_cls, |
| 157 | } |
| 158 | self.prom_metrics = connector_cls.build_prom_metrics( |
| 159 | vllm_config, |
| 160 | metric_types, |
| 161 | labelnames, |
| 162 | per_engine_labelvalues, |
| 163 | ) |
| 164 | |
| 165 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 166 | if self.prom_metrics is None: |
| 167 | return |
| 168 | self.prom_metrics.observe(transfer_stats_data, engine_idx) |