@@ -0,0 +1,640 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
|
||||
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Tutorial: Using the Load Balance Proxy Server Example
|
||||
#
|
||||
# This proxy server is designed to distribute requests between multiple
|
||||
# "prefiller" and "decoder" backend servers for large language model inference.
|
||||
# It is useful for scaling out inference workloads and balancing load across
|
||||
# multiple backend instances.
|
||||
#
|
||||
# Features:
|
||||
# - Load balances requests to multiple prefiller and decoder servers.
|
||||
# - Supports OpenAI-compatible /v1/completions and /v1/chat/completions endpoints.
|
||||
# - Streams responses from backend servers to clients.
|
||||
#
|
||||
# Prerequisites:
|
||||
# - Python 3.8+
|
||||
# - Install dependencies:
|
||||
# pip install fastapi<0.124.0 httpx uvicorn vllm
|
||||
#
|
||||
# Step 1: Start Your Backend Servers
|
||||
# ----------------------------------
|
||||
# You need to have at least one prefiller and one decoder backend running.
|
||||
# These can be mock servers or actual vLLM servers.
|
||||
#
|
||||
# For testing, you can use the provided mock server:
|
||||
#
|
||||
# vllm serve --host 0.0.0.0 --port 8100 ... # Prefiller 1
|
||||
# vllm serve --host 0.0.0.0 --port 8101 ... # Prefiller 2
|
||||
# vllm serve --host 0.0.0.0 --port 8200 ... # Decoder 1
|
||||
# vllm serve --host 0.0.0.0 --port 8201 ... # Decoder 2
|
||||
#
|
||||
# Step 2: Start the Proxy Server
|
||||
# ------------------------------
|
||||
# Run the proxy server, specifying the host/port for each prefiller and decoder:
|
||||
#
|
||||
# python load_balance_proxy_server_example.py \
|
||||
# --host 0.0.0.0 --port 9000 \
|
||||
# --prefiller-hosts 127.0.0.1 127.0.0.1 \
|
||||
# --prefiller-ports 8100 8101 \
|
||||
# --decoder-hosts 127.0.0.1 127.0.0.1 \
|
||||
# --decoder-ports 8200 8201
|
||||
#
|
||||
# This will start the proxy on port 9000, load balancing between two prefiller
|
||||
# and two decoder servers.
|
||||
#
|
||||
# Step 3: Send a Request to the Proxy
|
||||
# -----------------------------------
|
||||
# You can now send OpenAI-compatible requests to the proxy. For example:
|
||||
#
|
||||
# curl -X POST http://localhost:9000/v1/completions \
|
||||
# -H "Content-Type: application/json" \
|
||||
# -d '{
|
||||
# "model": "your-model",
|
||||
# "prompt": "The quick brown fox jumps over the lazy dog",
|
||||
# "max_tokens": 16
|
||||
# }'
|
||||
#
|
||||
# Or for chat completions:
|
||||
#
|
||||
# curl -X POST http://localhost:9000/v1/chat/completions \
|
||||
# -H "Content-Type: application/json" \
|
||||
# -d '{
|
||||
# "model": "your-model",
|
||||
# "messages": [{"role": "user", "content": "Hello!"}],
|
||||
# "max_tokens": 16
|
||||
# }'
|
||||
#
|
||||
# Step 4: Health Check
|
||||
# --------------------
|
||||
# To check if the proxy is running and see how many backend instances are
|
||||
# connected, use:
|
||||
#
|
||||
# curl http://localhost:9000/healthcheck
|
||||
#
|
||||
# This will return a JSON object with the status and the number of prefiller
|
||||
# and decoder instances.
|
||||
#
|
||||
# Notes:
|
||||
# - You can scale the number of prefiller and decoder servers as needed.
|
||||
# - The proxy will round-robin requests to balance load.
|
||||
# - For production, ensure your backend servers are robust and secure.
|
||||
#
|
||||
# For more details, see the code and comments in this file.
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import copy
|
||||
import functools
|
||||
import heapq
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import httpx
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
from vllm.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Add uvloop for 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"
|
||||
# Auto-completion for ipv6
|
||||
try:
|
||||
ip = ipaddress.ip_address(self.host)
|
||||
if isinstance(ip, ipaddress.IPv6Address):
|
||||
self.url = f"http://[{host}]:{port}/v1"
|
||||
except Exception:
|
||||
pass
|
||||
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.active_kv_cache = 0 # Only for prefiller
|
||||
self.active_requests = 0 # Number of active requests
|
||||
self.aborted_requests = set() # Track aborted requests
|
||||
# Removed individual server lock - will use global locks instead
|
||||
|
||||
|
||||
class ProxyState:
|
||||
def __init__(self, prefiller_instances, decoder_instances):
|
||||
self.prefillers: list[ServerState] = [ServerState(h, p) for h, p in prefiller_instances]
|
||||
self.decoders: list[ServerState] = [ServerState(h, p) for h, p in decoder_instances]
|
||||
self.req_to_prefiller = {}
|
||||
self.req_id_lock = asyncio.Lock()
|
||||
# Removed selection locks - no longer needed for synchronous methods
|
||||
|
||||
# Initialize priority queues for efficient server selection
|
||||
# Each entry is (priority_score, server_index, server_reference)
|
||||
# Lower priority score = higher priority (less loaded)
|
||||
self.prefiller_heap = [(0, i, server) for i, server in enumerate(self.prefillers)]
|
||||
self.decoder_heap = [(0, i, server) for i, server in enumerate(self.decoders)]
|
||||
heapq.heapify(self.prefiller_heap)
|
||||
heapq.heapify(self.decoder_heap)
|
||||
self.req_id_future = {}
|
||||
self.req_data_dict = {}
|
||||
|
||||
def _update_prefiller_priority(self, server_idx: int):
|
||||
"""Update the priority of a prefiller server in the heap."""
|
||||
server = self.prefillers[server_idx]
|
||||
# Priority based on active_tokens and active_kv_cache
|
||||
priority = server.active_tokens + server.active_kv_cache * 0.3
|
||||
# Remove old entry and add new one
|
||||
self.prefiller_heap = [(p, i, s) for p, i, s in self.prefiller_heap if i != server_idx]
|
||||
heapq.heappush(self.prefiller_heap, (priority, server_idx, server)) # type: ignore
|
||||
|
||||
def _update_decoder_priority(self, server_idx: int):
|
||||
"""Update the priority of a decoder server in the heap."""
|
||||
server = self.decoders[server_idx]
|
||||
priority = server.active_tokens
|
||||
# Remove old entry and add new one
|
||||
self.decoder_heap = [(p, i, s) for p, i, s in self.decoder_heap if i != server_idx]
|
||||
heapq.heappush(self.decoder_heap, (priority, server_idx, server)) # type: ignore
|
||||
|
||||
def abort_prefiller_request(self, server_idx: int, request_id): # Changed to synchronous
|
||||
"""
|
||||
Mark a request as aborted. This will helps to release kv cache in
|
||||
prefiller node.
|
||||
"""
|
||||
# No lock needed - atomic operation
|
||||
self.prefillers[server_idx].aborted_requests.add(request_id)
|
||||
|
||||
def acquire_aborted_prefiller_requests(self, server_idx: int): # Changed to synchronous
|
||||
"""
|
||||
Get the set of aborted requests and clear it.
|
||||
This is used to release kv cache in prefiller node.
|
||||
"""
|
||||
# No lock needed - atomic operation
|
||||
aborted_requests = self.prefillers[server_idx].aborted_requests.copy()
|
||||
self.prefillers[server_idx].aborted_requests.clear()
|
||||
return aborted_requests
|
||||
|
||||
async def next_req_id(self):
|
||||
async with self.req_id_lock:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
def select_prefiller(self, token_count): # Changed to synchronous
|
||||
# No lock needed - entire function is atomic
|
||||
if not self.prefiller_heap:
|
||||
raise RuntimeError("No prefiller servers available")
|
||||
|
||||
priority, chosen, server = heapq.heappop(self.prefiller_heap)
|
||||
|
||||
# Update the chosen server atomically
|
||||
self.prefillers[chosen].active_tokens += token_count
|
||||
self.prefillers[chosen].active_kv_cache += token_count
|
||||
|
||||
# Update priority and re-add to heap
|
||||
self._update_prefiller_priority(chosen)
|
||||
|
||||
return chosen
|
||||
|
||||
def release_prefiller(self, idx, token_count): # Changed to synchronous
|
||||
# No lock needed - atomic operation
|
||||
self.prefillers[idx].active_tokens -= token_count
|
||||
# Update priority queue after releasing
|
||||
self._update_prefiller_priority(idx)
|
||||
|
||||
def release_prefiller_kv(self, idx, token_count): # Changed to synchronous
|
||||
# No lock needed - atomic operation
|
||||
if self.prefillers[idx].active_kv_cache > 0:
|
||||
self.prefillers[idx].active_kv_cache -= token_count
|
||||
# Update priority queue after releasing
|
||||
self._update_prefiller_priority(idx)
|
||||
|
||||
def select_decoder(self, token_count): # Changed to synchronous
|
||||
# No lock needed - entire function is atomic
|
||||
if not self.decoder_heap:
|
||||
raise RuntimeError("No decoder servers available")
|
||||
|
||||
priority, chosen, server = heapq.heappop(self.decoder_heap)
|
||||
|
||||
# Update the chosen server atomically
|
||||
self.decoders[chosen].active_tokens += token_count
|
||||
|
||||
# Update priority and re-add to heap
|
||||
self._update_decoder_priority(chosen)
|
||||
|
||||
return chosen
|
||||
|
||||
def release_decoder(self, idx, token_count): # Changed to synchronous
|
||||
# No lock needed - atomic operation
|
||||
self.decoders[idx].active_tokens -= token_count
|
||||
# Update priority queue after releasing
|
||||
self._update_decoder_priority(idx)
|
||||
|
||||
# Omni_infer's calculate_input_scores function
|
||||
def calculate_prefill_scores(self, request_length: int) -> float:
|
||||
length_score = request_length / 4.0
|
||||
input_score = length_score * 0.0345 + 120.0745
|
||||
return input_score
|
||||
|
||||
def calculate_decode_scores(self, request_length: int) -> float:
|
||||
return request_length
|
||||
|
||||
|
||||
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("--prefiller-hosts", type=str, nargs="+", default=["localhost"])
|
||||
parser.add_argument("--prefiller-ports", type=int, nargs="+", default=[8001])
|
||||
parser.add_argument("--decoder-hosts", type=str, nargs="+", default=["localhost"])
|
||||
parser.add_argument("--decoder-ports", type=int, nargs="+", default=[8002])
|
||||
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"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
logger.info(
|
||||
"Decoder hosts will access Proxy host:port/metaserver, ensure that %s can access %s:%s/metaserver",
|
||||
set(args.decoder_hosts),
|
||||
args.host,
|
||||
args.port,
|
||||
)
|
||||
# Wildcard address is not allowed for layerwise connector
|
||||
if args.host in ["0.0.0.0", "::", "0:0:0:0:0:0:0:0"]:
|
||||
raise ValueError(
|
||||
f"Decoder hosts will access Proxy host:port/metaserver, to avoid configuration errors, "
|
||||
f"the Wildcard Address {args.host} is not allowed for proxy"
|
||||
)
|
||||
if len(args.prefiller_hosts) != len(args.prefiller_ports):
|
||||
raise ValueError("Number of prefiller hosts must match number of prefiller ports")
|
||||
if len(args.decoder_hosts) != len(args.decoder_ports):
|
||||
raise ValueError("Number of decoder hosts must match number of decoder ports")
|
||||
args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
|
||||
args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
|
||||
return args
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
global proxy_state
|
||||
proxy_state = ProxyState(global_args.prefiller_instances, global_args.decoder_instances)
|
||||
print(f"Initialized {len(proxy_state.prefillers)} prefill clients and {len(proxy_state.decoders)} decode clients.")
|
||||
yield
|
||||
for p in proxy_state.prefillers:
|
||||
await p.client.aclose()
|
||||
for d in proxy_state.decoders:
|
||||
await d.client.aclose()
|
||||
|
||||
|
||||
async def listen_for_disconnect(request: Request) -> None:
|
||||
"""Return if 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 send_request_to_service(
|
||||
client: httpx.AsyncClient,
|
||||
prefiller_id: int,
|
||||
endpoint: str,
|
||||
req_data: dict,
|
||||
request_id: str,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 0.2,
|
||||
):
|
||||
proxy_state.acquire_aborted_prefiller_requests(prefiller_id)
|
||||
req_data = req_data.copy()
|
||||
req_data["stream"] = False
|
||||
req_data["max_tokens"] = 1
|
||||
req_data["min_tokens"] = 1
|
||||
if "max_completion_tokens" in req_data:
|
||||
req_data["max_completion_tokens"] = 1
|
||||
if "stream_options" in req_data:
|
||||
del req_data["stream_options"]
|
||||
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
|
||||
last_exc = None
|
||||
for attempt in range(1, max_retries + 1):
|
||||
try:
|
||||
response = await client.post(endpoint, json=req_data, headers=headers)
|
||||
response.raise_for_status()
|
||||
if request_id in proxy_state.req_id_future:
|
||||
result_future = proxy_state.req_id_future[request_id]
|
||||
result_future.set_result(response.json()["kv_transfer_params"])
|
||||
return
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as e:
|
||||
logger.warning("Attempt %s failed for %s: %s", attempt, endpoint, e)
|
||||
last_exc = e
|
||||
if attempt < max_retries:
|
||||
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
|
||||
else:
|
||||
logger.error("All %s attempts failed for %s.", max_retries, endpoint)
|
||||
raise last_exc
|
||||
|
||||
|
||||
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):
|
||||
try:
|
||||
async with client.stream("POST", endpoint, json=req_data, headers=headers) as response:
|
||||
response.raise_for_status()
|
||||
first_chunk_sent = False
|
||||
async for chunk in response.aiter_bytes():
|
||||
first_chunk_sent = True
|
||||
yield chunk
|
||||
return # Success, exit after streaming
|
||||
except (httpx.RequestError, httpx.HTTPStatusError) as e:
|
||||
if attempt < max_retries:
|
||||
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, 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:
|
||||
# If any chunk has been sent, do not retry, just log and drop
|
||||
if "first_chunk_sent" in locals() and first_chunk_sent:
|
||||
logger.error("Streaming to client interrupted after response started: %s", e)
|
||||
return
|
||||
else:
|
||||
if attempt < max_retries:
|
||||
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
|
||||
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
|
||||
else:
|
||||
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
|
||||
raise e
|
||||
|
||||
|
||||
def get_api_request_id(api, req_id):
|
||||
if api == "/completions":
|
||||
return "cmpl-" + req_id + "-0"
|
||||
elif api == "/chat/completions":
|
||||
return "chatcmpl-" + req_id
|
||||
|
||||
|
||||
def get_origin_request_id(api, req_id):
|
||||
if api == "/completions":
|
||||
return req_id.replace("cmpl-", "")[:-2]
|
||||
elif api == "/chat/completions":
|
||||
return req_id.replace("chatcmpl-", "")
|
||||
|
||||
|
||||
async def _handle_completions(api: str, request: Request):
|
||||
try:
|
||||
req_data = await request.json()
|
||||
req_body = await request.body()
|
||||
request_length = len(req_body)
|
||||
request_id = await proxy_state.next_req_id()
|
||||
request_id_api = get_api_request_id(api, request_id)
|
||||
proxy_state.req_data_dict[request_id_api] = (copy.deepcopy(req_data), request_length, api)
|
||||
req_data["kv_transfer_params"] = {
|
||||
"do_remote_decode": False,
|
||||
"do_remote_prefill": True,
|
||||
"metaserver": f"http://{global_args.host}:{global_args.port}/v1/metaserver",
|
||||
}
|
||||
# Select decoder
|
||||
decoder_score = proxy_state.calculate_decode_scores(request_length)
|
||||
logger.debug("Decoder score: %f", decoder_score)
|
||||
# Use the prefiller's kv_transfer_params to select decoder
|
||||
decoder_idx = proxy_state.select_decoder(decoder_score)
|
||||
decoder = proxy_state.decoders[decoder_idx]
|
||||
# logger.debug("Using %s %s", prefiller.url, decoder.url)
|
||||
# Stream response from decoder
|
||||
released_kv = False
|
||||
|
||||
# Record request info for recompute
|
||||
stream_flag = bool(req_data.get("stream", False))
|
||||
chat_flag = "messages" in req_data
|
||||
if "prompt" in req_data:
|
||||
origin_prompt = req_data["prompt"]
|
||||
elif chat_flag:
|
||||
messages = req_data["messages"]
|
||||
origin_prompt = messages[0].get("content", "")
|
||||
if isinstance(origin_prompt, list):
|
||||
origin_prompt = origin_prompt[0].get("text", "")
|
||||
else:
|
||||
origin_prompt = ""
|
||||
# refer to vLLM sampling_params: max_token default value
|
||||
origin_max_tokens = req_data.get("max_tokens", 16)
|
||||
|
||||
async def generate_stream():
|
||||
nonlocal released_kv
|
||||
generated_token = ""
|
||||
released_kv = False
|
||||
retry_count = 0
|
||||
retry = True
|
||||
completion_tokens = 0
|
||||
# Only one await per chunk, minimal logic in loop
|
||||
try:
|
||||
while retry:
|
||||
retry = False
|
||||
async for chunk in stream_service_response_with_retry(
|
||||
decoder.client,
|
||||
api,
|
||||
req_data,
|
||||
request_id=request_id,
|
||||
max_retries=global_args.max_retries,
|
||||
base_delay=global_args.retry_delay,
|
||||
):
|
||||
try:
|
||||
chunk_str = chunk.decode("utf-8").strip()
|
||||
except UnicodeDecodeError:
|
||||
logger.debug("Skipping chunk: %s", chunk)
|
||||
yield chunk
|
||||
continue
|
||||
if not chunk_str:
|
||||
continue
|
||||
if chunk_str.startswith("data: "):
|
||||
chunk_str = chunk_str[len("data: ") :]
|
||||
try:
|
||||
chunk_json = json.loads(chunk_str)
|
||||
except json.JSONDecodeError:
|
||||
# if chunk is [done], skip it.
|
||||
logger.debug("Skipping chunk: %s", chunk_str)
|
||||
yield chunk
|
||||
continue
|
||||
choices = chunk_json.get("choices", [])
|
||||
if not choices:
|
||||
yield chunk
|
||||
continue
|
||||
|
||||
choice = choices[0]
|
||||
delta = choice.get("delta") or {}
|
||||
message = choice.get("message") or {}
|
||||
content = delta.get("content") or message.get("content") or choice.get("text") or ""
|
||||
generated_token += content
|
||||
|
||||
stop_reason = choice.get("stop_reason")
|
||||
usage = chunk_json.get("usage", {})
|
||||
completion_tokens = (
|
||||
(completion_tokens + 1)
|
||||
if stream_flag
|
||||
else (completion_tokens + usage.get("completion_tokens"))
|
||||
)
|
||||
if stop_reason == "recomputed":
|
||||
retry = True
|
||||
retry_count += 1
|
||||
if chat_flag:
|
||||
messages[0]["content"] = origin_prompt + generated_token
|
||||
else:
|
||||
req_data["prompt"] = origin_prompt + generated_token
|
||||
req_data["max_tokens"] = origin_max_tokens - completion_tokens + retry_count
|
||||
break
|
||||
if retry_count > 0 and not stream_flag:
|
||||
if chat_flag:
|
||||
choice["message"]["content"] = generated_token
|
||||
else:
|
||||
choice["text"] = generated_token
|
||||
chunk = json.dumps(chunk_json).encode("utf-8")
|
||||
yield chunk
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error during streaming from decoder %s: %s the aborted request %s "
|
||||
"will be routing to the target prefiller when new request is ready to dispatch to it",
|
||||
decoder.url,
|
||||
e,
|
||||
request_id,
|
||||
)
|
||||
finally:
|
||||
# After streaming done, release tokens
|
||||
proxy_state.release_decoder(decoder_idx, decoder_score)
|
||||
|
||||
if stream_flag:
|
||||
return StreamingResponse(generate_stream(), media_type="text/event-stream")
|
||||
else:
|
||||
return StreamingResponse(generate_stream(), media_type="application/json")
|
||||
except Exception as e:
|
||||
import traceback
|
||||
|
||||
exc_info = sys.exc_info()
|
||||
print(f"Error occurred in disagg prefill proxy server - {api} endpoint")
|
||||
print(e)
|
||||
print("".join(traceback.format_exception(*exc_info)))
|
||||
raise
|
||||
|
||||
|
||||
@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",
|
||||
"prefill_instances": len(proxy_state.prefillers),
|
||||
"decode_instances": len(proxy_state.decoders),
|
||||
}
|
||||
|
||||
|
||||
@app.post("/reset_prefix_cache")
|
||||
async def reset_prefix_cache(request: Request):
|
||||
params = dict(request.query_params)
|
||||
failures = []
|
||||
for client, base_url in [
|
||||
(s.client, f"http://{s.host}:{s.port}") for s in proxy_state.prefillers + proxy_state.decoders
|
||||
]:
|
||||
try:
|
||||
resp = await client.post(f"{base_url}/reset_prefix_cache", params=params)
|
||||
resp.raise_for_status()
|
||||
except Exception as e:
|
||||
logger.error("reset_prefix_cache failed for %s: %s", base_url, e)
|
||||
failures.append(base_url)
|
||||
if failures:
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
return JSONResponse(status_code=500, content={"failed": failures})
|
||||
from fastapi.responses import Response as FastAPIResponse
|
||||
|
||||
return FastAPIResponse(status_code=200)
|
||||
|
||||
|
||||
@app.post("/v1/metaserver")
|
||||
async def metaserver(request: Request):
|
||||
try:
|
||||
kv_transfer_params = await request.json()
|
||||
|
||||
request_id = kv_transfer_params["request_id"]
|
||||
assert request_id in proxy_state.req_data_dict
|
||||
req_data, request_length, api = proxy_state.req_data_dict[request_id]
|
||||
request_id = get_origin_request_id(api, request_id)
|
||||
req_data["kv_transfer_params"] = kv_transfer_params
|
||||
prefiller_score = proxy_state.calculate_prefill_scores(request_length)
|
||||
logger.debug("Request length: %s, Prefiller score: %s", request_length, prefiller_score)
|
||||
|
||||
# Select prefiller
|
||||
prefiller_idx = proxy_state.select_prefiller(prefiller_score)
|
||||
prefiller = proxy_state.prefillers[prefiller_idx]
|
||||
logger.debug("Using prefill prefiller.url=%r req_data=%r", prefiller.url, req_data)
|
||||
# Send request to prefiller
|
||||
await send_request_to_service(
|
||||
prefiller.client,
|
||||
prefiller_idx,
|
||||
api,
|
||||
req_data,
|
||||
request_id,
|
||||
max_retries=global_args.max_retries,
|
||||
base_delay=global_args.retry_delay,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Post metaserver failed with: %s", e)
|
||||
finally:
|
||||
proxy_state.release_prefiller(prefiller_idx, prefiller_score)
|
||||
proxy_state.release_prefiller_kv(prefiller_idx, prefiller_score)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
global global_args
|
||||
global_args = parse_args()
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host=global_args.host, port=global_args.port)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,30 +1,21 @@
|
||||
# Mooncake connector deployment Guide
|
||||
|
||||
## Environmental Dependencies
|
||||
|
||||
* Software:
|
||||
* Python >= 3.9, < 3.12
|
||||
* CANN >= 8.2.rc1
|
||||
* PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
|
||||
* vLLM (same version as vllm-ascend)
|
||||
* mooncake-transfer-engine reference documentation: https://github.com/kvcache-ai/Mooncake/blob/main/doc/zh/ascend_transport.md
|
||||
|
||||
The vllm version must be the same as the main branch of vllm-ascend, for example, 2025/07/30. The version is
|
||||
|
||||
* vllm: v0.10.1
|
||||
* vllm-ascend: v0.10.1rc1
|
||||
* vllm: v0.10.1
|
||||
* vllm-ascend: v0.10.1rc1
|
||||
|
||||
## run
|
||||
|
||||
### 1.Run `prefill` Node
|
||||
|
||||
```
|
||||
```shell
|
||||
bash run_prefill.sh
|
||||
```
|
||||
|
||||
Content of the run_prefill.sh script
|
||||
|
||||
```
|
||||
```shell
|
||||
export HCCL_EXEC_TIMEOUT=204
|
||||
export HCCL_CONNECT_TIMEOUT=120
|
||||
export HCCL_IF_IP=localhost
|
||||
@@ -53,9 +44,7 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
|
||||
"kv_role": "kv_producer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20001",
|
||||
"engine_id": "0",
|
||||
"kv_rank": 0,
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector",
|
||||
"kv_connector_extra_config": {
|
||||
"prefill": {
|
||||
"dp_size": 2,
|
||||
@@ -72,28 +61,27 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
|
||||
`HCCL_EXEC_TIMEOUT`, `HCCL_CONNECT_TIMEOUT`, and `HCCL_IF_IP` are hccl-related configurations.<br>
|
||||
Set `GLOO_SOCKET_IFNAME`, `TP_SOCKET_IFNAME`, and `HCCL_SOCKET_IFNAME` to the corresponding NIC.<br>
|
||||
`ASCEND_RT_VISIBLE_DEVICES` specifies the cards on which the node run resides. The total number of cards equals `dp_size*tp_size`.<br>
|
||||
`/xxxxx/DeepSeek-V2-Lite-Chat` is configured as a model that requires run.<br>
|
||||
`/xxxxx/DeepSeek-V2-Lite-Chat` is configured as a model that requires running.<br>
|
||||
`--host`: indicates the IP address of the node to be started.<br>
|
||||
`--port`: indicates the port to be started, which corresponds to the port in step 4.<br>
|
||||
`--seed`, --max-model-len, and --max-num-batched-tokens model basic configuration. Set this parameter based on the site requirements.<br>
|
||||
`--port`: indicates the port on which the prefill node will listen (e.g., 8100). This port is later referenced in step 3 when configuring the proxy server.<br>
|
||||
`--seed`: `--max-model-len`, and `--max-num-batched-tokens` are part of the model's basic configuration. Set this parameter based on the site requirements.<br>
|
||||
`--tensor-parallel-size`: specifies the TP size.<br>
|
||||
`--data-parallel-size`: indicates the DP size.<br>
|
||||
`--data-parallel-address`: indicates the IP address of the DP. Set this parameter to the IP address of the node.--data-parallel-rpc-port: indicates the RPC port for communication in the DP group.<br>
|
||||
`--trust-remote-code` can load the local model.<br>
|
||||
`--enforce-eager` Turn off the map mode<br>
|
||||
`--gpu-memory-utilization`: Percentage of video memory occupied by the card<br>
|
||||
`--kv-transfer-config`: follow kv_connector, kv_connector_module_path: mooncakeconnect, kv_buffer_device, and run on the NPU card. For kv_role, set kv_producer to the p node, kv_consumer to the d node, kv_parallel_size to 1, and kv_port to the port used by the node. For the p node, set engine_id and kv_rank to 0 and for the d node to 1. Configure the distributed parallel policy for the p and d nodes in the kv_connector_extra_config file based on --tensor-parallel-size and --data-parallel-size.<br>
|
||||
|
||||
`--kv-transfer-config`: follow kv_connector, kv_buffer_device, and run on the NPU card. For kv_role, set kv_producer to the p node, kv_consumer to the d node, kv_parallel_size to 1, and kv_port to the port used by the node. For the p node, set engine_id and kv_rank to 0 and for the d node to 1. Configure the distributed parallel policy for the p and d nodes in the kv_connector_extra_config file based on --tensor-parallel-size and --data-parallel-size.<br>
|
||||
|
||||
### 2. Run `decode` Node
|
||||
|
||||
```
|
||||
```shell
|
||||
bash run_decode.sh
|
||||
```
|
||||
|
||||
Content of the run_decode.sh script
|
||||
|
||||
```
|
||||
```shell
|
||||
export HCCL_EXEC_TIMEOUT=204
|
||||
export HCCL_CONNECT_TIMEOUT=120
|
||||
export HCCL_IF_IP=localhost
|
||||
@@ -122,9 +110,7 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
|
||||
"kv_role": "kv_consumer",
|
||||
"kv_parallel_size": 1,
|
||||
"kv_port": "20002",
|
||||
"engine_id": "1",
|
||||
"kv_rank": 1,
|
||||
"kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector",
|
||||
"kv_connector_extra_config": {
|
||||
"prefill": {
|
||||
"dp_size": 2,
|
||||
@@ -138,28 +124,27 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
|
||||
}'
|
||||
```
|
||||
|
||||
### 3. Start proxy_server. ###
|
||||
### 3. Start proxy_server
|
||||
|
||||
```
|
||||
```shell
|
||||
cd /vllm-ascend/examples/disaggregate_prefill_v1/
|
||||
python load_balance_proxy_server_example.py --host localhost --prefiller-hosts host1 host2 --prefiller-ports 8100 8101 --decoder-hosts host3 host4 --decoder-ports 8200 8201
|
||||
```
|
||||
|
||||
`--host`: indicates the active node. The value of localhost in the curl command delivered in step 5 must be the same as the host. The default port number for starting the service proxy is 8000.<br>
|
||||
`--prefiller-hosts`: Set this parameter to the IP addresses of all p nodes. In the xpyd scenario, add the IP addresses to the end of this configuration item and leave a blank space between the IP addresses.<br>
|
||||
`--prefiller-ports`: Set this parameter to the port number of all p nodes, which is the configuration of the port number for the vllm to start the service in step 3. Write the port number after the configuration in sequence and leave a blank space between the port number and the port number. The sequence must be one-to-one mapping to the IP address of --prefiller-hosts.<br>
|
||||
`--prefiller-ports`: Set this parameter to the port numbers of all prefill (P) nodes, which were defined in step 1 when starting the prefill nodes. Write the port number after the configuration in sequence and leave a blank space between the port number and the port number. The sequence must be one-to-one mapping to the IP address of --prefiller-hosts.<br>
|
||||
`--decoder-hosts`: Set this parameter to the IP addresses of all d nodes. In the xpyd scenario, add the IP addresses to the end of this configuration item and leave a blank space between the IP addresses.<br>
|
||||
`--decoder-ports`: Set this parameter to the port number of all d nodes, which is the configuration of the port number for the vllm to start the service in step 4. Set port to the end of the configuration, and leave a blank space between port and port. The sequence must be one-to-one mapping to the IP address of --decoder-hosts.<br>
|
||||
|
||||
|
||||
### 4. Run Inference
|
||||
|
||||
Set the IP address in the inference file to the actual IP address. Set the model variable to the path of the model. Ensure that the path is the same as that in the shell script.
|
||||
|
||||
```
|
||||
```shell
|
||||
curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
|
||||
"model": "model_path",
|
||||
"prompt": "Given the accelerating impacts of climate change—including rising sea levels, increasing frequency of extreme weather events, loss of biodiversity, and adverse effects on agriculture and human health—there is an urgent need for a robust, globally coordinated response. However, international efforts are complicated by a range of factors: economic disparities between high-income and low-income countries, differing levels of industrialization, varying access to clean energy technologies, and divergent political systems that influence climate policy implementation. In this context, how can global agreements like the Paris Accord be redesigned or strengthened to not only encourage but effectively enforce emission reduction targets? Furthermore, what mechanisms can be introduced to promote fair and transparent technology transfer, provide adequate financial support for climate adaptation in vulnerable regions, and hold nations accountable without exacerbating existing geopolitical tensions or disproportionately burdening those with historically lower emissions?",
|
||||
"max_tokens": 256
|
||||
}'
|
||||
```
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user