| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Stats and Prometheus metrics for the NIXL connector.""" |
| 4 | |
| 5 | import copy |
| 6 | from dataclasses import dataclass |
| 7 | from typing import TYPE_CHECKING, Any |
| 8 | |
| 9 | import numpy as np |
| 10 | |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( |
| 13 | KVConnectorPromMetrics, |
| 14 | KVConnectorStats, |
| 15 | PromMetric, |
| 16 | PromMetricT, |
| 17 | ) |
| 18 | from vllm.v1.metrics.utils import create_metric_per_engine |
| 19 | |
| 20 | if TYPE_CHECKING: |
| 21 | from vllm.distributed.nixl_utils import nixlXferTelemetry |
| 22 | |
| 23 | |
| 24 | @dataclass |
| 25 | class NixlKVConnectorStats(KVConnectorStats): |
| 26 | """Container for transfer performance metrics.""" |
| 27 | |
| 28 | def __post_init__(self): |
| 29 | if not self.data: |
| 30 | # Empty container init, no data is passed in. |
| 31 | self.reset() |
| 32 | |
| 33 | def reset(self): |
| 34 | # Must be serializable |
| 35 | self.data: dict[str, list[float | int]] = { |
| 36 | "transfer_duration": [], |
| 37 | "post_duration": [], |
| 38 | "bytes_transferred": [], |
| 39 | "num_descriptors": [], |
| 40 | "num_failed_transfers": [], |
| 41 | "num_failed_notifications": [], |
| 42 | "num_failed_handshakes": [], |
| 43 | "num_kv_expired_reqs": [], |
| 44 | "num_notifications_after_expiry": [], |
| 45 | } |
| 46 | |
| 47 | def record_transfer(self, res: "nixlXferTelemetry"): |
| 48 | # Keep metrics units consistent with rest of the code: time us->s |
| 49 | self.data["transfer_duration"].append(res.xferDuration / 1e6) |
| 50 | self.data["post_duration"].append(res.postDuration / 1e6) |
| 51 | self.data["bytes_transferred"].append(res.totalBytes) |
| 52 | self.data["num_descriptors"].append(res.descCount) |
| 53 | |
| 54 | def record_failed_transfer(self): |
| 55 | """Record a failed NIXL transfer operation.""" |
| 56 | self.data["num_failed_transfers"].append(1) |
| 57 | |
| 58 | def record_failed_notification(self): |
| 59 | """Record a failed NIXL notification (send_notif).""" |
| 60 | self.data["num_failed_notifications"].append(1) |
| 61 | |
| 62 | def record_failed_handshake(self): |
| 63 | """Record a failed NIXL handshake.""" |
| 64 | self.data["num_failed_handshakes"].append(1) |
| 65 | |
| 66 | def record_kv_expired_req(self): |
| 67 | """Record a request that had its KV blocks expire.""" |
| 68 | self.data["num_kv_expired_reqs"].append(1) |
| 69 | |
| 70 | def record_notification_after_expiry(self): |
| 71 | """Record a completion notification that arrived after its request expired.""" |
| 72 | self.data["num_notifications_after_expiry"].append(1) |
| 73 | |
| 74 | def clone_and_reset(self) -> "NixlKVConnectorStats": |
| 75 | old = copy.copy(self) |
| 76 | self.reset() |
| 77 | return old |
| 78 | |
| 79 | def is_empty(self) -> bool: |
| 80 | # Do not discard metrics update that are entirely failures related. |
| 81 | return ( |
| 82 | self.num_successful_transfers == 0 |
| 83 | and len(self.data["num_failed_transfers"]) == 0 |
| 84 | and len(self.data["num_failed_notifications"]) == 0 |
| 85 | and len(self.data["num_failed_handshakes"]) == 0 |
| 86 | and len(self.data["num_kv_expired_reqs"]) == 0 |
| 87 | and len(self.data["num_notifications_after_expiry"]) == 0 |
| 88 | ) |
| 89 | |
| 90 | def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: |
| 91 | if not other.is_empty(): |
| 92 | for k, v in other.data.items(): |
| 93 | accumulator = self.data[k] |
| 94 | assert isinstance(accumulator, list) |
| 95 | accumulator.extend(v) |
| 96 | return self |
| 97 | |
| 98 | def reduce(self) -> dict[str, int | float]: |
| 99 | # Compute compact representative stats suitable for CLI logging. |
| 100 | # Failure counts are reported on every interval: transfer, handshake |
| 101 | # and notification failures are grouped as sporadic |
| 102 | # lower-transport-layer events, while KV expiry is reported separately |
| 103 | # as it is an actionable autoscaler signal rather than a transport |
| 104 | # issue. Notifications arriving after expiry are surfaced too, as the |
| 105 | # reader may have consumed KV blocks that were already reused. |
| 106 | failure_counts = { |
| 107 | "Num failed transfers": len(self.data["num_failed_transfers"]) |
| 108 | + len(self.data["num_failed_handshakes"]) |
| 109 | + len(self.data["num_failed_notifications"]), |
| 110 | "Num KV expired reqs": len(self.data["num_kv_expired_reqs"]), |
| 111 | "Num notifs after expiry": len(self.data["num_notifications_after_expiry"]), |
| 112 | } |
| 113 | if self.num_successful_transfers == 0: |
| 114 | # Timing / throughput stats only cover successful transfers. If |
| 115 | # all requests in the interval were unsuccessful, the failure |
| 116 | # counts above still surface what happened. |
| 117 | return { |
| 118 | "Num successful transfers": 0, |
| 119 | "Avg xfer time (ms)": 0, |
| 120 | "P90 xfer time (ms)": 0, |
| 121 | "Avg post time (ms)": 0, |
| 122 | "P90 post time (ms)": 0, |
| 123 | "Avg MB per transfer": 0, |
| 124 | "Throughput (MB/s)": 0, |
| 125 | "Avg number of descriptors": 0, |
| 126 | **failure_counts, |
| 127 | } |
| 128 | |
| 129 | xfer_time = np.asarray(self.data["transfer_duration"]) |
| 130 | post_time = np.asarray(self.data["post_duration"]) |
| 131 | # Convert to MB for CLI logging. |
| 132 | mb = np.asarray(self.data["bytes_transferred"]) / 2**20 |
| 133 | descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32) |
| 134 | n = len(descs) |
| 135 | assert n == self.num_successful_transfers |
| 136 | |
| 137 | total_mb = mb.sum() |
| 138 | avg_mb = total_mb / n |
| 139 | |
| 140 | total_time_seconds = xfer_time.sum() |
| 141 | throughput_mb_s = total_mb / total_time_seconds |
| 142 | |
| 143 | return { |
| 144 | "Num successful transfers": n, |
| 145 | "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3), |
| 146 | "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), |
| 147 | "Avg post time (ms)": round(post_time.mean().item() * 1e3, 3), |
| 148 | "P90 post time (ms)": round(np.percentile(post_time, 90).item() * 1e3, 3), |
| 149 | "Avg MB per transfer": round(avg_mb.item(), 3), |
| 150 | "Throughput (MB/s)": round(throughput_mb_s.item(), 3), |
| 151 | "Avg number of descriptors": round(descs.mean().item(), 1), |
| 152 | **failure_counts, |
| 153 | } |
| 154 | |
| 155 | @property |
| 156 | def num_successful_transfers(self) -> int: |
| 157 | return len(self.data["transfer_duration"]) |
| 158 | |
| 159 | |
| 160 | class NixlPromMetrics(KVConnectorPromMetrics): |
| 161 | def __init__( |
| 162 | self, |
| 163 | vllm_config: VllmConfig, |
| 164 | metric_types: dict[type[PromMetric], type[PromMetricT]], |
| 165 | labelnames: list[str], |
| 166 | per_engine_labelvalues: dict[int, list[object]], |
| 167 | ): |
| 168 | super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues) |
| 169 | |
| 170 | buckets = [ |
| 171 | 0.001, |
| 172 | 0.005, |
| 173 | 0.01, |
| 174 | 0.025, |
| 175 | 0.05, |
| 176 | 0.075, |
| 177 | 0.1, |
| 178 | 0.2, |
| 179 | 0.3, |
| 180 | 0.5, |
| 181 | 0.75, |
| 182 | 1.0, |
| 183 | 5.0, |
| 184 | ] |
| 185 | nixl_histogram_xfer_time = self._histogram_cls( |
| 186 | name="vllm:nixl_xfer_time_seconds", |
| 187 | documentation="Histogram of transfer duration for NIXL KV Cache transfers.", |
| 188 | buckets=buckets[1:], |
| 189 | labelnames=labelnames, |
| 190 | ) |
| 191 | self.nixl_histogram_xfer_time = create_metric_per_engine( |
| 192 | nixl_histogram_xfer_time, self.per_engine_labelvalues |
| 193 | ) |
| 194 | nixl_histogram_post_time = self._histogram_cls( |
| 195 | name="vllm:nixl_post_time_seconds", |
| 196 | documentation="Histogram of transfer post time for NIXL KV" |
| 197 | " Cache transfers.", |
| 198 | buckets=buckets, |
| 199 | labelnames=labelnames, |
| 200 | ) |
| 201 | self.nixl_histogram_post_time = create_metric_per_engine( |
| 202 | nixl_histogram_post_time, self.per_engine_labelvalues |
| 203 | ) |
| 204 | # uniform 2kb to 16gb range |
| 205 | buckets = [2 ** (10 + i) for i in range(1, 25, 2)] |
| 206 | nixl_histogram_bytes_transferred = self._histogram_cls( |
| 207 | name="vllm:nixl_bytes_transferred", |
| 208 | documentation="Histogram of bytes transferred per NIXL KV Cache transfers.", |
| 209 | buckets=buckets, |
| 210 | labelnames=labelnames, |
| 211 | ) |
| 212 | self.nixl_histogram_bytes_transferred = create_metric_per_engine( |
| 213 | nixl_histogram_bytes_transferred, self.per_engine_labelvalues |
| 214 | ) |
| 215 | buckets = [ |
| 216 | 10, |
| 217 | 20, |
| 218 | 30, |
| 219 | 50, |
| 220 | 75, |
| 221 | 100, |
| 222 | 200, |
| 223 | 400, |
| 224 | 1000, |
| 225 | 2000, |
| 226 | 4000, |
| 227 | 10000, |
| 228 | 20000, |
| 229 | 50000, |
| 230 | ] |
| 231 | nixl_histogram_num_descriptors = self._histogram_cls( |
| 232 | name="vllm:nixl_num_descriptors", |
| 233 | documentation="Histogram of number of descriptors per NIXL" |
| 234 | " KV Cache transfers.", |
| 235 | buckets=buckets, |
| 236 | labelnames=labelnames, |
| 237 | ) |
| 238 | self.nixl_histogram_num_descriptors = create_metric_per_engine( |
| 239 | nixl_histogram_num_descriptors, self.per_engine_labelvalues |
| 240 | ) |
| 241 | counter_nixl_num_failed_transfers = self._counter_cls( |
| 242 | name="vllm:nixl_num_failed_transfers", |
| 243 | documentation="Number of failed NIXL KV Cache transfers, including" |
| 244 | " handshake and notification failures. NOTE: KV expiry is tracked" |
| 245 | " separately in vllm:nixl_num_kv_expired_reqs.", |
| 246 | labelnames=labelnames, |
| 247 | ) |
| 248 | self.counter_nixl_num_failed_transfers = create_metric_per_engine( |
| 249 | counter_nixl_num_failed_transfers, self.per_engine_labelvalues |
| 250 | ) |
| 251 | |
| 252 | counter_nixl_num_failed_notifications = self._counter_cls( |
| 253 | name="vllm:nixl_num_failed_notifications", |
| 254 | documentation="Number of failed NIXL KV Cache notifications. " |
| 255 | "Retained for compatibility; these failures are also included in " |
| 256 | "vllm:nixl_num_failed_transfers.", |
| 257 | labelnames=labelnames, |
| 258 | ) |
| 259 | self.counter_nixl_num_failed_notifications = create_metric_per_engine( |
| 260 | counter_nixl_num_failed_notifications, self.per_engine_labelvalues |
| 261 | ) |
| 262 | |
| 263 | counter_nixl_num_kv_expired_reqs = self._counter_cls( |
| 264 | name="vllm:nixl_num_kv_expired_reqs", |
| 265 | documentation="Number of requests that had their KV expire. " |
| 266 | "NOTE: This metric is tracked on the P instance.", |
| 267 | labelnames=labelnames, |
| 268 | ) |
| 269 | self.counter_nixl_num_kv_expired_reqs = create_metric_per_engine( |
| 270 | counter_nixl_num_kv_expired_reqs, self.per_engine_labelvalues |
| 271 | ) |
| 272 | |
| 273 | counter_nixl_num_notifications_after_expiry = self._counter_cls( |
| 274 | name="vllm:nixl_num_notifications_after_expiry", |
| 275 | documentation="Number of completion notifications for requests that " |
| 276 | "were no longer tracked, usually because their KV lease expired. The " |
| 277 | "KV blocks may have been reused before the transfer finished. Counted " |
| 278 | "per notification (one per remote rank), not per request.", |
| 279 | labelnames=labelnames, |
| 280 | ) |
| 281 | self.counter_nixl_num_notifications_after_expiry = create_metric_per_engine( |
| 282 | counter_nixl_num_notifications_after_expiry, self.per_engine_labelvalues |
| 283 | ) |
| 284 | |
| 285 | def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): |
| 286 | for prom_obj, list_item_key in zip( |
| 287 | [ |
| 288 | self.nixl_histogram_xfer_time, |
| 289 | self.nixl_histogram_post_time, |
| 290 | self.nixl_histogram_bytes_transferred, |
| 291 | self.nixl_histogram_num_descriptors, |
| 292 | ], |
| 293 | [ |
| 294 | "transfer_duration", |
| 295 | "post_duration", |
| 296 | "bytes_transferred", |
| 297 | "num_descriptors", |
| 298 | ], |
| 299 | ): |
| 300 | for list_item in transfer_stats_data[list_item_key]: |
| 301 | prom_obj[engine_idx].observe(list_item) |
| 302 | for counter_obj, counter_item_keys in zip( |
| 303 | [ |
| 304 | self.counter_nixl_num_failed_transfers, |
| 305 | self.counter_nixl_num_failed_notifications, |
| 306 | self.counter_nixl_num_kv_expired_reqs, |
| 307 | self.counter_nixl_num_notifications_after_expiry, |
| 308 | ], |
| 309 | [ |
| 310 | # Transfer, handshake and notification failures are grouped: |
| 311 | # all are sporadic lower-transport-layer events. KV expiry is |
| 312 | # reported separately as it signals autoscaler behavior, not |
| 313 | # transport health. |
| 314 | ( |
| 315 | "num_failed_transfers", |
| 316 | "num_failed_handshakes", |
| 317 | "num_failed_notifications", |
| 318 | ), |
| 319 | ("num_failed_notifications",), |
| 320 | ("num_kv_expired_reqs",), |
| 321 | ("num_notifications_after_expiry",), |
| 322 | ], |
| 323 | ): |
| 324 | for counter_item_key in counter_item_keys: |
| 325 | for list_item in transfer_stats_data[counter_item_key]: |
| 326 | counter_obj[engine_idx].inc(list_item) |