From e26f5b3bf1635ac263d971d26817d64676e30983 Mon Sep 17 00:00:00 2001 From: David Cheung Date: Tue, 13 Oct 2026 18:45:07 -0700 Subject: [PATCH] [Bugfix][NIXL] Fail pull reads whose producer lease has expired Problem (#61229): in pull mode D reads P's blocks without checking P's lease. Once P reaps an expired lease and reuses the blocks, D's READ still succeeds and D decodes another request's KV. Cause: only the bidirectional turn-2 readback was gated on a remote block expiry, and P exported remote_blocks_expiry_time only from D nodes. Fix: - P also exports remote_blocks_expiry_time in request_finished. - A pull D worker records, per (engine, remote request), the first and last time it sent a heartbeat to every agent of the engine and the largest gap between them. A record is dropped when the request's load finishes or fails on this worker, when a heartbeat set no longer lists it, and when its engine is evicted. - P extends a lease by kv_lease_duration * 2 // 3 per heartbeat and stamped the expiry E0 before D's first heartbeat, so (2 * (E0 - first) - 2) / 3 bounds P's extension from below, including the integer floor, without knowing P's config. If that bound is positive, the first heartbeat landed before E0 and no gap exceeds the bound, the deadline is max(E0, last + bound) - margin, with margin = min(5 s, bound / 4); otherwise it is E0 - 5 s. - D checks the deadline before issuing every pull READ (replacing the turn-2-only gate) and again when the READ completes, failing reads whose deadline passed in flight. Both use the KV-expired metric and the existing receive-failure path, so kv_load_failure_policy applies. For HMA no invalid blocks are reported (existing TODO), so recompute is not block-precise there. - A missing expiry skips the check. The check is conservative: a request whose first heartbeat leaves D within about 10 s of P's lease expiry gets no heartbeat credit and is failed if read after E0 - 5 s (kv_load_failure_policy applies; recompute avoids user-visible errors). Not covered: clock-offset error beyond the margin and heartbeats that were sent but lost still allow a stale read. Push mode is unchanged: there P drives the WRITEs, so it needs a check on the writer thread and a new wire message. Compatibility: no new kv_transfer_params key, no handshake change and no NIXL_CONNECTOR_VERSION bump. - New D with old P: P exports no expiry, so reads proceed as today. - Old D with new P: only an old D with bidirectional_kv_xfer enabled is affected. Its turn-2 gate now also gates turn-1 reads against E0 - 5 s without heartbeat credit, so upgrade D before P in that setup. Tests: worker tests on a fake clock for reads after E0 kept alive by heartbeats, failed heartbeat sends, a heartbeat gap longer than the extension, no heartbeats, a late first heartbeat, a lease expiring while the READ is in flight, small leases (1 s, 4 s) where the integer floor matters, and the heartbeat record dropped once the load finishes or fails; a P-node export test. The new expiry tests fail on the previous NIXL sources. Co-Authored-By: Claude Opus 5.5 Signed-off-by: David Cheung --- .../unit/test_bidirectional_kv_transfer.py | 18 ++ .../kv_connector/unit/test_nixl_connector.py | 2 + .../unit/test_nixl_connector_hma.py | 1 + .../kv_connector/unit/test_nixl_heartbeat.py | 174 +++++++++++++++++- .../unit/test_nixl_push_connector.py | 1 + .../kv_connector/v1/nixl/base_worker.py | 34 ++++ .../kv_connector/v1/nixl/pull_scheduler.py | 4 +- .../kv_connector/v1/nixl/pull_worker.py | 55 ++++-- 8 files changed, 273 insertions(+), 16 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py b/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py index 2aae7d71f515..969785036afe 100644 --- a/tests/v1/kv_connector/unit/test_bidirectional_kv_transfer.py +++ b/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.""" diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index dcbfa7eb0470..96a7be29aa12 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -2197,6 +2197,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 diff --git a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py index 469e0f5d8294..fb27b40732bf 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -56,6 +56,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 = () diff --git a/tests/v1/kv_connector/unit/test_nixl_heartbeat.py b/tests/v1/kv_connector/unit/test_nixl_heartbeat.py index f2d95daa4835..17f27d17f82b 100644 --- a/tests/v1/kv_connector/unit/test_nixl_heartbeat.py +++ b/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 diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index a36c7e5b2a71..462020c60163 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -788,6 +788,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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 778eb0f7560d..39509c5d3eb9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -655,6 +655,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 @@ -2859,6 +2864,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) @@ -2905,6 +2911,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: @@ -2991,6 +3001,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() @@ -3178,6 +3191,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. @@ -3200,15 +3220,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 @@ -3715,6 +3744,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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py index f283a40d3c2b..b49989655d42 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_scheduler.py +++ b/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. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py index 16a25149ea22..194879d5e799 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py @@ -28,8 +28,8 @@ logger = init_logger(__name__) -# Slack (seconds) subtracted from D's exported block-expiry deadline on the turn-2 -# readback, absorbing clock-offset error and read latency. +# Upper bound (seconds) on the slack subtracted from a remote block-expiry +# deadline, absorbing clock-offset error and read latency. _KV_BLOCKS_EXPIRY_SAFETY_MARGIN = 5.0 @@ -122,16 +122,49 @@ def start_load_kv(self, metadata: NixlConnectorMetadata): # requests sit in the D scheduler WAITING queue. self._send_heartbeats(metadata) - def _is_turn2_read_expired(self, meta: ReqMeta) -> bool: - """Whether D's cached blocks for this turn-2 readback have (nearly) expired.""" + def _is_lease_expired(self, meta: ReqMeta) -> bool: + """Whether the remote lease on this request's blocks has (nearly) expired. + + The exported expiry is extended by the heartbeats this worker sent, + each credited with a lower bound on the extension the remote grants. + """ assert meta.remote is not None - blocks_expiry_time = meta.remote.blocks_expiry_time - # Deadline may be absent (router may not forward it) -> read as usual. - if blocks_expiry_time is None or not meta.local_physical_block_ids: + remote = meta.remote + # Expiry may be absent (old remote, or router did not forward it). + if remote.blocks_expiry_time is None or not meta.local_physical_block_ids: + return False + clock_offset = self._engine_clock_offset.get(remote.engine_id) + if clock_offset is None: return False - clock_offset = self._engine_clock_offset[meta.remote.engine_id] - deadline = blocks_expiry_time - clock_offset - return time.perf_counter() + _KV_BLOCKS_EXPIRY_SAFETY_MARGIN >= deadline + expiry = remote.blocks_expiry_time - clock_offset + deadline = expiry - _KV_BLOCKS_EXPIRY_SAFETY_MARGIN + if sends := self._heartbeat_sends.get((remote.engine_id, remote.request_id)): + first, last, max_gap = sends + # The remote extends by kv_lease_duration * 2 // 3 per heartbeat and + # stamped the expiry before our first one, so this is a lower bound. + extension = (2 * (expiry - first) - 2) / 3 + margin = min(_KV_BLOCKS_EXPIRY_SAFETY_MARGIN, extension / 4) + if ( + extension > 0 + and first + margin < expiry + and max_gap + margin < extension + ): + deadline = max(expiry, last + extension) - margin + return time.perf_counter() >= deadline + + def _fail_expired_recvs(self, done_recving: set[str], failed: set[str]) -> None: + for req_id in done_recving - failed: + meta = self._recving_metadata.get(req_id) + if meta is None or not self._is_lease_expired(meta): + continue + logger.warning( + "Failing read for %s: remote KV lease expired during transfer.", + req_id, + ) + self.xfer_stats.record_kv_expired_req() + self._handle_failed_transfer( + req_id, None, failed, record_failed_transfer=False + ) def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): assert meta.remote is not None and self.transfer_topo is not None @@ -149,7 +182,7 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): # thread (this one), so we don't race on this structure. self._engine_last_active[engine_id] = time.perf_counter() - if self._bidirectional_kv_xfer_enabled and self._is_turn2_read_expired(meta): + if self._is_lease_expired(meta): logger.warning( "Declining expired remote read for %s from engine %s.", req_id,