| 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 | SecondaryTierManager, |
| 16 | TieringOffloadingMetrics, |
| 17 | TransferJob, |
| 18 | ) |
| 19 | |
| 20 | TierLabel = tuple[str] |
| 21 | PRIMARY_TIER_LABEL: TierLabel = ("0:primary",) |
| 22 | |
| 23 | |
| 24 | class _JobMetadataLike(Protocol): |
| 25 | @property |
| 26 | def transfer_job(self) -> TransferJob: ... |
| 27 | |
| 28 | @property |
| 29 | def tier_idx(self) -> int: ... |
| 30 | |
| 31 | |
| 32 | @dataclass(slots=True) |
| 33 | class _RequestMetricsState: |
| 34 | observed_lookups: dict[TierLabel, dict[OffloadKey, float | None]] | None = field( |
| 35 | default_factory=dict |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | @dataclass |
| 40 | class _TierState: |
| 41 | active_promotion_count: int = 0 |
| 42 | active_cascade_count: int = 0 |
| 43 | primary_write_chunk_count: int = 0 |
| 44 | primary_read_chunk_count: int = 0 |
| 45 | |
| 46 | |
| 47 | class TieringMetricsTracker: |
| 48 | def __init__( |
| 49 | self, |
| 50 | tier_types: list[str], |
| 51 | num_primary_chunks: int, |
| 52 | primary_chunk_size: int, |
| 53 | ) -> None: |
| 54 | self._tier_types = tier_types |
| 55 | self._num_primary_chunks = num_primary_chunks |
| 56 | self._primary_chunk_size = primary_chunk_size |
| 57 | self._request_states: dict[str, _RequestMetricsState] = {} |
| 58 | self._tier_states = [_TierState() for _ in tier_types] |
| 59 | self._stats = OffloadingConnectorStats() |
| 60 | |
| 61 | @functools.cache # noqa: B019 |
| 62 | def tier_label(self, tier_idx: int) -> TierLabel: |
| 63 | return (f"{tier_idx + 1}:{self._tier_types[tier_idx]}",) |
| 64 | |
| 65 | @property |
| 66 | def primary_tier_label(self) -> TierLabel: |
| 67 | return PRIMARY_TIER_LABEL |
| 68 | |
| 69 | def on_new_request(self, req_context: ReqContext) -> None: |
| 70 | self._request_states[req_context.req_id] = _RequestMetricsState() |
| 71 | |
| 72 | def on_request_allocated(self, req_context: ReqContext) -> None: |
| 73 | state = self._request_states.get(req_context.req_id) |
| 74 | if state is None: |
| 75 | return |
| 76 | state.observed_lookups = None |
| 77 | |
| 78 | def on_request_finished(self, req_context: ReqContext) -> None: |
| 79 | self._request_states.pop(req_context.req_id, None) |
| 80 | |
| 81 | def on_lookup( |
| 82 | self, |
| 83 | req_context: ReqContext, |
| 84 | key: OffloadKey, |
| 85 | tier_label: TierLabel, |
| 86 | result: LookupResult, |
| 87 | lookup_duration: float, |
| 88 | ) -> None: |
| 89 | state = self._request_states.get(req_context.req_id) |
| 90 | if state is None: |
| 91 | state = _RequestMetricsState() |
| 92 | self._request_states[req_context.req_id] = state |
| 93 | |
| 94 | if state.observed_lookups is not None and result in ( |
| 95 | LookupResult.HIT, |
| 96 | LookupResult.MISS, |
| 97 | ): |
| 98 | observed = state.observed_lookups.setdefault(tier_label, {}) |
| 99 | start_time = observed.get(key) |
| 100 | if key not in observed or start_time is not None: |
| 101 | observed[key] = None |
| 102 | self._observe_resolved_lookup( |
| 103 | tier_label, |
| 104 | result, |
| 105 | lookup_duration, |
| 106 | start_time, |
| 107 | ) |
| 108 | elif result is LookupResult.RETRY and state.observed_lookups is not None: |
| 109 | observed = state.observed_lookups.setdefault(tier_label, {}) |
| 110 | observed.setdefault(key, time.monotonic() - lookup_duration) |
| 111 | |
| 112 | def on_job_registered(self, job_metadata: _JobMetadataLike) -> None: |
| 113 | transfer_job = job_metadata.transfer_job |
| 114 | state = self._tier_states[job_metadata.tier_idx] |
| 115 | chunk_count = len(transfer_job.chunk_ids) |
| 116 | if transfer_job.is_promotion: |
| 117 | state.active_promotion_count += 1 |
| 118 | state.primary_write_chunk_count += chunk_count |
| 119 | else: |
| 120 | state.active_cascade_count += 1 |
| 121 | state.primary_read_chunk_count += chunk_count |
| 122 | |
| 123 | def on_job_finished( |
| 124 | self, job_metadata: _JobMetadataLike, result: JobResult |
| 125 | ) -> None: |
| 126 | self._observe_finished_job_stats(job_metadata, result) |
| 127 | self._decrement_tier_state(job_metadata) |
| 128 | |
| 129 | def on_promotion_allocation_failure(self) -> None: |
| 130 | self._stats.increase_counter( |
| 131 | TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES |
| 132 | ) |
| 133 | |
| 134 | def record_backpressure( |
| 135 | self, |
| 136 | tiers: list[SecondaryTierManager], |
| 137 | ) -> None: |
| 138 | for i, tier in enumerate(tiers): |
| 139 | detector = tier.bp_detector |
| 140 | if detector is None: |
| 141 | continue |
| 142 | label = self.tier_label(i) |
| 143 | ema = detector.stats.get("store_latency_ema") |
| 144 | if ema is not None: |
| 145 | self._stats.set_gauge( |
| 146 | TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA, |
| 147 | ema, |
| 148 | labelvalues=label, |
| 149 | ) |
| 150 | stores_dropped, blocks_dropped = detector.policy.pop_stores_dropped() |
| 151 | if stores_dropped > 0: |
| 152 | self._stats.increase_counter( |
| 153 | TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED, |
| 154 | stores_dropped, |
| 155 | labelvalues=label, |
| 156 | ) |
| 157 | if blocks_dropped > 0: |
| 158 | self._stats.increase_counter( |
| 159 | TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED, |
| 160 | blocks_dropped, |
| 161 | labelvalues=label, |
| 162 | ) |
| 163 | |
| 164 | def take_stats(self) -> OffloadingConnectorStats | None: |
| 165 | active_transfer_stats = OffloadingConnectorStats() |
| 166 | self._observe_active_transfer_stats(active_transfer_stats) |
| 167 | |
| 168 | stats = None |
| 169 | if not active_transfer_stats.is_empty(): |
| 170 | stats = active_transfer_stats |
| 171 | if not self._stats.is_empty(): |
| 172 | if stats is None: |
| 173 | stats = self._stats |
| 174 | else: |
| 175 | stats.aggregate(self._stats) |
| 176 | self._stats = OffloadingConnectorStats() |
| 177 | return stats |
| 178 | |
| 179 | def assert_idle(self) -> None: |
| 180 | assert all( |
| 181 | state.active_promotion_count == 0 |
| 182 | and state.active_cascade_count == 0 |
| 183 | and state.primary_write_chunk_count == 0 |
| 184 | and state.primary_read_chunk_count == 0 |
| 185 | for state in self._tier_states |
| 186 | ) |
| 187 | |
| 188 | def _decrement_tier_state(self, job_metadata: _JobMetadataLike) -> None: |
| 189 | transfer_job = job_metadata.transfer_job |
| 190 | state = self._tier_states[job_metadata.tier_idx] |
| 191 | chunk_count = len(transfer_job.chunk_ids) |
| 192 | if transfer_job.is_promotion: |
| 193 | assert state.active_promotion_count > 0 |
| 194 | state.active_promotion_count -= 1 |
| 195 | state.primary_write_chunk_count -= chunk_count |
| 196 | assert state.primary_write_chunk_count >= 0 |
| 197 | else: |
| 198 | assert state.active_cascade_count > 0 |
| 199 | state.active_cascade_count -= 1 |
| 200 | state.primary_read_chunk_count -= chunk_count |
| 201 | assert state.primary_read_chunk_count >= 0 |
| 202 | |
| 203 | def _observe_finished_job_stats( |
| 204 | self, |
| 205 | job_metadata: _JobMetadataLike, |
| 206 | completed_job: JobResult, |
| 207 | ) -> None: |
| 208 | transfer_job = job_metadata.transfer_job |
| 209 | labelvalues = self.tier_label(job_metadata.tier_idx) |
| 210 | completed_key_count = len(transfer_job.keys) |
| 211 | if not completed_job.success: |
| 212 | failure_metric = ( |
| 213 | TieringOffloadingMetrics.PROMOTION_JOB_FAILURES |
| 214 | if transfer_job.is_promotion |
| 215 | else TieringOffloadingMetrics.CASCADE_JOB_FAILURES |
| 216 | ) |
| 217 | self._stats.increase_counter(failure_metric, labelvalues=labelvalues) |
| 218 | if transfer_job.is_promotion and completed_job.successful_keys: |
| 219 | completed_key_count = len(completed_job.successful_keys) |
| 220 | else: |
| 221 | return |
| 222 | |
| 223 | bytes_metric = ( |
| 224 | TieringOffloadingMetrics.READ_BYTES |
| 225 | if transfer_job.is_promotion |
| 226 | else TieringOffloadingMetrics.WRITE_BYTES |
| 227 | ) |
| 228 | time_metric = ( |
| 229 | TieringOffloadingMetrics.READ_TIME |
| 230 | if transfer_job.is_promotion |
| 231 | else TieringOffloadingMetrics.WRITE_TIME |
| 232 | ) |
| 233 | transfer_size = completed_key_count * self._primary_chunk_size |
| 234 | self._stats.increase_counter(bytes_metric, transfer_size, labelvalues) |
| 235 | if completed_job.transfer_time is not None: |
| 236 | self._stats.increase_counter( |
| 237 | time_metric, completed_job.transfer_time, labelvalues |
| 238 | ) |
| 239 | |
| 240 | def _observe_active_transfer_stats(self, stats: OffloadingConnectorStats) -> None: |
| 241 | for tier_idx, state in enumerate(self._tier_states): |
| 242 | labelvalues = self.tier_label(tier_idx) |
| 243 | write_usage = ( |
| 244 | state.primary_write_chunk_count / self._num_primary_chunks |
| 245 | if self._num_primary_chunks > 0 |
| 246 | else 0.0 |
| 247 | ) |
| 248 | read_usage = ( |
| 249 | state.primary_read_chunk_count / self._num_primary_chunks |
| 250 | if self._num_primary_chunks > 0 |
| 251 | else 0.0 |
| 252 | ) |
| 253 | stats.set_gauge( |
| 254 | TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, |
| 255 | write_usage, |
| 256 | labelvalues, |
| 257 | ) |
| 258 | stats.set_gauge( |
| 259 | TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, |
| 260 | read_usage, |
| 261 | labelvalues, |
| 262 | ) |
| 263 | stats.set_gauge( |
| 264 | TieringOffloadingMetrics.ACTIVE_PROMOTION_JOBS, |
| 265 | state.active_promotion_count, |
| 266 | labelvalues, |
| 267 | ) |
| 268 | stats.set_gauge( |
| 269 | TieringOffloadingMetrics.ACTIVE_CASCADE_JOBS, |
| 270 | state.active_cascade_count, |
| 271 | labelvalues, |
| 272 | ) |
| 273 | |
| 274 | def _observe_resolved_lookup( |
| 275 | self, |
| 276 | tier_label: TierLabel, |
| 277 | result: LookupResult, |
| 278 | lookup_duration: float, |
| 279 | async_start_time: float | None, |
| 280 | ) -> None: |
| 281 | self._stats.increase_counter( |
| 282 | TieringOffloadingMetrics.CHUNK_QUERIES, |
| 283 | labelvalues=tier_label, |
| 284 | ) |
| 285 | if result is LookupResult.HIT: |
| 286 | self._stats.increase_counter( |
| 287 | TieringOffloadingMetrics.CHUNK_HITS, |
| 288 | labelvalues=tier_label, |
| 289 | ) |
| 290 | self._stats.observe_histogram( |
| 291 | TieringOffloadingMetrics.LOOKUP_SYNC_DELAY, |
| 292 | lookup_duration, |
| 293 | tier_label, |
| 294 | ) |
| 295 | if async_start_time is not None: |
| 296 | self._stats.observe_histogram( |
| 297 | TieringOffloadingMetrics.LOOKUP_ASYNC_DELAY, |
| 298 | time.monotonic() - async_start_time, |
| 299 | tier_label, |
| 300 | ) |