wehub-resource-sync
462 lines
14 KiB
Python
462 lines
14 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
import argparse
|
|
import asyncio
|
|
import itertools
|
|
import os
|
|
import threading
|
|
import time
|
|
import uuid
|
|
|
|
# Third Party
|
|
from fastapi import FastAPI, Request
|
|
from fastapi.responses import JSONResponse, StreamingResponse
|
|
import httpx
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
# Global dictionary to store asyncio.Events indexed by request ID.
|
|
# This is shared between the main proxy app and the telemetry app.
|
|
pending_requests: dict[str, asyncio.Event] = {}
|
|
pending_requests_lock = threading.Lock()
|
|
|
|
# Reference to the main event loop (set at startup)
|
|
main_event_loop: Optional[asyncio.AbstractEventLoop] = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""
|
|
Lifespan context manager to handle startup and shutdown events.
|
|
"""
|
|
global main_event_loop
|
|
main_event_loop = asyncio.get_running_loop()
|
|
|
|
# Startup: Initialize clients
|
|
pref_hosts = global_args.prefiller_host
|
|
pref_ports = global_args.prefiller_port
|
|
|
|
def pair_hosts_and_ports(hosts, ports, count=None):
|
|
"""
|
|
Flexible host-port pairing with expansion strategies.
|
|
"""
|
|
if not isinstance(hosts, list):
|
|
hosts = [hosts]
|
|
if not isinstance(ports, list):
|
|
ports = [ports]
|
|
if len(hosts) == 1 and len(ports) == 1:
|
|
if count is None or count <= 1:
|
|
return [(hosts[0], ports[0])]
|
|
else:
|
|
return [(hosts[0], ports[0] + i) for i in range(count)]
|
|
if len(hosts) == 1:
|
|
return [(hosts[0], p) for p in ports]
|
|
if len(ports) == 1:
|
|
return [(h, ports[0]) for h in hosts]
|
|
if len(hosts) != len(ports):
|
|
raise ValueError(
|
|
"Length mismatch between hosts and ports lists for pairing"
|
|
)
|
|
return list(zip(hosts, ports, strict=False))
|
|
|
|
prefill_pairs = pair_hosts_and_ports(
|
|
pref_hosts, pref_ports, global_args.num_prefillers
|
|
)
|
|
for host, port in prefill_pairs:
|
|
prefiller_base_url = f"http://{host}:{int(port)}"
|
|
prefill_client = httpx.AsyncClient(timeout=None, base_url=prefiller_base_url)
|
|
app.state.prefill_clients.append(ClientInfo(prefill_client))
|
|
|
|
# Build decoder clients
|
|
dec_hosts = global_args.decoder_host
|
|
dec_ports = global_args.decoder_port
|
|
|
|
decoder_pairs = pair_hosts_and_ports(dec_hosts, dec_ports, global_args.num_decoders)
|
|
|
|
for host, port in decoder_pairs:
|
|
decoder_base_url = f"http://{host}:{int(port)}"
|
|
decode_client = httpx.AsyncClient(timeout=None, base_url=decoder_base_url)
|
|
app.state.decode_clients.append(ClientInfo(decode_client))
|
|
|
|
app.state.total_clients = app.state.prefill_clients + app.state.decode_clients
|
|
|
|
yield
|
|
|
|
# Shutdown: Close clients
|
|
for client in app.state.prefill_clients:
|
|
await client.aclose()
|
|
for client in app.state.decode_clients:
|
|
await client.aclose()
|
|
|
|
|
|
# Main proxy FastAPI app
|
|
app = FastAPI(lifespan=lifespan)
|
|
|
|
|
|
def csv_ints(s):
|
|
return [int(x) for x in s.split(",")]
|
|
|
|
|
|
def csv_strs(s):
|
|
return [x.strip() for x in s.split(",")]
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("--port", type=int, default=8000)
|
|
parser.add_argument("--host", type=str, default="localhost")
|
|
parser.add_argument("--prefiller-host", type=csv_strs, default=["localhost"])
|
|
parser.add_argument("--prefiller-port", type=csv_ints, default=[8100])
|
|
parser.add_argument("--num-prefillers", type=int, default=1)
|
|
parser.add_argument("--decoder-host", type=csv_strs, default=["localhost"])
|
|
parser.add_argument("--decoder-port", type=csv_ints, default=[8200])
|
|
parser.add_argument("--num-decoders", type=int, default=1)
|
|
parser.add_argument("--telemetry-port", type=int, default=5768)
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
@dataclass
|
|
class ClientInfo:
|
|
client: httpx.AsyncClient
|
|
host: Optional[str] = None
|
|
|
|
async def aclose(self):
|
|
await self.client.aclose()
|
|
|
|
|
|
# Initialize state
|
|
app.state.prefill_clients = []
|
|
app.state.decode_clients = []
|
|
app.state.total_clients = []
|
|
|
|
|
|
def round_robin_pick_client(clients, idx):
|
|
return clients[idx % len(clients)]
|
|
|
|
|
|
round_robin_counter = itertools.count()
|
|
|
|
|
|
def round_robin_pick_clients() -> tuple[ClientInfo, ClientInfo]:
|
|
idx = next(round_robin_counter)
|
|
prefill_client = round_robin_pick_client(app.state.prefill_clients, idx)
|
|
decode_client = round_robin_pick_client(app.state.decode_clients, idx)
|
|
return prefill_client, decode_client
|
|
|
|
|
|
def _build_headers(**extra: str) -> dict[str, str]:
|
|
"""Build common HTTP headers, including auth if OPENAI_API_KEY is set."""
|
|
headers: dict[str, str] = {"Content-Type": "application/json"}
|
|
api_key = os.environ.get("OPENAI_API_KEY", "").strip()
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
headers.update(extra)
|
|
return headers
|
|
|
|
|
|
async def send_request_to_prefiller(
|
|
client: httpx.AsyncClient,
|
|
endpoint: str,
|
|
req_data: dict,
|
|
request_id: str,
|
|
):
|
|
"""
|
|
Send a request to prefiller with request-id header.
|
|
"""
|
|
headers = _build_headers(**{"X-Request-Id": request_id})
|
|
response = await client.post(endpoint, json=req_data, headers=headers)
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
|
|
async def send_request_to_decoder(
|
|
client: httpx.AsyncClient,
|
|
endpoint: str,
|
|
req_data: dict,
|
|
):
|
|
"""
|
|
Send a request to decoder service.
|
|
"""
|
|
headers = _build_headers()
|
|
response = await client.post(endpoint, json=req_data, headers=headers)
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
|
|
async def stream_service_response(
|
|
client: httpx.AsyncClient, endpoint: str, req_data: dict
|
|
):
|
|
"""
|
|
Asynchronously stream the response from a service.
|
|
"""
|
|
headers = _build_headers()
|
|
async with client.stream(
|
|
"POST", endpoint, json=req_data, headers=headers
|
|
) as response:
|
|
response.raise_for_status()
|
|
async for chunk in response.aiter_bytes():
|
|
yield chunk
|
|
|
|
|
|
def generate_request_id() -> str:
|
|
"""Generate a unique request ID."""
|
|
return str(uuid.uuid4())[:16]
|
|
|
|
|
|
def create_pending_request(request_id: str) -> asyncio.Event:
|
|
"""Create an asyncio.Event for a request and store it."""
|
|
event = asyncio.Event()
|
|
with pending_requests_lock:
|
|
pending_requests[request_id] = event
|
|
return event
|
|
|
|
|
|
def remove_pending_request(request_id: str):
|
|
"""Remove a pending request from the dictionary."""
|
|
with pending_requests_lock:
|
|
pending_requests.pop(request_id, None)
|
|
|
|
|
|
def notify_request(chatcmpl_request_id: str) -> bool:
|
|
"""
|
|
Notify a pending request that the KV store is complete.
|
|
Returns True if the request was found and notified.
|
|
"""
|
|
with pending_requests_lock:
|
|
# vLLM wraps the request ID as "chatcmpl-{uuid}-{suffix}",
|
|
# so strip the first and last segments to recover the original UUID.
|
|
request_id = "-".join(chatcmpl_request_id.split("-")[1:])
|
|
request_id = request_id[:16]
|
|
event = pending_requests.get(request_id, None)
|
|
if event:
|
|
# Schedule the event.set() on the main event loop
|
|
if main_event_loop is not None:
|
|
main_event_loop.call_soon_threadsafe(event.set)
|
|
return True
|
|
return False
|
|
|
|
|
|
async def _handle_disagg_request(request: Request, endpoint: str):
|
|
"""
|
|
Common handler for disaggregated prefill/decode requests.
|
|
|
|
Works for both /v1/completions and /v1/chat/completions — the only
|
|
difference is the *endpoint* path forwarded to the prefiller and decoder.
|
|
"""
|
|
try:
|
|
req_data = await request.json()
|
|
|
|
# Generate a random request ID
|
|
request_id = generate_request_id()
|
|
logger.info(f"Received {endpoint} request with generated ID: {request_id}")
|
|
|
|
# Pick prefill and decode clients
|
|
prefill_client, decode_client = round_robin_pick_clients()
|
|
|
|
# Create event for this request (signaled when KV store finishes)
|
|
event = create_pending_request(request_id)
|
|
|
|
try:
|
|
# Modify request for prefiller: set max_tokens=1
|
|
prefill_req_data = req_data.copy()
|
|
prefill_req_data["max_tokens"] = 1
|
|
if "max_completion_tokens" in prefill_req_data:
|
|
prefill_req_data["max_completion_tokens"] = 1
|
|
prefill_req_data["stream"] = False
|
|
prefill_req_data.pop("stream_options", None)
|
|
|
|
# Send to prefiller with X-Request-Id header (ignore output)
|
|
prefill_send_time = time.monotonic()
|
|
await send_request_to_prefiller(
|
|
prefill_client.client,
|
|
endpoint,
|
|
prefill_req_data,
|
|
request_id,
|
|
)
|
|
prefill_first_response_time = time.monotonic()
|
|
prefill_duration = prefill_first_response_time - prefill_send_time
|
|
logger.info(
|
|
f"Request {request_id}: prefill request"
|
|
f" duration = {prefill_duration:.4f}s"
|
|
)
|
|
|
|
# Wait for the event to be signaled (KV store finished)
|
|
await event.wait()
|
|
notify_time = time.monotonic()
|
|
notify_wait_duration = notify_time - prefill_first_response_time
|
|
logger.info(
|
|
f"Request {request_id}: finished saving KV caches after prefill"
|
|
f" response = {notify_wait_duration * 1000:.2f}ms"
|
|
)
|
|
logger.debug(f"Event signaled for {request_id}, forwarding to decoder")
|
|
|
|
finally:
|
|
# Clean up the pending request
|
|
remove_pending_request(request_id)
|
|
|
|
# Forward original request to decoder
|
|
is_stream = req_data.get("stream", False)
|
|
|
|
if is_stream:
|
|
|
|
async def generate_stream():
|
|
first_chunk = True
|
|
async for chunk in stream_service_response(
|
|
decode_client.client, endpoint, req_data
|
|
):
|
|
if first_chunk:
|
|
decode_first_response_time = time.monotonic()
|
|
latency = (
|
|
decode_first_response_time - prefill_first_response_time
|
|
)
|
|
logger.info(
|
|
f"Request {request_id}: latency between prefill first "
|
|
f"response and decode first response = "
|
|
f"{latency * 1000:.2f}ms"
|
|
)
|
|
first_chunk = False
|
|
yield chunk
|
|
|
|
return StreamingResponse(generate_stream(), media_type="application/json")
|
|
else:
|
|
response = await send_request_to_decoder(
|
|
decode_client.client, endpoint, req_data
|
|
)
|
|
return JSONResponse(content=response.json())
|
|
|
|
except Exception as e:
|
|
# Standard
|
|
import sys
|
|
import traceback
|
|
|
|
exc_info = sys.exc_info()
|
|
logger.error(f"Error in {endpoint} endpoint")
|
|
logger.error(str(e))
|
|
logger.error("".join(traceback.format_exception(*exc_info)))
|
|
raise
|
|
|
|
|
|
@app.post("/v1/completions")
|
|
async def handle_completions(request: Request):
|
|
"""Handle /v1/completions requests."""
|
|
return await _handle_disagg_request(request, "/v1/completions")
|
|
|
|
|
|
@app.post("/v1/chat/completions")
|
|
async def handle_chat_completions(request: Request):
|
|
"""Handle /v1/chat/completions requests."""
|
|
return await _handle_disagg_request(request, "/v1/chat/completions")
|
|
|
|
|
|
@app.get("/v1/models")
|
|
async def handle_models():
|
|
"""Handle /v1/models requests by forwarding to the first prefiller."""
|
|
try:
|
|
prefill_client = app.state.prefill_clients[0]
|
|
headers = _build_headers()
|
|
response = await prefill_client.client.get("/v1/models", headers=headers)
|
|
response.raise_for_status()
|
|
return JSONResponse(content=response.json())
|
|
except Exception as e:
|
|
logger.error(f"Error in /v1/models endpoint: {e}")
|
|
return JSONResponse(
|
|
content={"error": str(e)},
|
|
status_code=500,
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# Telemetry FastAPI app (runs on separate port)
|
|
# ============================================================================
|
|
|
|
telemetry_app = FastAPI()
|
|
|
|
|
|
@telemetry_app.post("/api/v1/telemetry")
|
|
async def handle_telemetry(request: Request):
|
|
"""
|
|
Handle telemetry POST requests.
|
|
|
|
Expected payload format (from FastAPIRequestTelemetry):
|
|
{
|
|
"event": "request_store_finished",
|
|
"request_ids_set": ["id1", "id2", ...],
|
|
"model_name": "...",
|
|
"world_size": N,
|
|
"kv_rank": K,
|
|
}
|
|
|
|
For each request ID in request_ids_set, signal the corresponding
|
|
event if it exists.
|
|
"""
|
|
try:
|
|
payload = await request.json()
|
|
|
|
event_type = payload.get("event")
|
|
request_ids = payload.get("request_ids_set", [])
|
|
|
|
for request_id in request_ids:
|
|
logger.info(
|
|
f"Received telemetry event: {event_type} for request: {request_id}"
|
|
)
|
|
|
|
notified_count = 0
|
|
for request_id in request_ids:
|
|
if notify_request(request_id):
|
|
notified_count += 1
|
|
logger.info(f"Notified request: {request_id}")
|
|
|
|
return JSONResponse(
|
|
content={
|
|
"status": "ok",
|
|
"notified": notified_count,
|
|
"total": len(request_ids),
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing telemetry: {e}")
|
|
return JSONResponse(
|
|
content={"status": "error", "message": str(e)},
|
|
status_code=500,
|
|
)
|
|
|
|
|
|
def run_telemetry_server(host: str, port: int):
|
|
"""Run the telemetry server in a separate thread."""
|
|
# Third Party
|
|
import uvicorn
|
|
|
|
uvicorn.run(telemetry_app, host=host, port=port, log_level="info")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
global global_args
|
|
global_args = parse_args()
|
|
|
|
# Third Party
|
|
import uvicorn
|
|
|
|
# Start telemetry server in a background thread
|
|
telemetry_thread = threading.Thread(
|
|
target=run_telemetry_server,
|
|
args=(global_args.host, global_args.telemetry_port),
|
|
daemon=True,
|
|
)
|
|
telemetry_thread.start()
|
|
logger.info(
|
|
f"Telemetry server started on {global_args.host}:{global_args.telemetry_port}"
|
|
)
|
|
|
|
# Run main proxy server
|
|
uvicorn.run(app, host=global_args.host, port=global_args.port)
|