[Bugfix][KV Offload][OBJ] Preserve job completion during cleanup (#49947)
Signed-off-by: MINJUN GIL <alswnsrlf12@naver.com> Co-authored-by: OpenAI Codex <codex@openai.com>
MINJUN GIL committed
Jul 28, 2026 at 13:10 UTC
52c3c4a42fd13b62ba985b9ceb9b9969964ee83e
2 files changed
+133
-7
tests/v1/kv_offload/tiering/test_obj_tier.py
+108
-3
@@ -35,6 +35,10 @@ from vllm.v1.kv_offload.config import (
35
OffloadingParallelConfig,
36
)
37
from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult
38
+from vllm.v1.kv_offload.tiering.manager import (
39
+ CPUPrimaryTierOffloadingManager,
40
+ TieringOffloadingManager,
41
+)
42
from vllm.v1.kv_offload.tiering.obj.config import ObjStoreConfig
43
from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager
44
@@ -220,12 +224,14 @@ def _make_events_spec(enable_kv_cache_events: bool) -> SimpleNamespace:
224
def _make_tier(
225
num_blocks: int = 4,
226
offloading_spec: SimpleNamespace = _OFFLOADING_SPEC,
227
+ primary_kv_view: memoryview | None = None,
228
**tier_kwargs,
229
) -> tuple[ObjectStoreSecondaryTierManager, MockNixlAgent]:
230
"""Create a tier backed by a fresh MockNixlAgent."""
231
mock_agent = MockNixlAgent()
227
- tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE)
228
- view = memoryview(tensor.numpy())
232
+ if primary_kv_view is None:
233
+ tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE)
234
+ primary_kv_view = memoryview(tensor.numpy())
235
with (
236
patch("vllm.v1.kv_offload.tiering.obj.manager.nixl_agent_config"),
237
patch(
@@ -235,7 +241,7 @@ def _make_tier(
241
):
242
tier = ObjectStoreSecondaryTierManager(
243
offloading_spec=offloading_spec,
238
- primary_kv_view=view,
244
+ primary_kv_view=primary_kv_view,
245
tier_type="obj",
246
store_config=_STORE_CONFIG,
247
prefix=_RUN_PREFIX,
@@ -449,6 +455,105 @@ class TestMockObjTierFailures:
455
assert not by_id[1].success
456
assert by_id[2].success
457
458
+ def test_release_xfer_failure_retries_without_losing_result(self, monkeypatch):
459
+ tier, agent = _make_tier(num_blocks=4)
460
+ agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed"))
461
+ release_xfer = MagicMock(
462
+ side_effect=[RuntimeError("transfer is still active"), None]
463
+ )
464
+ monkeypatch.setattr(agent, "release_xfer_handle", release_xfer)
465
+
466
+ tier.submit_store(make_job(1, [key(1)], [0]))
467
+
468
+ # The transfer handle could not be released safely, so the job must
469
+ # remain tracked and must not be finalized yet.
470
+ assert list(tier.get_finished_jobs()) == []
471
+ assert 1 in tier._transfers
472
+
473
+ # Cleanup is retried without polling again or changing the failure
474
+ # verdict. The completion is then returned exactly once.
475
+ results = list(tier.get_finished_jobs())
476
+ assert len(results) == 1
477
+ assert results[0].job_id == 1
478
+ assert not results[0].success
479
+ assert agent.check_xfer_state.call_count == 2
480
+ assert release_xfer.call_count == 2
481
+ assert not tier._transfers
482
+ assert list(tier.get_finished_jobs()) == []
483
+
484
+ @pytest.mark.parametrize(
485
+ "cleanup_method", ["release_dlist_handle", "deregister_memory"]
486
+ )
487
+ def test_post_transfer_cleanup_failure_does_not_lose_result(
488
+ self, monkeypatch, cleanup_method
489
+ ):
490
+ tier, agent = _make_tier(num_blocks=4)
491
+ monkeypatch.setattr(
492
+ agent,
493
+ cleanup_method,
494
+ MagicMock(side_effect=RuntimeError("cleanup failed")),
495
+ )
496
+
497
+ tier.submit_store(make_job(1, [key(1)], [0]))
498
+ results = list(tier.get_finished_jobs())
499
+
500
+ assert len(results) == 1
501
+ assert results[0].job_id == 1
502
+ assert results[0].success
503
+ assert not tier._transfers
504
+ assert list(tier.get_finished_jobs()) == []
505
+
506
+ def test_xfer_cleanup_retry_finalizes_parent_job_and_primary_pin(self, monkeypatch):
507
+ num_blocks = 4
508
+ tensor = torch.zeros((num_blocks, _BLOCK_ELEMENTS), dtype=_DTYPE)
509
+ primary_kv_view = memoryview(tensor.numpy())
510
+ mmap_region = MagicMock()
511
+ mmap_region.create_kv_memoryview.return_value = primary_kv_view
512
+ primary_tier = CPUPrimaryTierOffloadingManager(
513
+ num_blocks=num_blocks, mmap_region=mmap_region
514
+ )
515
+ obj_tier, agent = _make_tier(
516
+ num_blocks=num_blocks, primary_kv_view=primary_kv_view
517
+ )
518
+ manager = TieringOffloadingManager(
519
+ primary_tier=primary_tier, secondary_tiers=[obj_tier]
520
+ )
521
+
522
+ keys = [key(1)]
523
+ primary_result = primary_tier.prepare_store(keys, _CTX)
524
+ assert primary_result is not None
525
+ primary_tier.complete_store(keys, _CTX, success=True)
526
+ job = manager.create_store_job(keys, _CTX)
527
+ obj_tier.submit_store(job)
528
+
529
+ block = primary_tier._policy.get(keys[0])
530
+ assert block is not None
531
+ assert block.ref_cnt == 1
532
+ assert len(manager._transfer_jobs) == 1
533
+
534
+ agent.check_xfer_state = MagicMock(side_effect=RuntimeError("poll failed"))
535
+ release_xfer = MagicMock(
536
+ side_effect=[RuntimeError("transfer is still active"), None]
537
+ )
538
+ monkeypatch.setattr(agent, "release_xfer_handle", release_xfer)
539
+ schedule_context = ScheduleEndContext(new_req_ids=[], preempted_req_ids=())
540
+
541
+ manager.on_schedule_end(schedule_context)
542
+
543
+ assert len(manager._transfer_jobs) == 1
544
+ assert block.ref_cnt == 1
545
+ assert len(obj_tier._transfers) == 1
546
+ assert manager.has_pending_work()
547
+
548
+ manager.on_schedule_end(schedule_context)
549
+
550
+ assert manager._transfer_jobs == {}
551
+ assert block.ref_cnt == 0
552
+ assert obj_tier._transfers == {}
553
+ assert not manager.has_pending_work()
554
+ assert agent.check_xfer_state.call_count == 2
555
+ assert release_xfer.call_count == 2
556
+
557
558
class TestMockObjTierShutdown:
559
def test_shutdown_clears_in_flight_transfers(self):
vllm/v1/kv_offload/tiering/obj/manager.py
+25
-4
@@ -307,15 +307,36 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):
307
else:
308
if state == NIXL_PROC:
309
continue
310
- elif state == NIXL_DONE:
310
+ if state == NIXL_DONE:
311
success = True
312
else:
313
success = False
314
logger.warning("transfer failed job=%d state=%s", job_id, state)
315
+
316
+ try:
317
+ self._agent.release_xfer_handle(entry.xfer_handle)
318
+ except Exception as exc:
319
+ # Keep the entry until NIXL confirms that the transfer handle
320
+ # can be released. The transfer may still access primary-tier
321
+ # memory, so publishing its result would allow unsafe reuse.
322
+ logger.warning("release_xfer_handle failed for job %d: %s", job_id, exc)
323
+ continue
324
+
325
+ # Once the transfer handle is released, these remaining cleanup
326
+ # failures must not suppress the job completion. They can leak
327
+ # NIXL metadata, but cannot leave an active data transfer behind.
328
+ try:
329
+ self._agent.release_dlist_handle(entry.obj_handle)
330
+ except Exception as exc:
331
+ logger.warning(
332
+ "release_dlist_handle failed for job %d: %s", job_id, exc
333
+ )
334
+ try:
335
+ self._agent.deregister_memory(entry.files_desc)
336
+ except Exception as exc:
337
+ logger.warning("deregister_memory failed for job %d: %s", job_id, exc)
338
+
339
del self._transfers[job_id]
316
- self._agent.release_xfer_handle(entry.xfer_handle)
317
- self._agent.release_dlist_handle(entry.obj_handle)
318
- self._agent.deregister_memory(entry.files_desc)
340
self._pending_results.append(JobResult(job_id=job_id, success=success))
341
342
def get_finished_jobs(self) -> Iterable[JobResult]: