@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3acddf01e7

[Bugfix][ROCm] Keep Mooncake bootstrap ports bound during startup (#58967)

Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Andreas Karatzas committed Sep 29, 2026 at 04:14 UTC 3acddf01e709b13693d3a90d25fdb74d8124049c

5 files changed +226 -36

docs/features/mooncake_connector_usage.md

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

76 - **num_workers**: Size of thread pool for one prefiller worker to transfer KV caches by mooncake. (default 10)

77 - **mooncake_protocol**: Mooncake connector protocol. (default "rdma")

78 - **device_name**: Comma-separated whitelist of RDMA devices (e.g. `"mlx5_0,mlx5_1"`) to restrict topology discovery to. Empty discovers every device. Useful on hosts exposing a mix of InfiniBand and RoCE ports, where both peers must settle on the same link layer.

79 +- **bootstrap_server_address**: Optional `"host:port"` of an already running Mooncake bootstrap server. All prefiller ranks register with this endpoint instead of launching an embedded server. The caller owns the server lifetime and must configure the proxy to use the same endpoint. The example proxy assumes the bootstrap server is on the prefiller API host; a different host requires a proxy that supplies that address in `remote_bootstrap_addr`. Omit this option to use the embedded server and `VLLM_MOONCAKE_BOOTSTRAP_PORT`.

81 ## Example Scripts/Code

tests/v1/kv_connector/rocm_pd_accuracy_utils.py

+11 -2

@@ -158,8 +158,18 @@ def run_rocm_pd_accuracy(connector, model, prefill_tp, decode_tp, monkeypatch):

158 "VLLM_ROCM_USE_AITER": "1",

159 "VLLM_KV_CACHE_LAYOUT": "HND",

160 }

161 - bootstrap_port = get_open_port()

161 with contextlib.ExitStack() as stack:

162 + bootstrap_port = 0

163 + if connector == "mooncake":

164 + from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( # noqa: E501

165 + MooncakeBootstrapServer,

166 + )

167 +

168 + bootstrap = MooncakeBootstrapServer("127.0.0.1", 0)

169 + stack.callback(bootstrap.shutdown)

170 + bootstrap.start()

171 + bootstrap_port = bootstrap.port

172 + extra["bootstrap_server_address"] = f"127.0.0.1:{bootstrap_port}"

173 servers: list[RemoteOpenAIServer] = []

174 stack.callback(RemoteOpenAIServer.shutdown_many, servers)

175 for role, tp, assigned in (

@@ -176,7 +186,6 @@ def run_rocm_pd_accuracy(connector, model, prefill_tp, decode_tp, monkeypatch):

186 selector: ",".join(assigned),

187 "VLLM_PORT": str(get_open_port()),

188 "VLLM_NIXL_SIDE_CHANNEL_PORT": str(get_open_port()),

179 - "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(bootstrap_port),

189 }

190 servers.append(

191 RemoteOpenAIServer(

tests/v1/kv_connector/unit/test_mooncake_connector.py

+132 -15

@@ -3,6 +3,8 @@

4 import asyncio

5 import contextlib

6 +import errno

7 +import socket

8 import threading

9 import time

10 from types import SimpleNamespace

@@ -15,8 +17,11 @@ import zmq.asyncio

18 from tests.v1.attention.utils import dense_kv_cache_views

19 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.config import KVTransferConfig, set_current_vllm_config

21 +from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import (

22 + mooncake_connector,

23 + mooncake_utils,

24 +)

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

26 KVConnectorRole,

27 MooncakeConnector,

@@ -38,7 +43,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import

43 MooncakeBootstrapServer,

44 RegisterWorkerPayload,

45 )

41 -from vllm.utils.network_utils import get_open_port

46 from vllm.v1.kv_cache_interface import (

47 FullAttentionSpec,

48 KVCacheConfig,

@@ -535,15 +539,99 @@ def test_prompt_less_than_block_size():

539 @pytest.fixture

540 def bootstrap_server():

541 """Fixture to launch and cleanup a Mooncake Bootstrap HTTP Server."""

538 - port = get_open_port()

539 - server = MooncakeBootstrapServer("127.0.0.1", port)

540 - server.start()

541 - yield server

542 - server.shutdown()

542 + server = MooncakeBootstrapServer("127.0.0.1", 0)

543 + try:

544 + server.start()

545 + yield server

546 + finally:

547 + server.shutdown()

548 +

549 +

550 +def test_bootstrap_server_occupied_port_raises_to_caller():

551 + with socket.create_server(("127.0.0.1", 0)) as listener:

552 + server = MooncakeBootstrapServer("127.0.0.1", listener.getsockname()[1])

553 + try:

554 + with pytest.raises(OSError) as exc_info:

555 + server.start(timeout=0.05)

556 + assert exc_info.value.errno == errno.EADDRINUSE

557 + finally:

558 + server.shutdown()

559 +

560 +

561 +def test_bootstrap_server_retains_and_releases_assigned_port(bootstrap_server):

562 + port = bootstrap_server.port

563 + assert port > 0

564 + with (

565 + pytest.raises(OSError) as exc_info,

566 + socket.create_server(("127.0.0.1", port)),

567 + ):

568 + pass

569 + assert exc_info.value.errno == errno.EADDRINUSE

570 +

571 + bootstrap_server.shutdown()

572 + with socket.create_server(("127.0.0.1", port)):

573 + pass

574 +

575 +

576 +@pytest.mark.parametrize("failure", ["exit", "error", "timeout"])

577 +def test_bootstrap_server_startup_failure_releases_port(monkeypatch, failure):

578 + stop = threading.Event()

579 +

580 + def run_without_starting(server, *, sockets):

581 + if failure == "error":

582 + raise SystemExit(1)

583 + if failure == "timeout":

584 + while not server.should_exit and not stop.wait(0.01):

585 + pass

586 +

587 + monkeypatch.setattr(mooncake_utils.uvicorn.Server, "run", run_without_starting)

588 + server = MooncakeBootstrapServer("127.0.0.1", 0)

589 + try:

590 + error = TimeoutError if failure == "timeout" else RuntimeError

591 + with pytest.raises(error):

592 + server.start(timeout=0.05 if failure == "timeout" else 1.0)

593 + assert server.port > 0

594 + with socket.create_server(("127.0.0.1", server.port)):

595 + pass

596 + finally:

597 + stop.set()

598 + server.shutdown()

599 +

600 +

601 +@pytest.mark.parametrize("state", ["stopping", "exited"])

602 +def test_bootstrap_server_rejects_stale_started_state(monkeypatch, state):

603 + release = threading.Event()

604 +

605 + def run_until_released(server, *, sockets):

606 + server.started = True

607 + release.wait()

608 +

609 + monkeypatch.setattr(mooncake_utils.uvicorn.Server, "run", run_until_released)

610 + server = MooncakeBootstrapServer("127.0.0.1", 0)

611 + try:

612 + server.start()

613 + thread = server.server_thread

614 + assert thread is not None

615 + if state == "stopping":

616 + with patch.object(thread, "join", return_value=None):

617 + server.shutdown()

618 + else:

619 + release.set()

620 + thread.join(timeout=5)

621 + assert not thread.is_alive()

622 +

623 + with pytest.raises(RuntimeError, match="still starting or stopping"):

624 + server.start()

625 + finally:

626 + release.set()

627 + server.shutdown()

628

629

630 @pytest.mark.asyncio

546 -async def test_register_worker_recovers_from_slow_bootstrap_server(monkeypatch):

631 +@pytest.mark.parametrize("external", [False, True])

632 +async def test_register_worker_recovers_from_slow_bootstrap_server(

633 + monkeypatch, external

634 +):

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

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

637 original = MooncakeBootstrapServer.register_worker

@@ -562,20 +650,24 @@ async def test_register_worker_recovers_from_slow_bootstrap_server(monkeypatch):

650

651 monkeypatch.setattr(MooncakeBootstrapServer, "register_worker", slow_register)

652

565 - port = get_open_port()

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

653 monkeypatch.setenv("VLLM_MOONCAKE_CONNECTOR_TIMEOUT", "0.5")

654 monkeypatch.setattr(mooncake_connector, "_BOOTSTRAP_MAX_ATTEMPTS", 5)

655

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

656 + server = MooncakeBootstrapServer("127.0.0.1", 0)

657 server.start()

658 try:

659 + monkeypatch.setenv("VLLM_MOONCAKE_BOOTSTRAP_PORT", str(server.port))

660 worker = _make_local_register_worker_stub()

661 + if external:

662 + worker.vllm_config = _make_bootstrap_vllm_config(

663 + bootstrap_server_address=f"127.0.0.1:{server.port}"

664 + )

665 + assert not should_launch_bootstrap_server(worker.vllm_config)

666 await MooncakeConnectorWorker.register_worker_with_bootstrap(worker)

667 assert state["calls"] == 2

668

669 async with httpx.AsyncClient() as client:

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

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

671 assert response.status_code == 200

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

673 finally:

@@ -666,8 +758,16 @@ def _make_bootstrap_vllm_config(

758 data_parallel_rank_local: int = 0,

759 data_parallel_index: int = 0,

760 nnodes_within_dp: int = 1,

761 + bootstrap_server_address: str | None = None,

762 ) -> SimpleNamespace:

763 return SimpleNamespace(

764 + kv_transfer_config=KVTransferConfig(

765 + kv_connector="MooncakeConnector",

766 + kv_role="kv_producer",

767 + kv_connector_extra_config={

768 + "bootstrap_server_address": bootstrap_server_address

769 + },

770 + ),

771 parallel_config=SimpleNamespace(

772 local_engines_only=local_engines_only,

773 data_parallel_rank_local=data_parallel_rank_local,

@@ -675,13 +775,14 @@ def _make_bootstrap_vllm_config(

775 nnodes_within_dp=nnodes_within_dp,

776 master_addr="model-parallel-master",

777 data_parallel_master_ip="data-parallel-master",

678 - )

778 + ),

779 )

780

781

782 def _make_local_register_worker_stub():

783 return SimpleNamespace(

784 vllm_config=SimpleNamespace(

785 + kv_transfer_config=None,

786 parallel_config=SimpleNamespace(

787 local_engines_only=False,

788 data_parallel_rank_local=0,

@@ -689,7 +790,7 @@ def _make_local_register_worker_stub():

790 nnodes_within_dp=1,

791 master_addr="127.0.0.1",

792 data_parallel_master_ip="127.0.0.1",

692 - )

793 + ),

794 ),

795 hostname="127.0.0.1",

796 side_channel_port=1234,

@@ -889,6 +990,22 @@ def test_get_mooncake_bootstrap_addr_selects_expected_host(

990 )

991

992

993 +@pytest.mark.parametrize(

994 + ("address", "expected"),

995 + [("external:12345", ("external", 12345)), ("[::1]:12345", ("::1", 12345))],

996 +)

997 +def test_get_mooncake_bootstrap_addr_uses_external_endpoint(address, expected):

998 + config = _make_bootstrap_vllm_config(bootstrap_server_address=address)

999 + assert get_mooncake_bootstrap_addr(config) == expected

1000 +

1001 +

1002 +@pytest.mark.parametrize("address", [":12345", "host:0", "host:65536"])

1003 +def test_get_mooncake_bootstrap_addr_rejects_invalid_external_address(address):

1004 + config = _make_bootstrap_vllm_config(bootstrap_server_address=address)

1005 + with pytest.raises(ValueError, match="bootstrap_server_address"):

1006 + get_mooncake_bootstrap_addr(config)

1007 +

1008 +

1009 def test_scheduler_request_finished():

1010 """Tests the scheduler-side logic when a request finishes.

1011

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

+15

@@ -2329,6 +2329,10 @@ def _async_loop(loop: asyncio.AbstractEventLoop):

2329

2330

2331 def should_launch_bootstrap_server(vllm_config: VllmConfig) -> bool:

2332 + if (kv_config := vllm_config.kv_transfer_config) is not None and (

2333 + kv_config.get_from_extra_config("bootstrap_server_address", None) is not None

2334 + ):

2335 + return False

2336 assert (parallel_config := vllm_config.parallel_config)

2337 # Only the TP=0, PP=0 worker of the designated engine should launch it.

2338 if get_tensor_model_parallel_rank() != 0:

@@ -2351,6 +2355,17 @@ def get_mooncake_bootstrap_addr(vllm_config: VllmConfig) -> tuple[str, int]:

2355 This is only used by prefillers to register workers.

2356 Decoders should get addr from kv_transfer_params.

2357 """

2358 + if (kv_config := vllm_config.kv_transfer_config) is not None:

2359 + address = kv_config.get_from_extra_config("bootstrap_server_address", None)

2360 + if address is not None:

2361 + from vllm.utils.network_utils import split_host_port

2362 +

2363 + host, port = split_host_port(address)

2364 + if not host or not 1 <= port <= 65535:

2365 + raise ValueError(

2366 + "bootstrap_server_address must specify a host and port"

2367 + )

2368 + return host, port

2369 assert (parallel_config := vllm_config.parallel_config)

2370 if parallel_config.local_engines_only:

2371 # In hybrid or external LB mode, connect to local server.

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

+67 -19

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import socket

4 import threading

5 import time

6 from dataclasses import dataclass

@@ -32,9 +33,7 @@ class EngineEntry:

35 class MooncakeBootstrapServer:

35 - """A centralized server running on the global rank 0 prefiller worker.

36 - Prefiller workers register their connection info (IP, port, ranks) here.

37 - """

36 + """A centralized registry for prefiller connection info (IP, port, ranks)."""

38 def __init__(self, host: str, port: int):

39 self.workers: dict[int, EngineEntry] = {}

@@ -45,6 +44,8 @@ class MooncakeBootstrapServer:

44 self._register_routes()

45 self.server_thread: threading.Thread | None = None

46 self.server: uvicorn.Server | None = None

47 + self._socket: socket.socket | None = None

48 + self._startup_error: BaseException | None = None

50 def __del__(self):

51 self.shutdown()

@@ -54,27 +55,74 @@ class MooncakeBootstrapServer:

55 self.app.post("/register")(self.register_worker)

56 self.app.get("/query", response_model=dict[int, EngineEntry])(self.query)

57 - def start(self):

58 + def start(self, *, timeout: float = 30.0):

59 if self.server_thread:

59 - return

60 + if (

61 + self.server_thread.is_alive()

62 + and self.server is not None

63 + and self.server.started

64 + and not self.server.should_exit

65 + and self._socket is not None

66 + and self._socket.fileno() >= 0

67 + ):

68 + return

69 + raise RuntimeError(

70 + "Mooncake bootstrap server is still starting or stopping"

71 + )

61 - config = uvicorn.Config(app=self.app, host=self.host, port=self.port)

62 - self.server = uvicorn.Server(config=config)

63 - self.server_thread = threading.Thread(

64 - target=self.server.run, name="mooncake_bootstrap_server", daemon=True

65 - )

66 - self.server_thread.start()

67 - while not self.server.started:

68 - time.sleep(0.1) # Wait for the server to start

73 + try:

74 + # Bind in the caller so errors propagate, and retain the listener

75 + # when handing it to Uvicorn, including for automatically chosen ports.

76 + family = socket.AF_INET6 if ":" in self.host else socket.AF_INET

77 + self._socket = socket.create_server((self.host, self.port), family=family)

78 + self.port = self._socket.getsockname()[1]

79 + config = uvicorn.Config(app=self.app, host=self.host, port=self.port)

80 + server = self.server = uvicorn.Server(config=config)

81 + listener = self._socket

82 + self._startup_error = None

83 +

84 + def run():

85 + try:

86 + server.run(sockets=[listener])

87 + except BaseException as exc:

88 + self._startup_error = exc

89 +

90 + self.server_thread = threading.Thread(

91 + target=run, name="mooncake_bootstrap_server", daemon=True

92 + )

93 + self.server_thread.start()

94 + deadline = time.monotonic() + timeout

95 + while not server.started:

96 + if not self.server_thread.is_alive():

97 + raise RuntimeError(

98 + "Mooncake bootstrap server exited during startup"

99 + ) from self._startup_error

100 + if time.monotonic() >= deadline:

101 + raise TimeoutError(

102 + "Mooncake bootstrap server did not start in time"

103 + )

104 + time.sleep(0.01)

105 + except BaseException:

106 + self.shutdown()

107 + raise

108 logger.info("Mooncake Bootstrap Server started at %s:%d", self.host, self.port)

109

110 def shutdown(self):

72 - if self.server_thread is None or self.server is None or not self.server.started:

73 - return

74 -

75 - self.server.should_exit = True

76 - self.server_thread.join()

77 - logger.info("Mooncake Bootstrap Server stopped.")

111 + was_started = self.server is not None and self.server.started

112 + if self.server is not None:

113 + self.server.should_exit = True

114 + if self.server_thread is not None and self.server_thread.ident is not None:

115 + self.server_thread.join(timeout=5)

116 + if self._socket is not None:

117 + self._socket.close()

118 + self._socket = None

119 + if self.server_thread is not None and self.server_thread.is_alive():

120 + logger.warning("Mooncake bootstrap server did not stop in time")

121 + else:

122 + self.server_thread = None

123 + self.server = None

124 + if was_started:

125 + logger.info("Mooncake Bootstrap Server stopped.")

126

127 async def register_worker(self, payload: RegisterWorkerPayload):

128 """Handles registration of a prefiller worker."""