From 41726a18eef3b68cdecaa3c0778b1abd2422b45e Mon Sep 17 00:00:00 2001 From: Xinyi Zhang <64056861+Ann-1024@users.noreply.github.com> Date: Fri, 11 Apr 2025 19:07:31 +0800 Subject: [PATCH] Add disaggrated vllm benchmarks demo (#237) * add disaggregated proxy demo Co-authored-by: Siyu Liu Co-authored-by: Shangming Cai * add xpyd vllm benchmarks demo Co-authored-by: Siyu Liu --------- Co-authored-by: zhangxinyi Co-authored-by: Siyu Liu Co-authored-by: Shangming Cai --- benchmarks/xypd_benchmarks/proxy_demo.py | 447 ++++++++++++++++++ .../vllm-benchmarks/benchmarks.sh | 240 ++++++++++ .../vllm-benchmarks/mooncake.config | 7 + .../vllm-benchmarks/parse_results.py | 56 +++ 4 files changed, 750 insertions(+) create mode 100644 benchmarks/xypd_benchmarks/proxy_demo.py create mode 100644 benchmarks/xypd_benchmarks/vllm-benchmarks/benchmarks.sh create mode 100644 benchmarks/xypd_benchmarks/vllm-benchmarks/mooncake.config create mode 100644 benchmarks/xypd_benchmarks/vllm-benchmarks/parse_results.py diff --git a/benchmarks/xypd_benchmarks/proxy_demo.py b/benchmarks/xypd_benchmarks/proxy_demo.py new file mode 100644 index 00000000..76077930 --- /dev/null +++ b/benchmarks/xypd_benchmarks/proxy_demo.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +This file provides a disaggregated prefilling proxy demo to demonstrate an +example usage of XpYd disaggregated prefilling. +We can launch multiple vllm instances (2 for prefill and 2 for decode), and +launch this proxy demo through: + python3 examples/online_serving/disagg_examples/disagg_proxy_demo.py \ + --model $model_name \ + --prefill localhost:8100 localhost:8101 \ + --decode localhost:8200 localhost:8201 \ + --port 8000 +""" +import argparse +import ipaddress +import itertools +import json +import logging +import os +import sys +from abc import ABC, abstractmethod +from typing import Callable, Optional + +import aiohttp +import requests +import uvicorn +from fastapi import (APIRouter, Depends, FastAPI, Header, HTTPException, + Request, status) +from fastapi.responses import JSONResponse, StreamingResponse + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) +logger = logging.getLogger() +logging.basicConfig(level=logging.INFO) + + +class SchedulingPolicy(ABC): + + @abstractmethod + def schedule(self, cycler: itertools.cycle): + raise NotImplementedError("Scheduling Proxy is not set.") + + +class Proxy: + + def __init__( + self, + prefill_instances: list[str], + decode_instances: list[str], + model: str, + scheduling_policy: SchedulingPolicy, + custom_create_completion: Optional[Callable[[Request], + StreamingResponse]] = None, + custom_create_chat_completion: Optional[Callable[ + [Request], StreamingResponse]] = None, + ): + self.prefill_instances = prefill_instances + self.decode_instances = decode_instances + self.prefill_cycler = itertools.cycle(prefill_instances) + self.decode_cycler = itertools.cycle(decode_instances) + self.model = model + self.scheduling_policy = scheduling_policy + self.custom_create_completion = custom_create_completion + self.custom_create_chat_completion = custom_create_chat_completion + self.router = APIRouter() + self.setup_routes() + + def setup_routes(self): + self.router.post( + "/v1/completions", + dependencies=[ + Depends(self.validate_json_request) + ])(self.custom_create_completion if self. + custom_create_completion else self.create_completion) + self.router.post( + "/v1/chat/completions", + dependencies=[ + Depends(self.validate_json_request) + ])(self.custom_create_chat_completion if self. + custom_create_chat_completion else self.create_chat_completion) + self.router.get("/status", + response_class=JSONResponse)(self.get_status) + self.router.post("/instances/add", + dependencies=[Depends(self.api_key_authenticate) + ])(self.add_instance_endpoint) + + async def validate_json_request(self, raw_request: Request): + content_type = raw_request.headers.get("content-type", "").lower() + if content_type != "application/json": + raise HTTPException( + status_code=415, + detail= + "Unsupported Media Type: Only 'application/json' is allowed", + ) + + def api_key_authenticate(self, x_api_key: str = Header(...)): + expected_api_key = os.environ.get("ADMIN_API_KEY") + if not expected_api_key: + logger.error("ADMIN_API_KEY is not set in the environment.") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Server configuration error.", + ) + if x_api_key != expected_api_key: + logger.warning("Unauthorized access attempt with API Key: %s", + x_api_key) + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Forbidden: Invalid API Key.", + ) + + async def validate_instance(self, instance: str) -> bool: + url = f"http://{instance}/v1/models" + try: + async with aiohttp.ClientSession( + timeout=AIOHTTP_TIMEOUT) as client: + logger.info("Verifying %s ...", instance) + async with client.get(url) as response: + if response.status == 200: + data = await response.json() + if "data" in data and len(data["data"]) > 0: + model_cur = data["data"][0].get("id", "") + if model_cur == self.model: + logger.info("Instance: %s could be added.", + instance) + return True + else: + logger.warning("Mismatch model %s : %s != %s", + instance, model_cur, self.model) + return False + else: + return False + else: + return False + except aiohttp.ClientError as e: + logger.error(str(e)) + return False + except Exception as e: + logger.error(str(e)) + return False + + async def add_instance_endpoint(self, request: Request): + try: + data = await request.json() + logger.warning(str(data)) + instance_type = data.get("type") + instance = data.get("instance") + if instance_type not in ["prefill", "decode"]: + raise HTTPException(status_code=400, + detail="Invalid instance type.") + if not instance or ":" not in instance: + raise HTTPException(status_code=400, + detail="Invalid instance format.") + host, port_str = instance.split(":") + try: + if host != "localhost": + ipaddress.ip_address(host) + port = int(port_str) + if not (0 < port < 65536): + raise HTTPException(status_code=400, + detail="Invalid port number.") + except Exception as e: + raise HTTPException(status_code=400, + detail="Invalid instance address.") from e + + is_valid = await self.validate_instance(instance) + if not is_valid: + raise HTTPException(status_code=400, + detail="Instance validation failed.") + + if instance_type == "prefill": + if instance not in self.prefill_instances: + self.prefill_instances.append(instance) + self.prefill_cycler = itertools.cycle( + self.prefill_instances) + else: + raise HTTPException(status_code=400, + detail="Instance already exists.") + else: + if instance not in self.decode_instances: + self.decode_instances.append(instance) + self.decode_cycler = itertools.cycle(self.decode_instances) + else: + raise HTTPException(status_code=400, + detail="Instance already exists.") + + return JSONResponse(content={ + "message": + f"Added {instance} to {instance_type}_instances." + }) + except HTTPException as http_exc: + raise http_exc + except Exception as e: + logger.error("Error in add_instance_endpoint: %s", str(e)) + raise HTTPException(status_code=500, detail=str(e)) from e + + async def forward_request(self, url, data, use_chunked=True): + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}" + } + try: + async with session.post(url=url, json=data, + headers=headers) as response: + if 200 <= response.status < 300 or 400 <= response.status < 500: # noqa: E501 + if use_chunked: + async for chunk_bytes in response.content.iter_chunked( # noqa: E501 + 1024): + yield chunk_bytes + else: + content = await response.read() + yield content + else: + error_content = await response.text() + try: + error_content = json.loads(error_content) + except json.JSONDecodeError: + error_content = error_content + logger.error("Request failed with status %s: %s", + response.status, error_content) + raise HTTPException( + status_code=response.status, + detail= + f"Request failed with status {response.status}: " + f"{error_content}", + ) + except aiohttp.ClientError as e: + logger.error("ClientError occurred: %s", str(e)) + raise HTTPException( + status_code=502, + detail= + "Bad Gateway: Error communicating with upstream server.", + ) from e + except Exception as e: + logger.error("Unexpected error: %s", str(e)) + raise HTTPException(status_code=500, detail=str(e)) from e + + def schedule(self, cycler: itertools.cycle) -> str: + return self.scheduling_policy.schedule(cycler) + + async def get_status(self): + status = { + "prefill_node_count": len(self.prefill_instances), + "decode_node_count": len(self.decode_instances), + "prefill_nodes": self.prefill_instances, + "decode_nodes": self.decode_instances, + } + return status + + async def create_completion(self, raw_request: Request): + try: + request = await raw_request.json() + + kv_prepare_request = request.copy() + kv_prepare_request["max_tokens"] = 1 + + prefill_instance = self.schedule(self.prefill_cycler) + try: + async for _ in self.forward_request( + f"http://{prefill_instance}/v1/completions", + kv_prepare_request): + continue + except HTTPException as http_exc: + self.remove_instance_endpoint("prefill", prefill_instance) + raise http_exc + + # Perform kv recv and decoding stage + decode_instance = self.schedule(self.decode_cycler) + + try: + generator = self.forward_request( + f"http://{decode_instance}/v1/completions", request) + except HTTPException as http_exc: + self.remove_instance_endpoint("decode", decode_instance) + raise http_exc + response = StreamingResponse(generator) + return response + except Exception: + import sys + + exc_info = sys.exc_info() + print("Error occurred in disagg proxy server") + print(exc_info) + + async def create_chat_completion(self, raw_request: Request): + try: + request = await raw_request.json() + + # add params to request + kv_prepare_request = request.copy() + kv_prepare_request["max_tokens"] = 1 + + # prefill stage + prefill_instance = self.schedule(self.prefill_cycler) + try: + async for _ in self.forward_request( + f"http://{prefill_instance}/v1/chat/completions", + kv_prepare_request): + continue + except HTTPException as http_exc: + self.remove_instance_endpoint("prefill", prefill_instance) + raise http_exc + # Perform kv recv and decoding stage + decode_instance = self.schedule(self.decode_cycler) + + try: + generator = self.forward_request( + "http://" + decode_instance + "/v1/chat/completions", + request) + except HTTPException as http_exc: + self.remove_instance_endpoint("decode", decode_instance) + raise http_exc + response = StreamingResponse(content=generator) + return response + except Exception: + exc_info = sys.exc_info() + error_messages = [str(e) for e in exc_info if e] + print("Error occurred in disagg proxy server") + print(error_messages) + return StreamingResponse(content=iter(error_messages), + media_type="text/event-stream") + + def remove_instance_endpoint(self, instance_type, instance): + if (instance_type == "decode" and instance in self.decode_instances): + self.decode_instances.remove(instance) + self.decode_cycler = itertools.cycle(self.decode_instances) + if (instance_type == "prefill" and instance in self.decode_instances): + self.prefill_instances.remove(instance) + self.prefill_cycler = itertools.cycle(self.decode_instances) + + +class RoundRobinSchedulingPolicy(SchedulingPolicy): + + def __init__(self): + super().__init__() + + def schedule(self, cycler: itertools.cycle) -> str: + return next(cycler) + + +class ProxyServer: + + def __init__( + self, + args: argparse.Namespace, + scheduling_policy: Optional[SchedulingPolicy] = None, + create_completion: Optional[Callable[[Request], + StreamingResponse]] = None, + create_chat_completion: Optional[Callable[[Request], + StreamingResponse]] = None, + ): + self.validate_parsed_serve_args(args) + self.port = args.port + self.proxy_instance = Proxy( + prefill_instances=[] if args.prefill is None else args.prefill, + decode_instances=[] if args.decode is None else args.decode, + model=args.model, + scheduling_policy=(scheduling_policy if scheduling_policy + is not None else RoundRobinSchedulingPolicy()), + custom_create_completion=create_completion, + custom_create_chat_completion=create_chat_completion, + ) + + def validate_parsed_serve_args(self, args: argparse.Namespace): + if not args.prefill: + raise ValueError("Please specify at least one prefill node.") + if not args.decode: + raise ValueError("Please specify at least one decode node.") + self.validate_instances(args.prefill) + self.validate_instances(args.decode) + self.verify_model_config(args.prefill, args.model) + self.verify_model_config(args.decode, args.model) + + def validate_instances(self, instances: list): + for instance in instances: + if len(instance.split(":")) != 2: + raise ValueError(f"Invalid instance format: {instance}") + host, port = instance.split(":") + try: + if host != "localhost": + ipaddress.ip_address(host) + port = int(port) + if not (0 < port < 65536): + raise ValueError( + f"Invalid port number in instance: {instance}") + except Exception as e: + raise ValueError( + f"Invalid instance {instance}: {str(e)}") from e + + def verify_model_config(self, instances: list, model: str) -> None: + model_suffix = model.split("/")[-1] + for instance in instances: + try: + response = requests.get(f"http://{instance}/v1/models") + if response.status_code == 200: + model_cur = response.json()["data"][0]["id"] + model_cur_suffix = model_cur.split("/")[-1] + if model_cur_suffix != model_suffix: + raise ValueError( + f"{instance} serves a different model: " + f"{model_cur} != {model}") + else: + raise ValueError(f"Cannot get model id from {instance}!") + except requests.RequestException as e: + raise ValueError( + f"Error communicating with {instance}: {str(e)}") from e + + def run_server(self): + app = FastAPI() + app.include_router(self.proxy_instance.router) + config = uvicorn.Config(app, port=self.port, loop="uvloop") + server = uvicorn.Server(config) + server.run() + + +if __name__ == "__main__": + # Todo: allow more config + parser = argparse.ArgumentParser("vLLM disaggregated proxy server.") + parser.add_argument("--model", + "-m", + type=str, + required=True, + help="Model name") + + parser.add_argument( + "--prefill", + "-p", + type=str, + nargs="+", + help="List of prefill node URLs (host:port)", + ) + + parser.add_argument( + "--decode", + "-d", + type=str, + nargs="+", + help="List of decode node URLs (host:port)", + ) + + parser.add_argument( + "--port", + type=int, + default=8000, + help="Server port number", + ) + args = parser.parse_args() + proxy_server = ProxyServer(args=args) + proxy_server.run_server() \ No newline at end of file diff --git a/benchmarks/xypd_benchmarks/vllm-benchmarks/benchmarks.sh b/benchmarks/xypd_benchmarks/vllm-benchmarks/benchmarks.sh new file mode 100644 index 00000000..424799c4 --- /dev/null +++ b/benchmarks/xypd_benchmarks/vllm-benchmarks/benchmarks.sh @@ -0,0 +1,240 @@ +# the file provides is demo to test mooncake connector performance when work with vllm + +set -ex + +VLLM_SRC_PATH=${VLLM_SRC_PATH:-"vllm-src"} +MOONCAKE_CONFIG_PATH=${MOONCAKE_CONFIG_PATH:-"mooncake.json"} +MODEL=${MODEL:-"Qwen/Qwen2.5-7B-Instruct"} +DEMO_PATH=${DEMO_PATH:-"../proxy_demo.py"} +NUM_PREFILL=${NUM_PREFILL:-"4"} +NUM_DECODE=${NUM_DECODE:-"4"} +PREFIX_LEN=${PREFIX_LEN:-"50"} +NUM_FOLDS=${NUM_FOLDS:-"20"} +MASTER_PORT=${MASTER_PORT:-"10001"} +PREFILL_PORT_BASE=${PREFILL_PORT_BASE:-"8100"} +DECODE_PORT_BASE=${DECODE_PORT_BASE:-"8200"} +PROXY_PORT=${PROXY_PORT:-"8000"} +logs_root=${LOG_ROOT:-"logs"} +results_root=${RESULT_ROOT:-"results"} + +PROXY_ID=0 +CUDA_VISIBLE_ID=0 + +INPUT_LENS=(1024 4096) +OUTPUT_LENS=(6 256) + +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export VLLM_USE_V1=0 + +wait_for_server() { + # wait for vllm server to start + # return 1 if vllm server crashes + local port=$1 + timeout 1200 bash -c " + until curl -X POST -s http://localhost:${port}/v1/models > /dev/null; do + sleep 1 + done" && return 0 || return 1 +} + +get_related_pids() +{ + local pid=${1} + [ -z "$pid" ] && echo "" + ps -ef | grep "$pid" | grep -v 'grep' | awk -F ' ' '{print $2}' | tr '\n' ' ' +} + +destroy_vllm_engine() +{ + local port=$1 + local main_pid=$(ps -ef | grep 'vllm.entrypoints.openai.api_server' | grep "port=${port}" | awk -F ' ' '{print $2}') + if [ -n "${main_pid}" ]; then + local related_pids=$(get_related_pids "${main_pid}" | sed 's/^[ \t]*//;s/[ \t]*$//') + for pid in ${related_pids} + do + related_pids="${related_pids} $(get_related_pids $pid)" + done + if [ -n "$(echo "${related_pids}" | sed 's/^[ \t]*//;s/[ \t]*$//')" ];then + kill -9 ${related_pids} + fi + fi + sleep 5 +} + +kill_nodes() { + # kill all processes by port + lsof -t -i:$(PROXY_PORT) | xargs -r kill -9 + for ((i=0; i/dev/null; then + echo "Process with PID $pid has been terminated or does not exist." + break + else + echo "Sent termination signal to process with PID $pid, checking..." + sleep 1 + fi + done +} + +launch_nodes() { + nohup mooncake_master --port ${MASTER_PORT} > ${logs_root}/master.txt 2>&1 & + # launch prefill instance + for ((i=0; i ${logs_root}/prefill-${i}.txt 2>&1 & + echo "Launched node on port $PORT" + CUDA_VISIBLE_ID=$((CUDA_VISIBLE_ID + 1)) + done + # launch decode instance + for ((i=0; i ${logs_root}/decode-${i}.txt 2>&1 & + echo "Launched node on port $PORT" + CUDA_VISIBLE_ID=$((CUDA_VISIBLE_ID + 1)) + done + for ((i=0; i " + return 1 + fi + if [ "$PROXY_ID" -ne 0 ]; then + kill_process_by_pid $PROXY_ID + fi + num_prefill=$1 + num_decode=$2 + prefill_ports=() + for (( i=0; i&1 | tee ${logs_root}/proxy-${num_prefill}-${num_decode}.txt 2>&1 & + PROXY_ID=$! + echo "Launched disagg_proxy with PID: $PROXY_ID" + sleep 1 +} + +benchmark() { + dataset_name="random" + prefix_len=50 + file_prefix=$1 + shift + for max_concurrency in "$@"; do + num_prompts=$(( max_concurrency * NUM_FOLDS )) + for input_len in ${INPUT_LENS[@]}; do + for output_len in ${OUTPUT_LENS[@]}; do + input_len_name=$(printf %04d $input_len) + output_len_name=$(printf %04d $output_len) + max_concurrency_name=$(printf %03d $max_concurrency) + python3 $VLLM_SRC_PATH/benchmarks/benchmark_serving.py \ + --backend vllm \ + --model ${MODEL} \ + --dataset-name random \ + --random-input-len $input_len \ + --random-output-len $output_len \ + --random-prefix-len ${PREFIX_LEN} \ + --num-prompts $num_prompts \ + --max-concurrency=${max_concurrency} \ + --trust-remote-code \ + --ignore_eos \ + --port ${PROXY_PORT} \ + --save-result \ + --percentile-metrics="ttft,tpot,itl,e2el" \ + --result-dir=${results_root} \ + --result-filename=${file_prefix}-input-${input_len_name}-output-${output_len_name}-concurrency-${max_concurrency_name}-serving.json \ + 2>&1 | tee ${logs_root}/${file_prefix}-${input_len_name}-output-${output_len_name}-concurrency-${max_concurrency_name}-serving.txt + sleep 2 + done + done + done +} + +prepare_env(){ + (which wget && which curl) || (apt-get update && apt-get install -y wget curl) + (which git) || (apt-get -y install git) + (which socat) || (apt-get -y install socat) + pip install vllm + pip install quart httpx matplotlib aiohttp pandas datasets + if ! [ -d $VLLM_SRC_PATH ]; then + git clone https://github.com/vllm-project/vllm.git $VLLM_SRC_PATH + fi +} + +main() { + prepare_env + results_root=${results_root}-$(date "+%Y%m%d-%H:%M:%S") + logs_root=${logs_root}-$(date "+%Y%m%d-%H:%M:%S") + mkdir -p $results_root + mkdir -p $logs_root + echo "Results will be saved to $results_root" + echo "Logs will be saved to $logs_root" + + export VLLM_HOST_IP=$(hostname -I | awk '{print $1}') + kill_nodes + ## launch instances. + launch_nodes + + launch_disagg_proxy 1 1 + benchmark proxy-1-1 1 4 8 16 + + launch_disagg_proxy 2 1 + benchmark proxy-2-1 4 8 16 + + launch_disagg_proxy 2 2 + benchmark proxy-2-2 4 8 16 + + launch_disagg_proxy 2 4 + benchmark proxy-2-4 4 8 16 + + launch_disagg_proxy 4 4 + benchmark proxy-4-4 4 8 16 + + kill_nodes + + python3 parse_results.py $results_root $results_root/result.xlsx +} + +main "$@" \ No newline at end of file diff --git a/benchmarks/xypd_benchmarks/vllm-benchmarks/mooncake.config b/benchmarks/xypd_benchmarks/vllm-benchmarks/mooncake.config new file mode 100644 index 00000000..682bb3bb --- /dev/null +++ b/benchmarks/xypd_benchmarks/vllm-benchmarks/mooncake.config @@ -0,0 +1,7 @@ +{ + "local_hostname": "192.168.0.137", + "metadata_server": "etcd://192.168.0.137:2379", + "protocol": "rdma", + "device_name": "erdma_0", + "master_server_address": "192.168.0.137:10001" +} \ No newline at end of file diff --git a/benchmarks/xypd_benchmarks/vllm-benchmarks/parse_results.py b/benchmarks/xypd_benchmarks/vllm-benchmarks/parse_results.py new file mode 100644 index 00000000..5d268393 --- /dev/null +++ b/benchmarks/xypd_benchmarks/vllm-benchmarks/parse_results.py @@ -0,0 +1,56 @@ +import sys +import json +import os +import openpyxl + +global metrics +metrics = ['request_throughput', 'output_throughput', 'total_token_throughput',\ + 'mean_ttft_ms', 'median_ttft_ms', 'std_ttft_ms', 'p99_ttft_ms',\ + 'mean_tpot_ms', 'median_tpot_ms', 'std_tpot_ms', 'p99_tpot_ms',\ + 'mean_itl_ms', 'median_itl_ms', 'std_itl_ms', 'p99_itl_ms', \ + 'mean_e2el_ms', 'median_e2el_ms', 'std_e2el_ms', 'p99_e2el_ms' ] + +def parse_serving_throughput(path: str): + values=[] + with open(path, 'r') as f: + result = json.load(f) + for metric in metrics: + value = result[metric] + values.append(value) + return values + + +if __name__ == '__main__': + if len(sys.argv) != 3: + print("Usage: {} ".format(sys.argv[0])) + sys.exit(1) + result_path = sys.argv[1] + parsed_result_path = sys.argv[2] + col=1 + row=1 + workbook = openpyxl.Workbook() + sheet = workbook.active + config_names=['num_pserver','num_dserver','input_len','output_len','max_concurrency'] + for con in config_names: + sheet.cell(row,col,con) + row += 1 + for metric in metrics: + sheet.cell(row,col,metric) + row += 1 + files = os.listdir(result_path) + files.sort() + for file in files: + if file.endswith("json"): + configs=file.split('-') + col += 1 + sheet.cell(1,col,configs[1]) + sheet.cell(2,col,configs[2]) + sheet.cell(3,col,configs[4]) + sheet.cell(4,col,configs[6]) + sheet.cell(5,col,configs[8]) + results=parse_serving_throughput(os.path.join(result_path,file)) + row=5 + for result in results: + row+=1 + sheet.cell(row,col,result) + workbook.save(parsed_result_path) \ No newline at end of file