| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import json |
| 6 | import math |
| 7 | import os |
| 8 | from contextlib import contextmanager |
| 9 | from typing import Any |
| 10 | |
| 11 | |
| 12 | def redact_sensitive_namespace( |
| 13 | args: argparse.Namespace, fields: tuple[str, ...] |
| 14 | ) -> argparse.Namespace: |
| 15 | """Return a copy of CLI arguments with selected values redacted.""" |
| 16 | redacted_args = argparse.Namespace(**vars(args)) |
| 17 | for field in fields: |
| 18 | if getattr(redacted_args, field, None) is not None: |
| 19 | setattr(redacted_args, field, "***") |
| 20 | return redacted_args |
| 21 | |
| 22 | |
| 23 | def extract_field( |
| 24 | args: argparse.Namespace, extra_info: dict[str, Any], field_name: str |
| 25 | ) -> str: |
| 26 | if field_name in extra_info: |
| 27 | return extra_info[field_name] |
| 28 | |
| 29 | v: Any = args |
| 30 | # For example, args.compilation_config.mode |
| 31 | for nested_field in field_name.split("."): |
| 32 | if not hasattr(v, nested_field): |
| 33 | return "" |
| 34 | v = getattr(v, nested_field) |
| 35 | return v |
| 36 | |
| 37 | |
| 38 | def use_compile(args: argparse.Namespace, extra_info: dict[str, Any]) -> bool: |
| 39 | """Check if the benchmark is run with torch.compile.""" |
| 40 | return not ( |
| 41 | extract_field(args, extra_info, "compilation_config.mode") == "0" |
| 42 | or "eager" in getattr(args, "output_json", "") |
| 43 | or "eager" in getattr(args, "result_filename", "") |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | def convert_to_pytorch_benchmark_format( |
| 48 | args: argparse.Namespace, metrics: dict[str, list], extra_info: dict[str, Any] |
| 49 | ) -> list: |
| 50 | """Save the benchmark results in the format used by PyTorch OSS benchmark with |
| 51 | on metric per record |
| 52 | https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database |
| 53 | """ |
| 54 | records: list[Any] = [] |
| 55 | if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False): |
| 56 | return records |
| 57 | |
| 58 | for name, benchmark_values in metrics.items(): |
| 59 | if not isinstance(benchmark_values, list): |
| 60 | raise TypeError( |
| 61 | f"benchmark_values for metric '{name}' must be a list, " |
| 62 | f"but got {type(benchmark_values).__name__}" |
| 63 | ) |
| 64 | |
| 65 | record = { |
| 66 | "benchmark": { |
| 67 | "name": "vLLM benchmark", |
| 68 | "extra_info": { |
| 69 | "args": vars(args), |
| 70 | "compilation_config.mode": extract_field( |
| 71 | args, extra_info, "compilation_config.mode" |
| 72 | ), |
| 73 | "optimization_level": extract_field( |
| 74 | args, extra_info, "optimization_level" |
| 75 | ), |
| 76 | # A boolean field used by vLLM benchmark HUD dashboard |
| 77 | "use_compile": use_compile(args, extra_info), |
| 78 | }, |
| 79 | }, |
| 80 | "model": { |
| 81 | "name": args.model, |
| 82 | }, |
| 83 | "metric": { |
| 84 | "name": name, |
| 85 | "benchmark_values": benchmark_values, |
| 86 | "extra_info": extra_info, |
| 87 | }, |
| 88 | } |
| 89 | |
| 90 | tp = record["benchmark"]["extra_info"]["args"].get("tensor_parallel_size") |
| 91 | # Save tensor_parallel_size parameter if it's part of the metadata |
| 92 | if not tp and "tensor_parallel_size" in extra_info: |
| 93 | record["benchmark"]["extra_info"]["args"]["tensor_parallel_size"] = ( |
| 94 | extra_info["tensor_parallel_size"] |
| 95 | ) |
| 96 | |
| 97 | records.append(record) |
| 98 | |
| 99 | return records |
| 100 | |
| 101 | |
| 102 | class InfEncoder(json.JSONEncoder): |
| 103 | def clear_inf(self, o: Any): |
| 104 | if isinstance(o, dict): |
| 105 | return { |
| 106 | str(k) |
| 107 | if not isinstance(k, (str, int, float, bool, type(None))) |
| 108 | else k: self.clear_inf(v) |
| 109 | for k, v in o.items() |
| 110 | } |
| 111 | elif isinstance(o, list): |
| 112 | return [self.clear_inf(v) for v in o] |
| 113 | elif isinstance(o, float) and math.isinf(o): |
| 114 | return "inf" |
| 115 | return o |
| 116 | |
| 117 | def iterencode(self, o: Any, *args, **kwargs) -> Any: |
| 118 | return super().iterencode(self.clear_inf(o), *args, **kwargs) |
| 119 | |
| 120 | |
| 121 | def write_to_json(filename: str, records: list) -> None: |
| 122 | with open(filename, "w") as f: |
| 123 | json.dump( |
| 124 | records, |
| 125 | f, |
| 126 | cls=InfEncoder, |
| 127 | default=lambda o: f"<{type(o).__name__} is not JSON serializable>", |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | @contextmanager |
| 132 | def default_vllm_config(): |
| 133 | """Set a default VllmConfig for cases that directly test CustomOps or pathways |
| 134 | that use get_current_vllm_config() outside of a full engine context. |
| 135 | """ |
| 136 | from vllm.config import VllmConfig, set_current_vllm_config |
| 137 | |
| 138 | with set_current_vllm_config(VllmConfig()): |
| 139 | yield |