@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4ca13a8667

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[NIXL] Terminate handshake listener thread in shutdown (#26404)

Signed-off-by: Mark McLoughlin <markmc@redhat.com>

Mark McLoughlin committed Oct 22, 2025 at 15:59 UTC 4ca13a8667c7c177f99ed15ed3f6996dbb081b8a

2 files changed +52 -11

tests/v1/kv_connector/unit/test_nixl_connector.py

+22 -7

@@ -938,6 +938,13 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):

938 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

939 monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))

940

941 + def run_test_and_cleanup():

942 + llm = LLM(**llm_kwargs)

943 + try:

944 + _run_abort_timeout_test(llm, timeout)

945 + finally:

946 + llm.llm_engine.engine_core.shutdown()

947 +

948 # Build runtime_env only if we're using Ray

949 if distributed_executor_backend == "ray":

950 with _make_fake_nixl_pkg() as working_dir:

@@ -950,15 +957,16 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):

957 },

958 }

959 ray.init(runtime_env=runtime_env)

953 -

954 - _run_abort_timeout_test(llm_kwargs, timeout)

960 + try:

961 + run_test_and_cleanup()

962 + finally:

963 + ray.shutdown()

964 else:

956 - _run_abort_timeout_test(llm_kwargs, timeout)

965 + run_test_and_cleanup()

966

967

959 -def _run_abort_timeout_test(llm_kwargs: dict, timeout: int):

968 +def _run_abort_timeout_test(llm: LLM, timeout: int):

969 """Helper function to run the abort timeout test logic."""

961 - llm = LLM(**llm_kwargs)

970 remote_prefill_opts = {

971 "do_remote_decode": True,

972 "do_remote_prefill": False,

@@ -1042,7 +1050,7 @@ def test_register_kv_caches(dist_init):

1050 ),

1051 patch(

1052 "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Thread"

1045 - ),

1053 + ) as mock_thread,

1054 ): # noqa: E501

1055 # Create connector

1056 connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

@@ -1054,6 +1062,9 @@ def test_register_kv_caches(dist_init):

1062 mock_wrapper_instance = mock_nixl_wrapper.return_value

1063 connector.connector_worker.nixl_wrapper = mock_wrapper_instance

1064

1065 + # Reassure the shutdown() check that the thread is terminated

1066 + mock_thread.return_value.is_alive.return_value = False

1067 +

1068 # Execute register_kv_caches

1069 connector.register_kv_caches(kv_caches)

1070

@@ -1171,6 +1182,7 @@ def test_shutdown_cleans_up_resources(dist_init):

1182 with (

1183 patch.object(worker, "_handshake_initiation_executor") as mock_exec,

1184 patch.object(worker, "_nixl_handshake_listener_t") as mock_listener,

1185 + patch.object(worker, "_nixl_handshake_listener_stop_event") as mock_event,

1186 patch.object(nixl_wrapper, "release_xfer_handle") as mock_rel_xfer,

1187 patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel_dlist,

1188 patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem_agent,

@@ -1182,6 +1194,8 @@ def test_shutdown_cleans_up_resources(dist_init):

1194 worker._remote_agents = {"engine1": {0: "agent1"}}

1195 worker._registered_descs = ["desc1", "desc2"]

1196

1197 + mock_listener.is_alive.return_value = False

1198 +

1199 worker.shutdown()

1200

1201 # Test idempotency

@@ -1189,7 +1203,8 @@ def test_shutdown_cleans_up_resources(dist_init):

1203 worker.shutdown()

1204

1205 mock_exec.shutdown.assert_called_with(wait=False)

1192 - mock_listener.join.assert_called_once_with(timeout=0)

1206 + mock_event.set.assert_called_once()

1207 + mock_listener.join.assert_called_once_with(timeout=1.0)

1208

1209 mock_rel_xfer.assert_called_once_with(123)

1210 assert mock_rel_dlist.call_count == 2

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+30 -4

@@ -520,6 +520,8 @@ class NixlConnectorScheduler:

520 class NixlConnectorWorker:

521 """Implementation of Worker side methods"""

522

523 + _POLL_TIMEOUT = 0.1 # Handshake thread polls for stop event every 100ms

524 +

525 @dataclass

526 class TpKVTopology:

527 """

@@ -719,6 +721,7 @@ class NixlConnectorWorker:

721

722 # Background thread for handling new handshake requests.

723 self._nixl_handshake_listener_t: threading.Thread | None = None

724 + self._nixl_handshake_listener_stop_event: threading.Event | None = None

725 # Background thread for initializing new NIXL handshakes.

726 self._handshake_initiation_executor = ThreadPoolExecutor(

727 # NIXL is not guaranteed to be thread-safe, limit 1 worker.

@@ -773,6 +776,7 @@ class NixlConnectorWorker:

776 def _nixl_handshake_listener(

777 metadata: NixlAgentMetadata,

778 ready_event: threading.Event,

779 + stop_event: threading.Event,

780 base_port: int,

781 tp_rank: int,

782 ):

@@ -791,7 +795,14 @@ class NixlConnectorWorker:

795 logger.debug("Starting listening on path: %s", path)

796 with zmq_ctx(zmq.ROUTER, path) as sock:

797 ready_event.set()

794 - while True:

798 + poller = zmq.Poller()

799 + poller.register(sock, zmq.POLLIN)

800 + while not stop_event.is_set():

801 + events = dict(

802 + poller.poll(timeout=NixlConnectorWorker._POLL_TIMEOUT * 1000)

803 + )

804 + if sock not in events:

805 + continue

806 identity, _, msg = sock.recv_multipart()

807 if msg != GET_META_MSG:

808 logger.warning("Connection listener got unexpected message %s", msg)

@@ -1101,14 +1112,21 @@ class NixlConnectorWorker:

1112 attn_backend_name=self.backend_name,

1113 kv_cache_layout=self.kv_cache_layout,

1114 )

1104 - ready_event = threading.Event()

1115 + ready_event, stop_event = threading.Event(), threading.Event()

1116 self._nixl_handshake_listener_t = threading.Thread(

1117 target=self._nixl_handshake_listener,

1107 - args=(metadata, ready_event, self.side_channel_port, self.tp_rank),

1118 + args=(

1119 + metadata,

1120 + ready_event,

1121 + stop_event,

1122 + self.side_channel_port,

1123 + self.tp_rank,

1124 + ),

1125 daemon=True,

1126 name="nixl_handshake_listener",

1127 )

1128 self._nixl_handshake_listener_t.start()

1129 + self._nixl_handshake_listener_stop_event = stop_event

1130 ready_event.wait() # Wait for listener ZMQ socket to be ready.

1131

1132 def add_remote_agent(

@@ -1782,11 +1800,19 @@ class NixlConnectorWorker:

1800 self._invalid_block_ids = set()

1801 return result

1802

1803 + def __del__(self):

1804 + self.shutdown()

1805 +

1806 def shutdown(self):

1807 """Shutdown the connector worker."""

1808 self._handshake_initiation_executor.shutdown(wait=False)

1809 + if self._nixl_handshake_listener_stop_event is not None:

1810 + self._nixl_handshake_listener_stop_event.set()

1811 + self._nixl_handshake_listener_stop_event = None

1812 if self._nixl_handshake_listener_t is not None:

1789 - self._nixl_handshake_listener_t.join(timeout=0)

1813 + # Generous timeout to allow the thread to exit

1814 + self._nixl_handshake_listener_t.join(timeout=self._POLL_TIMEOUT * 10)

1815 + assert not self._nixl_handshake_listener_t.is_alive()

1816 self._nixl_handshake_listener_t = None

1817 for handles in self._recving_transfers.values():

1818 for handle, _ in handles: