| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for back-pressure detection in TieringOffloadingManager. |
| 4 | |
| 5 | These tests use a DelayedSecondaryTierManager that holds completed jobs |
| 6 | until explicitly released, allowing precise control over when the manager |
| 7 | observes store completions and their apparent latency. |
| 8 | """ |
| 9 | |
| 10 | import time |
| 11 | from collections.abc import Iterable |
| 12 | from unittest.mock import MagicMock, patch |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from vllm.v1.kv_offload.base import ( |
| 18 | LookupResult, |
| 19 | OffloadKey, |
| 20 | OffloadPolicy, |
| 21 | ReqContext, |
| 22 | RequestOffloadingContext, |
| 23 | ScheduleEndContext, |
| 24 | make_offload_key, |
| 25 | ) |
| 26 | from vllm.v1.kv_offload.tiering.backpressure import ( |
| 27 | DropAccountingPolicy, |
| 28 | DropStorePolicy, |
| 29 | EMABackpressureDetector, |
| 30 | ThrottledDropPolicy, |
| 31 | ) |
| 32 | from vllm.v1.kv_offload.tiering.base import ( |
| 33 | JobResult, |
| 34 | SecondaryTierManager, |
| 35 | TieringOffloadingMetrics, |
| 36 | TransferJob, |
| 37 | ) |
| 38 | from vllm.v1.kv_offload.tiering.manager import ( |
| 39 | CPUPrimaryTierOffloadingManager, |
| 40 | TieringOffloadingManager, |
| 41 | ) |
| 42 | |
| 43 | _BP_EMA_ALPHA = EMABackpressureDetector.DEFAULT_ALPHA |
| 44 | _BP_WARMUP = EMABackpressureDetector.DEFAULT_WARMUP_COMPLETIONS |
| 45 | # Water marks in s/MiB, scaled for the test's tiny 16-byte blocks. |
| 46 | _BP_HIGH_WATER_S = 1000.0 |
| 47 | _BP_LOW_WATER_S = 500.0 |
| 48 | |
| 49 | _CTX = ReqContext(req_id="test") |
| 50 | _MOCK_OFFLOADING_SPEC = MagicMock() |
| 51 | |
| 52 | |
| 53 | def _mock_mmap_region(num_blocks: int, row_bytes: int = 16): |
| 54 | mock = MagicMock() |
| 55 | view = memoryview(torch.zeros((num_blocks, row_bytes), dtype=torch.int8).numpy()) |
| 56 | mock.create_kv_memoryview.return_value = view |
| 57 | return mock |
| 58 | |
| 59 | |
| 60 | def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: |
| 61 | return [make_offload_key(str(i).encode(), 0) for i in int_ids] |
| 62 | |
| 63 | |
| 64 | class DelayedSecondaryTierManager(SecondaryTierManager): |
| 65 | """Secondary tier that holds completed store jobs until released. |
| 66 | |
| 67 | Jobs are stored immediately but their completion is not reported |
| 68 | by get_finished_jobs() until release_jobs() is called. This lets |
| 69 | tests control the apparent latency of store operations by adjusting |
| 70 | submit_time on the JobMetadata before triggering the completion poll. |
| 71 | """ |
| 72 | |
| 73 | def __init__( |
| 74 | self, offloading_spec, primary_kv_view, tier_type, backpressure_detector=None |
| 75 | ): |
| 76 | super().__init__( |
| 77 | offloading_spec, |
| 78 | primary_kv_view, |
| 79 | tier_type, |
| 80 | backpressure_detector=backpressure_detector, |
| 81 | ) |
| 82 | self.blocks: dict[OffloadKey, bool] = {} |
| 83 | self._held_jobs: list[JobResult] = [] |
| 84 | self._released_jobs: list[JobResult] = [] |
| 85 | self._request_policy: OffloadPolicy | None = None |
| 86 | |
| 87 | def lookup(self, key, req_context): |
| 88 | return LookupResult.HIT if key in self.blocks else LookupResult.MISS |
| 89 | |
| 90 | def submit_store(self, job_metadata: TransferJob) -> None: |
| 91 | for key in job_metadata.keys: |
| 92 | self.blocks[key] = True |
| 93 | self._held_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) |
| 94 | |
| 95 | def submit_load(self, job_metadata: TransferJob) -> None: |
| 96 | self._released_jobs.append(JobResult(job_id=job_metadata.job_id, success=True)) |
| 97 | |
| 98 | def get_finished_jobs(self) -> Iterable[JobResult]: |
| 99 | result = self._released_jobs |
| 100 | self._released_jobs = [] |
| 101 | return result |
| 102 | |
| 103 | def release_jobs(self): |
| 104 | self._released_jobs.extend(self._held_jobs) |
| 105 | self._held_jobs.clear() |
| 106 | |
| 107 | def on_new_request(self, req_context): |
| 108 | if self._request_policy is not None: |
| 109 | return RequestOffloadingContext(policy=self._request_policy) |
| 110 | return RequestOffloadingContext() |
| 111 | |
| 112 | def drain_jobs(self): |
| 113 | self.release_jobs() |
| 114 | |
| 115 | def has_pending_work(self): |
| 116 | return bool(self._held_jobs) |
| 117 | |
| 118 | def get_num_blocks(self): |
| 119 | return len(self.blocks) |
| 120 | |
| 121 | |
| 122 | class TestDefaultConfig: |
| 123 | def test_fs_gets_local_watermarks(self): |
| 124 | cfg = EMABackpressureDetector.default_config("fs") |
| 125 | assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S |
| 126 | assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S |
| 127 | |
| 128 | def test_obj_gets_network_watermarks(self): |
| 129 | cfg = EMABackpressureDetector.default_config("obj") |
| 130 | assert cfg["high_water_s"] == EMABackpressureDetector.NETWORK_HIGH_WATER_S |
| 131 | assert cfg["low_water_s"] == EMABackpressureDetector.NETWORK_LOW_WATER_S |
| 132 | |
| 133 | def test_p2p_gets_local_watermarks(self): |
| 134 | cfg = EMABackpressureDetector.default_config("p2p") |
| 135 | assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S |
| 136 | assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S |
| 137 | |
| 138 | def test_fs_with_remote_locality_gets_network_watermarks(self): |
| 139 | cfg = EMABackpressureDetector.default_config("fs", locality="REMOTE") |
| 140 | assert cfg["high_water_s"] == EMABackpressureDetector.NETWORK_HIGH_WATER_S |
| 141 | assert cfg["low_water_s"] == EMABackpressureDetector.NETWORK_LOW_WATER_S |
| 142 | |
| 143 | def test_fs_with_local_locality_gets_local_watermarks(self): |
| 144 | cfg = EMABackpressureDetector.default_config("fs", locality="LOCAL") |
| 145 | assert cfg["high_water_s"] == EMABackpressureDetector.LOCAL_HIGH_WATER_S |
| 146 | assert cfg["low_water_s"] == EMABackpressureDetector.LOCAL_LOW_WATER_S |
| 147 | |
| 148 | |
| 149 | class TestIdleDecay: |
| 150 | """Tests for EMA idle decay — recovery from pressure without completions.""" |
| 151 | |
| 152 | def _make_detector(self, **kwargs): |
| 153 | defaults = dict( |
| 154 | high_water_s=_BP_HIGH_WATER_S, |
| 155 | low_water_s=_BP_LOW_WATER_S, |
| 156 | decay_half_life_s=2.0, |
| 157 | ) |
| 158 | defaults.update(kwargs) |
| 159 | return EMABackpressureDetector(**defaults) |
| 160 | |
| 161 | def test_ema_decays_while_under_pressure(self): |
| 162 | bp = self._make_detector() |
| 163 | bp._completions = _BP_WARMUP |
| 164 | bp._ema = _BP_HIGH_WATER_S * 2 |
| 165 | bp._under_pressure = True |
| 166 | bp._last_update = time.monotonic() |
| 167 | |
| 168 | with patch("time.monotonic", return_value=bp._last_update + 4.1): |
| 169 | # ~2.05 half-lives → EMA * 0.24, well below low water (500) |
| 170 | assert bp.is_under_pressure() is False |
| 171 | assert bp.store_latency_ema < _BP_LOW_WATER_S |
| 172 | |
| 173 | def test_ema_stays_pressured_if_not_enough_decay(self): |
| 174 | bp = self._make_detector() |
| 175 | bp._completions = _BP_WARMUP |
| 176 | bp._ema = _BP_HIGH_WATER_S * 4 |
| 177 | bp._under_pressure = True |
| 178 | bp._last_update = time.monotonic() |
| 179 | |
| 180 | with patch("time.monotonic", return_value=bp._last_update + 2.0): |
| 181 | # 2s = 1 half-life → EMA * 0.5 = HIGH * 2, still above LOW |
| 182 | assert bp.is_under_pressure() is True |
| 183 | assert bp.store_latency_ema == pytest.approx(_BP_HIGH_WATER_S * 2, rel=0.01) |
| 184 | |
| 185 | def test_no_decay_before_warmup(self): |
| 186 | bp = self._make_detector() |
| 187 | bp._ema = _BP_HIGH_WATER_S * 2 |
| 188 | bp._under_pressure = True |
| 189 | bp._last_update = time.monotonic() |
| 190 | |
| 191 | with patch("time.monotonic", return_value=bp._last_update + 10.0): |
| 192 | # Still in warmup (_completions=0 < _warmup_completions), |
| 193 | # so decay should not be applied in is_under_pressure. |
| 194 | assert bp.is_under_pressure() is True |
| 195 | assert bp.store_latency_ema == _BP_HIGH_WATER_S * 2 |
| 196 | |
| 197 | def test_no_decay_applied_during_new_sample(self): |
| 198 | """on_store_completed applies the new sample directly without |
| 199 | idle decay — decay only runs in is_under_pressure().""" |
| 200 | bp = self._make_detector() |
| 201 | bp._completions = _BP_WARMUP |
| 202 | bp._ema = _BP_HIGH_WATER_S * 2 |
| 203 | bp._under_pressure = True |
| 204 | t0 = time.monotonic() |
| 205 | bp._last_update = t0 |
| 206 | |
| 207 | block_bytes = 16 |
| 208 | with patch("time.monotonic", return_value=t0 + 4.0): |
| 209 | bp.on_store_completed(0.0, block_bytes) |
| 210 | # No decay in on_store_completed: EMA = alpha*0 + (1-alpha)*2000 |
| 211 | expected = (1 - _BP_EMA_ALPHA) * _BP_HIGH_WATER_S * 2 |
| 212 | assert bp.store_latency_ema == pytest.approx(expected, rel=0.01) |
| 213 | |
| 214 | def test_reset_clears_last_update(self): |
| 215 | bp = self._make_detector() |
| 216 | bp._last_update = time.monotonic() |
| 217 | bp.reset() |
| 218 | assert bp._last_update == 0.0 |
| 219 | |
| 220 | |
| 221 | class TestBackpressure: |
| 222 | @pytest.fixture |
| 223 | def setup(self): |
| 224 | mock_region = _mock_mmap_region(20) |
| 225 | self.primary = CPUPrimaryTierOffloadingManager( |
| 226 | num_chunks=20, mmap_region=mock_region |
| 227 | ) |
| 228 | mock_view = mock_region.create_kv_memoryview() |
| 229 | self.tier = DelayedSecondaryTierManager( |
| 230 | offloading_spec=_MOCK_OFFLOADING_SPEC, |
| 231 | primary_kv_view=mock_view, |
| 232 | tier_type="delayed", |
| 233 | backpressure_detector=EMABackpressureDetector( |
| 234 | high_water_s=_BP_HIGH_WATER_S, |
| 235 | low_water_s=_BP_LOW_WATER_S, |
| 236 | ), |
| 237 | ) |
| 238 | self.manager = TieringOffloadingManager( |
| 239 | primary_tier=self.primary, |
| 240 | secondary_tiers=[self.tier], |
| 241 | ) |
| 242 | |
| 243 | def _start_request(self, ctx=_CTX): |
| 244 | if ctx.req_id not in self.manager._req_state: |
| 245 | self.manager.on_new_request(ctx) |
| 246 | |
| 247 | def _store_blocks(self, keys, ctx=_CTX): |
| 248 | self._start_request(ctx) |
| 249 | result = self.manager.prepare_store(keys, ctx) |
| 250 | assert result is not None |
| 251 | self.manager.complete_store(keys, ctx, success=True) |
| 252 | |
| 253 | def _simulate_on_schedule_end(self): |
| 254 | ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) |
| 255 | self.manager.on_schedule_end(ctx) |
| 256 | |
| 257 | def _backdate_held_jobs(self, age_s: float): |
| 258 | """Set submit_time on all in-flight store jobs so they appear |
| 259 | to have taken ``age_s`` seconds when completed.""" |
| 260 | self.manager._processed_jobs_this_step = False |
| 261 | now = time.monotonic() |
| 262 | for job_id, meta in self.manager._jobs.items(): |
| 263 | tj = meta.transfer_job |
| 264 | if not tj.is_promotion and tj.submit_time > 0: |
| 265 | tj.submit_time = now - age_s |
| 266 | |
| 267 | def test_ema_updates_on_store_completion(self, setup): |
| 268 | bp = self.tier.bp_detector |
| 269 | fast_latency = 0.01 |
| 270 | # EMA is in s/MiB; with 16-byte blocks, scale = MiB / block_bytes. |
| 271 | scale = EMABackpressureDetector._MIB / self.tier.block_size_bytes |
| 272 | |
| 273 | # During warmup, EMA stays at 0. |
| 274 | for i in range(_BP_WARMUP - 1): |
| 275 | keys = to_keys([300 + i]) |
| 276 | self._store_blocks(keys) |
| 277 | self._backdate_held_jobs(fast_latency) |
| 278 | self.tier.release_jobs() |
| 279 | self._simulate_on_schedule_end() |
| 280 | assert bp.store_latency_ema == 0.0 |
| 281 | |
| 282 | # Final warmup sample seeds EMA with the mean. |
| 283 | keys = to_keys([400]) |
| 284 | self._store_blocks(keys) |
| 285 | self._backdate_held_jobs(fast_latency) |
| 286 | self.tier.release_jobs() |
| 287 | self._simulate_on_schedule_end() |
| 288 | fast_s_per_mib = fast_latency * scale |
| 289 | assert bp.store_latency_ema == pytest.approx(fast_s_per_mib, rel=0.1) |
| 290 | |
| 291 | # After warmup, EMA updates normally. |
| 292 | new_latency = 0.05 |
| 293 | keys = to_keys([401]) |
| 294 | self._store_blocks(keys) |
| 295 | self._backdate_held_jobs(new_latency) |
| 296 | self.tier.release_jobs() |
| 297 | self._simulate_on_schedule_end() |
| 298 | new_s_per_mib = new_latency * scale |
| 299 | expected = _BP_EMA_ALPHA * new_s_per_mib + (1 - _BP_EMA_ALPHA) * fast_s_per_mib |
| 300 | assert bp.store_latency_ema == pytest.approx(expected, rel=0.1) |
| 301 | |
| 302 | def test_pressure_activates_above_high_water(self, setup): |
| 303 | bp = self.tier.bp_detector |
| 304 | |
| 305 | # Warm up the detector with fast completions first. |
| 306 | for i in range(_BP_WARMUP): |
| 307 | keys = to_keys([200 + i]) |
| 308 | self._store_blocks(keys) |
| 309 | self._backdate_held_jobs(0.001) |
| 310 | self.tier.release_jobs() |
| 311 | self._simulate_on_schedule_end() |
| 312 | |
| 313 | keys = to_keys([100]) |
| 314 | self._store_blocks(keys) |
| 315 | self._simulate_on_schedule_end() |
| 316 | |
| 317 | # 0.1s for a single 16-byte block → ~6553 s/MiB, well above high water. |
| 318 | self._backdate_held_jobs(0.1) |
| 319 | self.tier.release_jobs() |
| 320 | self._simulate_on_schedule_end() |
| 321 | |
| 322 | assert bp.store_latency_ema > _BP_HIGH_WATER_S |
| 323 | assert bp.is_under_pressure() is True |
| 324 | |
| 325 | def test_pressure_clears_below_low_water(self, setup): |
| 326 | bp = self.tier.bp_detector |
| 327 | bp.store_latency_ema = _BP_HIGH_WATER_S * 2 |
| 328 | bp._under_pressure = True |
| 329 | bp._completions = _BP_WARMUP |
| 330 | |
| 331 | # Drive several fast completions to bring EMA below low water. |
| 332 | block_id = 100 |
| 333 | for _ in range(50): |
| 334 | if bp.store_latency_ema < _BP_LOW_WATER_S: |
| 335 | break |
| 336 | keys = to_keys([block_id]) |
| 337 | block_id += 1 |
| 338 | self._store_blocks(keys) |
| 339 | # Pressure is on, so the cascade was skipped. Force-submit |
| 340 | # a store job manually to get a completion to feed the EMA. |
| 341 | job_meta = self.manager.create_store_job(keys, _CTX) |
| 342 | job_meta.submit_time = time.monotonic() - 0.001 |
| 343 | self.tier.submit_store(job_meta) |
| 344 | self.tier.release_jobs() |
| 345 | self._simulate_on_schedule_end() |
| 346 | else: |
| 347 | pytest.fail(f"EMA did not recover below low water: {bp.store_latency_ema}") |
| 348 | |
| 349 | assert bp.is_under_pressure() is False |
| 350 | |
| 351 | def test_hysteresis_prevents_oscillation(self, setup): |
| 352 | bp = self.tier.bp_detector |
| 353 | bp._completions = _BP_WARMUP |
| 354 | block_bytes = self.tier.block_size_bytes |
| 355 | |
| 356 | # Set EMA between low and high water marks (in s/MiB). |
| 357 | mid = (_BP_LOW_WATER_S + _BP_HIGH_WATER_S) / 2 |
| 358 | bp.store_latency_ema = mid |
| 359 | bp._last_update = time.monotonic() |
| 360 | |
| 361 | # Raw seconds that produce `mid` s/MiB for a given number of blocks. |
| 362 | def mid_latency_s(num_blocks: int) -> float: |
| 363 | return mid * (num_blocks * block_bytes / EMABackpressureDetector._MIB) |
| 364 | |
| 365 | # When not under pressure, mid-range EMA should not activate. |
| 366 | bp._under_pressure = False |
| 367 | keys = to_keys(range(2)) |
| 368 | self._store_blocks(keys) |
| 369 | self._backdate_held_jobs(mid_latency_s(len(keys))) |
| 370 | self.tier.release_jobs() |
| 371 | self._simulate_on_schedule_end() |
| 372 | assert bp.is_under_pressure() is False |
| 373 | |
| 374 | # When under pressure, mid-range EMA should not deactivate. |
| 375 | bp._under_pressure = True |
| 376 | bp._last_update = time.monotonic() |
| 377 | keys2 = to_keys(range(10, 12)) |
| 378 | # Force-submit since pressure is on (cascade would skip). |
| 379 | for k in keys2: |
| 380 | self.primary.prepare_store([k], _CTX) |
| 381 | self.primary.complete_store([k], _CTX) |
| 382 | job_meta = self.manager.create_store_job(keys2, _CTX) |
| 383 | job_meta.submit_time = time.monotonic() - mid_latency_s(len(keys2)) |
| 384 | self.tier.submit_store(job_meta) |
| 385 | self.tier.release_jobs() |
| 386 | self._simulate_on_schedule_end() |
| 387 | assert bp.is_under_pressure() is True |
| 388 | |
| 389 | def test_stores_skipped_under_pressure(self, setup): |
| 390 | bp = self.tier.bp_detector |
| 391 | bp._under_pressure = True |
| 392 | bp.store_latency_ema = _BP_HIGH_WATER_S * 4 |
| 393 | initial_blocks = self.tier.get_num_blocks() |
| 394 | |
| 395 | keys = to_keys(range(50, 53)) |
| 396 | self._store_blocks(keys) |
| 397 | |
| 398 | # Blocks should be in primary but NOT in secondary. |
| 399 | for k in keys: |
| 400 | assert self.primary.lookup(k, _CTX) is LookupResult.HIT |
| 401 | assert self.tier.get_num_blocks() == initial_blocks |
| 402 | |
| 403 | def test_stores_resume_after_pressure_clears(self, setup): |
| 404 | bp = self.tier.bp_detector |
| 405 | |
| 406 | # Start under pressure — stores skipped. |
| 407 | bp._under_pressure = True |
| 408 | bp.store_latency_ema = _BP_HIGH_WATER_S * 4 |
| 409 | keys1 = to_keys(range(60, 62)) |
| 410 | self._store_blocks(keys1) |
| 411 | assert all(k not in self.tier.blocks for k in keys1) |
| 412 | |
| 413 | # Clear pressure — next store should cascade. |
| 414 | bp._under_pressure = False |
| 415 | keys2 = to_keys(range(70, 72)) |
| 416 | self._store_blocks(keys2) |
| 417 | assert all(k in self.tier.blocks for k in keys2) |
| 418 | |
| 419 | def test_dropped_store_count_tracked(self, setup): |
| 420 | bp = self.tier.bp_detector |
| 421 | bp._under_pressure = True |
| 422 | bp.store_latency_ema = _BP_HIGH_WATER_S * 4 |
| 423 | policy = bp.policy |
| 424 | assert policy.pop_stores_dropped() == (0, 0) |
| 425 | |
| 426 | self._store_blocks(to_keys([80])) |
| 427 | assert policy._stores_dropped == 1 |
| 428 | assert policy._blocks_dropped == 1 |
| 429 | |
| 430 | self._store_blocks(to_keys([81, 82])) |
| 431 | assert policy._stores_dropped == 2 |
| 432 | assert policy._blocks_dropped == 3 |
| 433 | |
| 434 | def test_metrics_reported_via_get_stats(self, setup): |
| 435 | bp = self.tier.bp_detector |
| 436 | bp._under_pressure = True |
| 437 | bp.store_latency_ema = 2.5 |
| 438 | policy = bp.policy |
| 439 | policy._stores_dropped = 5 |
| 440 | policy._blocks_dropped = 12 |
| 441 | |
| 442 | stats = self.manager.get_stats() |
| 443 | assert stats is not None |
| 444 | reduced = stats.reduce() |
| 445 | |
| 446 | ema_key = TieringOffloadingMetrics.BACKPRESSURE_STORE_LATENCY_EMA |
| 447 | stores_key = TieringOffloadingMetrics.BACKPRESSURE_STORES_DROPPED |
| 448 | blocks_key = TieringOffloadingMetrics.BACKPRESSURE_BLOCKS_DROPPED |
| 449 | assert reduced[f"{ema_key}:('1:delayed',)"] == pytest.approx(2.5) |
| 450 | assert reduced[f"{stores_key}:('1:delayed',)"] == 5 |
| 451 | assert reduced[f"{blocks_key}:('1:delayed',)"] == 12 |
| 452 | |
| 453 | # Dropped counts reset after get_stats. |
| 454 | assert policy._stores_dropped == 0 |
| 455 | assert policy._blocks_dropped == 0 |
| 456 | |
| 457 | def test_reset_cache_clears_backpressure(self, setup): |
| 458 | bp = self.tier.bp_detector |
| 459 | bp.store_latency_ema = 5.0 |
| 460 | bp._under_pressure = True |
| 461 | |
| 462 | self.manager.reset_cache() |
| 463 | |
| 464 | assert bp.store_latency_ema == 0.0 |
| 465 | assert bp.is_under_pressure() is False |
| 466 | |
| 467 | def test_request_level_tier_respects_backpressure(self, setup): |
| 468 | """Request-level cascade also skips pressured tiers.""" |
| 469 | # Store blocks in primary first. |
| 470 | existing = to_keys(range(3)) |
| 471 | self._store_blocks(existing) |
| 472 | self.tier.release_jobs() |
| 473 | self._simulate_on_schedule_end() |
| 474 | self._simulate_on_schedule_end() |
| 475 | |
| 476 | # Make tier request-level for a new request. |
| 477 | self.tier._request_policy = OffloadPolicy.REQUEST_LEVEL |
| 478 | ctx = ReqContext(req_id="req_rl_bp") |
| 479 | self.manager.on_new_request(ctx) |
| 480 | |
| 481 | # Activate pressure. |
| 482 | bp = self.tier.bp_detector |
| 483 | bp._under_pressure = True |
| 484 | bp.store_latency_ema = _BP_HIGH_WATER_S * 4 |
| 485 | initial_blocks = self.tier.get_num_blocks() |
| 486 | |
| 487 | # prepare_store with existing + new blocks. |
| 488 | new = to_keys(range(10, 12)) |
| 489 | all_keys = existing + new |
| 490 | result = self.manager.prepare_store(all_keys, ctx) |
| 491 | assert result is not None |
| 492 | |
| 493 | # Existing blocks would normally cascade to request-level tier, |
| 494 | # but should be skipped under pressure. |
| 495 | assert self.tier.get_num_blocks() == initial_blocks |
| 496 | |
| 497 | def test_loads_continue_under_pressure(self, setup): |
| 498 | """Loads from a pressured tier still work.""" |
| 499 | blocks = to_keys(range(3)) |
| 500 | for b in blocks: |
| 501 | self.tier.blocks[b] = True |
| 502 | |
| 503 | bp = self.tier.bp_detector |
| 504 | bp._under_pressure = True |
| 505 | bp.store_latency_ema = _BP_HIGH_WATER_S * 4 |
| 506 | |
| 507 | # Lookups should still initiate promotion. |
| 508 | for b in blocks: |
| 509 | result = self.manager.lookup(b, _CTX) |
| 510 | assert result is LookupResult.HIT_PENDING |
| 511 | |
| 512 | |
| 513 | class TestThrottledDropPolicy: |
| 514 | """Tests for proportional throttling under pressure.""" |
| 515 | |
| 516 | def _make_detector(self, **kwargs): |
| 517 | defaults = dict( |
| 518 | high_water_s=_BP_HIGH_WATER_S, |
| 519 | low_water_s=_BP_LOW_WATER_S, |
| 520 | cooldown_s=0.0, |
| 521 | ) |
| 522 | defaults.update(kwargs) |
| 523 | return EMABackpressureDetector(**defaults) |
| 524 | |
| 525 | def test_no_drop_when_not_under_pressure(self): |
| 526 | bp = self._make_detector() |
| 527 | for _ in range(20): |
| 528 | assert bp.should_store(1) is True |
| 529 | |
| 530 | def test_no_drop_at_exactly_high_water(self): |
| 531 | bp = self._make_detector() |
| 532 | bp._under_pressure = True |
| 533 | bp.store_latency_ema = _BP_HIGH_WATER_S |
| 534 | for _ in range(20): |
| 535 | assert bp.should_store(1) is True |
| 536 | |
| 537 | def test_partial_drop_at_moderate_pressure(self): |
| 538 | bp = self._make_detector() |
| 539 | bp._under_pressure = True |
| 540 | bp.store_latency_ema = _BP_HIGH_WATER_S * 2 |
| 541 | results = [bp.should_store(1) for _ in range(20)] |
| 542 | assert any(results), "some stores should be allowed" |
| 543 | assert not all(results), "some stores should be dropped" |
| 544 | |
| 545 | def test_full_drop_at_severe_pressure(self): |
| 546 | bp = self._make_detector() |
| 547 | bp._under_pressure = True |
| 548 | bp.store_latency_ema = _BP_HIGH_WATER_S * 4 |
| 549 | for _ in range(20): |
| 550 | assert bp.should_store(1) is False |
| 551 | |
| 552 | def test_drop_rate_increases_with_severity(self): |
| 553 | """Higher EMA → higher drop rate.""" |
| 554 | drop_counts = [] |
| 555 | for multiplier in [1.5, 2.0, 2.5, 3.5]: |
| 556 | bp = self._make_detector() |
| 557 | bp._under_pressure = True |
| 558 | bp.store_latency_ema = _BP_HIGH_WATER_S * multiplier |
| 559 | drops = sum(1 for _ in range(100) if not bp.should_store(1)) |
| 560 | drop_counts.append(drops) |
| 561 | for i in range(len(drop_counts) - 1): |
| 562 | assert drop_counts[i] <= drop_counts[i + 1] |
| 563 | |
| 564 | def test_explicit_drop_store_policy_still_drops_all(self): |
| 565 | """DropStorePolicy (explicit) still drops everything under pressure.""" |
| 566 | bp = EMABackpressureDetector( |
| 567 | high_water_s=_BP_HIGH_WATER_S, |
| 568 | low_water_s=_BP_LOW_WATER_S, |
| 569 | policy=DropStorePolicy(), |
| 570 | cooldown_s=0.0, |
| 571 | ) |
| 572 | bp._under_pressure = True |
| 573 | bp.store_latency_ema = _BP_HIGH_WATER_S * 1.5 |
| 574 | for _ in range(20): |
| 575 | assert bp.should_store(1) is False |
| 576 | |
| 577 | |
| 578 | class TestSharedDropAccounting: |
| 579 | """Both policies share drop accounting via DropAccountingPolicy.""" |
| 580 | |
| 581 | @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy]) |
| 582 | def test_shared_base(self, policy_cls): |
| 583 | assert issubclass(policy_cls, DropAccountingPolicy) |
| 584 | |
| 585 | @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy]) |
| 586 | def test_pop_returns_and_resets(self, policy_cls): |
| 587 | policy = policy_cls() |
| 588 | policy.on_store_skipped(3) |
| 589 | policy.on_store_skipped(2) |
| 590 | assert policy.pop_stores_dropped() == (2, 5) |
| 591 | assert policy.pop_stores_dropped() == (0, 0) |
| 592 | |
| 593 | @pytest.mark.parametrize("policy_cls", [DropStorePolicy, ThrottledDropPolicy]) |
| 594 | def test_reset_clears_counters(self, policy_cls): |
| 595 | policy = policy_cls() |
| 596 | policy.on_store_skipped(4) |
| 597 | policy.reset() |
| 598 | assert policy.pop_stores_dropped() == (0, 0) |
| 599 | |
| 600 | def test_throttled_reset_also_clears_call_count(self): |
| 601 | policy = ThrottledDropPolicy() |
| 602 | policy._call_count = 7 |
| 603 | policy.reset() |
| 604 | assert policy._call_count == 0 |
| 605 | |
| 606 | |
| 607 | class TestHealthyBypass: |
| 608 | """Tests for the fast-path bypass when a tier is consistently healthy.""" |
| 609 | |
| 610 | def _make_detector(self, **kwargs): |
| 611 | defaults = dict( |
| 612 | high_water_s=_BP_HIGH_WATER_S, |
| 613 | low_water_s=_BP_LOW_WATER_S, |
| 614 | cooldown_s=0.0, |
| 615 | ) |
| 616 | defaults.update(kwargs) |
| 617 | return EMABackpressureDetector(**defaults) |
| 618 | |
| 619 | def test_not_healthy_initially(self): |
| 620 | bp = self._make_detector() |
| 621 | assert bp.is_healthy() is False |
| 622 | |
| 623 | def test_healthy_after_streak(self): |
| 624 | bp = self._make_detector() |
| 625 | bp._completions = _BP_WARMUP |
| 626 | bp._last_update = time.monotonic() |
| 627 | block_bytes = 16 |
| 628 | fast_s_per_mib = _BP_LOW_WATER_S * 0.5 |
| 629 | fast_elapsed = fast_s_per_mib * (block_bytes / EMABackpressureDetector._MIB) |
| 630 | for _ in range(EMABackpressureDetector._HEALTHY_THRESHOLD + _BP_WARMUP): |
| 631 | bp.on_store_completed(fast_elapsed, block_bytes) |
| 632 | assert bp.is_healthy() is True |
| 633 | |
| 634 | def test_healthy_resets_on_high_latency(self): |
| 635 | bp = self._make_detector() |
| 636 | bp._completions = _BP_WARMUP |
| 637 | bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD |
| 638 | bp._ema = _BP_LOW_WATER_S |
| 639 | bp._last_update = time.monotonic() |
| 640 | block_bytes = 16 |
| 641 | slow_s_per_mib = _BP_HIGH_WATER_S * 2 |
| 642 | slow_elapsed = slow_s_per_mib * (block_bytes / EMABackpressureDetector._MIB) |
| 643 | bp.on_store_completed(slow_elapsed, block_bytes) |
| 644 | assert bp.is_healthy() is False |
| 645 | |
| 646 | def test_healthy_bypasses_policy(self): |
| 647 | bp = self._make_detector() |
| 648 | bp._completions = _BP_WARMUP |
| 649 | bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD |
| 650 | assert bp.should_store(1) is True |
| 651 | |
| 652 | def test_reset_clears_healthy_streak(self): |
| 653 | bp = self._make_detector() |
| 654 | bp._healthy_streak = EMABackpressureDetector._HEALTHY_THRESHOLD |
| 655 | bp.reset() |
| 656 | assert bp._healthy_streak == 0 |
| 657 | assert bp.is_healthy() is False |
| 658 | |
| 659 | |
| 660 | class TestCooldown: |
| 661 | """Tests for the post-pressure cooldown period.""" |
| 662 | |
| 663 | def _make_detector(self, **kwargs): |
| 664 | defaults = dict( |
| 665 | high_water_s=_BP_HIGH_WATER_S, |
| 666 | low_water_s=_BP_LOW_WATER_S, |
| 667 | cooldown_s=1.0, |
| 668 | ) |
| 669 | defaults.update(kwargs) |
| 670 | return EMABackpressureDetector(**defaults) |
| 671 | |
| 672 | def test_cooldown_throttles_stores(self): |
| 673 | bp = self._make_detector() |
| 674 | bp._completions = _BP_WARMUP |
| 675 | bp._under_pressure = True |
| 676 | bp._ema = _BP_HIGH_WATER_S * 2 |
| 677 | bp._last_update = time.monotonic() |
| 678 | |
| 679 | with patch("time.monotonic") as mock_time: |
| 680 | t0 = bp._last_update + 10.0 |
| 681 | mock_time.return_value = t0 |
| 682 | assert bp.is_under_pressure() is False |
| 683 | assert bp._pressure_cleared_at > 0 |
| 684 | |
| 685 | mock_time.return_value = t0 + 0.1 |
| 686 | results = [bp.should_store(1) for _ in range(10)] |
| 687 | assert sum(results) == 5, "cooldown should allow ~50%" |
| 688 | |
| 689 | def test_no_cooldown_when_never_pressured(self): |
| 690 | bp = self._make_detector() |
| 691 | for _ in range(20): |
| 692 | assert bp.should_store(1) is True |
| 693 | |
| 694 | def test_no_cooldown_after_window_expires(self): |
| 695 | bp = self._make_detector() |
| 696 | bp._pressure_cleared_at = time.monotonic() - 2.0 |
| 697 | for _ in range(20): |
| 698 | assert bp.should_store(1) is True |
| 699 | |
| 700 | def test_reset_clears_cooldown(self): |
| 701 | bp = self._make_detector() |
| 702 | bp._pressure_cleared_at = time.monotonic() |
| 703 | bp._cooldown_count = 5 |
| 704 | bp.reset() |
| 705 | assert bp._pressure_cleared_at == 0.0 |
| 706 | assert bp._cooldown_count == 0 |