[Bugfix][NIXL] Fail pull reads whose producer lease has expired

tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py

@@ -1035,6 +1035,24 @@ def test_d_node_request_finished_exports_blocks_expiry_time():
     assert kv["remote_blocks_expiry_time"] > time.perf_counter()
 
 
+def test_p_node_request_finished_exports_blocks_expiry_time():
+    """P exports its lease expiry so D can decline reads after P reaps it,
+    with or without bidirectional transfer."""
+    vllm_config = create_vllm_config()
+    scheduler = create_scheduler(vllm_config)
+    BS = vllm_config.cache_config.block_size
+    req = create_request(
+        request_id=601, block_size=BS, num_tokens=int(BS * 2.5), do_remote_decode=True
+    )
+    scheduler.add_request(req)
+    lease_start = time.perf_counter()
+    so = scheduler.schedule()
+    eco = scheduler.update_from_output(so, create_model_runner_output(reqs=[req]))
+    kv = eco[0].outputs[0].kv_transfer_params
+    assert kv["do_remote_prefill"] is True
+    assert kv["remote_blocks_expiry_time"] >= lease_start + 30
+
+
 def test_handshake_listener_appends_perf_counter_frame():
     """The handshake reply carries a 2nd frame with the listener's live
     perf_counter."""

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -2288,6 +2288,8 @@ def recv_worker():
     worker._pending_recv_notifs = {}
     worker._reqs_to_send = {}
     worker._replicated_pcp_done_sending = set()
+    worker._engine_clock_offset = {}
+    worker._heartbeat_sends = {}
     worker._is_hma_required = False
     worker._has_mamba = False
     worker.use_host_buffer = False

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -57,6 +57,7 @@ def region_pull_worker():
     worker.enable_permute_local_kv = False
     worker.enable_heterogeneous_attn_post_process = False
     worker._engine_last_active = {}
+    worker._heartbeat_sends = {}
     worker._handshake_lock = threading.RLock()
     worker._handshake_futures = {}
     worker._transfer_layer_group_ids = ()

tests/v1/kv_connector/unit/test_nixl_heartbeat.py

@@ -3,13 +3,25 @@
 """Unit tests for the scheduler-driven heartbeat / lease-renewal system."""
 
 import time
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 
+from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
+    HeartbeatInfo,
+    NixlConnectorMetadata,
+    RemoteMeta,
+    ReqMeta,
+)
 from vllm.v1.outputs import KVConnectorOutput
 
-from .utils import create_request, make_nixl_scheduler
+from .utils import (
+    create_request,
+    create_vllm_config,
+    make_kv_cache_config,
+    make_nixl_scheduler,
+    maybe_update_block_size,
+)
 
 _ENGINE_A = "my-engine-id"
 
@@ -251,3 +263,161 @@ def test_reaper_reclaims_shorter_lease_behind_later_deadline(monkeypatch):
     assert done == {"Y"}
     assert list(w._reqs_to_send) == ["X"]
     assert w._reqs_to_process == {"X"}
+
+
+# ===================================================================
+# Worker: pull reads check the producer lease, crediting heartbeats
+# ===================================================================
+
+_REMOTE = "remote_engine"
+_LEASE = 30.0
+
+
+@pytest.fixture
+def pull_worker(dist_init, monkeypatch):
+    """Non-bidirectional D worker paired with P, on a fake clock that reads 0
+    when P stamps a 30 s lease."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker
+
+    from .test_nixl_connector import FakeNixlConnectorWorker, FakeNixlWrapper
+
+    monkeypatch.setattr(base_worker, "NixlWrapper", FakeNixlWrapper)
+    block_size = maybe_update_block_size(32)
+    vllm_config = create_vllm_config(block_size=block_size)
+    worker = FakeNixlConnectorWorker(
+        vllm_config,
+        vllm_config.kv_transfer_config.engine_id,
+        hand_shake_latency=0,
+        kv_cache_config=make_kv_cache_config(block_size=block_size, num_blocks=2),
+    )
+    assert worker.REMOTE_ENGINE_ID == _REMOTE
+    worker._remote_agents[_REMOTE], _ = worker._nixl_handshake(
+        "localhost", 1234, 1, _REMOTE
+    )
+    worker._engine_clock_offset[_REMOTE] = 0.0
+    clock = {"t": 0.0}
+    monkeypatch.setattr(base_worker.time, "perf_counter", lambda: clock["t"])
+    yield worker, clock
+    worker.shutdown()
+
+
+def _heartbeat_at(worker, clock, t: float) -> None:
+    clock["t"] = t
+    meta = NixlConnectorMetadata()
+    meta.heartbeat_by_engine = {
+        _REMOTE: HeartbeatInfo(
+            req_ids={"prefill-req"}, host="localhost", port=1234, tp_size=1
+        )
+    }
+    worker.start_load_kv(meta)
+
+
+def _read_at(worker, clock, t: float) -> None:
+    clock["t"] = t
+    meta = NixlConnectorMetadata()
+    meta.add_new_req_to_recv(
+        request_id="req",
+        local_block_ids=([0, 1],),
+        kv_transfer_params={
+            "remote_block_ids": ([0, 1],),
+            "remote_engine_id": _REMOTE,
+            "remote_request_id": "prefill-req",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "remote_blocks_expiry_time": _LEASE,
+        },
+        awaiting_kvs=True,
+    )
+    worker.start_load_kv(meta)
+
+
+_EVERY_5S = [5.0 * i for i in range(1, 9)]
+
+
+@pytest.mark.parametrize(
+    ("heartbeat_times", "send_fails", "read_time", "expect_expired"),
+    [
+        # Queued past the exported expiry, kept alive by a heartbeat every 5 s.
+        pytest.param(_EVERY_5S, False, 45.0, False, id="renewed"),
+        # Heartbeats that failed to send extended nothing.
+        pytest.param(_EVERY_5S, True, 45.0, True, id="send_failed"),
+        # A 30 s heartbeat gap exceeds P's 20 s extension: the lease lapsed.
+        pytest.param([5.0, 10.0, 40.0], False, 41.0, True, id="heartbeat_gap"),
+        pytest.param([], False, 31.0, True, id="no_heartbeat"),
+        # Heartbeats that start this close to the expiry are not credited, so
+        # the read is declined past expiry - 5 s.
+        pytest.param([20.0, 25.0], False, 26.0, True, id="late_first_heartbeat"),
+    ],
+)
+def test_pull_read_checks_producer_lease(
+    pull_worker, heartbeat_times, send_fails, read_time, expect_expired
+):
+    """D must not read P blocks whose lease P may have reaped and reused."""
+    worker, clock = pull_worker
+    with patch.object(
+        worker.nixl_wrapper,
+        "send_notif",
+        side_effect=RuntimeError("unreachable") if send_fails else None,
+    ):
+        for t in heartbeat_times:
+            _heartbeat_at(worker, clock, t)
+    _read_at(worker, clock, read_time)
+
+    result = worker.get_transfer_results()
+    assert result.finished_recving == {"req"}
+    assert result.failed_recving == ({"req"} if expect_expired else set())
+    assert worker.xfer_stats.data["num_kv_expired_reqs"] == (
+        [1] if expect_expired else []
+    )
+    assert worker.xfer_stats.data["num_failed_transfers"] == []
+
+
+@pytest.mark.parametrize("read_time", [15.0, 50.0], ids=["loaded", "expired"])
+def test_finished_load_drops_heartbeat_record(pull_worker, read_time):
+    """Heartbeat state does not outlive the request's load on this worker."""
+    worker, clock = pull_worker
+    _heartbeat_at(worker, clock, 5.0)
+    _read_at(worker, clock, read_time)
+    assert worker.get_transfer_results().finished_recving == {"req"}
+    assert not worker._heartbeat_sends
+
+
+def test_pull_read_completing_after_lease_expiry_fails(pull_worker):
+    """A read issued in time but finishing after the lease lapsed may hold
+    another request's KV, so its blocks are reported invalid."""
+    worker, clock = pull_worker
+    worker.nixl_wrapper._cycles_before_xfer_done = 1
+    _read_at(worker, clock, 10.0)
+    assert not worker.get_transfer_results().finished_recving
+
+    clock["t"] = _LEASE
+    result = worker.get_transfer_results()
+    assert result.finished_recving == result.failed_recving == {"req"}
+    assert worker.get_block_ids_with_load_errors() == {0, 1}
+    assert worker.xfer_stats.data["num_kv_expired_reqs"] == [1]
+
+
+@pytest.mark.parametrize("lease", [1, 4])
+def test_heartbeat_credit_never_outlives_producer_lease(pull_worker, lease):
+    """P extends by lease * 2 // 3 whole seconds; D's credited deadline must
+    still pass no later than P reaps the blocks."""
+    worker, clock = pull_worker
+    producer = _worker_stub()
+    producer._lease_extension = lease * 2 // 3
+    producer._reqs_to_send = {"prefill-req": float(lease)}
+    meta = ReqMeta(
+        local_block_ids=([0],),
+        local_physical_block_ids=([0],),
+        tp_size=1,
+        remote=RemoteMeta(
+            ([0],), "localhost", 1234, _REMOTE, "prefill-req", float(lease)
+        ),
+    )
+    for step in range(1, 80 * lease):
+        t = step / 20
+        if t <= 2 * lease and step % 2 == 0:
+            _heartbeat_at(worker, clock, t)
+            producer._handle_heartbeat("prefill-req")
+        clock["t"] = t
+        if t >= producer._reqs_to_send["prefill-req"]:
+            assert worker._is_lease_expired(meta), t

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -789,6 +789,7 @@ def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker:
     w._engine_ttl = engine_ttl
     w._engine_last_active = {}
     w._engine_clock_offset = {}
+    w._heartbeat_sends = {}
     w._handshake_lock = threading.RLock()
     # _cleanup_remote_engine touches these when reaping an engine.
     w.nixl_wrapper = MagicMock()

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -715,6 +715,11 @@ def __init__(
         self._invalid_engine_backoff: dict[EngineId, tuple[float, float]] = {}
         # Map of engine_id -> clock offset.
         self._engine_clock_offset: dict[EngineId, float] = {}
+        # Pull mode: (engine_id, remote request id) -> (first, last, max gap) of
+        # the perf_counter times heartbeats were sent to every agent of the engine.
+        self._heartbeat_sends: dict[
+            tuple[EngineId, ReqId], tuple[float, float, float]
+        ] = {}
 
         # Metadata.
         self.engine_id: EngineId = engine_id
@@ -3003,6 +3008,7 @@ def get_transfer_results(self) -> KVConnectorTransferResults:
         assert self.transfer_topo is not None
         done_sending = self._get_new_notifs()
         done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers)
+        self._fail_expired_recvs(done_recving, newly_failed)
         if newly_failed:
             self._recv_failures.update(newly_failed)
 
@@ -3049,6 +3055,10 @@ def get_transfer_results(self) -> KVConnectorTransferResults:
             # clean up metadata for completed requests
             meta = self._recving_metadata.pop(req_id, None)
             assert meta is not None, f"{req_id} not found in recving_metadata list"
+            if meta.remote is not None:
+                self._heartbeat_sends.pop(
+                    (meta.remote.engine_id, meta.remote.request_id), None
+                )
 
             # Skip KV sync and post-processing for failed requests
             if req_id in failed_recv_reqs:
@@ -3135,6 +3145,9 @@ def get_transfer_results(self) -> KVConnectorTransferResults:
             failed_recving=failed_recv_reqs,
         )
 
+    def _fail_expired_recvs(self, done_recving: set[str], failed: set[str]) -> None:
+        """Move finished receives whose remote KV lease expired into ``failed``."""
+
     def get_finished(self) -> tuple[set[str], set[str]]:
         """Compatibility wrapper for the legacy completion API."""
         results = self.get_transfer_results()
@@ -3322,6 +3335,13 @@ def _send_pending_recv_notifs(self, req_id: str) -> None:
 
     def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
         """Send heartbeats to remote engines, extending the lease on KV blocks."""
+        if metadata.heartbeat_by_engine:
+            self._heartbeat_sends = {
+                key: times
+                for key, times in self._heartbeat_sends.items()
+                if (hb := metadata.heartbeat_by_engine.get(key[0])) is not None
+                and key[1] in hb.req_ids
+            }
         for engine_id, hb_info in metadata.heartbeat_by_engine.items():
             # Proactive handshake (this request may still be in waiting queue) so
             # the **next** heartbeat for this remote can go through.
@@ -3344,15 +3364,24 @@ def _send_heartbeats(self, metadata: NixlConnectorMetadata) -> None:
             self._engine_last_active[engine_id] = time.perf_counter()
             # Build the heartbeat message: "HB:req1,req2,..."
             hb_msg = ("HB:" + ",".join(hb_info.req_ids)).encode()
+            all_sent = True
             for agent_name in self._remote_agents[engine_id].values():
                 try:
                     self.nixl_wrapper.send_notif(agent_name, notif_msg=hb_msg)
                 except Exception:
+                    all_sent = False
                     logger.debug(
                         "Failed to send heartbeat to engine %s",
                         engine_id,
                         exc_info=True,
                     )
+            if not all_sent or self._TRANSFER_MODE != "pull":
+                continue
+            now = time.perf_counter()
+            for req_id in hb_info.req_ids:
+                key = (engine_id, req_id)
+                first, last, max_gap = self._heartbeat_sends.get(key, (now, now, 0.0))
+                self._heartbeat_sends[key] = (first, now, max(max_gap, now - last))
 
     def get_mapped_blocks(
         self, block_ids: np.ndarray, block_size_ratio: int
@@ -3859,6 +3888,11 @@ def _cleanup_remote_engine(
 
         # Drop the cached clock offset; it is re-measured on the next handshake.
         self._engine_clock_offset.pop(engine_id, None)
+        self._heartbeat_sends = {
+            key: times
+            for key, times in self._heartbeat_sends.items()
+            if key[0] != engine_id
+        }
         # A just-completed handshake may not have recorded activity yet, so
         # tolerate a missing entry.
         last_active = self._engine_last_active.pop(engine_id, None)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py

@@ -264,9 +264,7 @@ def request_finished(
             self._reqs_need_send[request.request_id] = (
                 time.perf_counter() + request_kv_blocks_ttl
             )
-            if is_d_node:
-                # D blocks expiry time exported for the turn-2 readback.
-                blocks_expiry_time = self._reqs_need_send[request.request_id]
+            blocks_expiry_time = self._reqs_need_send[request.request_id]
             # NOTE HMA will "mark" empty/null blocks in groups with 0s (eg SWA ones),
             # trimming down after allocating for the whole sequence length. Empty
             # blocks are always at the start of the list.