From bee126165fb39e770030e9c3b4a43aee4acd3175 Mon Sep 17 00:00:00 2001 From: Zhewen Li Date: Tue, 5 May 2026 02:17:38 -0700 Subject: [PATCH] [P/D][Mooncake] Add KVConnectorStats for transfer observability (#40414) Signed-off-by: Zhewen Li Co-authored-by: Zhewen Li --- .../kv_connector/unit/test_mooncake_stats.py | 281 ++++++++++++++++++ .../v1/mooncake/mooncake_connector.py | 53 +++- .../kv_connector/v1/mooncake/stats.py | 146 +++++++++ 3 files changed, 477 insertions(+), 3 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_mooncake_stats.py create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py diff --git a/tests/v1/kv_connector/unit/test_mooncake_stats.py b/tests/v1/kv_connector/unit/test_mooncake_stats.py new file mode 100644 index 0000000000..a20fcb5053 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_stats.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from unittest.mock import MagicMock + +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import ( + MooncakeConnector, + MooncakeConnectorWorker, + SendBlockMeta, +) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.stats import ( + MooncakeKVConnectorStats, +) + + +def test_is_empty_on_fresh_stats(): + stats = MooncakeKVConnectorStats() + assert stats.is_empty() + assert stats.num_successful_transfers == 0 + + +def test_record_transfer_and_reduce(): + stats = MooncakeKVConnectorStats() + # 1 MB transfer in 1 ms -> 1000 MB/s throughput + stats.record_transfer(duration_s=0.001, total_bytes=1 * 2**20, num_descs=4) + # 2 MB transfer in 2 ms + stats.record_transfer(duration_s=0.002, total_bytes=2 * 2**20, num_descs=6) + assert not stats.is_empty() + assert stats.num_successful_transfers == 2 + + reduced = stats.reduce() + assert reduced["Num successful transfers"] == 2 + # avg = (1 + 2) / 2 = 1.5 ms + assert reduced["Avg xfer time (ms)"] == 1.5 + assert reduced["Avg MB per transfer"] == 1.5 + # 3 MB total / 3 ms total = 1000 MB/s + assert reduced["Throughput (MB/s)"] == 1000.0 + assert reduced["Avg number of descriptors"] == 5.0 + assert reduced["Num failed transfers"] == 0 + assert reduced["Num failed recvs"] == 0 + assert reduced["Num KV expired reqs"] == 0 + + +def test_record_failures_keeps_stats_non_empty(): + stats = MooncakeKVConnectorStats() + stats.record_failed_transfer() + stats.record_failed_recv() + stats.record_kv_expired_req() + assert not stats.is_empty() + + reduced = stats.reduce() + # No successful transfers -> latency/throughput all zero, but failure + # counters still surface. + assert reduced["Num successful transfers"] == 0 + assert reduced["Num failed transfers"] == 1 + assert reduced["Num failed recvs"] == 1 + assert reduced["Num KV expired reqs"] == 1 + + +def test_aggregate_sums_observations(): + a = MooncakeKVConnectorStats() + b = MooncakeKVConnectorStats() + a.record_transfer(duration_s=0.001, total_bytes=1 * 2**20, num_descs=1) + b.record_transfer(duration_s=0.002, total_bytes=2 * 2**20, num_descs=2) + b.record_failed_transfer() + + a.aggregate(b) + + assert a.num_successful_transfers == 2 + reduced = a.reduce() + assert reduced["Num successful transfers"] == 2 + assert reduced["Num failed transfers"] == 1 + + +def test_aggregate_with_empty_other_is_noop(): + a = MooncakeKVConnectorStats() + a.record_transfer(duration_s=0.001, total_bytes=1, num_descs=1) + b = MooncakeKVConnectorStats() + + a.aggregate(b) + + assert a.num_successful_transfers == 1 + + +def test_getstate_drops_lock_and_setstate_recreates_it(): + # KVConnectorStats subclasses must be picklable (worker→scheduler IPC), + # but threading.Lock isn't — so __getstate__ strips it and __setstate__ + # rebuilds a fresh per-process lock. + original = MooncakeKVConnectorStats() + original.record_transfer(duration_s=0.01, total_bytes=2048, num_descs=3) + + state = original.__getstate__() + assert "_lock" not in state + + rebuilt = MooncakeKVConnectorStats.__new__(MooncakeKVConnectorStats) + rebuilt.__setstate__(state) + assert rebuilt.data == original.data + # Lock works on the receiver side. + rebuilt.record_transfer(duration_s=0.02, total_bytes=4096, num_descs=5) + assert rebuilt.num_successful_transfers == 2 + + +def test_concurrent_writers_keep_row_lengths_aligned(): + # Multiple writers + a snapshot reader must never produce a snapshot + # with mismatched column lengths — reduce()'s + # len(descs) == num_successful_transfers assertion would fire. + stats = MooncakeKVConnectorStats() + stop = threading.Event() + writer_count = 4 + snapshots: list[MooncakeKVConnectorStats] = [] + + def writer(): + i = 0 + while not stop.is_set(): + stats.record_transfer( + duration_s=0.001 + i * 1e-9, + total_bytes=1024 + i, + num_descs=1 + (i % 8), + ) + i += 1 + + def snapper(): + while not stop.is_set(): + snap = stats.clone_and_reset() + if not snap.is_empty(): + # Force the same path the logger walks; reduce() will + # blow up on torn rows via its internal assert. + snap.reduce() + snapshots.append(snap) + + threads = [threading.Thread(target=writer) for _ in range(writer_count)] + snapshotter = threading.Thread(target=snapper) + for t in threads: + t.start() + snapshotter.start() + # Short fixed window — long enough to interleave thousands of ops. + threading.Event().wait(0.2) + stop.set() + for t in threads: + t.join() + snapshotter.join() + + # Final drain so we don't lose the in-flight tail. + final = stats.clone_and_reset() + if not final.is_empty(): + final.reduce() + snapshots.append(final) + + # Every snapshot's columns must have identical lengths (the invariant + # the lock protects), and the union must contain at least one row. + total_rows = 0 + for snap in snapshots: + n = len(snap.data["transfer_duration"]) + assert len(snap.data["bytes_transferred"]) == n + assert len(snap.data["num_descriptors"]) == n + total_rows += n + assert total_rows > 0 + + +def test_clone_and_reset_hands_off_old_data(): + stats = MooncakeKVConnectorStats() + stats.record_transfer(duration_s=0.001, total_bytes=1, num_descs=1) + stats.record_failed_recv() + + snapshot = stats.clone_and_reset() + + assert snapshot.num_successful_transfers == 1 + assert not snapshot.is_empty() + # Original is now empty. + assert stats.is_empty() + assert stats.num_successful_transfers == 0 + # Recording on the original does not mutate the snapshot. + stats.record_transfer(duration_s=0.005, total_bytes=2, num_descs=2) + assert snapshot.num_successful_transfers == 1 + + +def test_build_kv_connector_stats_none_returns_empty_instance(): + out = MooncakeConnector.build_kv_connector_stats() + assert isinstance(out, MooncakeKVConnectorStats) + assert out.is_empty() + + +def test_build_kv_connector_stats_with_data_round_trips(): + original = MooncakeKVConnectorStats() + original.record_transfer(duration_s=0.01, total_bytes=1024, num_descs=3) + original.record_failed_transfer() + + # Serialized form is the .data dict; build should reconstruct an instance + # that behaves the same. + rebuilt = MooncakeConnector.build_kv_connector_stats(data=original.data) + + assert isinstance(rebuilt, MooncakeKVConnectorStats) + assert rebuilt.num_successful_transfers == 1 + assert rebuilt.reduce()["Num failed transfers"] == 1 + + +def _bare_worker() -> MooncakeConnectorWorker: + """Construct a MooncakeConnectorWorker skipping __init__ (full init requires + a live TransferEngine). Only the attributes touched by the methods under + test are populated; role flags and async_zmq_ctx keep __del__'s shutdown + path a no-op.""" + worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) + worker.xfer_stats = MooncakeKVConnectorStats() + worker.engine = MagicMock() + worker.async_zmq_ctx = MagicMock() + worker.is_kv_consumer = True + worker.is_kv_producer = True + return worker + + +def test_send_blocks_records_success(): + worker = _bare_worker() + worker.engine.batch_transfer_sync_write.return_value = 0 + + ret = worker._send_blocks( + "host:1234", + src_ptrs=[0x1000, 0x2000], + dst_ptrs=[0x3000, 0x4000], + lengths=[1024, 2048], + ) + + assert ret == 0 + assert worker.xfer_stats.num_successful_transfers == 1 + data = worker.xfer_stats.data + assert data["bytes_transferred"] == [1024 + 2048] + assert data["num_descriptors"] == [2] + assert data["num_failed_transfers"] == [] + + +def test_send_blocks_records_failure(): + worker = _bare_worker() + worker.engine.batch_transfer_sync_write.return_value = 1 # non-zero = fail + + ret = worker._send_blocks("host:1234", [0x1000], [0x2000], [4096]) + + assert ret == 1 + assert worker.xfer_stats.num_successful_transfers == 0 + assert worker.xfer_stats.data["num_failed_transfers"] == [1] + + +def test_get_kv_connector_stats_returns_none_when_empty(): + worker = _bare_worker() + + assert worker.get_kv_connector_stats() is None + + +def test_get_kv_connector_stats_returns_and_resets(): + worker = _bare_worker() + worker.engine.batch_transfer_sync_write.return_value = 0 + worker._send_blocks("host:1234", [0x1000], [0x2000], [4096]) + + snapshot = worker.get_kv_connector_stats() + assert isinstance(snapshot, MooncakeKVConnectorStats) + assert snapshot.num_successful_transfers == 1 + + # Second call returns None because the worker's stats were reset. + assert worker.get_kv_connector_stats() is None + + +def test_expired_request_bumps_counter(): + import asyncio + + worker = _bare_worker() + worker.reqs_need_send = { + "tid1": SendBlockMeta( + p_req_id="req1", + transfer_id="tid1", + local_block_ids=[0, 1], + ready=asyncio.Event(), + expire_time=-1.0, # Already expired. + sending=0, + ), + } + worker.finished_sending_reqs = set() + + asyncio.run(worker.fetch_finished_sending_reqs()) + + assert worker.xfer_stats.data["num_kv_expired_reqs"] == [1] + # Expired transfer also cleaned out of reqs_need_send. + assert "tid1" not in worker.reqs_need_send diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py index 715fcbde16..608fd87847 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py @@ -31,10 +31,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorRole, SupportsHMA, ) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( MooncakeBootstrapServer, RegisterWorkerPayload, ) +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.stats import ( + MooncakeKVConnectorStats, +) from vllm.distributed.parallel_state import ( get_pp_group, get_tensor_model_parallel_rank, @@ -457,6 +461,25 @@ class MooncakeConnector(KVConnectorBase_V1, SupportsHMA): def wait_for_save(self): pass + def get_kv_connector_stats(self) -> KVConnectorStats | None: + """Return worker-local transfer stats since the last call. + + Note the P/D asymmetry: because Mooncake is P-push (P calls + batch_transfer_sync_write), P records successful transfer latency, + bytes, and descriptor counts, while D only records failures + (recv/ZMQ errors). Aggregated NIXL-style dashboards will find + successful-transfer metrics on the P worker, not D. + """ + if self.connector_worker is None: + return None + return self.connector_worker.get_kv_connector_stats() + + @classmethod + def build_kv_connector_stats( + cls, data: dict[str, Any] | None = None + ) -> KVConnectorStats | None: + return MooncakeKVConnectorStats(data=data or {}) + class MooncakeConnectorScheduler: """Implementation of Scheduler side methods""" @@ -816,6 +839,8 @@ class MooncakeConnectorWorker: self.finished_sending_reqs: set[ReqId] = set() self.finished_recving_reqs: set[ReqId] = set() + self.xfer_stats = MooncakeKVConnectorStats() + self.block_size = vllm_config.cache_config.block_size self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config @@ -1340,11 +1365,23 @@ class MooncakeConnectorWorker: ret_value = self.engine.batch_transfer_sync_write( remote_session, src_ptrs, dst_ptrs, lengths ) + duration = time.perf_counter() - start_time if ret_value == 0: - logger.debug( - "Sending to %s done, took %s", + self.xfer_stats.record_transfer( + duration_s=duration, + total_bytes=sum(lengths), + num_descs=len(src_ptrs), + ) + logger.debug("Sending to %s done, took %s", remote_session, duration) + else: + self.xfer_stats.record_failed_transfer() + logger.warning( + "Sending to %s failed (ret=%s) after %s (%d descriptors, %d bytes)", remote_session, - time.perf_counter() - start_time, + ret_value, + duration, + len(src_ptrs), + sum(lengths), ) return ret_value @@ -1445,6 +1482,7 @@ class MooncakeConnectorWorker: send_meta.p_req_id, envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT, ) + self.xfer_stats.record_kv_expired_req() finished_sending_reqs.add(send_meta.p_req_id) expired_transfer_id.append(transfer_id) @@ -1485,6 +1523,13 @@ class MooncakeConnectorWorker: return finished_sending_reqs or None, finished_recving_reqs or None + def get_kv_connector_stats(self) -> KVConnectorStats | None: + """Return transfer stats collected since the last call, or None + if nothing has been recorded in this interval.""" + if self.xfer_stats.is_empty(): + return None + return self.xfer_stats.clone_and_reset() + async def receive_kv_from_single_worker( self, worker_addr: str, @@ -1531,6 +1576,7 @@ class MooncakeConnectorWorker: req_ids, response.err_msg, ) + self.xfer_stats.record_failed_recv() return self.process_pulling_result(response, pull_metas) if response.status == MooncakeXferResponseStatus.FINISH: @@ -1539,6 +1585,7 @@ class MooncakeConnectorWorker: logger.debug("ZMQ context terminated, exiting Mooncake receiver thread.") except Exception as e: logger.error("MooncakeXferMetadata transfer failed for %s: %s", req_ids, e) + self.xfer_stats.record_failed_recv() return def process_pulling_result( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py new file mode 100644 index 0000000000..d177f55cc7 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/stats.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stats container for the Mooncake connector.""" + +import threading +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorStats, +) + +# TODO(mooncake-stats): add MooncakePromMetrics (mirror NixlPromMetrics) +# and wire it via MooncakeConnector.build_prom_metrics in a follow-up PR. + + +@dataclass +class MooncakeKVConnectorStats(KVConnectorStats): + """Container for Mooncake KV transfer performance metrics. + + `_lock` serializes record_* against clone_and_reset so each row's + appends are atomic and column lengths stay aligned. Writers run on + the sender pool / receiver loop / sender loop; reader runs on the + main worker thread. + """ + + def __post_init__(self): + self._lock = threading.Lock() + if not self.data: + self.reset() + + # threading.Lock is not picklable; strip it from the wire form and + # rebuild a fresh per-process lock on the receiver side. + def __getstate__(self) -> dict[str, Any]: + state = self.__dict__.copy() + state.pop("_lock", None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__.update(state) + self._lock = threading.Lock() + + def reset(self): + self.data: dict[str, list[float | int]] = { + "transfer_duration": [], + "bytes_transferred": [], + "num_descriptors": [], + "num_failed_transfers": [], + "num_failed_recvs": [], + "num_kv_expired_reqs": [], + } + + def record_transfer(self, duration_s: float, total_bytes: int, num_descs: int): + with self._lock: + self.data["transfer_duration"].append(duration_s) + self.data["bytes_transferred"].append(total_bytes) + self.data["num_descriptors"].append(num_descs) + + # Failure counters store a list of 1s so a future Prom counter can iterate + # with .inc(list_item), mirroring NIXL's NixlPromMetrics.observe. + def record_failed_transfer(self): + with self._lock: + self.data["num_failed_transfers"].append(1) + + def record_failed_recv(self): + with self._lock: + self.data["num_failed_recvs"].append(1) + + def record_kv_expired_req(self): + with self._lock: + self.data["num_kv_expired_reqs"].append(1) + + def clone_and_reset(self) -> "MooncakeKVConnectorStats": + # Copy lists under the lock for length alignment; return a fresh + # instance so the snapshot has its own _lock. + with self._lock: + snapshot_data: dict[str, list[float | int]] = { + k: list(v) for k, v in self.data.items() + } + self.reset() + return MooncakeKVConnectorStats(data=snapshot_data) + + def is_empty(self) -> bool: + return ( + self.num_successful_transfers == 0 + and len(self.data["num_failed_transfers"]) == 0 + and len(self.data["num_failed_recvs"]) == 0 + and len(self.data["num_kv_expired_reqs"]) == 0 + ) + + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: + if not other.is_empty(): + for k, v in other.data.items(): + accumulator = self.data[k] + assert isinstance(accumulator, list) + accumulator.extend(v) + return self + + def reduce(self) -> dict[str, int | float]: + num_failed_transfers = len(self.data["num_failed_transfers"]) + num_failed_recvs = len(self.data["num_failed_recvs"]) + num_kv_expired_reqs = len(self.data["num_kv_expired_reqs"]) + + if self.num_successful_transfers == 0: + return { + "Num successful transfers": 0, + "Avg xfer time (ms)": 0, + "P90 xfer time (ms)": 0, + "Avg MB per transfer": 0, + "Throughput (MB/s)": 0, + "Avg number of descriptors": 0, + "Num failed transfers": num_failed_transfers, + "Num failed recvs": num_failed_recvs, + "Num KV expired reqs": num_kv_expired_reqs, + } + + xfer_time = np.asarray(self.data["transfer_duration"]) + mb = np.asarray(self.data["bytes_transferred"]) / 2**20 + descs = np.asarray(self.data["num_descriptors"], dtype=np.uint32) + n = len(descs) + assert n == self.num_successful_transfers + + total_mb = mb.sum() + avg_mb = total_mb / n + total_time_seconds = xfer_time.sum() + throughput_mb_s = ( + total_mb / total_time_seconds if total_time_seconds > 0 else 0.0 + ) + + return { + "Num successful transfers": n, + "Avg xfer time (ms)": round(xfer_time.mean() * 1e3, 3), + "P90 xfer time (ms)": round(np.percentile(xfer_time, 90).item() * 1e3, 3), + "Avg MB per transfer": round(avg_mb, 3), + "Throughput (MB/s)": round(throughput_mb_s, 3), + "Avg number of descriptors": round(descs.mean(), 1), + "Num failed transfers": num_failed_transfers, + "Num failed recvs": num_failed_recvs, + "Num KV expired reqs": num_kv_expired_reqs, + } + + @property + def num_successful_transfers(self) -> int: + return len(self.data["transfer_duration"])