From 512a877511c7117683d7e3498feed7061c94f02b Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Wed, 30 Sep 2026 14:05:30 +0000 Subject: [PATCH 1/2] [MoRIIO] Keep discovery heartbeats independent of the worker GIL Extract the process-based discovery sender from #58968 onto main. Preserve the discovery payload and rank-zero startup guard, stop the helper on explicit shutdown or parent pipe EOF, and retain the GIL-starvation and parent-exit lifecycle regressions. Co-authored-by: Codex Signed-off-by: whx-sjtu --- .../unit/test_moriio_proxy_routing.py | 96 ++++++++++++++++++ .../v1/moriio/moriio_connector.py | 98 +++++-------------- .../v1/moriio/moriio_heartbeat.py | 86 ++++++++++++++++ 3 files changed, 209 insertions(+), 71 deletions(-) create mode 100644 vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py diff --git a/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py b/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py index 71ee0ef5d799..68066ede1c3d 100644 --- a/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py +++ b/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,93 @@ 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) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index b482e6a8c4a3..d56eb5eba99d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/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, @@ -1613,7 +1616,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. @@ -1643,10 +1646,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", @@ -1894,75 +1894,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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py new file mode 100644 index 000000000000..01a7f38c59f2 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Keep discovery heartbeats independent of the model worker's GIL.""" + +import json +import logging +import select +import subprocess +import sys +from typing import Any + +import msgpack +import zmq + +logger = logging.getLogger(__name__) + + +class MoRIIOHeartbeat: + def __init__( + self, + endpoint: str, + payload: dict[str, Any], + interval: float, + max_retries: int = 100, + ): + self._process = subprocess.Popen( + [sys.executable, __file__, endpoint, str(interval), str(max_retries)], + stdin=subprocess.PIPE, + text=True, + ) + assert self._process.stdin is not None + try: + self._process.stdin.write(json.dumps(payload) + "\n") + self._process.stdin.flush() + except BaseException: + self.shutdown() + raise + + def shutdown(self) -> None: + assert self._process.stdin is not None + try: + if not self._process.stdin.closed: + self._process.stdin.close() + except BrokenPipeError: + pass + finally: + try: + self._process.wait(timeout=5) + except subprocess.TimeoutExpired: + self._process.terminate() + try: + self._process.wait(timeout=1) + except subprocess.TimeoutExpired: + self._process.kill() + self._process.wait(timeout=1) + + +def main() -> None: + endpoint, interval_arg, max_retries_arg = sys.argv[1:] + initial = sys.stdin.readline() + if not initial: + return + payload = msgpack.dumps(json.loads(initial)) + interval = float(interval_arg) + max_retries = int(max_retries_arg) + failures = 0 + with zmq.Context() as context, context.socket(zmq.DEALER) as sock: + sock.setsockopt(zmq.LINGER, 0) + sock.setsockopt(zmq.SNDTIMEO, 1000) + sock.connect(endpoint) + while not select.select([sys.stdin], [], [], 0)[0]: + try: + sock.send(payload) + failures = 0 + except zmq.ZMQError: + failures += 1 + logger.exception("MoRIIO discovery heartbeat send failed") + if failures >= max_retries: + raise + # EOF also stops the child when the worker exits without cleanup. + if select.select([sys.stdin], [], [], interval)[0]: + break + + +if __name__ == "__main__": + main() From dc5f4f2ffad0efdbd03b1f55e1bb96c3d8deca94 Mon Sep 17 00:00:00 2001 From: whx-sjtu Date: Thu, 1 Oct 2026 04:59:50 +0000 Subject: [PATCH 2/2] [MoRIIO] Recover discovery heartbeats and report child exit Retry temporary ZMQ send timeouts for the worker lifetime. Surface unexpected heartbeat child termination, including exit code zero, from worker completion polling. Cover real router outage recovery and child exits. Co-authored-by: Codex Signed-off-by: whx-sjtu --- .../unit/test_moriio_proxy_routing.py | 34 +++++++++++++++++++ .../kv_connector/unit/test_moriio_tp_ack.py | 28 +++++++++++++++ .../v1/moriio/moriio_connector.py | 3 ++ .../v1/moriio/moriio_heartbeat.py | 10 ++++++ 4 files changed, 75 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py b/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py index 68066ede1c3d..e3e04aedb836 100644 --- a/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_proxy_routing.py @@ -299,3 +299,37 @@ def test_discovery_heartbeat_exits_when_worker_dies(): 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() diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py b/tests/v1/kv_connector/unit/test_moriio_tp_ack.py index a35f0a1169b6..34d193e272dd 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_ack.py +++ b/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): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index d56eb5eba99d..a68f76f0c4de 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -2549,6 +2549,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: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py index 01a7f38c59f2..87eefe8f4279 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py @@ -36,6 +36,13 @@ def __init__( self.shutdown() raise + def check_health(self) -> None: + returncode = self._process.poll() + if returncode is not None: + raise RuntimeError( + f"MoRIIO discovery heartbeat exited with code {returncode}" + ) + def shutdown(self) -> None: assert self._process.stdin is not None try: @@ -72,6 +79,9 @@ def main() -> None: try: sock.send(payload) failures = 0 + except zmq.Again: + failures = 0 + logger.warning("MoRIIO discovery heartbeat send failed; retrying") except zmq.ZMQError: failures += 1 logger.exception("MoRIIO discovery heartbeat send failed")