generate_vllm_benchmark_matrix.py

#!/usr/bin/env python

import glob
import json
import logging
import os
from argparse import Action, ArgumentParser, Namespace
from logging import warning
from typing import Any, Dict, List, Optional


logging.basicConfig(level=logging.INFO)
# Those are H100 runners from https://github.com/meta-pytorch/pytorch-gha-infra/blob/main/multi-tenant/inventory/manual_inventory
# while ROCm runner are provided by AMD
TP_TO_RUNNER_MAPPING = {
    1: [
        "mt-l-x86iamx-22-225-h100",
        "linux.rocm.gpu.gfx942.1",
        "linux.24xl.spr-metal",
        "linux.24xl.gnr",
        "linux.arm64.m8g.4xlarge",
        "linux.dgx.b200",
        "linux.hpu.gaudi3.8",
    ],
    # NB: There is no 2xH100 runner at the momement, so let's use the next one
    # in the list here which is 4xH100
    2: [
        "mt-l-x86iamx-88-900-h100-4",
        "linux.rocm.gpu.gfx942.2",
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
        "linux.dgx.b200.8",
    ],
    4: [
        "mt-l-x86iamx-88-900-h100-4",
        "linux.rocm.gpu.gfx942.4",
        "linux.hpu.gaudi3.8",
        "linux.dgx.b200.8",
    ],
    8: [
        "mt-l-bx86iamx-176-1800-h100-8",
        "linux.rocm.gpu.gfx942.8",
        "linux.dgx.b200.8",
        "linux.hpu.gaudi3.8",
    ],
}

# This mapping is needed to find out the platform of the runner
RUNNER_TO_PLATFORM_MAPPING = {
    "mt-l-x86iavx512-11-125-a100": "cuda",
    "mt-l-x86iamx-22-225-h100": "cuda",
    "mt-l-x86iamx-88-900-h100-4": "cuda",
    "mt-l-bx86iamx-176-1800-h100-8": "cuda",
    "linux.dgx.b200": "cuda",
    "linux.dgx.b200.8": "cuda",
    "linux.rocm.gpu.gfx942.1": "rocm",
    "linux.rocm.gpu.gfx942.2": "rocm",
    "linux.rocm.gpu.gfx942.4": "rocm",
    "linux.rocm.gpu.gfx942.8": "rocm",
    "linux.24xl.spr-metal": "cpu",
    "linux.24xl.gnr": "cpu",
    "linux.arm64.m8g.4xlarge": "arm64-cpu",
    "linux.hpu.gaudi3.8": "hpu",
}

# All the different names vLLM uses to refer to their benchmark configs
VLLM_BENCHMARK_CONFIGS_PARAMETER = set(
    [
        "parameters",
        "server_parameters",
        "common_parameters",
    ]
)

# Model and runner skip logic, for example, just need to run DeepSeek on b200
# and not h100. This also serves as a knob to tune CI behavior. TODO (huydhn):
# Figure out how to set this in the JSON benchmark configuration instead
PLATFORM_SKIPS = {
    "meta-llama/Llama-3.1-8B-Instruct": [
        "linux.24xl.spr-metal",  # Timed out
    ],
    # Already been covered in both A100 and H100
    "meta-llama/Meta-Llama-3.1-8B-Instruct": [
        "linux.dgx.b200",
    ],
    "meta-llama/Meta-Llama-3.1-70B-Instruct": [
        "linux.dgx.b200",
        "linux.rocm.gpu.gfx942",  # TODO: Fail on ROCm
        "linux.24xl.gnr",
    ],
    "mistralai/Mixtral-8x7B-Instruct-v0.1": [
        "linux.dgx.b200",
    ],
    "Qwen/Qwen3-8B": [
        "linux.dgx.b200",
        "linux.hpu.gaudi3.8",
    ],
    "google/gemma-3-4b-it": [
        "linux.dgx.b200",
        "linux.rocm.gpu.gfx942",  # TODO: Fail on ROCm
        "linux.hpu.gaudi3.8",
    ],
    # Run some bigger models on B200 to share the load
    "Qwen/Qwen3-30B-A3B": [
        "a100",
        "h100",
        "linux.rocm.gpu.gfx942",  # TODO: Fail on ROCm
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    "google/gemma-3-27b-it": [
        "a100",
        "h100",
        "linux.rocm.gpu.gfx942",  # TODO (huydhn): Fail on ROCm
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    "meta-llama/Llama-4-Scout-17B-16E-Instruct": [
        "a100",
        "h100",
        "linux.rocm.gpu.gfx942",  # TODO: Fail on ROCm
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": [
        "a100",
        "h100",
        "linux.rocm.gpu.gfx942",  # TODO (huydhn): Hang on ROCm
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    # Run gpt-oss on both H100 and B200
    "openai/gpt-oss-20b": [
        "a100",
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    "openai/gpt-oss-120b": [
        "a100",
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    # Deepseek can only run on B200
    "deepseek-ai/DeepSeek-V3.1": [
        "a100",
        "h100",
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    "deepseek-ai/DeepSeek-V3.2": [
        "a100",
        "h100",
        "linux.24xl.gnr",
        "linux.hpu.gaudi3.8",
    ],
    "deepseek-ai/DeepSeek-R1": [
        "a100",
        "linux.24xl.gnr",
        "h100",
        "linux.hpu.gaudi3.8",
    ],
}
# Lower case all the model names for consistency
PLATFORM_SKIPS = {k.lower(): v for k, v in PLATFORM_SKIPS.items()}


class ValidateDir(Action):
    def __call__(
        self,
        parser: ArgumentParser,
        namespace: Namespace,
        values: Any,
        option_string: Optional[str] = None,
    ) -> None:
        if os.path.isdir(values):
            setattr(namespace, self.dest, values)
            return

        parser.error(f"{values} is not a valid directory")


def parse_args() -> Any:
    parser = ArgumentParser("Generate vLLM benchmark CI matrix")

    parser.add_argument(
        "--benchmark-configs-dir",
        type=str,
        default="vllm-benchmarks/benchmarks",
        action=ValidateDir,
        help="the directory contains vLLM benchmark configs",
        required=True,
    )
    parser.add_argument(
        "--models",
        type=str,
        default="",
        help="the comma-separated list of models to benchmark",
    )
    parser.add_argument(
        "--runners",
        type=str,
        default="",
        help="the comma-separated list of runners to run the benchmark",
        required=True,
    )

    return parser.parse_args()


def set_output(name: str, val: Any) -> None:
    """
    Set the output value to be used by other GitHub jobs.

    Args:
        name (str): The name of the output variable.
        val (Any): The value to set for the output variable.

    Example:
        set_output("benchmark_matrix", {"include": [...]})
    """
    github_output = os.getenv("GITHUB_OUTPUT")

    if not github_output:
        print(f"::set-output name={name}::{val}")
        return

    with open(github_output, "a") as env:
        env.write(f"{name}={val}\n")


def generate_benchmark_matrix(
    benchmark_configs_dir: str, models: List[str], runners: List[str]
) -> Dict[str, Any]:
    """
    Parse all the JSON files in vLLM benchmark configs directory to get the
    model name and tensor parallel size (aka number of GPUs, CPU NUMA nodes - Intel
    or CPUs - ARM)"""
    benchmark_matrix: Dict[str, Any] = {
        "include": [],
    }

    platforms = set()
    if not runners:
        use_all_runners = True
        platforms = set(v for v in RUNNER_TO_PLATFORM_MAPPING.values())
    else:
        use_all_runners = False
        for k, v in RUNNER_TO_PLATFORM_MAPPING.items():
            for r in runners:
                if r.lower() in k:
                    platforms.add(v)

    # Gather all possible benchmarks
    for platform in sorted(platforms):
        selected_models = []
        # Only need to parse serving config because it has all the models and
        # their tensor_parallel_size field. The latter is used to find the runner
        # with the right capacity
        for file in glob.glob(f"{benchmark_configs_dir}/{platform}/*serving*.json"):
            with open(file) as f:
                try:
                    configs = json.load(f)
                except json.JSONDecodeError as e:
                    warning(f"Fail to load {file}: {e}")
                    continue

            for config in configs:
                param = list(VLLM_BENCHMARK_CONFIGS_PARAMETER & set(config.keys()))
                assert len(param) == 1

                benchmark_config = config[param[0]]
                if "model" not in benchmark_config:
                    warning(f"Model name is not set in {benchmark_config}, skipping...")
                    continue
                model = benchmark_config["model"].lower()

                # Dedup
                if model in selected_models:
                    continue
                # and only choose the selected model
                if models and model not in models:
                    continue
                selected_models.append(model)

                if "tensor_parallel_size" in benchmark_config:
                    tp = benchmark_config["tensor_parallel_size"]
                elif "tp" in benchmark_config:
                    tp = benchmark_config["tp"]
                else:
                    tp = 8
                assert tp in TP_TO_RUNNER_MAPPING

                for runner in TP_TO_RUNNER_MAPPING[tp]:
                    # Wrong platform
                    if (
                        runner not in RUNNER_TO_PLATFORM_MAPPING
                        or RUNNER_TO_PLATFORM_MAPPING[runner] != platform
                    ):
                        continue

                    found_runner = any([r and r.lower() in runner for r in runners])
                    if not found_runner and not use_all_runners:
                        continue

                    # Check the skip logic
                    if model in PLATFORM_SKIPS and any(
                        [r in runner for r in PLATFORM_SKIPS[model]]
                    ):
                        continue

                    benchmark_matrix["include"].append(
                        {
                            "runner": runner,
                            # I opt to return a comma-separated list of models here
                            # so that we could run multiple models on the same runner
                            "models": model,
                        }
                    )

    return benchmark_matrix


def main() -> None:
    args = parse_args()
    models = [m.strip().lower() for m in args.models.split(",") if m.strip()]
    runners = [m.strip().lower() for m in args.runners.split(",") if m.strip()]
    benchmark_matrix = generate_benchmark_matrix(
        args.benchmark_configs_dir,
        models,
        runners,
    )
    print(benchmark_matrix)
    set_output("benchmark_matrix", benchmark_matrix)


if __name__ == "__main__":
    main()