| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for FileSystemTierManager. |
| 4 | |
| 5 | These tests use real disk I/O to verify the filesystem tier implementation. |
| 6 | The tier manager writes KV cache blocks to disk and reads them back, verifying |
| 7 | data integrity throughout the process. |
| 8 | """ |
| 9 | |
| 10 | import mmap |
| 11 | import os |
| 12 | import threading |
| 13 | import time |
| 14 | from unittest.mock import MagicMock |
| 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 | OffloadingEvent, |
| 25 | OffloadingKVEventsConfig, |
| 26 | OffloadKey, |
| 27 | ReqContext, |
| 28 | ScheduleEndContext, |
| 29 | make_offload_key, |
| 30 | ) |
| 31 | from vllm.v1.kv_offload.config import ( |
| 32 | OffloadingCacheConfig, |
| 33 | OffloadingConfig, |
| 34 | OffloadingModelConfig, |
| 35 | OffloadingParallelConfig, |
| 36 | ) |
| 37 | from vllm.v1.kv_offload.tiering.base import TransferJob |
| 38 | from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory |
| 39 | from vllm.v1.kv_offload.tiering.fs.manager import ( |
| 40 | FileSystemTierManager, |
| 41 | ) |
| 42 | from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool |
| 43 | |
| 44 | # --------------------------------------------------------------------------- |
| 45 | # Helpers |
| 46 | # --------------------------------------------------------------------------- |
| 47 | |
| 48 | _NUM_BLOCKS = 8 |
| 49 | _BLOCK_ELEMENTS = 128 * mmap.PAGESIZE # 2MB per block for pagesize 4096. |
| 50 | _DTYPE: torch.dtype = torch.float32 |
| 51 | _CTX = ReqContext(req_id="test") |
| 52 | |
| 53 | |
| 54 | def _make_offloading_spec( |
| 55 | enable_kv_cache_events: bool = False, |
| 56 | *, |
| 57 | tp_size: int = 1, |
| 58 | rank: int = 0, |
| 59 | world_size: int | None = None, |
| 60 | replicated_layout: bool = False, |
| 61 | is_parallelism_agnostic: bool = False, |
| 62 | ) -> MagicMock: |
| 63 | """Mock spec with an explicit global KV events flag.""" |
| 64 | if world_size is None: |
| 65 | world_size = tp_size |
| 66 | spec = MagicMock() |
| 67 | spec.config = OffloadingConfig( |
| 68 | groups=(), |
| 69 | worker_kv_bytes_per_block=0, |
| 70 | enable_kv_cache_events=enable_kv_cache_events, |
| 71 | extra_config={}, |
| 72 | engine_id="test-engine", |
| 73 | model=OffloadingModelConfig(name="test-model", dtype="float32"), |
| 74 | cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), |
| 75 | parallel=OffloadingParallelConfig( |
| 76 | rank=rank, |
| 77 | world_size=world_size, |
| 78 | tp_size=tp_size, |
| 79 | pp_size=1, |
| 80 | pcp_size=1, |
| 81 | dcp_size=1, |
| 82 | data_parallel_index=0, |
| 83 | data_parallel_size=1, |
| 84 | data_parallel_rank_local=None, |
| 85 | is_parallelism_agnostic=is_parallelism_agnostic, |
| 86 | ), |
| 87 | replicated_layout=replicated_layout, |
| 88 | ) |
| 89 | spec.blocks_per_chunk = 1 |
| 90 | spec.kv_events_config = OffloadingKVEventsConfig( |
| 91 | enable_kv_cache_events=enable_kv_cache_events, |
| 92 | self_describing_kv_events=False, |
| 93 | ) |
| 94 | return spec |
| 95 | |
| 96 | |
| 97 | _MOCK_OFFLOADING_SPEC = _make_offloading_spec(enable_kv_cache_events=False) |
| 98 | |
| 99 | |
| 100 | def key(n: int) -> OffloadKey: |
| 101 | return make_offload_key(n.to_bytes(8, "big"), 0) |
| 102 | |
| 103 | |
| 104 | def make_job( |
| 105 | job_id: int, |
| 106 | keys: list[OffloadKey], |
| 107 | chunk_ids: list[int] | None = None, |
| 108 | is_promotion: bool = False, |
| 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.int64), |
| 116 | is_promotion=is_promotion, |
| 117 | req_context=_CTX, |
| 118 | ) |
| 119 | |
| 120 | |
| 121 | def drain(tier: FileSystemTierManager) -> list: |
| 122 | """Block until all in-flight jobs finish, then collect results.""" |
| 123 | tier.drain_jobs() |
| 124 | return list(tier.get_finished_jobs()) |
| 125 | |
| 126 | |
| 127 | def lookup_and_wait( |
| 128 | tier: FileSystemTierManager, |
| 129 | keys: list[OffloadKey], |
| 130 | ctx: ReqContext = _CTX, |
| 131 | timeout: float = 1.0, |
| 132 | ) -> list[LookupResult]: |
| 133 | """Perform a full async lookup cycle and return resolved results.""" |
| 134 | for k in keys: |
| 135 | tier.lookup(k, ctx) |
| 136 | tier.on_schedule_end(ScheduleEndContext(new_req_ids=[], preempted_req_ids=())) |
| 137 | deadline = time.monotonic() + timeout |
| 138 | while time.monotonic() < deadline: |
| 139 | if not tier._lookup_manager._pending_results.empty(): |
| 140 | break |
| 141 | time.sleep(0.01) |
| 142 | return [tier.lookup(k, ctx) for k in keys] |
| 143 | |
| 144 | |
| 145 | def _page_aligned_zero_tensor( |
| 146 | num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE |
| 147 | ) -> torch.Tensor: |
| 148 | page_size = mmap.PAGESIZE |
| 149 | dtype_num_bytes = torch.tensor([], dtype=dtype).element_size() |
| 150 | |
| 151 | num_bytes = num_blocks * block_elements * dtype_num_bytes |
| 152 | num_bytes_aligned = num_bytes + page_size |
| 153 | t = torch.zeros(num_bytes_aligned, dtype=torch.uint8) |
| 154 | |
| 155 | ptr = t.data_ptr() |
| 156 | alignment_offset = ptr % page_size |
| 157 | # Move tensor to next page regardless. |
| 158 | shift = page_size - alignment_offset |
| 159 | t = t[shift : shift + num_bytes] |
| 160 | return t.view(dtype).view(num_blocks, block_elements) |
| 161 | |
| 162 | |
| 163 | def _page_aligned_rand_tensor( |
| 164 | num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE |
| 165 | ) -> torch.Tensor: |
| 166 | rand_tensor = _page_aligned_zero_tensor(num_blocks, block_elements) |
| 167 | rand_tensor[:] = torch.rand(num_blocks, block_elements, dtype=dtype) |
| 168 | return rand_tensor |
| 169 | |
| 170 | |
| 171 | # --------------------------------------------------------------------------- |
| 172 | # Fixtures |
| 173 | # --------------------------------------------------------------------------- |
| 174 | |
| 175 | |
| 176 | @pytest.fixture |
| 177 | def fs_tier(tmp_path): |
| 178 | tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) |
| 179 | mock_view = memoryview(tensor.numpy()) |
| 180 | tier = FileSystemTierManager( |
| 181 | offloading_spec=_MOCK_OFFLOADING_SPEC, |
| 182 | primary_kv_view=mock_view, |
| 183 | tier_type="fs", |
| 184 | root_dir=str(tmp_path), |
| 185 | n_read_threads=4, |
| 186 | n_write_threads=4, |
| 187 | ) |
| 188 | yield tier, tensor |
| 189 | tier.shutdown() |
| 190 | |
| 191 | |
| 192 | @pytest.fixture |
| 193 | def fs_tier_with_events(tmp_path): |
| 194 | tensor = _page_aligned_zero_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) |
| 195 | mock_view = memoryview(tensor.numpy()) |
| 196 | tier = FileSystemTierManager( |
| 197 | offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), |
| 198 | primary_kv_view=mock_view, |
| 199 | tier_type="fs", |
| 200 | root_dir=str(tmp_path), |
| 201 | n_read_threads=4, |
| 202 | n_write_threads=4, |
| 203 | enable_kv_events=True, |
| 204 | locality="LOCAL", |
| 205 | ) |
| 206 | yield tier |
| 207 | tier.shutdown() |
| 208 | |
| 209 | |
| 210 | # --------------------------------------------------------------------------- |
| 211 | # Tests |
| 212 | # --------------------------------------------------------------------------- |
| 213 | |
| 214 | |
| 215 | def test_lookup_empty_tier(fs_tier): |
| 216 | tier, _ = fs_tier |
| 217 | results = lookup_and_wait(tier, [key(1), key(2)]) |
| 218 | assert results == [LookupResult.MISS, LookupResult.MISS] |
| 219 | |
| 220 | |
| 221 | def test_store_creates_file_and_lookup_succeeds(fs_tier): |
| 222 | tier, _ = fs_tier |
| 223 | job = make_job(1, [key(1)], [0]) |
| 224 | tier.submit_store(job) |
| 225 | results = drain(tier) |
| 226 | assert len(results) == 1 |
| 227 | assert results[0].success |
| 228 | assert lookup_and_wait(tier, [key(1)]) == [LookupResult.HIT] |
| 229 | dest = tier.file_mapper.get_file_name(key(1)) |
| 230 | assert os.path.exists(dest), f"Expected file at {dest}" |
| 231 | |
| 232 | |
| 233 | def test_store_then_load_roundtrip(fs_tier): |
| 234 | tier, _ = fs_tier |
| 235 | job_s = make_job(1, [key(1), key(2)], [0, 1]) |
| 236 | tier.submit_store(job_s) |
| 237 | store_results = drain(tier) |
| 238 | assert all(r.success for r in store_results) |
| 239 | |
| 240 | assert lookup_and_wait(tier, [key(1), key(2)]) == [ |
| 241 | LookupResult.HIT, |
| 242 | LookupResult.HIT, |
| 243 | ] |
| 244 | |
| 245 | job_l = make_job(2, [key(1), key(2)], [2, 3], is_promotion=True) |
| 246 | tier.submit_load(job_l) |
| 247 | load_results = drain(tier) |
| 248 | assert all(r.success for r in load_results) |
| 249 | # A successful load must NOT touch the file: the delete path fires only on |
| 250 | # a provable short read, so a good block stays on disk (guards against an |
| 251 | # over-eager delete regressing to upstream's delete-on-any-error). |
| 252 | for k in (key(1), key(2)): |
| 253 | assert os.path.exists(tier.file_mapper.get_file_name(k)) |
| 254 | # Blocks stay on disk after load |
| 255 | assert lookup_and_wait(tier, [key(1), key(2)]) == [ |
| 256 | LookupResult.HIT, |
| 257 | LookupResult.HIT, |
| 258 | ] |
| 259 | |
| 260 | |
| 261 | def test_invalid_path_raises_at_construction(): |
| 262 | """Construction must fail immediately when the config file cannot be written.""" |
| 263 | tensor = _page_aligned_zero_tensor(32, _BLOCK_ELEMENTS) |
| 264 | mock_view = memoryview(tensor.numpy()) |
| 265 | |
| 266 | with pytest.raises(OSError): |
| 267 | FileSystemTierManager( |
| 268 | offloading_spec=_MOCK_OFFLOADING_SPEC, |
| 269 | primary_kv_view=mock_view, |
| 270 | tier_type="fs", |
| 271 | root_dir="/dev/null/invalid_path", |
| 272 | ) |
| 273 | |
| 274 | |
| 275 | @pytest.mark.parametrize("locality", ["local", ""]) |
| 276 | def test_invalid_locality_raises_at_construction(tmp_path, locality): |
| 277 | tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) |
| 278 | |
| 279 | with pytest.raises(ValueError, match="Locality"): |
| 280 | FileSystemTierManager( |
| 281 | offloading_spec=_MOCK_OFFLOADING_SPEC, |
| 282 | primary_kv_view=memoryview(tensor.numpy()), |
| 283 | tier_type="fs", |
| 284 | root_dir=str(tmp_path), |
| 285 | locality=locality, |
| 286 | ) |
| 287 | |
| 288 | |
| 289 | def test_factory_forwards_locality_to_fs_tier(tmp_path): |
| 290 | tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) |
| 291 | tier = SecondaryTierFactory.create_secondary_tier( |
| 292 | { |
| 293 | "type": "fs", |
| 294 | "root_dir": str(tmp_path), |
| 295 | "n_read_threads": 1, |
| 296 | "n_write_threads": 1, |
| 297 | "locality": "LOCAL", |
| 298 | }, |
| 299 | memoryview(tensor.numpy()), |
| 300 | _MOCK_OFFLOADING_SPEC, |
| 301 | ) |
| 302 | try: |
| 303 | assert isinstance(tier, FileSystemTierManager) |
| 304 | assert tier.locality is Locality.LOCAL |
| 305 | finally: |
| 306 | tier.shutdown() |
| 307 | |
| 308 | |
| 309 | def test_failed_load_missing_file(fs_tier): |
| 310 | """Test that loading a block whose file does not exist results in a failed job.""" |
| 311 | tier, _ = fs_tier |
| 312 | job = make_job(1, [key(99)], [0], is_promotion=True) |
| 313 | tier.submit_load(job) |
| 314 | results = drain(tier) |
| 315 | assert len(results) == 1 |
| 316 | assert not results[0].success |
| 317 | |
| 318 | |
| 319 | def test_multiple_jobs_tracked_independently(fs_tier): |
| 320 | tier, _ = fs_tier |
| 321 | job1 = make_job(1, [key(1)], [0]) |
| 322 | job2 = make_job(2, [key(2)], [1]) |
| 323 | tier.submit_store(job1) |
| 324 | tier.submit_store(job2) |
| 325 | results = drain(tier) |
| 326 | job_ids = {r.job_id for r in results} |
| 327 | assert job_ids == {1, 2} |
| 328 | assert lookup_and_wait(tier, [key(1), key(2)]) == [ |
| 329 | LookupResult.HIT, |
| 330 | LookupResult.HIT, |
| 331 | ] |
| 332 | |
| 333 | |
| 334 | def test_multi_block_job_partial_failure(fs_tier): |
| 335 | """A load job where one block file is missing yields a single failed JobResult.""" |
| 336 | tier, _ = fs_tier |
| 337 | # Store two of three keys |
| 338 | tier.submit_store(make_job(1, [key(10), key(11)], [0, 1])) |
| 339 | assert all(r.success for r in drain(tier)) |
| 340 | |
| 341 | # Load all three — key(99) was never stored |
| 342 | tier.submit_load( |
| 343 | make_job(2, [key(10), key(11), key(99)], [0, 1, 2], is_promotion=True) |
| 344 | ) |
| 345 | results = drain(tier) |
| 346 | |
| 347 | assert len(results) == 1 |
| 348 | assert results[0].job_id == 2 |
| 349 | assert not results[0].success |
| 350 | |
| 351 | |
| 352 | def test_shutdown_discards_pending_tasks(fs_tier): |
| 353 | """Shutdown clears both queues and stops all worker threads without draining.""" |
| 354 | tier, _ = fs_tier |
| 355 | # Submit many tasks to ensure some remain pending |
| 356 | for i in range(10): |
| 357 | tier.submit_store(make_job(i, [key(i)], [i % 4])) |
| 358 | |
| 359 | # Shutdown immediately without draining |
| 360 | tier.shutdown() |
| 361 | |
| 362 | # Verify queues are cleared and threads stopped |
| 363 | assert len(tier._pool._load_q) == 0 |
| 364 | assert len(tier._pool._store_q) == 0 |
| 365 | assert all(not t.is_alive() for t in tier._pool._threads) |
| 366 | |
| 367 | |
| 368 | @pytest.mark.parametrize("batch_size", [0, 1, 2, 5]) |
| 369 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 370 | def test_store_load_data_integrity(fs_tier, monkeypatch, use_c_ext, batch_size): |
| 371 | """Data written by store must be exactly recovered by load, for batches |
| 372 | of any size -- including the empty batch.""" |
| 373 | import vllm.v1.kv_offload.tiering.fs.io as io_mod |
| 374 | |
| 375 | if use_c_ext and not io_mod._HAS_FSIO_C: |
| 376 | pytest.skip("fs_io_C extension not built") |
| 377 | monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) |
| 378 | |
| 379 | tier, tensor = fs_tier |
| 380 | # Populate tensor with random data |
| 381 | tensor[:] = _page_aligned_rand_tensor(_NUM_BLOCKS, _BLOCK_ELEMENTS) |
| 382 | |
| 383 | keys = [key(i) for i in range(batch_size)] |
| 384 | store_chunk_ids = list(range(batch_size)) |
| 385 | load_chunk_ids = list(range(_NUM_BLOCKS - batch_size, _NUM_BLOCKS)) |
| 386 | expected = tensor[:batch_size].clone() |
| 387 | |
| 388 | tier.submit_store(make_job(1, keys, store_chunk_ids)) |
| 389 | store_results = drain(tier) |
| 390 | assert len(store_results) == 1 |
| 391 | assert store_results[0].success |
| 392 | assert all(os.path.exists(tier.file_mapper.get_file_name(k)) for k in keys) |
| 393 | |
| 394 | # reset tensor to prove data is read from disk |
| 395 | tensor[:] = 0.0 |
| 396 | |
| 397 | # Load into a range disjoint by index from the store ids, to also |
| 398 | # exercise loading a chunk into a different id than it was stored from. |
| 399 | tier.submit_load(make_job(2, keys, load_chunk_ids, is_promotion=True)) |
| 400 | load_results = drain(tier) |
| 401 | assert len(load_results) == 1 |
| 402 | assert load_results[0].success |
| 403 | |
| 404 | for i, cid in enumerate(load_chunk_ids): |
| 405 | assert torch.allclose(tensor[cid], expected[i]), ( |
| 406 | f"Chunk {cid} data mismatch after store+load" |
| 407 | ) |
| 408 | |
| 409 | |
| 410 | def test_store_load_roundtrip_without_o_direct(tmp_path, monkeypatch): |
| 411 | """Buffered fallback must round-trip data when O_DIRECT is unsupported. |
| 412 | |
| 413 | Simulates filesystems (e.g. overlayfs, some NFS) that reject O_DIRECT by |
| 414 | forcing the capability probe to report it unavailable. |
| 415 | """ |
| 416 | monkeypatch.setattr( |
| 417 | "vllm.v1.kv_offload.tiering.fs.manager.probe_o_direct", |
| 418 | lambda _dir: False, |
| 419 | ) |
| 420 | tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) |
| 421 | tier = FileSystemTierManager( |
| 422 | offloading_spec=_MOCK_OFFLOADING_SPEC, |
| 423 | primary_kv_view=memoryview(tensor.numpy()), |
| 424 | tier_type="fs", |
| 425 | root_dir=str(tmp_path), |
| 426 | n_read_threads=4, |
| 427 | n_write_threads=4, |
| 428 | ) |
| 429 | try: |
| 430 | assert tier._use_o_direct is False |
| 431 | |
| 432 | keys = [key(0), key(1)] |
| 433 | expected = tensor[:2].clone() |
| 434 | tier.submit_store(make_job(1, keys, [0, 1])) |
| 435 | assert all(r.success for r in drain(tier)) |
| 436 | |
| 437 | tensor[:2] = 0.0 |
| 438 | tier.submit_load(make_job(2, keys, [2, 3], is_promotion=True)) |
| 439 | assert all(r.success for r in drain(tier)) |
| 440 | |
| 441 | for i, bid in enumerate([2, 3]): |
| 442 | assert torch.allclose(tensor[bid], expected[i]) |
| 443 | finally: |
| 444 | tier.shutdown() |
| 445 | |
| 446 | |
| 447 | def test_wait_idle_blocks_until_tasks_complete(): |
| 448 | """wait_idle must not return while a task is still in flight.""" |
| 449 | pool = DualQueueThreadPool(n_read_threads=1, n_write_threads=1) |
| 450 | gate = threading.Event() |
| 451 | pool.enqueue_store(job_id=1, n_tasks=1, tasks=[lambda: gate.wait(timeout=5.0)]) |
| 452 | |
| 453 | waiter = threading.Thread(target=pool.wait_idle) |
| 454 | waiter.start() |
| 455 | try: |
| 456 | waiter.join(timeout=0.2) |
| 457 | assert waiter.is_alive(), "wait_idle returned before task completed" |
| 458 | gate.set() |
| 459 | waiter.join(timeout=5.0) |
| 460 | assert not waiter.is_alive(), "wait_idle did not unblock" |
| 461 | finally: |
| 462 | gate.set() |
| 463 | pool.shutdown(wait=True) |
| 464 | waiter.join(timeout=5.0) |
| 465 | |
| 466 | |
| 467 | def test_batch_lookup_c_extension(tmp_path): |
| 468 | """Validates batch_lookup_C: empty, single, all-existing, all-missing, |
| 469 | mixed ordering, and input type validation.""" |
| 470 | try: |
| 471 | from vllm.fs_io_C import batch_lookup as batch_lookup_C |
| 472 | except ImportError: |
| 473 | pytest.skip("fs_io_C extension not built") |
| 474 | |
| 475 | # Setup |
| 476 | all_exist = [str(tmp_path / f"e{i}.bin") for i in range(3)] |
| 477 | for p in all_exist: |
| 478 | open(p, "w").close() |
| 479 | all_missing = [str(tmp_path / f"m{i}.bin") for i in range(3)] |
| 480 | |
| 481 | # Empty list |
| 482 | assert batch_lookup_C([]) == [] |
| 483 | |
| 484 | # Single existing / missing |
| 485 | assert batch_lookup_C([all_exist[0]]) == [True] |
| 486 | assert batch_lookup_C([all_missing[0]]) == [False] |
| 487 | |
| 488 | # All existing / all missing |
| 489 | assert batch_lookup_C(all_exist) == [True, True, True] |
| 490 | assert batch_lookup_C(all_missing) == [False, False, False] |
| 491 | |
| 492 | # Mixed — verifies index ordering is preserved |
| 493 | paths = [val for pair in zip(all_exist, all_missing) for val in pair] |
| 494 | assert batch_lookup_C(paths) == [True, False, True, False, True, False] |
| 495 | |
| 496 | # Input validation: non-list argument |
| 497 | with pytest.raises(TypeError): |
| 498 | batch_lookup_C(("/tmp/foo",)) |
| 499 | with pytest.raises(TypeError): |
| 500 | batch_lookup_C(None) |
| 501 | |
| 502 | # Input validation: non-str elements in list |
| 503 | with pytest.raises(TypeError): |
| 504 | batch_lookup_C([None]) |
| 505 | with pytest.raises(TypeError): |
| 506 | batch_lookup_C([b"/tmp/foo"]) |
| 507 | with pytest.raises(TypeError): |
| 508 | batch_lookup_C([42]) |
| 509 | with pytest.raises(TypeError): |
| 510 | batch_lookup_C([all_exist[0], None]) # valid first, invalid mid-list |
| 511 | |
| 512 | |
| 513 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 514 | def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): |
| 515 | import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod |
| 516 | |
| 517 | if use_c_ext and not mgr_mod._HAS_BATCH_LOOKUP_C: |
| 518 | pytest.skip("fs_io_C extension not built") |
| 519 | |
| 520 | monkeypatch.setattr(mgr_mod, "_HAS_BATCH_LOOKUP_C", use_c_ext) |
| 521 | |
| 522 | tier, _ = fs_tier |
| 523 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 524 | assert all(r.success for r in drain(tier)) |
| 525 | |
| 526 | results = lookup_and_wait(tier, [key(1), key(2)]) |
| 527 | assert results == [LookupResult.HIT, LookupResult.MISS] |
| 528 | |
| 529 | |
| 530 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 531 | def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext): |
| 532 | """Smoke test: a block id beyond the primary tensor's block count must |
| 533 | fail the job, for both the C extension and the Python fallback.""" |
| 534 | import vllm.v1.kv_offload.tiering.fs.io as io_mod |
| 535 | |
| 536 | if use_c_ext and not io_mod._HAS_FSIO_C: |
| 537 | pytest.skip("fs_io_C extension not built") |
| 538 | monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) |
| 539 | |
| 540 | tier, tensor = fs_tier |
| 541 | out_of_bounds_bid = tensor.shape[0] # one past the last valid block |
| 542 | |
| 543 | tier.submit_store(make_job(1, [key(1)], [out_of_bounds_bid])) |
| 544 | store_results = drain(tier) |
| 545 | assert len(store_results) == 1 |
| 546 | assert not store_results[0].success |
| 547 | |
| 548 | tier.submit_load(make_job(2, [key(1)], [out_of_bounds_bid], is_promotion=True)) |
| 549 | load_results = drain(tier) |
| 550 | assert len(load_results) == 1 |
| 551 | assert not load_results[0].success |
| 552 | |
| 553 | |
| 554 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 555 | def test_failed_load_corrects_verdict_and_removes_corrupt_file( |
| 556 | fs_tier, monkeypatch, use_c_ext |
| 557 | ): |
| 558 | """Failed-load livelock regression, covering the whole contract. |
| 559 | |
| 560 | A successful promotion leaves the cached HIT and the on-disk block intact. |
| 561 | A promotion that short-reads a truncated (corrupt) block fails, and in |
| 562 | get_finished_jobs() the tier removes the corrupt file (stores are atomic, |
| 563 | so a too-short file is genuine corruption) and marks the cached verdict |
| 564 | False. The SAME request's next lookup is then a MISS served from cache with |
| 565 | NO re-probe, so the scheduler cannot re-issue the doomed promotion. |
| 566 | """ |
| 567 | import vllm.v1.kv_offload.tiering.fs.io as io_mod |
| 568 | |
| 569 | if use_c_ext and not io_mod._HAS_FSIO_C: |
| 570 | pytest.skip("fs_io_C extension not built") |
| 571 | monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) |
| 572 | |
| 573 | tier, _ = fs_tier |
| 574 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 575 | assert all(r.success for r in drain(tier)) |
| 576 | path = tier.file_mapper.get_file_name(key(1)) |
| 577 | |
| 578 | ctx = ReqContext(req_id="livelock-req") |
| 579 | assert lookup_and_wait(tier, [key(1)], ctx=ctx) == [LookupResult.HIT] |
| 580 | |
| 581 | # A successful promotion must NOT touch the verdict or the file. |
| 582 | tier.submit_load(make_job(2, [key(1)], [0], is_promotion=True)) |
| 583 | results = drain(tier) |
| 584 | assert len(results) == 1 and results[0].success |
| 585 | assert tier.lookup(key(1), ctx) == LookupResult.HIT |
| 586 | assert os.path.exists(path) |
| 587 | |
| 588 | # Truncate below block_size so the next promotion short-reads. |
| 589 | with open(path, "wb") as f: |
| 590 | f.write(b"x" * 10) |
| 591 | tier.submit_load(make_job(3, [key(1)], [0], is_promotion=True)) |
| 592 | results = drain(tier) # get_finished_jobs() marks the verdict False here |
| 593 | assert len(results) == 1 and not results[0].success |
| 594 | |
| 595 | # Corrupt file removed; the SAME request now misses from cache, no re-probe. |
| 596 | assert not os.path.exists(path) |
| 597 | lm = tier._lookup_manager |
| 598 | assert tier.lookup(key(1), ctx) == LookupResult.MISS |
| 599 | assert lm._lookup_batch == [] |
| 600 | |
| 601 | # A FRESH request re-probes the tier (no cached verdict) and misses too, |
| 602 | # since the corrupt file is gone -- the real batch_lookup re-probe path. |
| 603 | fresh = ReqContext(req_id="fresh-after-short-read") |
| 604 | assert lookup_and_wait(tier, [key(1)], ctx=fresh) == [LookupResult.MISS] |
| 605 | |
| 606 | |
| 607 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 608 | def test_batched_partial_load_failure_keeps_loaded_blocks( |
| 609 | fs_tier, monkeypatch, use_c_ext |
| 610 | ): |
| 611 | """A batched promotion stops at the first bad block and reports how many |
| 612 | loaded before it (#50321). Corrupt the LAST block: the earlier blocks load |
| 613 | fine, so the job reports successful_keys for them and marks only the failed |
| 614 | tail a miss. The earlier keys stay HIT — including for the same request — |
| 615 | while the corrupt block stays a MISS (its file was removed).""" |
| 616 | import vllm.v1.kv_offload.tiering.fs.io as io_mod |
| 617 | |
| 618 | if use_c_ext and not io_mod._HAS_FSIO_C: |
| 619 | pytest.skip("fs_io_C extension not built") |
| 620 | monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) |
| 621 | |
| 622 | tier, _ = fs_tier |
| 623 | keys = [key(1), key(2), key(3)] # last one is the "bad" block |
| 624 | tier.submit_store(make_job(1, keys, [0, 1, 2])) |
| 625 | assert all(r.success for r in drain(tier)) |
| 626 | bad_path = tier.file_mapper.get_file_name(key(3)) |
| 627 | |
| 628 | ctx = ReqContext(req_id="batch-req") |
| 629 | assert lookup_and_wait(tier, keys, ctx=ctx) == [LookupResult.HIT] * 3 |
| 630 | |
| 631 | # Corrupt only the last block, then load the whole batch as one job. |
| 632 | with open(bad_path, "wb") as f: |
| 633 | f.write(b"x" * 10) |
| 634 | tier.submit_load(make_job(2, keys, [0, 1, 2], is_promotion=True)) |
| 635 | results = drain(tier) |
| 636 | # (a) the job fails but reports the two blocks that loaded before the bad one. |
| 637 | assert len(results) == 1 and not results[0].success |
| 638 | assert tuple(results[0].successful_keys) == (key(1), key(2)) |
| 639 | |
| 640 | # (b) Only the failed tail is a miss; the loaded blocks stay HIT on the same |
| 641 | # request, and nothing was re-probed. |
| 642 | lm = tier._lookup_manager |
| 643 | assert [tier.lookup(k, ctx) for k in keys] == [ |
| 644 | LookupResult.HIT, |
| 645 | LookupResult.HIT, |
| 646 | LookupResult.MISS, |
| 647 | ] |
| 648 | assert lm._lookup_batch == [] |
| 649 | |
| 650 | # (c) A fresh request re-probes: the loaded blocks are still on disk (HIT), |
| 651 | # only the corrupt block was removed (MISS). |
| 652 | tier.on_request_finished(ctx) |
| 653 | fresh = ReqContext(req_id="fresh-batch-req") |
| 654 | assert lookup_and_wait(tier, keys, ctx=fresh) == [ |
| 655 | LookupResult.HIT, |
| 656 | LookupResult.HIT, |
| 657 | LookupResult.MISS, |
| 658 | ] |
| 659 | |
| 660 | |
| 661 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 662 | def test_batched_load_first_block_fails_marks_whole_batch( |
| 663 | fs_tier, monkeypatch, use_c_ext |
| 664 | ): |
| 665 | """When the FIRST block fails, nothing loaded before it: the job reports no |
| 666 | successful_keys (None) and the whole batch is marked a miss for the |
| 667 | request.""" |
| 668 | import vllm.v1.kv_offload.tiering.fs.io as io_mod |
| 669 | |
| 670 | if use_c_ext and not io_mod._HAS_FSIO_C: |
| 671 | pytest.skip("fs_io_C extension not built") |
| 672 | monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) |
| 673 | |
| 674 | tier, _ = fs_tier |
| 675 | keys = [key(1), key(2), key(3)] # first one is the "bad" block |
| 676 | tier.submit_store(make_job(1, keys, [0, 1, 2])) |
| 677 | assert all(r.success for r in drain(tier)) |
| 678 | |
| 679 | ctx = ReqContext(req_id="batch-first-fail") |
| 680 | assert lookup_and_wait(tier, keys, ctx=ctx) == [LookupResult.HIT] * 3 |
| 681 | |
| 682 | with open(tier.file_mapper.get_file_name(key(1)), "wb") as f: |
| 683 | f.write(b"x" * 10) |
| 684 | tier.submit_load(make_job(2, keys, [0, 1, 2], is_promotion=True)) |
| 685 | results = drain(tier) |
| 686 | assert len(results) == 1 and not results[0].success |
| 687 | # Nothing loaded before the failure -> no partial success reported. |
| 688 | assert results[0].successful_keys is None |
| 689 | # The whole batch is a miss for this request. |
| 690 | assert [tier.lookup(k, ctx) for k in keys] == [LookupResult.MISS] * 3 |
| 691 | |
| 692 | |
| 693 | @pytest.mark.parametrize("use_c_ext", [True, False]) |
| 694 | def test_transient_load_failure_leaves_file(fs_tier, monkeypatch, use_c_ext): |
| 695 | """A transient host error (here ELOOP on open) is NOT a short read: the job |
| 696 | fails but the block file must survive untouched, on both the C and Python |
| 697 | paths. Deleting on a transient error would turn a passing hiccup into |
| 698 | permanent data loss.""" |
| 699 | import vllm.v1.kv_offload.tiering.fs.io as io_mod |
| 700 | |
| 701 | if use_c_ext and not io_mod._HAS_FSIO_C: |
| 702 | pytest.skip("fs_io_C extension not built") |
| 703 | monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) |
| 704 | |
| 705 | tier, _ = fs_tier |
| 706 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 707 | assert all(r.success for r in drain(tier)) |
| 708 | path = tier.file_mapper.get_file_name(key(1)) |
| 709 | with open(path, "rb") as f: |
| 710 | original = f.read() |
| 711 | |
| 712 | # Make open() fail with ELOOP (fd < 0) without truncating the block. Not |
| 713 | # chmod 000: CI runs as root, which bypasses permission bits, so open() |
| 714 | # would succeed and the load would not fail at all. |
| 715 | saved = path + ".saved" |
| 716 | loop = path + ".loop" |
| 717 | os.rename(path, saved) |
| 718 | os.symlink(loop, path) |
| 719 | os.symlink(path, loop) |
| 720 | |
| 721 | tier.submit_load(make_job(2, [key(1)], [0], is_promotion=True)) |
| 722 | results = drain(tier) |
| 723 | assert len(results) == 1 and not results[0].success |
| 724 | |
| 725 | # The path is left alone: a non-short-read error must not unlink. |
| 726 | assert os.path.lexists(path) |
| 727 | |
| 728 | os.unlink(path) |
| 729 | os.unlink(loop) |
| 730 | os.rename(saved, path) |
| 731 | with open(path, "rb") as f: |
| 732 | assert f.read() == original |
| 733 | |
| 734 | |
| 735 | # --------------------------------------------------------------------------- |
| 736 | # KV events |
| 737 | # --------------------------------------------------------------------------- |
| 738 | |
| 739 | |
| 740 | def test_successful_store_emits_stored_event(fs_tier_with_events): |
| 741 | """A completed store job emits one stored event with the job's keys.""" |
| 742 | tier = fs_tier_with_events |
| 743 | keys = [key(1), key(2)] |
| 744 | tier.submit_store(make_job(1, keys, [0, 1])) |
| 745 | assert all(r.success for r in drain(tier)) |
| 746 | |
| 747 | events = list(tier.take_events()) |
| 748 | assert len(events) == 1 |
| 749 | assert events[0].keys == keys |
| 750 | assert events[0].medium == Medium.STORAGE |
| 751 | assert events[0].locality is Locality.LOCAL |
| 752 | assert not events[0].removed |
| 753 | # take_events drains the buffer. |
| 754 | assert list(tier.take_events()) == [] |
| 755 | |
| 756 | |
| 757 | @pytest.mark.parametrize( |
| 758 | ("locality", "expected"), |
| 759 | [(None, None), ("REMOTE", Locality.REMOTE)], |
| 760 | ) |
| 761 | def test_store_event_uses_configured_locality(tmp_path, locality, expected): |
| 762 | tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) |
| 763 | locality_config = {} if locality is None else {"locality": locality} |
| 764 | tier = FileSystemTierManager( |
| 765 | offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), |
| 766 | primary_kv_view=memoryview(tensor.numpy()), |
| 767 | tier_type="fs", |
| 768 | root_dir=str(tmp_path), |
| 769 | enable_kv_events=True, |
| 770 | **locality_config, |
| 771 | ) |
| 772 | try: |
| 773 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 774 | assert all(r.success for r in drain(tier)) |
| 775 | |
| 776 | events = list(tier.take_events()) |
| 777 | assert len(events) == 1 |
| 778 | assert events[0].locality is expected |
| 779 | finally: |
| 780 | tier.shutdown() |
| 781 | |
| 782 | |
| 783 | def test_load_job_emits_no_event(fs_tier_with_events): |
| 784 | tier = fs_tier_with_events |
| 785 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 786 | results = drain(tier) |
| 787 | assert len(results) == 1 |
| 788 | assert results[0].success |
| 789 | list(tier.take_events()) |
| 790 | |
| 791 | tier.submit_load(make_job(2, [key(1)], [1], is_promotion=True)) |
| 792 | results = drain(tier) |
| 793 | assert len(results) == 1 |
| 794 | assert results[0].success |
| 795 | assert list(tier.take_events()) == [] |
| 796 | |
| 797 | |
| 798 | def test_mixed_job_results_emit_event_only_for_successful_job( |
| 799 | fs_tier_with_events, monkeypatch |
| 800 | ): |
| 801 | """With a failed and a successful store job in flight, exactly one event |
| 802 | is emitted and its keys belong to the successful job.""" |
| 803 | import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod |
| 804 | |
| 805 | tier = fs_tier_with_events |
| 806 | failing_path = tier.file_mapper.get_file_name(key(1)) |
| 807 | original_batch_store_block = mgr_mod.batch_store_block |
| 808 | |
| 809 | def flaky_batch_store_block(paths, *args, **kwargs): |
| 810 | if failing_path in paths: |
| 811 | raise OSError("injected store failure") |
| 812 | return original_batch_store_block(paths, *args, **kwargs) |
| 813 | |
| 814 | monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block) |
| 815 | |
| 816 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 817 | tier.submit_store(make_job(2, [key(2)], [1])) |
| 818 | results = drain(tier) |
| 819 | assert len(results) == 2 |
| 820 | by_id = {r.job_id: r for r in results} |
| 821 | assert not by_id[1].success |
| 822 | assert by_id[2].success |
| 823 | |
| 824 | events = list(tier.take_events()) |
| 825 | assert len(events) == 1 |
| 826 | assert events[0].keys == [key(2)] |
| 827 | |
| 828 | |
| 829 | def test_partially_failed_store_emits_no_event(fs_tier_with_events, monkeypatch): |
| 830 | """A store job with any failed block emits no event for the whole job.""" |
| 831 | import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod |
| 832 | |
| 833 | tier = fs_tier_with_events |
| 834 | failing_path = tier.file_mapper.get_file_name(key(2)) |
| 835 | original_batch_store_block = mgr_mod.batch_store_block |
| 836 | |
| 837 | def flaky_batch_store_block(paths, *args, **kwargs): |
| 838 | if failing_path in paths: |
| 839 | raise OSError("injected store failure") |
| 840 | return original_batch_store_block(paths, *args, **kwargs) |
| 841 | |
| 842 | monkeypatch.setattr(mgr_mod, "batch_store_block", flaky_batch_store_block) |
| 843 | |
| 844 | tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) |
| 845 | results = drain(tier) |
| 846 | assert len(results) == 1 |
| 847 | assert not results[0].success |
| 848 | assert list(tier.take_events()) == [] |
| 849 | assert tier._store_job_keys == {} |
| 850 | |
| 851 | |
| 852 | def test_events_disabled_by_default(fs_tier): |
| 853 | tier, _ = fs_tier |
| 854 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 855 | results = drain(tier) |
| 856 | assert len(results) == 1 |
| 857 | assert results[0].success |
| 858 | assert tier.events is None |
| 859 | assert tier._store_job_keys == {} |
| 860 | assert list(tier.take_events()) == [] |
| 861 | |
| 862 | |
| 863 | def test_events_require_global_kv_events_flag(tmp_path): |
| 864 | """Tier-level opt-in alone is not enough; the global flag gates events.""" |
| 865 | tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) |
| 866 | tier = FileSystemTierManager( |
| 867 | offloading_spec=_make_offloading_spec(enable_kv_cache_events=False), |
| 868 | primary_kv_view=memoryview(tensor.numpy()), |
| 869 | tier_type="fs", |
| 870 | root_dir=str(tmp_path), |
| 871 | enable_kv_events=True, |
| 872 | ) |
| 873 | try: |
| 874 | assert tier.events is None |
| 875 | tier.submit_store(make_job(1, [key(1)], [0])) |
| 876 | results = drain(tier) |
| 877 | assert len(results) == 1 |
| 878 | assert results[0].success |
| 879 | assert list(tier.take_events()) == [] |
| 880 | assert tier._store_job_keys == {} |
| 881 | finally: |
| 882 | tier.shutdown() |
| 883 | |
| 884 | |
| 885 | def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path): |
| 886 | """A GPU->CPU->fs cascade surfaces the tier-owned FS stored event via the |
| 887 | TieringOffloadingManager's aggregated take_events().""" |
| 888 | from vllm.v1.kv_offload.tiering.manager import ( |
| 889 | CPUPrimaryTierOffloadingManager, |
| 890 | TieringOffloadingManager, |
| 891 | ) |
| 892 | |
| 893 | tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) |
| 894 | view = memoryview(tensor.numpy()) |
| 895 | mock_region = MagicMock() |
| 896 | mock_region.create_kv_memoryview.return_value = view |
| 897 | primary = CPUPrimaryTierOffloadingManager(num_chunks=4, mmap_region=mock_region) |
| 898 | tier = FileSystemTierManager( |
| 899 | offloading_spec=_make_offloading_spec(enable_kv_cache_events=True), |
| 900 | primary_kv_view=primary.get_kv_memoryview(), |
| 901 | tier_type="fs", |
| 902 | root_dir=str(tmp_path), |
| 903 | enable_kv_events=True, |
| 904 | ) |
| 905 | manager = TieringOffloadingManager(primary_tier=primary, secondary_tiers=[tier]) |
| 906 | try: |
| 907 | keys = [key(1), key(2)] |
| 908 | manager.on_new_request(_CTX) |
| 909 | assert manager.prepare_store(keys, _CTX) is not None |
| 910 | manager.complete_store(keys, _CTX) # cascades to the fs tier |
| 911 | |
| 912 | events: list[OffloadingEvent] = [] |
| 913 | ctx = ScheduleEndContext(new_req_ids=[], preempted_req_ids=()) |
| 914 | deadline = time.monotonic() + 5.0 |
| 915 | while time.monotonic() < deadline and not events: |
| 916 | manager.on_schedule_end(ctx) |
| 917 | events.extend(manager.take_events()) |
| 918 | time.sleep(0.01) |
| 919 | |
| 920 | fs_events = [e for e in events if e.medium == Medium.STORAGE] |
| 921 | assert len(fs_events) == 1 |
| 922 | assert set(fs_events[0].keys) == set(keys) |
| 923 | assert not fs_events[0].removed |
| 924 | finally: |
| 925 | tier.shutdown() |
| 926 | |
| 927 | |
| 928 | def test_fs_tier_cross_tp_round_trip(tmp_path): |
| 929 | """TP=2 replicated writer and TP=4 reader share namespace and bytes.""" |
| 930 | root = str(tmp_path) |
| 931 | writer_tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS) |
| 932 | expected = writer_tensor[0].clone() |
| 933 | writer = FileSystemTierManager( |
| 934 | offloading_spec=_make_offloading_spec( |
| 935 | tp_size=2, world_size=2, rank=0, replicated_layout=True |
| 936 | ), |
| 937 | primary_kv_view=memoryview(writer_tensor.numpy()), |
| 938 | tier_type="fs", |
| 939 | root_dir=root, |
| 940 | n_read_threads=2, |
| 941 | n_write_threads=2, |
| 942 | ) |
| 943 | try: |
| 944 | writer.submit_store(make_job(1, [key(7)], [0])) |
| 945 | assert all(r.success for r in drain(writer)) |
| 946 | writer_base = writer.file_mapper.base_path |
| 947 | writer_path = writer.file_mapper.get_file_name(key(7)) |
| 948 | finally: |
| 949 | writer.shutdown() |
| 950 | |
| 951 | reader_tensor = _page_aligned_zero_tensor(4, _BLOCK_ELEMENTS) |
| 952 | reader = FileSystemTierManager( |
| 953 | offloading_spec=_make_offloading_spec( |
| 954 | tp_size=4, world_size=4, rank=3, replicated_layout=True |
| 955 | ), |
| 956 | primary_kv_view=memoryview(reader_tensor.numpy()), |
| 957 | tier_type="fs", |
| 958 | root_dir=root, |
| 959 | n_read_threads=2, |
| 960 | n_write_threads=2, |
| 961 | ) |
| 962 | try: |
| 963 | assert reader.file_mapper.base_path == writer_base |
| 964 | assert reader.file_mapper.get_file_name(key(7)) == writer_path |
| 965 | assert lookup_and_wait(reader, [key(7)]) == [LookupResult.HIT] |
| 966 | reader.submit_load(make_job(2, [key(7)], [1], is_promotion=True)) |
| 967 | assert all(r.success for r in drain(reader)) |
| 968 | assert torch.allclose(reader_tensor[1], expected) |
| 969 | finally: |
| 970 | reader.shutdown() |