LLM API 性能基准测试脚本
概述
本 Python 脚本旨在使用 asyncio
和 aiohttp
对 vLLM(或兼容 OpenAI API 格式的)服务器进行性能基准测试。它的核心目标是通过解决常见的基准测试陷阱,生成更真实、更稳健的性能评估结果。
该脚本通过测试服务器在不同输入提示长度和并发级别下的表现,帮助了解服务器在不同负载条件下的扩展性。
主要特点
此脚本包含多项特性,以确保基准测试结果更具参考价值:
-
通过多样化提示避免缓存 (Cache Avoidance):
- 针对每个测试的输入 Token 长度,脚本会生成一批独特的提示。
- 每个提示都包含一个随机的唯一标识符 (
[会话ID: ...]
),并结合了不同的主题句和随机元素。 - 在并发测试期间,尽可能为不同的并发请求分配不同的提示。
- 这显著降低了命中服务器端 KV 缓存优化的可能性,从而更好地衡量原始生成性能,而非缓存读取速度。
-
控制最小输出长度 (Controlled Output Length):
- 脚本通过增强提示(例如加入:“请针对上述内容提供详细分析和见解...”)明确指示语言模型生成详细、较长的回复。
- 利用
max_tokens
参数,并尝试使用min_tokens
参数(如果 API 端点支持),以鼓励模型产生足够长的输出。 - 这可以防止因模型快速生成极短、无意义回复而导致
Token/秒
指标虚高。测试重点在于衡量有意义的生成任务的吞吐量。
-
可变输入提示长度 (Variable Input Lengths):
- 基准测试可以配置为在一系列不同的输入提示 Token 长度(例如
[10, 100, 500, 1000, 2000]
)上运行。 - 测试不同的输入大小至关重要,因为 LLM 的性能(尤其是首个 Token 生成时间和整体吞吐量)会因所提供的上下文大小而显著变化。这避免了仅测试极短提示(例如 < 50 Tokens)可能带来的误导性结果。
- 基准测试可以配置为在一系列不同的输入提示 Token 长度(例如
-
异步并发请求: 使用
asyncio
和aiohttp
高效处理高并发请求。 -
可配置参数: 通过命令行参数轻松配置 API 端点、模型名称、API 密钥、Token 长度、并发级别和输出 Token 限制。
-
详细结果输出: 以多种格式保存结果:
- 原始 CSV (
benchmark_results.csv
) - 便于比较的透视表 (
benchmark_total_tokens_per_second.csv
,benchmark_avg_tokens_per_second.csv
) - 整体 JSON 摘要 (
benchmark_results.json
) - 用于细粒度分析和调试的单个请求详情 (
requests/
子目录中的 JSON 文件) - 测试配置 (
test_config.json
) - 生成的提示 (
prompts_length_*.txt
)
- 原始 CSV (
-
侧重输出的吞吐量指标: 主要基于生成的输出 Tokens 计算
tokens_per_second
,更清晰地衡量模型的生成速度。同时报告了总吞吐量(所有并发请求)和平均每请求吞吐量。 -
可选性能图表: 如果安装了
matplotlib
,则生成可视化性能与并发关系的图表。 -
分词器灵活性: 如果
tiktoken
可用,则使用它进行精确的 Token 计数;否则回退到基于字符的近似计数。
环境要求
- Python 3.7+
- 所需的 Python 库:
aiohttp
pandas
tiktoken
(可选,强烈推荐用于精确 Token 计数)matplotlib
(可选,用于生成性能图表)
安装
克隆仓库或下载脚本文件。然后安装必要的库:
pip install aiohttp pandas tiktoken matplotlib
或者,不安装可选库:
pip install aiohttp pandas
配置
您可以在脚本顶部直接修改配置参数,或者更灵活地通过命令行参数进行配置。关键参数包括:
API_BASE
: 您的 vLLM 或兼容 API 端点的基础 URL (例如http://localhost:8000/v1
)。MODEL_NAME
: 部署的模型名称 (例如DeepSeek-V3-0324
)。API_KEY
: 您的 API 密钥 (如果服务器需要)。
使用方法
在终端中运行脚本: python your_script_name.py
命令行参数:
--token-lengths
: 要测试的输入 Token 长度列表 (例如--token-lengths 100 500 1000
)。默认:[10, 100, 500, 1000, 2000]
--concurrency-levels
: 要测试的并发级别(同时请求数)列表 (例如--concurrency-levels 1 4 8 16
)。默认:[1, 2, 4, 8, 16, 32]
--max-output-tokens
: 每个请求生成的最大 Token 数。默认:500
--min-output-tokens
: 每个请求期望生成的最小 Token 数(用于提示增强,如果 API 支持也作为参数)。默认:300
--prompts-per-length
: 为每个 Token 长度生成多少个不同的提示。默认:50
--api-base
: 覆盖默认的 API 基础 URL。--model
: 覆盖默认的模型名称。--api-key
: 覆盖默认的 API 密钥。--result-dir
: 指定保存结果的目录 (否则将创建一个带时间戳的目录)。
示例:
python benchmark_vllm.py \
--token-lengths 100 500 1000 2000 \
--concurrency-levels 1 4 8 16 32 64 \
--max-output-tokens 1024 \
--min-output-tokens 500 \
--api-base http://10.10.10.10:8000/v1 \
--model DeepSeek-V3-0324 \
--api-key "your-optional-api-key" \
--result-dir ./my_benchmark_run_1
输出内容
脚本将创建一个结果目录(由 --result-dir
指定或名为 benchmark_results_YYYYMMDD_HHMMSS
)。在此目录中,您将找到:
benchmark_results.csv
: 每个输入长度和并发组合的详细结果。benchmark_total_tokens_per_second.csv
: 显示总输出 Tokens/秒 的透视表。benchmark_avg_tokens_per_second.csv
: 显示平均每请求输出 Tokens/秒 的透视表。benchmark_results.json
: JSON 格式的所有结果数据。test_config.json
: 本次基准测试运行使用的配置参数。tokens_per_second_metrics_info.txt
: 关于 Tokens/秒 指标如何计算的说明。prompts_length_*.txt
: 为每个输入长度生成并用于测试的实际提示文本。requests/
: 一个子目录,包含每个token_length_X
和concurrency_Y
的嵌套文件夹,其中包含每个 API 请求的详细日志(单个request_*.json
文件)。*.png
文件 (如果安装了matplotlib
): 可视化图表:- 总输出 Tokens/秒 vs. 并发数
- 平均每请求输出 Tokens/秒 vs. 并发数
- 成功率 vs. 并发数
关键指标说明
total_tokens_per_second
(总吞吐量): 计算方式为(所有成功的并发请求生成的output_tokens
之和)/(整个并发批次的总耗时)。这代表了服务器在该特定负载下的整体处理能力。avg_tokens_per_second
(平均每请求吞吐量): 计算方式为批次内每个成功请求的 (output_tokens
/elapsed_time
) 的平均值。这代表了在该负载下单个客户端请求所体验到的平均生成速度。- 注意: 默认情况下,这两个指标都侧重于 输出 Tokens,因为这通常反映了用户感知的生成速度。API 的原始响应(保存在单个请求的 JSON 文件
full_response
中)也包含输入 Token 计数 (usage.prompt_tokens
)。
注意事项
- Token 计数的准确性依赖于 API 服务器在响应的
usage
字段中正确返回prompt_tokens
和completion_tokens
。 min_tokens
参数的有效性取决于目标 API 端点是否支持它。提示增强是作为一种备用策略。- 如果未安装
tiktoken
,Token 计数将基于字符长度进行近似估算,这对于非英语文本或代码来说准确性较低。