基准测试和工作负载生成器

基准测试和工作负载生成器#

AIBrix 基准测试包含以下组件

  • 数据集（提示）生成

  • 工作负载生成

  • 基准测试客户端

  • 基准测试场景（正在进行中）

下图显示了 AIBrix 基准测试的端到端步骤。我们的组件在此图中以绿色矩形突出显示。

benchmark-component-doc

AIBrix 基准测试框架的核心是围绕一个清晰解耦的架构构建的：数据集生成、工作负载塑形和基准测试执行。每个组件都可以独立定制，这使得插入您自己的提示日志、流量跟踪或实验工作负载变得容易——无论您是在开发新的模型部署、扩展策略还是运行时优化。

端到端运行 AIBrix 基准测试#

注意

基准测试脚本 benchmark.py 执行所有步骤，直至 AIBrix 工作负载格式，并触发基准测试客户端，而无需为不同场景设置基准测试环境。它假设 AIBrix 已设置并期望一个完全响应的端点。

有关详细用法和如何配置参数，请参阅此 README 以获取深入解释。

首先，确保您已配置 API 密钥和端点，如下所示

export API_KEY="<your_api_key>"
kubectl -n envoy-gateway-system port-forward service/envoy-aibrix-system-aibrix-eg-903790dc 8888:80 &

要使用默认设置运行所有步骤，请尝试

python benchmark.py --stage all --config config.yaml

如果脚本成功完成，您应该会看到类似以下内容的输出

INFO:root:Running analysis with args: Namespace(trace='./output/client_output/output.jsonl', output='./output/trace_analysis', goodput_target='tpot:0.5')
WARNING:root:End-to-End Latency (s) Statistics: Average = 2.9710, Median = 3.0138, 99th Percentile = 3.3668
WARNING:root:Throughput (per request, toks/s) Statistics: Average = 43.2928, Median = 42.4732, 99th Percentile = 47.2407
WARNING:root:Tokens per Second Statistics: Average = 169.2524, Median = 165.2006, 99th Percentile = 259.2597
WARNING:root:Request Prompt Tokens Statistics: Average = 372.7500, Median = 371.0000, 99th Percentile = 670.6400
WARNING:root:Request Output Tokens Statistics: Average = 128.0000, Median = 128.0000, 99th Percentile = 128.0000
WARNING:root:Request Total Tokens Statistics: Average = 500.7500, Median = 499.0000, 99th Percentile = 798.6400
WARNING:root:Time to First Token (TTFT) Statistics: Average = 0.3053, Median = 0.2678, 99th Percentile = 0.5994
WARNING:root:Time per Output Token (TPOT) Statistics: Average = 0.0208, Median = 0.0216, 99th Percentile = 0.0217
WARNING:root:Errors Statistics: Average = 0.0000, Median = 0.0000, 99th Percentile = 0.0000
WARNING:root:Goodput (reqs/s) 1.0000
WARNING:root:Total Duration (s): 23.778891624999233
WARNING:root:Total tokens generated (toks): 4006
WARNING:root:Throughput (end-to-end, toks/s): 168.46874375710644
INFO:root:========== Benchmark Completed ==========

注意

所有配置文件都应在 .yaml 文件中指定。您可以在此处找到示例配置文件。

要覆盖默认配置，请使用 –override 标志。例如，要在数据集生成阶段覆盖目标模型和会话数，请使用

python benchmark.py --stage all \
    --config config.yaml \
    --override target_model="deepseek-llm-7b-chat" \
    --override dataset_configs.synthetic_multiturn.num_sessions=10

运行数据集生成器#

benchmark-component-dataset-generator

AIBrix 数据集生成器的目标是生成提示数据集或将现有数据集转换为工作负载生成器可以从中采样的格式。AIBrix 数据集生成器生成遵循特定应用程序模式（即缓存共享）的合成提示，或将时间序列跟踪（例如，ShareGPT 等开源 LLM 跟踪）转换为标准数据集格式。合成数据集需要采用以下两种格式之一

  • 纯格式（无会话）

{"prompt": "XXXX"}
{"prompt": "YYYY"}

  • 会话格式

{"session_id": 0, "prompts": ["XXX", "YYY"]}
{"session_id": 1, "prompts": ["AAA", "BBB", "CCC"]}

数据集生成器要么生成提示数据集，要么转换属于上述两种格式之一的现有数据集。

要运行数据集生成，请执行

python benchmark.py --stage dataset --config config.yaml

目前，我们支持四种类型的数据集

1. 受控合成共享 - 此类型允许用户生成一个缓存共享的纯格式数据集，其中包含受控的提示令牌长度和受控的前缀共享长度，以及受控数量的前缀（即会话）。要调整提示令牌长度和共享长度，请在配置文件中的`dataset_configs.synthetic_shared`下设置配置变量。

2. 多轮合成 - 多轮合成数据生成会生成一个会话格式数据集。每个会话 ID 都映射到每个会话的受控提示数量和受控提示长度。这些变量位于配置文件中的`dataset_configs.synthetic_multiturn`下。

3. ShareGPT - 此生成类型将 ShareGPT 数据集转换为包含 session_id、提示和完成的会话格式数据集。配置变量位于配置文件中的 dataset_configs.sharegpt 下。

4. 客户端跟踪 - 此生成类型将客户端输出转换为纯格式数据集。配置变量位于配置文件中的 dataset_configs.client_trace 下。

前两种类型生成合成提示，而后两种类型转换外部数据源或基准测试数据。

要设置要生成的数据集类型，请将配置文件中的环境变量 prompt_type 设置为以下值之一：`synthetic_multiturn`、`synthetic_shared`、`sharegpt`、`client_trace`。

有关数据集生成器的详细信息，请查看 dataset_generator 目录。

运行工作负载生成器#

benchmark-component-workload-generator

AIBrix 的工作负载生成器旨在执行工作负载塑形。工作负载生成器指定基准测试客户端要分派的时间和请求。工作负载生成器接受跟踪/指标文件（其中指定了时间、请求，或指定了 QPS/输入/输出量）或用户指定的静态或动态负载模式。工作负载生成器将根据数据集生成器部分中讨论的两种格式之一的数据集对工作负载进行采样。生成器目前支持四种类型的工作负载

1. “constant” 和 “synthetic” 工作负载类型

  • 工作负载生成器可以生成两种类型的合成负载模式，具有多个可以手动调整的工作负载配置（例如，流量/QPS 分布、输入请求令牌长度分布、输出令牌长度分布、最大并发会话数等）

    • 恒定负载 (constant)：平均负载（QPS/输入长度/输出长度）保持恒定，具有可控波动。配置变量位于配置文件中的 workload_configs.constant 下。

    • 合成波动负载 (synthetic)：负载（QPS/输入长度/输出长度）根据可配置参数波动。配置变量位于配置文件中的 workload_configs.synthetic 下。

2. “stat” 工作负载类型

  • 对于指标文件（例如，从 Grafana 面板导出的 .csv 文件），工作负载生成器将生成遵循文件中指定的收集到的时间序列指标的 QPS/输入长度/输出长度分布。工作负载中使用的实际提示将基于上一节中生成的合成数据集之一。配置变量位于配置文件中的 workload_configs.stat 下。

3. “azure” 工作负载类型

  • 对于跟踪（例如，Azure LLM 跟踪），提供了与请求相关的请求和时间戳，工作负载生成器将生成一个简单地根据时间戳重放请求的工作负载。配置变量位于配置文件中的 workload_configs.azure 下。

4. “mooncake” 工作负载类型

  • 对于 Mooncake LLM 跟踪，提供了请求输入/输出、缓存块 ID 以及与请求相关的时间戳，工作负载生成器将生成一个模拟遵循相同流量模式的请求的工作负载。配置变量位于配置文件中的 workload_configs.mooncake 下。

工作负载生成器可以通过以下方式运行

python benchmark.py --stage workload --config config.yaml

工作负载生成器将在 output/workload 目录中生成一个工作负载文件。该文件将如下所示

{
    "timestamp": 19,
    "requests":
    [
        {
            "prompt": "I need to understand data science ...",
            "prompt_length": 101,
            "output_length": null,
            "session_id": 0
        },
        {
            "prompt": "...",
            "prompt_length": "...",
            "output_length": "...",
            "session_id": "..."
        }
    ]
}

要选择不同的工作负载类型，请将配置文件中的环境变量 workload_type 设置为以下值之一：`constant`、`synthetic`、`stat`、`azure`、`mooncake`。

有关工作负载生成器的详细信息，请查看 workload_generator 目录。

运行基准测试客户端#

benchmark-component-workload-generator

基准测试客户端支持批量和流式传输模式。流式传输模式支持请求内指标，如 TTFT/TPOT。通过 config.yaml 或命令行参数配置端点和目标模型。

基准测试客户端可以使用以下方式运行

python benchmark.py --stage client --config config.yaml

运行分析#

使用以下方式对基准测试结果运行分析

python benchmark.py --stage analysis --config config.yaml

通过 config.yaml 或命令行参数配置路径和性能目标。