[Bugfix][KV Offload][P2P] Avoid leaking failed IDs for finished requests

tests/v1/kv_offload/tiering/p2p/test_manager.py

@@ -721,6 +721,19 @@ def test_timed_out_load_holds_chunks_until_abort_ack(monkeypatch):
 
 
 class TestOnRequestFinished:
+    def test_finished_inflight_load_does_not_leak_failed_id(self):
+        mgr, peer = _build_paired_managers()
+        params = _remote_prefiller_kv_params("B", 2, kv_request_id="req-finish")
+        ctx = _req_context(params)
+        mgr.on_new_request(ctx)
+        job = _job_metadata(7, kv_params=params)
+        mgr.submit_load(job)
+        mgr.on_request_finished(ctx)
+        results = list(mgr.get_finished_jobs())
+        assert len(results) == 1 and results[0].job_id == 7
+        assert not results[0].success
+        assert mgr._failed_req_ids == set()
+
     def _make_with_failed(self) -> P2PSecondaryTierManager:
         mgr = _make_manager()
         mgr._failed_req_ids = {"req-1"}
@@ -2022,16 +2035,24 @@ def poll(self_inner):
         assert mgr._kv_to_session["req-fast"] is sess
         assert sess.stores_added == []
 
-    def test_failed_load_records_kv_request_id(self):
-        """A LoadResult(success=False) from session.poll() must add its
-        kv_request_id to _failed_req_ids so future lookups return MISS."""
+    @pytest.mark.parametrize("request_finished", [False, True])
+    def test_failed_load_records_kv_request_id(self, request_finished):
+        """An unfinished failed load records its id so future lookups return MISS.
+        Finished requests still complete jobs without retaining their ids."""
         mgr = _make_manager()
         peer = "10.0.0.1:8000"
         sess = _FakeSession(
             peer_id=peer,
             alive=True,
             connected=True,
-            loads=[LoadResult(job_id=5, kv_request_id="req-5", success=False)],
+            loads=[
+                LoadResult(
+                    job_id=5,
+                    kv_request_id="req-5",
+                    success=False,
+                    request_finished=request_finished,
+                )
+            ],
         )
         mgr._sessions[peer] = sess
 
@@ -2044,7 +2065,7 @@ def poll(self_inner):
         mgr._poll_once()
 
         assert mgr._finished_jobs == [JobResult(job_id=5, success=False)]
-        assert "req-5" in mgr._failed_req_ids
+        assert ("req-5" in mgr._failed_req_ids) is not request_finished
 
 
 # ---------------------------------------------------------------------------

tests/v1/kv_offload/tiering/p2p/test_sessions.py

@@ -698,6 +698,7 @@ def test_finish_abort_ack_emits_one_failure(self):
         session.finish_request("req-1")
         (failure,) = session.poll().loads
         assert (failure.job_id, failure.success) == (8, False)
+        assert failure.request_finished
         assert not _released(failure)
         assert session._client.has_active_loads is True
 
@@ -750,6 +751,7 @@ def test_late_transfer_done_after_abort_ack_is_ignored(self):
         )
         (failure,) = session.poll().loads
         assert (failure.job_id, failure.success) == (10, False)
+        assert failure.request_finished
         assert _released(failure)
 
         conn.enqueue(
@@ -3017,6 +3019,71 @@ def test_pending_close_returns_pending_stores(self):
 
 
 class TestDisconnect:
+    @pytest.mark.parametrize(
+        "completion", ["abort_ack", "transfer_done", "timeout", "close"]
+    )
+    @pytest.mark.parametrize("finished", [False, True])
+    def test_late_failure_preserves_request_finish(
+        self, lookup_clock, completion, finished
+    ):
+        session, conn, _ = _make_session()
+        _activate(session, conn)
+        session.request_blocks(7, "req-7", [OffloadKey(b"key")], [0])
+        if finished:
+            session.finish_request("req-7")
+        if completion == "close":
+            result = session.close()
+            assert result.failed_jobs == [7]
+            assert result.failed_req_ids == ([] if finished else ["req-7"])
+            return
+        loads = []
+        if not finished and completion in ("timeout", "abort_ack"):
+            lookup_clock.now = _PD_LOAD_TIMEOUT_S + 1
+            loads.extend(session.poll().loads)
+        if completion == "timeout":
+            if finished:
+                lookup_clock.now = _ABORT_ACK_TIMEOUT_S + 1
+            else:
+                lookup_clock.now += _ABORT_ACK_TIMEOUT_S + 1
+        else:
+            terminal = AbortAckMsg if completion == "abort_ack" else TransferDoneMsg
+            msg = {
+                TYPE_KEY: terminal.TYPE,
+                TransferDoneMsg.KV_REQUEST_ID: "req-7",
+                TransferDoneMsg.ROUND_SEQ: 0,
+            }
+            if terminal is TransferDoneMsg:
+                msg[TransferDoneMsg.SUCCESS] = False
+            conn.enqueue(msg)
+        loads.extend(session.poll().loads)
+        assert len(loads) == 1
+        assert (loads[0].job_id, loads[0].kv_request_id, loads[0].success) == (
+            7,
+            "req-7",
+            False,
+        )
+        assert loads[0].request_finished is finished
+
+    def test_finish_marks_buffered_result_before_request_id_reuse(self):
+        session, conn, _ = _make_session()
+        _activate(session, conn)
+        session.request_blocks(7, "req-7", [OffloadKey(b"old")], [0])
+        session._client.on_transfer_done("req-7", success=False, round_seq=0)
+        session.finish_request("req-7")
+        session.request_blocks(8, "req-7", [OffloadKey(b"new")], [1])
+        loads = session.poll().loads
+        assert len(loads) == 1 and loads[0].job_id == 7 and not loads[0].success
+        assert loads[0].request_finished
+        conn.enqueue(
+            {
+                TYPE_KEY: TransferDoneMsg.TYPE,
+                TransferDoneMsg.KV_REQUEST_ID: "req-7",
+                TransferDoneMsg.ROUND_SEQ: 0,
+                TransferDoneMsg.SUCCESS: False,
+            }
+        )
+        assert session.poll().loads == [LoadResult(8, "req-7", success=False)]
+
     def test_disconnect_marks_session_dead(self):
         session, conn, _ = _make_session()
         _activate(session, conn)

tests/v1/kv_offload/tiering/test_p2p_lifecycle_e2e.py

@@ -0,0 +1,177 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import time
+
+import pytest
+
+
+@pytest.fixture
+def p2p_control_peer():
+    """A control-only peer: delay its handshake or fail fetches on demand."""
+    import queue
+    import socket
+    import threading
+
+    from vllm.v1.kv_offload.tiering.p2p.control.base import ControlConnection
+    from vllm.v1.kv_offload.tiering.p2p.control.zmq import ZmqTransport
+
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        port = sock.getsockname()[1]
+    handshake = threading.Event()
+    reply_to_fetch = threading.Event()
+    fetches: queue.Queue[dict] = queue.Queue()
+    ready = threading.Event()
+    stop = threading.Event()
+    errors = []
+
+    def serve():
+        transport = None
+        try:
+            transport = ZmqTransport(f"127.0.0.1:{port}", "127.0.0.1", port)
+            connections: list[ControlConnection] = []
+            acknowledged = set()
+            ready.set()
+            while not stop.is_set():
+                connections.extend(transport.poll())
+                for conn in connections:
+                    if handshake.is_set() and conn.peer_id not in acknowledged:
+                        conn.send(
+                            {"type": "connect_ack", "peer_id": f"127.0.0.1:{port}"}
+                        )
+                        acknowledged.add(conn.peer_id)
+                    for msg in conn.recv():
+                        kind = msg["type"]
+                        if kind == "fetch" and msg["keys"]:
+                            fetches.put(msg)
+                            if not reply_to_fetch.is_set():
+                                continue
+                            conn.send(
+                                {
+                                    "type": "transfer_done",
+                                    "kv_request_id": msg["kv_request_id"],
+                                    "round_seq": msg["round_seq"],
+                                    "success": False,
+                                }
+                            )
+                        elif kind == "abort_fetch":
+                            conn.send(
+                                {
+                                    "type": "abort_ack",
+                                    "kv_request_id": msg["kv_request_id"],
+                                    "round_seq": msg["round_seq"],
+                                }
+                            )
+                stop.wait(0.001)
+        except Exception as exc:
+            errors.append(exc)
+            ready.set()
+        finally:
+            if transport is not None:
+                transport.close()
+
+    thread = threading.Thread(target=serve, daemon=True)
+    thread.start()
+    assert ready.wait(10)
+    assert not errors
+    try:
+        yield port, handshake, reply_to_fetch, fetches
+    finally:
+        stop.set()
+        thread.join(timeout=10)
+        assert not thread.is_alive()
+        assert not errors
+
+
+def test_inference_reuses_a_finished_kv_request_id(
+    vllm_runner, monkeypatch, p2p_control_peer
+):
+    """A cancellation's late failure must not suppress a later request's fetch."""
+    from vllm import SamplingParams
+    from vllm.config import KVTransferConfig
+
+    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+    monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
+    monkeypatch.setenv("VLLM_P2P_HANDSHAKE_TIMEOUT_S", "120")
+    peer_port, handshake, reply_to_fetch, fetches = p2p_control_peer
+    import socket
+
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        local_port = sock.getsockname()[1]
+    config = KVTransferConfig(
+        kv_connector="OffloadingConnector",
+        kv_role="kv_both",
+        kv_load_failure_policy="recompute",
+        kv_connector_extra_config={
+            "spec_name": "TieringOffloadingSpec",
+            "cpu_bytes_to_use": 16 * 1024 * 1024,
+            "secondary_tiers": [
+                {"type": "p2p", "host": "127.0.0.1", "port": local_port}
+            ],
+        },
+    )
+    with vllm_runner(
+        "HuggingFaceTB/SmolLM2-135M",
+        skip_tokenizer_init=True,
+        dtype="float16",
+        max_model_len=256,
+        enforce_eager=True,
+        gpu_memory_utilization=0.35,
+        max_num_seqs=2,
+        max_num_batched_tokens=512,
+        enable_prefix_caching=True,
+        disable_log_stats=False,
+        kv_transfer_config=config,
+    ) as runner:
+        llm = runner.llm
+        params = SamplingParams(
+            temperature=0,
+            max_tokens=4,
+            ignore_eos=True,
+            extra_args={
+                "kv_transfer_params": {
+                    "remote_prefiller": {
+                        "remote_host": "127.0.0.1",
+                        "remote_port": peer_port,
+                        "kv_request_id": "retired",
+                    }
+                }
+            },
+        )
+        prompt = {"prompt_token_ids": [3 + i % 100 for i in range(192)]}
+
+        handshake.set()
+        request_ids = llm.enqueue([prompt], params, use_tqdm=False)
+        deadline = time.monotonic() + 30
+        while fetches.empty():
+            assert time.monotonic() < deadline, "no initial P2P fetch"
+            llm.llm_engine.step()
+        fetches.get_nowait()
+        llm.llm_engine.abort_request(request_ids, internal=True)
+        from vllm.v1.metrics.reader import Counter
+
+        deadline = time.monotonic() + 30
+        while True:
+            llm.generate(
+                [{"prompt_token_ids": [0]}],
+                SamplingParams(temperature=0, max_tokens=1),
+                use_tqdm=False,
+            )
+            failures = [
+                metric.value
+                for metric in llm.get_metrics()
+                if isinstance(metric, Counter)
+                and "promotion_job_failures" in metric.name
+            ]
+            if sum(failures) > 0:
+                break
+            assert time.monotonic() < deadline, (
+                f"cancelled load was not drained: {failures}"
+            )
+        reply_to_fetch.set()
+        second_prompt = {"prompt_token_ids": [103 + i % 100 for i in range(192)]}
+        result = llm.generate([second_prompt], params, use_tqdm=False)[0]
+        assert len(result.outputs[0].token_ids) == 4
+        assert not fetches.empty(), "late failure poisoned the reused KV request id"
+        assert fetches.get_nowait()["kv_request_id"] == "retired"

vllm/v1/kv_offload/tiering/p2p/manager.py

@@ -1135,7 +1135,7 @@ def _poll_once(self) -> None:
                         deferred_release=lr.deferred_release,
                     )
                 )
-                if not lr.success:
+                if not lr.success and not lr.request_finished:
                     self._failed_req_ids.add(lr.kv_request_id)
             for sr in result.stores:
                 self._finished_jobs.append(

vllm/v1/kv_offload/tiering/p2p/session/client.py

@@ -61,6 +61,7 @@ class _InboundLoadState:
     release: threading.Event | None = None
     # Its AbortAck timed out; only close() can release it now.
     ack_expired: bool = False
+    request_finished: bool = False
 
 
 @dataclass
@@ -109,6 +110,7 @@ class LoadResult(NamedTuple):
     # A failed load whose blocks the peer may still write: release them
     # once this is set.
     deferred_release: threading.Event | None = None
+    request_finished: bool = False
 
 
 class ClientCloseResult(NamedTuple):
@@ -211,6 +213,7 @@ def _abort_load(
                 kv_request_id=kv_request_id,
                 success=False,
                 deferred_release=load.release,
+                request_finished=load.request_finished,
             )
         )
 
@@ -307,10 +310,17 @@ def finish(self, kv_request_id: str) -> None:
         at once and kept until their AbortAck (or the session's close)
         releases their blocks.
         """
+        self._completed_loads = [
+            result._replace(request_finished=True)
+            if result.kv_request_id == kv_request_id
+            else result
+            for result in self._completed_loads
+        ]
         st = self._requests.get(kv_request_id)
         if st is None:
             return
         for round_seq, load in st.loads.items():
+            load.request_finished = True
             if load.aborted_at is None:
                 self._abort_load(kv_request_id, round_seq, load)
         if st.peer_lookup_open:
@@ -346,6 +356,7 @@ def on_transfer_done(
                     job_id=load.job_id,
                     kv_request_id=kv_request_id,
                     success=success,
+                    request_finished=load.request_finished,
                 )
             )
             self._on_load_terminal(kv_request_id, st)
@@ -569,7 +580,8 @@ def close(self) -> ClientCloseResult:
         failed_req_ids = [
             req_id
             for req_id, st in self._requests.items()
-            if st.loads or any(hit is None for hit in st.probes.values())
+            if any(not load.request_finished for load in st.loads.values())
+            or any(hit is None for hit in st.probes.values())
         ]
         self._requests.clear()
         self._flush_pending.clear()