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..e3e04aedb836 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,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() 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 ec24f0d199f4..1dc6f8459ab3 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, @@ -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: 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..87eefe8f4279 --- /dev/null +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_heartbeat.py @@ -0,0 +1,96 @@ +# 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 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: + 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.Again: + failures = 0 + logger.warning("MoRIIO discovery heartbeat send failed; retrying") + 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()