| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Routing and registration behaviour of the EPD proxy.""" |
| 4 | |
| 5 | import asyncio |
| 6 | from typing import Any |
| 7 | from unittest.mock import AsyncMock, Mock |
| 8 | |
| 9 | import msgspec |
| 10 | import pytest |
| 11 | from fastapi import HTTPException |
| 12 | from fastapi.testclient import TestClient |
| 13 | |
| 14 | from examples.disaggregated.disaggregated_encoder import ( |
| 15 | disagg_epd_proxy as proxy_module, |
| 16 | ) |
| 17 | from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( |
| 18 | EPDProxy, |
| 19 | EPDProxyConfig, |
| 20 | InstanceRecord, |
| 21 | InstanceRegistry, |
| 22 | InstanceRole, |
| 23 | build_app, |
| 24 | extract_mm_items, |
| 25 | ) |
| 26 | |
| 27 | ENCODE = InstanceRole.ENCODE |
| 28 | PREFILL = InstanceRole.PREFILL |
| 29 | DECODE = InstanceRole.DECODE |
| 30 | PD = InstanceRole.PREFILL_DECODE |
| 31 | |
| 32 | IMAGE_ITEM = {"type": "image_url", "image_url": {"url": "http://img/0.png"}} |
| 33 | |
| 34 | |
| 35 | @pytest.fixture |
| 36 | def proxy(): |
| 37 | registry = InstanceRegistry(probe_interval=0) |
| 38 | return EPDProxy(registry) |
| 39 | |
| 40 | |
| 41 | class TestRouting: |
| 42 | def test_no_decode_instance_is_service_unavailable(self, proxy): |
| 43 | """The proxy comes up before anything registers, so this is normal.""" |
| 44 | with pytest.raises(HTTPException) as excinfo: |
| 45 | proxy.route(num_items=0) |
| 46 | assert excinfo.value.status_code == 503 |
| 47 | |
| 48 | def test_media_without_an_encoder_is_service_unavailable(self, proxy): |
| 49 | proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) |
| 50 | assert proxy.route(num_items=0).decode.url == "http://d0:8000" |
| 51 | with pytest.raises(HTTPException) as excinfo: |
| 52 | proxy.route(num_items=1) |
| 53 | assert excinfo.value.status_code == 503 |
| 54 | |
| 55 | def test_prefill_is_optional(self, proxy): |
| 56 | """An E+PD deployment registers no prefill instance at all.""" |
| 57 | proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) |
| 58 | assert proxy.route(num_items=0).prefill is None |
| 59 | |
| 60 | def test_encoder_roster_tracks_new_instances(self, proxy): |
| 61 | proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) |
| 62 | for index in range(2): |
| 63 | proxy.registry.register(InstanceRecord(ENCODE, f"http://e{index}:8000")) |
| 64 | assert proxy.route(num_items=3).encoder_urls == [ |
| 65 | "http://e0:8000", |
| 66 | "http://e1:8000", |
| 67 | ] |
| 68 | proxy.registry.unregister("http://e0:8000") |
| 69 | assert proxy.route(num_items=3).encoder_urls == ["http://e1:8000"] |
| 70 | |
| 71 | @pytest.mark.parametrize("unregister", [False, True]) |
| 72 | def test_standalone_decode_always_requires_healthy_prefill(self, unregister): |
| 73 | app = build_app() |
| 74 | proxy = app.state.proxy |
| 75 | proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) |
| 76 | with pytest.raises(HTTPException) as excinfo: |
| 77 | proxy.route(num_items=0) |
| 78 | assert excinfo.value.status_code == 503 |
| 79 | proxy.registry.register(InstanceRecord(PREFILL, "http://p0:8000")) |
| 80 | route = proxy.route(num_items=0) |
| 81 | assert route.decode.url == "http://d0:8000" |
| 82 | assert route.consumer.url == "http://p0:8000" |
| 83 | if unregister: |
| 84 | proxy.registry.unregister("http://p0:8000") |
| 85 | else: |
| 86 | for _ in range(proxy.registry._fail_threshold): |
| 87 | proxy.registry._on_probe_failure(route.prefill, 0) |
| 88 | with pytest.raises(HTTPException) as excinfo: |
| 89 | proxy.route(num_items=0) |
| 90 | assert excinfo.value.status_code == 503 |
| 91 | |
| 92 | |
| 93 | class TestConsumerAddress: |
| 94 | """Which stage receives the embedding depends on the topology.""" |
| 95 | |
| 96 | def test_shared_storage_connectors_name_no_target(self, proxy): |
| 97 | """Rotate replicas without requiring a push endpoint.""" |
| 98 | proxy.registry.register(InstanceRecord(PD, "http://d0:8000", dp_size=2)) |
| 99 | route = proxy.route(num_items=0) |
| 100 | assert route.consumer_zmq is None |
| 101 | assert route.dp_rank == 0 |
| 102 | assert proxy.route(num_items=0).dp_rank == 1 |
| 103 | |
| 104 | def test_decode_is_the_consumer_when_prefill_is_not_split_out(self, proxy): |
| 105 | proxy.registry.register( |
| 106 | InstanceRecord(PD, "http://d0:8000", ec_zmq_addrs=["tcp://d0:20001"]) |
| 107 | ) |
| 108 | assert proxy.route(num_items=0).consumer_zmq == "tcp://d0:20001" |
| 109 | |
| 110 | def test_prefill_is_the_consumer_when_it_is_split_out(self, proxy): |
| 111 | """In E+P+D the prefill instance consumes the embedding, not decode. |
| 112 | |
| 113 | Reading the address out of whichever list it came from -- rather than |
| 114 | from the record that reported one -- sends the push to the wrong |
| 115 | instance in this topology. |
| 116 | """ |
| 117 | proxy.registry.register( |
| 118 | InstanceRecord(PREFILL, "http://p0:8000", ec_zmq_addrs=["tcp://p0:20001"]) |
| 119 | ) |
| 120 | proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) |
| 121 | assert proxy.route(num_items=0).consumer_zmq == "tcp://p0:20001" |
| 122 | |
| 123 | |
| 124 | class TestProxyApi: |
| 125 | @pytest.fixture |
| 126 | def client(self, monkeypatch): |
| 127 | monkeypatch.setenv("ADMIN_API_KEY", "test-key") |
| 128 | config = EPDProxyConfig(probe_interval=0) |
| 129 | with TestClient(build_app(config)) as client: |
| 130 | client.headers["X-API-Key"] = "test-key" |
| 131 | yield client |
| 132 | |
| 133 | def test_proxy_serves_before_anything_registers(self, client): |
| 134 | assert client.get("/health").status_code == 200 |
| 135 | assert client.get("/instances").json()["encode"]["live"] == [] |
| 136 | |
| 137 | def test_requests_are_refused_until_a_decoder_registers(self, client): |
| 138 | response = client.post( |
| 139 | "/v1/chat/completions", json={"model": "m", "messages": []} |
| 140 | ) |
| 141 | assert response.status_code == 503 |
| 142 | assert client.get("/v1/models").status_code == 503 |
| 143 | |
| 144 | def test_http_registration_keeps_consumer_replica_and_address_together( |
| 145 | self, client |
| 146 | ): |
| 147 | payload: dict[str, Any] = { |
| 148 | "role": "prefill_decode", |
| 149 | "url": "http://d0:8000/", |
| 150 | "dp_size": 2, |
| 151 | "ec_zmq_addrs": ["tcp://d0:20001", "tcp://d0:20003"], |
| 152 | } |
| 153 | assert client.post("/instances", json=payload).json() == {"registered": True} |
| 154 | assert client.post("/instances", json=payload).json() == {"registered": False} |
| 155 | proxy = client.app.state.proxy |
| 156 | for rank in [0, 1, 0]: |
| 157 | route = proxy.route(0) |
| 158 | assert route.dp_rank == rank |
| 159 | assert route.consumer_zmq == payload["ec_zmq_addrs"][rank] |
| 160 | assert client.delete("/instances", params={"url": payload["url"]}).json() == { |
| 161 | "removed": True |
| 162 | } |
| 163 | assert proxy.registry.pick(PD) is None |
| 164 | |
| 165 | def test_registration_rejects_invalid_topology(self, client): |
| 166 | cases: list[tuple[dict[str, Any], int]] = [ |
| 167 | ({"dp_size": 0}, 422), |
| 168 | ({"dp_size": 2, "ec_zmq_addrs": ["tcp://d0:20001"]}, 422), |
| 169 | ({"role": "decode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), |
| 170 | ({"role": "encode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), |
| 171 | ] |
| 172 | for fields, status in cases: |
| 173 | response = client.post( |
| 174 | "/instances", |
| 175 | json={"url": "http://d0:8000", "role": "prefill_decode", **fields}, |
| 176 | ) |
| 177 | assert response.status_code == status |
| 178 | |
| 179 | def test_registration_requires_admin_key(self, client): |
| 180 | client.headers.pop("X-API-Key") |
| 181 | assert ( |
| 182 | client.post( |
| 183 | "/instances", json={"role": "encode", "url": "http://e0:8000"} |
| 184 | ).status_code |
| 185 | == 403 |
| 186 | ) |
| 187 | assert ( |
| 188 | client.delete("/instances", params={"url": "http://e0:8000"}).status_code |
| 189 | == 403 |
| 190 | ) |
| 191 | |
| 192 | @pytest.mark.parametrize("split_role", ["prefill", "decode"]) |
| 193 | @pytest.mark.parametrize("combined_first", [False, True]) |
| 194 | def test_registration_rejects_mixed_topologies( |
| 195 | self, client, split_role, combined_first |
| 196 | ): |
| 197 | roles = ["prefill_decode", split_role] |
| 198 | if not combined_first: |
| 199 | roles.reverse() |
| 200 | first = {"role": roles[0], "url": "http://first:8000"} |
| 201 | second = {"role": roles[1], "url": "http://second:8000"} |
| 202 | assert client.post("/instances", json=first).status_code == 200 |
| 203 | assert client.post("/instances", json=second).status_code == 409 |
| 204 | # Health eviction must not silently permit a topology switch. |
| 205 | registry = client.app.state.registry |
| 206 | record = registry._live[first["url"]] |
| 207 | for _ in range(registry._fail_threshold): |
| 208 | registry._on_probe_failure(record, 0) |
| 209 | assert client.post("/instances", json=second).status_code == 409 |
| 210 | assert client.delete("/instances", params={"url": first["url"]}).json() == { |
| 211 | "removed": True |
| 212 | } |
| 213 | assert client.post("/instances", json=second).status_code == 200 |
| 214 | |
| 215 | |
| 216 | def test_extract_mm_items_finds_media_across_messages(): |
| 217 | req = { |
| 218 | "messages": [ |
| 219 | {"role": "user", "content": "plain text"}, |
| 220 | {"role": "user", "content": [IMAGE_ITEM, {"type": "text", "text": "hi"}]}, |
| 221 | {"role": "user", "content": [IMAGE_ITEM]}, |
| 222 | ] |
| 223 | } |
| 224 | assert extract_mm_items(req) == [IMAGE_ITEM, IMAGE_ITEM] |
| 225 | |
| 226 | |
| 227 | @pytest.mark.parametrize("push", [False, True]) |
| 228 | def test_encoder_handles_and_json_metadata_survive_rewrite(proxy, push, monkeypatch): |
| 229 | """Keep main's hash-keyed handles and per-item push IDs after the proxy move.""" |
| 230 | proxy.registry.register(InstanceRecord(ENCODE, "http://encoder")) |
| 231 | proxy.registry.register( |
| 232 | InstanceRecord( |
| 233 | PD, |
| 234 | "http://decode", |
| 235 | ec_zmq_addrs=["tcp://decode:14579"] if push else [], |
| 236 | ) |
| 237 | ) |
| 238 | handle = { |
| 239 | "metadata": {"image_grid_thw": [[1, 2, 3]]}, |
| 240 | "item_indices": [0, 1], |
| 241 | "peer_port": 5601, |
| 242 | } |
| 243 | response = Mock( |
| 244 | status=200, |
| 245 | read=AsyncMock( |
| 246 | return_value=msgspec.json.encode( |
| 247 | {"ec_transfer_params": {"engine-hash": handle}} |
| 248 | ) |
| 249 | ), |
| 250 | ) |
| 251 | session = Mock(post=AsyncMock(return_value=response)) |
| 252 | monkeypatch.setattr(proxy_module, "encode_session", session) |
| 253 | original = {"messages": [{"role": "user", "content": [IMAGE_ITEM, IMAGE_ITEM]}]} |
| 254 | route = proxy.route(2) |
| 255 | prepared, _, _ = asyncio.run( |
| 256 | proxy_module.prepare_for_decode( |
| 257 | original, "request", route.encoder_urls, None, route.consumer_zmq |
| 258 | ) |
| 259 | ) |
| 260 | assert "ec_transfer_params" not in original |
| 261 | assert prepared["ec_transfer_params"]["engine-hash"] == handle |
| 262 | items = prepared["ec_transfer_params"]["ec_items"] |
| 263 | assert [item["mm_hash"] for item in items] == ["engine-hash", "engine-hash"] |
| 264 | assert items[0]["transfer_id"] != items[1]["transfer_id"] |
| 265 | assert session.post.call_count == 1 |
| 266 | sent = msgspec.json.decode(session.post.call_args.kwargs["data"]).get( |
| 267 | "ec_transfer_params" |
| 268 | ) |
| 269 | if push: |
| 270 | assert sent == { |
| 271 | "consumer_zmq": "tcp://decode:14579", |
| 272 | "ec_items": [{"transfer_id": item["transfer_id"]} for item in items], |
| 273 | } |
| 274 | else: |
| 275 | assert sent is None |
| 276 | for index in range(2): |
| 277 | assert prepared["messages"][0]["content"][index]["image_embeds"] == { |
| 278 | "image_grid_thw": [1, 2, 3] |
| 279 | } |