2510 字13 分钟

LLM 推理性能测试：用 pytest 做自动化基准测试

2026-05-08技术分享#LLM / 性能测试 / pytest / vLLM

在 LLM 推理引擎的开发和优化过程中，性能测试是不可或缺的环节。传统的手工测试方式效率低、可重复性差，难以支撑大规模的性能基准测试需求。本文将介绍如何使用 pytest 7.0+ 构建一套完整的 LLM 推理性能自动化测试框架，实现从配置管理到报告生成的全流程自动化。

为什么选择 pytest？

pytest 是 Python 生态中最流行的测试框架之一，选择它作为性能测试框架有以下几个关键优势：

强大的参数化能力：支持多维度测试参数组合，轻松覆盖不同的推理场景
灵活的标记系统：可以按测试阶段、功能特性、运行平台等维度筛选测试
丰富的插件生态：pytest-html、pytest-xdist 等插件扩展功能
简洁的 fixture 机制：统一管理测试依赖和资源初始化

更重要的是，pytest 的设计理念与性能测试的需求高度契合——可配置、可扩展、可重复。

框架架构设计

整个测试框架分为五个核心模块：

1
├── config/               # YAML 配置管理
2
│   ├── base.yaml        # 基础配置
3
│   ├── ascend.yaml      # Ascend 平台配置
4
│   └── cuda.yaml        # CUDA 平台配置
5
├── tests/               # 测试用例
6
│   ├── conftest.py      # fixture 和装饰器定义
7
│   ├── test_latency.py  # 延迟测试
8
│   └── test_throughput.py # 吞吐量测试
9
├── collectors/          # 结果采集器
10
│   ├── db_exporter.py   # PostgreSQL 导出
11
│   └── html_reporter.py # HTML 报告生成
12
├── utils/               # 工具函数
13
│   └── vllm_client.py   # vLLM UC 集成
14
│   └── platform.py      # 平台适配层
15
└── pytest.ini           # pytest 配置

YAML 配置管理

配置管理是自动化测试的基础。我们使用 YAML 文件分层管理不同平台和场景的配置：

1
model:
2
  name: "Qwen/Qwen2.5-7B-Instruct"
3
  max_length: 4096
4

5
test_params:
6
  in_tokens: [128, 512, 1024, 2048]
7
  out_tokens: [64, 256, 512]
8
  concurrent: [1, 4, 8, 16]
9

10
thresholds:
11
  latency_p99: 500 # ms
12
  throughput_min: 50 # tokens/s
13

14
# config/ascend.yaml
15
extends: base.yaml
16
platform:
17
  type: "ascend"
18
  device: "npu"
19
  count: 8
20
  memory: 64 # GB
21

22
# config/cuda.yaml
23
extends: base.yaml
24
platform:
25
  type: "cuda"
26
  device: "gpu"
27
  count: 4
28
  memory: 80 # GB

通过 extends 机制实现配置继承，避免重复定义。测试启动时根据 --platform 参数自动加载对应配置：

1
import yaml
2
import pytest
3

4
def pytest_addoption(parser):
5
    parser.addoption("--platform", action="store", default="cuda")
6

7
@pytest.fixture(scope="session")
8
def config(request):
9
    platform = request.config.getoption("--platform")
10
    base = yaml.safe_load(open("config/base.yaml"))
11
    platform_cfg = yaml.safe_load(open(f"config/{platform}.yaml"))
12

13
    # 合并配置
14
    merged = {**base, **platform_cfg}
15
    return merged

多维标记系统

pytest 的标记系统让我们可以灵活筛选测试用例。我们定义了三类标记：

1
[pytest]
2
markers =
3
    stage_unit: 单元测试阶段
4
    stage_smoke: 冒烟测试阶段
5
    stage_regression: 回归测试阶段
6
    feature_latency: 延迟测试功能
7
    feature_throughput: 吞吐量测试功能
8
    platform_ascend: Ascend NPU 平台
9
    platform_cuda: CUDA GPU 平台

测试用例中使用多标记组合：

1
@pytest.mark.stage_smoke
2
@pytest.mark.feature_latency
3
@pytest.mark.platform_cuda
4
def test_single_request_latency(config, vllm_client):
5
    """单请求延迟测试 - CUDA 平台冒烟阶段"""
6
    prompt = generate_prompt(config["test_params"]["in_tokens"][0])
7

8
    result = vllm_client.generate(
9
        prompt,
10
        max_tokens=config["test_params"]["out_tokens"][0]
11
    )
12

13
    assert result["latency"] < config["thresholds"]["latency_p99"]

运行时通过 -m 参数筛选：

1
# 只运行冒烟测试
2
pytest -m "stage_smoke"
3

4
# 只运行 CUDA 平台的延迟测试
5
pytest -m "platform_cuda and feature_latency"
6

7
# 排除 Ascend 平台的测试
8
pytest -m "not platform_ascend"

参数化测试

性能测试的核心是覆盖多种参数组合。pytest 的参数化机制完美支持这一需求：

1
@pytest.mark.stage_regression
2
@pytest.mark.feature_throughput
3
@pytest.mark.parametrize("in_tokens", [128, 512, 1024, 2048])
4
@pytest.mark.parametrize("out_tokens", [64, 256, 512])
5
@pytest.mark.parametrize("concurrent", [1, 4, 8, 16])
6
def test_concurrent_throughput(
7
    config, vllm_client, in_tokens, out_tokens, concurrent
8
):
9
    """并发吞吐量测试 - 多参数组合"""
10
    prompts = [
11
        generate_prompt(in_tokens)
12
        for _ in range(concurrent)
13
    ]
14

15
    # 并发请求
16
    results = vllm_client.batch_generate(
17
        prompts,
18
        max_tokens=out_tokens,
19
        concurrent=concurrent
20
    )
21

22
    # 计算吞吐量
23
    total_tokens = sum(r["output_tokens"] for r in results)
24
    total_time = max(r["latency"] for r in results)
25
    throughput = total_tokens / total_time
26

27
    assert throughput >= config["thresholds"]["throughput_min"]

这个测试会自动生成 4 × 3 × 4 = 48 个测试用例，覆盖所有参数组合。每个组合独立运行，结果单独记录。

@export_vars 自动采集

性能测试不仅要验证结果，还要采集过程数据。我们设计了 @export_vars 装饰器，自动采集测试变量并导出到数据库：

1
import functools
2
import psycopg2
3
from datetime import datetime
4

5
export_registry = []
6

7
def export_vars(**var_defs):
8
    """
9
    装饰器：自动采集测试变量
10

11
    用法:
12
        @export_vars(
13
            latency="result['latency']",
14
            tokens="result['output_tokens']",
15
            throughput="throughput"
16
        )
17
        def test_xxx(...):
18
            ...
19
    """
20
    def decorator(func):
21
        @functools.wraps(func)
22
        def wrapper(*args, **kwargs):
23
            # 执行测试
24
            result = func(*args, **kwargs)
25

26
            # 采集变量
27
            exported = {}
28
            for var_name, var_path in var_defs.items():
29
                # 从函数局部变量中提取
30
                value = eval(var_path, func.__globals__, locals())
31
                exported[var_name] = value
32

33
            # 注册到全局采集器
34
            export_registry.append({
35
                "test_name": func.__name__,
36
                "timestamp": datetime.now(),
37
                "platform": kwargs.get("config", {}).get("platform", {}).get("type"),
38
                **exported
39
            })
40

41
            return result
42
        return wrapper
43
    return decorator

测试用例中使用装饰器：

1
@pytest.mark.stage_regression
2
@pytest.mark.feature_latency
3
@pytest.mark.parametrize("in_tokens", [128, 512])
4
@export_vars(
5
    latency="result['latency']",
6
    tokens="result['output_tokens']",
7
    ttft="result['time_to_first_token']",
8
    platform="config['platform']['type']"
9
)
10
def test_request_latency(config, vllm_client, in_tokens):
11
    prompt = generate_prompt(in_tokens)
12
    result = vllm_client.generate(prompt, max_tokens=256)
13
    assert result["latency"] < config["thresholds"]["latency_p99"]
14
    return result  # 装饰器需要返回结果

测试完成后，采集器自动将数据导出到 PostgreSQL：

1
@pytest.fixture(scope="session", autouse=True)
2
def db_exporter(config):
3
    """会话结束时导出所有采集数据到数据库"""
4

5
    yield  # 等待所有测试完成
6

7
    conn = psycopg2.connect(
8
        host=config["db"]["host"],
9
        database=config["db"]["name"],
10
        user=config["db"]["user"],
11
        password=config["db"]["password"]
12
    )
13

14
    cursor = conn.cursor()
15

16
    # 创建表（如果不存在）
17
    cursor.execute("""
18
        CREATE TABLE IF NOT EXISTS perf_results (
19
            test_name VARCHAR(255),
20
            timestamp TIMESTAMP,
21
            platform VARCHAR(50),
22
            latency FLOAT,
23
            tokens INT,
24
            throughput FLOAT,
25
            ttft FLOAT,
26
            in_tokens INT,
27
            out_tokens INT,
28
            concurrent INT
29
        )
30
    """)
31

32
    # 批量插入
33
    for record in export_registry:
34
        cursor.execute("""
35
            INSERT INTO perf_results VALUES (
36
                %(test_name)s, %(timestamp)s, %(platform)s,
37
                %(latency)s, %(tokens)s, %(throughput)s,
38
                %(ttft)s, %(in_tokens)s, %(out_tokens)s, %(concurrent)s
39
            )
40
        """, record)
41

42
    conn.commit()
43
    conn.close()

vLLM UC 集成

我们与 vLLM UC（Unified Client）深度集成，提供统一的推理接口：

1
import requests
2
import asyncio
3
import aiohttp
4

5
class vLLMClient:
6
    """vLLM UC 客户端封装"""
7

8
    def __init__(self, config):
9
        self.base_url = config["vllm"]["api_base"]
10
        self.model = config["model"]["name"]
11

12
    def generate(self, prompt, max_tokens=256):
13
        """单请求推理"""
14
        response = requests.post(
15
            f"{self.base_url}/generate",
16
            json={
17
                "model": self.model,
18
                "prompt": prompt,
19
                "max_tokens": max_tokens,
20
                "stream": False
21
            }
22
        )
23
        data = response.json()
24

25
        return {
26
            "text": data["text"],
27
            "latency": data["latency"],
28
            "output_tokens": data["output_tokens"],
29
            "time_to_first_token": data["ttft"]
30
        }
31

32
    async def batch_generate_async(
33
        self, prompts, max_tokens=256, concurrent=8
34
    ):
35
        """并发推理"""
36
        semaphore = asyncio.Semaphore(concurrent)
37

38
        async def _generate(session, prompt):
39
            async with semaphore:
40
                async with session.post(
41
                    f"{self.base_url}/generate",
42
                    json={
43
                        "model": self.model,
44
                        "prompt": prompt,
45
                        "max_tokens": max_tokens
46
                    }
47
                ) as resp:
48
                    return await resp.json()
49

50
        async with aiohttp.ClientSession() as session:
51
            results = await asyncio.gather(
52
                *[_generate(session, p) for p in prompts]
53
            )
54

55
        return results
56

57
    def batch_generate(self, prompts, max_tokens=256, concurrent=8):
58
        """同步包装"""
59
        return asyncio.run(
60
            self.batch_generate_async(prompts, max_tokens, concurrent)
61
        )

fixture 中初始化客户端：

1
@pytest.fixture(scope="session")
2
def vllm_client(config):
3
    """vLLM UC 客户端"""
4
    client = vLLMClient(config)
5

6
    # 健康检查
7
    try:
8
        requests.get(f"{client.base_url}/health")
9
    except Exception as e:
10
        pytest.fail(f"vLLM service not available: {e}")
11

12
    return client

HTML 报告生成

pytest-html 插件生成可视化报告，但我们进一步定制报告模板：

1
from pytest_html import HTMLReport
2

3
def pytest_configure(config):
4
    """注册自定义报告"""
5
    config.pluginmanager.register(HTMLReport(config))
6

7
class HTMLReport:
8
    """自定义 HTML 报告"""
9

10
    def pytest_sessionfinish(self, session):
11
        # 生成报告时注入自定义数据
12
        self.add_perf_summary(export_registry)
13

14
    def add_perf_summary(self, records):
15
        """添加性能摘要表"""
16
        html = """
17
        <h2>性能测试摘要</h2>
18
        <table>
19
            <tr>
20
                <th>测试名称</th>
21
                <th>平台</th>
22
                <th>平均延迟</th>
23
                <th>吞吐量</th>
24
                <th>P99 TTFT</th>
25
            </tr>
26
        """
27

28
        # 按测试聚合
29
        summary = self.aggregate_results(records)
30

31
        for test_name, data in summary.items():
32
            html += f"""
33
            <tr>
34
                <td>{test_name}</td>
35
                <td>{data['platform']}</td>
36
                <td>{data['avg_latency']:.2f} ms</td>
37
                <td>{data['avg_throughput']:.1f} tokens/s</td>
38
                <td>{data['p99_ttft']:.2f} ms</td>
39
            </tr>
40
            """
41

42
        html += "</table>"
43
        return html

运行测试时指定报告路径：

1
pytest \
2
    --platform cuda \
3
    -m "stage_regression" \
4
    --html=reports/cuda_regression.html \
5
    --self-contained-html

平台适配层

Ascend NPU 和 CUDA GPU 的差异通过适配层统一处理：

1
import os
2

3
class PlatformAdapter:
4
    """平台适配器"""
5

6
    def __init__(self, platform_type):
7
        self.platform = platform_type
8
        self.device_map = {
9
            "cuda": "cuda:0",
10
            "ascend": "npu:0"
11
        }
12

13
    def get_device_count(self):
14
        """获取设备数量"""
15
        if self.platform == "cuda":
16
            import torch
17
            return torch.cuda.device_count()
18
        elif self.platform == "ascend":
19
            return int(os.environ.get("ASCEND_DEVICE_COUNT", 8))
20

21
    def set_device(self, device_id):
22
        """设置当前设备"""
23
        if self.platform == "cuda":
24
            import torch
25
            torch.cuda.set_device(device_id)
26
        elif self.platform == "ascend":
27
            os.environ["ASCEND_DEVICE_ID"] = str(device_id)
28

29
    def get_memory_info(self):
30
        """获取内存信息"""
31
        if self.platform == "cuda":
32
            import torch
33
            total = torch.cuda.get_device_properties(0).total_memory
34
            used = torch.cuda.memory_allocated()
35
            return {"total": total, "used": used}
36
        elif self.platform == "ascend":
37
            # Ascend 内存查询接口
38
            return self._query_ascend_memory()

fixture 中根据平台初始化：

1
@pytest.fixture(scope="session")
2
def platform_adapter(config):
3
    """平台适配器"""
4
    platform_type = config["platform"]["type"]
5
    adapter = PlatformAdapter(platform_type)
6

7
    # 设置设备
8
    adapter.set_device(0)
9

10
    return adapter

完整运行示例

1
# CUDA 平台冒烟测试
2
pytest --platform cuda -m "stage_smoke" --html=reports/smoke.html
3

4
# Ascend 平台回归测试（并行）
5
pytest --platform ascend -m "stage_regression" \
6
    -n 4 --dist=loadscope \
7
    --html=reports/ascend_regression.html
8

9
# 全量测试（生成详细报告）
10
pytest --platform cuda \
11
    --html=reports/full.html \
12
    --self-contained-html \
13
    --json=reports/results.json

总结

这套 pytest 性能测试框架的核心优势：

特性	实现方式	价值
参数化测试	`@pytest.mark.parametrize`	自动覆盖多参数组合
多维标记	自定义 markers	按阶段/功能/平台筛选
变量采集	`@export_vars` 装饰器	自动导出到数据库
配置管理	YAML + extends	分层管理平台配置
平台适配	PlatformAdapter	统一 Ascend/CUDA 接口
报告生成	pytest-html + 定制	可视化性能摘要

通过这套框架，我们可以：

自动化执行 大规模性能基准测试（数百个参数组合）
持久化存储 测试结果到 PostgreSQL，支持历史对比
可视化报告 自动生成 HTML 报告，一目了然
跨平台支持 Ascend NPU 和 CUDA GPU 一套代码适配

最重要的是，这套框架可维护、可扩展。新增测试场景只需添加配置和测试用例，无需修改框架核心。这正是 pytest 设计理念的最佳实践。

参考资料

pytest 官方文档: https://docs.pytest.org/en/stable/
pytest-html 插件: https://pytest-html.readthedocs.io/
vLLM 官方文档: https://vllm.readthedocs.io/
PostgreSQL Python 驱动: https://www.psycopg.org/docs/