diff --git a/examples/disaggregated/disaggregated_encoder/README.md b/examples/disaggregated/disaggregated_encoder/README.md index 71a0d0e72e..ec10f23fc0 100644 --- a/examples/disaggregated/disaggregated_encoder/README.md +++ b/examples/disaggregated/disaggregated_encoder/README.md @@ -106,6 +106,82 @@ If you enable prefill instance (`--prefill-servers-urls` not disabled), you will | `--decode-servers-urls` | Comma-separated list of decode endpoints. Non-stream and stream paths both round-robin over this list. | | `--host`, `--port` | Bind address for the proxy itself (defaults: `0.0.0.0:8000`). | +### Dynamic registration + +Alternatively, let the external launcher register ready instances over HTTP. +No vLLM configuration changes or worker registration threads are required. +Set `ADMIN_API_KEY` on the proxy and supply it as `X-API-Key` for registration +and removal. These are trusted control-plane APIs; do not expose them publicly. + +```bash +export ADMIN_API_KEY="your-admin-key" +python disagg_epd_proxy.py --port 8000 --dynamic-registration +``` + +Omit the static server URL flags and register every stage. The roles determine +the topology: `encode` + `prefill_decode` gives E+PD; `encode` + `prefill` + +`decode` gives E+P+D. Standalone `decode` always requires an available `prefill`, +even when D registers first or all P instances go offline; otherwise requests +return `503`. The proxy rejects mixing combined PD with standalone P/D, +including unhealthy instances still in the registry. Explicitly remove the +old topology's P/D or PD registrations before switching topologies. + +After each instance is ready, the launcher registers its reachable HTTP URL: + +```bash +curl --fail-with-body http://proxy-host:8000/instances \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d '{"role":"encode","url":"http://e-host:8001"}' +curl --fail-with-body http://proxy-host:8000/instances \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d '{"role":"prefill_decode","url":"http://pd-host:8002"}' +``` + +For E+P+D, register P with `role: "prefill"` and D with `role: "decode"`. +Example and NIXL EC connectors need no additional registration fields. +For Mooncake, register `ec_zmq_addrs` on the **EC consumer** (`prefill_decode` +or `prefill`), using its configured `ec_ip` and `ec_port`. Keep the fixed-port +layout: supply the TP-rank-0 address for each DP replica, in DP-rank order; +replica `r` uses `ec_port + r * tensor_parallel_size`. The connector discovers +the remaining TP ranks itself. For example, DP=2, TP=2, `ec_port=19019`: + +```json +{ + "role": "prefill_decode", + "url": "http://pd-host:8002", + "dp_size": 2, + "ec_zmq_addrs": ["tcp://pd-host:19019", "tcp://pd-host:19021"] +} +``` + +The proxy selects one consumer replica and uses it for both the encoder push +and the consumer HTTP request. Standalone D does not need EC control addresses. +Port allocation and avoiding collisions remain the launcher's responsibility. + +Inspect or remove instances without restarting the proxy: + +```bash +curl http://proxy-host:8000/instances +curl --fail-with-body -X DELETE \ + 'http://proxy-host:8000/instances?url=http://e-host:8001' \ + -H "X-API-Key: $ADMIN_API_KEY" +``` + +Registration is idempotent. The proxy probes registered instances every +`--probe-interval` seconds (default 5), with `--probe-timeout` seconds per probe +(default 2). After `--fail-threshold` consecutive failures (default 3), it stops +sending new requests to that instance. Healthy instances rejoin automatically; +unreachable ones are forgotten after `--evicted-ttl` seconds (default 900; +`0` retains them indefinitely). Removal stops new routing; already routed +requests retain their selected endpoints, so drain requests before stopping +the instance. The launcher must re-register instances after a proxy restart. + +The example launch scripts also support this flow: export `ADMIN_API_KEY` and +set `DYNAMIC_REGISTRATION=1` when running `disagg_1e1pd_example.sh` or +`disagg_1e1p1d_example.sh`. Their default remains static routing. + +### Static configuration + The proxy batches images from the same user request assigned to the same encoder. Set the proxy environment variable `ENCODER_MAX_BATCH_SIZE` to limit the number of images per encoder subrequest. It defaults to `0` (unlimited); `1` sends each @@ -137,6 +213,6 @@ For E + P + D setup: ```bash $ python disagg_encoder_proxy.py \ --encode-servers-urls "http://e1:8001,http://e2:8001" \ - --prefill-servers-urls "http://p1:8003,http://p2:8004" \ + --prefill-servers-urls "http://p1:8003,http://p2:8004" \ --decode-servers-urls "http://d1:8005,http://d2:8006" ``` diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh index cb51333e90..57392a0408 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1p1d_example.sh @@ -189,17 +189,34 @@ wait_for_server "$DECODE_PORT" ############################################################################### # Proxy ############################################################################### +PROXY_ARGS=() +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}" + export ADMIN_API_KEY + PROXY_ARGS=(--dynamic-registration) +else + PROXY_ARGS=( + --encode-servers-urls "http://localhost:$ENCODE_PORT" + --prefill-servers-urls "http://localhost:$PREFILL_PORT" + --decode-servers-urls "http://localhost:$DECODE_PORT" + ) +fi python disagg_epd_proxy.py \ --host "0.0.0.0" \ --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "http://localhost:$PREFILL_PORT" \ - --decode-servers-urls "http://localhost:$DECODE_PORT" \ + "${PROXY_ARGS[@]}" \ >"${PROXY_LOG}" 2>&1 & PIDS+=($!) wait_for_server "$PROXY_PORT" +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + for endpoint in "encode:$ENCODE_PORT" "prefill:$PREFILL_PORT" "decode:$DECODE_PORT"; do + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}" + done +fi echo "All services are up!" ############################################################################### @@ -238,4 +255,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \ # cleanup echo "cleanup..." -cleanup \ No newline at end of file +cleanup diff --git a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh index 4bf60ea66d..44e43b8d7f 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh +++ b/examples/disaggregated/disaggregated_encoder/disagg_1e1pd_example.sh @@ -152,17 +152,33 @@ wait_for_server "$PREFILL_DECODE_PORT" ############################################################################### # Proxy ############################################################################### +PROXY_ARGS=() +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + : "${ADMIN_API_KEY:?Set ADMIN_API_KEY for dynamic registration}" + export ADMIN_API_KEY + PROXY_ARGS=(--dynamic-registration) +else + PROXY_ARGS=( + --encode-servers-urls "http://localhost:$ENCODE_PORT" + --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" + ) +fi python disagg_epd_proxy.py \ --host "0.0.0.0" \ --port "$PROXY_PORT" \ - --encode-servers-urls "http://localhost:$ENCODE_PORT" \ - --prefill-servers-urls "disable" \ - --decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \ + "${PROXY_ARGS[@]}" \ >"${PROXY_LOG}" 2>&1 & PIDS+=($!) wait_for_server "$PROXY_PORT" +if [[ "${DYNAMIC_REGISTRATION:-0}" == "1" ]]; then + for endpoint in "encode:$ENCODE_PORT" "prefill_decode:$PREFILL_DECODE_PORT"; do + curl --fail-with-body "http://127.0.0.1:$PROXY_PORT/instances" \ + -H "X-API-Key: $ADMIN_API_KEY" -H 'Content-Type: application/json' \ + -d "{\"role\":\"${endpoint%%:*}\",\"url\":\"http://127.0.0.1:${endpoint#*:}\"}" + done +fi echo "All services are up!" ############################################################################### @@ -201,4 +217,4 @@ curl http://127.0.0.1:"${PROXY_PORT}"/v1/chat/completions \ # cleanup echo "cleanup..." -cleanup \ No newline at end of file +cleanup diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index 3317b091fb..abae7eb65d 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -1,12 +1,13 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""disagg_encoder_proxy.py +"""OpenAI-compatible proxy for E+PD and E+P+D disaggregation. -Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two -clusters: - • encode (multimodal feature extraction) - • decode (language-model inference) +Use static server URLs, or --dynamic-registration for launcher-managed HTTP +registration of E, P/PD and D instances. Both modes route OpenAI-compatible +``/v1/chat/completions`` requests through encoder and inference clusters and +share encoder fan-out, metadata-only rewriting, prefill, and retry/streaming +forwarding. For MM input we: 1. Extract *every* image/audio/video item. @@ -20,22 +21,27 @@ from __future__ import annotations import argparse import asyncio +import enum import hashlib import itertools import json import logging import os import random +import secrets import time import uuid from collections.abc import AsyncIterator +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass, field from typing import Any import aiohttp import msgspec import uvicorn -from fastapi import FastAPI, HTTPException, Request +from fastapi import Depends, FastAPI, Header, HTTPException, Request from fastapi.responses import JSONResponse, Response, StreamingResponse +from pydantic import AnyHttpUrl, BaseModel, Field, model_validator ############################################################################### # FastAPI app & global state @@ -44,6 +50,340 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s") logger = logging.getLogger("proxy") +DEFAULT_PROBE_INTERVAL = 5.0 +DEFAULT_PROBE_TIMEOUT = 2.0 +DEFAULT_FAIL_THRESHOLD = 3 +# Stop probing an instance that has been down this long. 0 probes forever, +# which is what a cluster that restarts instances in place wants. +DEFAULT_EVICTED_TTL = 900.0 + + +class InstanceRole(str, enum.Enum): + ENCODE = "encode" + PREFILL = "prefill" + DECODE = "decode" + PREFILL_DECODE = "prefill_decode" + + +@dataclass +class InstanceRecord: + """One registered instance. + + Attributes: + role: Which stage this instance serves. + url: Base OpenAI-compatible URL, e.g. ``http://host:8000``. + ec_zmq_addrs: Mooncake TP-rank-0 control addresses, one per DP replica, + supplied by the launcher from the consumer's fixed port config. + dp_size: Data-parallel replicas behind `url`, so the proxy can pick + a replica and name the same one to both halves of a request. + + """ + + role: InstanceRole + url: str + ec_zmq_addrs: list[str] = field(default_factory=list) + dp_size: int = 1 + + +class InstanceRegistry: + def __init__( + self, + probe_interval: float = DEFAULT_PROBE_INTERVAL, + probe_timeout: float = DEFAULT_PROBE_TIMEOUT, + fail_threshold: int = DEFAULT_FAIL_THRESHOLD, + evicted_ttl: float = DEFAULT_EVICTED_TTL, + ): + self._probe_interval = probe_interval + self._probe_timeout = probe_timeout + self._fail_threshold = fail_threshold + self._evicted_ttl = evicted_ttl + + self._live: dict[str, InstanceRecord] = {} + self._evicted: dict[str, InstanceRecord] = {} + self._evicted_since: dict[str, float] = {} + self._fail_counts: dict[str, int] = {} + # One cursor per role, only ever incremented. Rebuilding it whenever + # the roster changes -- what an `itertools.cycle` over a mutable list + # forces -- restarts every fan-out at the first instance and hot-spots + # it after each registration. + self._cursors: dict[InstanceRole, int] = {role: 0 for role in InstanceRole} + self._replica_cursors: dict[str, int] = {} + self._probe_task: asyncio.Task | None = None + + def register(self, record: InstanceRecord) -> bool: + """Add or refresh an instance without overriding its health status.""" + roles = { + other.role + for other in itertools.chain(self._live.values(), self._evicted.values()) + } + split_roles = {InstanceRole.PREFILL, InstanceRole.DECODE} + if (record.role is InstanceRole.PREFILL_DECODE and roles & split_roles) or ( + record.role in split_roles and InstanceRole.PREFILL_DECODE in roles + ): + raise ValueError("Cannot mix prefill_decode with standalone prefill/decode") + key = record.url + previous = self._live.get(key) or self._evicted.get(key) + if previous is not None: + if previous == record: + return False + if previous.role is not record.role: + raise ValueError("Unregister the instance before changing its role") + target = self._live if key in self._live else self._evicted + target[key] = record + return False + self._live[key] = record + logger.info("Registered instance %s: %s", record.role.value, record.url) + return True + + def unregister(self, url: str) -> bool: + """Drop an instance for good, so a probe cannot bring it back.""" + key = url + if key not in self._live and key not in self._evicted: + return False + self._live.pop(key, None) + self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._fail_counts.pop(key, None) + self._replica_cursors.pop(key, None) + logger.info("Unregistered instance: %s", url) + return True + + def instances(self, role: InstanceRole) -> list[InstanceRecord]: + return [record for record in self._live.values() if record.role is role] + + def urls(self, role: InstanceRole) -> list[str]: + return [record.url for record in self.instances(role)] + + def pick(self, role: InstanceRole) -> InstanceRecord | None: + """Take the next instance of `role` in round-robin order.""" + picked = self.pick_many(role, 1) + return picked[0] if picked else None + + def pick_many(self, role: InstanceRole, count: int) -> list[InstanceRecord]: + """Take `count` instances, continuing the rotation across calls. + + A multimodal request fans out one encoder request per item, so the + assignment has to be contiguous with the previous request's rather + than restart at the first instance every time. + """ + alive = self.instances(role) + if not alive or count <= 0: + return [] + start = self._cursors[role] + self._cursors[role] = start + count + return [alive[(start + offset) % len(alive)] for offset in range(count)] + + def next_replica(self, record: InstanceRecord) -> int: + """Take the next data-parallel replica of `record`, round-robin. + + The encoder pushes to one replica's receive channel, so the request + has to run on that same replica; the caller names the rank to both + halves. + """ + if record.dp_size <= 1: + return 0 + cursor = self._replica_cursors.get(record.url, 0) + self._replica_cursors[record.url] = cursor + 1 + return cursor % record.dp_size + + def status(self) -> dict[str, Any]: + return { + role.value: { + "live": [record.url for record in self.instances(role)], + "evicted": [ + record.url + for record in self._evicted.values() + if record.role is role + ], + } + for role in InstanceRole + } + + def start_probing(self) -> None: + if self._probe_task is None and self._probe_interval > 0: + self._probe_task = asyncio.create_task(self._probe_loop()) + + async def stop_probing(self) -> None: + if self._probe_task is None: + return + self._probe_task.cancel() + with suppress(asyncio.CancelledError): + await self._probe_task + self._probe_task = None + + async def _probe_loop(self) -> None: + timeout = aiohttp.ClientTimeout(total=self._probe_timeout) + async with aiohttp.ClientSession(timeout=timeout) as session: + while True: + await asyncio.sleep(self._probe_interval) + try: + await self._probe_once(session) + except Exception: + logger.exception("EPD registry probe round failed") + + async def _probe_once(self, session: aiohttp.ClientSession) -> None: + targets = list(self._live.values()) + list(self._evicted.values()) + if not targets: + return + results = await asyncio.gather( + *(self._probe(session, record.url) for record in targets), + return_exceptions=True, + ) + now = time.monotonic() + for record, healthy in zip(targets, results): + current = self._live.get(record.url) or self._evicted.get(record.url) + if current is not record: + continue + if healthy is True: + self._on_probe_success(record) + else: + self._on_probe_failure(record, now) + self._drop_expired(now) + + async def _probe(self, session: aiohttp.ClientSession, url: str) -> bool: + async with session.get(f"{url}/health") as resp: + return resp.status == 200 + + def _on_probe_success(self, record: InstanceRecord) -> None: + key = record.url + self._fail_counts.pop(key, None) + if key in self._evicted: + self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._live[key] = record + logger.info( + "Instance %s (%s) is healthy again; routing resumed", + record.url, + record.role.value, + ) + + def _on_probe_failure(self, record: InstanceRecord, now: float) -> None: + key = record.url + if key not in self._live: + # Either already evicted, or unregistered while this probe was in + # flight. A round snapshots its targets and then awaits, so a + # removal inside that window would otherwise be undone by a result + # describing a registry that no longer exists. + return + failures = self._fail_counts.get(key, 0) + 1 + self._fail_counts[key] = failures + if failures < self._fail_threshold: + return + self._live.pop(key, None) + self._evicted[key] = record + self._evicted_since[key] = now + logger.warning( + "Instance %s (%s) failed %d consecutive probes; stopped routing " + "to it. It rejoins on its own once it responds again.", + record.url, + record.role.value, + failures, + ) + + def _drop_expired(self, now: float) -> None: + if self._evicted_ttl <= 0: + return + for key, since in list(self._evicted_since.items()): + if now - since < self._evicted_ttl: + continue + record = self._evicted.pop(key, None) + self._evicted_since.pop(key, None) + self._fail_counts.pop(key, None) + self._replica_cursors.pop(key, None) + if record is not None: + logger.warning("Instance %s stayed down; forgetting it", record.url) + + +class InstanceRegistration(BaseModel): + role: InstanceRole + url: AnyHttpUrl + ec_zmq_addrs: list[str] = Field(default_factory=list) + dp_size: int = Field(default=1, ge=1) + + @model_validator(mode="after") + def validate_consumer_addresses(self): + if self.ec_zmq_addrs and len(self.ec_zmq_addrs) != self.dp_size: + raise ValueError("Provide one Mooncake control address per DP replica") + return self + + +def require_admin_key(x_api_key: str = Header(default="")) -> None: + expected = os.getenv("ADMIN_API_KEY", "") + if not expected or not secrets.compare_digest(x_api_key, expected): + raise HTTPException(403, "Invalid admin API key") + + +@dataclass +class EPDProxyConfig: + probe_interval: float = 5.0 + probe_timeout: float = 2.0 + fail_threshold: int = 3 + evicted_ttl: float = 900.0 + + +@dataclass +class _Route: + """The instances one request was assigned to. + + Attributes: + consumer: The stage that receives the embedding, if any stage does. + consumer_zmq: The receive address named to the encoders. + dp_rank: Which replica of `consumer` was named, so the request can be + pinned to it. + + """ + + encoder_urls: list[str] + prefill: InstanceRecord | None + decode: InstanceRecord + consumer: InstanceRecord | None = None + consumer_zmq: str | None = None + dp_rank: int | None = None + + +class EPDProxy: + def __init__(self, registry: InstanceRegistry): + self.registry = registry + + # ---------------------------------------------------------------- # + # Routing # + # ---------------------------------------------------------------- # + def route(self, num_items: int) -> _Route: + decode = self.registry.pick(InstanceRole.PREFILL_DECODE) or self.registry.pick( + InstanceRole.DECODE + ) + if decode is None: + raise HTTPException( + status_code=503, detail="No decode instance is registered" + ) + prefill = ( + self.registry.pick(InstanceRole.PREFILL) + if decode.role is InstanceRole.DECODE + else None + ) + if decode.role is InstanceRole.DECODE and prefill is None: + raise HTTPException( + status_code=503, detail="No prefill instance is registered" + ) + encoders = self.registry.urls(InstanceRole.ENCODE) + if num_items and not encoders: + raise HTTPException( + status_code=503, detail="No encode instance is registered" + ) + route = _Route(encoder_urls=encoders, prefill=prefill, decode=decode) + self._name_consumer(route) + return route + + def _name_consumer(self, route: _Route) -> None: + """Pin the EC consumer replica and, for push connectors, its endpoint.""" + candidate = route.prefill or route.decode + route.consumer = candidate + rank = self.registry.next_replica(candidate) + route.dp_rank = rank if candidate.dp_size > 1 else None + if candidate.ec_zmq_addrs: + route.consumer_zmq = candidate.ec_zmq_addrs[rank] + + app = FastAPI() encode_session: aiohttp.ClientSession | None = None prefill_session: aiohttp.ClientSession | None = None @@ -448,6 +788,7 @@ async def maybe_prefill( req_data: dict, p_url: str, req_id: str, + dp_rank: int | None = None, ) -> dict: """- Do prefill-only task if p_url exist; - Return a new body carrying kv transfer params (for nixl connector) @@ -459,7 +800,7 @@ async def maybe_prefill( if p_url: logger.info("[%s] Processing through prefill: %s", req_id, p_url) - prefill_response = await process_prefill_stage(req_data, p_url, req_id) + prefill_response = await process_prefill_stage(req_data, p_url, req_id, dp_rank) # for nixl connector to facilitate kv transfer... prefill_response_json = msgspec.json.decode(await prefill_response.read()) kv_transfer_params = prefill_response_json.get("kv_transfer_params", {}) @@ -473,6 +814,7 @@ async def process_prefill_stage( req_data: dict, p_url: str, req_id: str, + dp_rank: int | None = None, ) -> dict: """Process request through Prefill stage and return kv_transfer_params""" logger.info("[%s] Sending prefill request to: %s", req_id, p_url) @@ -494,6 +836,8 @@ async def process_prefill_stage( del prefill_request["stream_options"] headers = {"x-request-id": req_id, "Content-Type": "application/json"} + if dp_rank is not None: + headers["X-data-parallel-rank"] = str(dp_rank) try: prefill_response = await prefill_session.post( f"{p_url}/v1/chat/completions", @@ -594,9 +938,7 @@ async def on_startup() -> None: ), ) encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) - if app.state.p_urls: - # only setup if prefill instance(s) exist - prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) + prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector) decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector) @@ -622,6 +964,7 @@ async def prepare_for_decode( e_urls: list[str], p_url: str, consumer_zmq: str | None, + prefill_dp_rank: int | None = None, ) -> tuple[dict, float, float]: """Encode, rewrite and prefill, returning the body to send to decode. @@ -642,7 +985,7 @@ async def prepare_for_decode( handles.update(ec_params) prepared = {**prepared, "ec_transfer_params": handles} _t2 = time.perf_counter() - prepared = await maybe_prefill(prepared, p_url, req_id) + prepared = await maybe_prefill(prepared, p_url, req_id, prefill_dp_rank) return prepared, _t1 - _t0, _t2 - _t1 @@ -654,12 +997,13 @@ async def forward_non_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, + prefill_dp_rank: int | None = None, ) -> Response: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() prepared, encode_s, rewrite_s = await prepare_for_decode( - req_data, req_id, e_urls, p_url, consumer_zmq + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank ) _t2 = time.perf_counter() @@ -733,12 +1077,13 @@ async def forward_stream( d_url: str, consumer_zmq: str | None, dp_rank: int | None = None, + prefill_dp_rank: int | None = None, ) -> AsyncIterator[bytes]: try: for attempt in range(DECODE_RETRIES + 1): _t0 = time.perf_counter() prepared, encode_s, rewrite_s = await prepare_for_decode( - req_data, req_id, e_urls, p_url, consumer_zmq + req_data, req_id, e_urls, p_url, consumer_zmq, prefill_dp_rank ) _t2 = time.perf_counter() @@ -976,10 +1321,147 @@ async def stop_profile(request: Request): return await _profile_cmd("stop", body, e_url, p_url, d_url) +def build_app(config: EPDProxyConfig | None = None) -> FastAPI: + config = config or EPDProxyConfig() + registry = InstanceRegistry( + probe_interval=config.probe_interval, + probe_timeout=config.probe_timeout, + fail_threshold=config.fail_threshold, + evicted_ttl=config.evicted_ttl, + ) + proxy = EPDProxy(registry) + + @asynccontextmanager + async def lifespan(app: FastAPI): + await on_startup() + try: + registry.start_probing() + yield + finally: + await registry.stop_probing() + await on_shutdown() + + app = FastAPI(lifespan=lifespan) + app.state.proxy = proxy + app.state.registry = registry + + @app.get("/instances") + async def list_instances(): + return registry.status() + + @app.post("/instances", dependencies=[Depends(require_admin_key)]) + async def register_instance(body: InstanceRegistration): + if body.ec_zmq_addrs and body.role not in ( + InstanceRole.PREFILL, + InstanceRole.PREFILL_DECODE, + ): + raise HTTPException(400, "Only EC consumers accept Mooncake addresses") + try: + created = registry.register( + InstanceRecord( + body.role, + str(body.url).rstrip("/"), + body.ec_zmq_addrs, + body.dp_size, + ) + ) + except ValueError as error: + raise HTTPException(409, str(error)) from error + return {"registered": created} + + @app.delete("/instances", dependencies=[Depends(require_admin_key)]) + async def unregister_instance(url: AnyHttpUrl): + return {"removed": registry.unregister(str(url).rstrip("/"))} + + @app.post("/v1/chat/completions") + async def chat_completions(request: Request): + req_data = msgspec.json.decode(await request.body()) + req_id = request.headers.get("x-request-id", str(uuid.uuid4())) + route = proxy.route(len(extract_mm_items(req_data))) + args = ( + req_data, + req_id, + route.encoder_urls, + route.prefill.url if route.prefill else None, + route.decode.url, + route.consumer_zmq, + ) + ranks = { + "dp_rank": route.dp_rank if route.prefill is None else None, + "prefill_dp_rank": route.dp_rank if route.prefill is not None else None, + } + if req_data.get("stream", False): + return StreamingResponse( + forward_stream(*args, **ranks), media_type="text/event-stream" + ) + return await forward_non_stream(*args, **ranks) + + @app.get("/v1/models") + async def list_models(): + decode = registry.pick(InstanceRole.PREFILL_DECODE) or registry.pick( + InstanceRole.DECODE + ) + if decode is None: + raise HTTPException( + status_code=503, detail="No decode instance is registered" + ) + async with decode_session.get(f"{decode.url}/v1/models") as resp: + resp.raise_for_status() + return await resp.json() + + @app.get("/health") + async def health(): + status = registry.status() + # An empty roster is not unhealthy: the proxy is meant to come up + # before anything registers with it. + return JSONResponse({"proxy": "healthy", "instances": status}) + + @app.post("/start_profile") + async def start_profile(request: Request): + return await _profile(registry, "start", await request.json()) + + @app.post("/stop_profile") + async def stop_profile(request: Request): + return await _profile(registry, "stop", await request.json()) + + return app + + +async def _profile(registry: InstanceRegistry, cmd: str, payload: dict) -> dict: + headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"} + targets = {role.value: registry.pick(role) for role in InstanceRole} + results = await asyncio.gather( + *( + _post_if_available( + decode_session, f"{record.url}/{cmd}_profile", payload, headers + ) + for record in targets.values() + if record is not None + ) + ) + reachable = [result for result in results if result is not None] + if not reachable: + raise HTTPException( + status_code=503, + detail="Profiling endpoints are disabled on every instance", + ) + live = [name for name, record in targets.items() if record is not None] + return dict(zip(live, results)) + + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--dynamic-registration", + action="store_true", + help="Enable launcher-managed HTTP registration for E+PD or E+P+D.", + ) + parser.add_argument("--probe-interval", type=float, default=DEFAULT_PROBE_INTERVAL) + parser.add_argument("--probe-timeout", type=float, default=DEFAULT_PROBE_TIMEOUT) + parser.add_argument("--fail-threshold", type=int, default=DEFAULT_FAIL_THRESHOLD) + parser.add_argument("--evicted-ttl", type=float, default=DEFAULT_EVICTED_TTL) parser.add_argument( "--log-requests", action="store_true", @@ -995,12 +1477,12 @@ if __name__ == "__main__": ) parser.add_argument( "--encode-servers-urls", - required=True, + default="", help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")', ) parser.add_argument( "--prefill-servers-urls", - required=True, + default="none", help=( 'Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") ' 'to enable E->P->D, set "disable" or "none" to enable E->PD' @@ -1008,7 +1490,7 @@ if __name__ == "__main__": ) parser.add_argument( "--decode-servers-urls", - required=True, + default="", help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")', ) parser.add_argument( @@ -1043,58 +1525,87 @@ if __name__ == "__main__": ) args = parser.parse_args() + if args.fail_threshold < 1: + parser.error("--fail-threshold must be at least 1") if args.log_requests: logging.getLogger().setLevel(logging.DEBUG) app.middleware("http")(log_requests) NO_REWRITE = args.no_rewrite DECODE_RETRIES = max(0, args.decode_retries) - app.state.e_urls = [ - u.strip() for u in args.encode_servers_urls.split(",") if u.strip() - ] - app.state.d_urls = [ - u.strip() for u in args.decode_servers_urls.split(",") if u.strip() - ] - app.state.d_ec_urls = [ - u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() - ] - if args.ec_consumer_dp_size < 1: - parser.error("--ec-consumer-dp-size must be at least 1") - app.state.ec_consumer_dp_size = args.ec_consumer_dp_size - app.state.replica_counter = itertools.count() - expected = len(app.state.d_urls) * args.ec_consumer_dp_size - if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: - parser.error( - "--ec-consumer-zmq-addrs must contain one address per consumer " - f"replica: expected {expected} " - f"({len(app.state.d_urls)} servers x {args.ec_consumer_dp_size} replicas), " - f"got {len(app.state.d_ec_urls)}" - ) - # handle prefill instances - if args.prefill_servers_urls.lower() in ("disable", "none", ""): - app.state.p_urls = [] - logger.info( - "Disaggregated prefill phase explicitly disabled by user. Running E + PD..." + if args.dynamic_registration: + if not os.getenv("ADMIN_API_KEY"): + parser.error("--dynamic-registration requires ADMIN_API_KEY") + if ( + args.encode_servers_urls + or args.prefill_servers_urls.lower() not in ("disable", "none", "") + or args.decode_servers_urls + or args.ec_consumer_zmq_addrs + ): + parser.error( + "With --dynamic-registration, instances register through /instances" + ) + app = build_app( + EPDProxyConfig( + probe_interval=args.probe_interval, + probe_timeout=args.probe_timeout, + fail_threshold=args.fail_threshold, + evicted_ttl=args.evicted_ttl, + ) ) + if args.log_requests: + app.middleware("http")(log_requests) else: - app.state.p_urls = [ - u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() + if not args.encode_servers_urls or not args.decode_servers_urls: + parser.error( + "Static routing requires --encode-servers-urls " + "and --decode-servers-urls" + ) + app.state.e_urls = [ + u.strip() for u in args.encode_servers_urls.split(",") if u.strip() ] - logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") - try: - validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) - except ValueError as exc: - parser.error(str(exc)) + app.state.d_urls = [ + u.strip() for u in args.decode_servers_urls.split(",") if u.strip() + ] + app.state.d_ec_urls = [ + u.strip() for u in args.ec_consumer_zmq_addrs.split(",") if u.strip() + ] + if args.ec_consumer_dp_size < 1: + parser.error("--ec-consumer-dp-size must be at least 1") + app.state.ec_consumer_dp_size = args.ec_consumer_dp_size + app.state.replica_counter = itertools.count() + expected = len(app.state.d_urls) * args.ec_consumer_dp_size + if app.state.d_ec_urls and len(app.state.d_ec_urls) != expected: + parser.error( + "--ec-consumer-zmq-addrs must contain one address per consumer " + f"replica: expected {expected} " + f"({len(app.state.d_urls)} servers x " + f"{args.ec_consumer_dp_size} replicas), " + f"got {len(app.state.d_ec_urls)}" + ) + # handle prefill instances + if args.prefill_servers_urls.lower() in ("disable", "none", ""): + app.state.p_urls = [] + logger.info("Disaggregated prefill disabled. Running E + PD...") + else: + app.state.p_urls = [ + u.strip() for u in args.prefill_servers_urls.split(",") if u.strip() + ] + logger.info("Disaggregated prefill phase is enabled. Running E + P + D...") + try: + validate_ec_consumer_routing(app.state.p_urls, app.state.d_ec_urls) + except ValueError as exc: + parser.error(str(exc)) - logger.info("Proxy listening on %s:%s", args.host, args.port) - logger.info("Encode servers: %s", app.state.e_urls) - logger.info("Prefill instances %s", app.state.p_urls) - logger.info("Decode servers: %s", app.state.d_urls) - if app.state.ec_consumer_dp_size > 1: - logger.info( - "EC consumer replicas per server: %d (control addresses: %s)", - app.state.ec_consumer_dp_size, - app.state.d_ec_urls, - ) + logger.info("Proxy listening on %s:%s", args.host, args.port) + logger.info("Encode servers: %s", app.state.e_urls) + logger.info("Prefill instances %s", app.state.p_urls) + logger.info("Decode servers: %s", app.state.d_urls) + if app.state.ec_consumer_dp_size > 1: + logger.info( + "EC consumer replicas per server: %d (control addresses: %s)", + app.state.ec_consumer_dp_size, + app.state.d_ec_urls, + ) uvicorn.run( app, diff --git a/tests/v1/ec_connector/unit/test_epd_proxy.py b/tests/v1/ec_connector/unit/test_epd_proxy.py new file mode 100644 index 0000000000..13496a7ffb --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_proxy.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Routing and registration behaviour of the EPD proxy.""" + +import asyncio +from typing import Any +from unittest.mock import AsyncMock, Mock + +import msgspec +import pytest +from fastapi import HTTPException +from fastapi.testclient import TestClient + +from examples.disaggregated.disaggregated_encoder import ( + disagg_epd_proxy as proxy_module, +) +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + EPDProxy, + EPDProxyConfig, + InstanceRecord, + InstanceRegistry, + InstanceRole, + build_app, + extract_mm_items, +) + +ENCODE = InstanceRole.ENCODE +PREFILL = InstanceRole.PREFILL +DECODE = InstanceRole.DECODE +PD = InstanceRole.PREFILL_DECODE + +IMAGE_ITEM = {"type": "image_url", "image_url": {"url": "http://img/0.png"}} + + +@pytest.fixture +def proxy(): + registry = InstanceRegistry(probe_interval=0) + return EPDProxy(registry) + + +class TestRouting: + def test_no_decode_instance_is_service_unavailable(self, proxy): + """The proxy comes up before anything registers, so this is normal.""" + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + + def test_media_without_an_encoder_is_service_unavailable(self, proxy): + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + assert proxy.route(num_items=0).decode.url == "http://d0:8000" + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=1) + assert excinfo.value.status_code == 503 + + def test_prefill_is_optional(self, proxy): + """An E+PD deployment registers no prefill instance at all.""" + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + assert proxy.route(num_items=0).prefill is None + + def test_encoder_roster_tracks_new_instances(self, proxy): + proxy.registry.register(InstanceRecord(PD, "http://d0:8000")) + for index in range(2): + proxy.registry.register(InstanceRecord(ENCODE, f"http://e{index}:8000")) + assert proxy.route(num_items=3).encoder_urls == [ + "http://e0:8000", + "http://e1:8000", + ] + proxy.registry.unregister("http://e0:8000") + assert proxy.route(num_items=3).encoder_urls == ["http://e1:8000"] + + @pytest.mark.parametrize("unregister", [False, True]) + def test_standalone_decode_always_requires_healthy_prefill(self, unregister): + app = build_app() + proxy = app.state.proxy + proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + proxy.registry.register(InstanceRecord(PREFILL, "http://p0:8000")) + route = proxy.route(num_items=0) + assert route.decode.url == "http://d0:8000" + assert route.consumer.url == "http://p0:8000" + if unregister: + proxy.registry.unregister("http://p0:8000") + else: + for _ in range(proxy.registry._fail_threshold): + proxy.registry._on_probe_failure(route.prefill, 0) + with pytest.raises(HTTPException) as excinfo: + proxy.route(num_items=0) + assert excinfo.value.status_code == 503 + + +class TestConsumerAddress: + """Which stage receives the embedding depends on the topology.""" + + def test_shared_storage_connectors_name_no_target(self, proxy): + """Rotate replicas without requiring a push endpoint.""" + proxy.registry.register(InstanceRecord(PD, "http://d0:8000", dp_size=2)) + route = proxy.route(num_items=0) + assert route.consumer_zmq is None + assert route.dp_rank == 0 + assert proxy.route(num_items=0).dp_rank == 1 + + def test_decode_is_the_consumer_when_prefill_is_not_split_out(self, proxy): + proxy.registry.register( + InstanceRecord(PD, "http://d0:8000", ec_zmq_addrs=["tcp://d0:20001"]) + ) + assert proxy.route(num_items=0).consumer_zmq == "tcp://d0:20001" + + def test_prefill_is_the_consumer_when_it_is_split_out(self, proxy): + """In E+P+D the prefill instance consumes the embedding, not decode. + + Reading the address out of whichever list it came from -- rather than + from the record that reported one -- sends the push to the wrong + instance in this topology. + """ + proxy.registry.register( + InstanceRecord(PREFILL, "http://p0:8000", ec_zmq_addrs=["tcp://p0:20001"]) + ) + proxy.registry.register(InstanceRecord(DECODE, "http://d0:8000")) + assert proxy.route(num_items=0).consumer_zmq == "tcp://p0:20001" + + +class TestProxyApi: + @pytest.fixture + def client(self, monkeypatch): + monkeypatch.setenv("ADMIN_API_KEY", "test-key") + config = EPDProxyConfig(probe_interval=0) + with TestClient(build_app(config)) as client: + client.headers["X-API-Key"] = "test-key" + yield client + + def test_proxy_serves_before_anything_registers(self, client): + assert client.get("/health").status_code == 200 + assert client.get("/instances").json()["encode"]["live"] == [] + + def test_requests_are_refused_until_a_decoder_registers(self, client): + response = client.post( + "/v1/chat/completions", json={"model": "m", "messages": []} + ) + assert response.status_code == 503 + assert client.get("/v1/models").status_code == 503 + + def test_http_registration_keeps_consumer_replica_and_address_together( + self, client + ): + payload: dict[str, Any] = { + "role": "prefill_decode", + "url": "http://d0:8000/", + "dp_size": 2, + "ec_zmq_addrs": ["tcp://d0:20001", "tcp://d0:20003"], + } + assert client.post("/instances", json=payload).json() == {"registered": True} + assert client.post("/instances", json=payload).json() == {"registered": False} + proxy = client.app.state.proxy + for rank in [0, 1, 0]: + route = proxy.route(0) + assert route.dp_rank == rank + assert route.consumer_zmq == payload["ec_zmq_addrs"][rank] + assert client.delete("/instances", params={"url": payload["url"]}).json() == { + "removed": True + } + assert proxy.registry.pick(PD) is None + + def test_registration_rejects_invalid_topology(self, client): + cases: list[tuple[dict[str, Any], int]] = [ + ({"dp_size": 0}, 422), + ({"dp_size": 2, "ec_zmq_addrs": ["tcp://d0:20001"]}, 422), + ({"role": "decode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), + ({"role": "encode", "ec_zmq_addrs": ["tcp://d0:20001"]}, 400), + ] + for fields, status in cases: + response = client.post( + "/instances", + json={"url": "http://d0:8000", "role": "prefill_decode", **fields}, + ) + assert response.status_code == status + + def test_registration_requires_admin_key(self, client): + client.headers.pop("X-API-Key") + assert ( + client.post( + "/instances", json={"role": "encode", "url": "http://e0:8000"} + ).status_code + == 403 + ) + assert ( + client.delete("/instances", params={"url": "http://e0:8000"}).status_code + == 403 + ) + + @pytest.mark.parametrize("split_role", ["prefill", "decode"]) + @pytest.mark.parametrize("combined_first", [False, True]) + def test_registration_rejects_mixed_topologies( + self, client, split_role, combined_first + ): + roles = ["prefill_decode", split_role] + if not combined_first: + roles.reverse() + first = {"role": roles[0], "url": "http://first:8000"} + second = {"role": roles[1], "url": "http://second:8000"} + assert client.post("/instances", json=first).status_code == 200 + assert client.post("/instances", json=second).status_code == 409 + # Health eviction must not silently permit a topology switch. + registry = client.app.state.registry + record = registry._live[first["url"]] + for _ in range(registry._fail_threshold): + registry._on_probe_failure(record, 0) + assert client.post("/instances", json=second).status_code == 409 + assert client.delete("/instances", params={"url": first["url"]}).json() == { + "removed": True + } + assert client.post("/instances", json=second).status_code == 200 + + +def test_extract_mm_items_finds_media_across_messages(): + req = { + "messages": [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [IMAGE_ITEM, {"type": "text", "text": "hi"}]}, + {"role": "user", "content": [IMAGE_ITEM]}, + ] + } + assert extract_mm_items(req) == [IMAGE_ITEM, IMAGE_ITEM] + + +@pytest.mark.parametrize("push", [False, True]) +def test_encoder_handles_and_json_metadata_survive_rewrite(proxy, push, monkeypatch): + """Keep main's hash-keyed handles and per-item push IDs after the proxy move.""" + proxy.registry.register(InstanceRecord(ENCODE, "http://encoder")) + proxy.registry.register( + InstanceRecord( + PD, + "http://decode", + ec_zmq_addrs=["tcp://decode:14579"] if push else [], + ) + ) + handle = { + "metadata": {"image_grid_thw": [[1, 2, 3]]}, + "item_indices": [0, 1], + "peer_port": 5601, + } + response = Mock( + status=200, + read=AsyncMock( + return_value=msgspec.json.encode( + {"ec_transfer_params": {"engine-hash": handle}} + ) + ), + ) + session = Mock(post=AsyncMock(return_value=response)) + monkeypatch.setattr(proxy_module, "encode_session", session) + original = {"messages": [{"role": "user", "content": [IMAGE_ITEM, IMAGE_ITEM]}]} + route = proxy.route(2) + prepared, _, _ = asyncio.run( + proxy_module.prepare_for_decode( + original, "request", route.encoder_urls, None, route.consumer_zmq + ) + ) + assert "ec_transfer_params" not in original + assert prepared["ec_transfer_params"]["engine-hash"] == handle + items = prepared["ec_transfer_params"]["ec_items"] + assert [item["mm_hash"] for item in items] == ["engine-hash", "engine-hash"] + assert items[0]["transfer_id"] != items[1]["transfer_id"] + assert session.post.call_count == 1 + sent = msgspec.json.decode(session.post.call_args.kwargs["data"]).get( + "ec_transfer_params" + ) + if push: + assert sent == { + "consumer_zmq": "tcp://decode:14579", + "ec_items": [{"transfer_id": item["transfer_id"]} for item in items], + } + else: + assert sent is None + for index in range(2): + assert prepared["messages"][0]["content"][index]["image_embeds"] == { + "image_grid_thw": [1, 2, 3] + } diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 0c48dc84d6..8db22fe1b9 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -8,6 +8,7 @@ change to them is what these tests catch. import asyncio import importlib.util +import sys from pathlib import Path from types import SimpleNamespace from typing import Any @@ -27,6 +28,7 @@ def proxy(): spec = importlib.util.spec_from_file_location("disagg_epd_proxy_retry", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module spec.loader.exec_module(module) return module @@ -48,7 +50,7 @@ def test_maybe_prefill_leaves_the_caller_body_untouched(proxy, monkeypatch): """ served = [{"remote_block_ids": [1, 2]}, {}] - async def _stage(req_data, p_url, req_id): + async def _stage(req_data, p_url, req_id, dp_rank=None): assert "kv_transfer_params" not in req_data return _Response(served.pop(0)) @@ -222,7 +224,7 @@ def test_a_decode_retry_does_not_inherit_the_previous_handles(proxy, monkeypatch _EncoderSession([{"encoder-side-hash": handle}, {}]), ) - async def _no_prefill(req_data, p_url, req_id): + async def _no_prefill(req_data, p_url, req_id, dp_rank=None): return req_data monkeypatch.setattr(proxy, "maybe_prefill", _no_prefill) @@ -280,6 +282,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("prefill", [False, True]) +@pytest.mark.parametrize("dynamic", [False, True]) @pytest.mark.parametrize( "request_options", [ @@ -293,7 +296,7 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch): ], ) async def test_http_roundtrip_preserves_payload_and_response_bytes( - proxy, monkeypatch, stream, prefill, request_options + proxy, monkeypatch, stream, prefill, dynamic, request_options ): """Preserve media and scheduling context across HTTP hops and retries.""" seen: dict[str, list[dict]] = {"encode": [], "prefill": [], "decode": []} @@ -307,6 +310,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( async def backend(request): stage = request.match_info["stage"] + if dynamic: + consumer = "prefill" if prefill else "decode" + expected_rank = "0" if stage == consumer else None + assert request.headers.get("X-data-parallel-rank") == expected_rank assert request.content_type == "application/json" body = await request.json() seen[stage].append(body) @@ -354,6 +361,10 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( monkeypatch.setattr(proxy.app.state, "ec_consumer_dp_size", 1, raising=False) monkeypatch.setattr(proxy, "DECODE_RETRIES", 1) monkeypatch.setattr(proxy, "NO_REWRITE", False) + target_app = proxy.app + if dynamic: + monkeypatch.setenv("ADMIN_API_KEY", "test-key") + target_app = proxy.build_app(proxy.EPDProxyConfig(probe_interval=0)) item = {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJj"}} body = { "model": "test", @@ -377,9 +388,37 @@ async def test_http_roundtrip_preserves_payload_and_response_bytes( await proxy.on_startup() try: async with httpx.AsyncClient( - transport=httpx.ASGITransport(app=proxy.app), base_url="http://proxy" + transport=httpx.ASGITransport(app=target_app), base_url="http://proxy" ) as client: + if dynamic: + registrations = [("encode", "encode")] + registrations += ( + [("prefill", "prefill"), ("decode", "decode")] + if prefill + else [("prefill_decode", "decode")] + ) + for role, stage in registrations: + consumer = "prefill" if prefill else "decode" + registered = await client.post( + "/instances", + headers={"X-API-Key": "test-key"}, + json={ + "role": role, + "url": str(server.make_url(f"/{stage}")), + "dp_size": 2 if stage == consumer else 1, + }, + ) + assert registered.status_code == 200 response = await client.post("/v1/chat/completions", json=body) + if dynamic: + removed = await client.delete( + "/instances", + headers={"X-API-Key": "test-key"}, + params={"url": str(server.make_url("/decode"))}, + ) + assert removed.json() == {"removed": True} + unavailable = await client.post("/v1/chat/completions", json=body) + assert unavailable.status_code == 503 assert response.status_code == 200 assert response.content == expected assert response.headers["content-type"].startswith( diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py index 558a80f9cf..8e1a730ca6 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py @@ -14,6 +14,7 @@ across calls so consecutive requests keep advancing through every URL. """ import importlib.util +import sys from collections import Counter from pathlib import Path @@ -27,6 +28,7 @@ def _load_proxy_module(): spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module spec.loader.exec_module(module) return module diff --git a/tests/v1/ec_connector/unit/test_epd_registry.py b/tests/v1/ec_connector/unit/test_epd_registry.py new file mode 100644 index 0000000000..be5a8e9b9b --- /dev/null +++ b/tests/v1/ec_connector/unit/test_epd_registry.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Dynamic registration and liveness for the EPD proxy's instance registry.""" + +from unittest.mock import patch + +import pytest + +from examples.disaggregated.disaggregated_encoder.disagg_epd_proxy import ( + InstanceRecord, + InstanceRegistry, + InstanceRole, +) + +ENCODE = InstanceRole.ENCODE +DECODE = InstanceRole.DECODE + + +@pytest.fixture +def registry(): + return InstanceRegistry(probe_interval=0, fail_threshold=3, evicted_ttl=60.0) + + +def _fill(registry, role, count, prefix="http://h"): + for index in range(count): + registry.register(InstanceRecord(role, f"{prefix}{index}:8000")) + + +async def _probe_round(registry, healthy: set[str], now: float = 0.0): + """Run one probe pass where only `healthy` URLs answer.""" + + async def fake_probe(self, session, url): + return url in healthy + + with ( + patch.object(InstanceRegistry, "_probe", fake_probe), + patch( + "examples.disaggregated.disaggregated_encoder.disagg_epd_proxy.time.monotonic", + return_value=now, + ), + ): + await registry._probe_once(session=None) + + +class TestRegistration: + def test_registering_is_idempotent_and_reports_novelty(self, registry): + record = InstanceRecord(ENCODE, "http://e0:8000") + assert registry.register(record) is True + assert registry.register(record) is False + assert registry.urls(ENCODE) == ["http://e0:8000"] + + def test_transfer_addresses_stay_with_the_registered_consumer(self, registry): + """Roster updates must not shift a consumer's transfer addresses.""" + registry.register(InstanceRecord(DECODE, "http://d0:8000")) + registry.register( + InstanceRecord( + DECODE, + "http://d0:8000", + ec_zmq_addrs=["tcp://d0:20001", "tcp://d0:20002"], + dp_size=2, + ) + ) + picked = registry.pick(DECODE) + assert picked.ec_zmq_addrs == ["tcp://d0:20001", "tcp://d0:20002"] + assert picked.dp_size == 2 + + def test_unregister_is_terminal(self, registry): + registry.register(InstanceRecord(ENCODE, "http://e0:8000")) + assert registry.unregister("http://e0:8000") is True + assert registry.unregister("http://e0:8000") is False + assert registry.pick(ENCODE) is None + + def test_round_robin_survives_registration(self, registry): + _fill(registry, ENCODE, 2) + assert [record.url for record in registry.pick_many(ENCODE, 2)] == [ + "http://h0:8000", + "http://h1:8000", + ] + registry.register(InstanceRecord(ENCODE, "http://h2:8000")) + assert registry.pick(ENCODE).url == "http://h2:8000" + + +class TestLiveness: + @pytest.mark.asyncio + async def test_identical_reregistration_preserves_inflight_probe_results( + self, registry + ): + """Repeated registrations must not suppress eviction or recovery.""" + url = "http://e0:8000" + registry.register(InstanceRecord(ENCODE, url)) + healthy = False + + async def probe(self, session, url): + registry.register(InstanceRecord(ENCODE, url)) + return healthy + + with patch.object(InstanceRegistry, "_probe", probe): + for _ in range(registry._fail_threshold): + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [], "evicted": [url]} + + healthy = True + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [url], "evicted": []} + + @pytest.mark.asyncio + async def test_consecutive_failures_stop_routing(self, registry): + registry.register(InstanceRecord(ENCODE, "http://e0:8000")) + registry.register(InstanceRecord(ENCODE, "http://e1:8000")) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy={"http://e1:8000"}) + assert registry.urls(ENCODE) == ["http://e1:8000"] + assert registry.status()["encode"]["evicted"] == ["http://e0:8000"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("role", [ENCODE, DECODE]) + async def test_a_recovered_instance_rejoins_without_re_registering( + self, registry, role + ): + """Restarting every encoder to recover from a blip is not acceptable.""" + registry.register(InstanceRecord(role, "http://e0:8000")) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + assert registry.urls(role) == [] + + await _probe_round(registry, healthy={"http://e0:8000"}, now=1) + assert registry.urls(role) == ["http://e0:8000"] + + @pytest.mark.asyncio + async def test_re_registration_does_not_override_failed_health_checks( + self, registry + ): + record = InstanceRecord(ENCODE, "http://e0:8000") + registry.register(record) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + registry.register(record) + assert registry.urls(ENCODE) == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize("healthy", [False, True]) + async def test_inflight_probe_cannot_restore_unregistered_instance( + self, registry, healthy + ): + record = InstanceRecord(ENCODE, "http://e0:8000") + registry.register(record) + for _ in range(registry._fail_threshold): + await _probe_round(registry, healthy=set()) + + async def probe(self, session, url): + registry.unregister(url) + return healthy + + with patch.object(InstanceRegistry, "_probe", probe): + await registry._probe_once(None) + assert registry.status()["encode"] == {"live": [], "evicted": []}