| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Stats container for the Mooncake connector.""" |
| 4 | |
| 5 | import threading |
| 6 | from dataclasses import dataclass |
| 7 | from typing import Any |
| 8 | |
| 9 | import numpy as np |
| 10 | |
| 11 | from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( |
| 12 | KVConnectorStats, |
| 13 | ) |
| 14 | |
| 15 | # TODO(mooncake-stats): add MooncakePromMetrics (mirror NixlPromMetrics) |
| 16 | # and wire it via MooncakeConnector.build_prom_metrics in a follow-up PR. |
| 17 | |
| 18 | |
| 19 | @dataclass |
| 20 | class MooncakeKVConnectorStats(KVConnectorStats): |
| 21 | """Container for Mooncake KV transfer performance metrics. |
| 22 | |
| 23 | `_lock` serializes record_* against clone_and_reset so each row's |
| 24 | appends are atomic and column lengths stay aligned. Writers run on |
| 25 | the sender pool / receiver loop / sender loop; reader runs on the |
| 26 | main worker thread. |
| 27 | """ |
| 28 | |
| 29 | def __post_init__(self): |
| 30 | self._lock = threading.Lock() |
| 31 | if not self.data: |
| 32 | self.reset() |
| 33 | |
| 34 | # threading.Lock is not picklable; strip it from the wire form and |
| 35 | # rebuild a fresh per-process lock on the receiver side. |
| 36 | def __getstate__(self) -> dict[str, Any]: |
| 37 | state = self.__dict__.copy() |
| 38 | state.pop("_lock", None) |
| 39 | return state |
| 40 | |
| 41 | def __setstate__(self, state: dict[str, Any]) -> None: |
| 42 | self.__dict__.update(state) |
| 43 | self._lock = threading.Lock() |
| 44 | |
| 45 | def reset(self): |
| 46 | self.data: dict[str, list[float | int]] = { |
| 47 | "transfer_duration": [], |
| 48 | "bytes_transferred": [], |
| 49 | "num_descriptors": [], |
| 50 | "num_failed_transfers": [], |
| 51 | "num_failed_recvs": [], |
| 52 | "num_kv_expired_reqs": [], |
| 53 | } |
| 54 | |
| 55 | def record_transfer(self, duration_s: float, total_bytes: int, num_descs: int): |
| 56 | with self._lock: |
| 57 | self.data["transfer_duration"].append(duration_s) |
| 58 | self.data["bytes_transferred"].append(total_bytes) |
| 59 | self.data["num_descriptors"].append(num_descs) |
| 60 | |
| 61 | # Failure counters store a list of 1s so a future Prom counter can iterate |
| 62 | # with .inc(list_item), mirroring NIXL's NixlPromMetrics.observe. |
| 63 | def record_failed_transfer(self): |
| 64 | with self._lock: |
| 65 | self.data["num_failed_transfers"].append(1) |
| 66 | |
| 67 | def record_failed_recv(self): |
| 68 | with self._lock: |
| 69 | self.data["num_failed_recvs"].append(1) |
| 70 | |
| 71 | def record_kv_expired_req(self): |
| 72 | with self._lock: |
| 73 | self.data["num_kv_expired_reqs"].append(1) |
| 74 | |
| 75 | def clone_and_reset(self) -> "MooncakeKVConnectorStats": |
| 76 | # Copy lists under the lock for length alignment; return a fresh |
| 77 | # instance so the snapshot has its own _lock. |
| 78 | with self._lock: |
| 79 | snapshot_data: dict[str, list[float | int]] = { |
| 80 | k: list(v) for k, v in self.data.items() |
| 81 | } |
| 82 | self.reset() |
| 83 | return MooncakeKVConnectorStats(data=snapshot_data) |
| 84 | |
| 85 | def is_empty(self) -> bool: |
| 86 | return ( |
| 87 | self.num_successful_transfers == 0 |
| 88 | and len(self.data["num_failed_transfers"]) == 0 |
| 89 | and len(self.data["num_failed_recvs"]) == 0 |
| 90 | and len(self.data["num_kv_expired_reqs"]) == 0 |
| 91 | ) |
| 92 | |
| 93 | def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: |
| 94 | if not other.is_empty(): |
| 95 | for k, v in other.data.items(): |
| 96 | accumulator = self.data[k] |
| 97 | assert isinstance(accumulator, list) |
| 98 | accumulator.extend(v) |
| 99 | return self |
| 100 | |
| 101 | def reduce(self) -> dict[str, int | float]: |
| 102 | num_failed_transfers = len(self.data["num_failed_transfers"]) |
| 103 | num_failed_recvs = len(self.data["num_failed_recvs"]) |
| 104 | num_kv_expired_reqs = len(self.data["num_kv_expired_reqs"]) |
| 105 | |
| 106 | if self.num_successful_transfers == 0: |
| 107 | return { |
| 108 | "Num successful transfers": 0, |
| 109 | "Avg xfer time (ms)": 0, |
| 110 | "P90 xfer time (ms)": 0, |
| 111 | "Avg MB per transfer": 0, |
| 112 | "Throughput (MB/s)": 0, |
| 113 | "Avg number of descriptors": 0, |
| 114 | "Num failed transfers": num_failed_transfers, |
| 115 | "Num failed recvs": num_failed_recvs, |
| 116 | "Num KV expired reqs": num_kv_expired_reqs, |
| 117 | } |
| 118 | |
| 119 | xfer_time = np.asarray(self.data["transfer_duration"]) |
| 120 | mb = np.asarray(self.data["bytes_transferred"]) / 2**20 |
| 121 | descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32) |
| 122 | n = len(descs) |
| 123 | assert n == self.num_successful_transfers |
| 124 | |
| 125 | total_mb = mb.sum() |
| 126 | avg_mb = total_mb / n |
| 127 | total_time_seconds = xfer_time.sum() |
| 128 | throughput_mb_s = ( |
| 129 | (total_mb / total_time_seconds).item() if total_time_seconds > 0 else 0.0 |
| 130 | ) |
| 131 | |
| 132 | return { |
| 133 | "Num successful transfers": n, |
| 134 | "Avg xfer time (ms)": round(xfer_time.mean().item() * 1e3, 3), |
| 135 | "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), |
| 136 | "Avg MB per transfer": round(avg_mb.item(), 3), |
| 137 | "Throughput (MB/s)": round(throughput_mb_s, 3), |
| 138 | "Avg number of descriptors": round(descs.mean().item(), 1), |
| 139 | "Num failed transfers": num_failed_transfers, |
| 140 | "Num failed recvs": num_failed_recvs, |
| 141 | "Num KV expired reqs": num_kv_expired_reqs, |
| 142 | } |
| 143 | |
| 144 | @property |
| 145 | def num_successful_transfers(self) -> int: |
| 146 | return len(self.data["transfer_duration"]) |