455
examples/dynamic_bucket_load_balancer/hybrid_proxy_server.py
Normal file
455
examples/dynamic_bucket_load_balancer/hybrid_proxy_server.py
Normal file
@@ -0,0 +1,455 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Dynamic bucketing-based hybrid load balance proxy server.
|
||||
# See README.md in this directory for the tutorial and usage.
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import functools
|
||||
import heapq
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from dynamic_bucket_load_balancer import DynamicBucketLoadBalancer, ServerInfo, Task
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
try:
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
except ImportError:
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Use uvloop for a faster event loop if available
|
||||
try:
|
||||
import uvloop
|
||||
|
||||
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
class ServerState:
|
||||
def __init__(self, host, port):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.url = f"http://{host}:{port}/v1"
|
||||
self.client = httpx.AsyncClient(
|
||||
timeout=None,
|
||||
base_url=self.url,
|
||||
limits=httpx.Limits(max_connections=100000, max_keepalive_connections=100000),
|
||||
)
|
||||
self.active_tokens = 0
|
||||
self.aborted_requests = set()
|
||||
|
||||
def __eq__(self, other):
|
||||
self_host = self.host.replace("localhost", "0.0.0.0").replace("127.0.0.1", "0.0.0.0")
|
||||
other_host = other.host.replace("localhost", "0.0.0.0").replace("127.0.0.1", "0.0.0.0")
|
||||
return self_host == other_host and str(self.port) == str(other.port)
|
||||
|
||||
def __hash__(self):
|
||||
self_host = self.host.replace("localhost", "0.0.0.0").replace("127.0.0.1", "0.0.0.0")
|
||||
return hash((self_host, str(self.port)))
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.host}:{self.port}"
|
||||
|
||||
|
||||
@dataclass(order=True)
|
||||
class ServerHeapItem:
|
||||
priority: float
|
||||
server_idx: int
|
||||
server: ServerState
|
||||
|
||||
|
||||
class ProxyState:
|
||||
def __init__(self, server_instances):
|
||||
self.infer_servers: list[ServerState] = [ServerState(h, p) for h, p in server_instances]
|
||||
self.req_id_lock = asyncio.Lock()
|
||||
|
||||
# Dynamic bucket load balancer
|
||||
self.bucket_load_balancer = None
|
||||
|
||||
if global_args.enable_dynamic_bucket:
|
||||
self.num_buckets = 2 # Two buckets (short/long) when dynamic bucketing is enabled
|
||||
self.server_group_threshold = global_args.server_group_threshold
|
||||
buckets = [(0, self.server_group_threshold), (self.server_group_threshold, global_args.max_request_tokens)]
|
||||
|
||||
self.bucket_load_balancer = DynamicBucketLoadBalancer(buckets=buckets)
|
||||
else:
|
||||
self.num_buckets = 1 # No bucketing by default
|
||||
|
||||
# Priority queue per group; smaller score = higher priority (lower load)
|
||||
server_heap_items = [ServerHeapItem(0.0, i, server) for i, server in enumerate(self.infer_servers)]
|
||||
self.server_heaps: list[list[ServerHeapItem]] = self._group_servers(server_heap_items, self.num_buckets)
|
||||
self.server_idx_to_group_idx = {}
|
||||
|
||||
# Heapify each group
|
||||
for idx, cur_heap in enumerate(self.server_heaps):
|
||||
for server_item in cur_heap:
|
||||
self.server_idx_to_group_idx[server_item.server_idx] = idx
|
||||
heapq.heapify(cur_heap)
|
||||
|
||||
logger.info(
|
||||
"Dynamic bucket enabled: %s, number of groups: %s",
|
||||
global_args.enable_dynamic_bucket,
|
||||
len(self.server_heaps),
|
||||
)
|
||||
for group_idx, cur_heap in enumerate(self.server_heaps):
|
||||
logger.info("Group %s: %s", group_idx, cur_heap)
|
||||
|
||||
@staticmethod
|
||||
def _group_servers(servers: list[ServerHeapItem], num_groups: int):
|
||||
"""
|
||||
Split servers into num_groups groups.
|
||||
|
||||
Args:
|
||||
servers (list): the server list to group.
|
||||
num_groups (int): the number of groups.
|
||||
|
||||
Returns:
|
||||
list[list]: the grouped server list.
|
||||
|
||||
Raises:
|
||||
ValueError: when num_groups <= 0.
|
||||
"""
|
||||
if num_groups <= 0:
|
||||
raise ValueError("Num of group is illegal")
|
||||
|
||||
if len(servers) < num_groups:
|
||||
raise ValueError("Number of servers must greater than or equal to number of groups")
|
||||
|
||||
n = len(servers)
|
||||
if n == 0:
|
||||
return [[] for _ in range(num_groups)]
|
||||
elif n == 1:
|
||||
return [servers]
|
||||
|
||||
base_size = n // num_groups
|
||||
remainder = n % num_groups
|
||||
|
||||
groups = []
|
||||
start_index = 0
|
||||
for i in range(num_groups):
|
||||
group_size = base_size + 1 if i < remainder else base_size
|
||||
end_index = start_index + group_size
|
||||
groups.append(servers[start_index:end_index])
|
||||
start_index = end_index
|
||||
|
||||
return groups
|
||||
|
||||
def _update_server_priority(self, server_idx: int):
|
||||
"""Update the priority of a server in the heap."""
|
||||
server = self.infer_servers[server_idx]
|
||||
priority = server.active_tokens
|
||||
# Remove the old entry, then add the new one
|
||||
group_idx = self.server_idx_to_group_idx[server_idx]
|
||||
|
||||
self.server_heaps[group_idx] = [
|
||||
server_heap_item
|
||||
for server_heap_item in self.server_heaps[group_idx]
|
||||
if server_heap_item.server_idx != server_idx
|
||||
]
|
||||
self.server_heaps[group_idx].append(ServerHeapItem(priority, server_idx, server))
|
||||
heapq.heapify(self.server_heaps[group_idx])
|
||||
|
||||
async def next_req_id(self):
|
||||
async with self.req_id_lock:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def select_server(self, token_count, group_idx: int):
|
||||
if not self.infer_servers:
|
||||
raise RuntimeError("No inference servers available")
|
||||
|
||||
server_heap_item: ServerHeapItem = heapq.heappop(self.server_heaps[group_idx])
|
||||
chosen_server_idx = server_heap_item.server_idx
|
||||
|
||||
# Update the chosen server (accumulate load)
|
||||
self.infer_servers[chosen_server_idx].active_tokens += token_count
|
||||
|
||||
# Update priority and re-add to the heap
|
||||
self._update_server_priority(chosen_server_idx)
|
||||
|
||||
return chosen_server_idx
|
||||
|
||||
def release_server(self, idx: int, token_count, req_id):
|
||||
self.infer_servers[idx].active_tokens -= token_count
|
||||
if global_args.enable_dynamic_bucket and req_id is not None and self.bucket_load_balancer is not None:
|
||||
self.bucket_load_balancer.release_task(req_id)
|
||||
# Update the priority queue after release
|
||||
self._update_server_priority(idx)
|
||||
|
||||
def calculate_request_score(self, request_length: int, max_tokens: int = 16, ignore_eos: bool = False) -> float:
|
||||
if ignore_eos:
|
||||
return request_length + max_tokens
|
||||
else:
|
||||
# Note that 0.5 is an empirical value here because we don't know
|
||||
# the actual number of tokens generated before EOS.
|
||||
return request_length + 0.5 * max_tokens
|
||||
|
||||
def calculate_request_tokens(self, request_length: int) -> float:
|
||||
return request_length / 4.0
|
||||
|
||||
def select_server_group(self, req_id: str, request_tokens, priority_score) -> tuple[int, Task | None]:
|
||||
"""Pick the best group given the request length and the current load of each group."""
|
||||
if global_args.enable_dynamic_bucket and self.bucket_load_balancer is not None:
|
||||
group_idx, task = self.bucket_load_balancer.dispatch_single_task(req_id, request_tokens, priority_score)
|
||||
return group_idx, task
|
||||
else:
|
||||
return 0, None
|
||||
|
||||
|
||||
proxy_state = None
|
||||
|
||||
|
||||
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("--server-hosts", type=str, nargs="+", default=["localhost"])
|
||||
parser.add_argument("--server-ports", type=int, nargs="+", default=[8001])
|
||||
parser.add_argument("--max-retries", type=int, default=3, help="Maximum number of retries for HTTP requests")
|
||||
parser.add_argument(
|
||||
"--retry-delay", type=float, default=0.001, help="Base delay (seconds) for exponential backoff retries"
|
||||
)
|
||||
|
||||
parser.add_argument("--server-group-threshold", type=int, default=32 * 1024, help="Threshold of server groups")
|
||||
parser.add_argument("--max-request-tokens", type=int, default=128 * 1024, help="Max tokens of request")
|
||||
parser.add_argument(
|
||||
"--enable-dynamic-bucket", action="store_true", default=False, help="Enable dynamic bucket load Balancer"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
if len(args.server_hosts) != len(args.server_ports):
|
||||
raise ValueError("Number of dp hosts must match number of dp ports")
|
||||
args.server_instances = list(zip(args.server_hosts, args.server_ports))
|
||||
return args
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
global proxy_state
|
||||
proxy_state = ProxyState(global_args.server_instances)
|
||||
logger.debug("Initialized %s dp server clients.", len(proxy_state.infer_servers))
|
||||
yield
|
||||
for p in proxy_state.infer_servers:
|
||||
await p.client.aclose()
|
||||
|
||||
|
||||
async def listen_for_disconnect(request: Request) -> None:
|
||||
"""Return when a disconnect message is received."""
|
||||
while True:
|
||||
message = await request.receive()
|
||||
if message["type"] == "http.disconnect":
|
||||
break
|
||||
|
||||
|
||||
def with_cancellation(handler_func):
|
||||
@functools.wraps(handler_func)
|
||||
async def wrapper(*args, **kwargs):
|
||||
request = kwargs["request"]
|
||||
handler_task = asyncio.create_task(handler_func(*args, **kwargs))
|
||||
cancellation_task = asyncio.create_task(listen_for_disconnect(request))
|
||||
done, pending = await asyncio.wait([handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if handler_task in done:
|
||||
return handler_task.result()
|
||||
return None
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
|
||||
async def stream_service_response_with_retry(
|
||||
client: httpx.AsyncClient,
|
||||
endpoint: str,
|
||||
req_data: dict,
|
||||
request_id: str,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 0.2,
|
||||
):
|
||||
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
|
||||
for attempt in range(1, max_retries + 1):
|
||||
# Reset per retry to avoid leaking a stale True from a previous iteration
|
||||
first_chunk_sent = False
|
||||
try:
|
||||
async with client.stream("POST", endpoint, json=req_data, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
async for chunk in response.aiter_bytes():
|
||||
first_chunk_sent = True
|
||||
yield chunk
|
||||
return # Success; exit after streaming completes
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as e:
|
||||
# After the first chunk is forwarded, retry is forbidden (would duplicate/corrupt the stream).
|
||||
if first_chunk_sent:
|
||||
logger.error("Streaming to client interrupted after response started: %s", str(e))
|
||||
return
|
||||
if attempt < max_retries:
|
||||
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, str(e))
|
||||
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
|
||||
else:
|
||||
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
|
||||
raise e
|
||||
except Exception as e:
|
||||
# Same guard as above for non-HTTP exceptions
|
||||
if first_chunk_sent:
|
||||
logger.error("Streaming to client interrupted after response started: %s", str(e))
|
||||
return
|
||||
if attempt < max_retries:
|
||||
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, str(e))
|
||||
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
|
||||
else:
|
||||
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
|
||||
raise e
|
||||
|
||||
|
||||
async def _select_instance(api: str, req_data: Any, request_length: int):
|
||||
# refer to vLLM sampling_params: max_token default value
|
||||
max_tokens = req_data.get("max_tokens", 16)
|
||||
ignore_eos = req_data.get("ignore_eos", False)
|
||||
priority_score = 0.0
|
||||
if global_args.enable_dynamic_bucket:
|
||||
priority_score = proxy_state.calculate_request_tokens(request_length)
|
||||
else:
|
||||
priority_score = proxy_state.calculate_request_score(
|
||||
request_length, max_tokens=max_tokens, ignore_eos=ignore_eos
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"Request length: %s, max tokens: %s, ignore_eos: %s, Priority score: %s",
|
||||
request_length,
|
||||
max_tokens,
|
||||
ignore_eos,
|
||||
priority_score,
|
||||
)
|
||||
request_id = await proxy_state.next_req_id()
|
||||
# Select server based on priority score
|
||||
request_tokens = proxy_state.calculate_request_tokens(request_length)
|
||||
group_idx, task = proxy_state.select_server_group(request_id, request_tokens, priority_score)
|
||||
|
||||
try:
|
||||
server_idx = proxy_state.select_server(priority_score, group_idx)
|
||||
except Exception:
|
||||
if global_args.enable_dynamic_bucket and task is not None and proxy_state.bucket_load_balancer is not None:
|
||||
proxy_state.bucket_load_balancer.release_task(task.id)
|
||||
raise
|
||||
|
||||
if global_args.enable_dynamic_bucket and task is not None:
|
||||
task.server_info = ServerInfo("DP", server_idx)
|
||||
|
||||
chosen_server = proxy_state.infer_servers[server_idx]
|
||||
logger.debug(
|
||||
"[group_idx=%s, server_idx=%s] Choose server %s to process request %s",
|
||||
group_idx,
|
||||
server_idx,
|
||||
chosen_server.url,
|
||||
request_id,
|
||||
)
|
||||
return InstanceInfo(
|
||||
request_id=request_id, server_idx=server_idx, priority_score=priority_score, server_state=chosen_server
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class InstanceInfo:
|
||||
request_id: str
|
||||
server_idx: int
|
||||
priority_score: float
|
||||
server_state: ServerState
|
||||
|
||||
|
||||
async def _handle_completions(api: str, request: Request):
|
||||
# streaming_started ensures release_server runs exactly once: in
|
||||
# generate_stream's finally on the normal path, or below if it never started.
|
||||
instance_info = None
|
||||
streaming_started = False
|
||||
try:
|
||||
req_data = await request.json()
|
||||
req_body = await request.body()
|
||||
request_length = len(req_body)
|
||||
instance_info = await _select_instance(api, req_data, request_length)
|
||||
|
||||
async def generate_stream():
|
||||
nonlocal instance_info
|
||||
try:
|
||||
async for chunk in stream_service_response_with_retry(
|
||||
instance_info.server_state.client, # type: ignore
|
||||
api,
|
||||
req_data,
|
||||
request_id=instance_info.request_id, # type: ignore
|
||||
max_retries=global_args.max_retries,
|
||||
base_delay=global_args.retry_delay,
|
||||
):
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error during streaming from server %s: %s, the aborted request is: %s.",
|
||||
instance_info.server_state.url, # type: ignore
|
||||
str(e),
|
||||
instance_info.request_id, # type: ignore
|
||||
)
|
||||
finally:
|
||||
# Release load after streaming completes
|
||||
proxy_state.release_server( # type: ignore
|
||||
instance_info.server_idx, # type: ignore
|
||||
instance_info.priority_score, # type: ignore
|
||||
instance_info.request_id, # type: ignore
|
||||
)
|
||||
|
||||
streaming_started = True
|
||||
return StreamingResponse(generate_stream(), media_type="application/json")
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
exc_info = sys.exc_info()
|
||||
print(f"Error occurred in external dp proxy server - {api} endpoint")
|
||||
print(e)
|
||||
print("".join(traceback.format_exception(*exc_info)))
|
||||
raise
|
||||
finally:
|
||||
# If streaming never started (client disconnect or selection error),
|
||||
# release here to avoid leaking active_tokens / the bucket task; the
|
||||
# normal path already released in generate_stream.
|
||||
if instance_info is not None and not streaming_started:
|
||||
proxy_state.release_server(instance_info.server_idx, instance_info.priority_score, instance_info.request_id)
|
||||
|
||||
|
||||
@app.post("/v1/completions")
|
||||
@with_cancellation
|
||||
async def handle_completions(request: Request):
|
||||
return await _handle_completions("/completions", request)
|
||||
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
@with_cancellation
|
||||
async def handle_chat_completions(request: Request):
|
||||
return await _handle_completions("/chat/completions", request)
|
||||
|
||||
|
||||
@app.get("/healthcheck")
|
||||
async def healthcheck():
|
||||
return {
|
||||
"status": "ok",
|
||||
"server_instances": len(proxy_state.infer_servers),
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
global global_args
|
||||
global_args = parse_args()
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host=global_args.host, port=global_args.port)
|
||||
Reference in New Issue
Block a user