@tduncan / Nvidia-Optimization-Stack-VLLM / commits / fff03267da

[Bugfix][NIXL] Release a dead peer's NIXL state without waiting for TTL (#50047)

Signed-off-by: Yannik Hinteregger <37209495+YannikHinteregger@users.noreply.github.com> Co-authored-by: xijiade.aihemaiti <3146335281@qq.com>

YannikHinteregger committed Sep 27, 2026 at 05:30 UTC fff03267dae61d03615225c2f48a06082d842e2d

5 files changed +237 -3

tests/v1/kv_connector/unit/test_nixl_connector.py

+175

@@ -7,6 +7,7 @@ import os

7 import queue

8 import tempfile

9 import textwrap

10 +import threading

11 import time

12 import uuid

13 from collections import defaultdict

@@ -49,6 +50,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (

50 NixlKVConnectorStats,

51 )

52 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (

53 + RemoteMeta,

54 + ReqMeta,

55 compute_nixl_compatibility_hash,

56 )

57 from vllm.distributed.kv_transfer.kv_transfer_state import (

@@ -2087,6 +2090,10 @@ def recv_worker():

2090 worker._recving_transfers = defaultdict(list)

2091 worker._failed_recv_reqs = queue.Queue()

2092 worker._recv_failures = set()

2093 + worker._handshake_lock = threading.RLock()

2094 + worker._handshake_futures = {}

2095 + worker._remote_agents = {}

2096 + worker._engine_by_address = {}

2097 worker._replicated_pcp_done_sending = set()

2098 worker._invalid_block_ids = queue.Queue()

2099 worker._pending_recv_notifs = {}

@@ -2681,6 +2688,174 @@ def test_engine_ttl_disabled(default_vllm_config, dist_init):

2688 assert engine_id in worker.dst_xfer_side_handles

2689

2690

2691 +@pytest.mark.cpu_test

2692 +class TestPeerReplacement:

2693 + @pytest.fixture(autouse=True)

2694 + def setup(self, monkeypatch):

2695 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw

2696 + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

2697 +

2698 + config = create_vllm_config(kv_connector_extra_config={"engine_ttl": 0})

2699 + config.kv_transfer_config.kv_buffer_device = "cpu"

2700 + platform = SimpleNamespace(

2701 + device_type="cpu",

2702 + discover_numa_topology=lambda: [],

2703 + get_nixl_memory_type=lambda: "DRAM",

2704 + is_rocm=lambda: False,

2705 + )

2706 + with (

2707 + patch.object(bw, "NixlWrapper", FakeNixlWrapper),

2708 + patch.object(bw, "current_platform", platform),

2709 + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0),

2710 + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1),

2711 + patch.object(

2712 + bw, "get_current_attn_backends", return_value=[FlashAttentionBackend]

2713 + ),

2714 + ):

2715 + self.worker = FakeNixlConnectorWorker(config, "local", hand_shake_latency=0)

2716 + self.nixl = self.worker.nixl_wrapper

2717 + self.transport = MagicMock(wraps=self.nixl)

2718 + self.worker.nixl_wrapper = self.transport

2719 + monkeypatch.setattr(

2720 + self.worker._handshake_initiation_executor,

2721 + "submit",

2722 + lambda *args: Future(),

2723 + )

2724 + self._connect("old")

2725 + yield

2726 + self.worker.shutdown()

2727 +

2728 + def _connect(self, engine_id="new", host="localhost", port=1234):

2729 + future = self.worker._ensure_handshake(engine_id, host, port, 1)

2730 + self.worker.REMOTE_ENGINE_ID = engine_id

2731 + self.nixl.REMOTE_AGENT_NAME = engine_id

2732 + future.set_result(self.worker._nixl_handshake(host, port, 1, engine_id))

2733 +

2734 + def _request(self, engine_id="new", blocks=(1,), awaiting_kvs=True):

2735 + metadata = NixlConnectorMetadata()

2736 + metadata.reqs_to_recv[f"{engine_id}-req"] = ReqMeta(

2737 + local_block_ids=(blocks,),

2738 + local_physical_block_ids=(blocks,),

2739 + tp_size=1,

2740 + remote=RemoteMeta(

2741 + ([2],), "localhost", 1234, engine_id, f"{engine_id}-prefill"

2742 + ),

2743 + awaiting_kvs=awaiting_kvs,

2744 + )

2745 + return metadata

2746 +

2747 + def test_cleans_old_peer_after_confirmed_handshake(self):

2748 + self._connect("healthy", host="other-host")

2749 + old_handle = self.worker.dst_xfer_side_handles["old"][0]

2750 + self.worker.start_load_kv(self._request())

2751 + self.worker.get_transfer_results()

2752 + self.transport.remove_remote_agent.assert_not_called()

2753 +

2754 + self._connect()

2755 + self.worker.get_transfer_results()

2756 + self.transport.release_dlist_handle.assert_called_once_with(old_handle)

2757 + self.transport.remove_remote_agent.assert_called_once_with("old")

2758 + assert set(self.worker._remote_agents) == {"new", "healthy"}

2759 + assert self.worker._engine_by_address == {

2760 + ("localhost", 1234): "new",

2761 + ("other-host", 1234): "healthy",

2762 + }

2763 + with pytest.raises(KeyError):

2764 + self.worker.transfer_topo.get_engine_info("old")

2765 +

2766 + self.worker.start_load_kv(NixlConnectorMetadata())

2767 + assert self.worker.get_transfer_results().finished_recving == {"new-req"}

2768 + self.transport.remove_remote_agent.assert_called_once_with("old")

2769 +

2770 + @pytest.mark.parametrize("host,port", [("other-host", 1234), ("localhost", 5678)])

2771 + def test_preserves_other_addresses(self, host, port):

2772 + self._connect(host=host, port=port)

2773 + self.worker.get_transfer_results()

2774 + assert "old" in self.worker._remote_agents

2775 + self.transport.remove_remote_agent.assert_not_called()

2776 +

2777 + def test_preserves_old_peer_when_handshake_fails(self):

2778 + self.worker.start_load_kv(self._request())

2779 + self.worker._handshake_futures["new"].set_exception(

2780 + RuntimeError("handshake failed")

2781 + )

2782 + assert self.worker.get_transfer_results().failed_recving == {"new-req"}

2783 + assert "old" in self.worker._remote_agents

2784 + assert self.worker._engine_by_address == {("localhost", 1234): "old"}

2785 + self.transport.remove_remote_agent.assert_not_called()

2786 +

2787 + @pytest.mark.parametrize(

2788 + "states,failed",

2789 + [

2790 + (["PROC", "DONE"], False),

2791 + (["ERR", "PROC", "DONE"], True),

2792 + (["ERR", "PROC", "ERR"], True),

2793 + ],

2794 + ids=["healthy", "failed-sibling-completes", "failed-sibling-fails"],

2795 + )

2796 + def test_waits_for_reads_before_cleanup_and_failure_reporting(self, states, failed):

2797 + self.worker._recving_metadata.update(self._request("old").reqs_to_recv)

2798 + self.worker._recving_transfers["old-req"] = [1, 2] if failed else [1]

2799 + self.transport.check_xfer_state.side_effect = states

2800 + self._connect()

2801 +

2802 + result = self.worker.get_transfer_results()

2803 + assert result.finished_recving == result.failed_recving == set()

2804 + assert self.worker.get_block_ids_with_load_errors() == set()

2805 + self.transport.remove_remote_agent.assert_not_called()

2806 +

2807 + result = self.worker.get_transfer_results()

2808 + assert result.finished_recving == {"old-req"}

2809 + assert result.failed_recving == ({"old-req"} if failed else set())

2810 + assert self.worker.get_block_ids_with_load_errors() == (

2811 + {1} if failed else set()

2812 + )

2813 + assert self.transport.release_xfer_handle.call_count == (2 if failed else 1)

2814 + self.transport.remove_remote_agent.assert_called_once_with("old")

2815 +

2816 + def test_waits_for_queued_notification(self):

2817 + metadata = self._request("old", blocks=())

2818 + self.worker._recving_metadata.update(metadata.reqs_to_recv)

2819 + self.worker._background_nixl_handshake(

2820 + "old-req", "old", metadata.reqs_to_recv["old-req"]

2821 + )

2822 + self._connect()

2823 + self.worker.get_transfer_results()

2824 + self.transport.remove_remote_agent.assert_not_called()

2825 +

2826 + self.worker.start_load_kv(NixlConnectorMetadata())

2827 + self.transport.send_notif.assert_called_once_with(

2828 + "old", notif_msg=b"old-prefill:1"

2829 + )

2830 + result = self.worker.get_transfer_results()

2831 + assert result.finished_recving == {"old-req"}

2832 + assert result.failed_recving == set()

2833 + self.transport.remove_remote_agent.assert_called_once_with("old")

2834 +

2835 + def test_waits_for_other_handshakes(self):

2836 + self.worker._ensure_handshake("other", "other-host", 1234, 1)

2837 + self._connect()

2838 + self.worker.get_transfer_results()

2839 + self.transport.remove_remote_agent.assert_not_called()

2840 +

2841 + self._connect("other", host="other-host")

2842 + self.worker.get_transfer_results()

2843 + self.transport.remove_remote_agent.assert_called_once_with("old")

2844 +

2845 + def test_rehandshakes_queued_request_for_released_engine(self):

2846 + # A request queued after its engine's handshake, but drained only after

2847 + # that engine was replaced and released, must not read from it.

2848 + metadata = self._request("old", blocks=(), awaiting_kvs=False)

2849 + self.worker._ready_requests.put(("old-req", metadata.reqs_to_recv["old-req"]))

2850 + self._connect()

2851 + self.worker.get_transfer_results()

2852 + self.transport.remove_remote_agent.assert_called_once_with("old")

2853 +

2854 + self.worker.start_load_kv(NixlConnectorMetadata())

2855 + assert "old" in self.worker._handshake_futures

2856 + self.transport.send_notif.assert_not_called()

2857 +

2858 +

2859 def test_transfer_topology_unregister():

2860 """TransferTopology.unregister_remote_engine removes the engine."""

2861 from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -4,6 +4,7 @@

5 import gc

6 import queue

7 +import threading

8 from types import SimpleNamespace

9 from unittest.mock import MagicMock, patch

@@ -54,6 +55,9 @@ def region_pull_worker():

55 worker.enable_permute_local_kv = False

56 worker.enable_heterogeneous_attn_post_process = False

57 worker._engine_last_active = {}

58 + worker._handshake_lock = threading.RLock()

59 + worker._handshake_futures = {}

60 + worker._transfer_layer_group_ids = ()

61 worker._bidirectional_kv_xfer_enabled = False

62 worker._recving_transfers = {}

63 worker.use_mla, worker._has_mamba = True, False

@@ -94,6 +98,7 @@ def region_pull_worker():

98 worker.dst_num_blocks = {"P": 100, "D": 100}

99 worker.dst_region_num_blocks = {"P": [100, 100], "D": [100, 100]}

100 worker._remote_agents = {"P": {(0, 0): "P-rank0"}}

101 + worker._engine_by_address = {("localhost", 1): "P"}

102 worker._read_blocks_mixed = MagicMock()

103 worker.nixl_wrapper = MagicMock()

104 return worker

tests/v1/kv_connector/unit/test_nixl_push_connector.py

@@ -395,6 +395,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

395 w._group_spec_types = (FullAttentionSpec,)

396 w._engine_ttl = 0.0

397 w._engine_last_active = {}

398 + w._engine_by_address = {}

399

400 # Track _do_start_push_kv invocations.

401 calls: list[tuple[str, Any, dict[str, Any]]] = []

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

+48 -2

@@ -839,6 +839,8 @@ class NixlBaseConnectorWorker:

839 self._engine_ttl: float = vllm_config.kv_transfer_config.get_from_extra_config(

840 "engine_ttl", 3600.0

841 )

842 + # Map of pull peer (host, port) -> engine_id currently serving it.

843 + self._engine_by_address: dict[tuple[str, int], EngineId] = {}

844

845 self.model_config = vllm_config.model_config

846

@@ -1300,6 +1302,7 @@ class NixlBaseConnectorWorker:

1302 self._remote_agents[eid] = remote_agents

1303 self._engine_clock_offset[eid] = clock_offset

1304 self._engine_last_active[eid] = time.perf_counter()

1305 + self._track_remote_engine_replacement(eid, host, port)

1306 except Exception as e:

1307 self._log_failure(

1308 failure_type="handshake_setup_failed",

@@ -2953,6 +2956,7 @@ class NixlBaseConnectorWorker:

2956 # Handle timeout to avoid stranding blocks on remote.

2957 self._reap_expired_send_leases(done_sending)

2958

2959 + self._cleanup_replaced_remote_engines()

2960 return KVConnectorTransferResults(

2961 finished_sending=done_sending,

2962 finished_recving=done_recving,

@@ -3541,14 +3545,53 @@ class NixlBaseConnectorWorker:

3545 and meta.remote is not None

3546 }

3547

3548 + def _track_remote_engine_replacement(

3549 + self, engine_id: EngineId, host: str, port: int

3550 + ) -> None:

3551 + """Record the engine currently serving a pull peer address."""

3552 + # TODO: Also handle push mode, which handshakes in both directions.

3553 + if self._TRANSFER_MODE != "pull":

3554 + return

3555 + previous = self._engine_by_address.get((host, port))

3556 + self._engine_by_address[(host, port)] = engine_id

3557 + if previous is None or previous == engine_id:

3558 + return

3559 + logger.info(

3560 + "Remote engine at %s:%d changed from %s to %s.",

3561 + host,

3562 + port,

3563 + previous,

3564 + engine_id,

3565 + )

3566 +

3567 + def _cleanup_replaced_remote_engines(self) -> None:

3568 + """Release replaced pull peers once requests and handshakes have drained."""

3569 + if self._TRANSFER_MODE != "pull":

3570 + return

3571 + with self._handshake_lock:

3572 + if self._handshake_futures:

3573 + return

3574 + replaced = self._remote_agents.keys() - self._engine_by_address.values()

3575 + if not replaced:

3576 + return

3577 + # _recving_metadata is only accessed from this (main) thread.

3578 + busy = {

3579 + meta.remote.engine_id

3580 + for meta in self._recving_metadata.values()

3581 + if meta.remote is not None

3582 + }

3583 + for engine_id in replaced - busy:

3584 + self._cleanup_remote_engine(engine_id, log_eviction=False)

3585 + logger.info("Released NIXL state for replaced remote engine %s.", engine_id)

3586 +

3587 def _cleanup_remote_engine(

3588 self, engine_id: EngineId, *, log_eviction: bool = True

3589 ) -> None:

3590 """Remove all state for a single remote engine.

3591

3592 Releases NIXL resources (dlist handles, remote agents) and clears

3550 - all per-engine data structures. Used by both TTL eviction and

3551 - shutdown.

3593 + all per-engine data structures. Used by TTL eviction, replaced peer

3594 + cleanup and shutdown.

3595 """

3596 assert engine_id in self._remote_agents

3597

@@ -3558,6 +3601,9 @@ class NixlBaseConnectorWorker:

3601 # Pop under the handshake lock; NIXL teardown stays outside it.

3602 with self._handshake_lock:

3603 agents = self._remote_agents.pop(engine_id)

3604 + for address, eid in list(self._engine_by_address.items()):

3605 + if eid == engine_id:

3606 + del self._engine_by_address[address]

3607 for agent_name in agents.values():

3608 self.nixl_wrapper.remove_remote_agent(agent_name)

3609

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

+8 -1

@@ -84,7 +84,14 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

85 # Start transfers for requests whose handshakes have now finished.

86 while not self._ready_requests.empty():

87 - self._read_blocks_for_req(*self._ready_requests.get_nowait())

87 + req_id, meta = self._ready_requests.get_nowait()

88 + assert meta.remote is not None

89 + if meta.remote.engine_id not in self._remote_agents:

90 + # The engine was released after its handshake completed, so

91 + # handshake again. This fails if the engine is gone.

92 + self._background_nixl_handshake(req_id, meta.remote.engine_id, meta)

93 + continue

94 + self._read_blocks_for_req(req_id, meta)

96 if self.pcp_rank > 0 and not self.pcp_dcp_sharded:

97 # Replicated-KV PCP: only PCP rank 0 serves the KV, so this rank