@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 386ac2573e

[Bugfix][KV Connector] Retry Mooncake bootstrap registration on timeout (reopens #55763) (#58919)

Chao Lei committed Sep 28, 2026 at 03:47 UTC 386ac2573e334a670a0406c9a2bca3a9981ac63d

5 files changed +254 -14

docs/features/mooncake_connector_usage.md

@@ -46,6 +46,13 @@ Now you can send requests to the proxy server through port 8000.

46 - For headless instances, must be the same as the master instance

47 - Each instance needs a unique port on its host; using the same port number across different hosts is fine

49 +- `VLLM_MOONCAKE_CONNECTOR_TIMEOUT`: Per-request timeout (in seconds) for MooncakeConnector HTTP calls to the bootstrap server, including prefiller worker registration at startup. (Optional)

50 + - Default: 30.0

51 + - Global rank 0 hosts the bootstrap server in the same process that mounts the Mooncake transfer engine's host segment, so a large `MOONCAKE_GLOBAL_SEGMENT_SIZE` can delay responses for several seconds during startup

52 + - Raise this if bootstrap server timeout warnings appear on hosts with very large host-memory segments

53 + - Bootstrap registration uses up to 3 attempts internally, with exponential backoff capped at 10 seconds; connection errors and timeouts are both retried

54 + - Once the attempts are exhausted the caller raises instead of blocking, so failures surface rather than hanging

55 +

56 - `WITH_NVIDIA_PEERMEM`: Selects how mooncake registers GPU memory for RDMA. Read by mooncake, not vLLM.

57 - Default: 1, which uses `ibv_reg_mr()` and requires the `nvidia-peermem` kernel module to be loaded

58 - Set to 0 to use the DMA-BUF path, which does not need that module. Required on hosts where `nvidia-peermem` is not loaded, such as GB200

tests/v1/kv_connector/unit/test_mooncake_connector.py

+183 -1

@@ -3,10 +3,12 @@

4 import asyncio

5 import contextlib

6 +import threading

7 import time

8 from types import SimpleNamespace

9 from unittest.mock import AsyncMock, MagicMock, patch

11 +import httpx

12 import pytest

13 import torch

14 import zmq.asyncio

@@ -14,6 +16,7 @@ import zmq.asyncio

16 from tests.v1.attention.utils import dense_kv_cache_views

17 from vllm import envs

18 from vllm.config import set_current_vllm_config

19 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import mooncake_connector

20 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector import (

21 KVConnectorRole,

22 MooncakeConnector,

@@ -33,6 +36,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector im

36 )

37 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import (

38 MooncakeBootstrapServer,

39 + RegisterWorkerPayload,

40 )

41 from vllm.utils.network_utils import get_open_port

42 from vllm.v1.kv_cache_interface import (

@@ -538,6 +542,46 @@ def bootstrap_server():

542 server.shutdown()

543

544

545 +@pytest.mark.asyncio

546 +async def test_register_worker_recovers_from_slow_bootstrap_server(monkeypatch):

547 + """End-to-end: a real HTTP server that responds slower than the configured

548 + timeout must be retried on the wire, not treated as fatal."""

549 + original = MooncakeBootstrapServer.register_worker

550 + state = {"done": False, "calls": 0}

551 +

552 + async def slow_register(self, payload: RegisterWorkerPayload):

553 + state["calls"] += 1

554 + if not state["done"]:

555 + state["done"] = True

556 + # Record the worker first, then stall past the client timeout, so

557 + # the retry exercises the duplicate-registration path.

558 + response = await original(self, payload)

559 + await asyncio.sleep(1.5)

560 + return response

561 + return await original(self, payload)

562 +

563 + monkeypatch.setattr(MooncakeBootstrapServer, "register_worker", slow_register)

564 +

565 + port = get_open_port()

566 + monkeypatch.setenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", str(port))

567 + monkeypatch.setenv("VLLM_MOONCAKE_CONNECTOR_TIMEOUT", "0.5")

568 + monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 5)

569 +

570 + server = MooncakeBootstrapServer("127.0.0.1", port)

571 + server.start()

572 + try:

573 + worker = _make_local_register_worker_stub()

574 + await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)

575 + assert state["calls"] == 2

576 +

577 + async with httpx.AsyncClient() as client:

578 + response = await client.get(f"http://127.0.0.1:{port}/query")

579 + assert response.status_code == 200

580 + assert response.json()["0"]["engine_id"] == "eng-1"

581 + finally:

582 + server.shutdown()

583 +

584 +

585 @pytest.mark.asyncio

586 async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):

587 """Tests the bootstrap server's api for worker registration and querying.

@@ -589,9 +633,16 @@ async def test_bootstrap_server(bootstrap_server: MooncakeBootstrapServer):

633 assert data["0"]["worker_addr"]["0"]["0"] == "tcp://1.1.1.1:1111"

634 assert data["0"]["worker_addr"]["0"]["1"] == "tcp://2.2.2.2:2222"

635

592 - # Test failure: re-registering the same worker

636 + # Re-registering the identical payload is idempotent.

637 async with httpx.AsyncClient() as client:

638 response = await client.post(f"{base_url}/register", json=payload1)

639 + assert response.status_code == 200

640 + assert response.json() == {"status": "ok"}

641 +

642 + # Test failure: same ranks, conflicting address

643 + payload1_conflict = dict(payload1, addr="tcp://9.9.9.9:9999")

644 + async with httpx.AsyncClient() as client:

645 + response = await client.post(f"{base_url}/register", json=payload1_conflict)

646 assert response.status_code == 400

647 assert "is already registered" in response.text

648

@@ -628,6 +679,137 @@ def _make_bootstrap_vllm_config(

679 )

680

681

682 +def _make_local_register_worker_stub():

683 + return SimpleNamespace(

684 + vllm_config=SimpleNamespace(

685 + parallel_config=SimpleNamespace(

686 + local_engines_only=False,

687 + data_parallel_rank_local=0,

688 + data_parallel_index=0,

689 + nnodes_within_dp=1,

690 + master_addr="127.0.0.1",

691 + data_parallel_master_ip="127.0.0.1",

692 + )

693 + ),

694 + hostname="127.0.0.1",

695 + side_channel_port=1234,

696 + engine_id="eng-1",

697 + dp_rank=0,

698 + tp_rank=0,

699 + pp_rank=0,

700 + )

701 +

702 +

703 +def _make_register_worker_stub():

704 + return SimpleNamespace(

705 + vllm_config=_make_bootstrap_vllm_config(),

706 + hostname="127.0.0.1",

707 + side_channel_port=1234,

708 + engine_id="eng-1",

709 + dp_rank=0,

710 + tp_rank=0,

711 + pp_rank=0,

712 + )

713 +

714 +

715 +class _FlakyAsyncClient:

716 + """httpx.AsyncClient stub that read-times-out `failures` times first."""

717 +

718 + def __init__(self, failures: int, calls: list[int], **kwargs):

719 + self.failures = failures

720 + self.calls = calls

721 +

722 + async def __aenter__(self):

723 + return self

724 +

725 + async def __aexit__(self, *exc_info):

726 + return False

727 +

728 + async def post(self, url, json=None):

729 + self.calls.append(1)

730 + if len(self.calls) <= self.failures:

731 + raise httpx.ReadTimeout("simulated slow bootstrap server")

732 + response = MagicMock()

733 + response.raise_for_status = MagicMock()

734 + return response

735 +

736 +

737 +@pytest.mark.asyncio

738 +async def test_register_worker_retries_on_read_timeout(monkeypatch):

739 + """A slow rank-0 bootstrap response must be retried, not treated as fatal."""

740 + monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 5)

741 + monkeypatch.setattr(asyncio, "sleep", AsyncMock())

742 +

743 + calls: list[int] = []

744 + monkeypatch.setattr(

745 + httpx,

746 + "AsyncClient",

747 + lambda **kwargs: _FlakyAsyncClient(2, calls),

748 + )

749 +

750 + worker = _make_register_worker_stub()

751 + await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)

752 +

753 + assert len(calls) == 3

754 +

755 +

756 +@pytest.mark.asyncio

757 +async def test_register_worker_raises_after_max_attempts(monkeypatch):

758 + """Terminal registration failure must raise, not loop forever."""

759 + monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 3)

760 + monkeypatch.setattr(asyncio, "sleep", AsyncMock())

761 +

762 + calls: list[int] = []

763 + monkeypatch.setattr(

764 + httpx,

765 + "AsyncClient",

766 + lambda **kwargs: _FlakyAsyncClient(99, calls),

767 + )

768 +

769 + worker = _make_register_worker_stub()

770 + with pytest.raises(RuntimeError, match="after 3 attempts"):

771 + await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)

772 +

773 + assert len(calls) == 3

774 +

775 +

776 +@pytest.mark.asyncio

777 +async def test_sender_listener_failure_propagates_to_caller(monkeypatch):

778 + """A terminal registration failure must surface via the Future rather than

779 + leaving register_kv_caches blocked on its ready event."""

780 + monkeypatch.setenv("VLLM_MOONCAKE_CONNECTOR_TIMEOUT", "0.1")

781 + monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 1)

782 +

783 + loop = asyncio.new_event_loop()

784 + thread = threading.Thread(target=loop.run_forever, daemon=True)

785 + thread.start()

786 +

787 + async def failing_listener(ready_event):

788 + raise RuntimeError("simulated terminal registration failure")

789 +

790 + try:

791 + ready_event = threading.Event()

792 + fut = asyncio.run_coroutine_threadsafe(failing_listener(ready_event), loop)

793 +

794 + deadline = time.monotonic() + 10.0

795 + raised = None

796 + while not ready_event.wait(timeout=0.1):

797 + if fut.done():

798 + try:

799 + fut.result()

800 + except RuntimeError as e:

801 + raised = e

802 + break

803 + assert time.monotonic() < deadline, "caller blocked past its deadline"

804 +

805 + assert raised is not None

806 + assert "simulated terminal registration failure" in str(raised)

807 + finally:

808 + loop.call_soon_threadsafe(loop.stop)

809 + thread.join(timeout=5)

810 + loop.close()

811 +

812 +

813 @pytest.mark.parametrize(

814 (

815 "tp_rank",

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

+51 -11

@@ -10,7 +10,7 @@ from collections.abc import Collection

10 from concurrent.futures import ThreadPoolExecutor

11 from dataclasses import dataclass

12 from enum import IntEnum

13 -from typing import TYPE_CHECKING, Any

13 +from typing import TYPE_CHECKING, Any, Final

15 import msgspec

16 import numpy as np

@@ -70,6 +70,8 @@ from vllm.v1.worker.utils import select_common_block_size

71 logger = init_logger(__name__)

73 +_BOOTSTRAP_MAX_ATTEMPTS: Final[int] = 3

74 +

75 try:

76 from mooncake.engine import TransferEngine

77 except ImportError:

@@ -1147,24 +1149,44 @@ class MooncakeConnectorWorker:

1149 pp_rank=self.pp_rank,

1150 addr=worker_addr,

1151 )

1150 - while True:

1152 + timeout = envs.VLLM_MOONCAKE_CONNECTOR_TIMEOUT

1153 + max_attempts = _BOOTSTRAP_MAX_ATTEMPTS

1154 + backoff = 1.0

1155 + for attempt in range(1, max_attempts + 1):

1156 try:

1152 - async with httpx.AsyncClient() as client:

1157 + async with httpx.AsyncClient(timeout=timeout) as client:

1158 response = await client.post(url, json=payload.model_dump())

1159 response.raise_for_status()

1160 logger.debug("Successfully registered with bootstrap server at %s", url)

1156 - break

1157 - except httpx.ConnectError:

1158 - # Bootstrap server not ready, wait for a while and retry.

1159 - await asyncio.sleep(1)

1161 + return

1162 + except (httpx.ConnectError, httpx.TimeoutException) as e:

1163 + # Rank 0 may be busy registering a large memory segment.

1164 + if attempt == max_attempts:

1165 + break

1166 + logger.warning(

1167 + "Bootstrap registration attempt %d/%d for %s failed with %s: %s. "

1168 + "Retrying in %.1fs.",

1169 + attempt,

1170 + max_attempts,

1171 + payload,

1172 + type(e).__name__,

1173 + e,

1174 + backoff,

1175 + )

1176 + await asyncio.sleep(backoff)

1177 + backoff = min(backoff * 2, 10.0)

1178 except Exception as e:

1179 err_msg = (

1180 e.response.text if isinstance(e, httpx.HTTPStatusError) else str(e)

1181 )

1164 - logger.error(

1182 + logger.exception(

1183 "Error registering %s with bootstrap server: %s", payload, err_msg

1184 )

1167 - raise e

1185 + raise

1186 + raise RuntimeError(

1187 + f"Failed to register {payload} with the Mooncake bootstrap server at "

1188 + f"{url} after {max_attempts} attempts."

1189 + )

1190

1191 async def _mooncake_sender_listener(self, ready_event: threading.Event):

1192 """Background thread that listens for Mooncake requests, dispatches them

@@ -1800,11 +1822,29 @@ class MooncakeConnectorWorker:

1822 if self.is_kv_consumer:

1823 return

1824

1825 + # httpx applies the timeout to each phase (connect, read, write, pool)

1826 + # rather than to the whole request, so budget for all four.

1827 + register_timeout = envs.VLLM_MOONCAKE_CONNECTOR_TIMEOUT

1828 + max_attempts = _BOOTSTRAP_MAX_ATTEMPTS

1829 + backoff_total = sum(min(2.0**i, 10.0) for i in range(max_attempts - 1))

1830 + ready_timeout = max_attempts * 4 * register_timeout + backoff_total + 30.0

1831 ready_event = threading.Event()

1804 - asyncio.run_coroutine_threadsafe(

1832 + fut = asyncio.run_coroutine_threadsafe(

1833 self._mooncake_sender_listener(ready_event), self.sender_loop

1834 )

1807 - ready_event.wait() # Wait for listener ZMQ socket to be ready.

1835 + deadline = time.monotonic() + ready_timeout

1836 + while not ready_event.wait(timeout=1.0):

1837 + if fut.done():

1838 + fut.result()

1839 + raise RuntimeError(

1840 + "Mooncake sender listener exited before becoming ready."

1841 + )

1842 + if time.monotonic() > deadline:

1843 + fut.cancel()

1844 + raise RuntimeError(

1845 + "Mooncake sender listener did not become ready within "

1846 + f"{ready_timeout:.0f}s."

1847 + )

1848

1849 async def fetch_finished_recving_reqs(self) -> set[ReqId]:

1850 finished_recving_reqs = self.finished_recving_reqs

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_utils.py

+7 -2

@@ -97,14 +97,19 @@ class MooncakeBootstrapServer:

97 dp_entry.worker_addr[payload.tp_rank] = {}

99 tp_entry = dp_entry.worker_addr[payload.tp_rank]

100 - if payload.pp_rank in tp_entry:

100 + existing = tp_entry.get(payload.pp_rank)

101 + if existing is not None:

102 + # A client timeout can fire after the server recorded the

103 + # registration so an identical retry must not be an error.

104 + if existing == payload.addr:

105 + return {"status": "ok"}

106 raise HTTPException(

107 status_code=400,

108 detail=(

109 f"Worker with dp_rank={payload.dp_rank}, "

110 f"tp_rank={payload.tp_rank}, pp_rank={payload.pp_rank} "

111 f"is already registered at "

107 - f"{tp_entry[payload.pp_rank]}, "

112 + f"{existing}, "

113 f"but still want to register at {payload.addr}"

114 ),

115 )

vllm/envs.py

@@ -260,6 +260,7 @@ if TYPE_CHECKING:

260 VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None

261 VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None

262 VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480

263 + VLLM_MOONCAKE_CONNECTOR_TIMEOUT: float = 30.0

264 VLLM_ENABLE_CUDAGRAPH_GC: bool = False

265 VLLM_LOOPBACK_IP: str = ""

266 VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True

@@ -1713,6 +1714,11 @@ environment_variables: dict[str, Callable[[], Any]] = {

1714 "VLLM_EC_SIDE_CHANNEL_PORT": lambda: int(

1715 os.getenv("VLLM_EC_SIDE_CHANNEL_PORT", "5601")

1716 ),

1717 + # Per-request timeout (seconds) when a prefiller worker registers with

1718 + # the Mooncake bootstrap server.

1719 + "VLLM_MOONCAKE_CONNECTOR_TIMEOUT": lambda: float(

1720 + os.getenv("VLLM_MOONCAKE_CONNECTOR_TIMEOUT", "30.0")

1721 + ),

1722 # Port used for Mooncake handshake between remote agents.

1723 "VLLM_MOONCAKE_BOOTSTRAP_PORT": lambda: int(

1724 os.getenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", "8998")