vLLM Api Client 学习教程:从零搭建推理服务
vLLM 在大语言模型推理加速领域表现极为出色,核心优势在于 KV 缓存内存近乎零浪费,彻底解决了传统方案中内存碎片导致的管理难题。这一设计显著提升了吞吐量与显存利用率,使 GPU 资源得到高效调度。
以下代码来自 examples/online_serving/api_client.py,完整演示了如何通过 HTTP 请求调用 vLLM 的 API 服务。需要说明的是,该 API 服务器主要用于演示场景和简单性能基准测试;生产环境建议使用 vllm serve 搭配 OpenAI 兼容客户端 API,这才是可靠的部署方案。
# SPDX-License-Identifier: Apache-2.0
"""vLLM API 客户端示例(基于 `vllm.entrypoints.api_server`)
注意:此 API 服务器仅限演示和基础性能基准测试,不适用于生产环境。
生产部署请采用 `vllm serve` 命令配合 OpenAI 客户端 API。
"""
import argparse
import json
from collections.abc import Iterable
import requests
def clear_line(n: int = 1) -> None:
LINE_UP = ' 33[1A'
LINE_CLEAR = 'x1b[2K'
for _ in range(n):
print(LINE_UP, end=LINE_CLEAR, flush=True)
def post_http_request(prompt: str,
api_url: str,
n: int = 1,
stream: bool = False) -> requests.Response:
headers = {"User-Agent": "Test Client"}
pload = {
"prompt": prompt,
"n": n,
"use_beam_search": True,
"temperature": 0.0,
"max_tokens": 16,
"stream": stream,
}
response = requests.post(api_url,
headers=headers,
json=pload,
stream=stream)
return response
def get_streaming_response(response: requests.Response) -> Iterable[list[str]]:
for chunk in response.iter_lines(chunk_size=8192,
decode_unicode=False,
delimiter=b"n"):
if chunk:
data = json.loads(chunk.decode("utf-8"))
output = data["text"]
yield output
def get_response(response: requests.Response) -> list[str]:
data = json.loads(response.content)
output = data["text"]
return output
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--n", type=int, default=4)
parser.add_argument("--prompt", type=str, default="San Francisco is a")
parser.add_argument("--stream", action="store_true")
args = parser.parse_args()
prompt = args.prompt
api_url = f"http://{args.host}:{args.port}/generate"
n = args.n
stream = args.stream
print(f"Prompt: {prompt!r}n", flush=True)
response = post_http_request(prompt, api_url, n, stream)
if stream:
num_printed_lines = 0
for h in get_streaming_response(response):
clear_line(num_printed_lines)
num_printed_lines = 0
for i, line in enumerate(h):
num_printed_lines += 1
print(f"Beam candidate {i}: {line!r}", flush=True)
else:
output = get_response(response)
for i, line in enumerate(output):
print(f"Beam candidate {i}: {line!r}", flush=True)
