| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import os |
| 6 | import time |
| 7 | from contextlib import asynccontextmanager |
| 8 | |
| 9 | import httpx |
| 10 | import numpy as np |
| 11 | from fastapi import FastAPI, Request |
| 12 | from fastapi.responses import StreamingResponse |
| 13 | |
| 14 | |
| 15 | @asynccontextmanager |
| 16 | async def lifespan(app: FastAPI): |
| 17 | """Lifespan context manager to handle startup and shutdown events.""" |
| 18 | # Startup: Initialize clients |
| 19 | prefiller_base_url = ( |
| 20 | f"http://{global_args.prefiller_host}:{global_args.prefiller_port}/v1" |
| 21 | ) |
| 22 | decoder_base_url = ( |
| 23 | f"http://{global_args.decoder_host}:{global_args.decoder_port}/v1" |
| 24 | ) |
| 25 | |
| 26 | app.state.prefill_client = httpx.AsyncClient( |
| 27 | timeout=None, |
| 28 | base_url=prefiller_base_url, |
| 29 | limits=httpx.Limits( |
| 30 | max_connections=None, |
| 31 | max_keepalive_connections=None, |
| 32 | ), |
| 33 | ) |
| 34 | app.state.decode_client = httpx.AsyncClient( |
| 35 | timeout=None, |
| 36 | base_url=decoder_base_url, |
| 37 | limits=httpx.Limits( |
| 38 | max_connections=None, |
| 39 | max_keepalive_connections=None, |
| 40 | ), |
| 41 | ) |
| 42 | |
| 43 | yield |
| 44 | |
| 45 | # Shutdown: Close clients |
| 46 | await app.state.prefill_client.aclose() |
| 47 | await app.state.decode_client.aclose() |
| 48 | |
| 49 | |
| 50 | # Update FastAPI app initialization to use lifespan |
| 51 | app = FastAPI(lifespan=lifespan) |
| 52 | |
| 53 | |
| 54 | class StatsCalculator: |
| 55 | def __init__(self): |
| 56 | self._stats = [] |
| 57 | self._last_log_time = time.time() |
| 58 | |
| 59 | def add(self, value): |
| 60 | self._stats.append(value) |
| 61 | if time.time() - self._last_log_time > 5: |
| 62 | self._log_stats() |
| 63 | self._last_log_time = time.time() |
| 64 | |
| 65 | def _log_stats(self): |
| 66 | # Print average, median, and 99th percentile |
| 67 | np_arr = np.array(self._stats) |
| 68 | output_str = ( |
| 69 | f"\nNum requests: {len(self._stats)}" |
| 70 | "\nPrefill node TTFT stats:" |
| 71 | f"\n - Average (ms): {np.mean(np_arr)}" |
| 72 | f"\n - Median (ms): {np.median(np_arr)}" |
| 73 | f"\n - 99th Percentile (ms): {np.percentile(np_arr, 99)}\n" |
| 74 | ) |
| 75 | print( |
| 76 | "===============================", |
| 77 | output_str, |
| 78 | "===============================", |
| 79 | ) |
| 80 | |
| 81 | |
| 82 | stats_calculator = StatsCalculator() |
| 83 | counter = 0 |
| 84 | |
| 85 | |
| 86 | def parse_args(): |
| 87 | parser = argparse.ArgumentParser() |
| 88 | |
| 89 | parser.add_argument("--port", type=int, default=8000) |
| 90 | parser.add_argument("--host", type=str, default="localhost") |
| 91 | parser.add_argument("--prefiller-host", type=str, default="localhost") |
| 92 | parser.add_argument("--prefiller-port", type=int, default=8100) |
| 93 | parser.add_argument("--decoder-host", type=str, default="localhost") |
| 94 | parser.add_argument("--decoder-port", type=int, default=8200) |
| 95 | args = parser.parse_args() |
| 96 | return args |
| 97 | |
| 98 | |
| 99 | # Initialize variables to hold the persistent clients |
| 100 | app.state.prefill_client = None |
| 101 | app.state.decode_client = None |
| 102 | |
| 103 | |
| 104 | async def send_request_to_service( |
| 105 | client: httpx.AsyncClient, endpoint: str, req_data: dict |
| 106 | ): |
| 107 | """Send a request to a service using a persistent client.""" |
| 108 | req_data = req_data.copy() |
| 109 | req_data["max_tokens"] = 1 |
| 110 | if "max_completion_tokens" in req_data: |
| 111 | req_data["max_completion_tokens"] = 1 |
| 112 | |
| 113 | headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} |
| 114 | response = await client.post(endpoint, json=req_data, headers=headers) |
| 115 | response.raise_for_status() |
| 116 | |
| 117 | # read/consume the response body to release the connection |
| 118 | # otherwise, it would http.ReadError |
| 119 | await response.aread() |
| 120 | |
| 121 | return response |
| 122 | |
| 123 | |
| 124 | async def stream_service_response( |
| 125 | client: httpx.AsyncClient, endpoint: str, req_data: dict |
| 126 | ): |
| 127 | """Asynchronously stream the response from a service using a persistent client.""" |
| 128 | headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} |
| 129 | async with client.stream( |
| 130 | "POST", endpoint, json=req_data, headers=headers |
| 131 | ) as response: |
| 132 | response.raise_for_status() |
| 133 | async for chunk in response.aiter_bytes(): |
| 134 | yield chunk |
| 135 | |
| 136 | |
| 137 | @app.post("/v1/completions") |
| 138 | async def handle_completions(request: Request): |
| 139 | global counter, stats_calculator |
| 140 | counter += 1 |
| 141 | |
| 142 | st = time.time() |
| 143 | try: |
| 144 | req_data = await request.json() |
| 145 | |
| 146 | # Send request to prefill service, ignore the response |
| 147 | await send_request_to_service( |
| 148 | app.state.prefill_client, "/completions", req_data |
| 149 | ) |
| 150 | |
| 151 | et = time.time() |
| 152 | stats_calculator.add(et - st) |
| 153 | |
| 154 | # Stream response from decode service |
| 155 | async def generate_stream(): |
| 156 | async for chunk in stream_service_response( |
| 157 | app.state.decode_client, "/completions", req_data |
| 158 | ): |
| 159 | yield chunk |
| 160 | |
| 161 | return StreamingResponse(generate_stream(), media_type="text/event-stream") |
| 162 | |
| 163 | except Exception as e: |
| 164 | import sys |
| 165 | import traceback |
| 166 | |
| 167 | exc_info = sys.exc_info() |
| 168 | print("Error occurred in disagg prefill proxy server - completions endpoint") |
| 169 | print(e) |
| 170 | print("".join(traceback.format_exception(*exc_info))) |
| 171 | raise |
| 172 | |
| 173 | |
| 174 | @app.post("/v1/chat/completions") |
| 175 | async def handle_chat_completions(request: Request): |
| 176 | global counter, stats_calculator |
| 177 | counter += 1 |
| 178 | |
| 179 | st = time.time() |
| 180 | try: |
| 181 | req_data = await request.json() |
| 182 | |
| 183 | # Send request to prefill service, ignore the response |
| 184 | await send_request_to_service( |
| 185 | app.state.prefill_client, "/chat/completions", req_data |
| 186 | ) |
| 187 | |
| 188 | et = time.time() |
| 189 | stats_calculator.add(et - st) |
| 190 | |
| 191 | # Stream response from decode service |
| 192 | async def generate_stream(): |
| 193 | async for chunk in stream_service_response( |
| 194 | app.state.decode_client, "/chat/completions", req_data |
| 195 | ): |
| 196 | yield chunk |
| 197 | |
| 198 | return StreamingResponse(generate_stream(), media_type="text/event-stream") |
| 199 | |
| 200 | except Exception as e: |
| 201 | import sys |
| 202 | import traceback |
| 203 | |
| 204 | exc_info = sys.exc_info() |
| 205 | print( |
| 206 | "Error occurred in disagg prefill proxy server - chat completions endpoint" |
| 207 | ) |
| 208 | print(e) |
| 209 | print("".join(traceback.format_exception(*exc_info))) |
| 210 | raise |
| 211 | |
| 212 | |
| 213 | if __name__ == "__main__": |
| 214 | global global_args |
| 215 | global_args = parse_args() |
| 216 | |
| 217 | import uvicorn |
| 218 | |
| 219 | uvicorn.run(app, host=global_args.host, port=global_args.port) |