|
| 1 | +# Copyright 2025- FlagOS Contributors |
| 2 | +# |
| 3 | +# Permission is hereby granted, free of charge, to any person obtaining a copy |
| 4 | +# of this software and associated documentation files (the "Software"), to deal |
| 5 | +# in the Software without restriction, including without limitation the rights |
| 6 | +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| 7 | +# copies of the Software, and to permit persons to whom the Software is |
| 8 | +# furnished to do so, subject to the following conditions: |
| 9 | +# |
| 10 | +# The above copyright notice and this permission notice shall be included in all |
| 11 | +# copies or substantial portions of the Software. |
| 12 | +# |
| 13 | +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR |
| 14 | +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, |
| 15 | +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE |
| 16 | +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER |
| 17 | +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, |
| 18 | +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE |
| 19 | +# SOFTWARE. |
| 20 | +""" |
| 21 | +vLLM 性能测试脚本 |
| 22 | +
|
| 23 | +功能： |
| 24 | +- 支持用户自定义 input_len、output_len、并发数 |
| 25 | +- 返回 vllm bench serve 的所有指标 |
| 26 | +""" |
| 27 | + |
| 28 | +import subprocess |
| 29 | +import sys |
| 30 | +import re |
| 31 | +import argparse |
| 32 | +from datetime import datetime |
| 33 | +from pathlib import Path |
| 34 | + |
| 35 | +sys.stdout.reconfigure(line_buffering=True) |
| 36 | +sys.stderr.reconfigure(line_buffering=True) |
| 37 | + |
| 38 | +# ============================================================================= |
| 39 | +# 服务配置（按需修改） |
| 40 | +# ============================================================================= |
| 41 | +SERVER_HOST = "127.0.0.1" |
| 42 | +SERVER_PORT = 8000 |
| 43 | +MODEL_NAME = "qwen36" |
| 44 | +TOKENIZER_PATH = "/root/flagrelease/qwen3.6/Qwen3.6-27B" |
| 45 | + |
| 46 | +# ============================================================================= |
| 47 | +# 默认测试参数 |
| 48 | +# ============================================================================= |
| 49 | +DEFAULT_INPUT_LEN = 4096 |
| 50 | +DEFAULT_OUTPUT_LEN = 1024 |
| 51 | +DEFAULT_CONCURRENCY = 64 |
| 52 | +WARMUP_ROUNDS = 4 |
| 53 | +TOTAL_ROUNDS = 10 |
| 54 | +OUTPUT_DIR = Path("./unit_test_output") |
| 55 | + |
| 56 | +# 输出解析正则表达式 |
| 57 | +METRIC_PATTERNS = { |
| 58 | + "Successful requests": r"Successful requests:\s+(\d+)", |
| 59 | + #"Failed requests": r"Failed requests:\s+(\d+)", |
| 60 | + #"Benchmark duration (s)": r"Benchmark duration \(s\):\s+([\d.]+)", |
| 61 | + #"Total input tokens": r"Total input tokens:\s+(\d+)", |
| 62 | + #"Total generated tokens": r"Total generated tokens:\s+(\d+)", |
| 63 | + #"Request throughput (req/s)": r"Request throughput \(req/s\):\s+([\d.]+)", |
| 64 | + "Output token throughput (tok/s)": r"Output token throughput \(tok/s\):\s+([\d.]+)", |
| 65 | + "Total token throughput (tok/s)": r"Total token throughput \(tok/s\):\s+([\d.]+)", |
| 66 | + #"Mean TTFT (ms)": r"Mean TTFT \(ms\):\s+([\d.]+)", |
| 67 | + "Median TTFT (ms)": r"Median TTFT \(ms\):\s+([\d.]+)", |
| 68 | + #"P99 TTFT (ms)": r"P99 TTFT \(ms\):\s+([\d.]+)", |
| 69 | + #"Mean TPOT (ms)": r"Mean TPOT \(ms\):\s+([\d.]+)", |
| 70 | + "Median TPOT (ms)": r"Median TPOT \(ms\):\s+([\d.]+)", |
| 71 | + #"P99 TPOT (ms)": r"P99 TPOT \(ms\):\s+([\d.]+)", |
| 72 | + #"Mean ITL (ms)": r"Mean ITL \(ms\):\s+([\d.]+)", |
| 73 | + "Median ITL (ms)": r"Median ITL \(ms\):\s+([\d.]+)", |
| 74 | + #"P99 ITL (ms)": r"P99 ITL \(ms\):\s+([\d.]+)", |
| 75 | +} |
| 76 | + |
| 77 | + |
| 78 | +def parse_output(output): |
| 79 | + """解析 vllm bench serve 的输出文本，提取所有指标""" |
| 80 | + metrics = {} |
| 81 | + for key, pattern in METRIC_PATTERNS.items(): |
| 82 | + match = re.search(pattern, output) |
| 83 | + if match: |
| 84 | + val = match.group(1) |
| 85 | + metrics[key] = float(val) if "." in val else int(val) |
| 86 | + else: |
| 87 | + metrics[key] = None |
| 88 | + return metrics |
| 89 | + |
| 90 | + |
| 91 | +def build_command(input_len, output_len, concurrency): |
| 92 | + """构建 vllm bench serve 命令""" |
| 93 | + return [ |
| 94 | + "vllm", |
| 95 | + "bench", |
| 96 | + "serve", |
| 97 | + "--host", |
| 98 | + SERVER_HOST, |
| 99 | + "--port", |
| 100 | + str(SERVER_PORT), |
| 101 | + "--model", |
| 102 | + MODEL_NAME, |
| 103 | + "--tokenizer", |
| 104 | + TOKENIZER_PATH, |
| 105 | + "--dataset-name", |
| 106 | + "random", |
| 107 | + "--random-input-len", |
| 108 | + str(input_len), |
| 109 | + "--random-output-len", |
| 110 | + str(output_len), |
| 111 | + "--endpoint", |
| 112 | + "/v1/completions", |
| 113 | + "--ignore-eos", |
| 114 | + "--trust-remote-code", |
| 115 | + "--num-prompts", |
| 116 | + str(concurrency), |
| 117 | + "--max-concurrency", |
| 118 | + str(concurrency), |
| 119 | + "--seed", |
| 120 | + "0", |
| 121 | + ] |
| 122 | + |
| 123 | + |
| 124 | +def run_single_test(round_num, input_len, output_len, concurrency): |
| 125 | + """执行单轮测试并返回所有解析后的指标""" |
| 126 | + #print(f"\n--- 第 {round_num}/{TOTAL_ROUNDS} 轮 ---") |
| 127 | + cmd = build_command(input_len, output_len, concurrency) |
| 128 | + #print(f" 执行命令: {' '.join(cmd)}") |
| 129 | + |
| 130 | + try: |
| 131 | + result = subprocess.run(cmd, capture_output=True, text=True, check=True) |
| 132 | + metrics = parse_output(result.stdout) |
| 133 | + #print(f" 第 {round_num} 轮完成") |
| 134 | + #for key, val in metrics.items(): |
| 135 | + # if val is not None: |
| 136 | + # print(f" {key}: {val}") |
| 137 | + return metrics |
| 138 | + except subprocess.CalledProcessError as e: |
| 139 | + print(f" 错误: 第 {round_num} 轮执行失败") |
| 140 | + print(f" 错误输出: {e.stderr[:20000]}") |
| 141 | + return None |
| 142 | + except FileNotFoundError: |
| 143 | + print(" 错误: vllm 命令未找到，请确认 vllm 已安装且在 PATH 中") |
| 144 | + return None |
| 145 | + except Exception as e: |
| 146 | + print(f" 错误: 第 {round_num} 轮发生未知错误: {e}") |
| 147 | + return None |
| 148 | + |
| 149 | + |
| 150 | +def compute_average(all_round_metrics): |
| 151 | + """对后面轮的数值型指标取平均""" |
| 152 | + last = [m for m in all_round_metrics[WARMUP_ROUNDS:] if m is not None] |
| 153 | + if not last: |
| 154 | + return None |
| 155 | + |
| 156 | + avg = {} |
| 157 | + for key in METRIC_PATTERNS: |
| 158 | + values = [m[key] for m in last if m.get(key) is not None] |
| 159 | + if values: |
| 160 | + avg[key] = sum(values) / len(values) |
| 161 | + else: |
| 162 | + avg[key] = None |
| 163 | + return avg |
| 164 | + |
| 165 | + |
| 166 | +def print_and_save_results(all_round_metrics, avg_metrics, input_len, output_len, concurrency): |
| 167 | + """打印每轮完整结果 + 后4轮平均值，并保存到文件""" |
| 168 | + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) |
| 169 | + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
| 170 | + filename = f"perf_results_in{input_len}_out{output_len}_c{concurrency}_{timestamp}.txt" |
| 171 | + filepath = OUTPUT_DIR / filename |
| 172 | + |
| 173 | + lines = [] |
| 174 | + |
| 175 | + def out(text=""): |
| 176 | + print(text) |
| 177 | + lines.append(text) |
| 178 | + |
| 179 | + #out("=" * 70) |
| 180 | + #out("vLLM 性能测试结果") |
| 181 | + #out(f" 服务: {SERVER_HOST}:{SERVER_PORT}") |
| 182 | + #out(f" 模型: {MODEL_NAME}") |
| 183 | + #out(f" 输入长度: {input_len}, 输出长度: {output_len}, 并发数: {concurrency}") |
| 184 | + #out(f" 总轮数: {TOTAL_ROUNDS} (前{WARMUP_ROUNDS}轮为warmup，此后取平均)") |
| 185 | + #out(f" 执行时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") |
| 186 | + #out("=" * 70) |
| 187 | + |
| 188 | + for i, metrics in enumerate(all_round_metrics): |
| 189 | + #round_num = i + 1 |
| 190 | + #label = "（warmup，不计入平均）" if round_num <= WARMUP_ROUNDS else "" |
| 191 | + #out(f"\n--- 第 {round_num} 轮 {label} ---") |
| 192 | + if metrics is None: |
| 193 | + out(" [FAILED]") |
| 194 | + continue |
| 195 | + #for key, val in metrics.items(): |
| 196 | + # if val is not None: |
| 197 | + # out(f" {key}: {val}") |
| 198 | + |
| 199 | + out("\n" + "=" * 70) |
| 200 | + out("结果平均值") |
| 201 | + out("=" * 70) |
| 202 | + if avg_metrics is None: |
| 203 | + out(" 无有效数据，无法计算平均值") |
| 204 | + else: |
| 205 | + for key, val in avg_metrics.items(): |
| 206 | + if val is not None: |
| 207 | + out(f" {key}: {val:.2f}") |
| 208 | + |
| 209 | + with open(filepath, "w", encoding="utf-8") as f: |
| 210 | + f.write("\n".join(lines) + "\n") |
| 211 | + |
| 212 | + print(f"\n结果已保存至: {filepath}") |
| 213 | + |
| 214 | + |
| 215 | +def main(): |
| 216 | + parser = argparse.ArgumentParser(description="vLLM 性能测试脚本") |
| 217 | + parser.add_argument("--input-len", type=int, default=None, help=f"输入长度 (默认: {DEFAULT_INPUT_LEN})") |
| 218 | + parser.add_argument("--output-len", type=int, default=None, help=f"输出长度 (默认: {DEFAULT_OUTPUT_LEN})") |
| 219 | + parser.add_argument("--concurrency", type=int, default=None, help=f"并发数 (默认: {DEFAULT_CONCURRENCY})") |
| 220 | + parser.add_argument("--dry-run", action="store_true", help="仅打印命令，不执行") |
| 221 | + args = parser.parse_args() |
| 222 | + |
| 223 | + input_len = args.input_len if args.input_len is not None else DEFAULT_INPUT_LEN |
| 224 | + output_len = args.output_len if args.output_len is not None else DEFAULT_OUTPUT_LEN |
| 225 | + concurrency = args.concurrency if args.concurrency is not None else DEFAULT_CONCURRENCY |
| 226 | + |
| 227 | + print("=" * 70) |
| 228 | + print("vLLM 性能测试") |
| 229 | + #print(f" 服务: {SERVER_HOST}:{SERVER_PORT}") |
| 230 | + #print(f" 模型: {MODEL_NAME}") |
| 231 | + print(f" 输入长度: {input_len}, 输出长度: {output_len}, 并发数: {concurrency}") |
| 232 | + print(f" 总轮数: {TOTAL_ROUNDS} (前{WARMUP_ROUNDS}轮warmup，此后取平均)") |
| 233 | + #print("=" * 70) |
| 234 | + |
| 235 | + if args.dry_run: |
| 236 | + print("[DRY RUN MODE]") |
| 237 | + cmd = build_command(input_len, output_len, concurrency) |
| 238 | + print(f" 命令: {' '.join(cmd)}") |
| 239 | + print(f" 将执行 {TOTAL_ROUNDS} 轮") |
| 240 | + print("\n[DRY RUN] 脚本验证完成，未实际执行测试。") |
| 241 | + return |
| 242 | + |
| 243 | + all_round_metrics = [] |
| 244 | + for round_num in range(1, TOTAL_ROUNDS + 1): |
| 245 | + metrics = run_single_test(round_num, input_len, output_len, concurrency) |
| 246 | + if metrics is None: |
| 247 | + sys.exit(1) |
| 248 | + all_round_metrics.append(metrics) |
| 249 | + |
| 250 | + avg_metrics = compute_average(all_round_metrics) |
| 251 | + print_and_save_results(all_round_metrics, avg_metrics, input_len, output_len, concurrency) |
| 252 | + |
| 253 | + |
| 254 | +if __name__ == "__main__": |
| 255 | + main() |
0 commit comments