diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 1f9cfc7b251d..866568eca91c 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -6,6 +6,7 @@ import os import tempfile import textwrap +import threading import time import uuid from collections import defaultdict @@ -18,6 +19,7 @@ import pytest import ray import torch +import zmq from tests.v1.attention.utils import dense_kv_cache_views from vllm import LLM @@ -44,6 +46,7 @@ NixlKVConnectorStats, ) from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import ( + GET_META_MSG, compute_nixl_compatibility_hash, ) from vllm.distributed.kv_transfer.kv_transfer_state import ( @@ -3382,6 +3385,95 @@ def test_mla_broadcast_notif_uses_remote_request_id( ) +def _start_handshake_listener(port: int) -> tuple[threading.Thread, threading.Event]: + encoded_data = {(0, 0): b"fake-metadata-blob"} + ready_event = threading.Event() + stop_event = threading.Event() + thread = threading.Thread( + target=NixlConnectorScheduler._nixl_handshake_listener, + args=(encoded_data, ready_event, stop_event, "127.0.0.1", port), + daemon=True, + ) + thread.start() + assert ready_event.wait(timeout=5), "listener did not start" + return thread, stop_event + + +def _ask_handshake_listener(port: int, timeout_ms: int = 2000): + """Send a real GET_META_MSG request, like a decode worker would.""" + ctx = zmq.Context() + sock = ctx.socket(zmq.REQ) + sock.setsockopt(zmq.RCVTIMEO, timeout_ms) + sock.setsockopt(zmq.LINGER, 0) + sock.connect(f"tcp://127.0.0.1:{port}") + sock.send(msgspec.msgpack.encode((GET_META_MSG, 0, 0))) + try: + reply = sock.recv_multipart() + except zmq.Again: + reply = None + sock.close() + ctx.term() + return reply + + +def test_handshake_listener_survives_bad_envelope(): + """A peer that doesn't speak the REQ/ROUTER envelope (e.g. a bare + DEALER socket, which doesn't auto-add the empty delimiter frame the + way REQ does) must not kill the listener thread. Reproduces the + production crash: ValueError: not enough values to unpack.""" + port = 15700 + thread, stop_event = _start_handshake_listener(port) + try: + # Sanity check: the listener is healthy before the bad message. + reply = _ask_handshake_listener(port) + assert reply is not None and reply[0] == b"fake-metadata-blob" + + ctx = zmq.Context() + stray = ctx.socket(zmq.DEALER) + stray.connect(f"tcp://127.0.0.1:{port}") + stray.send(b"whatever") + time.sleep(0.3) + stray.setsockopt(zmq.LINGER, 0) + stray.close() + ctx.term() + time.sleep(0.3) + + assert thread.is_alive(), "listener died on a bad envelope" + + reply = _ask_handshake_listener(port) + assert reply is not None and reply[0] == b"fake-metadata-blob" + finally: + stop_event.set() + thread.join(timeout=3) + + +def test_handshake_listener_survives_malformed_content(): + """A well-formed REQ envelope carrying malformed msgpack content must + also not kill the listener thread.""" + port = 15701 + thread, stop_event = _start_handshake_listener(port) + try: + ctx = zmq.Context() + sock = ctx.socket(zmq.REQ) + sock.setsockopt(zmq.RCVTIMEO, 1000) + sock.setsockopt(zmq.LINGER, 0) + sock.connect(f"tcp://127.0.0.1:{port}") + sock.send(b"not valid msgpack") + with contextlib.suppress(zmq.Again): + sock.recv_multipart() + sock.close() + ctx.term() + time.sleep(0.3) + + assert thread.is_alive(), "listener died on malformed content" + + reply = _ask_handshake_listener(port) + assert reply is not None and reply[0] == b"fake-metadata-blob" + finally: + stop_event.set() + thread.join(timeout=3) + + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py index ddfcf35e033b..2ad11eedfd01 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py @@ -356,28 +356,65 @@ def _nixl_handshake_listener( ready_event.set() while True: try: - identity, _, msg = sock.recv_multipart() + # Left unpacked here: a peer that isn't a well-behaved + # REQ socket (e.g. a bare DEALER) can send a different + # number of frames, and unpacking that must not be + # allowed to kill this thread either. + frames = sock.recv_multipart() except zmq.Again: if stop_event.is_set(): break continue - # Decode (GET_META_MSG, pp_rank, tp_rank). - msg, target_pp_rank, target_tp_rank = msgspec.msgpack.decode(msg) - logger.debug( - "Received message for pp rank %s, tp rank %s", - target_pp_rank, - target_tp_rank, - ) - if msg != GET_META_MSG: - logger.warning("Connection listener got unexpected message %s", msg) - # Echo our perf_counter so P can estimate the clock offset. - # perf_counter is only comparable within a process, so this - # listener must run in the same process that stamps the block - # expiry deadline (`_reqs_need_send`). - ts = msgspec.msgpack.encode(time.perf_counter()) - sock.send_multipart( - (identity, b"", encoded_data[(target_pp_rank, target_tp_rank)], ts) - ) + try: + identity, _, msg = frames + # Decode (GET_META_MSG, pp_rank, tp_rank). + decoded_msg, target_pp_rank, target_tp_rank = ( + msgspec.msgpack.decode(msg) + ) + logger.debug( + "Received message for pp rank %s, tp rank %s", + target_pp_rank, + target_tp_rank, + ) + if decoded_msg != GET_META_MSG: + logger.warning( + "Connection listener got unexpected message %s", + decoded_msg, + ) + # Echo our perf_counter so P can estimate the clock offset. + # perf_counter is only comparable within a process, so this + # listener must run in the same process that stamps the + # block expiry deadline (`_reqs_need_send`). + ts = msgspec.msgpack.encode(time.perf_counter()) + sock.send_multipart( + ( + identity, + b"", + encoded_data[(target_pp_rank, target_tp_rank)], + ts, + ) + ) + except ( + msgspec.DecodeError, + ValueError, + TypeError, + KeyError, + zmq.ZMQError, + ): + # A single malformed or unexpected request (e.g. a stray + # non-NIXL client, one that doesn't speak the REQ/ROUTER + # envelope, or a request for a pp/tp rank this engine + # doesn't serve) must not take down the listener: every + # remote's handshake depends on this thread staying alive + # for the lifetime of the process. Anything outside this + # list is a genuine bug and should still crash the thread + # loudly rather than be swallowed here. + logger.exception( + "Error handling NIXL handshake request " + "(raw=%.200r); ignoring it and continuing to listen on %s", + frames[-1], + path, + ) def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int: """D-side only. Returns N-1 for Mamba models since the decoder 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 19b35b8253cf..c9eec18ed3ad 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 @@ -764,10 +764,24 @@ def _nixl_handshake( (GET_META_MSG, remote_pp_rank, remote_rank) ) # Set receive timeout to 5 seconds to avoid hanging on dead server - sock.setsockopt(zmq.RCVTIMEO, 5000) # milliseconds + handshake_timeout_ms = 5000 + sock.setsockopt(zmq.RCVTIMEO, handshake_timeout_ms) start_time = time.perf_counter() sock.send(msg) - reply_parts = sock.recv_multipart() + try: + reply_parts = sock.recv_multipart() + except zmq.error.Again as e: + # ZMQ hides the underlying connection state, so this + # timeout is all we get; point at the probe script + # instead of a bare "Resource temporarily unavailable". + raise RuntimeError( + f"No reply from NIXL handshake listener at {path} " + f"(pp_rank={remote_pp_rank}, tp_rank={remote_rank}) " + f"after {handshake_timeout_ms}ms. Probe it with " + "tests/v1/kv_connector/nixl_integration/" + f"nixl_side_channel_probe.py --host {host} --port {port} " + f"--pp-rank {remote_pp_rank} --rank {remote_rank}." + ) from e recv_time = time.perf_counter() assert len(reply_parts) == 2 handshake_bytes = reply_parts[0]