采用vllm bench serve 压测各类模型,可以了解相关的模型性能,帮助业务更好的去选型模型。
自动化 vLLM 基准测试脚本,脚本分为几个部分
1、vllm bench serve 压测命令
2、支持多并发、多输入/输出长度组合
3、解析 vLLM bench serve 输出(包含 TTFT、TPOT、ITL 等完整性能指标)
4、 生成 CSV 报告

此脚本是python 脚本,是运行在部署模型的机器上。

#!/usr/bin/env python3
"""
说明num-prompts 自动设置为 max-concurrency 的 3 倍,如果需求稳定且大量的数据平均,增大这个值的设置。
"""

import subprocess
import re
import os
import time
from datetime import datetime
import pandas as pd

# ==================== 配置参数(按需修改)====================
VLLM_HOST = "172.16.98.78"
VLLM_PORT = 8008
MODEL_PATH = "/data/"
SERVED_MODEL_NAME = "DeepSeek-V3.2-W8A8"
REQUEST_RATE = 1.0             # 请求速率 (RPS)
TEMPERATURE = 0
TOKENIZER_MODE = "deepseek_v32"
TRUST_REMOTE_CODE = True

# 测试范围
CONCURRENCIES = [1, 2, 4, 8, 16]
INPUT_LENS = [1024, 2048, 4096, 8192, 16384]
OUTPUT_LENS = [1024, 2048, 4096, 8192, 16384]

# 输出目录
OUTPUT_DIR = "./concurrency_test"
os.makedirs(OUTPUT_DIR, exist_ok=True)
# ============================================================

def run_bench(concurrency, input_len, output_len):
    """执行一次 vLLM bench serve 测试,返回解析后的指标字典"""
    # num-prompts 设置为 max-concurrency 的 3 倍
    num_prompts = 3 * concurrency

    cmd = [
        "vllm", "bench", "serve",
        "--backend", "vllm",
        "--dataset-name", "random",
        "--host", VLLM_HOST,
        "--port", str(VLLM_PORT),
        "--model", MODEL_PATH,
        "--served-model-name", SERVED_MODEL_NAME,
        "--num-prompts", str(num_prompts),     # 动态计算:3 * max-concurrency
        "--random-input-len", str(input_len),
        "--random-output-len", str(output_len),
        "--max-concurrency", str(concurrency),
        "--request-rate", str(REQUEST_RATE),
        "--metric-percentiles", "95,99",
        "--temperature", str(TEMPERATURE),
        "--save-result",
        "--result-dir", OUTPUT_DIR,
        "--tokenizer-mode", TOKENIZER_MODE
    ]
    if TRUST_REMOTE_CODE:
        cmd.append("--trust-remote-code")

    print(f"\n{'='*60}")
    print(f"测试组合: 并发={concurrency} | 输入={input_len} | 输出={output_len} | 请求数={num_prompts}")
    print(f"命令: {' '.join(cmd)}")
    print(f"{'='*60}")

    # 执行命令并捕获输出
    try:
        proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
        output = proc.stdout
    except subprocess.CalledProcessError as e:
        print(f"命令执行失败,返回码 {e.returncode}")
        print(f"stderr: {e.stderr}")
        return None

    # 解析关键指标(现在包含 TTFT)
    metrics = {
        "concurrency": concurrency,
        "input_len": input_len,
        "output_len": output_len,
        "num_prompts": num_prompts,
        "Output token throughput (tok/s)": None,
        "Mean TTFT (ms)": None,
        "P95 TTFT (ms)": None,
        "P99 TTFT (ms)": None,
        "Mean TPOT (ms)": None,
        "P95 TPOT (ms)": None,
        "P99 TPOT (ms)": None,
        "Mean ITL (ms)": None,
        "P95 ITL (ms)": None,
        "P99 ITL (ms)": None,
    }

    patterns = {
        "Output token throughput (tok/s)": r"Output token throughput \(tok/s\):\s+([\d\.]+)",
        "Mean TTFT (ms)": r"Mean TTFT \(ms\):\s+([\d\.]+)",
        "P95 TTFT (ms)": r"P95 TTFT \(ms\):\s+([\d\.]+)",
        "P99 TTFT (ms)": r"P99 TTFT \(ms\):\s+([\d\.]+)",
        "Mean TPOT (ms)": r"Mean TPOT \(ms\):\s+([\d\.]+)",
        "P95 TPOT (ms)": r"P95 TPOT \(ms\):\s+([\d\.]+)",
        "P99 TPOT (ms)": r"P99 TPOT \(ms\):\s+([\d\.]+)",
        "Mean ITL (ms)": r"Mean ITL \(ms\):\s+([\d\.]+)",
        "P95 ITL (ms)": r"P95 ITL \(ms\):\s+([\d\.]+)",
        "P99 ITL (ms)": r"P99 ITL \(ms\):\s+([\d\.]+)",
    }

    for key, pattern in patterns.items():
        match = re.search(pattern, output)
        if match:
            metrics[key] = float(match.group(1))
        else:
            print(f"警告: 未找到指标 {key}")

    # 保存原始输出以备调试
    log_file = os.path.join(OUTPUT_DIR, f"log_c{concurrency}_i{input_len}_o{output_len}.txt")
    with open(log_file, "w") as f:
        f.write(output)

    return metrics

def main():
    results = []
    total_tests = len(CONCURRENCIES) * len(INPUT_LENS) * len(OUTPUT_LENS)
    test_idx = 0

    for concurrency in CONCURRENCIES:
        for output_len in OUTPUT_LENS:
            for input_len in INPUT_LENS:
                test_idx += 1
                print(f"\n进度: {test_idx}/{total_tests}")
                metrics = run_bench(concurrency, input_len, output_len)
                if metrics:
                    results.append(metrics)
                time.sleep(1)  # 避免服务过载

    # 生成 DataFrame 并保存为 CSV
    df = pd.DataFrame(results)
    columns = [
        "concurrency", "input_len", "output_len", "num_prompts",
        "Output token throughput (tok/s)",
        "Mean TTFT (ms)", "P95 TTFT (ms)", "P99 TTFT (ms)",
        "Mean TPOT (ms)", "P95 TPOT (ms)", "P99 TPOT (ms)",
        "Mean ITL (ms)", "P95 ITL (ms)", "P99 ITL (ms)"
    ]
    df = df[columns]

    csv_path = os.path.join(OUTPUT_DIR, f"benchmark_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv")
    df.to_csv(csv_path, index=False)

    print(f"\n 所有测试完成!结果已保存至: {csv_path}")

if __name__ == "__main__":
    main()

测试部分结果如下:
在这里插入图片描述

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐