@tduncan / Nvidia-Optimization-Stack-VLLM / commits / cd1daf58a9

[Metrics][KV Offload] Add Prometheus metrics for SimpleCPUOffloadConnector (#57251)

Signed-off-by: Vincent <vincexxchan@gmail.com>

mevince committed Sep 26, 2026 at 17:17 UTC cd1daf58a98aad8c7241377edbe91a19aa38f71a
6 files changed +903 -6
docs/mkdocs/gen_files/generate_metrics.py
+4
@@ -23,6 +23,10 @@ METRIC_SOURCE_FILES = [
23 "path": "vllm/distributed/kv_transfer/kv_connector/v1/nixl/stats.py",
24 "key": "metrics-nixl",
25 },
26 + {
27 + "path": "vllm/v1/simple_kv_offload/metrics.py",
28 + "key": "metrics-simple-kv-offload",
29 + },
30 {"path": "vllm/v1/metrics/perf.py", "key": "metrics-mfu"},
31 ]
32
docs/usage/metrics.md
+24
@@ -45,6 +45,30 @@ The following metrics are exposed:
45
46 --8<-- "gen:metrics-nixl"
47
48 +## Simple CPU Offload Connector Metrics
49 +
50 +These metrics are exposed when the `SimpleCPUOffloadConnector` KV connector
51 +is configured (e.g. `--kv-transfer-config='{"kv_connector":
52 +"SimpleCPUOffloadConnector", "kv_role": "kv_both", "kv_connector_extra_config":
53 +{"kv_offload_backend": "disk", "disk_path": "/mnt/nvme/kv"}}'`). They are
54 +updated once per engine step.
55 +
56 +Caveats to keep in mind when interpreting them:
57 +
58 +- A "completed" store means the write syscalls returned; it is **not**
59 + fsync-durable, and the disk backend's file is process-lifetime scratch,
60 + unlinked at startup and shutdown.
61 +- `save_outcomes_total` classifies eager-mode boundary hand-off stores;
62 + lazy-mode stores are not classified.
63 +- `used_blocks` counts blocks pinned by in-flight transfers or cache hits;
64 + warm cached blocks that are evictable are not counted. Use the
65 + `capacity_blocks` label of `simple_kv_offload_info` as the denominator.
66 +- Counters and gauges are quantized to engine steps. They are reported by the
67 + scheduler process and reflect engine-wide logical block counts, not values
68 + pooled from individual tensor-parallel workers.
69 +
70 +--8<-- "gen:metrics-simple-kv-offload"
71 +
72 ## Model Flops Utilization (MFU) Performance Metrics
73
74 These metrics are available via `--enable-mfu-metrics`:
tests/v1/kv_connector/unit/test_simple_cpu_offload_metrics.py new
+466
@@ -0,0 +1,466 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for SimpleCPUOffloadConnector metrics."""
4 +
5 +from __future__ import annotations
6 +
7 +import dataclasses
8 +from types import SimpleNamespace
9 +from typing import Any
10 +
11 +import pytest
12 +from prometheus_client import Counter, Gauge, Histogram
13 +
14 +from tests.v1.kv_connector.unit.test_simple_cpu_offload_connector import (
15 + _make_connector,
16 +)
17 +from tests.v1.simple_kv_offload.test_scheduler import (
18 + make_request,
19 + make_scheduler,
20 + make_scheduler_output,
21 +)
22 +from vllm.distributed.kv_transfer.kv_connector.v1.simple_cpu_offload_connector import (
23 + SimpleCPUOffloadConnector,
24 +)
25 +from vllm.v1.outputs import KVConnectorOutput
26 +from vllm.v1.simple_kv_offload.manager import (
27 + BoundaryStoreStats,
28 + LoadRequestState,
29 + TransferMeta,
30 +)
31 +from vllm.v1.simple_kv_offload.metrics import (
32 + INFO_LABELS,
33 + OUTCOME_TO_FIELD,
34 + MetricName,
35 + SimpleCPUOffloadStats,
36 + _MetricType,
37 + _StatsKey,
38 +)
39 +
40 +SAVE_OUTCOMES = MetricName.SAVE_OUTCOMES
41 +LOAD_BLOCKS = MetricName.LOAD_BLOCKS
42 +USED_BLOCKS = MetricName.USED_BLOCKS
43 +PENDING_STORE_BLOCKS = MetricName.PENDING_STORE_BLOCKS
44 +INFO = MetricName.INFO
45 +
46 +
47 +class _FakeMetric:
48 + def __init__(self, **kwargs: Any):
49 + self.kwargs = kwargs
50 + self.children: list[_FakeMetric] = []
51 + self.observed: list[int | float] = []
52 + self.increments: list[int | float] = []
53 + self.set_values: list[int | float] = []
54 + self.labelvalues: tuple[object, ...] = ()
55 +
56 + def labels(self, *labelvalues: object) -> _FakeMetric:
57 + child = _FakeMetric(**self.kwargs)
58 + child.labelvalues = labelvalues
59 + self.children.append(child)
60 + return child
61 +
62 + def observe(self, value: int | float) -> None:
63 + self.observed.append(value)
64 +
65 + def inc(self, value: int | float) -> None:
66 + self.increments.append(value)
67 +
68 + def set(self, value: int | float) -> None:
69 + self.set_values.append(value)
70 +
71 +
72 +def _fake_metric_types() -> dict[type, type]:
73 + return {Gauge: _FakeMetric, Counter: _FakeMetric, Histogram: _FakeMetric}
74 +
75 +
76 +def _make_prom_metrics(per_engine: dict[int, list[object]] | None = None) -> Any:
77 + from vllm.v1.simple_kv_offload.metrics import SimpleCPUOffloadPromMetrics
78 +
79 + return SimpleCPUOffloadPromMetrics(
80 + vllm_config=SimpleNamespace(kv_transfer_config=None),
81 + metric_types=_fake_metric_types(),
82 + labelnames=["model_name", "engine"],
83 + per_engine_labelvalues=per_engine or {0: ["model", "0"]},
84 + )
85 +
86 +
87 +def _stats_payload(types: dict, data: dict) -> dict[str, Any]:
88 + return {_StatsKey.TYPES: types, _StatsKey.DATA: data}
89 +
90 +
91 +def _values(stats: SimpleCPUOffloadStats, name: str) -> dict:
92 + return stats.data[_StatsKey.DATA][name]
93 +
94 +
95 +# ---------------------------------------------------------------------------
96 +# Stats class
97 +# ---------------------------------------------------------------------------
98 +
99 +
100 +def test_build_kv_connector_stats_none_and_empty() -> None:
101 + stats = SimpleCPUOffloadConnector.build_kv_connector_stats(data=None)
102 + assert isinstance(stats, SimpleCPUOffloadStats)
103 + assert stats.is_empty()
104 +
105 + stats = SimpleCPUOffloadConnector.build_kv_connector_stats(data={})
106 + assert isinstance(stats, SimpleCPUOffloadStats)
107 + assert stats.is_empty()
108 +
109 +
110 +def test_stats_round_trip_through_serialized_dict() -> None:
111 + stats = SimpleCPUOffloadStats()
112 + stats.increase_counter(SAVE_OUTCOMES, 3, ("stored",))
113 + stats.increase_counter(LOAD_BLOCKS, 4)
114 + stats.set_gauge(USED_BLOCKS, 7)
115 + stats.set_gauge(INFO, 1, ("cpu", "false", "false", "64"))
116 +
117 + rebuilt = SimpleCPUOffloadConnector.build_kv_connector_stats(data=stats.to_dict())
118 + assert isinstance(rebuilt, SimpleCPUOffloadStats)
119 + assert _values(rebuilt, SAVE_OUTCOMES) == {("stored",): 3}
120 + assert _values(rebuilt, LOAD_BLOCKS) == {(): 4}
121 + assert _values(rebuilt, USED_BLOCKS) == {(): 7}
122 + assert _values(rebuilt, INFO) == {("cpu", "false", "false", "64"): 1}
123 +
124 +
125 +def test_stats_aggregate_sums_counters_and_keeps_latest_gauge() -> None:
126 + stats1 = SimpleCPUOffloadStats()
127 + stats1.increase_counter(SAVE_OUTCOMES, 3, ("stored",))
128 + stats1.increase_counter(SAVE_OUTCOMES, 1, ("dropped_cpu_full",))
129 + stats1.set_gauge(USED_BLOCKS, 7)
130 +
131 + stats2 = SimpleCPUOffloadStats()
132 + stats2.increase_counter(SAVE_OUTCOMES, 4, ("stored",))
133 + stats2.set_gauge(USED_BLOCKS, 2)
134 +
135 + result = stats1.aggregate(stats2)
136 + assert result is stats1
137 + outcomes = _values(stats1, SAVE_OUTCOMES)
138 + assert outcomes == {("stored",): 7, ("dropped_cpu_full",): 1}
139 + assert _values(stats1, USED_BLOCKS) == {(): 2}
140 +
141 +
142 +def test_stats_aggregate_merges_types_from_other() -> None:
143 + stats1 = SimpleCPUOffloadStats()
144 + stats1.set_gauge(USED_BLOCKS, 1)
145 +
146 + stats2 = SimpleCPUOffloadStats()
147 + stats2.increase_counter(LOAD_BLOCKS, 2)
148 +
149 + stats1.aggregate(stats2)
150 + assert stats1.data[_StatsKey.TYPES][LOAD_BLOCKS] == _MetricType.COUNTER
151 + assert _values(stats1, LOAD_BLOCKS) == {(): 2}
152 +
153 +
154 +def test_stats_reduce_and_reset() -> None:
155 + stats = SimpleCPUOffloadStats()
156 + stats.increase_counter(SAVE_OUTCOMES, 3, ("stored",))
157 + stats.set_gauge(PENDING_STORE_BLOCKS, 4)
158 + stats.increase_counter(LOAD_BLOCKS, 9)
159 + stats.set_gauge(INFO, 1, ("cpu", "false", "false", "64"))
160 +
161 + reduced = stats.reduce()
162 + assert reduced[f"{SAVE_OUTCOMES}:{('stored',)}"] == 3
163 + assert reduced[PENDING_STORE_BLOCKS] == 4
164 + assert reduced[LOAD_BLOCKS] == 9
165 + assert not any(key.startswith(INFO) for key in reduced)
166 +
167 + stats.reset()
168 + assert stats.is_empty()
169 +
170 +
171 +def test_outcome_labels_match_boundary_store_stats_fields() -> None:
172 + """Every outcome label maps 1:1 to a BoundaryStoreStats counter field."""
173 + boundary_fields = {
174 + f.name for f in dataclasses.fields(BoundaryStoreStats) if f.name != "published"
175 + }
176 + assert set(OUTCOME_TO_FIELD.values()) == boundary_fields
177 + assert len(set(OUTCOME_TO_FIELD.values())) == len(OUTCOME_TO_FIELD)
178 +
179 +
180 +# ---------------------------------------------------------------------------
181 +# Prom metrics class
182 +# ---------------------------------------------------------------------------
183 +
184 +
185 +def test_prom_metrics_registers_tier1_metrics() -> None:
186 + prom = _make_prom_metrics()
187 + defs = prom._defs
188 + assert set(defs) == {
189 + SAVE_OUTCOMES,
190 + LOAD_BLOCKS,
191 + USED_BLOCKS,
192 + PENDING_STORE_BLOCKS,
193 + INFO,
194 + }
195 + # The prom class registers literal names (required by the docs
196 + # generator); they must stay in sync with the emission-side constants.
197 + assert {m.kwargs["name"] for m in defs.values()} == {
198 + SAVE_OUTCOMES,
199 + LOAD_BLOCKS,
200 + USED_BLOCKS,
201 + PENDING_STORE_BLOCKS,
202 + INFO,
203 + }
204 + for metric in defs.values():
205 + assert metric.kwargs["documentation"]
206 + assert metric.kwargs["labelnames"][:2] == ["model_name", "engine"]
207 +
208 + assert defs[SAVE_OUTCOMES].kwargs["labelnames"] == [
209 + "model_name",
210 + "engine",
211 + "outcome",
212 + ]
213 + assert defs[LOAD_BLOCKS].kwargs["labelnames"] == ["model_name", "engine"]
214 + assert defs[INFO].kwargs["labelnames"] == [
215 + "model_name",
216 + "engine",
217 + *INFO_LABELS,
218 + ]
219 + assert defs[USED_BLOCKS].kwargs["labelnames"] == ["model_name", "engine"]
220 +
221 + # mostrecent: multiprocess scrapes expose one series per label set, not
222 + # one pid-tagged series per writer process.
223 + for gauge in (USED_BLOCKS, PENDING_STORE_BLOCKS, INFO):
224 + assert defs[gauge].kwargs["multiprocess_mode"] == "mostrecent"
225 +
226 +
227 +def test_prom_metrics_observe_routes_each_series() -> None:
228 + prom = _make_prom_metrics()
229 + prom.observe(
230 + _stats_payload(
231 + types={
232 + SAVE_OUTCOMES: _MetricType.COUNTER,
233 + LOAD_BLOCKS: _MetricType.COUNTER,
234 + USED_BLOCKS: _MetricType.GAUGE,
235 + PENDING_STORE_BLOCKS: _MetricType.GAUGE,
236 + INFO: _MetricType.GAUGE,
237 + },
238 + data={
239 + SAVE_OUTCOMES: {("stored",): 3},
240 + LOAD_BLOCKS: {(): 4},
241 + USED_BLOCKS: {(): 5},
242 + PENDING_STORE_BLOCKS: {(): 1},
243 + INFO: {("cpu", "false", "false", "64"): 1},
244 + },
245 + )
246 + )
247 +
248 + assert prom._metrics[(0, SAVE_OUTCOMES, ("stored",))].increments == [3]
249 + assert prom._metrics[(0, LOAD_BLOCKS, ())].increments == [4]
250 + assert prom._metrics[(0, USED_BLOCKS, ())].set_values == [5]
251 + assert prom._metrics[(0, PENDING_STORE_BLOCKS, ())].set_values == [1]
252 + assert prom._metrics[(0, INFO, ("cpu", "false", "false", "64"))].set_values == [1]
253 + assert prom._metrics[(0, SAVE_OUTCOMES, ("stored",))].labelvalues == (
254 + "model",
255 + "0",
256 + "stored",
257 + )
258 + assert prom._metrics[(0, INFO, ("cpu", "false", "false", "64"))].labelvalues == (
259 + "model",
260 + "0",
261 + "cpu",
262 + "false",
263 + "false",
264 + "64",
265 + )
266 +
267 +
268 +def test_prom_metrics_reuses_bound_children() -> None:
269 + prom = _make_prom_metrics()
270 + payload = _stats_payload(
271 + types={SAVE_OUTCOMES: _MetricType.COUNTER},
272 + data={SAVE_OUTCOMES: {("stored",): 1}},
273 + )
274 + prom.observe(payload)
275 + prom.observe(payload)
276 +
277 + child = prom._metrics[(0, SAVE_OUTCOMES, ("stored",))]
278 + assert child.increments == [1, 1]
279 + assert len(prom._defs[SAVE_OUTCOMES].children) == 1
280 +
281 +
282 +def test_prom_metrics_rejects_unknown_metric() -> None:
283 + prom = _make_prom_metrics()
284 + with pytest.raises(AssertionError, match="Unknown"):
285 + prom.observe(
286 + _stats_payload(
287 + types={"vllm:not_a_metric": _MetricType.COUNTER},
288 + data={"vllm:not_a_metric": {(): 1}},
289 + )
290 + )
291 +
292 +
293 +def test_prom_metrics_rejects_wrong_label_count() -> None:
294 + prom = _make_prom_metrics()
295 + with pytest.raises(AssertionError, match="labels"):
296 + prom.observe(
297 + _stats_payload(
298 + types={SAVE_OUTCOMES: _MetricType.COUNTER},
299 + data={SAVE_OUTCOMES: {("stored", "extra"): 1}},
300 + )
301 + )
302 +
303 +
304 +def test_prom_metrics_routes_per_engine() -> None:
305 + prom = _make_prom_metrics(per_engine={0: ["model", "0"], 1: ["model", "1"]})
306 + prom.observe(
307 + _stats_payload(
308 + types={USED_BLOCKS: _MetricType.GAUGE}, data={USED_BLOCKS: {(): 3}}
309 + ),
310 + engine_idx=1,
311 + )
312 + assert (0, USED_BLOCKS, ()) not in prom._metrics
313 + assert prom._metrics[(1, USED_BLOCKS, ())].set_values == [3]
314 +
315 +
316 +def test_connector_build_prom_metrics() -> None:
317 + from vllm.v1.simple_kv_offload.metrics import SimpleCPUOffloadPromMetrics
318 +
319 + prom = SimpleCPUOffloadConnector.build_prom_metrics(
320 + vllm_config=SimpleNamespace(kv_transfer_config=None),
321 + metric_types=_fake_metric_types(),
322 + labelnames=["model_name", "engine"],
323 + per_engine_labelvalues={0: ["model", "0"]},
324 + )
325 + assert isinstance(prom, SimpleCPUOffloadPromMetrics)
326 +
327 +
328 +# ---------------------------------------------------------------------------
329 +# Manager: interval stats
330 +# ---------------------------------------------------------------------------
331 +
332 +
333 +def test_manager_get_stats_reports_outcome_deltas_per_interval() -> None:
334 + fixture = make_scheduler()
335 + sched = fixture.scheduler
336 +
337 + sched.boundary_store_stats.stored += 3
338 + sched.boundary_store_stats.dropped_cpu_full += 2
339 + sched.boundary_store_stats.skipped_in_flight += 1
340 +
341 + stats = sched.get_stats()
342 + assert _values(stats, SAVE_OUTCOMES) == {
343 + ("stored",): 3,
344 + ("dropped_cpu_full",): 2,
345 + ("skipped_in_flight",): 1,
346 + }
347 +
348 + # Deltas are per-interval: a second drain reports no new outcomes.
349 + stats2 = sched.get_stats()
350 + assert SAVE_OUTCOMES not in stats2.data[_StatsKey.DATA]
351 +
352 +
353 +def test_manager_reset_preserves_undrained_outcome_deltas() -> None:
354 + fixture = make_scheduler()
355 + sched = fixture.scheduler
356 +
357 + sched.boundary_store_stats.stored += 3
358 + sched.reset()
359 +
360 + stats = sched.get_stats()
361 + assert _values(stats, SAVE_OUTCOMES) == {("stored",): 3}
362 +
363 +
364 +def test_manager_get_stats_gauges_and_info_labels() -> None:
365 + fixture = make_scheduler(num_cpu_blocks=8)
366 + sched = fixture.scheduler
367 +
368 + sched._store_event_to_blocks[0] = TransferMeta([10], [11])
369 + sched._pending_finished_stores.append(TransferMeta([12], [13]))
370 + sched._abandoned_store_event_to_blocks[1] = TransferMeta([14], [15])
371 +
372 + stats = sched.get_stats()
373 + assert _values(stats, PENDING_STORE_BLOCKS) == {(): 3}
374 + used = _values(stats, USED_BLOCKS)[()]
375 + assert used == sched.num_cpu_blocks - sched.cpu_block_pool.get_num_free_blocks()
376 +
377 + info = _values(stats, INFO)
378 + assert info == {("cpu", "false", "false", str(sched.num_cpu_blocks)): 1}
379 +
380 + # Gauges are re-reported every drain.
381 + stats2 = sched.get_stats()
382 + assert _values(stats2, PENDING_STORE_BLOCKS) == {(): 3}
383 +
384 +
385 +def test_manager_info_labels_reflect_disk_mode() -> None:
386 + connector = _make_connector(
387 + extra_config={
388 + "kv_offload_backend": "disk",
389 + "disk_path": "/tmp/fake-disk-metrics",
390 + "disk_capacity_bytes": 1024**3,
391 + "use_page_cache": True,
392 + }
393 + )
394 + sched = connector.scheduler_manager
395 + assert sched is not None
396 + assert sched._info_labelvalues[0] == "disk"
397 + assert sched._info_labelvalues[1] == "true"
398 + stats = sched.get_stats()
399 + assert set(_values(stats, INFO)) == {sched._info_labelvalues}
400 +
401 +
402 +def test_manager_info_labels_page_cache_false_for_cpu_backend() -> None:
403 + """CPU backend ignores use_page_cache; the info gauge must not claim it."""
404 + connector = _make_connector(extra_config={"use_page_cache": True})
405 + sched = connector.scheduler_manager
406 + assert sched is not None
407 + assert sched._info_labelvalues[:2] == ("cpu", "false")
408 + stats = sched.get_stats()
409 + assert set(_values(stats, INFO)) == {sched._info_labelvalues}
410 +
411 +
412 +def test_manager_counts_load_blocks_completed() -> None:
413 + fixture = make_scheduler(num_cpu_blocks=8)
414 + sched = fixture.scheduler
415 +
416 + request = make_request(num_blocks=2)
417 + gpu_blocks = fixture.gpu_block_pool.get_new_blocks(2)
418 + cpu_blocks = sched.cpu_block_pool.get_new_blocks(2)
419 + sched._reqs_to_load[request.request_id] = LoadRequestState(
420 + request=request,
421 + transfer_meta=TransferMeta(
422 + [b.block_id for b in gpu_blocks], [b.block_id for b in cpu_blocks]
423 + ),
424 + )
425 +
426 + sched.build_connector_meta(make_scheduler_output({}))
427 + stats = sched.get_stats()
428 + assert LOAD_BLOCKS not in stats.data[_StatsKey.DATA]
429 +
430 + sched.update_connector_output(
431 + KVConnectorOutput(finished_recving={request.request_id})
432 + )
433 + stats = sched.get_stats()
434 + assert _values(stats, LOAD_BLOCKS) == {(): 2}
435 +
436 +
437 +def test_manager_abandoned_load_counts_as_completed_after_reset() -> None:
438 + """Loads abandoned by reset() that still finish on the worker are
439 + counted as loaded blocks.
440 + """
441 + fixture = make_scheduler(num_cpu_blocks=8)
442 + sched = fixture.scheduler
443 +
444 + request = make_request(num_blocks=2)
445 + gpu_blocks = fixture.gpu_block_pool.get_new_blocks(2)
446 + cpu_blocks = sched.cpu_block_pool.get_new_blocks(2)
447 + sched._reqs_to_load[request.request_id] = LoadRequestState(
448 + request=request,
449 + transfer_meta=TransferMeta(
450 + [b.block_id for b in gpu_blocks], [b.block_id for b in cpu_blocks]
451 + ),
452 + )
453 +
454 + sched.build_connector_meta(make_scheduler_output({}))
455 + stats = sched.get_stats()
456 + assert LOAD_BLOCKS not in stats.data[_StatsKey.DATA]
457 +
458 + assert sched.reset() is False
459 + assert set(sched._abandoned_reqs_to_load) == {request.request_id}
460 +
461 + sched.update_connector_output(
462 + KVConnectorOutput(finished_recving={request.request_id})
463 + )
464 + assert sched._abandoned_reqs_to_load == {}
465 + stats = sched.get_stats()
466 + assert _values(stats, LOAD_BLOCKS) == {(): 2}
vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
+43
@@ -15,6 +15,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
15 KVConnectorRole,
16 SupportsHMA,
17 )
18 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
19 + KVConnectorPromMetrics,
20 + KVConnectorStats,
21 + PromMetric,
22 + PromMetricT,
23 +)
24 from vllm.logger import init_logger
25 from vllm.v1.core.sched.output import SchedulerOutput
26 from vllm.v1.outputs import KVConnectorOutput
@@ -25,6 +31,10 @@ from vllm.v1.simple_kv_offload.manager import (
31 from vllm.v1.simple_kv_offload.metadata import (
32 SimpleCPUOffloadMetadata,
33 )
34 +from vllm.v1.simple_kv_offload.metrics import (
35 + SimpleCPUOffloadPromMetrics,
36 + SimpleCPUOffloadStats,
37 +)
38 from vllm.v1.simple_kv_offload.worker import (
39 SimpleCPUOffloadWorker,
40 )
@@ -162,6 +172,7 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
172 hash_block_size=hash_block_size,
173 lazy_offload=lazy_offload,
174 disk_capacity_bytes=disk_capacity_bytes if disk_mode else 0,
175 + use_page_cache=use_page_cache if disk_mode else False,
176 )
177 elif role == KVConnectorRole.WORKER:
178 self.worker_handler = SimpleCPUOffloadWorker(
@@ -307,6 +318,38 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA):
318 return self.scheduler_manager.get_boundary_store_stats()
319 return None
320
321 + def get_kv_connector_stats(self) -> KVConnectorStats | None:
322 + # Worker-side callers may invoke this hook on connectors constructed
323 + # via __new__ (as in test_worker.py), which lack scheduler_manager.
324 + scheduler_manager: SimpleCPUOffloadScheduler | None = getattr(
325 + self, "scheduler_manager", None
326 + )
327 + if scheduler_manager is not None:
328 + return scheduler_manager.get_stats()
329 + return None
330 +
331 + @classmethod
332 + def build_kv_connector_stats(
333 + cls, data: dict[str, Any] | None = None
334 + ) -> KVConnectorStats | None:
335 + return (
336 + SimpleCPUOffloadStats(data=data)
337 + if data is not None
338 + else SimpleCPUOffloadStats()
339 + )
340 +
341 + @classmethod
342 + def build_prom_metrics(
343 + cls,
344 + vllm_config: VllmConfig,
345 + metric_types: dict[type[PromMetric], type[PromMetricT]],
346 + labelnames: list[str],
347 + per_engine_labelvalues: dict[int, list[object]],
348 + ) -> KVConnectorPromMetrics:
349 + return SimpleCPUOffloadPromMetrics(
350 + vllm_config, metric_types, labelnames, per_engine_labelvalues
351 + )
352 +
353 def take_events(self) -> Iterable[KVCacheEvent]:
354 if self.scheduler_manager is not None:
355 return self.scheduler_manager.take_events()
vllm/v1/simple_kv_offload/manager.py
+82 -6
@@ -4,7 +4,7 @@
4
5 import contextlib
6 from collections.abc import Iterable
7 -from dataclasses import dataclass, field, replace
7 +from dataclasses import dataclass, field, fields, replace
8 from enum import Enum, auto
9 from typing import TYPE_CHECKING, Any
10
@@ -48,6 +48,11 @@ from vllm.v1.simple_kv_offload.metadata import (
48 SimpleCPUOffloadMetadata,
49 SimpleCPUOffloadWorkerMetadata,
50 )
51 +from vllm.v1.simple_kv_offload.metrics import (
52 + OUTCOME_TO_FIELD,
53 + MetricName,
54 + SimpleCPUOffloadStats,
55 +)
56
57 if TYPE_CHECKING:
58 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -118,7 +123,8 @@ class BoundaryStoreStats:
123 symptom is a lower cache hit rate with nothing in the logs, so each decline
124 reason is counted and exposed through
125 ``SimpleCPUOffloadConnector.get_boundary_store_stats()``. Reset by
121 - ``reset()``; not otherwise cleared.
126 + ``reset()``, which carries any not-yet-drained deltas; not otherwise
127 + cleared.
128 """
129
130 published: int = 0
@@ -143,6 +149,7 @@ class SimpleCPUOffloadScheduler:
149 hash_block_size: int,
150 lazy_offload: bool = False,
151 disk_capacity_bytes: int = 0,
152 + use_page_cache: bool = False,
153 ):
154 self.vllm_config = vllm_config
155 self.kv_cache_config = kv_cache_config
@@ -169,6 +176,13 @@ class SimpleCPUOffloadScheduler:
176 self.cpu_kv_cache_config.prefix_cacheable_group_ids
177 )
178 self.kv_event_medium = MEDIUM_STORAGE if disk_capacity_bytes > 0 else MEDIUM_CPU
179 + # Must track kv_event_medium.
180 + self._info_labelvalues: tuple[str, ...] = (
181 + "disk" if disk_capacity_bytes > 0 else "cpu",
182 + str(use_page_cache).lower(),
183 + str(lazy_offload).lower(),
184 + str(self.num_cpu_blocks),
185 + )
186 # Find the full attention kv group for prefix cache matching.
187 self.fa_gidx = -1
188 for g_idx, g in enumerate(self.cpu_kv_cache_config.kv_cache_groups):
@@ -251,6 +265,9 @@ class SimpleCPUOffloadScheduler:
265 self._load_event_counter: int = 0
266 self._store_event_counter: int = 0
267 self.boundary_store_stats = BoundaryStoreStats()
268 + # Interval stats state drained by get_stats()
269 + self._boundary_stats_snapshot = BoundaryStoreStats()
270 + self._interval_load_blocks_completed = 0
271
272 # For TP/PP: track partial store completions across steps.
273 # Events must be reported by all world_size workers before considered complete.
@@ -951,6 +968,42 @@ class SimpleCPUOffloadScheduler:
968 def get_boundary_store_stats(self) -> BoundaryStoreStats:
969 return replace(self.boundary_store_stats)
970
971 + def get_stats(self) -> SimpleCPUOffloadStats:
972 + """Drain per-step stats for the connector's stats hooks."""
973 + stats = SimpleCPUOffloadStats()
974 +
975 + current = self.boundary_store_stats
976 + for outcome, field_name in OUTCOME_TO_FIELD.items():
977 + delta = getattr(current, field_name) - getattr(
978 + self._boundary_stats_snapshot, field_name
979 + )
980 + if delta > 0:
981 + stats.increase_counter(MetricName.SAVE_OUTCOMES, delta, (outcome,))
982 + self._boundary_stats_snapshot = replace(current)
983 +
984 + if self._interval_load_blocks_completed:
985 + stats.increase_counter(
986 + MetricName.LOAD_BLOCKS,
987 + self._interval_load_blocks_completed,
988 + )
989 + self._interval_load_blocks_completed = 0
990 +
991 + stats.set_gauge(
992 + MetricName.USED_BLOCKS,
993 + self.num_cpu_blocks - self.cpu_block_pool.get_num_free_blocks(),
994 + )
995 + pending = sum(
996 + len(t.cpu_block_ids)
997 + for t in (
998 + *self._store_event_to_blocks.values(),
999 + *self._pending_finished_stores,
1000 + *self._abandoned_store_event_to_blocks.values(),
1001 + )
1002 + )
1003 + stats.set_gauge(MetricName.PENDING_STORE_BLOCKS, pending)
1004 + stats.set_gauge(MetricName.INFO, 1, self._info_labelvalues)
1005 + return stats
1006 +
1007 def update_connector_output(self, connector_output: KVConnectorOutput) -> None:
1008 """Handle async transfer completions from worker.
1009
@@ -961,7 +1014,9 @@ class SimpleCPUOffloadScheduler:
1014 """
1015 # --- Load completions ---
1016 for req_id in list(connector_output.finished_recving or []):
964 - self._cleanup_load_request(req_id)
1017 + completed_blocks = self._cleanup_load_request(req_id)
1018 + if completed_blocks:
1019 + self._interval_load_blocks_completed += completed_blocks
1020
1021 # --- Store completions ---
1022 meta = connector_output.kv_connector_worker_meta
@@ -1266,18 +1321,22 @@ class SimpleCPUOffloadScheduler:
1321 if blocks_to_free:
1322 self.cpu_block_pool.free_blocks(blocks_to_free)
1323
1269 - def _cleanup_load_request(self, req_id: str) -> None:
1324 + def _cleanup_load_request(self, req_id: str) -> int:
1325 """Release all load resources for a request.
1326
1327 Shared between request_finished() and update_connector_output() paths.
1328 Removes the request from _reqs_to_load, cleans up event mappings,
1329 and frees CPU/GPU touch refs.
1330 +
1331 + Returns the number of blocks in the load if it had been issued to
1332 + the worker, else 0 (never-issued loads are not counted as completed
1333 + by the caller).
1334 """
1335 state = self._reqs_to_load.pop(req_id, None)
1336 if state is None:
1337 state = self._abandoned_reqs_to_load.pop(req_id, None)
1338 if state is None:
1280 - return
1339 + return 0
1340 # Remove from load event mapping (only this req, not whole event)
1341 if state.load_event is not None:
1342 reqs = self._load_event_to_reqs.get(state.load_event)
@@ -1300,6 +1359,10 @@ class SimpleCPUOffloadScheduler:
1359 for bid in state.transfer_meta.gpu_block_ids
1360 )
1361
1362 + if state.load_event is None:
1363 + return 0
1364 + return len(state.transfer_meta.gpu_block_ids)
1365 +
1366 def _cleanup_store_request(self, req_id: str) -> None:
1367 """Release store metadata for a request.
1368
@@ -1359,7 +1422,20 @@ class SimpleCPUOffloadScheduler:
1422 if event_idx in self._abandoned_store_event_to_blocks
1423 }
1424 self._cursor = None
1362 - self.boundary_store_stats = BoundaryStoreStats()
1425 + # Seed the fresh counters with outcomes since the last drain so a
1426 + # mid-interval reset() does not drop them from the next report.
1427 + undrained = BoundaryStoreStats(
1428 + **{
1429 + f.name: max(
1430 + 0,
1431 + getattr(self.boundary_store_stats, f.name)
1432 + - getattr(self._boundary_stats_snapshot, f.name),
1433 + )
1434 + for f in fields(BoundaryStoreStats)
1435 + }
1436 + )
1437 + self.boundary_store_stats = undrained
1438 + self._boundary_stats_snapshot = BoundaryStoreStats()
1439 # NOTE: _load_event_counter / _store_event_counter are not
1440 # reset as they are monotonic and must stay ahead of the workers
1441 # high-water marks to avoid event index collisions
vllm/v1/simple_kv_offload/metrics.py new
+284
@@ -0,0 +1,284 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Prometheus metrics for SimpleCPUOffloadConnector.
4 +
5 +The stats payload uses the same self-describing wire format as the
6 +offloading connector, so it survives IPC serialization without the
7 +metric classes on the receiving side.
8 +"""
9 +
10 +from dataclasses import dataclass
11 +from typing import Any
12 +
13 +from vllm.config import VllmConfig
14 +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
15 + KVConnectorPromMetrics,
16 + KVConnectorStats,
17 + PromMetric,
18 + PromMetricT,
19 +)
20 +
21 +
22 +class MetricName:
23 + """Metric names for SimpleCPUOffloadConnector."""
24 +
25 + SAVE_OUTCOMES = "vllm:simple_kv_offload_save_outcomes_total"
26 + LOAD_BLOCKS = "vllm:simple_kv_offload_load_blocks_total"
27 + USED_BLOCKS = "vllm:simple_kv_offload_used_blocks"
28 + PENDING_STORE_BLOCKS = "vllm:simple_kv_offload_pending_store_blocks"
29 + INFO = "vllm:simple_kv_offload_info"
30 +
31 +
32 +# outcome label value -> BoundaryStoreStats field name.
33 +# ``published`` is the sum of all outcomes and is not labeled.
34 +OUTCOME_TO_FIELD = {
35 + "stored": "stored",
36 + "dropped_cpu_full": "dropped_cpu_full",
37 + "dropped_request_gone": "dropped_request_gone",
38 + "dropped_null_block": "dropped_null_block",
39 + "dropped_not_hashed": "dropped_not_hashed",
40 + "skipped_already_cached": "skipped_already_cached",
41 + "skipped_in_flight": "skipped_in_flight",
42 +}
43 +
44 +INFO_LABELS = ["backend", "page_cache", "lazy_offload", "capacity_blocks"]
45 +
46 +
47 +class _MetricType:
48 + """Type tags embedded in the serialized stats payload."""
49 +
50 + COUNTER = "counter"
51 + GAUGE = "gauge"
52 +
53 +
54 +class _StatsKey:
55 + """Top-level keys in the serialized stats dict."""
56 +
57 + # Maps metric name -> _MetricType value
58 + TYPES = "types"
59 + # Maps metric name -> {label values tuple -> observed value}
60 + DATA = "data"
61 +
62 +
63 +@dataclass
64 +class SimpleCPUOffloadStats(KVConnectorStats):
65 + """Stats payload for SimpleCPUOffloadConnector.
66 +
67 + The ``data`` dict is structured as::
68 +
69 + {
70 + _StatsKey.TYPES: {name: _MetricType.*, ...},
71 + _StatsKey.DATA: {name: {labelvalues: value, ...}, ...},
72 + }
73 +
74 + Counter values are summed when aggregating, gauge values keep the
75 + latest snapshot, and unlabeled metrics use ``()`` as their
76 + labelvalues tuple.
77 + """
78 +
79 + def __post_init__(self):
80 + if _StatsKey.DATA not in self.data:
81 + self.reset()
82 +
83 + def reset(self):
84 + self.data: dict[str, Any] = {
85 + _StatsKey.TYPES: {},
86 + _StatsKey.DATA: {},
87 + }
88 +
89 + @property
90 + def _types(self) -> dict[str, str]:
91 + return self.data[_StatsKey.TYPES]
92 +
93 + @property
94 + def _values(self) -> dict[str, Any]:
95 + return self.data[_StatsKey.DATA]
96 +
97 + def aggregate(self, other: "KVConnectorStats") -> "KVConnectorStats":
98 + if other.is_empty():
99 + return self
100 + assert isinstance(other, SimpleCPUOffloadStats)
101 + for key, other_label_values in other._values.items():
102 + type_str = other._types.get(key)
103 + if type_str is None:
104 + raise AssertionError(f"Unknown simple KV offload stats key: {key}")
105 + self._types.setdefault(key, type_str)
106 + current_label_values = self._values.setdefault(key, {})
107 + for labelvalues, value in other_label_values.items():
108 + if type_str == _MetricType.COUNTER:
109 + assert isinstance(value, int | float)
110 + current_label_values[labelvalues] = (
111 + current_label_values.get(labelvalues, 0) + value
112 + )
113 + elif type_str == _MetricType.GAUGE:
114 + assert isinstance(value, int | float)
115 + current_label_values[labelvalues] = value
116 + else:
117 + raise AssertionError(
118 + f"Unknown metric type '{type_str}' for key: {key}"
119 + )
120 + return self
121 +
122 + def reduce(self) -> dict[str, int | float]:
123 + """Reduce the interval observations to a flat summary for logging."""
124 + return_dict: dict[str, int | float] = {}
125 + for key, label_value_map in self._values.items():
126 + if key == MetricName.INFO:
127 + # Constant config fingerprint; exclude from the log summary.
128 + continue
129 + type_str = self._types.get(key)
130 + if type_str is None:
131 + raise AssertionError(f"Unknown simple KV offload stats key: {key}")
132 + for labelvalues, value in label_value_map.items():
133 + key_with_labels = f"{key}:{labelvalues}" if labelvalues else key
134 + if type_str in (_MetricType.COUNTER, _MetricType.GAUGE):
135 + assert isinstance(value, int | float)
136 + return_dict[key_with_labels] = value
137 + else:
138 + raise AssertionError(
139 + f"Unknown metric type '{type_str}' for key: {key}"
140 + )
141 + return return_dict
142 +
143 + def is_empty(self) -> bool:
144 + return not self.data.get(_StatsKey.DATA)
145 +
146 + def increase_counter(
147 + self,
148 + counter_name: str,
149 + counter_increase_value: int | float = 1,
150 + labelvalues: tuple[str, ...] = (),
151 + ) -> None:
152 + """Increase a counter on the stats payload."""
153 + self._types.setdefault(counter_name, _MetricType.COUNTER)
154 + counter_values = self._values.setdefault(counter_name, {})
155 + counter_values[labelvalues] = (
156 + counter_values.get(labelvalues, 0) + counter_increase_value
157 + )
158 +
159 + def set_gauge(
160 + self,
161 + gauge_name: str,
162 + gauge_value: int | float,
163 + labelvalues: tuple[str, ...] = (),
164 + ) -> None:
165 + """Set a gauge snapshot on the stats payload."""
166 + self._types.setdefault(gauge_name, _MetricType.GAUGE)
167 + gauge_values = self._values.setdefault(gauge_name, {})
168 + gauge_values[labelvalues] = gauge_value
169 +
170 +
171 +class SimpleCPUOffloadPromMetrics(KVConnectorPromMetrics):
172 + """Prometheus registration and recording for SimpleCPUOffloadConnector."""
173 +
174 + def __init__(
175 + self,
176 + vllm_config: VllmConfig,
177 + metric_types: dict[type[PromMetric], type[PromMetricT]],
178 + labelnames: list[str],
179 + per_engine_labelvalues: dict[int, list[object]],
180 + ):
181 + super().__init__(vllm_config, metric_types, labelnames, per_engine_labelvalues)
182 + # (engine_idx, metric_name, labelvalues) -> bound child metric
183 + self._metrics: dict[tuple[int, str, tuple[str, ...]], PromMetricT] = {}
184 +
185 + self._counter_save_outcomes = self._counter_cls(
186 + name="vllm:simple_kv_offload_save_outcomes_total",
187 + documentation=(
188 + "Store admission decisions in SimpleCPUOffloadConnector, by "
189 + "outcome. Outcomes classify eager-mode boundary hand-off "
190 + "stores; lazy-mode stores are not classified."
191 + ),
192 + labelnames=self._labelnames + ["outcome"],
193 + )
194 + self._counter_load_blocks = self._counter_cls(
195 + name="vllm:simple_kv_offload_load_blocks_total",
196 + documentation=(
197 + "KV blocks restored to GPU from the offload pool, counted "
198 + "when a load finishes. This is prefill work avoided "
199 + "through offload cache hits."
200 + ),
201 + labelnames=self._labelnames,
202 + )
203 + self._gauge_used_blocks = self._gauge_cls(
204 + name="vllm:simple_kv_offload_used_blocks",
205 + documentation=(
206 + "Offload-pool blocks currently pinned by in-flight "
207 + "transfers or cache hits (capacity minus free). Evictable "
208 + "cached blocks are not counted; capacity_blocks is on "
209 + "vllm:simple_kv_offload_info."
210 + ),
211 + labelnames=self._labelnames,
212 + multiprocess_mode="mostrecent",
213 + )
214 + self._gauge_pending_store_blocks = self._gauge_cls(
215 + name="vllm:simple_kv_offload_pending_store_blocks",
216 + documentation=(
217 + "Offload blocks in pending stores: queued for dispatch, "
218 + "issued to workers, or awaiting release after a cache "
219 + "reset. A persistently growing value indicates a stuck "
220 + "transfer."
221 + ),
222 + labelnames=self._labelnames,
223 + multiprocess_mode="mostrecent",
224 + )
225 + self._gauge_info = self._gauge_cls(
226 + name="vllm:simple_kv_offload_info",
227 + documentation=(
228 + "SimpleCPUOffloadConnector deployment facts. Value is "
229 + "always 1. backend is cpu or disk; page_cache and "
230 + "lazy_offload are true/false; capacity_blocks is the "
231 + "offload pool size."
232 + ),
233 + labelnames=self._labelnames + INFO_LABELS,
234 + multiprocess_mode="mostrecent",
235 + )
236 +
237 + self._defs: dict[str, PromMetricT] = {
238 + MetricName.SAVE_OUTCOMES: self._counter_save_outcomes,
239 + MetricName.LOAD_BLOCKS: self._counter_load_blocks,
240 + MetricName.USED_BLOCKS: self._gauge_used_blocks,
241 + MetricName.PENDING_STORE_BLOCKS: self._gauge_pending_store_blocks,
242 + MetricName.INFO: self._gauge_info,
243 + }
244 + self._num_label_values: dict[str, int] = {
245 + MetricName.SAVE_OUTCOMES: 1,
246 + MetricName.LOAD_BLOCKS: 0,
247 + MetricName.USED_BLOCKS: 0,
248 + MetricName.PENDING_STORE_BLOCKS: 0,
249 + MetricName.INFO: len(INFO_LABELS),
250 + }
251 +
252 + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0):
253 + """Observe transfer statistics."""
254 + metric_types = transfer_stats_data.get(_StatsKey.TYPES, {})
255 + metric_data = transfer_stats_data.get(_StatsKey.DATA, {})
256 + for key, label_value_map in metric_data.items():
257 + prom_metric = self._defs.get(key)
258 + if prom_metric is None:
259 + raise AssertionError(f"Unknown simple KV offload stats key: {key}")
260 + type_str = metric_types.get(key)
261 + for labelvalues, value in label_value_map.items():
262 + if len(labelvalues) != self._num_label_values[key]:
263 + raise AssertionError(
264 + f"Metric {key} expects "
265 + f"{self._num_label_values[key]} labels, "
266 + f"got {len(labelvalues)}"
267 + )
268 + child = self._metrics.get((engine_idx, key, labelvalues))
269 + if child is None:
270 + engine_labelvalues = self.per_engine_labelvalues[engine_idx]
271 + child = prom_metric.labels(
272 + *(engine_labelvalues + list(labelvalues))
273 + )
274 + self._metrics[(engine_idx, key, labelvalues)] = child
275 + if type_str == _MetricType.COUNTER:
276 + assert isinstance(value, int | float)
277 + child.inc(value)
278 + elif type_str == _MetricType.GAUGE:
279 + assert isinstance(value, int | float)
280 + child.set(value)
281 + else:
282 + raise AssertionError(
283 + f"Unknown metric type '{type_str}' for key: {key}"
284 + )