@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70fc7bd695

[Feature][EPD] Support dynamic EPD (EC-connector) proxy (#54176)

Signed-off-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: OpenAI Codex <noreply@openai.com>

Tianyu Guo committed Sep 20, 2026 at 12:24 UTC 70fc7bd695e6c9161ec7f3ba67b73e243189d571
8 files changed +1169 -73
examples/disaggregated/disaggregated_encoder/README.md
+77 -1
@@ -106,6 +106,82 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will
106 | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. |
107 | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). |
108
109 +### Dynamic registration
110 +
111 +Alternatively, let the external launcher register ready instances over HTTP.
112 +No vLLM configuration changes or worker registration threads are required.
113 +Set `ADMIN_API_KEY` on the proxy and supply it as `X-API-Key` for registration
114 +and removal. These are trusted control-plane APIs; do not expose them publicly.
115 +
116 +```bash
117 +export ADMIN_API_KEY="your-admin-key"
118 +python disagg_epd_proxy.py --port 8000 --dynamic-registration
119 +```
120 +
121 +Omit the static server URL flags and register every stage. The roles determine
122 +the topology: `encode` + `prefill_decode` gives E+PD; `encode` + `prefill` +
123 +`decode` gives E+P+D. Standalone `decode` always requires an available `prefill`,
124 +even when D registers first or all P instances go offline; otherwise requests
125 +return `503`. The proxy rejects mixing combined PD with standalone P/D,
126 +including unhealthy instances still in the registry. Explicitly remove the
127 +old topology's P/D or PD registrations before switching topologies.
128 +
129 +After each instance is ready, the launcher registers its reachable HTTP URL:
130 +
131 +```bash
132 +curl --fail-with-body http://proxy-host:8000/instances \
133 + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
134 + -d '{"role":"encode","url":"http://e-host:8001"}'
135 +curl --fail-with-body http://proxy-host:8000/instances \
136 + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
137 + -d '{"role":"prefill_decode","url":"http://pd-host:8002"}'
138 +```
139 +
140 +For E+P+D, register P with `role: "prefill"` and D with `role: "decode"`.
141 +Example and NIXL EC connectors need no additional registration fields.
142 +For Mooncake, register `ec_zmq_addrs` on the **EC consumer** (`prefill_decode`
143 +or `prefill`), using its configured `ec_ip` and `ec_port`. Keep the fixed-port
144 +layout: supply the TP-rank-0 address for each DP replica, in DP-rank order;
145 +replica `r` uses `ec_port + r * tensor_parallel_size`. The connector discovers
146 +the remaining TP ranks itself. For example, DP=2, TP=2, `ec_port=19019`:
147 +
148 +```json
149 +{
150 + "role": "prefill_decode",
151 + "url": "http://pd-host:8002",
152 + "dp_size": 2,
153 + "ec_zmq_addrs": ["tcp://pd-host:19019", "tcp://pd-host:19021"]
154 +}
155 +```
156 +
157 +The proxy selects one consumer replica and uses it for both the encoder push
158 +and the consumer HTTP request. Standalone D does not need EC control addresses.
159 +Port allocation and avoiding collisions remain the launcher's responsibility.
160 +
161 +Inspect or remove instances without restarting the proxy:
162 +
163 +```bash
164 +curl http://proxy-host:8000/instances
165 +curl --fail-with-body -X DELETE \
166 + 'http://proxy-host:8000/instances?url=http://e-host:8001' \
167 + -H "X-API-Key: $ADMIN_API_KEY"
168 +```
169 +
170 +Registration is idempotent. The proxy probes registered instances every
171 +`--probe-interval` seconds (default 5), with `--probe-timeout` seconds per probe
172 +(default 2). After `--fail-threshold` consecutive failures (default 3), it stops
173 +sending new requests to that instance. Healthy instances rejoin automatically;
174 +unreachable ones are forgotten after `--evicted-ttl` seconds (default 900;
175 +`0` retains them indefinitely). Removal stops new routing; already routed
176 +requests retain their selected endpoints, so drain requests before stopping
177 +the instance. The launcher must re-register instances after a proxy restart.
178 +
179 +The example launch scripts also support this flow: export `ADMIN_API_KEY` and
180 +set `DYNAMIC_REGISTRATION=1` when running `disagg_1e1pd_example.sh` or
181 +`disagg_1e1p1d_example.sh`. Their default remains static routing.
182 +
183 +### Static configuration
184 +
185 The proxy batches images from the same user request assigned to the same encoder.
186 Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number
187 of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each
@@ -137,6 +213,6 @@ For E + P + D setup:
213 ```bash
214 $ python disagg_encoder_proxy.py \
215 --encode-servers-urls "http://e1:8001,http://e2:8001" \
140 - --prefill-servers-urls "http://p1:8003,http://p2:8004" \
216 + --prefill-servers-urls "http://p1:8003,http://p2:8004" \
217 --decode-servers-urls "http://d1:8005,http://d2:8006"
218 ```
examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh
+21 -4
@@ -189,17 +189,34 @@ wait_for_server "$DECODE_PORT"
189 ###############################################################################
190 # Proxy
191 ###############################################################################
192 +PROXY_ARGS=()
193 +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
194 + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}"
195 + export ADMIN_API_KEY
196 + PROXY_ARGS=(--dynamic-registration)
197 +else
198 + PROXY_ARGS=(
199 + --encode-servers-urls "http://localhost:$ENCODE_PORT"
200 + --prefill-servers-urls "http://localhost:$PREFILL_PORT"
201 + --decode-servers-urls "http://localhost:$DECODE_PORT"
202 + )
203 +fi
204 python disagg_epd_proxy.py \
205 --host "0.0.0.0" \
206 --port "$PROXY_PORT" \
195 - --encode-servers-urls "http://localhost:$ENCODE_PORT" \
196 - --prefill-servers-urls "http://localhost:$PREFILL_PORT" \
197 - --decode-servers-urls "http://localhost:$DECODE_PORT" \
207 + "${PROXY_ARGS[@]}" \
208 >"${PROXY_LOG}" 2>&1 &
209
210 PIDS+=($!)
211
212 wait_for_server "$PROXY_PORT"
213 +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
214 + for endpoint in "encode:$ENCODE_PORT" "prefill:$PREFILL_PORT" "decode:$DECODE_PORT"; do
215 + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \
216 + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
217 + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}"
218 + done
219 +fi
220 echo "All services are up!"
221
222 ###############################################################################
@@ -238,4 +255,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
255
256 # cleanup
257 echo "cleanup..."
241 -cleanup
\ No newline at end of file
258 +cleanup
examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh
+20 -4
@@ -152,17 +152,33 @@ wait_for_server "$PREFILL_DECODE_PORT"
152 ###############################################################################
153 # Proxy
154 ###############################################################################
155 +PROXY_ARGS=()
156 +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
157 + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}"
158 + export ADMIN_API_KEY
159 + PROXY_ARGS=(--dynamic-registration)
160 +else
161 + PROXY_ARGS=(
162 + --encode-servers-urls "http://localhost:$ENCODE_PORT"
163 + --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT"
164 + )
165 +fi
166 python disagg_epd_proxy.py \
167 --host "0.0.0.0" \
168 --port "$PROXY_PORT" \
158 - --encode-servers-urls "http://localhost:$ENCODE_PORT" \
159 - --prefill-servers-urls "disable" \
160 - --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
169 + "${PROXY_ARGS[@]}" \
170 >"${PROXY_LOG}" 2>&1 &
171
172 PIDS+=($!)
173
174 wait_for_server "$PROXY_PORT"
175 +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then
176 + for endpoint in "encode:$ENCODE_PORT" "prefill_decode:$PREFILL_DECODE_PORT"; do
177 + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \
178 + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \
179 + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}"
180 + done
181 +fi
182 echo "All services are up!"
183
184 ###############################################################################
@@ -201,4 +217,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \
217
218 # cleanup
219 echo "cleanup..."
204 -cleanup
\ No newline at end of file
220 +cleanup
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+571 -60
@@ -1,12 +1,13 @@
1 #!/usr/bin/env python3
2 # SPDX-License-Identifier: Apache-2.0
3 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 -"""disagg_encoder_proxy.py
4 +"""OpenAI-compatible proxy for E+PD and E+P+D disaggregation.
5
6 -Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two
7 -clusters:
8 - • encode (multimodal feature extraction)
9 - • decode (language-model inference)
6 +Use static server URLs, or --dynamic-registration for launcher-managed HTTP
7 +registration of E, P/PD and D instances. Both modes route OpenAI-compatible
8 +``/v1/chat/completions`` requests through encoder and inference clusters and
9 +share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming
10 +forwarding.
11
12 For MM input we:
13 1. Extract *every* image/audio/video item.
@@ -20,22 +21,27 @@ from __future__ import annotations
21
22 import argparse
23 import asyncio
24 +import enum
25 import hashlib
26 import itertools
27 import json
28 import logging
29 import os
30 import random
31 +import secrets
32 import time
33 import uuid
34 from collections.abc import AsyncIterator
35 +from contextlib import asynccontextmanager, suppress
36 +from dataclasses import dataclass, field
37 from typing import Any
38
39 import aiohttp
40 import msgspec
41 import uvicorn
37 -from fastapi import FastAPI, HTTPException, Request
42 +from fastapi import Depends, FastAPI, Header, HTTPException, Request
43 from fastapi.responses import JSONResponse, Response, StreamingResponse
44 +from pydantic import AnyHttpUrl, BaseModel, Field, model_validator
45
46 ###############################################################################
47 # FastAPI app & global state
@@ -44,6 +50,340 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse
50 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
51 logger = logging.getLogger("proxy")
52
53 +DEFAULT_PROBE_INTERVAL = 5.0
54 +DEFAULT_PROBE_TIMEOUT = 2.0
55 +DEFAULT_FAIL_THRESHOLD = 3
56 +# Stop probing an instance that has been down this long. 0 probes forever,
57 +# which is what a cluster that restarts instances in place wants.
58 +DEFAULT_EVICTED_TTL = 900.0
59 +
60 +
61 +class InstanceRole(str, enum.Enum):
62 + ENCODE = "encode"
63 + PREFILL = "prefill"
64 + DECODE = "decode"
65 + PREFILL_DECODE = "prefill_decode"
66 +
67 +
68 +@dataclass
69 +class InstanceRecord:
70 + """One registered instance.
71 +
72 + Attributes:
73 + role: Which stage this instance serves.
74 + url: Base OpenAI-compatible URL, e.g. ``http://host:8000``.
75 + ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica,
76 + supplied by the launcher from the consumer's fixed port config.
77 + dp_size: Data-parallel replicas behind `url`, so the proxy can pick
78 + a replica and name the same one to both halves of a request.
79 +
80 + """
81 +
82 + role: InstanceRole
83 + url: str
84 + ec_zmq_addrs: list[str] = field(default_factory=list)
85 + dp_size: int = 1
86 +
87 +
88 +class InstanceRegistry:
89 + def __init__(
90 + self,
91 + probe_interval: float = DEFAULT_PROBE_INTERVAL,
92 + probe_timeout: float = DEFAULT_PROBE_TIMEOUT,
93 + fail_threshold: int = DEFAULT_FAIL_THRESHOLD,
94 + evicted_ttl: float = DEFAULT_EVICTED_TTL,
95 + ):
96 + self._probe_interval = probe_interval
97 + self._probe_timeout = probe_timeout
98 + self._fail_threshold = fail_threshold
99 + self._evicted_ttl = evicted_ttl
100 +
101 + self._live: dict[str, InstanceRecord] = {}
102 + self._evicted: dict[str, InstanceRecord] = {}
103 + self._evicted_since: dict[str, float] = {}
104 + self._fail_counts: dict[str, int] = {}
105 + # One cursor per role, only ever incremented. Rebuilding it whenever
106 + # the roster changes -- what an `itertools.cycle` over a mutable list
107 + # forces -- restarts every fan-out at the first instance and hot-spots
108 + # it after each registration.
109 + self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole}
110 + self._replica_cursors: dict[str, int] = {}
111 + self._probe_task: asyncio.Task | None = None
112 +
113 + def register(self, record: InstanceRecord) -> bool:
114 + """Add or refresh an instance without overriding its health status."""
115 + roles = {
116 + other.role
117 + for other in itertools.chain(self._live.values(), self._evicted.values())
118 + }
119 + split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE}
120 + if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or (
121 + record.role in split_roles and InstanceRole.PREFILL_DECODE in roles
122 + ):
123 + raise ValueError("Cannot mix prefill_decode with standalone prefill/decode")
124 + key = record.url
125 + previous = self._live.get(key) or self._evicted.get(key)
126 + if previous is not None:
127 + if previous == record:
128 + return False
129 + if previous.role is not record.role:
130 + raise ValueError("Unregister the instance before changing its role")
131 + target = self._live if key in self._live else self._evicted
132 + target[key] = record
133 + return False
134 + self._live[key] = record
135 + logger.info("Registered instance %s: %s", record.role.value, record.url)
136 + return True
137 +
138 + def unregister(self, url: str) -> bool:
139 + """Drop an instance for good, so a probe cannot bring it back."""
140 + key = url
141 + if key not in self._live and key not in self._evicted:
142 + return False
143 + self._live.pop(key, None)
144 + self._evicted.pop(key, None)
145 + self._evicted_since.pop(key, None)
146 + self._fail_counts.pop(key, None)
147 + self._replica_cursors.pop(key, None)
148 + logger.info("Unregistered instance: %s", url)
149 + return True
150 +
151 + def instances(self, role: InstanceRole) -> list[InstanceRecord]:
152 + return [record for record in self._live.values() if record.role is role]
153 +
154 + def urls(self, role: InstanceRole) -> list[str]:
155 + return [record.url for record in self.instances(role)]
156 +
157 + def pick(self, role: InstanceRole) -> InstanceRecord | None:
158 + """Take the next instance of `role` in round-robin order."""
159 + picked = self.pick_many(role, 1)
160 + return picked[0] if picked else None
161 +
162 + def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]:
163 + """Take `count` instances, continuing the rotation across calls.
164 +
165 + A multimodal request fans out one encoder request per item, so the
166 + assignment has to be contiguous with the previous request's rather
167 + than restart at the first instance every time.
168 + """
169 + alive = self.instances(role)
170 + if not alive or count <= 0:
171 + return []
172 + start = self._cursors[role]
173 + self._cursors[role] = start + count
174 + return [alive[(start + offset) % len(alive)] for offset in range(count)]
175 +
176 + def next_replica(self, record: InstanceRecord) -> int:
177 + """Take the next data-parallel replica of `record`, round-robin.
178 +
179 + The encoder pushes to one replica's receive channel, so the request
180 + has to run on that same replica; the caller names the rank to both
181 + halves.
182 + """
183 + if record.dp_size <= 1:
184 + return 0
185 + cursor = self._replica_cursors.get(record.url, 0)
186 + self._replica_cursors[record.url] = cursor + 1
187 + return cursor % record.dp_size
188 +
189 + def status(self) -> dict[str, Any]:
190 + return {
191 + role.value: {
192 + "live": [record.url for record in self.instances(role)],
193 + "evicted": [
194 + record.url
195 + for record in self._evicted.values()
196 + if record.role is role
197 + ],
198 + }
199 + for role in InstanceRole
200 + }
201 +
202 + def start_probing(self) -> None:
203 + if self._probe_task is None and self._probe_interval > 0:
204 + self._probe_task = asyncio.create_task(self._probe_loop())
205 +
206 + async def stop_probing(self) -> None:
207 + if self._probe_task is None:
208 + return
209 + self._probe_task.cancel()
210 + with suppress(asyncio.CancelledError):
211 + await self._probe_task
212 + self._probe_task = None
213 +
214 + async def _probe_loop(self) -> None:
215 + timeout = aiohttp.ClientTimeout(total=self._probe_timeout)
216 + async with aiohttp.ClientSession(timeout=timeout) as session:
217 + while True:
218 + await asyncio.sleep(self._probe_interval)
219 + try:
220 + await self._probe_once(session)
221 + except Exception:
222 + logger.exception("EPD registry probe round failed")
223 +
224 + async def _probe_once(self, session: aiohttp.ClientSession) -> None:
225 + targets = list(self._live.values()) + list(self._evicted.values())
226 + if not targets:
227 + return
228 + results = await asyncio.gather(
229 + *(self._probe(session, record.url) for record in targets),
230 + return_exceptions=True,
231 + )
232 + now = time.monotonic()
233 + for record, healthy in zip(targets, results):
234 + current = self._live.get(record.url) or self._evicted.get(record.url)
235 + if current is not record:
236 + continue
237 + if healthy is True:
238 + self._on_probe_success(record)
239 + else:
240 + self._on_probe_failure(record, now)
241 + self._drop_expired(now)
242 +
243 + async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool:
244 + async with session.get(f"{url}/health") as resp:
245 + return resp.status == 200
246 +
247 + def _on_probe_success(self, record: InstanceRecord) -> None:
248 + key = record.url
249 + self._fail_counts.pop(key, None)
250 + if key in self._evicted:
251 + self._evicted.pop(key, None)
252 + self._evicted_since.pop(key, None)
253 + self._live[key] = record
254 + logger.info(
255 + "Instance %s (%s) is healthy again; routing resumed",
256 + record.url,
257 + record.role.value,
258 + )
259 +
260 + def _on_probe_failure(self, record: InstanceRecord, now: float) -> None:
261 + key = record.url
262 + if key not in self._live:
263 + # Either already evicted, or unregistered while this probe was in
264 + # flight. A round snapshots its targets and then awaits, so a
265 + # removal inside that window would otherwise be undone by a result
266 + # describing a registry that no longer exists.
267 + return
268 + failures = self._fail_counts.get(key, 0) + 1
269 + self._fail_counts[key] = failures
270 + if failures < self._fail_threshold:
271 + return
272 + self._live.pop(key, None)
273 + self._evicted[key] = record
274 + self._evicted_since[key] = now
275 + logger.warning(
276 + "Instance %s (%s) failed %d consecutive probes; stopped routing "
277 + "to it. It rejoins on its own once it responds again.",
278 + record.url,
279 + record.role.value,
280 + failures,
281 + )
282 +
283 + def _drop_expired(self, now: float) -> None:
284 + if self._evicted_ttl <= 0:
285 + return
286 + for key, since in list(self._evicted_since.items()):
287 + if now - since < self._evicted_ttl:
288 + continue
289 + record = self._evicted.pop(key, None)
290 + self._evicted_since.pop(key, None)
291 + self._fail_counts.pop(key, None)
292 + self._replica_cursors.pop(key, None)
293 + if record is not None:
294 + logger.warning("Instance %s stayed down; forgetting it", record.url)
295 +
296 +
297 +class InstanceRegistration(BaseModel):
298 + role: InstanceRole
299 + url: AnyHttpUrl
300 + ec_zmq_addrs: list[str] = Field(default_factory=list)
301 + dp_size: int = Field(default=1, ge=1)
302 +
303 + @model_validator(mode="after")
304 + def validate_consumer_addresses(self):
305 + if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size:
306 + raise ValueError("Provide one Mooncake control address per DP replica")
307 + return self
308 +
309 +
310 +def require_admin_key(x_api_key: str = Header(default="")) -> None:
311 + expected = os.getenv("ADMIN_API_KEY", "")
312 + if not expected or not secrets.compare_digest(x_api_key, expected):
313 + raise HTTPException(403, "Invalid admin API key")
314 +
315 +
316 +@dataclass
317 +class EPDProxyConfig:
318 + probe_interval: float = 5.0
319 + probe_timeout: float = 2.0
320 + fail_threshold: int = 3
321 + evicted_ttl: float = 900.0
322 +
323 +
324 +@dataclass
325 +class _Route:
326 + """The instances one request was assigned to.
327 +
328 + Attributes:
329 + consumer: The stage that receives the embedding, if any stage does.
330 + consumer_zmq: The receive address named to the encoders.
331 + dp_rank: Which replica of `consumer` was named, so the request can be
332 + pinned to it.
333 +
334 + """
335 +
336 + encoder_urls: list[str]
337 + prefill: InstanceRecord | None
338 + decode: InstanceRecord
339 + consumer: InstanceRecord | None = None
340 + consumer_zmq: str | None = None
341 + dp_rank: int | None = None
342 +
343 +
344 +class EPDProxy:
345 + def __init__(self, registry: InstanceRegistry):
346 + self.registry = registry
347 +
348 + # ---------------------------------------------------------------- #
349 + # Routing #
350 + # ---------------------------------------------------------------- #
351 + def route(self, num_items: int) -> _Route:
352 + decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick(
353 + InstanceRole.DECODE
354 + )
355 + if decode is None:
356 + raise HTTPException(
357 + status_code=503, detail="No decode instance is registered"
358 + )
359 + prefill = (
360 + self.registry.pick(InstanceRole.PREFILL)
361 + if decode.role is InstanceRole.DECODE
362 + else None
363 + )
364 + if decode.role is InstanceRole.DECODE and prefill is None:
365 + raise HTTPException(
366 + status_code=503, detail="No prefill instance is registered"
367 + )
368 + encoders = self.registry.urls(InstanceRole.ENCODE)
369 + if num_items and not encoders:
370 + raise HTTPException(
371 + status_code=503, detail="No encode instance is registered"
372 + )
373 + route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode)
374 + self._name_consumer(route)
375 + return route
376 +
377 + def _name_consumer(self, route: _Route) -> None:
378 + """Pin the EC consumer replica and, for push connectors, its endpoint."""
379 + candidate = route.prefill or route.decode
380 + route.consumer = candidate
381 + rank = self.registry.next_replica(candidate)
382 + route.dp_rank = rank if candidate.dp_size > 1 else None
383 + if candidate.ec_zmq_addrs:
384 + route.consumer_zmq = candidate.ec_zmq_addrs[rank]
385 +
386 +
387 app = FastAPI()
388 encode_session: aiohttp.ClientSession | None = None
389 prefill_session: aiohttp.ClientSession | None = None
@@ -448,6 +788,7 @@ async def maybe_prefill(
788 req_data: dict,
789 p_url: str,
790 req_id: str,
791 + dp_rank: int | None = None,
792 ) -> dict:
793 """- Do prefill-only task if p_url exist;
794 - Return a new body carrying kv transfer params (for nixl connector)
@@ -459,7 +800,7 @@ async def maybe_prefill(
800 if p_url:
801 logger.info("[%s] Processing through prefill: %s", req_id, p_url)
802
462 - prefill_response = await process_prefill_stage(req_data, p_url, req_id)
803 + prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank)
804 # for nixl connector to facilitate kv transfer...
805 prefill_response_json = msgspec.json.decode(await prefill_response.read())
806 kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
@@ -473,6 +814,7 @@ async def process_prefill_stage(
814 req_data: dict,
815 p_url: str,
816 req_id: str,
817 + dp_rank: int | None = None,
818 ) -> dict:
819 """Process request through Prefill stage and return kv_transfer_params"""
820 logger.info("[%s] Sending prefill request to: %s", req_id, p_url)
@@ -494,6 +836,8 @@ async def process_prefill_stage(
836 del prefill_request["stream_options"]
837
838 headers = {"x-request-id": req_id, "Content-Type": "application/json"}
839 + if dp_rank is not None:
840 + headers["X-data-parallel-rank"] = str(dp_rank)
841 try:
842 prefill_response = await prefill_session.post(
843 f"{p_url}/v1/chat/completions",
@@ -594,9 +938,7 @@ async def on_startup() -> None:
938 ),
939 )
940 encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
597 - if app.state.p_urls:
598 - # only setup if prefill instance(s) exist
599 - prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
941 + prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
942 decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
943
944
@@ -622,6 +964,7 @@ async def prepare_for_decode(
964 e_urls: list[str],
965 p_url: str,
966 consumer_zmq: str | None,
967 + prefill_dp_rank: int | None = None,
968 ) -> tuple[dict, float, float]:
969 """Encode, rewrite and prefill, returning the body to send to decode.
970
@@ -642,7 +985,7 @@ async def prepare_for_decode(
985 handles.update(ec_params)
986 prepared = {**prepared, "ec_transfer_params": handles}
987 _t2 = time.perf_counter()
645 - prepared = await maybe_prefill(prepared, p_url, req_id)
988 + prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank)
989 return prepared, _t1 - _t0, _t2 - _t1
990
991
@@ -654,12 +997,13 @@ async def forward_non_stream(
997 d_url: str,
998 consumer_zmq: str | None,
999 dp_rank: int | None = None,
1000 + prefill_dp_rank: int | None = None,
1001 ) -> Response:
1002 try:
1003 for attempt in range(DECODE_RETRIES + 1):
1004 _t0 = time.perf_counter()
1005 prepared, encode_s, rewrite_s = await prepare_for_decode(
662 - req_data, req_id, e_urls, p_url, consumer_zmq
1006 + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank
1007 )
1008 _t2 = time.perf_counter()
1009
@@ -733,12 +1077,13 @@ async def forward_stream(
1077 d_url: str,
1078 consumer_zmq: str | None,
1079 dp_rank: int | None = None,
1080 + prefill_dp_rank: int | None = None,
1081 ) -> AsyncIterator[bytes]:
1082 try:
1083 for attempt in range(DECODE_RETRIES + 1):
1084 _t0 = time.perf_counter()
1085 prepared, encode_s, rewrite_s = await prepare_for_decode(
741 - req_data, req_id, e_urls, p_url, consumer_zmq
1086 + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank
1087 )
1088 _t2 = time.perf_counter()
1089
@@ -976,10 +1321,147 @@ async def stop_profile(request: Request):
1321 return await _profile_cmd("stop", body, e_url, p_url, d_url)
1322
1323
1324 +def build_app(config: EPDProxyConfig | None = None) -> FastAPI:
1325 + config = config or EPDProxyConfig()
1326 + registry = InstanceRegistry(
1327 + probe_interval=config.probe_interval,
1328 + probe_timeout=config.probe_timeout,
1329 + fail_threshold=config.fail_threshold,
1330 + evicted_ttl=config.evicted_ttl,
1331 + )
1332 + proxy = EPDProxy(registry)
1333 +
1334 + @asynccontextmanager
1335 + async def lifespan(app: FastAPI):
1336 + await on_startup()
1337 + try:
1338 + registry.start_probing()
1339 + yield
1340 + finally:
1341 + await registry.stop_probing()
1342 + await on_shutdown()
1343 +
1344 + app = FastAPI(lifespan=lifespan)
1345 + app.state.proxy = proxy
1346 + app.state.registry = registry
1347 +
1348 + @app.get("/instances")
1349 + async def list_instances():
1350 + return registry.status()
1351 +
1352 + @app.post("/instances", dependencies=[Depends(require_admin_key)])
1353 + async def register_instance(body: InstanceRegistration):
1354 + if body.ec_zmq_addrs and body.role not in (
1355 + InstanceRole.PREFILL,
1356 + InstanceRole.PREFILL_DECODE,
1357 + ):
1358 + raise HTTPException(400, "Only EC consumers accept Mooncake addresses")
1359 + try:
1360 + created = registry.register(
1361 + InstanceRecord(
1362 + body.role,
1363 + str(body.url).rstrip("/"),
1364 + body.ec_zmq_addrs,
1365 + body.dp_size,
1366 + )
1367 + )
1368 + except ValueError as error:
1369 + raise HTTPException(409, str(error)) from error
1370 + return {"registered": created}
1371 +
1372 + @app.delete("/instances", dependencies=[Depends(require_admin_key)])
1373 + async def unregister_instance(url: AnyHttpUrl):
1374 + return {"removed": registry.unregister(str(url).rstrip("/"))}
1375 +
1376 + @app.post("/v1/chat/completions")
1377 + async def chat_completions(request: Request):
1378 + req_data = msgspec.json.decode(await request.body())
1379 + req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
1380 + route = proxy.route(len(extract_mm_items(req_data)))
1381 + args = (
1382 + req_data,
1383 + req_id,
1384 + route.encoder_urls,
1385 + route.prefill.url if route.prefill else None,
1386 + route.decode.url,
1387 + route.consumer_zmq,
1388 + )
1389 + ranks = {
1390 + "dp_rank": route.dp_rank if route.prefill is None else None,
1391 + "prefill_dp_rank": route.dp_rank if route.prefill is not None else None,
1392 + }
1393 + if req_data.get("stream", False):
1394 + return StreamingResponse(
1395 + forward_stream(*args, **ranks), media_type="text/event-stream"
1396 + )
1397 + return await forward_non_stream(*args, **ranks)
1398 +
1399 + @app.get("/v1/models")
1400 + async def list_models():
1401 + decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick(
1402 + InstanceRole.DECODE
1403 + )
1404 + if decode is None:
1405 + raise HTTPException(
1406 + status_code=503, detail="No decode instance is registered"
1407 + )
1408 + async with decode_session.get(f"{decode.url}/v1/models") as resp:
1409 + resp.raise_for_status()
1410 + return await resp.json()
1411 +
1412 + @app.get("/health")
1413 + async def health():
1414 + status = registry.status()
1415 + # An empty roster is not unhealthy: the proxy is meant to come up
1416 + # before anything registers with it.
1417 + return JSONResponse({"proxy": "healthy", "instances": status})
1418 +
1419 + @app.post("/start_profile")
1420 + async def start_profile(request: Request):
1421 + return await _profile(registry, "start", await request.json())
1422 +
1423 + @app.post("/stop_profile")
1424 + async def stop_profile(request: Request):
1425 + return await _profile(registry, "stop", await request.json())
1426 +
1427 + return app
1428 +
1429 +
1430 +async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict:
1431 + headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
1432 + targets = {role.value: registry.pick(role) for role in InstanceRole}
1433 + results = await asyncio.gather(
1434 + *(
1435 + _post_if_available(
1436 + decode_session, f"{record.url}/{cmd}_profile", payload, headers
1437 + )
1438 + for record in targets.values()
1439 + if record is not None
1440 + )
1441 + )
1442 + reachable = [result for result in results if result is not None]
1443 + if not reachable:
1444 + raise HTTPException(
1445 + status_code=503,
1446 + detail="Profiling endpoints are disabled on every instance",
1447 + )
1448 + live = [name for name, record in targets.items() if record is not None]
1449 + return dict(zip(live, results))
1450 +
1451 +
1452 if __name__ == "__main__":
1453 parser = argparse.ArgumentParser()
1454 parser.add_argument("--host", default="0.0.0.0")
1455 parser.add_argument("--port", type=int, default=8000)
1456 + parser.add_argument(
1457 + "--dynamic-registration",
1458 + action="store_true",
1459 + help="Enable launcher-managed HTTP registration for E+PD or E+P+D.",
1460 + )
1461 + parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL)
1462 + parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT)
1463 + parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD)
1464 + parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL)
1465 parser.add_argument(
1466 "--log-requests",
1467 action="store_true",
@@ -995,12 +1477,12 @@ if __name__ == "__main__":
1477 )
1478 parser.add_argument(
1479 "--encode-servers-urls",
998 - required=True,
1480 + default="",
1481 help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")',
1482 )
1483 parser.add_argument(
1484 "--prefill-servers-urls",
1003 - required=True,
1485 + default="none",
1486 help=(
1487 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") '
1488 'to enable E->P->D, set "disable" or "none" to enable E->PD'
@@ -1008,7 +1490,7 @@ if __name__ == "__main__":
1490 )
1491 parser.add_argument(
1492 "--decode-servers-urls",
1011 - required=True,
1493 + default="",
1494 help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")',
1495 )
1496 parser.add_argument(
@@ -1043,58 +1525,87 @@ if __name__ == "__main__":
1525 )
1526
1527 args = parser.parse_args()
1528 + if args.fail_threshold < 1:
1529 + parser.error("--fail-threshold must be at least 1")
1530 if args.log_requests:
1531 logging.getLogger().setLevel(logging.DEBUG)
1532 app.middleware("http")(log_requests)
1533 NO_REWRITE = args.no_rewrite
1534 DECODE_RETRIES = max(0, args.decode_retries)
1051 - app.state.e_urls = [
1052 - u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
1053 - ]
1054 - app.state.d_urls = [
1055 - u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
1056 - ]
1057 - app.state.d_ec_urls = [
1058 - u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
1059 - ]
1060 - if args.ec_consumer_dp_size < 1:
1061 - parser.error("--ec-consumer-dp-size must be at least 1")
1062 - app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
1063 - app.state.replica_counter = itertools.count()
1064 - expected = len(app.state.d_urls) * args.ec_consumer_dp_size
1065 - if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
1066 - parser.error(
1067 - "--ec-consumer-zmq-addrs must contain one address per consumer "
1068 - f"replica: expected {expected} "
1069 - f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), "
1070 - f"got {len(app.state.d_ec_urls)}"
1071 - )
1072 - # handle prefill instances
1073 - if args.prefill_servers_urls.lower() in ("disable", "none", ""):
1074 - app.state.p_urls = []
1075 - logger.info(
1076 - "Disaggregated prefill phase explicitly disabled by user. Running E + PD..."
1535 + if args.dynamic_registration:
1536 + if not os.getenv("ADMIN_API_KEY"):
1537 + parser.error("--dynamic-registration requires ADMIN_API_KEY")
1538 + if (
1539 + args.encode_servers_urls
1540 + or args.prefill_servers_urls.lower() not in ("disable", "none", "")
1541 + or args.decode_servers_urls
1542 + or args.ec_consumer_zmq_addrs
1543 + ):
1544 + parser.error(
1545 + "With --dynamic-registration, instances register through /instances"
1546 + )
1547 + app = build_app(
1548 + EPDProxyConfig(
1549 + probe_interval=args.probe_interval,
1550 + probe_timeout=args.probe_timeout,
1551 + fail_threshold=args.fail_threshold,
1552 + evicted_ttl=args.evicted_ttl,
1553 + )
1554 )
1555 + if args.log_requests:
1556 + app.middleware("http")(log_requests)
1557 else:
1079 - app.state.p_urls = [
1080 - u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
1558 + if not args.encode_servers_urls or not args.decode_servers_urls:
1559 + parser.error(
1560 + "Static routing requires --encode-servers-urls "
1561 + "and --decode-servers-urls"
1562 + )
1563 + app.state.e_urls = [
1564 + u.strip() for u in args.encode_servers_urls.split(",") if u.strip()
1565 ]
1082 - logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
1083 - try:
1084 - validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
1085 - except ValueError as exc:
1086 - parser.error(str(exc))
1087 -
1088 - logger.info("Proxy listening on %s:%s", args.host, args.port)
1089 - logger.info("Encode servers: %s", app.state.e_urls)
1090 - logger.info("Prefill instances %s", app.state.p_urls)
1091 - logger.info("Decode servers: %s", app.state.d_urls)
1092 - if app.state.ec_consumer_dp_size > 1:
1093 - logger.info(
1094 - "EC consumer replicas per server: %d (control addresses: %s)",
1095 - app.state.ec_consumer_dp_size,
1096 - app.state.d_ec_urls,
1097 - )
1566 + app.state.d_urls = [
1567 + u.strip() for u in args.decode_servers_urls.split(",") if u.strip()
1568 + ]
1569 + app.state.d_ec_urls = [
1570 + u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip()
1571 + ]
1572 + if args.ec_consumer_dp_size < 1:
1573 + parser.error("--ec-consumer-dp-size must be at least 1")
1574 + app.state.ec_consumer_dp_size = args.ec_consumer_dp_size
1575 + app.state.replica_counter = itertools.count()
1576 + expected = len(app.state.d_urls) * args.ec_consumer_dp_size
1577 + if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected:
1578 + parser.error(
1579 + "--ec-consumer-zmq-addrs must contain one address per consumer "
1580 + f"replica: expected {expected} "
1581 + f"({len(app.state.d_urls)} servers x "
1582 + f"{args.ec_consumer_dp_size} replicas), "
1583 + f"got {len(app.state.d_ec_urls)}"
1584 + )
1585 + # handle prefill instances
1586 + if args.prefill_servers_urls.lower() in ("disable", "none", ""):
1587 + app.state.p_urls = []
1588 + logger.info("Disaggregated prefill disabled. Running E + PD...")
1589 + else:
1590 + app.state.p_urls = [
1591 + u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()
1592 + ]
1593 + logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
1594 + try:
1595 + validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls)
1596 + except ValueError as exc:
1597 + parser.error(str(exc))
1598 +
1599 + logger.info("Proxy listening on %s:%s", args.host, args.port)
1600 + logger.info("Encode servers: %s", app.state.e_urls)
1601 + logger.info("Prefill instances %s", app.state.p_urls)
1602 + logger.info("Decode servers: %s", app.state.d_urls)
1603 + if app.state.ec_consumer_dp_size > 1:
1604 + logger.info(
1605 + "EC consumer replicas per server: %d (control addresses: %s)",
1606 + app.state.ec_consumer_dp_size,
1607 + app.state.d_ec_urls,
1608 + )
1609
1610 uvicorn.run(
1611 app,
tests/v1/ec_connector/unit/test_epd_proxy.py new
+279
@@ -0,0 +1,279 @@
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 + }
tests/v1/ec_connector/unit/test_epd_proxy_retry.py
+43 -4
@@ -8,6 +8,7 @@ change to them is what these tests catch.
8
9 import asyncio
10 import importlib.util
11 +import sys
12 from pathlib import Path
13 from types import SimpleNamespace
14 from typing import Any
@@ -27,6 +28,7 @@ def proxy():
28 spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path)
29 assert spec is not None and spec.loader is not None
30 module = importlib.util.module_from_spec(spec)
31 + sys.modules[spec.name] = module
32 spec.loader.exec_module(module)
33 return module
34
@@ -48,7 +50,7 @@ def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch):
50 """
51 served = [{"remote_block_ids": [1, 2]}, {}]
52
51 - async def _stage(req_data, p_url, req_id):
53 + async def _stage(req_data, p_url, req_id, dp_rank=None):
54 assert "kv_transfer_params" not in req_data
55 return _Response(served.pop(0))
56
@@ -222,7 +224,7 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch
224 _EncoderSession([{"encoder-side-hash": handle}, {}]),
225 )
226
225 - async def _no_prefill(req_data, p_url, req_id):
227 + async def _no_prefill(req_data, p_url, req_id, dp_rank=None):
228 return req_data
229
230 monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill)
@@ -280,6 +282,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
282 @pytest.mark.asyncio
283 @pytest.mark.parametrize("stream", [False, True])
284 @pytest.mark.parametrize("prefill", [False, True])
285 +@pytest.mark.parametrize("dynamic", [False, True])
286 @pytest.mark.parametrize(
287 "request_options",
288 [
@@ -293,7 +296,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
296 ],
297 )
298 async def test_http_roundtrip_preserves_payload_and_response_bytes(
296 - proxy, monkeypatch, stream, prefill, request_options
299 + proxy, monkeypatch, stream, prefill, dynamic, request_options
300 ):
301 """Preserve media and scheduling context across HTTP hops and retries."""
302 seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []}
@@ -307,6 +310,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
310
311 async def backend(request):
312 stage = request.match_info["stage"]
313 + if dynamic:
314 + consumer = "prefill" if prefill else "decode"
315 + expected_rank = "0" if stage == consumer else None
316 + assert request.headers.get("X-data-parallel-rank") == expected_rank
317 assert request.content_type == "application/json"
318 body = await request.json()
319 seen[stage].append(body)
@@ -354,6 +361,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
361 monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False)
362 monkeypatch.setattr(proxy, "DECODE_RETRIES", 1)
363 monkeypatch.setattr(proxy, "NO_REWRITE", False)
364 + target_app = proxy.app
365 + if dynamic:
366 + monkeypatch.setenv("ADMIN_API_KEY", "test-key")
367 + target_app = proxy.build_app(proxy.EPDProxyConfig(probe_interval=0))
368 item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}}
369 body = {
370 "model": "test",
@@ -377,9 +388,37 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes(
388 await proxy.on_startup()
389 try:
390 async with httpx.AsyncClient(
380 - transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy"
391 + transport=httpx.ASGITransport(app=target_app), base_url="http://proxy"
392 ) as client:
393 + if dynamic:
394 + registrations = [("encode", "encode")]
395 + registrations += (
396 + [("prefill", "prefill"), ("decode", "decode")]
397 + if prefill
398 + else [("prefill_decode", "decode")]
399 + )
400 + for role, stage in registrations:
401 + consumer = "prefill" if prefill else "decode"
402 + registered = await client.post(
403 + "/instances",
404 + headers={"X-API-Key": "test-key"},
405 + json={
406 + "role": role,
407 + "url": str(server.make_url(f"/{stage}")),
408 + "dp_size": 2 if stage == consumer else 1,
409 + },
410 + )
411 + assert registered.status_code == 200
412 response = await client.post("/v1/chat/completions", json=body)
413 + if dynamic:
414 + removed = await client.delete(
415 + "/instances",
416 + headers={"X-API-Key": "test-key"},
417 + params={"url": str(server.make_url("/decode"))},
418 + )
419 + assert removed.json() == {"removed": True}
420 + unavailable = await client.post("/v1/chat/completions", json=body)
421 + assert unavailable.status_code == 503
422 assert response.status_code == 200
423 assert response.content == expected
424 assert response.headers["content-type"].startswith(
tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
+2
@@ -14,6 +14,7 @@ across calls so consecutive requests keep advancing through every URL.
14 """
15
16 import importlib.util
17 +import sys
18 from collections import Counter
19 from pathlib import Path
20
@@ -27,6 +28,7 @@ def _load_proxy_module():
28 spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path)
29 assert spec is not None and spec.loader is not None
30 module = importlib.util.module_from_spec(spec)
31 + sys.modules[spec.name] = module
32 spec.loader.exec_module(module)
33 return module
34
tests/v1/ec_connector/unit/test_epd_registry.py new
+156
@@ -0,0 +1,156 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Dynamic registration and liveness for the EPD proxy's instance registry."""
4 +
5 +from unittest.mock import patch
6 +
7 +import pytest
8 +
9 +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import (
10 + InstanceRecord,
11 + InstanceRegistry,
12 + InstanceRole,
13 +)
14 +
15 +ENCODE = InstanceRole.ENCODE
16 +DECODE = InstanceRole.DECODE
17 +
18 +
19 +@pytest.fixture
20 +def registry():
21 + return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0)
22 +
23 +
24 +def _fill(registry, role, count, prefix="http://h"):
25 + for index in range(count):
26 + registry.register(InstanceRecord(role, f"{prefix}{index}:8000"))
27 +
28 +
29 +async def _probe_round(registry, healthy: set[str], now: float = 0.0):
30 + """Run one probe pass where only `healthy` URLs answer."""
31 +
32 + async def fake_probe(self, session, url):
33 + return url in healthy
34 +
35 + with (
36 + patch.object(InstanceRegistry, "_probe", fake_probe),
37 + patch(
38 + "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic",
39 + return_value=now,
40 + ),
41 + ):
42 + await registry._probe_once(session=None)
43 +
44 +
45 +class TestRegistration:
46 + def test_registering_is_idempotent_and_reports_novelty(self, registry):
47 + record = InstanceRecord(ENCODE, "http://e0:8000")
48 + assert registry.register(record) is True
49 + assert registry.register(record) is False
50 + assert registry.urls(ENCODE) == ["http://e0:8000"]
51 +
52 + def test_transfer_addresses_stay_with_the_registered_consumer(self, registry):
53 + """Roster updates must not shift a consumer's transfer addresses."""
54 + registry.register(InstanceRecord(DECODE, "http://d0:8000"))
55 + registry.register(
56 + InstanceRecord(
57 + DECODE,
58 + "http://d0:8000",
59 + ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"],
60 + dp_size=2,
61 + )
62 + )
63 + picked = registry.pick(DECODE)
64 + assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"]
65 + assert picked.dp_size == 2
66 +
67 + def test_unregister_is_terminal(self, registry):
68 + registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
69 + assert registry.unregister("http://e0:8000") is True
70 + assert registry.unregister("http://e0:8000") is False
71 + assert registry.pick(ENCODE) is None
72 +
73 + def test_round_robin_survives_registration(self, registry):
74 + _fill(registry, ENCODE, 2)
75 + assert [record.url for record in registry.pick_many(ENCODE, 2)] == [
76 + "http://h0:8000",
77 + "http://h1:8000",
78 + ]
79 + registry.register(InstanceRecord(ENCODE, "http://h2:8000"))
80 + assert registry.pick(ENCODE).url == "http://h2:8000"
81 +
82 +
83 +class TestLiveness:
84 + @pytest.mark.asyncio
85 + async def test_identical_reregistration_preserves_inflight_probe_results(
86 + self, registry
87 + ):
88 + """Repeated registrations must not suppress eviction or recovery."""
89 + url = "http://e0:8000"
90 + registry.register(InstanceRecord(ENCODE, url))
91 + healthy = False
92 +
93 + async def probe(self, session, url):
94 + registry.register(InstanceRecord(ENCODE, url))
95 + return healthy
96 +
97 + with patch.object(InstanceRegistry, "_probe", probe):
98 + for _ in range(registry._fail_threshold):
99 + await registry._probe_once(None)
100 + assert registry.status()["encode"] == {"live": [], "evicted": [url]}
101 +
102 + healthy = True
103 + await registry._probe_once(None)
104 + assert registry.status()["encode"] == {"live": [url], "evicted": []}
105 +
106 + @pytest.mark.asyncio
107 + async def test_consecutive_failures_stop_routing(self, registry):
108 + registry.register(InstanceRecord(ENCODE, "http://e0:8000"))
109 + registry.register(InstanceRecord(ENCODE, "http://e1:8000"))
110 + for _ in range(registry._fail_threshold):
111 + await _probe_round(registry, healthy={"http://e1:8000"})
112 + assert registry.urls(ENCODE) == ["http://e1:8000"]
113 + assert registry.status()["encode"]["evicted"] == ["http://e0:8000"]
114 +
115 + @pytest.mark.asyncio
116 + @pytest.mark.parametrize("role", [ENCODE, DECODE])
117 + async def test_a_recovered_instance_rejoins_without_re_registering(
118 + self, registry, role
119 + ):
120 + """Restarting every encoder to recover from a blip is not acceptable."""
121 + registry.register(InstanceRecord(role, "http://e0:8000"))
122 + for _ in range(registry._fail_threshold):
123 + await _probe_round(registry, healthy=set())
124 + assert registry.urls(role) == []
125 +
126 + await _probe_round(registry, healthy={"http://e0:8000"}, now=1)
127 + assert registry.urls(role) == ["http://e0:8000"]
128 +
129 + @pytest.mark.asyncio
130 + async def test_re_registration_does_not_override_failed_health_checks(
131 + self, registry
132 + ):
133 + record = InstanceRecord(ENCODE, "http://e0:8000")
134 + registry.register(record)
135 + for _ in range(registry._fail_threshold):
136 + await _probe_round(registry, healthy=set())
137 + registry.register(record)
138 + assert registry.urls(ENCODE) == []
139 +
140 + @pytest.mark.asyncio
141 + @pytest.mark.parametrize("healthy", [False, True])
142 + async def test_inflight_probe_cannot_restore_unregistered_instance(
143 + self, registry, healthy
144 + ):
145 + record = InstanceRecord(ENCODE, "http://e0:8000")
146 + registry.register(record)
147 + for _ in range(registry._fail_threshold):
148 + await _probe_round(registry, healthy=set())
149 +
150 + async def probe(self, session, url):
151 + registry.unregister(url)
152 + return healthy
153 +
154 + with patch.object(InstanceRegistry, "_probe", probe):
155 + await registry._probe_once(None)
156 + assert registry.status()["encode"] == {"live": [], "evicted": []}