| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Mock-based unit tests for ObjectStoreSecondaryTierManager. |
| 4 | |
| 5 | These tests replace the NIXL backend with an in-memory mock so they run |
| 6 | without S3 credentials or a live object store. They verify the manager's |
| 7 | state machine: job submission, transfer completion polling, and lookup. |
| 8 | """ |
| 9 | |
| 10 | import time |
| 11 | import uuid |
| 12 | from collections.abc import Callable |
| 13 | from types import SimpleNamespace |
| 14 | from unittest.mock import MagicMock, patch |
| 15 | |
| 16 | import numpy as np |
| 17 | import pytest |
| 18 | import torch |
| 19 | |
| 20 | from vllm.v1.kv_offload.base import ( |
| 21 | Locality, |
| 22 | LookupResult, |
| 23 | Medium, |
| 24 | OffloadingKVEventsConfig, |
| 25 | OffloadKey, |
| 26 | ReqContext, |
| 27 | ScheduleEndContext, |
| 28 | make_offload_key, |
| 29 | ) |
| 30 | from vllm.v1.kv_offload.config import ( |
| 31 | OffloadingCacheConfig, |
| 32 | OffloadingConfig, |
| 33 | OffloadingModelConfig, |
| 34 | OffloadingParallelConfig, |
| 35 | ) |
| 36 | from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob |
| 37 | from vllm.v1.kv_offload.tiering.manager import ( |
| 38 | CPUPrimaryTierOffloadingManager, |
| 39 | TieringOffloadingManager, |
| 40 | ) |
| 41 | from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig |
| 42 | from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager |
| 43 | |
| 44 | # --------------------------------------------------------------------------- |
| 45 | # Shared stubs |
| 46 | # --------------------------------------------------------------------------- |
| 47 | |
| 48 | |
| 49 | def _make_offloading_config( |
| 50 | enable_kv_cache_events: bool, |
| 51 | *, |
| 52 | tp_size: int = 1, |
| 53 | rank: int = 0, |
| 54 | world_size: int | None = None, |
| 55 | replicated_layout: bool = False, |
| 56 | is_parallelism_agnostic: bool = False, |
| 57 | ) -> OffloadingConfig: |
| 58 | if world_size is None: |
| 59 | world_size = tp_size |
| 60 | return OffloadingConfig( |
| 61 | groups=(), |
| 62 | worker_kv_bytes_per_block=0, |
| 63 | enable_kv_cache_events=enable_kv_cache_events, |
| 64 | extra_config={}, |
| 65 | engine_id="test-engine", |
| 66 | model=OffloadingModelConfig(name="test/model", dtype="float16"), |
| 67 | cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), |
| 68 | parallel=OffloadingParallelConfig( |
| 69 | rank=rank, |
| 70 | world_size=world_size, |
| 71 | tp_size=tp_size, |
| 72 | pp_size=1, |
| 73 | pcp_size=1, |
| 74 | dcp_size=1, |
| 75 | data_parallel_index=0, |
| 76 | data_parallel_size=1, |
| 77 | data_parallel_rank_local=None, |
| 78 | is_parallelism_agnostic=is_parallelism_agnostic, |
| 79 | ), |
| 80 | replicated_layout=replicated_layout, |
| 81 | ) |
| 82 | |
| 83 | |
| 84 | _OFFLOADING_SPEC = SimpleNamespace( |
| 85 | config=_make_offloading_config(enable_kv_cache_events=False), |
| 86 | ) |
| 87 | |
| 88 | _STORE_CONFIG = { |
| 89 | "bucket": "mock-bucket", |
| 90 | "endpoint_override": "mock:9000", |
| 91 | "access_key": "mock-access", |
| 92 | "secret_key": "mock-secret", |
| 93 | } |
| 94 | |
| 95 | _BLOCK_ELEMENTS = 256 |
| 96 | _DTYPE = torch.float32 |
| 97 | _RUN_PREFIX = f"test/{uuid.uuid4().hex[:8]}" |
| 98 | _CTX = ReqContext(req_id="test-req") |
| 99 | |
| 100 | |
| 101 | def key(n: int) -> OffloadKey: |
| 102 | return make_offload_key(n.to_bytes(8, "big"), 0) |
| 103 | |
| 104 | |
| 105 | def make_job( |
| 106 | job_id: int, |
| 107 | keys: list[OffloadKey], |
| 108 | chunk_ids: list[int] | None = None, |
| 109 | ) -> TransferJob: |
| 110 | if chunk_ids is None: |
| 111 | chunk_ids = list(range(len(keys))) |
| 112 | return TransferJob( |
| 113 | job_id=job_id, |
| 114 | keys=keys, |
| 115 | chunk_ids=np.array(chunk_ids, dtype=np.int32), |
| 116 | is_promotion=False, |
| 117 | req_context=_CTX, |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | # --------------------------------------------------------------------------- |
| 122 | # Mock NIXL agent |
| 123 | # --------------------------------------------------------------------------- |
| 124 | |
| 125 | |
| 126 | class MockNixlAgent: |
| 127 | """In-memory NIXL agent. Tracks stored object keys and simulates async |
| 128 | transfers: transfer() returns PROC, check_xfer_state() returns DONE and |
| 129 | commits the write to the in-memory key set. |
| 130 | |
| 131 | The four methods overridden by tests (register_memory, make_prepped_xfer, |
| 132 | check_xfer_state, query_memory) are stored as Callable instance attributes |
| 133 | so mypy allows reassignment in tests. |
| 134 | """ |
| 135 | |
| 136 | # Callable attributes — tests may reassign these on instances. |
| 137 | register_memory: Callable |
| 138 | make_prepped_xfer: Callable |
| 139 | check_xfer_state: Callable |
| 140 | query_memory: Callable |
| 141 | |
| 142 | def __init__(self): |
| 143 | self._stored_obj_keys: set[str] = set() |
| 144 | # handle_id -> (op, [obj_keys]) |
| 145 | self._pending: dict[int, tuple[str, list[str]]] = {} |
| 146 | self._handle_counter = 0 |
| 147 | self._last_obj_keys: list[str] = [] |
| 148 | # Bind default implementations as instance attributes. |
| 149 | self.register_memory = self._register_memory |
| 150 | self.make_prepped_xfer = self._make_prepped_xfer |
| 151 | self.check_xfer_state = self._check_xfer_state |
| 152 | self.query_memory = self._query_memory |
| 153 | |
| 154 | def create_backend(self, backend_type, params): |
| 155 | pass |
| 156 | |
| 157 | def _register_memory(self, descs, mem_type=None, backends=None): |
| 158 | mock = MagicMock() |
| 159 | mock.trim.return_value = MagicMock() |
| 160 | # Capture obj_keys from OBJ 4-tuples: (addr, len, dev_id, obj_key) |
| 161 | if mem_type == "OBJ" and descs: |
| 162 | self._last_obj_keys = [d[3] for d in descs if d[3]] |
| 163 | return mock |
| 164 | |
| 165 | def deregister_memory(self, desc): |
| 166 | pass |
| 167 | |
| 168 | def prep_xfer_dlist(self, agent_name, descs, mem_type=None, backends=None): |
| 169 | return MagicMock() |
| 170 | |
| 171 | def _make_prepped_xfer( |
| 172 | self, |
| 173 | op, |
| 174 | local_handle, |
| 175 | local_indices, |
| 176 | remote_handle, |
| 177 | remote_indices, |
| 178 | notif_msg=b"", |
| 179 | backends=None, |
| 180 | skip_desc_merge=False, |
| 181 | ): |
| 182 | handle = MagicMock() |
| 183 | handle._id = self._handle_counter |
| 184 | self._pending[self._handle_counter] = (op, list(self._last_obj_keys)) |
| 185 | self._handle_counter += 1 |
| 186 | return handle |
| 187 | |
| 188 | def transfer(self, handle): |
| 189 | return "PROC" |
| 190 | |
| 191 | def _check_xfer_state(self, handle): |
| 192 | entry = self._pending.pop(handle._id, None) |
| 193 | if entry: |
| 194 | op, obj_keys = entry |
| 195 | if op == "WRITE": |
| 196 | self._stored_obj_keys.update(obj_keys) |
| 197 | return "DONE" |
| 198 | |
| 199 | def release_xfer_handle(self, handle): |
| 200 | pass |
| 201 | |
| 202 | def release_dlist_handle(self, handle): |
| 203 | pass |
| 204 | |
| 205 | def get_xfer_telemetry(self, handle): |
| 206 | return SimpleNamespace(xferDuration=1000) |
| 207 | |
| 208 | def _query_memory(self, queries, mem_type, agent_name): |
| 209 | return [object() if q[3] in self._stored_obj_keys else None for q in queries] |
| 210 | |
| 211 | |
| 212 | # --------------------------------------------------------------------------- |
| 213 | # Fixture |
| 214 | # --------------------------------------------------------------------------- |
| 215 | |
| 216 | |
| 217 | def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace: |
| 218 | """Offloading spec stub with an explicit global KV events flag.""" |
| 219 | return SimpleNamespace( |
| 220 | config=_make_offloading_config(enable_kv_cache_events), |
| 221 | kv_events_config=OffloadingKVEventsConfig( |
| 222 | enable_kv_cache_events=enable_kv_cache_events, |
| 223 | self_describing_kv_events=False, |
| 224 | ), |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | def _make_tier( |
| 229 | num_blocks: int = 4, |
| 230 | offloading_spec: SimpleNamespace = _OFFLOADING_SPEC, |
| 231 | primary_kv_view: memoryview | None = None, |
| 232 | **tier_kwargs, |
| 233 | ) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]: |
| 234 | """Create a tier backed by a fresh MockNixlAgent.""" |
| 235 | mock_agent = MockNixlAgent() |
| 236 | if primary_kv_view is None: |
| 237 | tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) |
| 238 | primary_kv_view = memoryview(tensor.numpy()) |
| 239 | with ( |
| 240 | patch("vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config"), |
| 241 | patch( |
| 242 | "vllm.v1.kv_offload.tiering.obj.manager.nixl_agent", |
| 243 | return_value=mock_agent, |
| 244 | ), |
| 245 | ): |
| 246 | tier = ObjectStoreSecondaryTierManager( |
| 247 | offloading_spec=offloading_spec, |
| 248 | primary_kv_view=primary_kv_view, |
| 249 | tier_type="obj", |
| 250 | store_config=_STORE_CONFIG, |
| 251 | prefix=_RUN_PREFIX, |
| 252 | **tier_kwargs, |
| 253 | ) |
| 254 | return tier, mock_agent |
| 255 | |
| 256 | |
| 257 | def drain( |
| 258 | tier: ObjectStoreSecondaryTierManager, max_rounds: int = 20 |
| 259 | ) -> list[JobResult]: |
| 260 | """Poll get_finished_jobs() until all in-flight jobs resolve.""" |
| 261 | results: list[JobResult] = [] |
| 262 | for _ in range(max_rounds): |
| 263 | results.extend(tier.get_finished_jobs()) |
| 264 | if not tier._transfers: |
| 265 | break |
| 266 | return results |
| 267 | |
| 268 | |
| 269 | def lookup_and_wait( |
| 270 | tier: ObjectStoreSecondaryTierManager, |
| 271 | keys: list[OffloadKey], |
| 272 | ctx: ReqContext = _CTX, |
| 273 | timeout: float = 1.0, |
| 274 | ) -> list[bool]: |
| 275 | """Perform a full async lookup cycle and return resolved results.""" |
| 276 | for k in keys: |
| 277 | tier.lookup(k, ctx) |
| 278 | tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=())) |
| 279 | deadline = time.monotonic() + timeout |
| 280 | while time.monotonic() < deadline: |
| 281 | if not tier._lookup_manager._pending_results.empty(): |
| 282 | break |
| 283 | time.sleep(0.01) |
| 284 | return [tier.lookup(k, ctx) for k in keys] |
| 285 | |
| 286 | |
| 287 | # --------------------------------------------------------------------------- |
| 288 | # Tests |
| 289 | # --------------------------------------------------------------------------- |
| 290 | |
| 291 | |
| 292 | @pytest.mark.parametrize("locality", ["local", ""]) |
| 293 | def test_invalid_locality_raises_at_construction(locality): |
| 294 | with pytest.raises(ValueError, match="Locality"): |
| 295 | _make_tier(locality=locality) |
| 296 | |
| 297 | |
| 298 | class TestMockObjTierBasic: |
| 299 | def setup_method(self): |
| 300 | self.tier, self.agent = _make_tier(num_blocks=4) |
| 301 | |
| 302 | def test_lookup_empty_tier(self): |
| 303 | assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.MISS] |
| 304 | |
| 305 | def test_store_and_lookup(self): |
| 306 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 307 | results = drain(self.tier) |
| 308 | assert len(results) == 1 |
| 309 | assert results[0].success |
| 310 | assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.HIT] |
| 311 | |
| 312 | def test_lookup_unrelated_key_returns_false(self): |
| 313 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 314 | drain(self.tier) |
| 315 | assert lookup_and_wait(self.tier, [key(999)]) == [LookupResult.MISS] |
| 316 | |
| 317 | def test_store_then_load_roundtrip(self): |
| 318 | self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) |
| 319 | results = drain(self.tier) |
| 320 | assert results[0].success |
| 321 | |
| 322 | self.tier.submit_load(make_job(2, [key(1), key(2)], [0, 1])) |
| 323 | results = drain(self.tier) |
| 324 | assert len(results) == 1 |
| 325 | assert results[0].success |
| 326 | |
| 327 | def test_multiple_jobs_tracked_independently(self): |
| 328 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 329 | self.tier.submit_store(make_job(2, [key(2)], [1])) |
| 330 | results = drain(self.tier) |
| 331 | assert len(results) == 2 |
| 332 | assert all(r.success for r in results) |
| 333 | |
| 334 | def test_failed_transfer_reported(self): |
| 335 | self.agent.check_xfer_state = lambda h: "ERR" |
| 336 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 337 | results = drain(self.tier) |
| 338 | assert len(results) == 1 |
| 339 | assert not results[0].success |
| 340 | |
| 341 | def test_failed_load_marks_verdict_negative(self): |
| 342 | """Regression for the failed-load livelock on the obj tier: a |
| 343 | cached HIT must not survive a failed load of the same key. On the |
| 344 | failed promotion the tier marks the verdict False from |
| 345 | get_finished_jobs() (drained here) on the scheduler thread; otherwise |
| 346 | the scheduler would re-issue the same doomed promotion every step for |
| 347 | the life of the request. The mark is served from cache with no |
| 348 | re-probe, so even though the mock object is still 'present' the SAME |
| 349 | request now resolves to MISS.""" |
| 350 | ctx = ReqContext(req_id="obj-livelock") |
| 351 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 352 | assert all(r.success for r in drain(self.tier)) |
| 353 | # Cache a positive verdict: the object is present, so lookup is a HIT. |
| 354 | assert lookup_and_wait(self.tier, [key(1)], ctx=ctx) == [LookupResult.HIT] |
| 355 | |
| 356 | # The promotion the HIT triggered fails. |
| 357 | self.agent.check_xfer_state = lambda h: "ERR" |
| 358 | self.tier.submit_load(make_job(2, [key(1)], [0])) |
| 359 | results = drain(self.tier) |
| 360 | assert len(results) == 1 and not results[0].success |
| 361 | |
| 362 | # After the failed promotion the SAME request's lookup must resolve to |
| 363 | # MISS (verdict marked False) instead of serving the stale HIT — even |
| 364 | # though the object itself is still present in the mock store. |
| 365 | assert lookup_and_wait(self.tier, [key(1)], ctx=ctx) == [LookupResult.MISS] |
| 366 | |
| 367 | def test_pending_transfer_not_returned_until_done(self): |
| 368 | # First poll returns PROC; second poll returns DONE. |
| 369 | call_count = [0] |
| 370 | original = self.agent.check_xfer_state |
| 371 | |
| 372 | def delayed(h): |
| 373 | call_count[0] += 1 |
| 374 | return "PROC" if call_count[0] == 1 else original(h) |
| 375 | |
| 376 | self.agent.check_xfer_state = delayed |
| 377 | |
| 378 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 379 | assert list(self.tier.get_finished_jobs()) == [] |
| 380 | results = list(self.tier.get_finished_jobs()) |
| 381 | assert len(results) == 1 |
| 382 | assert results[0].success |
| 383 | |
| 384 | def test_drain_jobs_polls_until_transfers_complete(self): |
| 385 | """drain_jobs must keep polling check_xfer_state until every |
| 386 | in-flight transfer finishes. A buggy implementation that only |
| 387 | polled once would return with _transfers still populated. |
| 388 | """ |
| 389 | call_count = [0] |
| 390 | original = self.agent.check_xfer_state |
| 391 | |
| 392 | def delayed(h): |
| 393 | call_count[0] += 1 |
| 394 | # Stay in PROC for the first 2 polls, then DONE. |
| 395 | return "PROC" if call_count[0] < 3 else original(h) |
| 396 | |
| 397 | self.agent.check_xfer_state = delayed |
| 398 | |
| 399 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 400 | assert self.tier._transfers # in flight |
| 401 | |
| 402 | self.tier.drain_jobs() |
| 403 | |
| 404 | assert not self.tier._transfers # fully drained |
| 405 | assert call_count[0] >= 3 # polled past the initial PROC responses |
| 406 | # Result is buffered for the next get_finished_jobs() call. |
| 407 | results = list(self.tier.get_finished_jobs()) |
| 408 | assert len(results) == 1 |
| 409 | assert results[0].success |
| 410 | |
| 411 | |
| 412 | class TestMockObjTierMultiBlock: |
| 413 | def test_store_multiple_blocks(self): |
| 414 | tier, _ = _make_tier(num_blocks=8) |
| 415 | keys = [key(i) for i in range(8)] |
| 416 | tier.submit_store(make_job(1, keys, list(range(8)))) |
| 417 | results = drain(tier) |
| 418 | assert len(results) == 1 |
| 419 | assert results[0].success |
| 420 | assert lookup_and_wait(tier, keys) == [LookupResult.HIT] * 8 |
| 421 | |
| 422 | def test_partial_block_lookup(self): |
| 423 | tier, _ = _make_tier(num_blocks=4) |
| 424 | tier.submit_store(make_job(1, [key(0), key(1)], [0, 1])) |
| 425 | drain(tier) |
| 426 | assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [ |
| 427 | LookupResult.HIT, |
| 428 | LookupResult.HIT, |
| 429 | LookupResult.MISS, |
| 430 | ] |
| 431 | |
| 432 | |
| 433 | class TestMockObjTierFailures: |
| 434 | def test_lookup_exception_returns_false(self): |
| 435 | tier, agent = _make_tier(num_blocks=4) |
| 436 | agent.query_memory = lambda *a, **k: (_ for _ in ()).throw( |
| 437 | RuntimeError("backend error") |
| 438 | ) |
| 439 | assert lookup_and_wait(tier, [key(1)]) == [LookupResult.MISS] |
| 440 | |
| 441 | def test_submit_store_register_memory_failure_reported_in_get_finished(self): |
| 442 | tier, agent = _make_tier(num_blocks=4) |
| 443 | agent.register_memory = lambda *a, **k: None |
| 444 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 445 | results = list(tier.get_finished_jobs()) |
| 446 | assert len(results) == 1 |
| 447 | assert results[0].job_id == 1 |
| 448 | assert not results[0].success |
| 449 | |
| 450 | def test_submit_load_register_memory_failure_reported_in_get_finished(self): |
| 451 | tier, agent = _make_tier(num_blocks=4) |
| 452 | agent.register_memory = lambda *a, **k: None |
| 453 | tier.submit_load(make_job(2, [key(1)], [0])) |
| 454 | results = list(tier.get_finished_jobs()) |
| 455 | assert len(results) == 1 |
| 456 | assert results[0].job_id == 2 |
| 457 | assert not results[0].success |
| 458 | |
| 459 | def test_submit_store_make_prepped_xfer_failure_reported_in_get_finished(self): |
| 460 | tier, agent = _make_tier(num_blocks=4) |
| 461 | agent.make_prepped_xfer = lambda *a, **k: None |
| 462 | tier.submit_store(make_job(3, [key(1)], [0])) |
| 463 | results = list(tier.get_finished_jobs()) |
| 464 | assert len(results) == 1 |
| 465 | assert results[0].job_id == 3 |
| 466 | assert not results[0].success |
| 467 | |
| 468 | def test_failure_and_success_both_returned_by_get_finished(self): |
| 469 | # One job fails at submission, another succeeds in flight. |
| 470 | tier, agent = _make_tier(num_blocks=4) |
| 471 | original_register = agent.register_memory |
| 472 | call_count = [0] |
| 473 | |
| 474 | def register_once_fail(*a, **k): |
| 475 | call_count[0] += 1 |
| 476 | return None if call_count[0] == 1 else original_register(*a, **k) |
| 477 | |
| 478 | agent.register_memory = register_once_fail |
| 479 | |
| 480 | tier.submit_store(make_job(1, [key(1)], [0])) # fails immediately |
| 481 | tier.submit_store(make_job(2, [key(2)], [1])) # succeeds |
| 482 | results = drain(tier) |
| 483 | assert len(results) == 2 |
| 484 | by_id = {r.job_id: r for r in results} |
| 485 | assert not by_id[1].success |
| 486 | assert by_id[2].success |
| 487 | |
| 488 | def test_release_xfer_failure_retries_without_losing_result(self, monkeypatch): |
| 489 | tier, agent = _make_tier(num_blocks=4) |
| 490 | agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) |
| 491 | release_xfer = MagicMock( |
| 492 | side_effect=[RuntimeError("transfer is still active"), None] |
| 493 | ) |
| 494 | monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) |
| 495 | |
| 496 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 497 | |
| 498 | # The transfer handle could not be released safely, so the job must |
| 499 | # remain tracked and must not be finalized yet. |
| 500 | assert list(tier.get_finished_jobs()) == [] |
| 501 | assert 1 in tier._transfers |
| 502 | |
| 503 | # Cleanup is retried without polling again or changing the failure |
| 504 | # verdict. The completion is then returned exactly once. |
| 505 | results = list(tier.get_finished_jobs()) |
| 506 | assert len(results) == 1 |
| 507 | assert results[0].job_id == 1 |
| 508 | assert not results[0].success |
| 509 | assert agent.check_xfer_state.call_count == 2 |
| 510 | assert release_xfer.call_count == 2 |
| 511 | assert not tier._transfers |
| 512 | assert list(tier.get_finished_jobs()) == [] |
| 513 | |
| 514 | @pytest.mark.parametrize( |
| 515 | "cleanup_method", ["release_dlist_handle", "deregister_memory"] |
| 516 | ) |
| 517 | def test_post_transfer_cleanup_failure_does_not_lose_result( |
| 518 | self, monkeypatch, cleanup_method |
| 519 | ): |
| 520 | tier, agent = _make_tier(num_blocks=4) |
| 521 | monkeypatch.setattr( |
| 522 | agent, |
| 523 | cleanup_method, |
| 524 | MagicMock(side_effect=RuntimeError("cleanup failed")), |
| 525 | ) |
| 526 | |
| 527 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 528 | results = list(tier.get_finished_jobs()) |
| 529 | |
| 530 | assert len(results) == 1 |
| 531 | assert results[0].job_id == 1 |
| 532 | assert results[0].success |
| 533 | assert not tier._transfers |
| 534 | assert list(tier.get_finished_jobs()) == [] |
| 535 | |
| 536 | def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypatch): |
| 537 | num_blocks = 4 |
| 538 | tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE) |
| 539 | primary_kv_view = memoryview(tensor.numpy()) |
| 540 | mmap_region = MagicMock() |
| 541 | mmap_region.create_kv_memoryview.return_value = primary_kv_view |
| 542 | primary_tier = CPUPrimaryTierOffloadingManager( |
| 543 | num_chunks=num_blocks, mmap_region=mmap_region |
| 544 | ) |
| 545 | obj_tier, agent = _make_tier( |
| 546 | num_blocks=num_blocks, primary_kv_view=primary_kv_view |
| 547 | ) |
| 548 | manager = TieringOffloadingManager( |
| 549 | primary_tier=primary_tier, secondary_tiers=[obj_tier] |
| 550 | ) |
| 551 | |
| 552 | keys = [key(1)] |
| 553 | primary_result = primary_tier.prepare_store(keys, _CTX) |
| 554 | assert primary_result is not None |
| 555 | primary_tier.complete_store(keys, _CTX, success=True) |
| 556 | job = manager.create_store_job(keys, _CTX) |
| 557 | obj_tier.submit_store(job) |
| 558 | |
| 559 | block = primary_tier._policy.get(keys[0]) |
| 560 | assert block is not None |
| 561 | assert block.ref_cnt == 1 |
| 562 | assert len(manager._transfer_jobs) == 1 |
| 563 | |
| 564 | agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed")) |
| 565 | release_xfer = MagicMock( |
| 566 | side_effect=[RuntimeError("transfer is still active"), None] |
| 567 | ) |
| 568 | monkeypatch.setattr(agent, "release_xfer_handle", release_xfer) |
| 569 | schedule_context = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) |
| 570 | |
| 571 | manager.on_schedule_end(schedule_context) |
| 572 | |
| 573 | assert len(manager._transfer_jobs) == 1 |
| 574 | assert block.ref_cnt == 1 |
| 575 | assert len(obj_tier._transfers) == 1 |
| 576 | assert manager.has_pending_work() |
| 577 | |
| 578 | manager.on_schedule_end(schedule_context) |
| 579 | |
| 580 | assert manager._transfer_jobs == {} |
| 581 | assert block.ref_cnt == 0 |
| 582 | assert obj_tier._transfers == {} |
| 583 | assert not manager.has_pending_work() |
| 584 | assert agent.check_xfer_state.call_count == 2 |
| 585 | assert release_xfer.call_count == 2 |
| 586 | |
| 587 | |
| 588 | class TestMockObjTierShutdown: |
| 589 | def test_shutdown_clears_in_flight_transfers(self): |
| 590 | tier, agent = _make_tier(num_blocks=4) |
| 591 | # Keep transfer in flight by never completing it |
| 592 | agent.check_xfer_state = lambda h: "PROC" |
| 593 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 594 | assert len(tier._transfers) == 1 |
| 595 | tier.shutdown() |
| 596 | assert len(tier._transfers) == 0 |
| 597 | assert tier._dram_prepped_handle is None |
| 598 | assert tier._primary_reg is None |
| 599 | |
| 600 | def test_shutdown_idempotent(self): |
| 601 | tier, _ = _make_tier(num_blocks=4) |
| 602 | tier.shutdown() |
| 603 | tier.shutdown() # must not raise |
| 604 | |
| 605 | |
| 606 | class TestObjTierKVEvents: |
| 607 | def setup_method(self): |
| 608 | self.tier, self.agent = _make_tier( |
| 609 | offloading_spec=_make_events_spec(enable_kv_cache_events=True), |
| 610 | enable_kv_events=True, |
| 611 | locality="REMOTE", |
| 612 | ) |
| 613 | |
| 614 | def test_successful_store_emits_stored_event(self): |
| 615 | """A completed store transfer emits one stored event with the job's keys.""" |
| 616 | keys = [key(1), key(2)] |
| 617 | self.tier.submit_store(make_job(1, keys, [0, 1])) |
| 618 | assert all(r.success for r in drain(self.tier)) |
| 619 | |
| 620 | events = list(self.tier.take_events()) |
| 621 | assert len(events) == 1 |
| 622 | assert events[0].keys == keys |
| 623 | assert events[0].medium == Medium.STORAGE |
| 624 | assert events[0].locality is Locality.REMOTE |
| 625 | assert not events[0].removed |
| 626 | # take_events drains the buffer. |
| 627 | assert list(self.tier.take_events()) == [] |
| 628 | |
| 629 | @pytest.mark.parametrize( |
| 630 | ("locality", "expected"), |
| 631 | [(None, None), ("LOCAL", Locality.LOCAL)], |
| 632 | ) |
| 633 | def test_store_event_uses_configured_locality(self, locality, expected): |
| 634 | locality_config = {} if locality is None else {"locality": locality} |
| 635 | tier, _ = _make_tier( |
| 636 | offloading_spec=_make_events_spec(enable_kv_cache_events=True), |
| 637 | enable_kv_events=True, |
| 638 | **locality_config, |
| 639 | ) |
| 640 | try: |
| 641 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 642 | assert all(r.success for r in drain(tier)) |
| 643 | |
| 644 | events = list(tier.take_events()) |
| 645 | assert len(events) == 1 |
| 646 | assert events[0].locality is expected |
| 647 | finally: |
| 648 | tier.shutdown() |
| 649 | |
| 650 | def test_mixed_job_results_emit_event_only_for_successful_job(self): |
| 651 | """With a failed and a successful store job resolving in the same |
| 652 | poll, exactly one event is emitted and its keys belong to the |
| 653 | successful job.""" |
| 654 | original = self.agent.check_xfer_state |
| 655 | self.agent.check_xfer_state = lambda h: "ERR" if h._id == 0 else original(h) |
| 656 | self.tier.submit_store(make_job(1, [key(1)], [0])) # handle 0: fails |
| 657 | self.tier.submit_store(make_job(2, [key(2)], [1])) # handle 1: succeeds |
| 658 | results = drain(self.tier) |
| 659 | by_id = {r.job_id: r for r in results} |
| 660 | assert not by_id[1].success |
| 661 | assert by_id[2].success |
| 662 | |
| 663 | events = list(self.tier.take_events()) |
| 664 | assert len(events) == 1 |
| 665 | assert events[0].keys == [key(2)] |
| 666 | assert self.tier._store_job_keys == {} |
| 667 | |
| 668 | def test_load_job_emits_no_event(self): |
| 669 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 670 | results = drain(self.tier) |
| 671 | assert len(results) == 1 |
| 672 | assert results[0].success |
| 673 | list(self.tier.take_events()) |
| 674 | |
| 675 | self.tier.submit_load(make_job(2, [key(1)], [0])) |
| 676 | results = drain(self.tier) |
| 677 | assert len(results) == 1 |
| 678 | assert results[0].success |
| 679 | assert list(self.tier.take_events()) == [] |
| 680 | |
| 681 | def test_failed_transfer_emits_no_event(self): |
| 682 | self.agent.check_xfer_state = lambda h: "ERR" |
| 683 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 684 | results = drain(self.tier) |
| 685 | assert not results[0].success |
| 686 | assert list(self.tier.take_events()) == [] |
| 687 | assert self.tier._store_job_keys == {} |
| 688 | |
| 689 | def test_submission_failure_emits_no_event(self): |
| 690 | self.agent.make_prepped_xfer = lambda *a, **k: None |
| 691 | self.tier.submit_store(make_job(1, [key(1)], [0])) |
| 692 | results = list(self.tier.get_finished_jobs()) |
| 693 | assert not results[0].success |
| 694 | assert list(self.tier.take_events()) == [] |
| 695 | assert self.tier._store_job_keys == {} |
| 696 | |
| 697 | def test_events_disabled_by_default(self): |
| 698 | tier, _ = _make_tier() |
| 699 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 700 | results = drain(tier) |
| 701 | assert len(results) == 1 |
| 702 | assert results[0].success |
| 703 | assert tier.events is None |
| 704 | assert tier._store_job_keys == {} |
| 705 | assert list(tier.take_events()) == [] |
| 706 | |
| 707 | def test_events_require_global_kv_events_flag(self): |
| 708 | """Tier-level opt-in alone is not enough; the global flag gates events.""" |
| 709 | tier, _ = _make_tier( |
| 710 | offloading_spec=_make_events_spec(enable_kv_cache_events=False), |
| 711 | enable_kv_events=True, |
| 712 | ) |
| 713 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 714 | results = drain(tier) |
| 715 | assert len(results) == 1 |
| 716 | assert results[0].success |
| 717 | assert tier.events is None |
| 718 | assert tier._store_job_keys == {} |
| 719 | assert list(tier.take_events()) == [] |
| 720 | |
| 721 | |
| 722 | class TestObjStoreConfig: |
| 723 | def test_explicit_credentials_included(self): |
| 724 | cfg = ObjStoreConfig( |
| 725 | bucket="b", |
| 726 | endpoint_override="ep", |
| 727 | access_key="ak", |
| 728 | secret_key="sk", |
| 729 | ) |
| 730 | params = cfg.to_nixl_params() |
| 731 | assert params["access_key"] == "ak" |
| 732 | assert params["secret_key"] == "sk" |
| 733 | |
| 734 | def test_credentials_omitted_when_empty(self): |
| 735 | cfg = ObjStoreConfig(bucket="b", endpoint_override="ep") |
| 736 | params = cfg.to_nixl_params() |
| 737 | assert "access_key" not in params |
| 738 | assert "secret_key" not in params |
| 739 | assert "session_token" not in params |
| 740 | assert "region" not in params |
| 741 | assert params["bucket"] == "b" |
| 742 | assert params["endpoint_override"] == "ep" |
| 743 | |
| 744 | def test_session_token_and_region_included(self): |
| 745 | cfg = ObjStoreConfig( |
| 746 | bucket="b", |
| 747 | endpoint_override="ep", |
| 748 | access_key="ak", |
| 749 | secret_key="sk", |
| 750 | session_token="tok", |
| 751 | region="us-east-1", |
| 752 | ) |
| 753 | params = cfg.to_nixl_params() |
| 754 | assert params["session_token"] == "tok" |
| 755 | assert params["region"] == "us-east-1" |
| 756 | |
| 757 | def test_ca_bundle_included_when_set(self): |
| 758 | cfg = ObjStoreConfig( |
| 759 | bucket="b", |
| 760 | endpoint_override="ep", |
| 761 | ca_bundle="/path/to/ca.pem", |
| 762 | ) |
| 763 | params = cfg.to_nixl_params() |
| 764 | assert params["ca_bundle"] == "/path/to/ca.pem" |
| 765 | assert "access_key" not in params |
| 766 | |
| 767 | |
| 768 | def test_obj_tier_replicated_layout_collapses_mapper_identity(): |
| 769 | """TP=2 and TP=4 replicated configs share the obj FileMapper namespace.""" |
| 770 | tp2_spec = SimpleNamespace( |
| 771 | config=_make_offloading_config( |
| 772 | False, tp_size=2, world_size=2, rank=1, replicated_layout=True |
| 773 | ), |
| 774 | kv_events_config=OffloadingKVEventsConfig( |
| 775 | enable_kv_cache_events=False, |
| 776 | self_describing_kv_events=False, |
| 777 | ), |
| 778 | ) |
| 779 | tp4_spec = SimpleNamespace( |
| 780 | config=_make_offloading_config( |
| 781 | False, tp_size=4, world_size=4, rank=3, replicated_layout=True |
| 782 | ), |
| 783 | kv_events_config=OffloadingKVEventsConfig( |
| 784 | enable_kv_cache_events=False, |
| 785 | self_describing_kv_events=False, |
| 786 | ), |
| 787 | ) |
| 788 | tp2_tier, _ = _make_tier(offloading_spec=tp2_spec) |
| 789 | tp4_tier, _ = _make_tier(offloading_spec=tp4_spec) |
| 790 | try: |
| 791 | assert tp2_tier._file_mapper.base_path == tp4_tier._file_mapper.base_path |
| 792 | assert tp2_tier._file_mapper.rank == 0 |
| 793 | assert tp4_tier._file_mapper.rank == 0 |
| 794 | assert tp2_tier._file_mapper.get_run_config() == ( |
| 795 | tp4_tier._file_mapper.get_run_config() |
| 796 | ) |
| 797 | finally: |
| 798 | tp2_tier.shutdown() |
| 799 | tp4_tier.shutdown() |