[Bugfix][MoRIIO] Keep discovery heartbeats running while workers hold the GIL

tests/v1/kv_connector/unit/test_moriio_proxy_routing.py

@@ -22,14 +22,20 @@
 """
 
 import contextlib
+import ctypes
 import importlib.util
+import os
+import subprocess
 import sys
+import time
 import types
 from collections import Counter
 from pathlib import Path
 from typing import cast
 
+import msgpack
 import pytest
+import zmq
 
 _MISSING = object()
 
@@ -203,3 +209,127 @@ def test_consecutive_requests_alternate_instances(route):
     # than filling one instance's ranks before moving on.
     insts = [inst for inst, _ in _route_n(route, _instances(3, 8), 3)]
     assert insts == [0, 1, 2]
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="GIL probe uses libc.usleep")
+def test_discovery_heartbeat_progresses_while_worker_holds_gil():
+    from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import (
+        MoRIIOHeartbeat,
+    )
+
+    payload = {
+        "type": "P",
+        "http_address": "127.0.0.1:8000",
+        "zmq_address": "host:127.0.0.1,handshake:6301,notify:61005",
+        "dp_size": 1,
+        "tp_size": 8,
+        "transfer_mode": "READ",
+    }
+    with zmq.Context() as context, context.socket(zmq.ROUTER) as socket:
+        port = socket.bind_to_random_port("tcp://127.0.0.1")
+        heartbeat = MoRIIOHeartbeat(f"tcp://127.0.0.1:{port}", payload, 0.02)
+        try:
+            assert socket.poll(5000), "Discovery did not register"
+            assert msgpack.unpackb(socket.recv_multipart()[1]) == payload
+            # PyDLL retains the GIL throughout the native call.
+            ctypes.PyDLL(None).usleep(300_000)
+            received = 0
+            while socket.poll(0):
+                assert msgpack.unpackb(socket.recv_multipart()[1]) == payload
+                received += 1
+            assert received >= 3, "Discovery stalled while the worker held the GIL"
+        finally:
+            heartbeat.shutdown()
+            heartbeat.shutdown()
+        while socket.poll(0):
+            socket.recv_multipart()
+        assert not socket.poll(100), "Heartbeats continued after shutdown"
+
+
+@pytest.mark.skipif(sys.platform != "linux", reason="Checks child exit via procfs")
+def test_discovery_heartbeat_exits_when_worker_dies():
+    from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_heartbeat
+
+    owner_code = """
+import runpy
+import sys
+Heartbeat = runpy.run_path(sys.argv[1])["MoRIIOHeartbeat"]
+heartbeat = Heartbeat(sys.argv[2], {"type": "D"}, 0.02)
+print(heartbeat._process.pid, flush=True)
+sys.stdin.read()
+"""
+    with zmq.Context() as context, context.socket(zmq.ROUTER) as socket:
+        port = socket.bind_to_random_port("tcp://127.0.0.1")
+        owner = subprocess.Popen(
+            [
+                sys.executable,
+                "-c",
+                owner_code,
+                moriio_heartbeat.__file__,
+                f"tcp://127.0.0.1:{port}",
+            ],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            text=True,
+        )
+        child_pid = None
+        child_exited = False
+        try:
+            assert socket.poll(5000), "Discovery did not register"
+            assert owner.stdout is not None
+            child_pid = int(owner.stdout.readline())
+            owner.kill()
+            owner.wait(timeout=5)
+            stat = Path(f"/proc/{child_pid}/stat")
+            deadline = time.monotonic() + 5
+            while stat.exists():
+                # A terminated child may await reaping by the container's init.
+                try:
+                    if stat.read_text().split(") ")[1].startswith("Z "):
+                        break
+                except FileNotFoundError:
+                    break
+                assert time.monotonic() < deadline, "Orphan heartbeat survived worker"
+                time.sleep(0.01)
+            child_exited = True
+        finally:
+            if owner.poll() is None:
+                owner.kill()
+            owner.communicate(timeout=5)
+            if child_pid is not None and not child_exited:
+                with contextlib.suppress(ProcessLookupError):
+                    os.kill(child_pid, 9)
+
+
+def test_discovery_heartbeat_recovers_after_router_outage(capfd):
+    from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import (
+        MoRIIOHeartbeat,
+    )
+
+    payload = {"type": "P", "http_address": "127.0.0.1:8000"}
+    with zmq.Context() as context:
+        with context.socket(zmq.ROUTER) as socket:
+            port = socket.bind_to_random_port("tcp://127.0.0.1")
+            endpoint = f"tcp://127.0.0.1:{port}"
+            heartbeat = MoRIIOHeartbeat(endpoint, payload, 0.001, max_retries=1)
+            try:
+                assert socket.poll(5000), "Discovery did not register"
+                assert msgpack.unpackb(socket.recv_multipart()[1]) == payload
+            except BaseException:
+                heartbeat.shutdown()
+                raise
+        try:
+            # Fill the real DEALER queue while discovery is unavailable.
+            deadline = time.monotonic() + 10
+            errors = ""
+            while "heartbeat send failed" not in errors:
+                errors += capfd.readouterr().err
+                assert time.monotonic() < deadline, "No send timeout during outage"
+                time.sleep(0.01)
+            with context.socket(zmq.ROUTER) as socket:
+                socket.bind(endpoint)
+                assert socket.poll(5000), "Discovery did not recover after outage"
+                assert msgpack.unpackb(socket.recv_multipart()[1]) == payload
+                assert heartbeat._process.poll() is None
+        finally:
+            heartbeat.shutdown()

tests/v1/kv_connector/unit/test_moriio_tp_ack.py

@@ -269,6 +269,34 @@ def shutdown(self):
     assert worker._completed_consumer_notifications == {"tx-fanin"}
 
 
+@pytest.mark.parametrize("exit_kind", ["killed", "eof"])
+def test_worker_reports_unexpected_heartbeat_exit(exit_kind):
+    from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import (
+        MoRIIOHeartbeat,
+    )
+
+    worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker)
+    worker.is_producer = True
+    worker.mode = MoRIIOMode.READ
+    worker.moriio_wrapper = SimpleNamespace(
+        pop_finished_req_ids=lambda: [], shutdown=lambda: None
+    )
+    worker._pending_unmapped_acks = []
+    worker._heartbeat = MoRIIOHeartbeat("tcp://127.0.0.1:1", {"type": "P"}, 0.02)
+    process = worker._heartbeat._process
+    try:
+        assert worker.get_finished() == (set(), set())
+        if exit_kind == "killed":
+            process.kill()
+        else:
+            process.stdin.close()
+        returncode = process.wait(timeout=5)
+        with pytest.raises(RuntimeError, match=f"heartbeat.*{returncode}"):
+            worker.get_finished()
+    finally:
+        worker.shutdown()
+
+
 def test_read_completion_sends_structured_release_with_consumer_tp_size():
     class DoneStatus:
         def Succeeded(self):

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -59,6 +59,9 @@
     MoRIIOWrapper,
     MoRIIOWriter,
 )
+from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_heartbeat import (
+    MoRIIOHeartbeat,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_layout import (
     LayerTransferGeometry,
     MambaOffsetTemplate,
@@ -1600,7 +1603,7 @@ def __init__(
 
         self.moriio_engine = None
         self._handle_request_thread = None
-        self._ping_thread = None
+        self._heartbeat: MoRIIOHeartbeat | None = None
         self._writer = MoRIIOWriter(self)
         # Completions that arrived before transfer_id_to_request_id was populated.
         # Retried each step until the mapping is established.
@@ -1630,10 +1633,7 @@ def __init__(
         )
 
         if self._rank == 0 and self.moriio_config.proxy_ip:
-            self._ping_thread = threading.Thread(
-                target=self._ping, args=(self.zmq_context,), daemon=True
-            )
-            self._ping_thread.start()
+            self._heartbeat = self._start_heartbeat()
 
         logger.info(
             "Initializing MoRIIO Engine, engine = %s, role = %s",
@@ -1881,75 +1881,31 @@ def _get_built_session(self, remote_engine_id):
             remote_engine_id
         ]
 
-    def _ping(self, zmq_context):
-        # Use host:port format for http_address (compatible with official router)
-        http_address = f"{self.request_address}"
-        # Include host so the router embeds it in the request_id; the connector
-        # on the other side parses host/ports from there.
-        zmq_address = (
-            f"host:{self.local_ip},"
-            f"handshake:{self.handshake_port},"
-            f"notify:{self.notify_port}"
+    def _start_heartbeat(self) -> MoRIIOHeartbeat:
+        payload = {
+            "type": "P" if self.is_producer else "D",
+            "http_address": self.request_address,
+            "zmq_address": (
+                f"host:{self.local_ip},handshake:{self.handshake_port},"
+                f"notify:{self.notify_port}"
+            ),
+            "dp_size": self.moriio_config.dp_size,
+            "tp_size": self.moriio_config.tp_size,
+            "transfer_mode": self.mode.name,
+        }
+        return MoRIIOHeartbeat(
+            f"tcp://{self.proxy_ip}:{self.proxy_ping_port}",
+            payload,
+            MoRIIOConstants.PING_INTERVAL,
+            MoRIIOConstants.MAX_PING_RETRIES,
         )
-        role = "P" if self.is_producer else "D"
-
-        retry_count = 0
-        index = 1
-        with zmq_context.socket(zmq.DEALER) as sock:
-            sock.connect(f"tcp://{self.proxy_ip}:{self.proxy_ping_port}")
-
-            while True:
-                try:
-                    data = {
-                        "type": role,  # "P" or "D"
-                        "http_address": http_address,
-                        "zmq_address": zmq_address,
-                        # dp_size/tp_size are not used by the official vLLM router
-                        # (routing operates at the http_address level); they are
-                        # consumed only by the toy proxy server.
-                        "dp_size": self.moriio_config.dp_size,
-                        "tp_size": self.moriio_config.tp_size,
-                        # transfer_mode is included so the router can distinguish
-                        # READ (prefill-then-decode, sequential) from WRITE (concurrent)
-                        # scheduling.
-                        "transfer_mode": self.mode.name,
-                    }
-
-                    sock.send(msgpack.dumps(data))
-                    # logger.debug(f"Successfully sent ping message #{index}")
-                    retry_count = 0
-
-                except ConnectionRefusedError:
-                    logger.info(
-                        "Connection refused: %s:%s -> %s:%s",
-                        self.local_ip,
-                        self.local_ping_port,
-                        self.proxy_ip,
-                        self.proxy_ping_port,
-                    )
-                    retry_count += 1
-
-                except OSError as e:
-                    logger.info("OS error when sending ping: %s", e)
-                    retry_count += 1
-
-                except Exception as e:
-                    logger.info("Unexpected error when sending ping: %s", e)
-                    retry_count += 1
-                    if retry_count >= MoRIIOConstants.MAX_PING_RETRIES:
-                        logger.error(
-                            "Max retries (%s) exceeded. Stopping ping loop.",
-                            MoRIIOConstants.MAX_PING_RETRIES,
-                        )
-                        raise RuntimeError(
-                            f"Ping failed after {retry_count} retries"
-                        ) from e
-
-                finally:
-                    time.sleep(MoRIIOConstants.PING_INTERVAL)
-                    index += 1
 
     def shutdown(self):
+        heartbeat = getattr(self, "_heartbeat", None)
+        if heartbeat is not None:
+            heartbeat.shutdown()
+            self._heartbeat = None
+
         if hasattr(self, "moriio_wrapper") and self.moriio_wrapper:
             self.moriio_wrapper.shutdown()
 
@@ -2580,6 +2536,9 @@ def get_finished(self) -> tuple[set[str], set[str]]:
         The scheduler process (via the MultiprocExecutor) will use this output
         to track which workers are done.
         """
+        heartbeat = getattr(self, "_heartbeat", None)
+        if heartbeat is not None:
+            heartbeat.check_health()
         done_sending, done_recving = set(), set()
 
         if self.is_producer: