From da91d03e13cc2daa5a3b744d0bafc006e270366e Mon Sep 17 00:00:00 2001 From: Chaemin Lim Date: Fri, 3 Jul 2026 02:25:37 +0000 Subject: [PATCH 1/2] =?UTF-8?q?[Core][KV-transfer]=20MoRIIO:=20multi-node?= =?UTF-8?q?=20TP=20prefill=E2=86=92decode=20dispatch=20via=20published=20h?= =?UTF-8?q?ost=20list?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In a multi-node TP prefill deployment, each decode worker must open the RDMA handshake and post-transfer notify against the prefill node that actually owns its KV slice, not the prefill head. Add a published node-host list (kv_connector_extra_config["moriio_node_hosts"]) with parsing, trusted-host validation, and rank->host selection (remote_hosts[tp_rank // ranks_per_node], DP-rank aware). Falls back to single-host behaviour when the list is unset or length <= 1, preserving TP=8 / monolithic setups. Signed-off-by: Chaemin Lim Co-authored-by: Edwin Lim --- .../moriio_toy_proxy_server.py | 14 + .../unit/test_moriio_multinode_hosts.py | 534 ++++++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 148 ++++- .../v1/moriio/moriio_connector.py | 103 +++- 4 files changed, 782 insertions(+), 17 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index 07a462711d22..24637cf07014 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -75,6 +75,7 @@ def _listen_for_register(hostname, port): "dp_size": data["dp_size"], "tp_size": data["tp_size"], "transfer_mode": data["transfer_mode"], + "node_hosts": data.get("node_hosts", []), } # zmq_address format: "host:IP,handshake:PORT,notify:PORT" # Stored verbatim; embedded into the request_id by handle_request. @@ -277,6 +278,11 @@ async def handle_request(api: str, request: Request): decode_instance_endpoint["tp_size"] ) req_data_to_prefill["kv_transfer_params"]["transfer_id"] = transfer_id + decode_hosts = decode_instance_endpoint.get("node_hosts") or [] + if decode_hosts: + req_data_to_prefill["kv_transfer_params"]["remote_hosts"] = list( + decode_hosts + ) prefill_request_url = prefill_instance_endpoint["request_address"] + api send_prefill_task = asyncio.create_task( @@ -308,6 +314,13 @@ async def handle_request(api: str, request: Request): "remote_block_ids" ] req_data["kv_transfer_params"]["transfer_id"] = prefill_kv["transfer_id"] + req_data["kv_transfer_params"]["remote_hosts"] = prefill_kv.get( + "remote_hosts" + ) + + prefill_hosts = prefill_instance_endpoint.get("node_hosts") or [] + if prefill_hosts: + req_data["kv_transfer_params"]["remote_hosts"] = list(prefill_hosts) req_data["kv_transfer_params"]["remote_dp_size"] = prefill_instance_endpoint[ "dp_size" @@ -315,6 +328,7 @@ async def handle_request(api: str, request: Request): req_data["kv_transfer_params"]["remote_tp_size"] = prefill_instance_endpoint[ "tp_size" ] + req_data["kv_transfer_params"]["tp_size"] = prefill_instance_endpoint["tp_size"] if selected_prefill_dp_rank is not None: req_data["kv_transfer_params"]["remote_dp_rank"] = selected_prefill_dp_rank diff --git a/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py b/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py new file mode 100644 index 000000000000..758bdcc6005b --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py @@ -0,0 +1,534 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import copy +from types import SimpleNamespace + +import pytest + +from vllm.config import KVTransferConfig +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( + MoRIIOConnectorMetadata, + MoRIIOConstants, + MoRIIOMode, + get_moriio_node_hosts, + get_moriio_trusted_remote_hosts, + validate_moriio_remote_host, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorScheduler, + MoRIIOConnectorWorker, + _pick_remote_rank_host, +) + + +@pytest.mark.parametrize( + ("node_hosts", "expected"), + [ + ("prefill-a, prefill-b, ", ["prefill-a", "prefill-b"]), + (["prefill-a", " prefill-b ", ""], ["prefill-a", "prefill-b"]), + (None, ["local-host"]), + ], +) +def test_moriio_node_hosts_come_from_explicit_config_or_local_fallback( + node_hosts, expected +): + config = KVTransferConfig( + kv_connector_extra_config={} + if node_hosts is None + else {"node_hosts": node_hosts} + ) + + assert get_moriio_node_hosts(config, "local-host") == expected + + +def test_moriio_trusted_remote_hosts_use_explicit_config_or_node_hosts(): + config = KVTransferConfig( + kv_connector_extra_config={ + "trusted_remote_hosts": "peer-a, peer-b", + "node_hosts": ["local-a"], + } + ) + assert get_moriio_trusted_remote_hosts(config, ["local-a"]) == frozenset( + {"peer-a", "peer-b"} + ) + + config = KVTransferConfig(kv_connector_extra_config={}) + assert get_moriio_trusted_remote_hosts(config, ["local-a"]) == frozenset( + {"local-a"} + ) + + +@pytest.mark.parametrize( + ("tp_rank", "expected"), + [ + (0, "prefill-a"), + (7, "prefill-a"), + (8, "prefill-b"), + (15, "prefill-b"), + ], +) +def test_remote_rank_host_uses_tp_rank_for_multi_host_tp(tp_rank, expected): + assert ( + _pick_remote_rank_host( + "fallback", + ["prefill-a", "prefill-b"], + tp_size=16, + tp_rank=tp_rank, + ) + == expected + ) + + +@pytest.mark.parametrize( + ("remote_dp_rank", "expected"), + [ + (0, "prefill-a"), + (7, "prefill-a"), + (8, "prefill-b"), + (15, "prefill-b"), + ], +) +def test_remote_rank_host_uses_dp_rank_when_remote_dp_spans_hosts( + remote_dp_rank, expected +): + assert ( + _pick_remote_rank_host( + "fallback", + ["prefill-a", "prefill-b"], + tp_size=1, + tp_rank=0, + remote_dp_size=16, + remote_dp_rank=remote_dp_rank, + ) + == expected + ) + + +def test_connector_metadata_uses_remote_hosts_for_plain_request_id(): + metadata = MoRIIOConnectorMetadata(trusted_remote_hosts={"prefill-a", "prefill-b"}) + + metadata.add_new_req( + request_id="plain-request-id", + local_block_ids=[1], + kv_transfer_params={ + "transfer_id": "tx-1", + "remote_block_ids": [2], + "remote_engine_id": "prefill-engine", + "remote_hosts": ["prefill-a", "prefill-b"], + "remote_tp_size": 16, + "remote_dp_size": 16, + "remote_dp_rank": 8, + }, + ) + + req_meta = metadata.reqs_to_recv["plain-request-id"] + assert req_meta.remote_host == "prefill-a" + assert req_meta.remote_handshake_port == int(MoRIIOConstants.DEFAULT_HANDSHAKE_PORT) + assert req_meta.remote_notify_port == int(MoRIIOConstants.DEFAULT_NOTIFY_PORT) + assert req_meta.remote_hosts == ["prefill-a", "prefill-b"] + assert req_meta.tp_size == 16 + assert req_meta.remote_dp_size == 16 + assert req_meta.remote_dp_rank == 8 + + +def test_connector_metadata_rejects_untrusted_remote_hosts(): + metadata = MoRIIOConnectorMetadata(trusted_remote_hosts={"prefill-a"}) + + with pytest.raises(ValueError, match="untrusted remote_hosts"): + metadata.add_new_req( + request_id="plain-request-id", + local_block_ids=[1], + kv_transfer_params={ + "transfer_id": "tx-1", + "remote_block_ids": [2], + "remote_engine_id": "prefill-engine", + "remote_hosts": ["prefill-a", "metadata.example"], + "remote_tp_size": 16, + }, + ) + + +def test_scheduler_rejects_untrusted_remote_hosts_before_notify(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.mode = MoRIIOMode.WRITE + scheduler.tp_size = 2 + scheduler.dp_rank = 0 + scheduler.transfer_id_to_request_id = {} + scheduler.request_id_to_transfer_id = {} + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + sent = [] + monkeypatch.setattr( + scheduler, + "send_notify_block", + lambda **kwargs: sent.append(kwargs), + ) + request = SimpleNamespace( + request_id=( + "___prefill_addr_host:prefill-a,handshake:6301,notify:61005" + "___decode_addr_host:decode-a,handshake:6301,notify:61005" + "_0123456789abcdef0123456789abcdef" + ), + kv_transfer_params={ + "transfer_id": "tx-1", + "do_remote_prefill": True, + "remote_hosts": ["prefill-a", "notify.example"], + }, + ) + + with pytest.raises(ValueError, match="untrusted remote_hosts"): + scheduler.update_state_after_alloc( + request, + SimpleNamespace(get_block_ids=lambda: ([1, 2],)), + 0, + ) + + assert sent == [] + + +@pytest.mark.asyncio +async def test_toy_proxy_forwards_node_hosts_to_prefill_and_decode(monkeypatch): + from examples.disaggregated.disaggregated_serving import ( + moriio_toy_proxy_server as proxy, + ) + + captured_prefill = {} + captured_decode = {} + prefill_hosts = ["prefill-a", "prefill-b"] + decode_hosts = ["decode-a", "decode-b"] + + class FakeRequest: + async def get_json(self): + return {"model": "test-model", "prompt": "hello", "max_tokens": 2} + + async def fake_send_request_to_prefill( + endpoint, req_data, request_id, selected_prefill_dp_rank + ): + captured_prefill.update( + { + "endpoint": endpoint, + "request_id": request_id, + "selected_prefill_dp_rank": selected_prefill_dp_rank, + "req_data": copy.deepcopy(req_data), + } + ) + return { + "kv_transfer_params": { + "remote_engine_id": "prefill-engine", + "remote_block_ids": [1, 2], + "transfer_id": req_data["kv_transfer_params"]["transfer_id"], + "remote_hosts": ["unused-read-mode-host"], + } + } + + async def fake_start_decode_request(endpoint, req_data, request_id): + captured_decode.update( + { + "endpoint": endpoint, + "request_id": request_id, + "req_data": copy.deepcopy(req_data), + } + ) + return object(), object() + + def fake_stream_decode_response(session, response, request_id): + async def stream(): + yield b"ok" + + return stream() + + async def fake_make_response(response): + return response + + monkeypatch.setattr( + proxy, + "prefill_instances", + [ + { + "request_address": "http://prefill.example/v1", + "zmq_address": "host:prefill-a,handshake:6301,notify:61005", + "dp_size": 1, + "tp_size": 16, + "node_hosts": prefill_hosts, + } + ], + ) + monkeypatch.setattr( + proxy, + "decode_instances", + [ + { + "request_address": "http://decode.example/v1", + "zmq_address": "host:decode-a,handshake:6301,notify:61005", + "dp_size": 1, + "tp_size": 16, + "node_hosts": decode_hosts, + } + ], + ) + monkeypatch.setattr(proxy, "TRANSFER_TYPE", "READ") + monkeypatch.setattr(proxy, "request_nums", 0) + monkeypatch.setattr(proxy, "send_request_to_prefill", fake_send_request_to_prefill) + monkeypatch.setattr(proxy, "start_decode_request", fake_start_decode_request) + monkeypatch.setattr(proxy, "stream_decode_response", fake_stream_decode_response) + monkeypatch.setattr(proxy, "make_response", fake_make_response) + + await proxy.handle_request("/chat/completions", FakeRequest()) + + assert captured_prefill["endpoint"] == ( + "http://prefill.example/v1/chat/completions" + ) + assert ( + captured_prefill["req_data"]["kv_transfer_params"]["remote_hosts"] + == decode_hosts + ) + assert captured_decode["endpoint"] == "http://decode.example/v1/chat/completions" + assert ( + captured_decode["req_data"]["kv_transfer_params"]["remote_hosts"] + == prefill_hosts + ) + assert captured_decode["req_data"]["kv_transfer_params"]["tp_size"] == 16 + + +def _multi_host_recv_meta(): + """Decode-side ReqMeta for a 2-node prefill peer (TP=16, DP=16).""" + metadata = MoRIIOConnectorMetadata(trusted_remote_hosts={"prefill-a", "prefill-b"}) + metadata.add_new_req( + request_id="glue-multihost", + local_block_ids=[1], + kv_transfer_params={ + "transfer_id": "tx-1", + "remote_block_ids": [2], + "remote_engine_id": "prefill-engine", + "remote_hosts": ["prefill-a", "prefill-b"], + "remote_tp_size": 16, + "remote_dp_size": 16, + }, + ) + return metadata.reqs_to_recv["glue-multihost"] + + +@pytest.mark.parametrize( + ("tp_rank", "expected"), + [ + (0, "prefill-a"), + (8, "prefill-b"), + (15, "prefill-b"), + ], +) +def test_worker_pick_remote_host_routes_by_tp_rank(tp_rank, expected): + # Proves _pick_remote_host feeds self.tp_rank + meta.remote_hosts into the + # pure picker (TP branch), so each decode worker targets its own peer node. + meta = _multi_host_recv_meta() + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.tp_rank = tp_rank + assert worker._pick_remote_host(meta) == expected + + +@pytest.mark.parametrize( + ("dp_rank", "expected"), + [ + (0, "prefill-a"), + (8, "prefill-b"), + (15, "prefill-b"), + ], +) +def test_worker_pick_host_for_dp_rank_routes_by_dp_rank(dp_rank, expected): + # Proves _pick_host_for_dp_rank forwards meta.remote_dp_size + the dp_rank + # arg (DP branch). tp_rank is pinned to 0 so only dp_rank can move the host. + meta = _multi_host_recv_meta() + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.tp_rank = 0 + assert worker._pick_host_for_dp_rank(meta, dp_rank) == expected + + +@pytest.mark.parametrize("tp_rank", [0, 7]) +def test_worker_pick_remote_host_single_host_passthrough(tp_rank): + # Proves the TP<=8 / single-host path is untouched: no host list means + # _pick_remote_host returns meta.remote_host unchanged for any tp_rank. + metadata = MoRIIOConnectorMetadata() + metadata.add_new_req( + request_id="glue-singlehost", + local_block_ids=[1], + kv_transfer_params={ + "transfer_id": "tx-1", + "remote_block_ids": [2], + "remote_engine_id": "prefill-engine", + "remote_host": "prefill-solo", + "remote_handshake_port": 6301, + "remote_notify_port": 61005, + "remote_tp_size": 8, + }, + ) + meta = metadata.reqs_to_recv["glue-singlehost"] + assert meta.remote_hosts is None + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.tp_rank = tp_rank + assert worker._pick_remote_host(meta) == "prefill-solo" + + +def test_validate_moriio_remote_host(): + trusted = frozenset({"prefill-a"}) + key = "kv_transfer_params['remote_host']" + + # (a) value=None: singular path is optional, so a missing host is a no-op. + validate_moriio_remote_host(None, trusted, key) + # (b) no trust list configured: default single-host flow stays unenforced. + validate_moriio_remote_host("anything.example", (), key) + # (c) host present in the trust list is accepted. + validate_moriio_remote_host("prefill-a", trusted, key) + # (d) host absent from a non-empty trust list is rejected. + with pytest.raises(ValueError, match="untrusted host"): + validate_moriio_remote_host("evil.example", trusted, key) + + +def test_connector_metadata_rejects_untrusted_remote_host_singular(): + metadata = MoRIIOConnectorMetadata(trusted_remote_hosts={"prefill-a"}) + + # Ports are supplied so the request_id parse branch is skipped: without the + # singular-host guard the untrusted host would be accepted outright. + with pytest.raises(ValueError, match="untrusted host"): + metadata.add_new_req( + request_id="plain-request-id", + local_block_ids=[1], + kv_transfer_params={ + "transfer_id": "tx-1", + "remote_block_ids": [2], + "remote_engine_id": "prefill-engine", + "remote_host": "evil.example", + "remote_handshake_port": 6301, + "remote_notify_port": 61005, + "remote_tp_size": 16, + }, + ) + + +def test_connector_metadata_accepts_trusted_remote_host_singular(): + metadata = MoRIIOConnectorMetadata(trusted_remote_hosts={"prefill-a"}) + + metadata.add_new_req( + request_id="plain-request-id", + local_block_ids=[1], + kv_transfer_params={ + "transfer_id": "tx-1", + "remote_block_ids": [2], + "remote_engine_id": "prefill-engine", + "remote_host": "prefill-a", + "remote_handshake_port": 6301, + "remote_notify_port": 61005, + "remote_tp_size": 16, + }, + ) + + # A trusted singular host survives validation and is registered as the peer. + assert metadata.reqs_to_recv["plain-request-id"].remote_host == "prefill-a" + + +def test_scheduler_release_write_rejects_untrusted_remote_host(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler.tp_size = 1 + sent = [] + monkeypatch.setattr( + scheduler, + "_send_transfer_release", + lambda *args, **kwargs: sent.append((args, kwargs)), + ) + + # remote_notify_port present -> request_id parse is skipped; the guard must + # fail closed before any release dial reaches the untrusted host. + with pytest.raises(ValueError, match="untrusted host"): + scheduler._release_write_prefill_blocks( + "req", + { + "transfer_id": "tx", + "remote_host": "evil.example", + "remote_notify_port": 6100, + }, + ) + + assert sent == [] + + +def test_scheduler_release_write_sends_for_trusted_remote_host(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler.tp_size = 1 + sent = [] + monkeypatch.setattr( + scheduler, + "_send_transfer_release", + lambda transfer_id, host, port: sent.append((transfer_id, host, port)), + ) + + scheduler._release_write_prefill_blocks( + "req", + {"transfer_id": "tx", "remote_host": "prefill-a", "remote_notify_port": 6100}, + ) + + # A trusted host reaches the release dial with the resolved transfer target. + assert sent == [("tx", "prefill-a", 6100)] + + +def test_scheduler_release_write_none_host_uses_request_id(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler.tp_size = 1 + sent = [] + monkeypatch.setattr( + scheduler, + "_send_transfer_release", + lambda transfer_id, host, port: sent.append((transfer_id, host, port)), + ) + + request_id = ( + "___prefill_addr_host:prefill-a,handshake:6301,notify:6100" + "___decode_addr_host:decode-a,handshake:6301,notify:6100" + "_0123456789abcdef0123456789abcdef" + ) + + # remote_host omitted: the guard returns early (no false reject) and the + # notify address is recovered from the request_id, then dialed. + scheduler._release_write_prefill_blocks(request_id, {"transfer_id": "tx"}) + + assert sent == [("tx", "prefill-a", 6100)] + + +def test_scheduler_update_state_rejects_untrusted_remote_host_singular(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.mode = MoRIIOMode.WRITE + scheduler.tp_size = 2 + scheduler.dp_rank = 0 + scheduler.transfer_id_to_request_id = {} + scheduler.request_id_to_transfer_id = {} + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + sent = [] + monkeypatch.setattr( + scheduler, + "send_notify_block", + lambda **kwargs: sent.append(kwargs), + ) + request = SimpleNamespace( + request_id=( + "___prefill_addr_host:prefill-a,handshake:6301,notify:61005" + "___decode_addr_host:decode-a,handshake:6301,notify:61005" + "_0123456789abcdef0123456789abcdef" + ), + kv_transfer_params={ + "transfer_id": "tx-1", + "do_remote_prefill": True, + "remote_host": "evil.example", + "remote_notify_port": 61005, + }, + ) + + # remote_notify_port present -> parse skipped; the singular host must be + # rejected before any notify is emitted to the untrusted peer. + with pytest.raises(ValueError, match="untrusted host"): + scheduler.update_state_after_alloc( + request, + SimpleNamespace(get_block_ids=lambda: ([1, 2],)), + 0, + ) + + assert sent == [] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 7b8ab566058a..e803bed89ff5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -4,7 +4,7 @@ import os import threading import time -from collections.abc import Iterator +from collections.abc import Collection, Iterator from dataclasses import dataclass from typing import TYPE_CHECKING, Any, NamedTuple @@ -203,6 +203,82 @@ def resolve_host_ip(extra_config: dict) -> str: return extra_config.get("host_ip") or get_ip() +def _normalize_node_hosts(value: Any, config_key: str) -> list[str]: + if value is None: + return [] + if isinstance(value, str): + return [host.strip() for host in value.split(",") if host.strip()] + try: + return [str(host).strip() for host in value if str(host).strip()] + except TypeError as exc: + raise ValueError( + f"{config_key} must be a comma-separated string or iterable of hosts" + ) from exc + + +def get_moriio_node_hosts( + kv_transfer_config: KVTransferConfig, default_host: str +) -> list[str]: + extra_config = kv_transfer_config.kv_connector_extra_config + node_hosts = _normalize_node_hosts( + extra_config.get("node_hosts"), + "kv_connector_extra_config['node_hosts']", + ) + if node_hosts: + return node_hosts + + return [default_host] + + +def get_moriio_trusted_remote_hosts( + kv_transfer_config: KVTransferConfig, + node_hosts: Collection[str], +) -> frozenset[str]: + extra_config = kv_transfer_config.kv_connector_extra_config + trusted_remote_hosts = _normalize_node_hosts( + extra_config.get("trusted_remote_hosts"), + "kv_connector_extra_config['trusted_remote_hosts']", + ) + return frozenset(trusted_remote_hosts or node_hosts) + + +def validate_moriio_remote_hosts( + value: Any, + trusted_remote_hosts: Collection[str], + config_key: str, +) -> list[str]: + remote_hosts = _normalize_node_hosts(value, config_key) + if not remote_hosts: + return [] + if not trusted_remote_hosts: + raise ValueError( + f"{config_key} requires kv_connector_extra_config['trusted_remote_hosts']" + ) + untrusted_hosts = sorted(set(remote_hosts).difference(trusted_remote_hosts)) + if untrusted_hosts: + raise ValueError( + f"{config_key} contains untrusted remote_hosts: {untrusted_hosts}" + ) + return remote_hosts + + +def validate_moriio_remote_host( + value: str | None, + trusted_remote_hosts: Collection[str], + config_key: str, +) -> None: + """Reject a directly-supplied peer host that is not trusted. + + Mirrors ``validate_moriio_remote_hosts`` for the single-host path. Only + enforced when ``trusted_remote_hosts`` is configured, so the default + single-host flow is unchanged when no trust list is set. + """ + if value is None or not trusted_remote_hosts: + return + if value not in trusted_remote_hosts: + raise ValueError(f"{config_key} contains untrusted host: {value!r}") + + _DEPRECATED_ENV_VARS: dict[str, str] = { "VLLM_MORIIO_CONNECTOR_READ_MODE": "read_mode", "VLLM_MORIIO_QP_PER_TRANSFER": "qp_per_transfer", @@ -244,6 +320,7 @@ class MoRIIOConfig: post_batch_size: int = -1 num_workers: int = 1 backend: str = "rdma" + node_hosts: list[str] = field(default_factory=list) @classmethod def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": @@ -301,8 +378,10 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": extra_config.get("defer_timeout", MoRIIOConstants.DEFAULT_DEFER_TIMEOUT) ) + local_ip = resolve_host_ip(extra_config) + return cls( - local_ip=resolve_host_ip(extra_config), + local_ip=local_ip, local_kv_port=get_open_port(), proxy_ip=extra_config["proxy_ip"], local_ping_port=get_open_port(), @@ -319,6 +398,7 @@ def from_vllm_config(cls, vllm_config: VllmConfig) -> "MoRIIOConfig": post_batch_size=int(extra_config.get("post_batch_size", -1)), num_workers=int(extra_config.get("num_workers", 1)), backend=backend, + node_hosts=get_moriio_node_hosts(kv_transfer_config, local_ip), transfer_timeout=transfer_timeout, defer_timeout=defer_timeout, ) @@ -422,14 +502,24 @@ class ReqMeta: remote_engine_id: str tp_size: int remote_dp_size: int + remote_dp_rank: int = 0 + # Ordered list of all prefill-instance host IPs for multi-node TP. + # Each decode worker picks remote_hosts[tp_rank // ranks_per_node] as its + # actual peer host for handshake + post-transfer notify. None or len<=1 + # falls back to single-host behaviour (remote_host). + remote_hosts: list[str] | None = None class MoRIIOConnectorMetadata(KVConnectorMetadata): - def __init__(self): + def __init__( + self, + trusted_remote_hosts: Collection[str] = (), + ): self.reqs_to_recv: dict[ReqId, ReqMeta] = {} self.reqs_to_save: dict[ReqId, ReqMeta] = {} self.reqs_to_send: dict[ReqId, float] = {} self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} + self.trusted_remote_hosts = frozenset(trusted_remote_hosts) def __repr__(self): return ( @@ -449,8 +539,18 @@ def add_new_req( transfer_id = kv_transfer_params["transfer_id"] remote_host = kv_transfer_params.get("remote_host") + validate_moriio_remote_host( + remote_host, + self.trusted_remote_hosts, + "kv_transfer_params['remote_host']", + ) remote_handshake_port = kv_transfer_params.get("remote_handshake_port") remote_notify_port = kv_transfer_params.get("remote_notify_port") + remote_hosts = validate_moriio_remote_hosts( + kv_transfer_params.get("remote_hosts"), + self.trusted_remote_hosts, + "kv_transfer_params['remote_hosts']", + ) if ( remote_host is None or remote_handshake_port is None @@ -459,22 +559,48 @@ def add_new_req( # Parse host/ports from the request_id. The router embeds both # zmq_addresses in PD request IDs, but WRITE decode requests may carry # a plain request ID and get the remote address via kv_transfer_params. - peer_zmq = get_peer_zmq_from_request_id(request_id, is_producer=write_mode) - remote_host, remote_handshake_port, remote_notify_port = ( - parse_moriio_zmq_address(peer_zmq) - ) - + try: + peer_zmq = get_peer_zmq_from_request_id( + request_id, is_producer=write_mode + ) + remote_host, remote_handshake_port, remote_notify_port = ( + parse_moriio_zmq_address(peer_zmq) + ) + except ValueError: + # Normalize remote_hosts: callers may pass a list (per-rank + # host vector from the proxy) or a single host string from + # older proxy versions. A bare string would silently slice + # into "1." for "172.30.0.1" if we just did remote_hosts[0]. + if not remote_hosts: + raise ValueError( + f"MoRIIO add_new_req: could not resolve peer host/ports " + f"for {request_id!r}; neither request_id parse nor " + f"kv_transfer_params.remote_hosts provided them" + ) from None + remote_host = remote_hosts[0] + remote_handshake_port = int(MoRIIOConstants.DEFAULT_HANDSHAKE_PORT) + remote_notify_port = int(MoRIIOConstants.DEFAULT_NOTIFY_PORT) + + # If remote block metadata is absent, use empty defaults so the request + # remains a no-op. _req = ReqMeta( transfer_id=transfer_id, local_block_ids=local_block_ids, - remote_block_ids=kv_transfer_params["remote_block_ids"], - remote_engine_id=kv_transfer_params["remote_engine_id"], + remote_block_ids=kv_transfer_params.get("remote_block_ids", []), + remote_engine_id=kv_transfer_params.get("remote_engine_id", ""), remote_host=remote_host, remote_port=int(remote_handshake_port), remote_handshake_port=int(remote_handshake_port), remote_notify_port=int(remote_notify_port), - tp_size=kv_transfer_params.get("tp_size", 1), + # Callers may forward `remote_tp_size` without `tp_size`. + # Prefer `tp_size`, then fall back to `remote_tp_size`, then 1. + tp_size=kv_transfer_params.get( + "tp_size", + kv_transfer_params.get("remote_tp_size", 1), + ), remote_dp_size=kv_transfer_params.get("remote_dp_size", 1), + remote_dp_rank=int(kv_transfer_params.get("remote_dp_rank", 0) or 0), + remote_hosts=remote_hosts or None, ) if write_mode: self.reqs_to_save[request_id] = _req 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 172836f5cc3f..363a85cf2680 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 @@ -37,12 +37,16 @@ TransferId, WriteTask, get_moriio_mode, + get_moriio_node_hosts, + get_moriio_trusted_remote_hosts, get_peer_zmq_from_request_id, get_port_offset, get_role, parse_moriio_zmq_address, resolve_host_ip, set_role, + validate_moriio_remote_host, + validate_moriio_remote_hosts, zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( @@ -186,6 +190,28 @@ def resolve_moriio_transfer_ack( return transfer_id +def _pick_remote_rank_host( + default_host: str, + remote_hosts: list[str] | None, + tp_size: int, + tp_rank: int, + remote_dp_size: int = 1, + remote_dp_rank: int = 0, +) -> str: + if remote_hosts and len(remote_hosts) > 1: + n_hosts = len(remote_hosts) + if int(remote_dp_size) >= n_hosts: + dp_per_node = max(1, int(remote_dp_size) // n_hosts) + node_idx = int(remote_dp_rank) // dp_per_node + if 0 <= node_idx < n_hosts: + return remote_hosts[node_idx] + ranks_per_node = max(1, int(tp_size) // n_hosts) + node_idx = int(tp_rank) // ranks_per_node + if 0 <= node_idx < n_hosts: + return remote_hosts[node_idx] + return default_host + + class MoRIIOConnector(KVConnectorBase_V1): def __init__( self, @@ -364,10 +390,22 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.host_ip = resolve_host_ip( self.kv_transfer_config.kv_connector_extra_config ) + # Multi-node TP: node_hosts holds the ordered host IPs in this + # engine's TP group (rank 0 first). Surfaced to the peer side via + # request_finished's kv_transfer_params so workers can dial the rank + # owner. Single-node TP falls back to [host_ip]. + self.node_hosts = get_moriio_node_hosts(self.kv_transfer_config, self.host_ip) + self.trusted_remote_hosts = get_moriio_trusted_remote_hosts( + self.kv_transfer_config, self.node_hosts + ) self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[ "handshake_port" ] - logger.info("Initializing MoRIIO Scheduler engine_id = %s", engine_id) + logger.info( + "Initializing MoRIIO Scheduler engine_id = %s node_hosts = %s", + engine_id, + self.node_hosts, + ) self.side_notify_port = self.kv_transfer_config.kv_connector_extra_config[ "notify_port" @@ -509,6 +547,11 @@ def _release_write_prefill_blocks(self, request_id: ReqId, params: dict[str, Any remote_dp_rank = params.get("remote_dp_rank", 0) remote_host = params.get("remote_host") + validate_moriio_remote_host( + remote_host, + self.trusted_remote_hosts, + "kv_transfer_params['remote_host']", + ) remote_notify_port = params.get("remote_notify_port") if remote_host is None or remote_notify_port is None: try: @@ -581,6 +624,11 @@ def update_state_after_alloc( remote_dp_rank = request.kv_transfer_params.get("remote_dp_rank", 0) remote_host = request.kv_transfer_params.get("remote_host") + validate_moriio_remote_host( + remote_host, + self.trusted_remote_hosts, + "kv_transfer_params['remote_host']", + ) remote_notify_port = request.kv_transfer_params.get( "remote_notify_port" ) @@ -595,16 +643,30 @@ def update_state_after_alloc( block_ids = blocks.get_block_ids()[0] + remote_hosts = validate_moriio_remote_hosts( + request.kv_transfer_params.get("remote_hosts"), + self.trusted_remote_hosts, + "kv_transfer_params['remote_hosts']", + ) + for tp_index in range(self.tp_size): target_port = remote_notify_port + get_port_offset( remote_dp_rank, tp_index ) + target_host = _pick_remote_rank_host( + remote_host, + remote_hosts, + self.tp_size, + tp_index, + int(request.kv_transfer_params.get("remote_dp_size", 1)), + int(remote_dp_rank), + ) self.send_notify_block( req_id=request.request_id, transfer_id=request.kv_transfer_params["transfer_id"], block_notify_list=block_ids, - host=remote_host, + host=target_host, port=target_port, ) @@ -616,7 +678,7 @@ def build_connector_meta( self, scheduler_output: SchedulerOutput, ) -> KVConnectorMetadata: - meta = MoRIIOConnectorMetadata() + meta = MoRIIOConnectorMetadata(self.trusted_remote_hosts) meta.transfer_id_to_request_id = self.transfer_id_to_request_id if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.PRODUCER: @@ -764,6 +826,10 @@ def request_finished( remote_dp_size=self.vllm_config.parallel_config.data_parallel_size, tp_size=self.vllm_config.parallel_config.tensor_parallel_size, transfer_id=params["transfer_id"], + # Multi-node TP: list of all prefill-instance host IPs in this + # engine's TP group (rank 0 first). Decode workers use this to + # pick the correct producer host per their tp_rank. + remote_hosts=self.node_hosts, ) def update_connector_output(self, connector_output: KVConnectorOutput) -> None: @@ -870,6 +936,7 @@ def __init__( self.http_port = self.moriio_config.http_port self.handshake_port = self.moriio_config.handshake_port self.notify_port = self.moriio_config.notify_port + self.node_hosts = self.moriio_config.node_hosts self.zmq_context = zmq.Context() self.metadata_address = ( @@ -1133,6 +1200,7 @@ def _ping(self, zmq_context): # READ (prefill-then-decode, sequential) from WRITE (concurrent) # scheduling. "transfer_mode": self.mode.name, + "node_hosts": list(self.node_hosts), } sock.send(msgpack.dumps(data)) @@ -1324,6 +1392,22 @@ def _moriio_handshake( def _remote_tp_rank(self, remote_tp_size: int) -> int: return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size) + def _pick_remote_host(self, meta: ReqMeta) -> str: + """Resolve the per-worker peer host for multi-node TP prefill-decode.""" + return _pick_remote_rank_host( + meta.remote_host, meta.remote_hosts, int(meta.tp_size), self.tp_rank + ) + + def _pick_host_for_dp_rank(self, meta: ReqMeta, dp_rank: int) -> str: + return _pick_remote_rank_host( + meta.remote_host, + meta.remote_hosts, + int(meta.tp_size), + self.tp_rank, + int(meta.remote_dp_size), + dp_rank, + ) + def _background_moriio_handshake( self, req_id: ReqId, remote_engine_id: EngineId, meta: ReqMeta ): @@ -1332,7 +1416,6 @@ def _background_moriio_handshake( if remote_engine_id is not None: fut = self._handshake_futures.get(remote_engine_id) if fut is None: - host = meta.remote_host port = int(meta.remote_handshake_port) tp_size = int(meta.tp_size) remote_dp_size = int(meta.remote_dp_size) @@ -1349,6 +1432,7 @@ def request_ready(_f: Future[Any], entry=(req_id, meta)): for cur_dp_rank in range(remote_dp_size): dp_engine_id = self.get_engine_name_with_dp(remote_engine_id, cur_dp_rank) + host = self._pick_host_for_dp_rank(meta, cur_dp_rank) future = self._handshake_initiation_executor.submit( self._moriio_handshake, host, port, tp_size, dp_engine_id, cur_dp_rank ) @@ -1793,13 +1877,20 @@ def _read_blocks_for_req(self, req_id: str, meta: ReqMeta): meta.remote_engine_id, req_id, ) + # Multi-node TP: remote_host here is used by _read_blocks to record the + # post-transfer notify callback address (so prefill can free its KV + # blocks). For multi-node prefill the notify must go to the prefill + # node that actually owns this worker's KV slice, not the prefill + # head. _pick_remote_host returns meta.remote_host unchanged for + # single-host setups, preserving TP=8 / monolithic behaviour. + actual_remote_host = self._pick_remote_host(meta) self._read_blocks( request_id=req_id, transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, local_block_ids=meta.local_block_ids, remote_block_ids=meta.remote_block_ids, - remote_host=meta.remote_host, + remote_host=actual_remote_host, remote_notify_port=meta.remote_notify_port, remote_tp_size=meta.tp_size, ) @@ -1814,7 +1905,7 @@ def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_lay layer_name=layer_name, kv_layer=kv_layer, remote_notify_port=meta.remote_notify_port, - remote_ip=meta.remote_host, + remote_ip=self._pick_remote_host(meta), ) def merge_contiguous_blocks( From fa888fec2e1d2f7928cad4788fb94f9106e29f4b Mon Sep 17 00:00:00 2001 From: Chaemin Lim Date: Tue, 7 Jul 2026 07:41:15 +0000 Subject: [PATCH 2/2] [Core][KV-transfer] MoRIIO: validate request_id-derived peer host X-Request-Id is client-controllable, so the peer host parsed from the request_id is now checked against trusted_remote_hosts too, not just the directly-supplied remote_host / remote_hosts. Validation is opt-in: when trusted_remote_hosts is not configured the request_id-derived host is left unvalidated (default flow unchanged); when it is configured, the host must be in the allowlist (the configured peers plus this instance's own node_hosts). Signed-off-by: Chaemin Lim --- .../unit/test_moriio_multinode_hosts.py | 215 ++++++++++++++++++ .../kv_connector/v1/moriio/moriio_common.py | 33 +++ .../v1/moriio/moriio_connector.py | 18 +- 3 files changed, 265 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py b/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py index 758bdcc6005b..b82cc97ab4db 100644 --- a/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py +++ b/tests/v1/kv_connector/unit/test_moriio_multinode_hosts.py @@ -12,6 +12,7 @@ MoRIIOConstants, MoRIIOMode, get_moriio_node_hosts, + get_moriio_request_id_trusted_hosts, get_moriio_trusted_remote_hosts, validate_moriio_remote_host, ) @@ -157,6 +158,7 @@ def test_scheduler_rejects_untrusted_remote_hosts_before_notify(monkeypatch): scheduler.transfer_id_to_request_id = {} scheduler.request_id_to_transfer_id = {} scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset() sent = [] monkeypatch.setattr( scheduler, @@ -453,6 +455,7 @@ def test_scheduler_release_write_rejects_untrusted_remote_host(monkeypatch): def test_scheduler_release_write_sends_for_trusted_remote_host(monkeypatch): scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset() scheduler.tp_size = 1 sent = [] monkeypatch.setattr( @@ -473,6 +476,7 @@ def test_scheduler_release_write_sends_for_trusted_remote_host(monkeypatch): def test_scheduler_release_write_none_host_uses_request_id(monkeypatch): scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset() scheduler.tp_size = 1 sent = [] monkeypatch.setattr( @@ -532,3 +536,214 @@ def test_scheduler_update_state_rejects_untrusted_remote_host_singular(monkeypat ) assert sent == [] + + +def _request_id_with_prefill_host(host): + """Router request_id whose embedded *prefill* peer host is ``host``. + + Consumer-side resolution (``is_producer=False``) matches ``_PREFILL_ZMQ_RE`` + and parses the prefill zmq_address, so ``host`` becomes the + request_id-derived ``remote_host`` (notify port 6100). + """ + return ( + f"___prefill_addr_host:{host},handshake:6301,notify:6100" + "___decode_addr_host:decode-a,handshake:6301,notify:6100" + "_0123456789abcdef0123456789abcdef" + ) + + +def test_request_id_trusted_hosts_empty_when_unconfigured(): + # Opt-in: without an explicit trusted_remote_hosts the allowlist is empty, + # which makes the request_id-derived validate a no-op (default flow intact). + config = KVTransferConfig(kv_connector_extra_config={}) + assert get_moriio_request_id_trusted_hosts(config, ["local-1"]) == frozenset() + + +def test_request_id_trusted_hosts_union_node_hosts_when_configured(): + # Configured: explicit peers UNION this instance's own node_hosts (the local + # host is trivially safe to accept from a request_id). + config = KVTransferConfig( + kv_connector_extra_config={"trusted_remote_hosts": "prefill-a"} + ) + assert get_moriio_request_id_trusted_hosts(config, ["local-1"]) == frozenset( + {"prefill-a", "local-1"} + ) + + +def test_add_new_req_rejects_untrusted_request_id_host(): + # No direct remote_host -> host is resolved from the client-controllable + # request_id; with a configured allowlist an untrusted host must be rejected. + metadata = MoRIIOConnectorMetadata(request_id_trusted_hosts={"prefill-a"}) + + with pytest.raises(ValueError, match="untrusted host"): + metadata.add_new_req( + request_id=_request_id_with_prefill_host("evil.example"), + local_block_ids=[1], + kv_transfer_params={"transfer_id": "tx-1"}, + ) + + assert metadata.reqs_to_recv == {} + + +def test_add_new_req_accepts_trusted_request_id_host(): + metadata = MoRIIOConnectorMetadata(request_id_trusted_hosts={"prefill-a"}) + request_id = _request_id_with_prefill_host("prefill-a") + + metadata.add_new_req( + request_id=request_id, + local_block_ids=[1], + kv_transfer_params={"transfer_id": "tx-1", "remote_block_ids": [2]}, + ) + + # A request_id host inside the allowlist survives and is registered. + assert metadata.reqs_to_recv[request_id].remote_host == "prefill-a" + + +def test_add_new_req_skips_request_id_validation_when_unconfigured(): + # No-regression: unconfigured allowlist (empty) -> the request_id-derived + # host is NOT validated, so an arbitrary embedded host flows through exactly + # as it did before the opt-in guard existed. + metadata = MoRIIOConnectorMetadata() + request_id = _request_id_with_prefill_host("evil.example") + + metadata.add_new_req( + request_id=request_id, + local_block_ids=[1], + kv_transfer_params={"transfer_id": "tx-1"}, + ) + + assert metadata.reqs_to_recv[request_id].remote_host == "evil.example" + + +def test_scheduler_release_rejects_untrusted_request_id_host(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset({"prefill-a"}) + scheduler.tp_size = 1 + sent = [] + monkeypatch.setattr( + scheduler, + "_send_transfer_release", + lambda *args, **kwargs: sent.append((args, kwargs)), + ) + + # remote_host omitted -> host recovered from the request_id; an untrusted + # host must fail closed before any release dial leaves the box. + with pytest.raises(ValueError, match="untrusted host"): + scheduler._release_write_prefill_blocks( + _request_id_with_prefill_host("evil.example"), + {"transfer_id": "tx"}, + ) + + assert sent == [] + + +def test_scheduler_release_accepts_trusted_request_id_host(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset({"prefill-a"}) + scheduler.tp_size = 1 + sent = [] + monkeypatch.setattr( + scheduler, + "_send_transfer_release", + lambda transfer_id, host, port: sent.append((transfer_id, host, port)), + ) + + scheduler._release_write_prefill_blocks( + _request_id_with_prefill_host("prefill-a"), + {"transfer_id": "tx"}, + ) + + # A request_id host inside the allowlist reaches the release dial. + assert sent == [("tx", "prefill-a", 6100)] + + +def test_scheduler_release_skips_request_id_validation_when_unconfigured(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + # Effective direct-host trust list is populated (its unconfigured default is + # node_hosts), but the request_id allowlist is empty. The request_id host + # must be checked against the EMPTY allowlist, not the effective one -- else + # the opt-in default flow would false-reject legitimate peers. + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset() + scheduler.tp_size = 1 + sent = [] + monkeypatch.setattr( + scheduler, + "_send_transfer_release", + lambda transfer_id, host, port: sent.append((transfer_id, host, port)), + ) + + scheduler._release_write_prefill_blocks( + _request_id_with_prefill_host("evil.example"), + {"transfer_id": "tx"}, + ) + + # Unconfigured request_id trust -> arbitrary embedded host flows through. + assert sent == [("tx", "evil.example", 6100)] + + +def test_scheduler_update_state_rejects_untrusted_request_id_host(monkeypatch): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.mode = MoRIIOMode.WRITE + scheduler.tp_size = 2 + scheduler.dp_rank = 0 + scheduler.transfer_id_to_request_id = {} + scheduler.request_id_to_transfer_id = {} + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset({"prefill-a"}) + sent = [] + monkeypatch.setattr( + scheduler, + "send_notify_block", + lambda **kwargs: sent.append(kwargs), + ) + request = SimpleNamespace( + request_id=_request_id_with_prefill_host("evil.example"), + kv_transfer_params={"transfer_id": "tx-1", "do_remote_prefill": True}, + ) + + # remote_host omitted -> resolved from request_id; the untrusted host must be + # rejected before any notify reaches the peer. + with pytest.raises(ValueError, match="untrusted host"): + scheduler.update_state_after_alloc( + request, + SimpleNamespace(get_block_ids=lambda: ([1, 2],)), + 0, + ) + + assert sent == [] + + +def test_scheduler_update_state_skips_request_id_validation_when_unconfigured( + monkeypatch, +): + scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) + scheduler.mode = MoRIIOMode.WRITE + scheduler.tp_size = 1 + scheduler.dp_rank = 0 + scheduler.transfer_id_to_request_id = {} + scheduler.request_id_to_transfer_id = {} + scheduler.trusted_remote_hosts = frozenset({"prefill-a"}) + scheduler._request_id_trusted_hosts = frozenset() + sent = [] + monkeypatch.setattr( + scheduler, + "send_notify_block", + lambda **kwargs: sent.append(kwargs), + ) + request = SimpleNamespace( + request_id=_request_id_with_prefill_host("evil.example"), + kv_transfer_params={"transfer_id": "tx-1", "do_remote_prefill": True}, + ) + + # No-regression: unconfigured request_id trust -> arbitrary embedded host is + # accepted and the notify path proceeds to the resolved peer. + scheduler.update_state_after_alloc( + request, + SimpleNamespace(get_block_ids=lambda: ([1, 2],)), + 0, + ) + + assert [kw["host"] for kw in sent] == ["evil.example"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index e803bed89ff5..5be4af452a73 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -279,6 +279,27 @@ def validate_moriio_remote_host( raise ValueError(f"{config_key} contains untrusted host: {value!r}") +def get_moriio_request_id_trusted_hosts( + kv_transfer_config: KVTransferConfig, + node_hosts: Collection[str], +) -> frozenset[str]: + """Allowlist for peer hosts parsed from the request_id. + + The request_id can carry a client-supplied ``X-Request-Id``, so a host + parsed from it is only trusted when the operator explicitly set + ``trusted_remote_hosts``. Returns an empty set when unconfigured (the + default flow is left unvalidated); otherwise the configured peers plus this + instance's own ``node_hosts`` (the local host is trivially safe). + """ + explicit = _normalize_node_hosts( + kv_transfer_config.kv_connector_extra_config.get("trusted_remote_hosts"), + "kv_connector_extra_config['trusted_remote_hosts']", + ) + if not explicit: + return frozenset() + return frozenset(explicit) | frozenset(node_hosts) + + _DEPRECATED_ENV_VARS: dict[str, str] = { "VLLM_MORIIO_CONNECTOR_READ_MODE": "read_mode", "VLLM_MORIIO_QP_PER_TRANSFER": "qp_per_transfer", @@ -514,12 +535,14 @@ class MoRIIOConnectorMetadata(KVConnectorMetadata): def __init__( self, trusted_remote_hosts: Collection[str] = (), + request_id_trusted_hosts: Collection[str] = (), ): self.reqs_to_recv: dict[ReqId, ReqMeta] = {} self.reqs_to_save: dict[ReqId, ReqMeta] = {} self.reqs_to_send: dict[ReqId, float] = {} self.transfer_id_to_request_id: dict[TransferId, ReqId] = {} self.trusted_remote_hosts = frozenset(trusted_remote_hosts) + self._request_id_trusted_hosts = frozenset(request_id_trusted_hosts) def __repr__(self): return ( @@ -581,6 +604,16 @@ def add_new_req( remote_handshake_port = int(MoRIIOConstants.DEFAULT_HANDSHAKE_PORT) remote_notify_port = int(MoRIIOConstants.DEFAULT_NOTIFY_PORT) + # request_id-derived peer host: the request_id can carry a client + # X-Request-Id, so validate it against the operator's allowlist when + # configured (empty set = unconfigured = no-op, keeping the default + # flow unchanged). Direct/plural hosts were already checked above. + validate_moriio_remote_host( + remote_host, + self._request_id_trusted_hosts, + "request_id-derived remote_host", + ) + # If remote block metadata is absent, use empty defaults so the request # remains a no-op. _req = ReqMeta( 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 363a85cf2680..ec915be4cdeb 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 @@ -38,6 +38,7 @@ WriteTask, get_moriio_mode, get_moriio_node_hosts, + get_moriio_request_id_trusted_hosts, get_moriio_trusted_remote_hosts, get_peer_zmq_from_request_id, get_port_offset, @@ -398,6 +399,9 @@ def __init__(self, vllm_config: VllmConfig, engine_id: str): self.trusted_remote_hosts = get_moriio_trusted_remote_hosts( self.kv_transfer_config, self.node_hosts ) + self._request_id_trusted_hosts = get_moriio_request_id_trusted_hosts( + self.kv_transfer_config, self.node_hosts + ) self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[ "handshake_port" ] @@ -565,6 +569,11 @@ def _release_write_prefill_blocks(self, request_id: ReqId, params: dict[str, Any ) return + validate_moriio_remote_host( + remote_host, + self._request_id_trusted_hosts, + "request_id-derived remote_host", + ) remote_notify_port = int(remote_notify_port) for tp_index in range(self.tp_size): target_port = remote_notify_port + get_port_offset(remote_dp_rank, tp_index) @@ -639,6 +648,11 @@ def update_state_after_alloc( remote_host, _, remote_notify_port = parse_moriio_zmq_address( peer_zmq ) + validate_moriio_remote_host( + remote_host, + self._request_id_trusted_hosts, + "request_id-derived remote_host", + ) remote_notify_port = int(remote_notify_port) block_ids = blocks.get_block_ids()[0] @@ -678,7 +692,9 @@ def build_connector_meta( self, scheduler_output: SchedulerOutput, ) -> KVConnectorMetadata: - meta = MoRIIOConnectorMetadata(self.trusted_remote_hosts) + meta = MoRIIOConnectorMetadata( + self.trusted_remote_hosts, self._request_id_trusted_hosts + ) meta.transfer_id_to_request_id = self.transfer_id_to_request_id if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.PRODUCER: