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 f90de6297a7b..2fc4f73f4d31 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/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 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 e431beba691e..13b0fa108109 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/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 = () 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 d92857c34186..42a1854f1234 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/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() 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 4e86000bb11b..547a24e6e803 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 @@ -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) 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,