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

Add tiering offloading metrics (#48798)

Signed-off-by: srinivas_oo7 <sklinkedin0120@gmail.com> Co-authored-by: srinivas_oo7 <sklinkedin0120@gmail.com>

Srinivas Krovvidi committed Aug 9, 2026 at 21:23 UTC a123159f7ab0dbacb4d8f45cdeec3cb366982e50
16 files changed +1031 -237
tests/v1/kv_connector/unit/offloading_connector/test_metrics.py
+34 -4
@@ -25,6 +25,7 @@ from vllm.v1.kv_offload.base import (
25 OffloadingHistogramMetadata,
26 )
27 from vllm.v1.kv_offload.factory import OffloadingSpecFactory
28 +from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics
29
30 LOAD_BYTES = _TransferMetricName.LOAD_BYTES
31 LOAD_TIME = _TransferMetricName.LOAD_TIME
@@ -101,10 +102,15 @@ class _FakeMetric:
102
103
104 class _FakeVllmConfig:
104 - def __init__(self, store_threshold: int = 2):
105 - self.kv_transfer_config = SimpleNamespace(
106 - kv_connector_extra_config={"store_threshold": store_threshold}
107 - )
105 + def __init__(
106 + self,
107 + store_threshold: int = 2,
108 + extra_config: dict[str, Any] | None = None,
109 + ):
110 + config = {"store_threshold": store_threshold}
111 + if extra_config is not None:
112 + config.update(extra_config)
113 + self.kv_transfer_config = SimpleNamespace(kv_connector_extra_config=config)
114
115
116 def _spec_cls_with_metric_definitions(
@@ -679,6 +685,30 @@ def test_prom_metrics_uses_configured_manager_metrics():
685 assert STORES_SKIPPED not in prom_metrics._offloading_metric_metadata
686
687
688 +def test_prom_metrics_registers_tiering_metrics_from_spec():
689 + prom_metrics = OffloadPromMetrics(
690 + vllm_config=_FakeVllmConfig(
691 + store_threshold=0,
692 + extra_config={
693 + "spec_name": "TieringOffloadingSpec",
694 + "secondary_tiers": [],
695 + },
696 + ), # type: ignore[arg-type]
697 + metric_types={
698 + Gauge: _FakeMetric,
699 + Counter: _FakeMetric,
700 + Histogram: _FakeMetric,
701 + },
702 + labelnames=["model_name", "engine"],
703 + per_engine_labelvalues={0: ["model", "0"]},
704 + )
705 +
706 + metric = prom_metrics._offloading_metric_defs[
707 + TieringOffloadingMetrics.BLOCK_QUERIES
708 + ]
709 + assert metric.kwargs["labelnames"] == ["model_name", "engine", "tier"]
710 +
711 +
712 def test_aggregate_into_empty_stats():
713 """Aggregating non-empty stats into a fresh (empty) stats object works."""
714 empty = OffloadingConnectorStats()
tests/v1/kv_offload/tiering/p2p/test_manager.py
+3 -3
@@ -16,7 +16,7 @@ import numpy as np
16 import pytest
17
18 from vllm.v1.kv_offload.base import LookupResult, ReqContext, ScheduleEndContext
19 -from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
19 +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
20 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
21 from vllm.v1.kv_offload.tiering.p2p.manager import (
22 _UNBOUND_STORE_TIMEOUT_S,
@@ -86,12 +86,12 @@ def _job_metadata(
86 keys: list[bytes] | None = None,
87 block_ids: list[int] | None = None,
88 kv_params: dict | None = None,
89 -) -> JobMetadata:
89 +) -> TransferJob:
90 if keys is None:
91 keys = [b"key1"]
92 if block_ids is None:
93 block_ids = list(range(len(keys)))
94 - return JobMetadata(
94 + return TransferJob(
95 job_id=job_id,
96 keys=keys,
97 block_ids=np.array(block_ids),
tests/v1/kv_offload/tiering/p2p/test_sessions.py
+3 -3
@@ -24,7 +24,7 @@ from vllm.v1.kv_offload.base import (
24 ReqContext,
25 RequestOffloadingContext,
26 )
27 -from vllm.v1.kv_offload.tiering.base import JobMetadata
27 +from vllm.v1.kv_offload.tiering.base import TransferJob
28 from vllm.v1.kv_offload.tiering.p2p.session import (
29 LoadResult,
30 P2PSession,
@@ -256,13 +256,13 @@ class FakeParent:
256 self,
257 keys: Sequence[OffloadKey],
258 ctx: ReqContext,
259 - ) -> JobMetadata:
259 + ) -> TransferJob:
260 keys_list = list(keys)
261 self.calls.append(("create_store_job", tuple(keys_list), ctx.req_id))
262 block_ids = np.array([self.stored[k] for k in keys_list], dtype=np.int32)
263 job_id = self._next_job_id
264 self._next_job_id += 1
265 - return JobMetadata(
265 + return TransferJob(
266 job_id=job_id,
267 keys=keys_list,
268 block_ids=block_ids,
tests/v1/kv_offload/tiering/test_fs_tier.py
+3 -3
@@ -35,7 +35,7 @@ from vllm.v1.kv_offload.config import (
35 OffloadingModelConfig,
36 OffloadingParallelConfig,
37 )
38 -from vllm.v1.kv_offload.tiering.base import JobMetadata
38 +from vllm.v1.kv_offload.tiering.base import TransferJob
39 from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
40 from vllm.v1.kv_offload.tiering.fs.manager import (
41 FileSystemTierManager,
@@ -105,10 +105,10 @@ def make_job(
105 keys: list[OffloadKey],
106 block_ids: list[int] | None = None,
107 is_promotion: bool = False,
108 -) -> JobMetadata:
108 +) -> TransferJob:
109 if block_ids is None:
110 block_ids = list(range(len(keys)))
111 - return JobMetadata(
111 + return TransferJob(
112 job_id=job_id,
113 keys=keys,
114 block_ids=np.array(block_ids, dtype=np.int64),
tests/v1/kv_offload/tiering/test_metrics.py new
+226
@@ -0,0 +1,226 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterable
5 +
6 +import numpy as np
7 +import pytest
8 +
9 +from vllm.v1.kv_offload.base import LookupResult, ReqContext, make_offload_key
10 +from vllm.v1.kv_offload.tiering.base import (
11 + JobResult,
12 + TieringOffloadingMetrics,
13 + TransferJob,
14 +)
15 +from vllm.v1.kv_offload.tiering.manager import JobMetadata
16 +from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker
17 +
18 +_CTX = ReqContext(req_id="test")
19 +
20 +
21 +def to_keys(int_ids: Iterable[int]):
22 + return [make_offload_key(str(i).encode(), 0) for i in int_ids]
23 +
24 +
25 +def test_tiering_metrics_tracker_records_lookup_metrics():
26 + tracker = TieringMetricsTracker(
27 + tier_types=["fs", "p2p"],
28 + num_primary_blocks=5,
29 + primary_block_size=16,
30 + )
31 + tracker.on_new_request(_CTX)
32 +
33 + key = to_keys([1])[0]
34 + tracker.on_lookup(
35 + _CTX,
36 + key,
37 + tracker.primary_tier_label,
38 + LookupResult.MISS,
39 + lookup_duration=0.0,
40 + )
41 + tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.MISS, 0.0)
42 + tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0)
43 + tracker.on_lookup(_CTX, key, tracker.tier_label(1), LookupResult.HIT, 0.0)
44 +
45 + stats = tracker.take_stats()
46 + assert stats is not None
47 + values = stats.data["data"]
48 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1
49 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1
50 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("2:p2p",)] == 1
51 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("2:p2p",)] == 1
52 + assert ("1:fs",) not in values[TieringOffloadingMetrics.BLOCK_HITS]
53 +
54 +
55 +def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation():
56 + tracker = TieringMetricsTracker(
57 + tier_types=["fs"],
58 + num_primary_blocks=5,
59 + primary_block_size=16,
60 + )
61 + tracker.on_new_request(_CTX)
62 + key = to_keys([1])[0]
63 +
64 + tracker.on_lookup(
65 + _CTX,
66 + key,
67 + tracker.primary_tier_label,
68 + LookupResult.MISS,
69 + lookup_duration=0.0,
70 + )
71 + tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0)
72 +
73 + stats = tracker.take_stats()
74 + assert stats is not None
75 + values = stats.data["data"]
76 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("0:primary",)] == 1
77 + assert values[TieringOffloadingMetrics.BLOCK_QUERIES][("1:fs",)] == 1
78 + assert values[TieringOffloadingMetrics.BLOCK_HITS][("1:fs",)] == 1
79 +
80 + tracker.on_request_allocated(_CTX)
81 + tracker.on_lookup(
82 + _CTX,
83 + key,
84 + tracker.primary_tier_label,
85 + LookupResult.MISS,
86 + lookup_duration=0.0,
87 + )
88 + tracker.on_lookup(_CTX, key, tracker.tier_label(0), LookupResult.HIT, 0.0)
89 +
90 + stats = tracker.take_stats()
91 + assert stats is not None
92 + values = stats.data["data"]
93 + assert TieringOffloadingMetrics.BLOCK_QUERIES not in values
94 + assert TieringOffloadingMetrics.BLOCK_HITS not in values
95 +
96 +
97 +def test_tiering_metrics_tracker_records_finished_job_metrics():
98 + tracker = TieringMetricsTracker(
99 + tier_types=["fs"],
100 + num_primary_blocks=5,
101 + primary_block_size=16,
102 + )
103 + cascade_key, promotion_key, failed_cascade_key, failed_promotion_key = to_keys(
104 + range(4)
105 + )
106 + jobs = [
107 + JobMetadata(
108 + TransferJob(0, [cascade_key], np.array([0]), False, _CTX),
109 + 0,
110 + ),
111 + JobMetadata(
112 + TransferJob(1, [promotion_key], np.array([1]), True, _CTX),
113 + 0,
114 + ),
115 + JobMetadata(
116 + TransferJob(2, [failed_cascade_key], np.array([2]), False, _CTX),
117 + 0,
118 + ),
119 + JobMetadata(
120 + TransferJob(3, [failed_promotion_key], np.array([3]), True, _CTX),
121 + 0,
122 + ),
123 + ]
124 + for job in jobs:
125 + tracker.on_job_registered(job)
126 +
127 + tracker.on_job_finished(
128 + jobs[0], JobResult(job_id=0, success=True, transfer_time=0.5)
129 + )
130 + tracker.on_job_finished(
131 + jobs[1], JobResult(job_id=1, success=True, transfer_time=0.25)
132 + )
133 + tracker.on_job_finished(jobs[2], JobResult(job_id=2, success=False))
134 + tracker.on_job_finished(jobs[3], JobResult(job_id=3, success=False))
135 +
136 + stats = tracker.take_stats()
137 + assert stats is not None
138 + values = stats.data["data"]
139 + label = ("1:fs",)
140 + assert values[TieringOffloadingMetrics.WRITE_BYTES][label] == 16
141 + assert values[TieringOffloadingMetrics.WRITE_TIME][label] == 0.5
142 + assert values[TieringOffloadingMetrics.READ_BYTES][label] == 16
143 + assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.25
144 + assert values[TieringOffloadingMetrics.CASCADE_JOB_FAILURES][label] == 1
145 + assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1
146 + tracker.assert_idle()
147 +
148 +
149 +def test_tiering_metrics_tracker_records_partial_promotion_success_bytes():
150 + tracker = TieringMetricsTracker(
151 + tier_types=["fs"],
152 + num_primary_blocks=5,
153 + primary_block_size=16,
154 + )
155 + keys = to_keys(range(3))
156 + job = JobMetadata(
157 + TransferJob(0, keys, np.array([0, 1, 2]), True, _CTX),
158 + 0,
159 + )
160 + tracker.on_job_registered(job)
161 +
162 + tracker.on_job_finished(
163 + job,
164 + JobResult(
165 + job_id=0,
166 + success=False,
167 + successful_keys=(keys[0], keys[2]),
168 + transfer_time=0.5,
169 + ),
170 + )
171 +
172 + stats = tracker.take_stats()
173 + assert stats is not None
174 + values = stats.data["data"]
175 + label = ("1:fs",)
176 + assert values[TieringOffloadingMetrics.READ_BYTES][label] == 32
177 + assert values[TieringOffloadingMetrics.READ_TIME][label] == 0.5
178 + assert values[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES][label] == 1
179 + tracker.assert_idle()
180 +
181 +
182 +def test_tiering_metrics_tracker_reports_active_job_and_primary_usage_gauges():
183 + tracker = TieringMetricsTracker(
184 + tier_types=["fs", "p2p"],
185 + num_primary_blocks=6,
186 + primary_block_size=16,
187 + )
188 + fs_job = JobMetadata(
189 + TransferJob(0, to_keys([0, 1]), np.array([0, 1]), False, _CTX),
190 + 0,
191 + )
192 + p2p_job = JobMetadata(
193 + TransferJob(1, to_keys([2, 3, 4]), np.array([2, 3, 4]), True, _CTX),
194 + 1,
195 + )
196 + tracker.on_job_registered(fs_job)
197 + tracker.on_job_registered(p2p_job)
198 +
199 + stats = tracker.take_stats()
200 + assert stats is not None
201 + values = stats.data["data"]
202 + fs_label = ("1:fs",)
203 + p2p_label = ("2:p2p",)
204 + assert values[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC][
205 + fs_label
206 + ] == pytest.approx(2 / 6)
207 + assert values[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS][fs_label] == 1
208 + assert values[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC][
209 + p2p_label
210 + ] == pytest.approx(3 / 6)
211 + assert values[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS][p2p_label] == 1
212 +
213 +
214 +def test_tiering_metrics_tracker_records_promotion_allocation_failures():
215 + tracker = TieringMetricsTracker(
216 + tier_types=["fs"],
217 + num_primary_blocks=1,
218 + primary_block_size=16,
219 + )
220 +
221 + tracker.on_promotion_allocation_failure()
222 +
223 + stats = tracker.take_stats()
224 + assert stats is not None
225 + values = stats.data["data"]
226 + assert values[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES][()] == 1
tests/v1/kv_offload/tiering/test_obj_tier.py
+6 -3
@@ -34,7 +34,7 @@ from vllm.v1.kv_offload.config import (
34 OffloadingModelConfig,
35 OffloadingParallelConfig,
36 )
37 -from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
37 +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
38 from vllm.v1.kv_offload.tiering.manager import (
39 CPUPrimaryTierOffloadingManager,
40 TieringOffloadingManager,
@@ -105,10 +105,10 @@ def make_job(
105 job_id: int,
106 keys: list[OffloadKey],
107 block_ids: list[int] | None = None,
108 -) -> JobMetadata:
108 +) -> TransferJob:
109 if block_ids is None:
110 block_ids = list(range(len(keys)))
111 - return JobMetadata(
111 + return TransferJob(
112 job_id=job_id,
113 keys=keys,
114 block_ids=np.array(block_ids, dtype=np.int64),
@@ -201,6 +201,9 @@ class MockNixlAgent:
201 def release_dlist_handle(self, handle):
202 pass
203
204 + def get_xfer_telemetry(self, handle):
205 + return SimpleNamespace(xferDuration=1000)
206 +
207 def _query_memory(self, queries, mem_type, agent_name):
208 return [object() if q[3] in self._stored_obj_keys else None for q in queries]
209
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+149 -63
@@ -29,6 +29,7 @@ from vllm.v1.kv_offload.base import (
29 Medium,
30 OffloadingCounterMetadata,
31 OffloadingEvent,
32 + OffloadingGaugeMetadata,
33 OffloadKey,
34 OffloadPolicy,
35 ReqContext,
@@ -39,10 +40,10 @@ from vllm.v1.kv_offload.base import (
40 make_offload_key,
41 )
42 from vllm.v1.kv_offload.tiering.base import (
42 - JobMetadata,
43 JobResult,
44 SecondaryTierManager,
45 TieringOffloadingMetrics,
46 + TransferJob,
47 )
48 from vllm.v1.kv_offload.tiering.example.manager import ExampleSecondaryTierManager
49 from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory
@@ -85,6 +86,16 @@ def count_hits(manager, keys: list[OffloadKey]) -> int | None:
86 return count
87
88
89 +def histogram_count_key(metric_name: str, labelvalues: tuple[str, ...]) -> str:
90 + return f"{metric_name}:{labelvalues}_count"
91 +
92 +
93 +def has_histogram_count(reduced: dict[str, int | float], metric_name: str) -> bool:
94 + return any(
95 + key.startswith(f"{metric_name}:") and key.endswith("_count") for key in reduced
96 + )
97 +
98 +
99 class MetricsSecondaryTierManager(SecondaryTierManager):
100 """Test-only secondary tier that declares and emits one labeled metric."""
101
@@ -102,18 +113,26 @@ class MetricsSecondaryTierManager(SecondaryTierManager):
113 def __init__(self, *args, **kwargs):
114 super().__init__(*args, **kwargs)
115 self.stats: OffloadingConnectorStats | None = None
116 + self.lookup_result = LookupResult.MISS
117 + self.finished_jobs: list[JobResult] = []
118 + self.submitted_loads: list[TransferJob] = []
119 + self.submitted_stores: list[TransferJob] = []
120
106 - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:
107 - return False
121 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
122 + return self.lookup_result
123
109 - def submit_store(self, job_metadata: JobMetadata) -> None:
124 + def submit_store(self, job_metadata: TransferJob) -> None:
125 + self.submitted_stores.append(job_metadata)
126 return
127
112 - def submit_load(self, job_metadata: JobMetadata) -> None:
128 + def submit_load(self, job_metadata: TransferJob) -> None:
129 + self.submitted_loads.append(job_metadata)
130 return
131
132 def get_finished_jobs(self) -> Iterable[JobResult]:
116 - return ()
133 + results = self.finished_jobs
134 + self.finished_jobs = []
135 + return results
136
137 def drain_jobs(self) -> None:
138 return
@@ -141,6 +160,24 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch):
160 metadata = metrics[MetricsSecondaryTierManager.MY_TIER_METRIC]
161 assert metadata.documentation == "Number of bytes served by the test tier."
162 assert metadata.labelnames == ("tier",)
163 + assert metrics[TieringOffloadingMetrics.BLOCK_QUERIES].labelnames == ("tier",)
164 + assert metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY].labelnames == ("tier",)
165 + assert metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY].labelnames == ("tier",)
166 + assert metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES].labelnames == (
167 + "tier",
168 + )
169 + assert isinstance(
170 + metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES],
171 + OffloadingCounterMetadata,
172 + )
173 + for metric_name in (
174 + TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC,
175 + TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC,
176 + TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS,
177 + TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS,
178 + ):
179 + assert isinstance(metrics[metric_name], OffloadingGaugeMetadata)
180 + assert metrics[metric_name].labelnames == ("tier",)
181
182
183 def test_tiering_manager_aggregates_secondary_stats():
@@ -261,12 +298,15 @@ class TestTieringOffloadingManager:
298 self._start_request()
299 # Register an in-flight promotion job for tier1 by hand.
300 job_id = self.manager._next_job_id()
264 - self.manager._transfer_jobs[job_id] = JobMetadata(
265 - job_id=job_id,
266 - keys=to_keys([1, 2]),
267 - block_ids=[0, 1],
268 - is_promotion=True,
269 - req_context=_CTX,
301 + self.manager._register_job(
302 + TransferJob(
303 + job_id=job_id,
304 + keys=to_keys([1, 2]),
305 + block_ids=[0, 1],
306 + is_promotion=True,
307 + req_context=_CTX,
308 + ),
309 + 0,
310 )
311 failed = JobResult(job_id=job_id, success=False)
312 with (
@@ -285,12 +325,15 @@ class TestTieringOffloadingManager:
325
326 self._start_request()
327 job_id = self.manager._next_job_id()
288 - self.manager._transfer_jobs[job_id] = JobMetadata(
289 - job_id=job_id,
290 - keys=to_keys([1]),
291 - block_ids=[0],
292 - is_promotion=True,
293 - req_context=_CTX,
328 + self.manager._register_job(
329 + TransferJob(
330 + job_id=job_id,
331 + keys=to_keys([1]),
332 + block_ids=[0],
333 + is_promotion=True,
334 + req_context=_CTX,
335 + ),
336 + 0,
337 )
338 ok = JobResult(job_id=job_id, success=True)
339 with (
@@ -473,7 +516,7 @@ class TestTieringOffloadingManager:
516 for block in blocks:
517 self.secondary_tier1.blocks[block] = True
518
476 - def submit_partial(job_metadata: JobMetadata) -> None:
519 + def submit_partial(job_metadata: TransferJob) -> None:
520 successful_keys = (
521 None
522 if successful_indices is None
@@ -500,7 +543,7 @@ class TestTieringOffloadingManager:
543 ] == expected_results
544
545 def test_lookup_reports_sync_delay_for_resolved_lookups(self, manager_setup):
503 - """Resolved lookups report one sync delay sample on allocation."""
546 + """Resolved lookups report one sync delay sample per tier and block."""
547 self._start_request()
548 blocks = to_keys(range(2))
549
@@ -508,22 +551,39 @@ class TestTieringOffloadingManager:
551 for block in blocks:
552 assert self.manager.lookup(block, _CTX) is LookupResult.MISS
553
511 - stats = self.manager.get_stats()
512 - if stats is not None:
513 - assert f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count" not in (
514 - stats.reduce()
515 - )
516 -
517 - self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id])
518 -
554 stats = self.manager.get_stats()
555 assert stats is not None
556 reduced = stats.reduce()
522 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1
523 - assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in reduced
557 + assert (
558 + reduced[
559 + histogram_count_key(
560 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("0:primary",)
561 + )
562 + ]
563 + == 2
564 + )
565 + assert (
566 + reduced[
567 + histogram_count_key(
568 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("1:example",)
569 + )
570 + ]
571 + == 2
572 + )
573 + assert (
574 + reduced[
575 + histogram_count_key(
576 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("2:example",)
577 + )
578 + ]
579 + == 2
580 + )
581 + assert not has_histogram_count(
582 + reduced, TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY
583 + )
584
525 - def test_lookup_reports_async_delay_across_promotion(self, manager_setup):
526 - """A new request reports async delay at schedule end."""
585 + def test_lookup_does_not_report_async_delay_for_promotion(self, manager_setup):
586 + """Promotion time is not included in lookup async delay."""
587 self._start_request()
588 block = to_keys(range(1))[0]
589 self.secondary_tier1.blocks[block] = True
@@ -531,13 +591,30 @@ class TestTieringOffloadingManager:
591 # First lookup finds the block in a secondary tier and defers.
592 assert self.manager.lookup(block, _CTX) is LookupResult.RETRY
593
534 - # The first scheduler step reports async delay for the new request.
594 + # Promotion is not treated as unresolved secondary lookup time.
595 self._simulate_on_schedule_end(new_req_ids=[_CTX.req_id])
596 stats = self.manager.get_stats()
597 assert stats is not None
598 reduced = stats.reduce()
539 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1
540 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count"] == 1
599 + assert (
600 + reduced[
601 + histogram_count_key(
602 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("0:primary",)
603 + )
604 + ]
605 + == 1
606 + )
607 + assert (
608 + reduced[
609 + histogram_count_key(
610 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("1:example",)
611 + )
612 + ]
613 + == 1
614 + )
615 + assert not has_histogram_count(
616 + reduced, TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY
617 + )
618
619 # Promotion completes on the next scheduler step.
620 self._simulate_on_schedule_end()
@@ -547,38 +624,51 @@ class TestTieringOffloadingManager:
624
625 stats = self.manager.get_stats()
626 if stats is not None:
550 - assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in (
551 - stats.reduce()
627 + assert not has_histogram_count(
628 + stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY
629 )
630
554 - def test_lookup_reports_async_delay_on_request_finish(self, manager_setup):
555 - """Never-allocated lookup delays flush at teardown."""
631 + def test_lookup_reports_async_delay_when_deferred_lookup_resolves(
632 + self, manager_setup
633 + ):
634 + """Async delay is observed when a RETRY lookup later resolves."""
635 ctx = ReqContext(req_id="req_lookup_finish")
636 self._start_request(ctx)
637 block = to_keys(range(1))[0]
559 - self.secondary_tier1.blocks[block] = True
638 + self.secondary_tier1.lookup = MagicMock(
639 + side_effect=[LookupResult.RETRY, LookupResult.HIT]
640 + )
641
561 - # Lookup finds the block in a secondary tier and defers.
642 + # First lookup is deferred by the secondary tier.
643 assert self.manager.lookup(block, ctx) is LookupResult.RETRY
644
564 - self._simulate_on_schedule_end()
645 stats = self.manager.get_stats()
646 if stats is not None:
567 - assert f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count" not in (
568 - stats.reduce()
569 - )
570 - assert f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count" not in (
571 - stats.reduce()
647 + assert not has_histogram_count(
648 + stats.reduce(), TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY
649 )
650
574 - # Request finishes before the deferred lookup is ever resolved.
575 - self.manager.on_request_finished(ctx)
651 + assert self.manager.lookup(block, ctx) is LookupResult.RETRY
652
653 stats = self.manager.get_stats()
654 assert stats is not None
655 reduced = stats.reduce()
580 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_SYNC_DELAY}_count"] == 1
581 - assert reduced[f"{TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY}_count"] == 1
656 + assert (
657 + reduced[
658 + histogram_count_key(
659 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, ("1:example",)
660 + )
661 + ]
662 + == 1
663 + )
664 + assert (
665 + reduced[
666 + histogram_count_key(
667 + TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, ("1:example",)
668 + )
669 + ]
670 + == 1
671 + )
672
673 def test_partial_lookup(self, manager_setup):
674 """Test lookup with partial hits."""
@@ -938,9 +1028,7 @@ class TestTieringOffloadingManager:
1028 ctx = ReqContext(req_id="req_policy_lifecycle_2")
1029 result = self.manager.on_new_request(ctx)
1030 assert result.policy == OffloadPolicy.REQUEST_LEVEL
941 - assert self.manager._req_state[ctx.req_id].request_level_tiers == {
942 - self.secondary_tier1
943 - }
1031 + assert self.manager._req_state[ctx.req_id].request_level_tiers == {0}
1032
1033 # Cleanup
1034 self.manager.on_request_finished(ctx)
@@ -996,13 +1084,13 @@ class TestTieringOffloadingManager:
1084 primary tier; pending submissions are dropped without being sent
1085 to the secondary tier. Active request state is retained."""
1086 # Cascade — populates primary blocks and leaves cascade jobs
999 - # in _transfer_jobs (the synchronous example tier has already
1087 + # in _jobs (the synchronous example tier has already
1088 # queued completions); reset_cache's drain loop will pick them up.
1089 blocks = to_keys(range(3))
1090 self._start_request()
1091 self.manager.prepare_store(blocks, _CTX)
1092 self.manager.complete_store(blocks, _CTX, success=True)
1005 - assert self.manager._transfer_jobs
1093 + assert self.manager._jobs
1094
1095 # Pending promotion submission (deferred — no on_schedule_end after
1096 # the lookup that staged it).
@@ -1020,9 +1108,7 @@ class TestTieringOffloadingManager:
1108 )
1109 rl_ctx = ReqContext(req_id="rl")
1110 self.manager.on_new_request(rl_ctx)
1023 - assert self.manager._req_state[rl_ctx.req_id].request_level_tiers == {
1024 - self.secondary_tier1
1025 - }
1111 + assert self.manager._req_state[rl_ctx.req_id].request_level_tiers == {0}
1112
1113 # Mark this step as already polled (reset_cache must clear it).
1114 self.manager._processed_jobs_this_step = True
@@ -1035,7 +1121,7 @@ class TestTieringOffloadingManager:
1121 self.manager.reset_cache()
1122
1123 # Orchestrator state cleared.
1038 - assert self.manager._transfer_jobs == {}
1124 + assert self.manager._jobs == {}
1125 assert self.manager._pending_load_submissions == {}
1126 assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id}
1127 assert self.manager._processed_jobs_this_step is False
@@ -1063,18 +1149,18 @@ class TestTieringOffloadingManager:
1149 wraps=self.secondary_tier2.drain_jobs
1150 )
1151
1066 - # Drive a cascade so a job lands in _transfer_jobs.
1152 + # Drive a cascade so a job lands in _jobs.
1153 blocks = to_keys(range(3))
1154 self._start_request()
1155 self.manager.prepare_store(blocks, _CTX)
1156 self.manager.complete_store(blocks, _CTX, success=True)
1071 - assert self.manager._transfer_jobs
1157 + assert self.manager._jobs
1158
1159 self.manager.reset_cache()
1160
1161 self.secondary_tier1.drain_jobs.assert_called_once()
1162 self.secondary_tier2.drain_jobs.assert_called_once()
1077 - assert self.manager._transfer_jobs == {}
1163 + assert self.manager._jobs == {}
1164
1165 @pytest.mark.parametrize(
1166 "load_tier_filter",
vllm/v1/kv_offload/tiering/base.py
+20 -4
@@ -39,10 +39,25 @@ class TieringOffloadingMetrics:
39
40 LOOKUP_SYNC_DELAY = "vllm:kv_offload_tiering_lookup_sync_delay_seconds"
41 LOOKUP_ASYNC_DELAY = "vllm:kv_offload_tiering_lookup_async_delay_seconds"
42 + READ_BYTES = "vllm:kv_offload_tiering_read_bytes"
43 + READ_TIME = "vllm:kv_offload_tiering_read_time"
44 + WRITE_BYTES = "vllm:kv_offload_tiering_write_bytes"
45 + WRITE_TIME = "vllm:kv_offload_tiering_write_time"
46 + PROMOTION_JOB_FAILURES = "vllm:kv_offload_tiering_promotion_job_failures"
47 + CASCADE_JOB_FAILURES = "vllm:kv_offload_tiering_cascade_job_failures"
48 + BLOCK_QUERIES = "vllm:kv_offload_tiering_block_queries"
49 + BLOCK_HITS = "vllm:kv_offload_tiering_block_hits"
50 + PRIMARY_WRITE_USAGE_PERC = "vllm:kv_offload_tiering_primary_write_usage_perc"
51 + PRIMARY_READ_USAGE_PERC = "vllm:kv_offload_tiering_primary_read_usage_perc"
52 + PROMOTION_ALLOCATION_FAILURES = (
53 + "vllm:kv_offload_tiering_promotion_allocation_failures"
54 + )
55 + ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs"
56 + ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs"
57
58
59 @dataclass
45 -class JobMetadata:
60 +class TransferJob:
61 """Metadata for an in-flight async transfer job."""
62
63 job_id: JobId
@@ -63,6 +78,7 @@ class JobResult:
78 # keys that were successfully loaded. None means all keys share the fate
79 # indicated by `success`. Must be a subset of the job's original keys.
80 successful_keys: Collection[OffloadKey] | None = None
81 + transfer_time: float | None = None
82
83
84 class ParentManager(ABC):
@@ -96,7 +112,7 @@ class ParentManager(ABC):
112 self,
113 keys: Collection[OffloadKey],
114 req_context: ReqContext,
99 - ) -> JobMetadata: ...
115 + ) -> TransferJob: ...
116
117 @abstractmethod
118 def on_request_finished(self, req_context: ReqContext) -> None: ...
@@ -153,7 +169,7 @@ class SecondaryTierManager(ABC):
169 pass
170
171 @abstractmethod
156 - def submit_store(self, job_metadata: JobMetadata) -> None:
172 + def submit_store(self, job_metadata: TransferJob) -> None:
173 """
174 Submit an async job to store blocks from the primary tier to this
175 secondary tier.
@@ -181,7 +197,7 @@ class SecondaryTierManager(ABC):
197 pass
198
199 @abstractmethod
184 - def submit_load(self, job_metadata: JobMetadata) -> None:
200 + def submit_load(self, job_metadata: TransferJob) -> None:
201 """
202 Submit an async job to load blocks from this secondary tier to the
203 primary tier.
vllm/v1/kv_offload/tiering/example/manager.py
+19 -5
@@ -23,9 +23,9 @@ from vllm.v1.kv_offload.base import (
23 RequestOffloadingContext,
24 )
25 from vllm.v1.kv_offload.tiering.base import (
26 - JobMetadata,
26 JobResult,
27 SecondaryTierManager,
28 + TransferJob,
29 )
30
31 logger = logging.getLogger(__name__)
@@ -73,6 +73,8 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
73
74 # Completed jobs waiting to be retrieved by get_finished_jobs()
75 self.completed_jobs: list[JobResult] = []
76 + assert primary_kv_view.strides is not None
77 + self._block_size = primary_kv_view.strides[0]
78
79 @override
80 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
@@ -89,7 +91,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
91 return LookupResult.HIT if key in self.blocks else LookupResult.MISS
92
93 @override
92 - def submit_store(self, job_metadata: JobMetadata) -> None:
94 + def submit_store(self, job_metadata: TransferJob) -> None:
95 """
96 Submit a job to store blocks from primary tier to this tier.
97
@@ -106,10 +108,16 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
108
109 for key in keys:
110 self.blocks[key] = True
109 - self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
111 + self.completed_jobs.append(
112 + JobResult(
113 + job_id=job_metadata.job_id,
114 + success=True,
115 + transfer_time=0.0,
116 + )
117 + )
118
119 @override
112 - def submit_load(self, job_metadata: JobMetadata) -> None:
120 + def submit_load(self, job_metadata: TransferJob) -> None:
121 """
122 Submit a job to load blocks from this tier to primary tier.
123
@@ -131,7 +139,13 @@ class ExampleSecondaryTierManager(SecondaryTierManager):
139 )
140 return
141
134 - self.completed_jobs.append(JobResult(job_id=job_metadata.job_id, success=True))
142 + self.completed_jobs.append(
143 + JobResult(
144 + job_id=job_metadata.job_id,
145 + success=True,
146 + transfer_time=0.0,
147 + )
148 + )
149
150 @override
151 def get_finished_jobs(self) -> Iterable[JobResult]:
vllm/v1/kv_offload/tiering/fs/manager.py
+16 -8
@@ -43,11 +43,11 @@ from vllm.v1.kv_offload.file_mapper import FileMapper
43 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
44 from vllm.v1.kv_offload.tiering.base import (
45 JobId,
46 - JobMetadata,
46 JobResult,
47 RequestOffloadingContext,
48 ScheduleEndContext,
49 SecondaryTierManager,
50 + TransferJob,
51 )
52 from vllm.v1.kv_offload.tiering.fs.io import (
53 batch_load_block,
@@ -214,12 +214,13 @@ class FileSystemTierManager(SecondaryTierManager):
214 return LookupResult.HIT if result else LookupResult.MISS
215
216 @override
217 - def submit_store(self, job_metadata: JobMetadata) -> None:
217 + def submit_store(self, job_metadata: TransferJob) -> None:
218 + keys = list(job_metadata.keys)
219 if self.events is not None:
219 - self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
220 + self._store_job_keys[job_metadata.job_id] = keys
221 task = functools.partial(
222 batch_store_block,
222 - [self.file_mapper.get_file_name(key) for key in job_metadata.keys],
223 + [self.file_mapper.get_file_name(key) for key in keys],
224 self._primary_kv_view,
225 [int(bid) * self._block_size for bid in job_metadata.block_ids],
226 self._block_size,
@@ -228,7 +229,7 @@ class FileSystemTierManager(SecondaryTierManager):
229 self._pool.enqueue_store(job_metadata.job_id, 1, [task])
230
231 @override
231 - def submit_load(self, job_metadata: JobMetadata) -> None:
232 + def submit_load(self, job_metadata: TransferJob) -> None:
233 job_id = job_metadata.job_id
234 # Track this load's keys so a failed promotion can mark only its failed
235 # keys as a miss (see get_finished_jobs).
@@ -271,7 +272,7 @@ class FileSystemTierManager(SecondaryTierManager):
272 """Collect finished jobs; a failed promotion marks only its failed keys
273 as a miss here (scheduler thread)."""
274 results = []
274 - for job_id, success in self._pool.get_finished():
275 + for job_id, success, transfer_time in self._pool.get_finished():
276 if self.events is not None:
277 keys = self._store_job_keys.pop(job_id, None)
278 if success and keys:
@@ -298,10 +299,17 @@ class FileSystemTierManager(SecondaryTierManager):
299 job_id=job_id,
300 success=False,
301 successful_keys=tuple(successful) if successful else None,
302 + transfer_time=transfer_time,
303 )
304 )
303 - else:
304 - results.append(JobResult(job_id=job_id, success=success))
305 + continue
306 + results.append(
307 + JobResult(
308 + job_id=job_id,
309 + success=success,
310 + transfer_time=transfer_time,
311 + )
312 + )
313 return results
314
315 @override
vllm/v1/kv_offload/tiering/fs/thread_pool.py
+25 -8
@@ -9,6 +9,7 @@ Thread pool:
9 """
10
11 import threading
12 +import time
13 from collections import deque
14 from collections.abc import Callable, Iterable
15
@@ -25,26 +26,37 @@ class JobState:
26 Each task calls task_done(success) when it finishes.
27 """
28
28 - __slots__ = ("_job_id", "_n_tasks", "_completed", "_success", "_lock")
29 + __slots__ = (
30 + "_job_id",
31 + "_n_tasks",
32 + "_completed",
33 + "_success",
34 + "_transfer_time",
35 + "_lock",
36 + )
37
38 def __init__(self, job_id: JobId, n_tasks: int) -> None:
39 self._job_id: JobId = job_id
40 self._n_tasks = n_tasks
41 self._completed = 0
42 self._success = True
43 + self._transfer_time = 0.0
44 self._lock = threading.Lock()
45
46 @property
47 def job_id(self) -> JobId:
48 return self._job_id
49
41 - def task_done(self, success: bool) -> tuple[bool, bool]:
50 + def task_done(
51 + self, success: bool, transfer_time: float
52 + ) -> tuple[bool, bool, float]:
53 """Returns if job completed and success flag"""
54 with self._lock:
55 self._completed += 1
56 + self._transfer_time += transfer_time
57 if not success:
58 self._success = False
47 - return self._completed == self._n_tasks, self._success
59 + return self._completed == self._n_tasks, self._success, self._transfer_time
60
61
62 class DualQueueThreadPool:
@@ -67,7 +79,7 @@ class DualQueueThreadPool:
79 self._condition = threading.Condition(threading.Lock())
80 self._stop = False
81 self._threads: list[threading.Thread] = []
70 - self._finished_q: deque[tuple[JobId, bool]] = deque()
82 + self._finished_q: deque[tuple[JobId, bool, float]] = deque()
83 self._inflight_jobs = 0 # guarded by _condition
84
85 for i in range(n_read_threads):
@@ -118,7 +130,7 @@ class DualQueueThreadPool:
130 self._store_q.append((fn, state))
131 self._condition.notify(n_tasks)
132
121 - def get_finished(self) -> list[tuple[JobId, bool]]:
133 + def get_finished(self) -> list[tuple[JobId, bool, float]]:
134 # No lock needed: deque is thread-safe for concurrent append/popleft,
135 # and the manager is the sole popper.
136 jobs = []
@@ -163,18 +175,23 @@ class DualQueueThreadPool:
175 secondary = self._store_q if load_priority else self._load_q
176 task, state = primary.popleft() if primary else secondary.popleft()
177 try:
178 + start_time = time.monotonic()
179 task()
167 - job_finished, success = state.task_done(True)
180 + transfer_time = time.monotonic() - start_time
181 + job_finished, success, total_time = state.task_done(True, transfer_time)
182 except Exception as exc:
183 + transfer_time = time.monotonic() - start_time
184 logger.error(
185 "Job %s block I/O failed: %s",
186 state.job_id,
187 exc,
188 )
174 - job_finished, success = state.task_done(False)
189 + job_finished, success, total_time = state.task_done(
190 + False, transfer_time
191 + )
192
193 if job_finished:
194 with self._condition:
178 - self._finished_q.append((state.job_id, success))
195 + self._finished_q.append((state.job_id, success, total_time))
196 self._inflight_jobs -= 1
197 self._condition.notify_all()
vllm/v1/kv_offload/tiering/manager.py
+132 -119
@@ -23,6 +23,7 @@ Key Design Principles:
23 import time
24 from collections.abc import Collection, Iterable, Sequence
25 from dataclasses import dataclass, field
26 +from typing import NamedTuple
27
28 import numpy as np
29 from typing_extensions import override
@@ -48,12 +49,12 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager
49 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
50 from vllm.v1.kv_offload.tiering.base import (
51 JobId,
51 - JobMetadata,
52 JobResult,
53 ParentManager,
54 SecondaryTierManager,
55 - TieringOffloadingMetrics,
55 + TransferJob,
56 )
57 +from vllm.v1.kv_offload.tiering.metrics import TieringMetricsTracker
58
59 logger = init_logger(__name__)
60
@@ -72,11 +73,12 @@ class RequestState:
73 req_context: ReqContext
74 pending_primary_stores: int = 0
75 is_finished: bool = False
75 - request_level_tiers: set[SecondaryTierManager] | None = None
76 - sync_lookup_delay: float = 0.0
77 - # time.monotonic() of this request's first deferred secondary-tier lookup;
78 - # None once consumed (observed) or while no secondary lookup is pending.
79 - secondary_lookup_start_time: float | None = None
76 + request_level_tiers: set[int] | None = None
77 +
78 +
79 +class JobMetadata(NamedTuple):
80 + transfer_job: TransferJob
81 + tier_idx: int
82
83
84 class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
@@ -132,31 +134,33 @@ class CPUPrimaryTierOffloadingManager(CPUOffloadingManager):
134
135 class _SecondaryTierFacingParent(ParentManager):
136 """Wrapper that implements ParentManager by delegating to the
135 - TieringOffloadingManager with exclude_tier set to the origin tier."""
137 + TieringOffloadingManager with exclude_tier_idx set to the origin tier."""
138
137 - __slots__ = ("_m", "_origin")
139 + __slots__ = ("_m", "_origin_idx")
140
141 def __init__(
142 self,
143 manager: "TieringOffloadingManager",
142 - tier: SecondaryTierManager,
144 + tier_idx: int,
145 ):
146 self._m = manager
145 - self._origin = tier
147 + self._origin_idx = tier_idx
148
149 def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
148 - return self._m.on_new_request(req_context, exclude_tier=self._origin)
150 + return self._m.on_new_request(req_context, exclude_tier_idx=self._origin_idx)
151
152 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
151 - return self._m.lookup(key, req_context, exclude_tier=self._origin)
153 + return self._m.lookup(key, req_context, exclude_tier_idx=self._origin_idx)
154
155 def create_store_job(
156 self, keys: Collection[OffloadKey], req_context: ReqContext
155 - ) -> JobMetadata:
156 - return self._m.create_store_job(keys, req_context)
157 + ) -> TransferJob:
158 + return self._m.create_store_job(keys, req_context, self._origin_idx)
159
160 def on_request_finished(self, req_context: ReqContext) -> None:
159 - return self._m.on_request_finished(req_context, exclude_tier=self._origin)
161 + return self._m.on_request_finished(
162 + req_context, exclude_tier_idx=self._origin_idx
163 + )
164
165
166 class TieringOffloadingManager(OffloadingManager):
@@ -193,18 +197,23 @@ class TieringOffloadingManager(OffloadingManager):
197
198 self._job_id_counter: int = 0
199 # Job tracking: maps job_id to metadata for all in-flight transfers.
196 - # JobMetadata.is_promotion distinguishes direction:
200 + # TransferJob.is_promotion distinguishes direction:
201 # True: secondary → primary (promotion)
202 # False: primary → secondary (cascade)
199 - self._transfer_jobs: dict[JobId, JobMetadata] = {}
203 + self._jobs: dict[JobId, JobMetadata] = {}
204 + primary_view = self.primary_tier.get_kv_memoryview()
205 + assert primary_view.strides is not None
206 + self._metrics = TieringMetricsTracker(
207 + tier_types=[tier.tier_type for tier in self.secondary_tiers],
208 + num_primary_blocks=self.primary_tier._num_blocks,
209 + primary_block_size=primary_view.strides[0],
210 + )
211
212 # Pending promotion requests accumulated during lookup() calls; flushed
213 # as one batched submit_load() per (tier, request) in on_schedule_end().
203 - # Outer key: tier. Inner key: req_context.req_id — the same ReqContext
214 + # Outer key: tier index. Inner key: req_context.req_id — the same ReqContext
215 # object is reused for all block lookups of a given request per engine step.
205 - self._pending_load_submissions: dict[
206 - SecondaryTierManager, dict[str, PendingPromotion]
207 - ] = {}
216 + self._pending_load_submissions: dict[int, dict[str, PendingPromotion]] = {}
217
218 # Gate for once-per-step execution of _maybe_process_finished_jobs().
219 # Reset at the end of each step in on_schedule_end().
@@ -217,13 +226,13 @@ class TieringOffloadingManager(OffloadingManager):
226
227 # Cached ParentManager wrappers for each secondary tier.
228 self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = {
220 - tier: _SecondaryTierFacingParent(self, tier)
221 - for tier in self.secondary_tiers
229 + tier: _SecondaryTierFacingParent(self, tier_idx)
230 + for tier_idx, tier in enumerate(self.secondary_tiers)
231 }
232
224 - # Buffers manager-level observations (e.g. lookup delay) between
225 - # get_stats() calls; merged in and reset each time get_stats() runs.
226 - self._stats = OffloadingConnectorStats()
233 + @property
234 + def _transfer_jobs(self) -> dict[JobId, JobMetadata]:
235 + return self._jobs
236
237 def _next_job_id(self) -> JobId:
238 """Generate a unique job ID for async transfer tracking."""
@@ -231,6 +240,14 @@ class TieringOffloadingManager(OffloadingManager):
240 self._job_id_counter += 1
241 return job_id
242
243 + def _register_job(self, transfer_job: TransferJob, tier_idx: int) -> None:
244 + job_metadata = JobMetadata(transfer_job, tier_idx)
245 + self._jobs[transfer_job.job_id] = job_metadata
246 + self._metrics.on_job_registered(job_metadata)
247 +
248 + def _pop_job(self, job_id: JobId) -> JobMetadata | None:
249 + return self._jobs.pop(job_id, None)
250 +
251 def _maybe_process_finished_jobs(self):
252 """
253 Poll secondary tiers for completed jobs (at most once per step).
@@ -247,13 +264,14 @@ class TieringOffloadingManager(OffloadingManager):
264 def _complete_promotion(
265 self, job_metadata: JobMetadata, completed_job: JobResult
266 ) -> None:
267 + transfer_job = job_metadata.transfer_job
268 successful_keys = completed_job.successful_keys
269 failed_keys: Collection[OffloadKey]
270 if completed_job.success:
253 - successful_keys = job_metadata.keys
271 + successful_keys = transfer_job.keys
272 failed_keys = ()
273 elif successful_keys:
256 - failed_keys_set = set(job_metadata.keys)
274 + failed_keys_set = set(transfer_job.keys)
275 assert failed_keys_set.issuperset(successful_keys), (
276 f"Finished promotion job_id {completed_job.job_id} "
277 "reported unknown successful keys"
@@ -262,18 +280,18 @@ class TieringOffloadingManager(OffloadingManager):
280 failed_keys = failed_keys_set
281 else:
282 successful_keys = ()
265 - failed_keys = job_metadata.keys
283 + failed_keys = transfer_job.keys
284
285 if successful_keys:
286 self.primary_tier.complete_write(
287 successful_keys,
270 - job_metadata.req_context,
288 + transfer_job.req_context,
289 True,
290 )
291 if failed_keys:
292 self.primary_tier.complete_write(
293 failed_keys,
276 - job_metadata.req_context,
294 + transfer_job.req_context,
295 False,
296 )
297
@@ -291,13 +309,19 @@ class TieringOffloadingManager(OffloadingManager):
309 for i, tier in enumerate(self.secondary_tiers):
310 for completed_job in tier.get_finished_jobs():
311 job_id = completed_job.job_id
294 - job_metadata = self._transfer_jobs.pop(job_id, None)
312 + job_metadata = self._pop_job(job_id)
313 assert job_metadata is not None, (
314 f"Finished job_id {job_id} from tier #{i}"
297 - f" ({tier.tier_type}) not in _transfer_jobs"
315 + f" ({tier.tier_type}) not in _jobs"
316 )
317 + assert job_metadata.tier_idx == i, (
318 + f"Finished job_id {job_id} reported by tier #{i}"
319 + f" but belongs to tier #{job_metadata.tier_idx}"
320 + )
321 + transfer_job = job_metadata.transfer_job
322 + self._metrics.on_job_finished(job_metadata, completed_job)
323
300 - if job_metadata.is_promotion:
324 + if transfer_job.is_promotion:
325 # secondary→primary transfer (promotion) completed.
326 # Make blocks available in primary tier.
327 self._complete_promotion(job_metadata, completed_job)
@@ -305,7 +329,7 @@ class TieringOffloadingManager(OffloadingManager):
329 # primary→secondary transfer completed.
330 # Decrement ref_cnt on primary blocks.
331 self.primary_tier.complete_read(
308 - job_metadata.keys, job_metadata.req_context
332 + transfer_job.keys, transfer_job.req_context
333 )
334
335 @override
@@ -314,7 +338,7 @@ class TieringOffloadingManager(OffloadingManager):
338 key: OffloadKey,
339 req_context: ReqContext,
340 *,
317 - exclude_tier: SecondaryTierManager | None = None,
341 + exclude_tier_idx: int | None = None,
342 ) -> LookupResult:
343 """
344 Check whether a single block is offloaded and ready.
@@ -343,73 +367,58 @@ class TieringOffloadingManager(OffloadingManager):
367 # in time for a promotion this lookup may initiate.
368 self._maybe_process_finished_jobs()
369
346 - req_state = self._req_state.get(req_context.req_id)
347 -
370 + start_time = time.monotonic()
371 primary_hit = self.primary_tier.lookup(key, req_context)
372 + lookup_duration = time.monotonic() - start_time
373 + self._metrics.on_lookup(
374 + req_context,
375 + key,
376 + self._metrics.primary_tier_label,
377 + primary_hit,
378 + lookup_duration,
379 + )
380 if primary_hit is LookupResult.HIT:
381 return LookupResult.HIT
382 if primary_hit is LookupResult.HIT_PENDING:
383 return LookupResult.HIT_PENDING
384
354 - lookup_start = time.monotonic()
385 any_retry = False
356 - for tier in self.secondary_tiers:
357 - if tier is exclude_tier:
386 + for i, tier in enumerate(self.secondary_tiers):
387 + if i == exclude_tier_idx:
388 continue
389 if not req_context.load_tier_filter.allows(tier.medium, tier.locality):
390 continue
391 + labelvalues = self._metrics.tier_label(i)
392 + start_time = time.monotonic()
393 result = tier.lookup(key, req_context)
394 + lookup_duration = time.monotonic() - start_time
395 if result is LookupResult.HIT:
363 - promoted = self._initiate_promotion(tier, key, req_context)
364 - self._accumulate_lookup_sync_delay(req_state, lookup_start)
365 - if (
366 - req_state is not None
367 - and promoted
368 - and req_state.secondary_lookup_start_time is None
369 - ):
370 - req_state.secondary_lookup_start_time = lookup_start
396 + self._metrics.on_lookup(
397 + req_context,
398 + key,
399 + labelvalues,
400 + result,
401 + lookup_duration,
402 + )
403 + promoted = self._initiate_promotion(i, key, req_context)
404 return LookupResult.MISS if not promoted else LookupResult.RETRY
405 if result is LookupResult.RETRY:
406 any_retry = True
407 + self._metrics.on_lookup(
408 + req_context,
409 + key,
410 + labelvalues,
411 + result,
412 + lookup_duration,
413 + )
414
375 - self._accumulate_lookup_sync_delay(req_state, lookup_start)
415 if any_retry:
377 - if req_state is not None and req_state.secondary_lookup_start_time is None:
378 - req_state.secondary_lookup_start_time = lookup_start
416 return LookupResult.RETRY
417 return LookupResult.MISS
418
382 - def _accumulate_lookup_sync_delay(
383 - self, req_state: RequestState | None, start_time: float
384 - ) -> None:
385 - """Accumulate secondary-tier lookup time until allocation or finish."""
386 - if req_state is not None:
387 - req_state.sync_lookup_delay += time.monotonic() - start_time
388 -
389 - def _maybe_observe_lookup_sync_delay(self, req_state: RequestState) -> None:
390 - delay = req_state.sync_lookup_delay
391 - if delay == 0:
392 - return
393 - req_state.sync_lookup_delay = 0.0
394 - self._stats.observe_histogram(
395 - TieringOffloadingMetrics.LOOKUP_SYNC_DELAY,
396 - delay,
397 - )
398 -
399 - def _maybe_observe_lookup_async_delay(self, req_state: RequestState) -> None:
400 - """Flush a pending deferred secondary-tier lookup timer, if any."""
401 - start_time = req_state.secondary_lookup_start_time
402 - if start_time is None:
403 - return
404 - req_state.secondary_lookup_start_time = None
405 - self._stats.observe_histogram(
406 - TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY,
407 - time.monotonic() - start_time,
408 - )
409 -
419 def _initiate_promotion(
420 self,
412 - tier: SecondaryTierManager,
421 + tier_idx: int,
422 key: OffloadKey,
423 req_context: ReqContext,
424 ) -> bool:
@@ -423,7 +432,7 @@ class TieringOffloadingManager(OffloadingManager):
432 batched job.
433
434 Args:
426 - tier: The secondary tier to promote from
435 + tier_idx: The secondary tier index to promote from
436 key: Block to promote
437 req_context: Per-request context forwarded to primary.prepare_write().
438
@@ -439,13 +448,14 @@ class TieringOffloadingManager(OffloadingManager):
448 if primary_write_result is None:
449 # Primary tier is full; caller should treat the block as unavailable
450 # rather than retrying indefinitely.
451 + self._metrics.on_promotion_allocation_failure()
452 return False
453
454 store_spec = primary_write_result.store_spec
455 assert isinstance(store_spec, CPULoadStoreSpec)
456 # Defer submit_load to on_schedule_end(). Group by (tier, request) so
457 # each request's blocks are submitted as one batched job per tier.
448 - tier_pending = self._pending_load_submissions.setdefault(tier, {})
458 + tier_pending = self._pending_load_submissions.setdefault(tier_idx, {})
459 ctx_id = req_context.req_id
460 if ctx_id not in tier_pending:
461 tier_pending[ctx_id] = PendingPromotion(
@@ -465,17 +475,18 @@ class TieringOffloadingManager(OffloadingManager):
475 if not self._pending_load_submissions:
476 return
477
468 - for tier, pending_by_ctx in self._pending_load_submissions.items():
478 + for tier_idx, pending_by_ctx in self._pending_load_submissions.items():
479 + tier = self.secondary_tiers[tier_idx]
480 for entry in pending_by_ctx.values():
481 job_id = self._next_job_id()
471 - job_metadata = JobMetadata(
482 + job_metadata = TransferJob(
483 job_id=job_id,
484 keys=entry.keys,
485 block_ids=np.array(entry.block_ids, dtype=np.int64),
486 is_promotion=True,
487 req_context=entry.req_context,
488 )
478 - self._transfer_jobs[job_id] = job_metadata
489 + self._register_job(job_metadata, tier_idx)
490 tier.submit_load(job_metadata)
491
492 self._pending_load_submissions.clear()
@@ -594,7 +605,7 @@ class TieringOffloadingManager(OffloadingManager):
605 self,
606 keys: Sequence[OffloadKey],
607 req_context: ReqContext,
597 - request_level_tiers: set[SecondaryTierManager],
608 + request_level_tiers: set[int],
609 ) -> None:
610 """
611 For tiers that requested request-level policy, submit_store() for
@@ -609,8 +620,9 @@ class TieringOffloadingManager(OffloadingManager):
620 if not ready_keys:
621 return
622
612 - for tier in request_level_tiers:
613 - job_metadata = self.create_store_job(ready_keys, req_context)
623 + for tier_idx in request_level_tiers:
624 + job_metadata = self.create_store_job(ready_keys, req_context, tier_idx)
625 + tier = self.secondary_tiers[tier_idx]
626 tier.submit_store(job_metadata)
627
628 @override
@@ -647,8 +659,8 @@ class TieringOffloadingManager(OffloadingManager):
659 # LoadStoreSpec AND to increment ref_cnt (protecting blocks from
660 # eviction during the async transfer). One prepare_read() call per
661 # secondary tier.
650 - for tier in self.secondary_tiers:
651 - job_metadata = self.create_store_job(keys, req_context)
662 + for tier_idx, tier in enumerate(self.secondary_tiers):
663 + job_metadata = self.create_store_job(keys, req_context, tier_idx)
664 tier.submit_store(job_metadata)
665
666 # Note: The async transfers are now in flight. Their completion is
@@ -663,12 +675,13 @@ class TieringOffloadingManager(OffloadingManager):
675 self,
676 keys: Collection[OffloadKey],
677 req_context: ReqContext,
666 - ) -> JobMetadata:
678 + tier_idx: int = 0,
679 + ) -> TransferJob:
680 """Pin blocks in the primary tier and create a tracked store job.
681
682 Calls prepare_read() to increment ref_cnt (protecting blocks
683 from eviction during the async transfer), allocates a job ID,
671 - and registers the job in _transfer_jobs.
684 + and registers the job in _jobs.
685
686 The caller is responsible for the actual data transfer and
687 reporting completion via get_finished_jobs().
@@ -676,14 +689,14 @@ class TieringOffloadingManager(OffloadingManager):
689 primary_blocks_spec = self.primary_tier.prepare_read(keys, req_context)
690 assert isinstance(primary_blocks_spec, CPULoadStoreSpec)
691 job_id = self._next_job_id()
679 - job_metadata = JobMetadata(
692 + job_metadata = TransferJob(
693 job_id=job_id,
694 keys=keys,
695 block_ids=primary_blocks_spec.block_ids,
696 is_promotion=False,
697 req_context=req_context,
698 )
686 - self._transfer_jobs[job_id] = job_metadata
699 + self._register_job(job_metadata, tier_idx)
700 return job_metadata
701
702 @override
@@ -691,7 +704,7 @@ class TieringOffloadingManager(OffloadingManager):
704 self,
705 req_context: ReqContext,
706 *,
694 - exclude_tier: SecondaryTierManager | None = None,
707 + exclude_tier_idx: int | None = None,
708 ) -> RequestOffloadingContext:
709 """
710 Query each secondary tier for its offload policy preference.
@@ -700,14 +713,15 @@ class TieringOffloadingManager(OffloadingManager):
713 Only stores REQUEST_LEVEL tier decisions for use in prepare_store.
714 """
715 state = RequestState(req_context=req_context)
703 - for tier in self.secondary_tiers:
704 - if tier is exclude_tier:
716 + self._metrics.on_new_request(req_context)
717 + for tier_idx, tier in enumerate(self.secondary_tiers):
718 + if tier_idx == exclude_tier_idx:
719 continue
720 tier_ctx = tier.on_new_request(req_context)
721 if tier_ctx.policy == OffloadPolicy.REQUEST_LEVEL:
722 if state.request_level_tiers is None:
723 state.request_level_tiers = set()
710 - state.request_level_tiers.add(tier)
724 + state.request_level_tiers.add(tier_idx)
725 self._req_state[req_context.req_id] = state
726
727 policy = (
@@ -722,17 +736,17 @@ class TieringOffloadingManager(OffloadingManager):
736 self,
737 req_context: ReqContext,
738 *,
725 - exclude_tier: SecondaryTierManager | None = None,
739 + exclude_tier_idx: int | None = None,
740 ) -> None:
741 self.primary_tier.on_request_finished(req_context)
742 state = self._req_state[req_context.req_id]
743 state.is_finished = True
730 - self._maybe_finalize_request(req_context.req_id, exclude_tier)
744 + self._maybe_finalize_request(req_context.req_id, exclude_tier_idx)
745
746 def _maybe_finalize_request(
747 self,
748 req_id: str,
735 - exclude_tier: SecondaryTierManager | None = None,
749 + exclude_tier_idx: int | None = None,
750 ) -> None:
751 """Finalize secondary tiers once no more store cascades can be submitted.
752
@@ -746,12 +760,11 @@ class TieringOffloadingManager(OffloadingManager):
760 if state.pending_primary_stores != 0:
761 return
762
749 - for tier in self.secondary_tiers:
750 - if tier is exclude_tier:
763 + for tier_idx, tier in enumerate(self.secondary_tiers):
764 + if tier_idx == exclude_tier_idx:
765 continue
766 tier.on_request_finished(state.req_context)
753 - self._maybe_observe_lookup_sync_delay(state)
754 - self._maybe_observe_lookup_async_delay(state)
767 + self._metrics.on_request_finished(state.req_context)
768 del self._req_state[req_id]
769
770 @override
@@ -782,15 +795,14 @@ class TieringOffloadingManager(OffloadingManager):
795 state = self._req_state.get(req_id)
796 if state is None:
797 continue
785 - self._maybe_observe_lookup_sync_delay(state)
786 - self._maybe_observe_lookup_async_delay(state)
798 + self._metrics.on_request_allocated(state.req_context)
799
800 @override
801 def has_pending_work(self) -> bool:
802 # In-flight primary<->secondary transfers (pending promotions are
803 # translated to transfer jobs in on_schedule_end), plus any work the
804 # secondary tiers themselves still have outstanding.
793 - return bool(self._transfer_jobs) or any(
805 + return bool(self._jobs) or any(
806 tier.has_pending_work() for tier in self.secondary_tiers
807 )
808
@@ -825,11 +837,13 @@ class TieringOffloadingManager(OffloadingManager):
837 # All tier I/O has stopped; consume their completion notifications
838 # so manager bookkeeping is consistent before the primary reset.
839 self._process_finished_jobs()
840 + assert not self._jobs
841
842 # Deferred promotion submissions reserve primary slots that the
843 # reset below invalidates; their submit_load() has not yet been
844 # called so no tier I/O is touching that memory.
845 self._pending_load_submissions.clear()
846 + self._metrics.assert_idle()
847
848 finished_req_ids = []
849 for req_id, state in self._req_state.items():
@@ -838,8 +852,7 @@ class TieringOffloadingManager(OffloadingManager):
852 continue
853 for tier in self.secondary_tiers:
854 tier.on_request_finished(state.req_context)
841 - self._maybe_observe_lookup_sync_delay(state)
842 - self._maybe_observe_lookup_async_delay(state)
855 + self._metrics.on_request_finished(state.req_context)
856 finished_req_ids.append(req_id)
857
858 self.primary_tier.reset_cache()
@@ -855,6 +868,13 @@ class TieringOffloadingManager(OffloadingManager):
868 if stats is not None and stats.is_empty():
869 stats = None
870
871 + metrics_stats = self._metrics.take_stats()
872 + if metrics_stats is not None:
873 + if stats is None:
874 + stats = metrics_stats
875 + else:
876 + stats.aggregate(metrics_stats)
877 +
878 for tier in self.secondary_tiers:
879 tier_stats = tier.get_stats()
880 if tier_stats is None or tier_stats.is_empty():
@@ -864,13 +884,6 @@ class TieringOffloadingManager(OffloadingManager):
884 else:
885 stats.aggregate(tier_stats)
886
867 - if not self._stats.is_empty():
868 - if stats is None:
869 - stats = self._stats
870 - else:
871 - stats.aggregate(self._stats)
872 - self._stats = OffloadingConnectorStats()
873 -
887 return stats
888
889 @override
vllm/v1/kv_offload/tiering/metrics.py new
+269
@@ -0,0 +1,269 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import functools
5 +import time
6 +from dataclasses import dataclass, field
7 +from typing import Protocol
8 +
9 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
10 + OffloadingConnectorStats,
11 +)
12 +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext
13 +from vllm.v1.kv_offload.tiering.base import (
14 + JobResult,
15 + TieringOffloadingMetrics,
16 + TransferJob,
17 +)
18 +
19 +TierLabel = tuple[str]
20 +PRIMARY_TIER_LABEL: TierLabel = ("0:primary",)
21 +
22 +
23 +class _JobMetadataLike(Protocol):
24 + @property
25 + def transfer_job(self) -> TransferJob: ...
26 +
27 + @property
28 + def tier_idx(self) -> int: ...
29 +
30 +
31 +@dataclass(slots=True)
32 +class _RequestMetricsState:
33 + observed_lookups: dict[TierLabel, dict[OffloadKey, float | None]] | None = field(
34 + default_factory=dict
35 + )
36 +
37 +
38 +@dataclass
39 +class _TierState:
40 + active_promotion_count: int = 0
41 + active_cascade_count: int = 0
42 + primary_write_block_count: int = 0
43 + primary_read_block_count: int = 0
44 +
45 +
46 +class TieringMetricsTracker:
47 + def __init__(
48 + self,
49 + tier_types: list[str],
50 + num_primary_blocks: int,
51 + primary_block_size: int,
52 + ) -> None:
53 + self._tier_types = tier_types
54 + self._num_primary_blocks = num_primary_blocks
55 + self._primary_block_size = primary_block_size
56 + self._request_states: dict[str, _RequestMetricsState] = {}
57 + self._tier_states = [_TierState() for _ in tier_types]
58 + self._stats = OffloadingConnectorStats()
59 +
60 + @functools.cache # noqa: B019
61 + def tier_label(self, tier_idx: int) -> TierLabel:
62 + return (f"{tier_idx + 1}:{self._tier_types[tier_idx]}",)
63 +
64 + @property
65 + def primary_tier_label(self) -> TierLabel:
66 + return PRIMARY_TIER_LABEL
67 +
68 + def on_new_request(self, req_context: ReqContext) -> None:
69 + self._request_states[req_context.req_id] = _RequestMetricsState()
70 +
71 + def on_request_allocated(self, req_context: ReqContext) -> None:
72 + state = self._request_states.get(req_context.req_id)
73 + if state is None:
74 + return
75 + state.observed_lookups = None
76 +
77 + def on_request_finished(self, req_context: ReqContext) -> None:
78 + self._request_states.pop(req_context.req_id, None)
79 +
80 + def on_lookup(
81 + self,
82 + req_context: ReqContext,
83 + key: OffloadKey,
84 + tier_label: TierLabel,
85 + result: LookupResult,
86 + lookup_duration: float,
87 + ) -> None:
88 + state = self._request_states.get(req_context.req_id)
89 + if state is None:
90 + state = _RequestMetricsState()
91 + self._request_states[req_context.req_id] = state
92 +
93 + if state.observed_lookups is not None and result in (
94 + LookupResult.HIT,
95 + LookupResult.MISS,
96 + ):
97 + observed = state.observed_lookups.setdefault(tier_label, {})
98 + start_time = observed.get(key)
99 + if key not in observed or start_time is not None:
100 + observed[key] = None
101 + self._observe_resolved_lookup(
102 + tier_label,
103 + result,
104 + lookup_duration,
105 + start_time,
106 + )
107 + elif result is LookupResult.RETRY and state.observed_lookups is not None:
108 + observed = state.observed_lookups.setdefault(tier_label, {})
109 + observed.setdefault(key, time.monotonic() - lookup_duration)
110 +
111 + def on_job_registered(self, job_metadata: _JobMetadataLike) -> None:
112 + transfer_job = job_metadata.transfer_job
113 + state = self._tier_states[job_metadata.tier_idx]
114 + block_count = len(transfer_job.block_ids)
115 + if transfer_job.is_promotion:
116 + state.active_promotion_count += 1
117 + state.primary_write_block_count += block_count
118 + else:
119 + state.active_cascade_count += 1
120 + state.primary_read_block_count += block_count
121 +
122 + def on_job_finished(
123 + self, job_metadata: _JobMetadataLike, result: JobResult
124 + ) -> None:
125 + self._observe_finished_job_stats(job_metadata, result)
126 + self._decrement_tier_state(job_metadata)
127 +
128 + def on_promotion_allocation_failure(self) -> None:
129 + self._stats.increase_counter(
130 + TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES
131 + )
132 +
133 + def take_stats(self) -> OffloadingConnectorStats | None:
134 + active_transfer_stats = OffloadingConnectorStats()
135 + self._observe_active_transfer_stats(active_transfer_stats)
136 +
137 + stats = None
138 + if not active_transfer_stats.is_empty():
139 + stats = active_transfer_stats
140 + if not self._stats.is_empty():
141 + if stats is None:
142 + stats = self._stats
143 + else:
144 + stats.aggregate(self._stats)
145 + self._stats = OffloadingConnectorStats()
146 + return stats
147 +
148 + def assert_idle(self) -> None:
149 + assert all(
150 + state.active_promotion_count == 0
151 + and state.active_cascade_count == 0
152 + and state.primary_write_block_count == 0
153 + and state.primary_read_block_count == 0
154 + for state in self._tier_states
155 + )
156 +
157 + def _decrement_tier_state(self, job_metadata: _JobMetadataLike) -> None:
158 + transfer_job = job_metadata.transfer_job
159 + state = self._tier_states[job_metadata.tier_idx]
160 + block_count = len(transfer_job.block_ids)
161 + if transfer_job.is_promotion:
162 + assert state.active_promotion_count > 0
163 + state.active_promotion_count -= 1
164 + state.primary_write_block_count -= block_count
165 + assert state.primary_write_block_count >= 0
166 + else:
167 + assert state.active_cascade_count > 0
168 + state.active_cascade_count -= 1
169 + state.primary_read_block_count -= block_count
170 + assert state.primary_read_block_count >= 0
171 +
172 + def _observe_finished_job_stats(
173 + self,
174 + job_metadata: _JobMetadataLike,
175 + completed_job: JobResult,
176 + ) -> None:
177 + transfer_job = job_metadata.transfer_job
178 + labelvalues = self.tier_label(job_metadata.tier_idx)
179 + completed_key_count = len(transfer_job.keys)
180 + if not completed_job.success:
181 + failure_metric = (
182 + TieringOffloadingMetrics.PROMOTION_JOB_FAILURES
183 + if transfer_job.is_promotion
184 + else TieringOffloadingMetrics.CASCADE_JOB_FAILURES
185 + )
186 + self._stats.increase_counter(failure_metric, labelvalues=labelvalues)
187 + if transfer_job.is_promotion and completed_job.successful_keys:
188 + completed_key_count = len(completed_job.successful_keys)
189 + else:
190 + return
191 +
192 + bytes_metric = (
193 + TieringOffloadingMetrics.READ_BYTES
194 + if transfer_job.is_promotion
195 + else TieringOffloadingMetrics.WRITE_BYTES
196 + )
197 + time_metric = (
198 + TieringOffloadingMetrics.READ_TIME
199 + if transfer_job.is_promotion
200 + else TieringOffloadingMetrics.WRITE_TIME
201 + )
202 + transfer_size = completed_key_count * self._primary_block_size
203 + self._stats.increase_counter(bytes_metric, transfer_size, labelvalues)
204 + if completed_job.transfer_time is not None:
205 + self._stats.increase_counter(
206 + time_metric, completed_job.transfer_time, labelvalues
207 + )
208 +
209 + def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None:
210 + for tier_idx, state in enumerate(self._tier_states):
211 + labelvalues = self.tier_label(tier_idx)
212 + write_usage = (
213 + state.primary_write_block_count / self._num_primary_blocks
214 + if self._num_primary_blocks > 0
215 + else 0.0
216 + )
217 + read_usage = (
218 + state.primary_read_block_count / self._num_primary_blocks
219 + if self._num_primary_blocks > 0
220 + else 0.0
221 + )
222 + stats.set_gauge(
223 + TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC,
224 + write_usage,
225 + labelvalues,
226 + )
227 + stats.set_gauge(
228 + TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC,
229 + read_usage,
230 + labelvalues,
231 + )
232 + stats.set_gauge(
233 + TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS,
234 + state.active_promotion_count,
235 + labelvalues,
236 + )
237 + stats.set_gauge(
238 + TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS,
239 + state.active_cascade_count,
240 + labelvalues,
241 + )
242 +
243 + def _observe_resolved_lookup(
244 + self,
245 + tier_label: TierLabel,
246 + result: LookupResult,
247 + lookup_duration: float,
248 + async_start_time: float | None,
249 + ) -> None:
250 + self._stats.increase_counter(
251 + TieringOffloadingMetrics.BLOCK_QUERIES,
252 + labelvalues=tier_label,
253 + )
254 + if result is LookupResult.HIT:
255 + self._stats.increase_counter(
256 + TieringOffloadingMetrics.BLOCK_HITS,
257 + labelvalues=tier_label,
258 + )
259 + self._stats.observe_histogram(
260 + TieringOffloadingMetrics.LOOKUP_SYNC_DELAY,
261 + lookup_duration,
262 + tier_label,
263 + )
264 + if async_start_time is not None:
265 + self._stats.observe_histogram(
266 + TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY,
267 + time.monotonic() - async_start_time,
268 + tier_label,
269 + )
vllm/v1/kv_offload/tiering/obj/manager.py
+20 -5
@@ -22,11 +22,11 @@ from vllm.v1.kv_offload.file_mapper import FileMapper
22 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager
23 from vllm.v1.kv_offload.tiering.base import (
24 JobId,
25 - JobMetadata,
25 JobResult,
26 RequestOffloadingContext,
27 ScheduleEndContext,
28 SecondaryTierManager,
29 + TransferJob,
30 )
31 from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig
32
@@ -267,7 +267,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
267 self._pending_results.append(JobResult(job_id=job_id, success=False))
268 return
269
270 - self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle)
270 + self._transfers[job_id] = TransferEntry(
271 + xfer_handle=xfer_handle,
272 + files_desc=files_desc,
273 + obj_handle=obj_handle,
274 + )
275
276 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
277 result = self._lookup_manager.lookup(key, req_context)
@@ -275,7 +279,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
279 return LookupResult.RETRY
280 return LookupResult.HIT if result else LookupResult.MISS
281
278 - def submit_store(self, job_metadata: JobMetadata) -> None:
282 + def submit_store(self, job_metadata: TransferJob) -> None:
283 if self.events is not None:
284 self._store_job_keys[job_metadata.job_id] = list(job_metadata.keys)
285 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
@@ -283,7 +287,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
287 job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_WRITE
288 )
289
286 - def submit_load(self, job_metadata: JobMetadata) -> None:
290 + def submit_load(self, job_metadata: TransferJob) -> None:
291 self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys)
292 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)
293 self._submit_transfer(
@@ -317,6 +321,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
321 success = False
322 logger.warning("transfer failed job=%d state=%s", job_id, state)
323
324 + transfer_time = None
325 + if success:
326 + telemetry = self._agent.get_xfer_telemetry(entry.xfer_handle)
327 + transfer_time = telemetry.xferDuration / 1e6
328 +
329 try:
330 self._agent.release_xfer_handle(entry.xfer_handle)
331 except Exception as exc:
@@ -341,7 +350,13 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
350 logger.warning("deregister_memory failed for job %d: %s", job_id, exc)
351
352 del self._transfers[job_id]
344 - self._pending_results.append(JobResult(job_id=job_id, success=success))
353 + self._pending_results.append(
354 + JobResult(
355 + job_id=job_id,
356 + success=success,
357 + transfer_time=transfer_time,
358 + )
359 + )
360
361 def get_finished_jobs(self) -> Iterable[JobResult]:
362 """Poll transfers; a failed promotion marks its cached verdicts False
vllm/v1/kv_offload/tiering/p2p/manager.py
+3 -3
@@ -28,9 +28,9 @@ from vllm.v1.kv_offload.base import (
28 )
29 from vllm.v1.kv_offload.file_mapper import FileMapper
30 from vllm.v1.kv_offload.tiering.base import (
31 - JobMetadata,
31 JobResult,
32 SecondaryTierManager,
33 + TransferJob,
34 )
35 from vllm.v1.kv_offload.tiering.p2p.control import ControlTransport, ZmqTransport
36 from vllm.v1.kv_offload.tiering.p2p.data import DataTransport, NixlTransport
@@ -411,7 +411,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
411 return
412
413 @override
414 - def submit_store(self, job_metadata: JobMetadata) -> None:
414 + def submit_store(self, job_metadata: TransferJob) -> None:
415 job_id = job_metadata.job_id
416 keys = list(job_metadata.keys)
417 block_ids = job_metadata.block_ids
@@ -470,7 +470,7 @@ class P2PSecondaryTierManager(SecondaryTierManager):
470 )
471
472 @override
473 - def submit_load(self, job_metadata: JobMetadata) -> None:
473 + def submit_load(self, job_metadata: TransferJob) -> None:
474 job_id = job_metadata.job_id
475 keys = list(job_metadata.keys)
476 block_ids = job_metadata.block_ids
vllm/v1/kv_offload/tiering/spec.py
+103 -6
@@ -61,6 +61,8 @@ from typing_extensions import override
61 from vllm.logger import init_logger
62 from vllm.v1.kv_offload.base import (
63 CanonicalKVCaches,
64 + OffloadingCounterMetadata,
65 + OffloadingGaugeMetadata,
66 OffloadingHistogramMetadata,
67 OffloadingManager,
68 OffloadingMetricMetadata,
@@ -103,10 +105,10 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
105 metrics[TieringOffloadingMetrics.LOOKUP_SYNC_DELAY] = (
106 OffloadingHistogramMetadata(
107 documentation=(
106 - "Histogram of total blocking time spent querying secondary "
107 - "tiers for a request, accumulated from first lookup until "
108 - "the request is allocated or finishes, in seconds."
108 + "Histogram of blocking time spent in a per-block tier lookup "
109 + "that resolved as a hit or miss, labeled by tier, in seconds."
110 ),
111 + labelnames=("tier",),
112 buckets=(
113 0.00001,
114 0.00005,
@@ -125,10 +127,11 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
127 metrics[TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY] = (
128 OffloadingHistogramMetadata(
129 documentation=(
128 - "Histogram of wall-clock time from a request's first deferred "
129 - "secondary-tier lookup until the request is allocated or "
130 - "finishes, in seconds."
130 + "Histogram of wall-clock time from a per-block tier lookup "
131 + "first returning retry until that same tier lookup resolves "
132 + "as a hit or miss, labeled by tier, in seconds."
133 ),
134 + labelnames=("tier",),
135 buckets=(
136 0.0001,
137 0.0005,
@@ -144,6 +147,100 @@ class TieringOffloadingSpec(CPUOffloadingSpec):
147 ),
148 )
149 )
150 + metrics[TieringOffloadingMetrics.READ_BYTES] = OffloadingCounterMetadata(
151 + documentation=(
152 + "Total bytes read from secondary tiers into the primary tier, "
153 + "labeled by tier."
154 + ),
155 + labelnames=("tier",),
156 + )
157 + metrics[TieringOffloadingMetrics.READ_TIME] = OffloadingCounterMetadata(
158 + documentation=(
159 + "Total time spent reading from secondary tiers into the primary "
160 + "tier, in seconds, labeled by tier."
161 + ),
162 + labelnames=("tier",),
163 + )
164 + metrics[TieringOffloadingMetrics.WRITE_BYTES] = OffloadingCounterMetadata(
165 + documentation=(
166 + "Total bytes written from the primary tier to secondary tiers, "
167 + "labeled by tier."
168 + ),
169 + labelnames=("tier",),
170 + )
171 + metrics[TieringOffloadingMetrics.WRITE_TIME] = OffloadingCounterMetadata(
172 + documentation=(
173 + "Total time spent writing from the primary tier to secondary "
174 + "tiers, in seconds, labeled by tier."
175 + ),
176 + labelnames=("tier",),
177 + )
178 + metrics[TieringOffloadingMetrics.PROMOTION_JOB_FAILURES] = (
179 + OffloadingCounterMetadata(
180 + documentation=(
181 + "Number of failed secondary-tier promotion jobs, labeled by tier."
182 + ),
183 + labelnames=("tier",),
184 + )
185 + )
186 + metrics[TieringOffloadingMetrics.CASCADE_JOB_FAILURES] = (
187 + OffloadingCounterMetadata(
188 + documentation=(
189 + "Number of failed secondary-tier cascade jobs, labeled by tier."
190 + ),
191 + labelnames=("tier",),
192 + )
193 + )
194 + metrics[TieringOffloadingMetrics.BLOCK_QUERIES] = OffloadingCounterMetadata(
195 + documentation=(
196 + "Number of block lookup queries sent to a tier, labeled by tier."
197 + ),
198 + labelnames=("tier",),
199 + )
200 + metrics[TieringOffloadingMetrics.BLOCK_HITS] = OffloadingCounterMetadata(
201 + documentation="Number of block lookup hits in a tier, labeled by tier.",
202 + labelnames=("tier",),
203 + )
204 + metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES] = (
205 + OffloadingCounterMetadata(
206 + documentation=(
207 + "Number of promotion attempts that failed because the "
208 + "primary tier could not allocate space."
209 + ),
210 + )
211 + )
212 + metrics[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC] = (
213 + OffloadingGaugeMetadata(
214 + documentation=(
215 + "Current fraction of primary-tier space used by writes from "
216 + "secondary tiers, labeled by tier."
217 + ),
218 + labelnames=("tier",),
219 + )
220 + )
221 + metrics[TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC] = (
222 + OffloadingGaugeMetadata(
223 + documentation=(
224 + "Current fraction of primary-tier space used by reads to "
225 + "secondary tiers, labeled by tier."
226 + ),
227 + labelnames=("tier",),
228 + )
229 + )
230 + metrics[TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS] = (
231 + OffloadingGaugeMetadata(
232 + documentation=(
233 + "Number of active secondary-tier promotion jobs, labeled by tier."
234 + ),
235 + labelnames=("tier",),
236 + )
237 + )
238 + metrics[TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS] = OffloadingGaugeMetadata(
239 + documentation=(
240 + "Number of active secondary-tier cascade jobs, labeled by tier."
241 + ),
242 + labelnames=("tier",),
243 + )
244 secondary_tier_configs = extra_config.get("secondary_tiers", [])
245 if not isinstance(secondary_tier_configs, list):
246 raise ValueError("secondary_tiers must be a list of tier configurations")