| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import time |
| 6 | from abc import ABC, abstractmethod |
| 7 | |
| 8 | |
| 9 | class BackpressurePolicy(ABC): |
| 10 | """Decides what action to take when a detector signals pressure.""" |
| 11 | |
| 12 | @abstractmethod |
| 13 | def should_store(self, detector: BackpressureDetector) -> bool: ... |
| 14 | |
| 15 | @abstractmethod |
| 16 | def on_store_skipped(self, num_blocks: int) -> None: ... |
| 17 | |
| 18 | @abstractmethod |
| 19 | def pop_stores_dropped(self) -> tuple[int, int]: |
| 20 | """Return and reset (stores_dropped, blocks_dropped).""" |
| 21 | ... |
| 22 | |
| 23 | @abstractmethod |
| 24 | def reset(self) -> None: ... |
| 25 | |
| 26 | |
| 27 | class DropAccountingPolicy(BackpressurePolicy): |
| 28 | """Base policy that tracks dropped store/block counters. |
| 29 | |
| 30 | Subclasses implement ``should_store``; the shared drop accounting |
| 31 | (``on_store_skipped``/``pop_stores_dropped``/``reset``) lives here so |
| 32 | every policy reports these counters identically. Subclasses that add |
| 33 | their own state override ``reset`` and call ``super().reset()``. |
| 34 | """ |
| 35 | |
| 36 | def __init__(self): |
| 37 | self._stores_dropped: int = 0 |
| 38 | self._blocks_dropped: int = 0 |
| 39 | |
| 40 | def on_store_skipped(self, num_blocks: int) -> None: |
| 41 | self._stores_dropped += 1 |
| 42 | self._blocks_dropped += num_blocks |
| 43 | |
| 44 | def pop_stores_dropped(self) -> tuple[int, int]: |
| 45 | stores, blocks = self._stores_dropped, self._blocks_dropped |
| 46 | self._stores_dropped = 0 |
| 47 | self._blocks_dropped = 0 |
| 48 | return stores, blocks |
| 49 | |
| 50 | def reset(self) -> None: |
| 51 | self._stores_dropped = 0 |
| 52 | self._blocks_dropped = 0 |
| 53 | |
| 54 | |
| 55 | class DropStorePolicy(DropAccountingPolicy): |
| 56 | """Silently drop stores to pressured tiers.""" |
| 57 | |
| 58 | def should_store(self, detector) -> bool: |
| 59 | return not detector.is_under_pressure() |
| 60 | |
| 61 | |
| 62 | class ThrottledDropPolicy(DropAccountingPolicy): |
| 63 | """Drop stores proportionally to pressure severity. |
| 64 | |
| 65 | Instead of dropping all stores when pressure is detected, the drop |
| 66 | rate scales with how far the EMA exceeds the high watermark:: |
| 67 | |
| 68 | overshoot = (ema - high_water) / high_water |
| 69 | drop_rate = min(1.0, overshoot / ramp_factor) |
| 70 | |
| 71 | With the default ``ramp_factor=2.0``, stores ramp from 0% drop at |
| 72 | the high watermark to 100% drop at 3x the high watermark. This |
| 73 | keeps secondary tiers populated under moderate pressure, reducing |
| 74 | cache misses (and therefore recomputation / ITL regression) while |
| 75 | still protecting TTFT under severe congestion. |
| 76 | """ |
| 77 | |
| 78 | def __init__(self, ramp_factor: float = 2.0): |
| 79 | super().__init__() |
| 80 | self._call_count: int = 0 |
| 81 | self._ramp_factor = ramp_factor |
| 82 | |
| 83 | def should_store(self, detector: BackpressureDetector) -> bool: |
| 84 | # Proportional throttling needs the EMA and high watermark, which |
| 85 | # only the EMA detector exposes. |
| 86 | assert isinstance(detector, EMABackpressureDetector), ( |
| 87 | "ThrottledDropPolicy requires an EMABackpressureDetector" |
| 88 | ) |
| 89 | if not detector.is_under_pressure(): |
| 90 | return True |
| 91 | ema = detector.store_latency_ema |
| 92 | high = detector._high |
| 93 | if high <= 0: |
| 94 | return False |
| 95 | overshoot = max(0.0, ema - high) / high |
| 96 | drop_rate = min(1.0, overshoot / self._ramp_factor) |
| 97 | if drop_rate <= 0: |
| 98 | return True |
| 99 | if drop_rate >= 1.0: |
| 100 | return False |
| 101 | self._call_count += 1 |
| 102 | period = max(2, round(1.0 / drop_rate)) |
| 103 | return (self._call_count % period) != 0 |
| 104 | |
| 105 | def reset(self) -> None: |
| 106 | super().reset() |
| 107 | self._call_count = 0 |
| 108 | |
| 109 | |
| 110 | class BackpressureDetector(ABC): |
| 111 | """Observes store completion signals and determines pressure state.""" |
| 112 | |
| 113 | def __init__( |
| 114 | self, |
| 115 | policy: BackpressurePolicy | None = None, |
| 116 | ): |
| 117 | self._policy = policy or ThrottledDropPolicy() |
| 118 | |
| 119 | @property |
| 120 | def policy(self) -> BackpressurePolicy: |
| 121 | return self._policy |
| 122 | |
| 123 | @abstractmethod |
| 124 | def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: ... |
| 125 | |
| 126 | @abstractmethod |
| 127 | def is_under_pressure(self) -> bool: ... |
| 128 | |
| 129 | @abstractmethod |
| 130 | def reset(self) -> None: ... |
| 131 | |
| 132 | def update(self, submit_time: float, num_bytes: int) -> None: |
| 133 | """Update pressure state from a completed store job. |
| 134 | |
| 135 | Args: |
| 136 | submit_time: ``time.monotonic()`` when the job was submitted. |
| 137 | num_bytes: Total bytes written (num_blocks * block_size_bytes). |
| 138 | |
| 139 | """ |
| 140 | if num_bytes <= 0: |
| 141 | return |
| 142 | elapsed = time.monotonic() - submit_time |
| 143 | self.on_store_completed(elapsed, num_bytes) |
| 144 | |
| 145 | def should_store(self, num_blocks: int) -> bool: |
| 146 | """Check policy and record skip if rejected.""" |
| 147 | if not self._policy.should_store(self): |
| 148 | self._policy.on_store_skipped(num_blocks) |
| 149 | return False |
| 150 | return True |
| 151 | |
| 152 | @classmethod |
| 153 | def default_config( |
| 154 | cls, tier_type: str, *, locality: str | None = None |
| 155 | ) -> dict | None: |
| 156 | """Return default constructor kwargs for ``tier_type``, or None. |
| 157 | |
| 158 | Args: |
| 159 | tier_type: The tier type string (e.g. ``"fs"``, ``"obj"``). |
| 160 | locality: Optional locality hint from the tier config |
| 161 | (``"LOCAL"`` or ``"REMOTE"``). When ``"REMOTE"``, tiers |
| 162 | that would normally get local-storage watermarks (e.g. |
| 163 | ``"fs"``) receive network watermarks instead. |
| 164 | |
| 165 | """ |
| 166 | return None |
| 167 | |
| 168 | @property |
| 169 | def stats(self) -> dict[str, float]: |
| 170 | return {} |
| 171 | |
| 172 | |
| 173 | class EMABackpressureDetector(BackpressureDetector): |
| 174 | """EMA of store latency normalized by transfer size. |
| 175 | |
| 176 | The EMA tracks seconds per megabyte (s/MiB) so that the metric is |
| 177 | comparable regardless of how many blocks are in a job or how large |
| 178 | each block is. Water marks are in the same unit. |
| 179 | |
| 180 | Default water marks are derived from fio benchmarks on the WDC H100 |
| 181 | cluster. Two presets are provided: |
| 182 | |
| 183 | LOCAL (NVMe/SSD): NVMe sustains ~5 GB/s writes. |
| 184 | high=0.005 s/MiB (~200 MB/s) catches severe congestion; |
| 185 | low=0.001 s/MiB (~1 GB/s) requires meaningful recovery. |
| 186 | |
| 187 | NETWORK (CephFS, object store, ``obj``tiers, or any tier |
| 188 | with ``"locality": "REMOTE"``): CephFS sustains ~1.5 GB/s. |
| 189 | high=0.020 s/MiB (~50 MB/s); low=0.005 s/MiB (~200 MB/s). |
| 190 | |
| 191 | ``obj`` tiers get NETWORK defaults automatically. An |
| 192 | ``fs`` tier defaults to LOCAL (appropriate for NVMe/SSD); set |
| 193 | ``"locality": "REMOTE"`` in the tier config for network-backed |
| 194 | filesystems like CephFS. |
| 195 | """ |
| 196 | |
| 197 | _MIB = 1 << 20 |
| 198 | |
| 199 | # EMA smoothing factor: higher values (→1) react faster to latency |
| 200 | # spikes but are noisier; lower values (→0) smooth more but lag. |
| 201 | DEFAULT_ALPHA = 0.15 |
| 202 | |
| 203 | # Number of store completions to collect before seeding the EMA. |
| 204 | # During warmup, pressure is never signalled; the EMA is initialized |
| 205 | # to the mean of the warmup samples to avoid cold-start false positives. |
| 206 | DEFAULT_WARMUP_COMPLETIONS = 10 |
| 207 | |
| 208 | # When under pressure and no completions arrive, the EMA decays |
| 209 | # toward zero so the detector can eventually recover. The half-life |
| 210 | # controls how quickly: after this many seconds of silence the EMA |
| 211 | # halves. With a 2 s half-life and high_water=0.005, an EMA frozen |
| 212 | # at 0.010 reaches low_water=0.001 in ~6.6 s. |
| 213 | DEFAULT_DECAY_HALF_LIFE_S = 2.0 |
| 214 | |
| 215 | # After pressure clears, stores are throttled to 50% for this many |
| 216 | # seconds to prevent a burst of stores from re-overwhelming the tier. |
| 217 | DEFAULT_COOLDOWN_S = 1.0 |
| 218 | |
| 219 | # Number of consecutive sub-low-watermark completions before the |
| 220 | # detector enters a "healthy" fast-path that bypasses policy checks. |
| 221 | _HEALTHY_THRESHOLD = 10 |
| 222 | |
| 223 | LOCAL_HIGH_WATER_S = 0.010 |
| 224 | LOCAL_LOW_WATER_S = 0.002 |
| 225 | |
| 226 | NETWORK_HIGH_WATER_S = 0.050 |
| 227 | NETWORK_LOW_WATER_S = 0.010 |
| 228 | |
| 229 | _NETWORK_TIER_TYPES = frozenset({"obj"}) |
| 230 | |
| 231 | @classmethod |
| 232 | def default_config( |
| 233 | cls, tier_type: str, *, locality: str | None = None |
| 234 | ) -> dict | None: |
| 235 | is_remote = locality is not None and locality.upper() == "REMOTE" |
| 236 | if is_remote or tier_type in cls._NETWORK_TIER_TYPES: |
| 237 | return { |
| 238 | "high_water_s": cls.NETWORK_HIGH_WATER_S, |
| 239 | "low_water_s": cls.NETWORK_LOW_WATER_S, |
| 240 | } |
| 241 | return { |
| 242 | "high_water_s": cls.LOCAL_HIGH_WATER_S, |
| 243 | "low_water_s": cls.LOCAL_LOW_WATER_S, |
| 244 | } |
| 245 | |
| 246 | def __init__( |
| 247 | self, |
| 248 | high_water_s: float, |
| 249 | low_water_s: float, |
| 250 | alpha: float = DEFAULT_ALPHA, |
| 251 | warmup_completions: int = DEFAULT_WARMUP_COMPLETIONS, |
| 252 | decay_half_life_s: float = DEFAULT_DECAY_HALF_LIFE_S, |
| 253 | cooldown_s: float = DEFAULT_COOLDOWN_S, |
| 254 | policy: BackpressurePolicy | None = None, |
| 255 | ): |
| 256 | super().__init__(policy=policy) |
| 257 | if not (0 < alpha <= 1): |
| 258 | raise ValueError(f"alpha must be in (0, 1], got {alpha}") |
| 259 | if low_water_s > high_water_s: |
| 260 | raise ValueError( |
| 261 | f"low_water_s ({low_water_s}) must be <= high_water_s ({high_water_s})" |
| 262 | ) |
| 263 | |
| 264 | self._alpha = alpha |
| 265 | self._high = high_water_s |
| 266 | self._low = low_water_s |
| 267 | self._warmup_completions = warmup_completions |
| 268 | self._decay_half_life_s = decay_half_life_s |
| 269 | self._cooldown_s = cooldown_s |
| 270 | # Exponential moving average of store latency in seconds. |
| 271 | self._ema: float = 0.0 |
| 272 | # Whether the tier is currently considered under pressure. |
| 273 | self._under_pressure: bool = False |
| 274 | # Number of store completions observed (for warmup gating). |
| 275 | self._completions: int = 0 |
| 276 | # Samples collected during warmup; used to seed the EMA. |
| 277 | self._warmup_samples: list[float] = [] |
| 278 | # Monotonic timestamp of the last EMA update; used for idle decay. |
| 279 | self._last_update: float = 0.0 |
| 280 | # Consecutive completions with EMA below low watermark. |
| 281 | self._healthy_streak: int = 0 |
| 282 | # Monotonic timestamp when pressure was last cleared (for cooldown). |
| 283 | self._pressure_cleared_at: float = 0.0 |
| 284 | self._cooldown_count: int = 0 |
| 285 | |
| 286 | def _apply_idle_decay(self) -> None: |
| 287 | """Decay the EMA toward zero based on elapsed idle time. |
| 288 | |
| 289 | Uses the formula ``ema * 0.5^(dt / half_life)`` so the EMA |
| 290 | halves every ``decay_half_life_s`` seconds of silence. This |
| 291 | lets the detector recover from pressure when no stores are |
| 292 | attempted (and therefore no completions arrive to update the |
| 293 | EMA organically). |
| 294 | """ |
| 295 | now = time.monotonic() |
| 296 | dt = now - self._last_update |
| 297 | if dt <= 0 or self._decay_half_life_s <= 0: |
| 298 | return |
| 299 | self._ema *= 0.5 ** (dt / self._decay_half_life_s) |
| 300 | self._last_update = now |
| 301 | |
| 302 | def on_store_completed(self, elapsed_s: float, num_bytes: int) -> None: |
| 303 | now = time.monotonic() |
| 304 | s_per_mib = elapsed_s / (num_bytes / self._MIB) |
| 305 | self._completions += 1 |
| 306 | if self._completions <= self._warmup_completions: |
| 307 | self._warmup_samples.append(s_per_mib) |
| 308 | if self._completions == self._warmup_completions: |
| 309 | self._ema = sum(self._warmup_samples) / len(self._warmup_samples) |
| 310 | self._last_update = now |
| 311 | return |
| 312 | self._ema = self._alpha * s_per_mib + (1 - self._alpha) * self._ema |
| 313 | self._last_update = now |
| 314 | if self._ema > self._high: |
| 315 | self._under_pressure = True |
| 316 | self._pressure_cleared_at = 0.0 |
| 317 | self._healthy_streak = 0 |
| 318 | elif self._ema < self._low: |
| 319 | if self._under_pressure: |
| 320 | self._pressure_cleared_at = now |
| 321 | self._cooldown_count = 0 |
| 322 | self._under_pressure = False |
| 323 | self._healthy_streak += 1 |
| 324 | else: |
| 325 | self._healthy_streak = 0 |
| 326 | |
| 327 | def is_under_pressure(self) -> bool: |
| 328 | if self._under_pressure and self._completions >= self._warmup_completions: |
| 329 | self._apply_idle_decay() |
| 330 | if self._ema < self._low: |
| 331 | self._under_pressure = False |
| 332 | self._pressure_cleared_at = time.monotonic() |
| 333 | self._cooldown_count = 0 |
| 334 | return self._under_pressure |
| 335 | |
| 336 | def is_healthy(self) -> bool: |
| 337 | """True when the tier has been consistently fast. |
| 338 | |
| 339 | After ``_HEALTHY_THRESHOLD`` consecutive store completions with |
| 340 | EMA below the low watermark (and no active pressure), the |
| 341 | detector enters a healthy fast-path: ``should_store()`` returns |
| 342 | ``True`` immediately, skipping the policy check entirely. |
| 343 | """ |
| 344 | return ( |
| 345 | not self._under_pressure and self._healthy_streak >= self._HEALTHY_THRESHOLD |
| 346 | ) |
| 347 | |
| 348 | def should_store(self, num_blocks: int) -> bool: |
| 349 | under_pressure = self.is_under_pressure() |
| 350 | if self.is_healthy(): |
| 351 | return True |
| 352 | if ( |
| 353 | not under_pressure |
| 354 | and self._pressure_cleared_at > 0 |
| 355 | and time.monotonic() - self._pressure_cleared_at < self._cooldown_s |
| 356 | ): |
| 357 | self._cooldown_count += 1 |
| 358 | if self._cooldown_count % 2 == 0: |
| 359 | self._policy.on_store_skipped(num_blocks) |
| 360 | return False |
| 361 | return super().should_store(num_blocks) |
| 362 | |
| 363 | def reset(self) -> None: |
| 364 | self._ema = 0.0 |
| 365 | self._under_pressure = False |
| 366 | self._completions = 0 |
| 367 | self._warmup_samples.clear() |
| 368 | self._last_update = 0.0 |
| 369 | self._healthy_streak = 0 |
| 370 | self._pressure_cleared_at = 0.0 |
| 371 | self._cooldown_count = 0 |
| 372 | self._policy.reset() |
| 373 | |
| 374 | @property |
| 375 | def store_latency_ema(self) -> float: |
| 376 | return self._ema |
| 377 | |
| 378 | @store_latency_ema.setter |
| 379 | def store_latency_ema(self, value: float) -> None: |
| 380 | self._ema = value |
| 381 | |
| 382 | @property |
| 383 | def stats(self) -> dict[str, float]: |
| 384 | return {"store_latency_ema": self._ema} |