A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 22:00 UTC
[Mooncake] Add metrics for MooncakeStoreConnector operations (#43392)
Dao007forever committed
May 23, 2026 at 13:34 UTC
819c610f9bd0f473573c4a9826b229aebad8b301
5 files changed
+507
-4
tests/v1/kv_connector/unit/test_mooncake_store_connector.py
+46
@@ -14,6 +14,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store import (
14
from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
15
MooncakeStoreConnectorMetadata,
16
)
17
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
18
+ MooncakeStoreConnectorStats,
19
+)
20
from vllm.v1.kv_cache_interface import (
21
FullAttentionSpec,
22
KVCacheConfig,
@@ -129,6 +132,49 @@ def test_get_kv_connector_kv_cache_events_returns_none_when_empty():
132
assert connector.get_kv_connector_kv_cache_events() is None
133
134
135
+def test_get_kv_connector_stats_delegates_to_worker():
136
+ vllm_config = _make_vllm_config()
137
+ kv_cache_config = _make_kv_cache_config()
138
+ expected_stats = MooncakeStoreConnectorStats()
139
+ expected_stats.record_operation("save_put", 0.01, 2, num_bytes=1024)
140
+
141
+ with (
142
+ set_current_vllm_config(vllm_config),
143
+ patch(
144
+ "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
145
+ "connector.MooncakeStoreWorker"
146
+ ) as mock_worker_cls,
147
+ ):
148
+ connector = mooncake_store_connector.MooncakeStoreConnector(
149
+ vllm_config, KVConnectorRole.WORKER, kv_cache_config
150
+ )
151
+
152
+ mock_worker_cls.return_value.get_kv_connector_stats.return_value = expected_stats
153
+ stats = connector.get_kv_connector_stats()
154
+
155
+ assert stats is expected_stats
156
+ mock_worker_cls.return_value.get_kv_connector_stats.assert_called_once_with()
157
+
158
+
159
+def test_build_kv_connector_stats_reconstructs_mooncake_stats():
160
+ stats = mooncake_store_connector.MooncakeStoreConnector.build_kv_connector_stats(
161
+ {
162
+ "save_put": [
163
+ {
164
+ "duration_seconds": 0.02,
165
+ "num_keys": 4,
166
+ "num_bytes": 2048,
167
+ "status": "ok",
168
+ "num_failed_keys": 0,
169
+ }
170
+ ]
171
+ }
172
+ )
173
+
174
+ assert isinstance(stats, MooncakeStoreConnectorStats)
175
+ assert stats.data["save_put"][0]["num_bytes"] == 2048
176
+
177
+
178
def test_get_kv_connector_kv_cache_events_wraps_worker_events():
179
vllm_config = _make_vllm_config()
180
kv_cache_config = _make_kv_cache_config()
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+102
@@ -28,6 +28,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
28
LoadSpec,
29
ReqMeta,
30
)
31
+from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import (
32
+ MooncakeStoreConnectorStats,
33
+)
34
35
36
def _default_send_coord() -> mooncake_store_worker.MooncakeStoreCoordinator:
@@ -417,6 +420,24 @@ def test_store_sending_thread_skips_request_during_cpu_pressure():
420
assert store.batch_put_from_multi_buffers.call_count == 3
421
422
423
+def test_store_sending_thread_records_mooncake_metrics():
424
+ store = MagicMock()
425
+ store.batch_is_exist.return_value = [0, 0]
426
+ store.batch_put_from_multi_buffers.return_value = [256, 256]
427
+ thread = _make_store_sending_thread(store)
428
+ stats = MooncakeStoreConnectorStats()
429
+ thread._record_operation_cb = stats.record_operation
430
+
431
+ thread.add_stored_request("req-a")
432
+ thread._handle_request(_make_store_req("req-a", [b"a0", b"a1"]))
433
+
434
+ assert len(stats.data["save_exists"]) == 1
435
+ assert stats.data["save_exists"][0]["num_keys"] == 2
436
+ assert len(stats.data["save_put"]) == 1
437
+ assert stats.data["save_put"][0]["num_bytes"] == 512
438
+ assert stats.data["save_put"][0]["status"] == "ok"
439
+
440
+
441
def test_store_sending_thread_only_skips_on_no_available_handle():
442
store = MagicMock()
443
store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)
@@ -548,6 +569,29 @@ def test_recv_thread_logs_tier_summary_when_enabled(monkeypatch, caplog_vllm):
569
)
570
571
572
+def test_recv_thread_records_partial_failure_metrics(monkeypatch):
573
+ monkeypatch.delenv("VLLM_MOONCAKE_STORE_TIER_LOG", raising=False)
574
+ store = MagicMock()
575
+ store.batch_get_into_multi_buffers.return_value = [256, -10]
576
+ thread = _make_store_recving_thread(store, disk_offload_buffer_budget_bytes=None)
577
+ stats = MooncakeStoreConnectorStats()
578
+ thread._record_operation_cb = stats.record_operation
579
+
580
+ req = _make_load_req(
581
+ "req-a",
582
+ [b"a0", b"a1"],
583
+ token_len=32,
584
+ )
585
+
586
+ thread._handle_request(req)
587
+
588
+ assert len(stats.data["load_get"]) == 1
589
+ assert stats.data["load_get"][0]["num_keys"] == 2
590
+ assert stats.data["load_get"][0]["num_bytes"] == 512
591
+ assert stats.data["load_get"][0]["status"] == "partial_failure"
592
+ assert stats.data["load_get"][0]["num_failed_keys"] == 1
593
+
594
+
595
def test_recv_thread_uses_ratio_scaled_budget_for_first_pass_split():
596
store = MagicMock()
597
store.batch_get_into_multi_buffers.side_effect = [
@@ -1025,6 +1069,8 @@ def _make_bare_worker(
1069
1070
worker.disk_offload_buffer_budget_bytes = None
1071
worker.store_replicate_config = SimpleNamespace()
1072
+ worker._kv_connector_stats_lock = threading.Lock()
1073
+ worker.kv_connector_stats = MooncakeStoreConnectorStats()
1074
1075
spec = FullAttentionSpec(
1076
block_size=block_size, num_kv_heads=8, head_size=64, dtype=None
@@ -1373,3 +1419,59 @@ def test_topology_embedded_cpu_only(tmp_path, monkeypatch):
1419
assert w.store_replicate_config.preferred_segment == ""
1420
# No disk budget — enable_offload was absent (defaults to False).
1421
assert w.disk_offload_buffer_budget_bytes is None
1422
+
1423
+
1424
+# ---------------------------------------------------------------------------
1425
+# Stats/metrics tests (PR-35 port)
1426
+# ---------------------------------------------------------------------------
1427
+
1428
+
1429
+def test_mooncake_store_stats_aggregate_reduce():
1430
+ stats = MooncakeStoreConnectorStats()
1431
+ stats.record_operation("save_put", 0.01, 2, num_bytes=128)
1432
+ other = MooncakeStoreConnectorStats()
1433
+ other.record_operation(
1434
+ "save_put",
1435
+ 0.03,
1436
+ 1,
1437
+ num_bytes=64,
1438
+ status="error",
1439
+ num_failed_keys=1,
1440
+ )
1441
+
1442
+ reduced = stats.aggregate(other).reduce()
1443
+
1444
+ assert reduced["save_put_count"] == 2
1445
+ assert reduced["save_put_total_keys"] == 3
1446
+ assert reduced["save_put_total_bytes"] == 192
1447
+ assert reduced["save_put_failed_keys"] == 1
1448
+ assert reduced["save_put_error_count"] == 1
1449
+
1450
+
1451
+def test_worker_get_kv_connector_stats_resets_after_read():
1452
+ worker = _make_bare_worker()
1453
+ worker._record_kv_connector_operation(
1454
+ "save_put",
1455
+ 0.01,
1456
+ 2,
1457
+ num_bytes=128,
1458
+ )
1459
+
1460
+ stats = worker.get_kv_connector_stats()
1461
+
1462
+ assert isinstance(stats, MooncakeStoreConnectorStats)
1463
+ assert stats.data["save_put"][0]["num_bytes"] == 128
1464
+ assert worker.get_kv_connector_stats() is None
1465
+
1466
+
1467
+def test_lookup_records_mooncake_metrics():
1468
+ worker = _make_bare_worker()
1469
+ worker.store.batch_is_exist.return_value = [1, 1]
1470
+
1471
+ result = worker.lookup(32, [b"a0", b"a1"])
1472
+ stats = worker.get_kv_connector_stats()
1473
+
1474
+ assert result == 32
1475
+ assert isinstance(stats, MooncakeStoreConnectorStats)
1476
+ assert len(stats.data["lookup_exists"]) == 1
1477
+ assert stats.data["lookup_exists"][0]["num_keys"] == 2
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py
+34
@@ -28,6 +28,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
28
KVConnectorRole,
29
SupportsHMA,
30
)
31
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
32
+ KVConnectorPromMetrics,
33
+ KVConnectorStats,
34
+ PromMetric,
35
+ PromMetricT,
36
+)
37
from vllm.forward_context import ForwardContext
38
from vllm.logger import init_logger
39
from vllm.v1.attention.backend import AttentionMetadata
@@ -38,6 +44,7 @@ from vllm.v1.outputs import KVConnectorOutput
44
from vllm.v1.request import Request
45
46
from .data import MooncakeStoreConnectorMetadata
47
+from .metrics import MooncakeStoreConnectorStats, MooncakeStorePromMetrics
48
from .scheduler import MooncakeStoreScheduler
49
from .worker import MooncakeStoreWorker
50
@@ -274,3 +281,30 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
281
kv_events = MooncakeStoreKVEvents(num_workers=1)
282
kv_events.add_events(events)
283
return kv_events
284
+
285
+ def get_kv_connector_stats(self) -> KVConnectorStats | None:
286
+ if self.connector_worker is None:
287
+ return None
288
+ return self.connector_worker.get_kv_connector_stats()
289
+
290
+ @classmethod
291
+ def build_kv_connector_stats(
292
+ cls, data: dict[str, Any] | None = None
293
+ ) -> KVConnectorStats | None:
294
+ return (
295
+ MooncakeStoreConnectorStats(data=data)
296
+ if data is not None
297
+ else MooncakeStoreConnectorStats()
298
+ )
299
+
300
+ @classmethod
301
+ def build_prom_metrics(
302
+ cls,
303
+ vllm_config: VllmConfig,
304
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
305
+ labelnames: list[str],
306
+ per_engine_labelvalues: dict[int, list[object]],
307
+ ) -> KVConnectorPromMetrics:
308
+ return MooncakeStorePromMetrics(
309
+ vllm_config, metric_types, labelnames, per_engine_labelvalues
310
+ )
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/metrics.py
new
+189
@@ -0,0 +1,189 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Per-operation telemetry for MooncakeStoreConnector.
4
+
5
+Records one row per Mooncake RPC (``save_exists``, ``save_put``, ``load_get``,
6
+``lookup_exists``) with duration, key/byte counts, status, and failed-key
7
+count. Exposed to the logger via ``KVConnectorLogging`` and to Prometheus
8
+via ``MooncakeStorePromMetrics``.
9
+"""
10
+
11
+from dataclasses import dataclass
12
+from statistics import fmean
13
+from typing import Any
14
+
15
+from vllm.config import VllmConfig
16
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
17
+ KVConnectorPromMetrics,
18
+ KVConnectorStats,
19
+ PromMetric,
20
+ PromMetricT,
21
+)
22
+
23
+
24
+def _nearest_rank_percentile(values: list[float], percentile: float) -> float:
25
+ if not values:
26
+ return 0.0
27
+ sorted_values = sorted(values)
28
+ rank = max(
29
+ 0, min(len(sorted_values) - 1, int(percentile * len(sorted_values) - 1e-12))
30
+ )
31
+ return sorted_values[rank]
32
+
33
+
34
+@dataclass
35
+class MooncakeStoreConnectorStats(KVConnectorStats):
36
+ """Serializable Mooncake store communication telemetry."""
37
+
38
+ def __post_init__(self):
39
+ if not self.data:
40
+ self.reset()
41
+
42
+ def reset(self):
43
+ self.data: dict[str, list[dict[str, int | float | str]]] = {}
44
+
45
+ def is_empty(self) -> bool:
46
+ return not self.data
47
+
48
+ def aggregate(self, other: KVConnectorStats) -> KVConnectorStats:
49
+ if other.is_empty():
50
+ return self
51
+ for operation, records in other.data.items():
52
+ self.data.setdefault(operation, []).extend(records)
53
+ return self
54
+
55
+ def reduce(self) -> dict[str, int | float]:
56
+ reduced: dict[str, int | float] = {}
57
+ for operation, records in sorted(self.data.items()):
58
+ if not records:
59
+ continue
60
+ durations = [float(record["duration_seconds"]) for record in records]
61
+ reduced[f"{operation}_count"] = len(records)
62
+ reduced[f"{operation}_avg_ms"] = round(fmean(durations) * 1e3, 3)
63
+ reduced[f"{operation}_p90_ms"] = round(
64
+ _nearest_rank_percentile(durations, 0.9) * 1e3, 3
65
+ )
66
+ reduced[f"{operation}_total_keys"] = sum(
67
+ int(record["num_keys"]) for record in records
68
+ )
69
+ reduced[f"{operation}_total_bytes"] = sum(
70
+ int(record["num_bytes"]) for record in records
71
+ )
72
+ reduced[f"{operation}_failed_keys"] = sum(
73
+ int(record["num_failed_keys"]) for record in records
74
+ )
75
+ reduced[f"{operation}_error_count"] = sum(
76
+ 1 for record in records if record["status"] == "error"
77
+ )
78
+ return reduced
79
+
80
+ def record_operation(
81
+ self,
82
+ operation: str,
83
+ duration_seconds: float,
84
+ num_keys: int,
85
+ *,
86
+ num_bytes: int = 0,
87
+ status: str = "ok",
88
+ num_failed_keys: int = 0,
89
+ ) -> None:
90
+ self.data.setdefault(operation, []).append(
91
+ {
92
+ "duration_seconds": duration_seconds,
93
+ "num_keys": num_keys,
94
+ "num_bytes": num_bytes,
95
+ "status": status,
96
+ "num_failed_keys": num_failed_keys,
97
+ }
98
+ )
99
+
100
+
101
+class MooncakeStorePromMetrics(KVConnectorPromMetrics):
102
+ """Prometheus metrics for Mooncake store communication."""
103
+
104
+ def __init__(
105
+ self,
106
+ vllm_config: VllmConfig,
107
+ metric_types: dict[type[PromMetric], type[PromMetricT]],
108
+ labelnames: list[str],
109
+ per_engine_labelvalues: dict[int, list[object]],
110
+ ):
111
+ super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
112
+ metric_labelnames = labelnames + ["operation", "status"]
113
+ self._metric_cache: dict[tuple[int, str, str], dict[str, PromMetric]] = {}
114
+
115
+ self._histogram_operation_time = self._histogram_cls(
116
+ name="vllm:mooncake_store_operation_time_seconds",
117
+ documentation="Histogram of Mooncake store communication time.",
118
+ buckets=[
119
+ 1e-3,
120
+ 5e-3,
121
+ 1e-2,
122
+ 5e-2,
123
+ 1e-1,
124
+ 2e-1,
125
+ 3e-1,
126
+ 4e-1,
127
+ 5e-1,
128
+ 7.5e-1,
129
+ 1.0,
130
+ 1.5,
131
+ 2.0,
132
+ 3.0,
133
+ 4.0,
134
+ ],
135
+ labelnames=metric_labelnames,
136
+ )
137
+ self._counter_operation_calls = self._counter_cls(
138
+ name="vllm:mooncake_store_operation_total",
139
+ documentation="Number of Mooncake store communication operations.",
140
+ labelnames=metric_labelnames,
141
+ )
142
+ self._counter_operation_keys = self._counter_cls(
143
+ name="vllm:mooncake_store_operation_keys_total",
144
+ documentation="Number of Mooncake store keys touched by operations.",
145
+ labelnames=metric_labelnames,
146
+ )
147
+ self._counter_operation_bytes = self._counter_cls(
148
+ name="vllm:mooncake_store_operation_bytes_total",
149
+ documentation="Number of bytes transferred by Mooncake store operations.",
150
+ labelnames=metric_labelnames,
151
+ )
152
+ self._counter_failed_keys = self._counter_cls(
153
+ name="vllm:mooncake_store_operation_failed_keys_total",
154
+ documentation="Number of Mooncake store keys that failed in operations.",
155
+ labelnames=metric_labelnames,
156
+ )
157
+
158
+ def _get_metrics(
159
+ self,
160
+ engine_idx: int,
161
+ operation: str,
162
+ status: str,
163
+ ) -> dict[str, PromMetric]:
164
+ cache_key = (engine_idx, operation, status)
165
+ if cache_key not in self._metric_cache:
166
+ label_values = self.per_engine_labelvalues[engine_idx] + [operation, status]
167
+ self._metric_cache[cache_key] = {
168
+ "time": self._histogram_operation_time.labels(*label_values),
169
+ "calls": self._counter_operation_calls.labels(*label_values),
170
+ "keys": self._counter_operation_keys.labels(*label_values),
171
+ "bytes": self._counter_operation_bytes.labels(*label_values),
172
+ "failed_keys": self._counter_failed_keys.labels(*label_values),
173
+ }
174
+ return self._metric_cache[cache_key]
175
+
176
+ def observe(self, transfer_stats_data: dict[str, Any] | None, engine_idx: int = 0):
177
+ if not transfer_stats_data:
178
+ return
179
+ for operation, records in transfer_stats_data.items():
180
+ assert isinstance(records, list)
181
+ for record in records:
182
+ assert isinstance(record, dict)
183
+ status = str(record["status"])
184
+ metrics = self._get_metrics(engine_idx, operation, status)
185
+ metrics["time"].observe(float(record["duration_seconds"]))
186
+ metrics["calls"].inc()
187
+ metrics["keys"].inc(int(record["num_keys"]))
188
+ metrics["bytes"].inc(int(record["num_bytes"]))
189
+ metrics["failed_keys"].inc(int(record["num_failed_keys"]))
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+136
-4
@@ -16,7 +16,9 @@ import os
16
import queue
17
import socket
18
import threading
19
+import time
20
from collections import defaultdict
21
+from collections.abc import Callable
22
from dataclasses import dataclass
23
from typing import Any, Literal
24
@@ -58,6 +60,8 @@ from vllm.v1.core.kv_cache_utils import (
60
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheGroupSpec
61
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
62
63
+from .metrics import MooncakeStoreConnectorStats
64
+
65
logger = init_logger(__name__)
66
67
DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024 # 4 GiB
@@ -178,6 +182,10 @@ def _estimate_disk_offload_staging_bytes(size_list: list[int]) -> int:
182
return _align_up(data_size, _DIRECT_IO_ALIGNMENT) + _DIRECT_IO_PADDING_BYTES
183
184
185
+def _sum_batch_bytes(sizes: list[list[int]]) -> int:
186
+ return sum(sum(size) for size in sizes)
187
+
188
+
189
def _get_usable_disk_offload_buffer_budget_bytes(raw_budget_bytes: int) -> int:
190
return max(1, int(raw_budget_bytes * envs.VLLM_MOONCAKE_DISK_STAGING_USABLE_RATIO))
191
@@ -337,6 +345,7 @@ class KVTransferThread(threading.Thread):
345
tp_rank: int,
346
ready_event: threading.Event,
347
name: str,
348
+ record_operation: Callable[..., None] | None = None,
349
):
350
super().__init__(daemon=True, name=name)
351
self.store = store
@@ -344,6 +353,7 @@ class KVTransferThread(threading.Thread):
353
self.block_size = block_size
354
self.tp_rank = tp_rank
355
self.token_databases = token_databases
356
+ self._record_operation_cb = record_operation
357
self.done_task_lock = threading.Lock()
358
self.request_queue: queue.Queue[Any] = queue.Queue()
359
self.finished_requests: set[str] = set()
@@ -379,6 +389,27 @@ class KVTransferThread(threading.Thread):
389
def _handle_request(self, req_meta: Any):
390
pass
391
392
+ def _record_operation(
393
+ self,
394
+ operation: str,
395
+ start_time: float,
396
+ num_keys: int,
397
+ *,
398
+ num_bytes: int = 0,
399
+ status: str = "ok",
400
+ num_failed_keys: int = 0,
401
+ ) -> None:
402
+ if self._record_operation_cb is None:
403
+ return
404
+ self._record_operation_cb(
405
+ operation=operation,
406
+ duration_seconds=time.perf_counter() - start_time,
407
+ num_keys=num_keys,
408
+ num_bytes=num_bytes,
409
+ status=status,
410
+ num_failed_keys=num_failed_keys,
411
+ )
412
+
413
def update_kv_event(self, events: list[BlockStored]):
414
with self.kv_event_lock:
415
self.kv_events.extend(events)
@@ -405,6 +436,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
436
ready_event: threading.Event,
437
enable_kv_event: bool = False,
438
replicate_config: Any = None,
439
+ record_operation: Callable[..., None] | None = None,
440
):
441
super().__init__(
442
store,
@@ -413,6 +445,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
445
tp_rank,
446
ready_event,
447
name="KVCacheStoreSendingThread",
448
+ record_operation=record_operation,
449
)
450
self.put_step = put_step
451
self.coord = coord
@@ -520,7 +553,23 @@ class KVCacheStoreSendingThread(KVTransferThread):
553
return
554
555
# Check which blocks already exist (dedup)
523
- exists_states = self.store.batch_is_exist(keys)
556
+ save_exists_start = time.perf_counter()
557
+ try:
558
+ exists_states = self.store.batch_is_exist(keys)
559
+ except Exception:
560
+ self._record_operation(
561
+ "save_exists",
562
+ save_exists_start,
563
+ len(keys),
564
+ status="error",
565
+ num_failed_keys=len(keys),
566
+ )
567
+ raise
568
+ self._record_operation(
569
+ "save_exists",
570
+ save_exists_start,
571
+ len(keys),
572
+ )
573
missing_indices = [i for i, exists in enumerate(exists_states) if exists != 1]
574
575
if not missing_indices:
@@ -574,6 +623,8 @@ class KVCacheStoreSendingThread(KVTransferThread):
623
if current_event is not None:
624
current_event.synchronize()
625
626
+ batch_bytes = _sum_batch_bytes(sizes)
627
+ put_start = time.perf_counter()
628
try:
629
res = self.store.batch_put_from_multi_buffers(
630
keys,
@@ -582,9 +633,15 @@ class KVCacheStoreSendingThread(KVTransferThread):
633
self.replicate_config,
634
)
635
failed = [i for i, v in enumerate(res) if v < 0]
636
+ self._record_operation(
637
+ "save_put",
638
+ put_start,
639
+ len(keys),
640
+ num_bytes=batch_bytes,
641
+ status="partial_failure" if failed else "ok",
642
+ num_failed_keys=len(failed),
643
+ )
644
if failed:
586
- # Compute total bytes attempted for this batch
587
- total_bytes = sum(sum(s) if isinstance(s, list) else s for s in sizes)
645
failed_codes = set(res[i] for i in failed)
646
logger.warning(
647
"batch_put failed: %d/%d keys failed "
@@ -593,7 +650,7 @@ class KVCacheStoreSendingThread(KVTransferThread):
650
len(failed),
651
len(keys),
652
failed_codes,
596
- total_bytes,
653
+ batch_bytes,
654
len(keys),
655
keys[0] if keys else "N/A",
656
)
@@ -614,6 +671,14 @@ class KVCacheStoreSendingThread(KVTransferThread):
671
"successful store batch"
672
)
673
except Exception as e:
674
+ self._record_operation(
675
+ "save_put",
676
+ put_start,
677
+ len(keys),
678
+ num_bytes=batch_bytes,
679
+ status="error",
680
+ num_failed_keys=len(keys),
681
+ )
682
logger.error("Failed to put key %s, error: %s", keys, e)
683
684
if self.enable_kv_event and stored_events:
@@ -635,6 +700,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
700
tp_rank: int,
701
ready_event: threading.Event,
702
disk_offload_buffer_budget_bytes: int | None = None,
703
+ record_operation: Callable[..., None] | None = None,
704
):
705
super().__init__(
706
store,
@@ -643,6 +709,7 @@ class KVCacheStoreRecvingThread(KVTransferThread):
709
tp_rank,
710
ready_event,
711
name="KVCacheStoreRecvingThread",
712
+ record_operation=record_operation,
713
)
714
self.disk_offload_buffer_budget_bytes = disk_offload_buffer_budget_bytes
715
self.usable_disk_offload_buffer_budget_bytes = (
@@ -721,12 +788,16 @@ class KVCacheStoreRecvingThread(KVTransferThread):
788
return
789
790
current_batch_keys: list[str] = key_list_c
791
+ batch_bytes = 0
792
try:
793
for batch_keys, batch_addrs, batch_sizes in load_batches:
794
current_batch_keys = batch_keys
795
+ batch_bytes = _sum_batch_bytes(batch_sizes)
796
tiers_by_key: dict[str, str] | None = None
797
if envs.VLLM_MOONCAKE_STORE_TIER_LOG:
798
tiers_by_key = _get_replica_tiers_by_key(self.store, batch_keys)
799
+ # Reset so the recorded RPC duration excludes tier lookup.
800
+ load_get_start = time.perf_counter()
801
res = self.store.batch_get_into_multi_buffers(
802
batch_keys, batch_addrs, batch_sizes
803
)
@@ -739,6 +810,14 @@ class KVCacheStoreRecvingThread(KVTransferThread):
810
for key, value in zip(batch_keys, res, strict=True)
811
if value < 0
812
]
813
+ self._record_operation(
814
+ "load_get",
815
+ load_get_start,
816
+ len(batch_keys),
817
+ num_bytes=batch_bytes,
818
+ status="partial_failure" if failed else "ok",
819
+ num_failed_keys=len(failed),
820
+ )
821
if failed:
822
logger.warning(
823
"Failed to get %d Mooncake keys from sub-batch "
@@ -749,6 +828,14 @@ class KVCacheStoreRecvingThread(KVTransferThread):
828
)
829
break
830
except Exception as e:
831
+ self._record_operation(
832
+ "load_get",
833
+ load_get_start,
834
+ len(current_batch_keys),
835
+ num_bytes=batch_bytes,
836
+ status="error",
837
+ num_failed_keys=len(current_batch_keys),
838
+ )
839
logger.warning(
840
"Failed to get Mooncake sub-batch %s, error: %s",
841
current_batch_keys[:3],
@@ -921,6 +1008,8 @@ class MooncakeStoreWorker:
1008
self.kv_send_thread: KVCacheStoreSendingThread | None = None
1009
self.kv_recv_thread: KVCacheStoreRecvingThread | None = None
1010
self.finished_store_req: set[str] = set()
1011
+ self._kv_connector_stats_lock = threading.Lock()
1012
+ self.kv_connector_stats = MooncakeStoreConnectorStats()
1013
1014
self._kv_cache_config = kv_cache_config
1015
# Single-group + PCP/DCP > 1: scale the lone group's spec.block_size to
@@ -1057,6 +1146,7 @@ class MooncakeStoreWorker:
1146
ready_event_sending,
1147
self.enable_kv_events,
1148
self.store_replicate_config,
1149
+ record_operation=self._record_kv_connector_operation,
1150
)
1151
self.kv_send_thread.start()
1152
@@ -1069,6 +1159,7 @@ class MooncakeStoreWorker:
1159
self.tp_rank,
1160
ready_event_recving,
1161
disk_offload_buffer_budget_bytes=self.disk_offload_buffer_budget_bytes,
1162
+ record_operation=self._record_kv_connector_operation,
1163
)
1164
self.kv_recv_thread.start()
1165
ready_event_recving.wait()
@@ -1155,6 +1246,34 @@ class MooncakeStoreWorker:
1246
)
1247
return done_sending, done_recving
1248
1249
+ def _record_kv_connector_operation(
1250
+ self,
1251
+ operation: str,
1252
+ duration_seconds: float,
1253
+ num_keys: int,
1254
+ *,
1255
+ num_bytes: int = 0,
1256
+ status: str = "ok",
1257
+ num_failed_keys: int = 0,
1258
+ ) -> None:
1259
+ with self._kv_connector_stats_lock:
1260
+ self.kv_connector_stats.record_operation(
1261
+ operation=operation,
1262
+ duration_seconds=duration_seconds,
1263
+ num_keys=num_keys,
1264
+ num_bytes=num_bytes,
1265
+ status=status,
1266
+ num_failed_keys=num_failed_keys,
1267
+ )
1268
+
1269
+ def get_kv_connector_stats(self) -> MooncakeStoreConnectorStats | None:
1270
+ with self._kv_connector_stats_lock:
1271
+ if self.kv_connector_stats.is_empty():
1272
+ return None
1273
+ kv_connector_stats = self.kv_connector_stats
1274
+ self.kv_connector_stats = MooncakeStoreConnectorStats()
1275
+ return kv_connector_stats
1276
+
1277
def _get_and_clear_finished_sending(
1278
self,
1279
finished_req_ids: set[str],
@@ -1216,9 +1335,22 @@ class MooncakeStoreWorker:
1335
if not candidate_keys:
1336
return 0
1337
1338
+ lookup_start = time.perf_counter()
1339
try:
1340
res = self.store.batch_is_exist(candidate_keys)
1341
+ self._record_kv_connector_operation(
1342
+ "lookup_exists",
1343
+ time.perf_counter() - lookup_start,
1344
+ len(candidate_keys),
1345
+ )
1346
except Exception as e:
1347
+ self._record_kv_connector_operation(
1348
+ "lookup_exists",
1349
+ time.perf_counter() - lookup_start,
1350
+ len(candidate_keys),
1351
+ status="error",
1352
+ num_failed_keys=len(candidate_keys),
1353
+ )
1354
logger.error("Remote connection failed in lookup: %s", e)
1355
return 0
1356