diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 39f8f88d5c..d8b6fcb5a0 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -7,6 +7,7 @@ import os import queue import tempfile import textwrap +import threading import time import uuid from collections import defaultdict @@ -49,6 +50,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import ( NixlKVConnectorStats, ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + RemoteMeta, + ReqMeta, compute_nixl_compatibility_hash, ) from vllm.distributed.kv_transfer.kv_transfer_state import ( @@ -2087,6 +2090,10 @@ def recv_worker(): worker._recving_transfers = defaultdict(list) worker._failed_recv_reqs = queue.Queue() worker._recv_failures = set() + worker._handshake_lock = threading.RLock() + worker._handshake_futures = {} + worker._remote_agents = {} + worker._engine_by_address = {} worker._replicated_pcp_done_sending = set() worker._invalid_block_ids = queue.Queue() worker._pending_recv_notifs = {} @@ -2681,6 +2688,174 @@ def test_engine_ttl_disabled(default_vllm_config, dist_init): assert engine_id in worker.dst_xfer_side_handles +@pytest.mark.cpu_test +class TestPeerReplacement: + @pytest.fixture(autouse=True) + def setup(self, monkeypatch): + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + + config = create_vllm_config(kv_connector_extra_config={"engine_ttl": 0}) + config.kv_transfer_config.kv_buffer_device = "cpu" + platform = SimpleNamespace( + device_type="cpu", + discover_numa_topology=lambda: [], + get_nixl_memory_type=lambda: "DRAM", + is_rocm=lambda: False, + ) + with ( + patch.object(bw, "NixlWrapper", FakeNixlWrapper), + patch.object(bw, "current_platform", platform), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object( + bw, "get_current_attn_backends", return_value=[FlashAttentionBackend] + ), + ): + self.worker = FakeNixlConnectorWorker(config, "local", hand_shake_latency=0) + self.nixl = self.worker.nixl_wrapper + self.transport = MagicMock(wraps=self.nixl) + self.worker.nixl_wrapper = self.transport + monkeypatch.setattr( + self.worker._handshake_initiation_executor, + "submit", + lambda *args: Future(), + ) + self._connect("old") + yield + self.worker.shutdown() + + def _connect(self, engine_id="new", host="localhost", port=1234): + future = self.worker._ensure_handshake(engine_id, host, port, 1) + self.worker.REMOTE_ENGINE_ID = engine_id + self.nixl.REMOTE_AGENT_NAME = engine_id + future.set_result(self.worker._nixl_handshake(host, port, 1, engine_id)) + + def _request(self, engine_id="new", blocks=(1,), awaiting_kvs=True): + metadata = NixlConnectorMetadata() + metadata.reqs_to_recv[f"{engine_id}-req"] = ReqMeta( + local_block_ids=(blocks,), + local_physical_block_ids=(blocks,), + tp_size=1, + remote=RemoteMeta( + ([2],), "localhost", 1234, engine_id, f"{engine_id}-prefill" + ), + awaiting_kvs=awaiting_kvs, + ) + return metadata + + def test_cleans_old_peer_after_confirmed_handshake(self): + self._connect("healthy", host="other-host") + old_handle = self.worker.dst_xfer_side_handles["old"][0] + self.worker.start_load_kv(self._request()) + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_not_called() + + self._connect() + self.worker.get_transfer_results() + self.transport.release_dlist_handle.assert_called_once_with(old_handle) + self.transport.remove_remote_agent.assert_called_once_with("old") + assert set(self.worker._remote_agents) == {"new", "healthy"} + assert self.worker._engine_by_address == { + ("localhost", 1234): "new", + ("other-host", 1234): "healthy", + } + with pytest.raises(KeyError): + self.worker.transfer_topo.get_engine_info("old") + + self.worker.start_load_kv(NixlConnectorMetadata()) + assert self.worker.get_transfer_results().finished_recving == {"new-req"} + self.transport.remove_remote_agent.assert_called_once_with("old") + + @pytest.mark.parametrize("host,port", [("other-host", 1234), ("localhost", 5678)]) + def test_preserves_other_addresses(self, host, port): + self._connect(host=host, port=port) + self.worker.get_transfer_results() + assert "old" in self.worker._remote_agents + self.transport.remove_remote_agent.assert_not_called() + + def test_preserves_old_peer_when_handshake_fails(self): + self.worker.start_load_kv(self._request()) + self.worker._handshake_futures["new"].set_exception( + RuntimeError("handshake failed") + ) + assert self.worker.get_transfer_results().failed_recving == {"new-req"} + assert "old" in self.worker._remote_agents + assert self.worker._engine_by_address == {("localhost", 1234): "old"} + self.transport.remove_remote_agent.assert_not_called() + + @pytest.mark.parametrize( + "states,failed", + [ + (["PROC", "DONE"], False), + (["ERR", "PROC", "DONE"], True), + (["ERR", "PROC", "ERR"], True), + ], + ids=["healthy", "failed-sibling-completes", "failed-sibling-fails"], + ) + def test_waits_for_reads_before_cleanup_and_failure_reporting(self, states, failed): + self.worker._recving_metadata.update(self._request("old").reqs_to_recv) + self.worker._recving_transfers["old-req"] = [1, 2] if failed else [1] + self.transport.check_xfer_state.side_effect = states + self._connect() + + result = self.worker.get_transfer_results() + assert result.finished_recving == result.failed_recving == set() + assert self.worker.get_block_ids_with_load_errors() == set() + self.transport.remove_remote_agent.assert_not_called() + + result = self.worker.get_transfer_results() + assert result.finished_recving == {"old-req"} + assert result.failed_recving == ({"old-req"} if failed else set()) + assert self.worker.get_block_ids_with_load_errors() == ( + {1} if failed else set() + ) + assert self.transport.release_xfer_handle.call_count == (2 if failed else 1) + self.transport.remove_remote_agent.assert_called_once_with("old") + + def test_waits_for_queued_notification(self): + metadata = self._request("old", blocks=()) + self.worker._recving_metadata.update(metadata.reqs_to_recv) + self.worker._background_nixl_handshake( + "old-req", "old", metadata.reqs_to_recv["old-req"] + ) + self._connect() + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_not_called() + + self.worker.start_load_kv(NixlConnectorMetadata()) + self.transport.send_notif.assert_called_once_with( + "old", notif_msg=b"old-prefill:1" + ) + result = self.worker.get_transfer_results() + assert result.finished_recving == {"old-req"} + assert result.failed_recving == set() + self.transport.remove_remote_agent.assert_called_once_with("old") + + def test_waits_for_other_handshakes(self): + self.worker._ensure_handshake("other", "other-host", 1234, 1) + self._connect() + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_not_called() + + self._connect("other", host="other-host") + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_called_once_with("old") + + def test_rehandshakes_queued_request_for_released_engine(self): + # A request queued after its engine's handshake, but drained only after + # that engine was replaced and released, must not read from it. + metadata = self._request("old", blocks=(), awaiting_kvs=False) + self.worker._ready_requests.put(("old-req", metadata.reqs_to_recv["old-req"])) + self._connect() + self.worker.get_transfer_results() + self.transport.remove_remote_agent.assert_called_once_with("old") + + self.worker.start_load_kv(NixlConnectorMetadata()) + assert "old" in self.worker._handshake_futures + self.transport.send_notif.assert_not_called() + + def test_transfer_topology_unregister(): """TransferTopology.unregister_remote_engine removes the engine.""" from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend 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 3be1730ec3..125dfa6143 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector_hma.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector_hma.py @@ -4,6 +4,7 @@ import gc import queue +import threading from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -54,6 +55,9 @@ def region_pull_worker(): worker.enable_permute_local_kv = False worker.enable_heterogeneous_attn_post_process = False worker._engine_last_active = {} + worker._handshake_lock = threading.RLock() + worker._handshake_futures = {} + worker._transfer_layer_group_ids = () worker._bidirectional_kv_xfer_enabled = False worker._recving_transfers = {} worker.use_mla, worker._has_mamba = True, False @@ -94,6 +98,7 @@ def region_pull_worker(): worker.dst_num_blocks = {"P": 100, "D": 100} worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]} worker._remote_agents = {"P": {(0, 0): "P-rank0"}} + worker._engine_by_address = {("localhost", 1): "P"} worker._read_blocks_mixed = MagicMock() worker.nixl_wrapper = MagicMock() return worker 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 231343de37..4d8ead28f3 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -395,6 +395,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): w._group_spec_types = (FullAttentionSpec,) w._engine_ttl = 0.0 w._engine_last_active = {} + w._engine_by_address = {} # Track _do_start_push_kv invocations. calls: list[tuple[str, Any, dict[str, Any]]] = [] 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 b630dedfc3..f4a643d2b7 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 @@ -839,6 +839,8 @@ class NixlBaseConnectorWorker: self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config( "engine_ttl", 3600.0 ) + # Map of pull peer (host, port) -> engine_id currently serving it. + self._engine_by_address: dict[tuple[str, int], EngineId] = {} self.model_config = vllm_config.model_config @@ -1300,6 +1302,7 @@ class NixlBaseConnectorWorker: self._remote_agents[eid] = remote_agents self._engine_clock_offset[eid] = clock_offset self._engine_last_active[eid] = time.perf_counter() + self._track_remote_engine_replacement(eid, host, port) except Exception as e: self._log_failure( failure_type="handshake_setup_failed", @@ -2953,6 +2956,7 @@ class NixlBaseConnectorWorker: # Handle timeout to avoid stranding blocks on remote. self._reap_expired_send_leases(done_sending) + self._cleanup_replaced_remote_engines() return KVConnectorTransferResults( finished_sending=done_sending, finished_recving=done_recving, @@ -3541,14 +3545,53 @@ class NixlBaseConnectorWorker: and meta.remote is not None } + def _track_remote_engine_replacement( + self, engine_id: EngineId, host: str, port: int + ) -> None: + """Record the engine currently serving a pull peer address.""" + # TODO: Also handle push mode, which handshakes in both directions. + if self._TRANSFER_MODE != "pull": + return + previous = self._engine_by_address.get((host, port)) + self._engine_by_address[(host, port)] = engine_id + if previous is None or previous == engine_id: + return + logger.info( + "Remote engine at %s:%d changed from %s to %s.", + host, + port, + previous, + engine_id, + ) + + def _cleanup_replaced_remote_engines(self) -> None: + """Release replaced pull peers once requests and handshakes have drained.""" + if self._TRANSFER_MODE != "pull": + return + with self._handshake_lock: + if self._handshake_futures: + return + replaced = self._remote_agents.keys() - self._engine_by_address.values() + if not replaced: + return + # _recving_metadata is only accessed from this (main) thread. + busy = { + meta.remote.engine_id + for meta in self._recving_metadata.values() + if meta.remote is not None + } + for engine_id in replaced - busy: + self._cleanup_remote_engine(engine_id, log_eviction=False) + logger.info("Released NIXL state for replaced remote engine %s.", engine_id) + def _cleanup_remote_engine( self, engine_id: EngineId, *, log_eviction: bool = True ) -> None: """Remove all state for a single remote engine. Releases NIXL resources (dlist handles, remote agents) and clears - all per-engine data structures. Used by both TTL eviction and - shutdown. + all per-engine data structures. Used by TTL eviction, replaced peer + cleanup and shutdown. """ assert engine_id in self._remote_agents @@ -3558,6 +3601,9 @@ class NixlBaseConnectorWorker: # Pop under the handshake lock; NIXL teardown stays outside it. with self._handshake_lock: agents = self._remote_agents.pop(engine_id) + for address, eid in list(self._engine_by_address.items()): + if eid == engine_id: + del self._engine_by_address[address] for agent_name in agents.values(): self.nixl_wrapper.remove_remote_agent(agent_name) 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 6539b7e57e..11b7215950 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 @@ -84,7 +84,14 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker): # Start transfers for requests whose handshakes have now finished. while not self._ready_requests.empty(): - self._read_blocks_for_req(*self._ready_requests.get_nowait()) + req_id, meta = self._ready_requests.get_nowait() + assert meta.remote is not None + if meta.remote.engine_id not in self._remote_agents: + # The engine was released after its handshake completed, so + # handshake again. This fails if the engine is gone. + self._background_nixl_handshake(req_id, meta.remote.engine_id, meta) + continue + self._read_blocks_for_req(req_id, meta) if self.pcp_rank > 0 and not self.pcp_dcp_sharded: # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank