[Bugfix][KV Connector][Mooncake] Wait for all peers before completing failed receives to avoid potential KV block write-after-free

tests/v1/kv_connector/unit/test_mooncake_connector.py

@@ -1595,6 +1595,152 @@ def test_pull_completion_requires_local_blocks(
     assert worker.finished_recving_reqs == expected_finished
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("terminal", ["error", "exception", "finish"])
+async def test_peer_failure_only_settles_its_pending_requests(terminal):
+    """Partial and duplicate replies must not release another peer's targets."""
+    with mooncake_register_worker(_make_test_kv_cache_config()) as (worker, _):
+        metas = {
+            f"d-req-{i}": PullReqMeta(
+                d_req_id=f"d-req-{i}",
+                transfer_id=f"xfer-{i}",
+                local_block_ids=[[i]],
+                remote_engine_id="p-engine",
+                remote_bootstrap_addr="http://bootstrap:33333",
+                pull_tasks_count=2,
+            )
+            for i in range(1, 4)
+        }
+        requests = list(metas.values())
+        partial = MooncakeXferResponse(
+            status=MooncakeXferResponseStatus.CONTINUE,
+            ok_reqs=["d-req-1"],
+            err_reqs=["d-req-2"],
+            err_msg="one request failed",
+        )
+        end = (
+            RuntimeError("connection lost")
+            if terminal == "exception"
+            else worker._encoder.encode(
+                MooncakeXferResponse(
+                    status=(
+                        MooncakeXferResponseStatus.FINISH
+                        if terminal == "finish"
+                        else MooncakeXferResponseStatus.ERROR
+                    ),
+                    err_msg="peer finished",
+                )
+            )
+        )
+        sock = mooncake_connector.make_zmq_socket.return_value.__enter__.return_value
+        sock.recv.side_effect = [worker._encoder.encode(partial)] * 2 + [end]
+        await worker.receive_kv_from_single_worker("tcp://producer:1234", metas)
+        assert [meta.pull_tasks_count for meta in requests] == [1, 1, 1]
+        assert [meta.failed for meta in requests] == [False, True, True]
+        assert worker.finished_recving_reqs == set()
+        assert worker.get_block_ids_with_load_errors() == {2, 3}
+        worker.process_pulling_result(
+            MooncakeXferResponse(
+                status=MooncakeXferResponseStatus.FINISH, ok_reqs=list(metas)
+            ),
+            metas,
+        )
+        assert worker.finished_recving_reqs == {meta.d_req_id for meta in requests}
+        assert all(meta.pull_tasks_count == 0 for meta in requests)
+        worker.shutdown()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("first_peer_succeeds", [True, False])
+async def test_expired_peer_finish_settles_pending_receive(first_peer_succeeds):
+    """An expired peer's omitted result must not strand the decoder's blocks."""
+    config = _make_test_kv_cache_config()
+    with mooncake_register_worker(config, kv_role="kv_producer") as (producer, _):
+        try:
+            producer.kv_caches_base_addr = [0x1000]
+            producer.block_len_per_layer = [256]
+            producer.kv_block_len_per_layer = [256]
+            producer.registered_layer_names = ["model.layers.0.self_attn"]
+            producer.registered_layer_indices = [0]
+            send_meta = SendBlockMeta(
+                p_req_id="p-req-1",
+                transfer_id="xfer-req-1",
+                local_block_ids=[[1]],
+                ready=asyncio.Event(),
+                expire_time=time.perf_counter() - 1,
+            )
+            send_meta.ready.set()
+            producer.reqs_need_send[send_meta.transfer_id] = send_meta
+
+            async def expire_before_send():
+                assert await producer.fetch_finished_sending_reqs() == {"p-req-1"}
+
+            metadata = _xfer_meta(
+                [
+                    _region(
+                        0x2000,
+                        row_offset=-1,
+                        layer_name=producer.registered_layer_names[0],
+                    )
+                ],
+                {"d-req-1": (send_meta.transfer_id, [[100]])},
+            )
+            sock = AsyncMock(spec=zmq.asyncio.Socket, send_multipart=AsyncMock())
+            with (
+                patch.object(send_meta.ready, "wait", side_effect=expire_before_send),
+                patch.object(producer, "_send_blocks") as send,
+            ):
+                await producer.send_kv_to_decode(b"consumer", sock, metadata)
+                send.assert_not_called()
+            assert send_meta.transfer_id not in producer.reqs_need_send
+            sock.send_multipart.assert_awaited_once()
+            _, payload = sock.send_multipart.call_args.args[0]
+            response = producer._xfer_resp_decoder.decode(payload)
+            assert response.status == MooncakeXferResponseStatus.FINISH
+            assert not response.ok_reqs
+            assert not response.err_reqs
+        finally:
+            producer.shutdown()
+            producer.is_kv_consumer = True
+
+    with mooncake_register_worker(config) as (worker, _):
+        try:
+            meta = PullReqMeta(
+                d_req_id="d-req-1",
+                transfer_id="xfer-req-1",
+                local_block_ids=[[100]],
+                remote_engine_id="p-engine",
+                remote_bootstrap_addr="http://bootstrap:33333",
+                pull_tasks_count=2,
+            )
+            worker.process_pulling_result(
+                MooncakeXferResponse(
+                    status=MooncakeXferResponseStatus.FINISH,
+                    ok_reqs=[meta.d_req_id] if first_peer_succeeds else None,
+                    err_reqs=None if first_peer_succeeds else [meta.d_req_id],
+                ),
+                {meta.d_req_id: meta},
+            )
+            assert meta.pull_tasks_count == 1
+            assert worker.finished_recving_reqs == set()
+            sock = (
+                mooncake_connector.make_zmq_socket.return_value.__enter__.return_value
+            )
+            sock.recv.return_value = payload
+            await worker.receive_kv_from_single_worker(
+                "tcp://producer:1234", {meta.d_req_id: meta}
+            )
+            assert meta.pull_tasks_count == 0
+            assert meta.failed
+            assert worker.finished_recving_reqs == {meta.d_req_id}
+            assert worker.get_block_ids_with_load_errors() == {100}
+            stats = worker.get_kv_connector_stats()
+            assert stats is not None
+            assert stats.reduce()["Num failed recvs"] == 1
+        finally:
+            worker.shutdown()
+
+
 @pytest.mark.asyncio
 async def test_kv_consumuer(monkeypatch):
     """Simulates a Consumer Worker (Decoder) initiating a pull from a Producer.

tests/v1/kv_connector/unit/test_mooncake_connector_hma.py

@@ -19,6 +19,8 @@
     MooncakeConnectorMetadata,
     MooncakeConnectorScheduler,
     MooncakeXferMetadata,
+    MooncakeXferResponse,
+    MooncakeXferResponseStatus,
     PullReqMeta,
     SendBlockMeta,
     TransferRegion,
@@ -553,3 +555,36 @@ def test_worker_failed_recv_reports_block_ids_without_hma():
     results = worker.get_transfer_results()
     assert results.failed_recving == set()
     assert results.finished_recving == {"d-req-1"}
+
+
+@pytest.mark.parametrize("swa_enabled", [False, True])
+@pytest.mark.parametrize("outcomes", [(False, True), (True, False), (False, False)])
+def test_failed_receive_waits_for_every_peer(swa_enabled, outcomes):
+    """A failed request keeps its destinations until the last peer completes."""
+    worker = _make_kv_consumer_worker(swa_enabled=swa_enabled)
+    meta = _make_pull_meta("d-req-1", [[1, 2], [3, 4]])
+    meta.pull_tasks_count = 2
+    error_reports = 0
+    try:
+        for index, successful in enumerate(outcomes):
+            response = MooncakeXferResponse(
+                status=MooncakeXferResponseStatus.FINISH,
+                ok_reqs=["d-req-1"] if successful else None,
+                err_reqs=None if successful else ["d-req-1"],
+                err_msg=None if successful else "peer failed",
+            )
+            worker.process_pulling_result(response, {"d-req-1": meta})
+            result = worker.get_transfer_results()
+            invalid = worker.get_block_ids_with_load_errors()
+            error_reports += bool(result.failed_recving or invalid)
+            assert meta.pull_tasks_count == 1 - index
+            assert result.finished_recving == ({"d-req-1"} if index else set())
+            assert meta.failed == (not all(outcomes[: index + 1]))
+            if result.failed_recving:
+                assert result.failed_recving == {"d-req-1"}
+            if invalid:
+                assert invalid == {1, 2, 3, 4}
+        assert error_reports == 1
+        assert worker.get_transfer_results().finished_recving == set()
+    finally:
+        worker.shutdown()

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

@@ -668,8 +668,8 @@ class PullReqMeta:
     expire_time: float = float("inf")
     # Designed for one D pairing to multiple P
     pull_tasks_count: int = 0
-    # Set once any worker reports a failure, so a success from another worker
-    # for the same request is not counted afterwards.
+    # A failure is sticky; still settle later peer results before releasing
+    # destination blocks to prevent write-after-free from delayed peer writes.
     failed: bool = False
 
 
@@ -2399,6 +2399,8 @@ async def receive_kv_from_single_worker(
         worker_addr: str,
         pull_metas: dict[ReqId, PullReqMeta],
     ):
+        # Each peer settles its own requests, sharing the per-request counters.
+        pull_metas = pull_metas.copy()
         req_ids = set(pull_metas)
         include_layout = (
             not self._kv_layout_id or worker_addr not in self._acked_layout_peers
@@ -2459,6 +2461,12 @@ async def receive_kv_from_single_worker(
                             self._acked_layout_peers.add(worker_addr)
                         self.process_pulling_result(response, pull_metas)
                         if status == MooncakeXferResponseStatus.FINISH:
+                            # P may omit expired requests from its final response.
+                            self._handle_failed_recv(
+                                pull_metas,
+                                list(pull_metas),
+                                "peer finished without reporting a transfer result",
+                            )
                             return
             except zmq.ContextTerminated:
                 logger.debug(
@@ -2478,8 +2486,11 @@ def _handle_failed_recv(
         """Report a failed remote KV load so the scheduler can fail or recompute it."""
         failed: list[ReqId] = []
         for req_id in req_ids:
-            pull_meta = pull_metas.get(req_id)
-            if pull_meta is None or pull_meta.failed:
+            pull_meta = pull_metas.pop(req_id, None)
+            if pull_meta is None:
+                continue
+            if pull_meta.failed:
+                self._complete_pull_task(pull_meta)
                 continue
             pull_meta.failed = True
             failed.append(req_id)
@@ -2493,16 +2504,24 @@ def _handle_failed_recv(
                 # abort_immediately request just to run request_finished). No D
                 # request is waiting on a load, and reporting one here would trip
                 # the scheduler's `assert req_id in self.requests`.
+                self._complete_pull_task(pull_meta)
                 continue
             if self._is_hma_required:
                 self._failed_recv_reqs.put(pull_meta.d_req_id)
             else:
                 self._invalid_block_ids.put(invalid)
-            self.finished_recving_reqs.add(pull_meta.d_req_id)
+            self._complete_pull_task(pull_meta)
 
         if failed:
             logger.error("pulling kv_caches for %s failed: %s", failed, reason)
 
+    def _complete_pull_task(self, pull_meta: PullReqMeta) -> None:
+        # Bootstrap failures have no peer tasks to settle.
+        if pull_meta.pull_tasks_count > 0:
+            pull_meta.pull_tasks_count -= 1
+        if pull_meta.pull_tasks_count == 0 and any(pull_meta.local_block_ids):
+            self.finished_recving_reqs.add(pull_meta.d_req_id)
+
     def get_block_ids_with_load_errors(self) -> set[int]:
         """Drain the blocks whose remote KV load failed since the last call."""
         result: set[int] = set()
@@ -2518,27 +2537,20 @@ def process_pulling_result(
         response: MooncakeXferResponse,
         pull_metas: dict[ReqId, PullReqMeta],
     ):
+        if response.err_reqs:
+            self._handle_failed_recv(
+                pull_metas, response.err_reqs, response.err_msg or "unknown error"
+            )
         ok_reqs: list[ReqId] = response.ok_reqs or []
 
         for req_id in ok_reqs:
-            pull_meta = pull_metas[req_id]
-            if pull_meta.failed:
-                continue
-            # No race because we are in async loop.
-            pull_meta.pull_tasks_count -= 1
-            # Empty pulls only release the producer's blocks; the consumer
-            # did not enter WAITING_FOR_REMOTE_KVS.
-            if pull_meta.pull_tasks_count == 0 and any(pull_meta.local_block_ids):
-                self.finished_recving_reqs.add(pull_meta.d_req_id)
+            pull_meta = pull_metas.pop(req_id, None)
+            if pull_meta is not None:
+                self._complete_pull_task(pull_meta)
 
         if ok_reqs:
             logger.debug("pulling kv_caches for %s finished", ok_reqs)
 
-        if response.err_reqs:
-            self._handle_failed_recv(
-                pull_metas, response.err_reqs, response.err_msg or "unknown error"
-            )
-
     async def _connect_to_prefiller_bootstrap(self, remote_bootstrap_addr: str):
         url = remote_bootstrap_addr + "/query"
         max_attempts = _BOOTSTRAP_MAX_ATTEMPTS