@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 226 lines 7.13 KB
Raw
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_chunks=5,
29 primary_chunk_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.CHUNK_QUERIES][("0:primary",)] == 1
49 assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("1:fs",)] == 1
50 assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("2:p2p",)] == 1
51 assert values[TieringOffloadingMetrics.CHUNK_HITS][("2:p2p",)] == 1
52 assert ("1:fs",) not in values[TieringOffloadingMetrics.CHUNK_HITS]
53
54
55 def test_tiering_metrics_tracker_stops_lookup_metrics_after_allocation():
56 tracker = TieringMetricsTracker(
57 tier_types=["fs"],
58 num_primary_chunks=5,
59 primary_chunk_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.CHUNK_QUERIES][("0:primary",)] == 1
77 assert values[TieringOffloadingMetrics.CHUNK_QUERIES][("1:fs",)] == 1
78 assert values[TieringOffloadingMetrics.CHUNK_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.CHUNK_QUERIES not in values
94 assert TieringOffloadingMetrics.CHUNK_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_chunks=5,
101 primary_chunk_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_chunks=5,
153 primary_chunk_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_chunks=6,
186 primary_chunk_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_chunks=1,
218 primary_chunk_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