[Bugfix][NIXL] Harden handshake listener against bad requests

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,

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

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]