Implementation:Pytorch Serve Auto Benchmark

Knowledge Sources	Pytorch_Serve
Domains	Benchmarking, DevOps
Last Updated	2026-02-13 18:52 GMT

Overview

Auto_Benchmark is the automated benchmark orchestration module for TorchServe. It reads a YAML configuration file defining models and benchmark parameters, installs TorchServe into a clean environment, and sequentially runs benchmark-ab.py for each model configuration. The module produces aggregated performance metrics and reports across all benchmarked models.

Description

The auto_benchmark.py script provides end-to-end automation of the TorchServe benchmarking workflow. It handles environment setup, TorchServe installation, model registration, benchmark execution, and report generation in a single orchestrated pipeline.

Key Responsibilities

• Configuration Loading: Parses a YAML configuration file using ruamel.yaml to extract model definitions, benchmark parameters, and environment settings
• Environment Setup: Prepares the benchmark environment by installing dependencies and configuring TorchServe via benchmark_env_setup()
• TorchServe Installation: Automates the installation of TorchServe and its dependencies through install_torchserve()
• Benchmark Execution: Iterates over each model configuration and invokes benchmark-ab.py through run_benchmark()
• Report Generation: Collects metrics from individual benchmark runs and aggregates them into a consolidated report

Key Class: BenchmarkConfig

The BenchmarkConfig class (lines 19-122) encapsulates all configuration parameters for a benchmark run. It holds model paths, batch sizes, worker counts, concurrency levels, and other settings parsed from the YAML configuration file.

Usage

from benchmarks.auto_benchmark import BenchmarkConfig, main

Run from the command line:

python benchmarks/auto_benchmark.py --config benchmarks/config.yaml

Code Reference

Source Location

File	Lines	Repository
benchmarks/auto_benchmark.py	L1-327	pytorch/serve

Signature

class BenchmarkConfig:
    """
    Encapsulates benchmark configuration loaded from YAML.

    Attributes:
        model_name (str): Name of the model to benchmark.
        model_url (str): URL or path to the model archive.
        batch_size (int): Number of requests per batch.
        workers (int): Number of TorchServe workers.
        concurrency (int): Number of concurrent benchmark clients.
        requests (int): Total number of requests to send.
        backend (str): Backend engine to use (e.g., 'eager', 'torchscript').
        config (dict): Raw YAML configuration dictionary.
    """

    def __init__(self, config_dict):
        """
        Initialize BenchmarkConfig from a configuration dictionary.

        Args:
            config_dict (dict): Parsed YAML configuration for a single benchmark run.
        """
        ...


def main():
    """
    Main entry point for automated benchmarking.

    Parses command-line arguments, loads YAML config, sets up the
    benchmark environment, installs TorchServe, and runs benchmarks
    for each model configuration.

    Command-line Args:
        --config (str): Path to the YAML configuration file.
        --skip-install (bool): Skip TorchServe installation if already present.

    Returns:
        None
    """
    ...

Import

import argparse
import subprocess
from ruamel.yaml import YAML

from benchmarks.auto_benchmark import BenchmarkConfig
from benchmarks.auto_benchmark import benchmark_env_setup
from benchmarks.auto_benchmark import install_torchserve
from benchmarks.auto_benchmark import run_benchmark
from benchmarks.auto_benchmark import main

I/O Contract

Function / Class	Input	Output	Notes
BenchmarkConfig(config_dict)	config_dict: dict parsed from YAML	BenchmarkConfig instance with model, batch, worker, and concurrency attributes	Lines 19-122
benchmark_env_setup()	None (reads environment variables)	None (configures environment, installs dependencies)	Prepares the benchmarking environment
install_torchserve()	None	None (installs TorchServe via subprocess)	Runs installation scripts for TorchServe and model-archiver
run_benchmark(config)	config: BenchmarkConfig instance	dict with benchmark metrics (throughput, latency, error rate)	Invokes benchmark-ab.py as a subprocess
main()	Command-line args: --config, --skip-install	None (writes reports to filesystem)	Orchestrates the full benchmark pipeline

Configuration File Format

# Example benchmark configuration (config.yaml)
benchmark:
  models:
    - model_name: resnet-18
      url: https://torchserve.pytorch.org/mar_files/resnet-18.mar
      batch_size: 1
      workers: 4
      concurrency: 10
      requests: 1000
    - model_name: vgg16
      url: https://torchserve.pytorch.org/mar_files/vgg16.mar
      batch_size: 4
      workers: 2
      concurrency: 20
      requests: 5000

Usage Examples

Example 1: Running automated benchmarks from the command line

# Run the full benchmark suite with a YAML config
python benchmarks/auto_benchmark.py --config benchmarks/config.yaml

# Skip TorchServe installation (already installed)
python benchmarks/auto_benchmark.py --config benchmarks/config.yaml --skip-install

Example 2: Using BenchmarkConfig programmatically

from ruamel.yaml import YAML
from benchmarks.auto_benchmark import BenchmarkConfig, run_benchmark

yaml = YAML()
with open("benchmarks/config.yaml", "r") as f:
    config_data = yaml.load(f)

for model_config in config_data["benchmark"]["models"]:
    bc = BenchmarkConfig(model_config)
    print(f"Benchmarking {bc.model_name} with batch_size={bc.batch_size}")
    metrics = run_benchmark(bc)
    print(f"  Throughput: {metrics['throughput']} req/s")
    print(f"  Latency p99: {metrics['latency_p99']} ms")

Example 3: Benchmark pipeline execution flow

# The main() function executes the following pipeline:

# 1. Parse command-line arguments
args = argparse.ArgumentParser()
args.add_argument("--config", required=True)
args.add_argument("--skip-install", action="store_true")
parsed = args.parse_args()

# 2. Load YAML configuration
yaml = YAML()
with open(parsed.config, "r") as f:
    config = yaml.load(f)

# 3. Set up benchmark environment
benchmark_env_setup()

# 4. Install TorchServe (unless --skip-install)
if not parsed.skip_install:
    install_torchserve()

# 5. Run benchmarks for each model
for model_cfg in config["benchmark"]["models"]:
    bc = BenchmarkConfig(model_cfg)
    run_benchmark(bc)

# 6. Reports are written to the benchmark output directory

Related Pages

• Principle:Pytorch_Serve_Automated_Benchmarking -- The principle of automated performance benchmarking for TorchServe models
• Implementation:Pytorch_Serve_Benchmark -- JMeter-based benchmarking framework used alongside auto_benchmark