| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import threading |
| 5 | from unittest.mock import MagicMock |
| 6 | |
| 7 | import numpy as np |
| 8 | |
| 9 | from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import ( |
| 10 | MooncakeConnector, |
| 11 | MooncakeConnectorWorker, |
| 12 | SendBlockMeta, |
| 13 | ) |
| 14 | from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.stats import ( |
| 15 | MooncakeKVConnectorStats, |
| 16 | ) |
| 17 | |
| 18 | |
| 19 | def test_is_empty_on_fresh_stats(): |
| 20 | stats = MooncakeKVConnectorStats() |
| 21 | assert stats.is_empty() |
| 22 | assert stats.num_successful_transfers == 0 |
| 23 | |
| 24 | |
| 25 | def test_record_transfer_and_reduce(): |
| 26 | stats = MooncakeKVConnectorStats() |
| 27 | # 1 MB transfer in 1 ms -> 1000 MB/s throughput |
| 28 | stats.record_transfer(duration_s=0.001, total_bytes=1 * 2**20, num_descs=4) |
| 29 | # 2 MB transfer in 2 ms |
| 30 | stats.record_transfer(duration_s=0.002, total_bytes=2 * 2**20, num_descs=6) |
| 31 | assert not stats.is_empty() |
| 32 | assert stats.num_successful_transfers == 2 |
| 33 | |
| 34 | reduced = stats.reduce() |
| 35 | assert reduced["Num successful transfers"] == 2 |
| 36 | # avg = (1 + 2) / 2 = 1.5 ms |
| 37 | assert reduced["Avg xfer time (ms)"] == 1.5 |
| 38 | assert reduced["Avg MB per transfer"] == 1.5 |
| 39 | # 3 MB total / 3 ms total = 1000 MB/s |
| 40 | assert reduced["Throughput (MB/s)"] == 1000.0 |
| 41 | assert reduced["Avg number of descriptors"] == 5.0 |
| 42 | assert reduced["Num failed transfers"] == 0 |
| 43 | assert reduced["Num failed recvs"] == 0 |
| 44 | assert reduced["Num KV expired reqs"] == 0 |
| 45 | # Reduced values must be plain Python scalars so CLI logging renders |
| 46 | # them without numpy reprs (eg np.float64(...)). |
| 47 | assert all(not isinstance(v, np.generic) for v in reduced.values()) |
| 48 | |
| 49 | |
| 50 | def test_record_failures_keeps_stats_non_empty(): |
| 51 | stats = MooncakeKVConnectorStats() |
| 52 | stats.record_failed_transfer() |
| 53 | stats.record_failed_recv() |
| 54 | stats.record_kv_expired_req() |
| 55 | assert not stats.is_empty() |
| 56 | |
| 57 | reduced = stats.reduce() |
| 58 | # No successful transfers -> latency/throughput all zero, but failure |
| 59 | # counters still surface. |
| 60 | assert reduced["Num successful transfers"] == 0 |
| 61 | assert reduced["Num failed transfers"] == 1 |
| 62 | assert reduced["Num failed recvs"] == 1 |
| 63 | assert reduced["Num KV expired reqs"] == 1 |
| 64 | |
| 65 | |
| 66 | def test_aggregate_sums_observations(): |
| 67 | a = MooncakeKVConnectorStats() |
| 68 | b = MooncakeKVConnectorStats() |
| 69 | a.record_transfer(duration_s=0.001, total_bytes=1 * 2**20, num_descs=1) |
| 70 | b.record_transfer(duration_s=0.002, total_bytes=2 * 2**20, num_descs=2) |
| 71 | b.record_failed_transfer() |
| 72 | |
| 73 | a.aggregate(b) |
| 74 | |
| 75 | assert a.num_successful_transfers == 2 |
| 76 | reduced = a.reduce() |
| 77 | assert reduced["Num successful transfers"] == 2 |
| 78 | assert reduced["Num failed transfers"] == 1 |
| 79 | |
| 80 | |
| 81 | def test_aggregate_with_empty_other_is_noop(): |
| 82 | a = MooncakeKVConnectorStats() |
| 83 | a.record_transfer(duration_s=0.001, total_bytes=1, num_descs=1) |
| 84 | b = MooncakeKVConnectorStats() |
| 85 | |
| 86 | a.aggregate(b) |
| 87 | |
| 88 | assert a.num_successful_transfers == 1 |
| 89 | |
| 90 | |
| 91 | def test_getstate_drops_lock_and_setstate_recreates_it(): |
| 92 | # KVConnectorStats subclasses must be picklable (worker→scheduler IPC), |
| 93 | # but threading.Lock isn't — so __getstate__ strips it and __setstate__ |
| 94 | # rebuilds a fresh per-process lock. |
| 95 | original = MooncakeKVConnectorStats() |
| 96 | original.record_transfer(duration_s=0.01, total_bytes=2048, num_descs=3) |
| 97 | |
| 98 | state = original.__getstate__() |
| 99 | assert "_lock" not in state |
| 100 | |
| 101 | rebuilt = MooncakeKVConnectorStats.__new__(MooncakeKVConnectorStats) |
| 102 | rebuilt.__setstate__(state) |
| 103 | assert rebuilt.data == original.data |
| 104 | # Lock works on the receiver side. |
| 105 | rebuilt.record_transfer(duration_s=0.02, total_bytes=4096, num_descs=5) |
| 106 | assert rebuilt.num_successful_transfers == 2 |
| 107 | |
| 108 | |
| 109 | def test_concurrent_writers_keep_row_lengths_aligned(): |
| 110 | # Multiple writers + a snapshot reader must never produce a snapshot |
| 111 | # with mismatched column lengths — reduce()'s |
| 112 | # len(descs) == num_successful_transfers assertion would fire. |
| 113 | stats = MooncakeKVConnectorStats() |
| 114 | stop = threading.Event() |
| 115 | writer_count = 4 |
| 116 | snapshots: list[MooncakeKVConnectorStats] = [] |
| 117 | |
| 118 | def writer(): |
| 119 | i = 0 |
| 120 | while not stop.is_set(): |
| 121 | stats.record_transfer( |
| 122 | duration_s=0.001 + i * 1e-9, |
| 123 | total_bytes=1024 + i, |
| 124 | num_descs=1 + (i % 8), |
| 125 | ) |
| 126 | i += 1 |
| 127 | |
| 128 | def snapper(): |
| 129 | while not stop.is_set(): |
| 130 | snap = stats.clone_and_reset() |
| 131 | if not snap.is_empty(): |
| 132 | # Force the same path the logger walks; reduce() will |
| 133 | # blow up on torn rows via its internal assert. |
| 134 | snap.reduce() |
| 135 | snapshots.append(snap) |
| 136 | |
| 137 | threads = [threading.Thread(target=writer) for _ in range(writer_count)] |
| 138 | snapshotter = threading.Thread(target=snapper) |
| 139 | for t in threads: |
| 140 | t.start() |
| 141 | snapshotter.start() |
| 142 | # Short fixed window — long enough to interleave thousands of ops. |
| 143 | threading.Event().wait(0.2) |
| 144 | stop.set() |
| 145 | for t in threads: |
| 146 | t.join() |
| 147 | snapshotter.join() |
| 148 | |
| 149 | # Final drain so we don't lose the in-flight tail. |
| 150 | final = stats.clone_and_reset() |
| 151 | if not final.is_empty(): |
| 152 | final.reduce() |
| 153 | snapshots.append(final) |
| 154 | |
| 155 | # Every snapshot's columns must have identical lengths (the invariant |
| 156 | # the lock protects), and the union must contain at least one row. |
| 157 | total_rows = 0 |
| 158 | for snap in snapshots: |
| 159 | n = len(snap.data["transfer_duration"]) |
| 160 | assert len(snap.data["bytes_transferred"]) == n |
| 161 | assert len(snap.data["num_descriptors"]) == n |
| 162 | total_rows += n |
| 163 | assert total_rows > 0 |
| 164 | |
| 165 | |
| 166 | def test_clone_and_reset_hands_off_old_data(): |
| 167 | stats = MooncakeKVConnectorStats() |
| 168 | stats.record_transfer(duration_s=0.001, total_bytes=1, num_descs=1) |
| 169 | stats.record_failed_recv() |
| 170 | |
| 171 | snapshot = stats.clone_and_reset() |
| 172 | |
| 173 | assert snapshot.num_successful_transfers == 1 |
| 174 | assert not snapshot.is_empty() |
| 175 | # Original is now empty. |
| 176 | assert stats.is_empty() |
| 177 | assert stats.num_successful_transfers == 0 |
| 178 | # Recording on the original does not mutate the snapshot. |
| 179 | stats.record_transfer(duration_s=0.005, total_bytes=2, num_descs=2) |
| 180 | assert snapshot.num_successful_transfers == 1 |
| 181 | |
| 182 | |
| 183 | def test_build_kv_connector_stats_none_returns_empty_instance(): |
| 184 | out = MooncakeConnector.build_kv_connector_stats() |
| 185 | assert isinstance(out, MooncakeKVConnectorStats) |
| 186 | assert out.is_empty() |
| 187 | |
| 188 | |
| 189 | def test_build_kv_connector_stats_with_data_round_trips(): |
| 190 | original = MooncakeKVConnectorStats() |
| 191 | original.record_transfer(duration_s=0.01, total_bytes=1024, num_descs=3) |
| 192 | original.record_failed_transfer() |
| 193 | |
| 194 | # Serialized form is the .data dict; build should reconstruct an instance |
| 195 | # that behaves the same. |
| 196 | rebuilt = MooncakeConnector.build_kv_connector_stats(data=original.data) |
| 197 | |
| 198 | assert isinstance(rebuilt, MooncakeKVConnectorStats) |
| 199 | assert rebuilt.num_successful_transfers == 1 |
| 200 | assert rebuilt.reduce()["Num failed transfers"] == 1 |
| 201 | |
| 202 | |
| 203 | def _bare_worker() -> MooncakeConnectorWorker: |
| 204 | """Construct a MooncakeConnectorWorker skipping __init__ (full init requires |
| 205 | a live TransferEngine). Only the attributes touched by the methods under |
| 206 | test are populated; role flags and async_zmq_ctx keep __del__'s shutdown |
| 207 | path a no-op.""" |
| 208 | worker = MooncakeConnectorWorker.__new__(MooncakeConnectorWorker) |
| 209 | worker.xfer_stats = MooncakeKVConnectorStats() |
| 210 | worker.engine = MagicMock() |
| 211 | worker.async_zmq_ctx = MagicMock() |
| 212 | worker.is_kv_consumer = True |
| 213 | worker.is_kv_producer = True |
| 214 | return worker |
| 215 | |
| 216 | |
| 217 | def test_send_blocks_records_success(): |
| 218 | worker = _bare_worker() |
| 219 | worker.engine.batch_transfer_sync_write.return_value = 0 |
| 220 | |
| 221 | ret = worker._send_blocks( |
| 222 | "host:1234", |
| 223 | src_ptrs=[0x1000, 0x2000], |
| 224 | dst_ptrs=[0x3000, 0x4000], |
| 225 | lengths=[1024, 2048], |
| 226 | ) |
| 227 | |
| 228 | assert ret == 0 |
| 229 | assert worker.xfer_stats.num_successful_transfers == 1 |
| 230 | data = worker.xfer_stats.data |
| 231 | assert data["bytes_transferred"] == [1024 + 2048] |
| 232 | assert data["num_descriptors"] == [2] |
| 233 | assert data["num_failed_transfers"] == [] |
| 234 | |
| 235 | |
| 236 | def test_send_blocks_records_failure(): |
| 237 | worker = _bare_worker() |
| 238 | worker.engine.batch_transfer_sync_write.return_value = 1 # non-zero = fail |
| 239 | |
| 240 | ret = worker._send_blocks("host:1234", [0x1000], [0x2000], [4096]) |
| 241 | |
| 242 | assert ret == 1 |
| 243 | assert worker.xfer_stats.num_successful_transfers == 0 |
| 244 | assert worker.xfer_stats.data["num_failed_transfers"] == [1] |
| 245 | |
| 246 | |
| 247 | def test_get_kv_connector_stats_returns_none_when_empty(): |
| 248 | worker = _bare_worker() |
| 249 | |
| 250 | assert worker.get_kv_connector_stats() is None |
| 251 | |
| 252 | |
| 253 | def test_get_kv_connector_stats_returns_and_resets(): |
| 254 | worker = _bare_worker() |
| 255 | worker.engine.batch_transfer_sync_write.return_value = 0 |
| 256 | worker._send_blocks("host:1234", [0x1000], [0x2000], [4096]) |
| 257 | |
| 258 | snapshot = worker.get_kv_connector_stats() |
| 259 | assert isinstance(snapshot, MooncakeKVConnectorStats) |
| 260 | assert snapshot.num_successful_transfers == 1 |
| 261 | |
| 262 | # Second call returns None because the worker's stats were reset. |
| 263 | assert worker.get_kv_connector_stats() is None |
| 264 | |
| 265 | |
| 266 | def test_expired_request_bumps_counter(): |
| 267 | import asyncio |
| 268 | |
| 269 | worker = _bare_worker() |
| 270 | worker.reqs_need_send = { |
| 271 | "tid1": SendBlockMeta( |
| 272 | p_req_id="req1", |
| 273 | transfer_id="tid1", |
| 274 | local_block_ids=[0, 1], |
| 275 | ready=asyncio.Event(), |
| 276 | expire_time=-1.0, # Already expired. |
| 277 | sending=0, |
| 278 | ), |
| 279 | } |
| 280 | worker.finished_sending_reqs = set() |
| 281 | |
| 282 | asyncio.run(worker.fetch_finished_sending_reqs()) |
| 283 | |
| 284 | assert worker.xfer_stats.data["num_kv_expired_reqs"] == [1] |
| 285 | # Expired transfer also cleaned out of reqs_need_send. |
| 286 | assert "tid1" not in worker.reqs_need_send |