| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """This file provides a disaggregated prefilling proxy demo to demonstrate an |
| 4 | example usage of XpYd disaggregated prefilling. |
| 5 | We can launch multiple vllm instances (2 for prefill and 2 for decode), and |
| 6 | launch this proxy demo through: |
| 7 | python3 examples/disaggregated/disaggregated_serving/disagg_proxy_demo.py \ |
| 8 | --model $model_name \ |
| 9 | --prefill localhost:8100 localhost:8101 \ |
| 10 | --decode localhost:8200 localhost:8201 \ |
| 11 | --port 8000 |
| 12 | |
| 13 | Note: This demo will be removed once the PDController implemented in PR 15343 |
| 14 | (https://github.com/vllm-project/vllm/pull/15343) supports XpYd. |
| 15 | """ |
| 16 | |
| 17 | import argparse |
| 18 | import ipaddress |
| 19 | import itertools |
| 20 | import json |
| 21 | import logging |
| 22 | import os |
| 23 | import sys |
| 24 | from abc import ABC, abstractmethod |
| 25 | from collections.abc import Callable |
| 26 | |
| 27 | import aiohttp |
| 28 | import requests |
| 29 | import uvicorn |
| 30 | from fastapi import APIRouter, Depends, FastAPI, Header, HTTPException, Request, status |
| 31 | from fastapi.responses import JSONResponse, StreamingResponse |
| 32 | |
| 33 | AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) |
| 34 | logger = logging.getLogger() |
| 35 | logging.basicConfig(level=logging.INFO) |
| 36 | |
| 37 | |
| 38 | class SchedulingPolicy(ABC): |
| 39 | @abstractmethod |
| 40 | def schedule(self, cycler: itertools.cycle): |
| 41 | raise NotImplementedError("Scheduling Proxy is not set.") |
| 42 | |
| 43 | |
| 44 | class Proxy: |
| 45 | def __init__( |
| 46 | self, |
| 47 | prefill_instances: list[str], |
| 48 | decode_instances: list[str], |
| 49 | model: str, |
| 50 | scheduling_policy: SchedulingPolicy, |
| 51 | custom_create_completion: Callable[[Request], StreamingResponse] | None = None, |
| 52 | custom_create_chat_completion: Callable[[Request], StreamingResponse] |
| 53 | | None = None, |
| 54 | ): |
| 55 | self.prefill_instances = prefill_instances |
| 56 | self.decode_instances = decode_instances |
| 57 | self.prefill_cycler = itertools.cycle(prefill_instances) |
| 58 | self.decode_cycler = itertools.cycle(decode_instances) |
| 59 | self.model = model |
| 60 | self.scheduling_policy = scheduling_policy |
| 61 | self.custom_create_completion = custom_create_completion |
| 62 | self.custom_create_chat_completion = custom_create_chat_completion |
| 63 | self.router = APIRouter() |
| 64 | self.setup_routes() |
| 65 | |
| 66 | def setup_routes(self): |
| 67 | self.router.post( |
| 68 | "/v1/completions", dependencies=[Depends(self.validate_json_request)] |
| 69 | )( |
| 70 | self.custom_create_completion |
| 71 | if self.custom_create_completion |
| 72 | else self.create_completion |
| 73 | ) |
| 74 | self.router.post( |
| 75 | "/v1/chat/completions", dependencies=[Depends(self.validate_json_request)] |
| 76 | )( |
| 77 | self.custom_create_chat_completion |
| 78 | if self.custom_create_chat_completion |
| 79 | else self.create_chat_completion |
| 80 | ) |
| 81 | self.router.get("/status", response_class=JSONResponse)(self.get_status) |
| 82 | self.router.post( |
| 83 | "/instances/add", dependencies=[Depends(self.api_key_authenticate)] |
| 84 | )(self.add_instance_endpoint) |
| 85 | |
| 86 | async def validate_json_request(self, raw_request: Request): |
| 87 | content_type = raw_request.headers.get("content-type", "").lower() |
| 88 | if content_type != "application/json": |
| 89 | raise HTTPException( |
| 90 | status_code=415, |
| 91 | detail="Unsupported Media Type: Only 'application/json' is allowed", |
| 92 | ) |
| 93 | |
| 94 | def api_key_authenticate(self, x_api_key: str = Header(...)): |
| 95 | expected_api_key = os.environ.get("ADMIN_API_KEY") |
| 96 | if not expected_api_key: |
| 97 | logger.error("ADMIN_API_KEY is not set in the environment.") |
| 98 | raise HTTPException( |
| 99 | status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, |
| 100 | detail="Server configuration error.", |
| 101 | ) |
| 102 | if x_api_key != expected_api_key: |
| 103 | logger.warning("Unauthorized access attempt with API Key: %s", x_api_key) |
| 104 | raise HTTPException( |
| 105 | status_code=status.HTTP_403_FORBIDDEN, |
| 106 | detail="Forbidden: Invalid API Key.", |
| 107 | ) |
| 108 | |
| 109 | async def validate_instance(self, instance: str) -> bool: |
| 110 | url = f"http://{instance}/v1/models" |
| 111 | try: |
| 112 | async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as client: |
| 113 | logger.info("Verifying %s ...", instance) |
| 114 | async with client.get(url) as response: |
| 115 | if response.status == 200: |
| 116 | data = await response.json() |
| 117 | if "data" in data and len(data["data"]) > 0: |
| 118 | model_cur = data["data"][0].get("id", "") |
| 119 | if model_cur == self.model: |
| 120 | logger.info("Instance: %s could be added.", instance) |
| 121 | return True |
| 122 | else: |
| 123 | logger.warning( |
| 124 | "Mismatch model %s : %s != %s", |
| 125 | instance, |
| 126 | model_cur, |
| 127 | self.model, |
| 128 | ) |
| 129 | return False |
| 130 | else: |
| 131 | return False |
| 132 | else: |
| 133 | return False |
| 134 | except aiohttp.ClientError as e: |
| 135 | logger.error(str(e)) |
| 136 | return False |
| 137 | except Exception as e: |
| 138 | logger.error(str(e)) |
| 139 | return False |
| 140 | |
| 141 | async def add_instance_endpoint(self, request: Request): |
| 142 | try: |
| 143 | data = await request.json() |
| 144 | logger.warning(str(data)) |
| 145 | instance_type = data.get("type") |
| 146 | instance = data.get("instance") |
| 147 | if instance_type not in ["prefill", "decode"]: |
| 148 | raise HTTPException(status_code=400, detail="Invalid instance type.") |
| 149 | if not instance or ":" not in instance: |
| 150 | raise HTTPException(status_code=400, detail="Invalid instance format.") |
| 151 | host, port_str = instance.split(":") |
| 152 | try: |
| 153 | if host != "localhost": |
| 154 | ipaddress.ip_address(host) |
| 155 | port = int(port_str) |
| 156 | if not (0 < port < 65536): |
| 157 | raise HTTPException(status_code=400, detail="Invalid port number.") |
| 158 | except Exception as e: |
| 159 | raise HTTPException( |
| 160 | status_code=400, detail="Invalid instance address." |
| 161 | ) from e |
| 162 | |
| 163 | is_valid = await self.validate_instance(instance) |
| 164 | if not is_valid: |
| 165 | raise HTTPException( |
| 166 | status_code=400, detail="Instance validation failed." |
| 167 | ) |
| 168 | |
| 169 | if instance_type == "prefill": |
| 170 | if instance not in self.prefill_instances: |
| 171 | self.prefill_instances.append(instance) |
| 172 | self.prefill_cycler = itertools.cycle(self.prefill_instances) |
| 173 | else: |
| 174 | raise HTTPException( |
| 175 | status_code=400, detail="Instance already exists." |
| 176 | ) |
| 177 | else: |
| 178 | if instance not in self.decode_instances: |
| 179 | self.decode_instances.append(instance) |
| 180 | self.decode_cycler = itertools.cycle(self.decode_instances) |
| 181 | else: |
| 182 | raise HTTPException( |
| 183 | status_code=400, detail="Instance already exists." |
| 184 | ) |
| 185 | |
| 186 | return JSONResponse( |
| 187 | content={"message": f"Added {instance} to {instance_type}_instances."} |
| 188 | ) |
| 189 | except HTTPException as http_exc: |
| 190 | raise http_exc |
| 191 | except Exception as e: |
| 192 | logger.error("Error in add_instance_endpoint: %s", str(e)) |
| 193 | raise HTTPException(status_code=500, detail=str(e)) from e |
| 194 | |
| 195 | async def forward_request(self, url, data, use_chunked=True): |
| 196 | async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: |
| 197 | headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} |
| 198 | try: |
| 199 | async with session.post( |
| 200 | url=url, json=data, headers=headers |
| 201 | ) as response: |
| 202 | if 200 <= response.status < 300 or 400 <= response.status < 500: |
| 203 | if use_chunked: |
| 204 | async for chunk_bytes in response.content.iter_chunked( |
| 205 | 1024 |
| 206 | ): |
| 207 | yield chunk_bytes |
| 208 | else: |
| 209 | content = await response.read() |
| 210 | yield content |
| 211 | else: |
| 212 | error_content = await response.text() |
| 213 | try: |
| 214 | error_content = json.loads(error_content) |
| 215 | except json.JSONDecodeError: |
| 216 | error_content = error_content |
| 217 | logger.error( |
| 218 | "Request failed with status %s: %s", |
| 219 | response.status, |
| 220 | error_content, |
| 221 | ) |
| 222 | raise HTTPException( |
| 223 | status_code=response.status, |
| 224 | detail=f"Request failed with status {response.status}: " |
| 225 | f"{error_content}", |
| 226 | ) |
| 227 | except aiohttp.ClientError as e: |
| 228 | logger.error("ClientError occurred: %s", str(e)) |
| 229 | raise HTTPException( |
| 230 | status_code=502, |
| 231 | detail="Bad Gateway: Error communicating with upstream server.", |
| 232 | ) from e |
| 233 | except Exception as e: |
| 234 | logger.error("Unexpected error: %s", str(e)) |
| 235 | raise HTTPException(status_code=500, detail=str(e)) from e |
| 236 | |
| 237 | def schedule(self, cycler: itertools.cycle) -> str: |
| 238 | return self.scheduling_policy.schedule(cycler) |
| 239 | |
| 240 | async def get_status(self): |
| 241 | status = { |
| 242 | "prefill_node_count": len(self.prefill_instances), |
| 243 | "decode_node_count": len(self.decode_instances), |
| 244 | "prefill_nodes": self.prefill_instances, |
| 245 | "decode_nodes": self.decode_instances, |
| 246 | } |
| 247 | return status |
| 248 | |
| 249 | async def create_completion(self, raw_request: Request): |
| 250 | try: |
| 251 | request = await raw_request.json() |
| 252 | |
| 253 | kv_prepare_request = request.copy() |
| 254 | kv_prepare_request["max_tokens"] = 1 |
| 255 | |
| 256 | prefill_instance = self.schedule(self.prefill_cycler) |
| 257 | try: |
| 258 | async for _ in self.forward_request( |
| 259 | f"http://{prefill_instance}/v1/completions", kv_prepare_request |
| 260 | ): |
| 261 | continue |
| 262 | except HTTPException as http_exc: |
| 263 | self.remove_instance_endpoint("prefill", prefill_instance) |
| 264 | raise http_exc |
| 265 | |
| 266 | # Perform kv recv and decoding stage |
| 267 | decode_instance = self.schedule(self.decode_cycler) |
| 268 | |
| 269 | try: |
| 270 | generator = self.forward_request( |
| 271 | f"http://{decode_instance}/v1/completions", request |
| 272 | ) |
| 273 | except HTTPException as http_exc: |
| 274 | self.remove_instance_endpoint("decode", decode_instance) |
| 275 | raise http_exc |
| 276 | response = StreamingResponse(generator) |
| 277 | return response |
| 278 | except Exception: |
| 279 | import sys |
| 280 | |
| 281 | exc_info = sys.exc_info() |
| 282 | print("Error occurred in disagg proxy server") |
| 283 | print(exc_info) |
| 284 | |
| 285 | async def create_chat_completion(self, raw_request: Request): |
| 286 | try: |
| 287 | request = await raw_request.json() |
| 288 | |
| 289 | # add params to request |
| 290 | kv_prepare_request = request.copy() |
| 291 | kv_prepare_request["max_tokens"] = 1 |
| 292 | if "max_completion_tokens" in kv_prepare_request: |
| 293 | kv_prepare_request["max_completion_tokens"] = 1 |
| 294 | |
| 295 | # prefill stage |
| 296 | prefill_instance = self.schedule(self.prefill_cycler) |
| 297 | try: |
| 298 | async for _ in self.forward_request( |
| 299 | f"http://{prefill_instance}/v1/chat/completions", kv_prepare_request |
| 300 | ): |
| 301 | continue |
| 302 | except HTTPException as http_exc: |
| 303 | self.remove_instance_endpoint("prefill", prefill_instance) |
| 304 | raise http_exc |
| 305 | # Perform kv recv and decoding stage |
| 306 | decode_instance = self.schedule(self.decode_cycler) |
| 307 | |
| 308 | try: |
| 309 | generator = self.forward_request( |
| 310 | "http://" + decode_instance + "/v1/chat/completions", request |
| 311 | ) |
| 312 | except HTTPException as http_exc: |
| 313 | self.remove_instance_endpoint("decode", decode_instance) |
| 314 | raise http_exc |
| 315 | response = StreamingResponse(content=generator) |
| 316 | return response |
| 317 | except Exception: |
| 318 | exc_info = sys.exc_info() |
| 319 | error_messages = [str(e) for e in exc_info if e] |
| 320 | print("Error occurred in disagg proxy server") |
| 321 | print(error_messages) |
| 322 | return StreamingResponse( |
| 323 | content=iter(error_messages), media_type="text/event-stream" |
| 324 | ) |
| 325 | |
| 326 | def remove_instance_endpoint(self, instance_type, instance): |
| 327 | if instance_type == "decode" and instance in self.decode_instances: |
| 328 | self.decode_instances.remove(instance) |
| 329 | self.decode_cycler = itertools.cycle(self.decode_instances) |
| 330 | if instance_type == "prefill" and instance in self.prefill_instances: |
| 331 | self.prefill_instances.remove(instance) |
| 332 | self.prefill_cycler = itertools.cycle(self.prefill_instances) |
| 333 | |
| 334 | |
| 335 | class RoundRobinSchedulingPolicy(SchedulingPolicy): |
| 336 | def __init__(self): |
| 337 | super().__init__() |
| 338 | |
| 339 | def schedule(self, cycler: itertools.cycle) -> str: |
| 340 | return next(cycler) |
| 341 | |
| 342 | |
| 343 | class ProxyServer: |
| 344 | def __init__( |
| 345 | self, |
| 346 | args: argparse.Namespace, |
| 347 | scheduling_policy: SchedulingPolicy | None = None, |
| 348 | create_completion: Callable[[Request], StreamingResponse] | None = None, |
| 349 | create_chat_completion: Callable[[Request], StreamingResponse] | None = None, |
| 350 | ): |
| 351 | self.validate_parsed_serve_args(args) |
| 352 | self.port = args.port |
| 353 | self.proxy_instance = Proxy( |
| 354 | prefill_instances=[] if args.prefill is None else args.prefill, |
| 355 | decode_instances=[] if args.decode is None else args.decode, |
| 356 | model=args.model, |
| 357 | scheduling_policy=( |
| 358 | scheduling_policy |
| 359 | if scheduling_policy is not None |
| 360 | else RoundRobinSchedulingPolicy() |
| 361 | ), |
| 362 | custom_create_completion=create_completion, |
| 363 | custom_create_chat_completion=create_chat_completion, |
| 364 | ) |
| 365 | |
| 366 | def validate_parsed_serve_args(self, args: argparse.Namespace): |
| 367 | if not args.prefill: |
| 368 | raise ValueError("Please specify at least one prefill node.") |
| 369 | if not args.decode: |
| 370 | raise ValueError("Please specify at least one decode node.") |
| 371 | self.validate_instances(args.prefill) |
| 372 | self.validate_instances(args.decode) |
| 373 | self.verify_model_config(args.prefill, args.model) |
| 374 | self.verify_model_config(args.decode, args.model) |
| 375 | |
| 376 | def validate_instances(self, instances: list): |
| 377 | for instance in instances: |
| 378 | if len(instance.split(":")) != 2: |
| 379 | raise ValueError(f"Invalid instance format: {instance}") |
| 380 | host, port = instance.split(":") |
| 381 | try: |
| 382 | if host != "localhost": |
| 383 | ipaddress.ip_address(host) |
| 384 | port = int(port) |
| 385 | if not (0 < port < 65536): |
| 386 | raise ValueError(f"Invalid port number in instance: {instance}") |
| 387 | except Exception as e: |
| 388 | raise ValueError(f"Invalid instance {instance}: {str(e)}") from e |
| 389 | |
| 390 | def verify_model_config(self, instances: list, model: str) -> None: |
| 391 | model_suffix = model.split("/")[-1] |
| 392 | for instance in instances: |
| 393 | try: |
| 394 | response = requests.get(f"http://{instance}/v1/models") |
| 395 | if response.status_code == 200: |
| 396 | model_cur = response.json()["data"][0]["id"] |
| 397 | model_cur_suffix = model_cur.split("/")[-1] |
| 398 | if model_cur_suffix != model_suffix: |
| 399 | raise ValueError( |
| 400 | f"{instance} serves a different model: " |
| 401 | f"{model_cur} != {model}" |
| 402 | ) |
| 403 | else: |
| 404 | raise ValueError(f"Cannot get model id from {instance}!") |
| 405 | except requests.RequestException as e: |
| 406 | raise ValueError( |
| 407 | f"Error communicating with {instance}: {str(e)}" |
| 408 | ) from e |
| 409 | |
| 410 | def run_server(self): |
| 411 | app = FastAPI() |
| 412 | app.include_router(self.proxy_instance.router) |
| 413 | config = uvicorn.Config(app, port=self.port, loop="uvloop") |
| 414 | server = uvicorn.Server(config) |
| 415 | server.run() |
| 416 | |
| 417 | |
| 418 | def parse_args(): |
| 419 | # Todo: allow more config |
| 420 | parser = argparse.ArgumentParser("vLLM disaggregated proxy server.") |
| 421 | parser.add_argument("--model", "-m", type=str, required=True, help="Model name") |
| 422 | |
| 423 | parser.add_argument( |
| 424 | "--prefill", |
| 425 | "-p", |
| 426 | type=str, |
| 427 | nargs="+", |
| 428 | help="List of prefill node URLs (host:port)", |
| 429 | ) |
| 430 | |
| 431 | parser.add_argument( |
| 432 | "--decode", |
| 433 | "-d", |
| 434 | type=str, |
| 435 | nargs="+", |
| 436 | help="List of decode node URLs (host:port)", |
| 437 | ) |
| 438 | |
| 439 | parser.add_argument( |
| 440 | "--port", |
| 441 | type=int, |
| 442 | default=8000, |
| 443 | help="Server port number", |
| 444 | ) |
| 445 | return parser.parse_args() |
| 446 | |
| 447 | |
| 448 | if __name__ == "__main__": |
| 449 | args = parse_args() |
| 450 | proxy_server = ProxyServer(args=args) |
| 451 | proxy_server.run_server() |