335 lines
15 KiB
Python
335 lines
15 KiB
Python
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2023 The vLLM team.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import filelock
|
|
import huggingface_hub
|
|
import pandas as pd
|
|
import regex as re
|
|
from modelscope import snapshot_download # type: ignore
|
|
|
|
BENCHMARK_HOME = os.getenv("BENCHMARK_HOME", os.path.abspath("./benchmark"))
|
|
DATASET_CONF_DIR = os.path.join(BENCHMARK_HOME, "ais_bench", "benchmark", "configs", "datasets")
|
|
REQUEST_CONF_DIR = os.path.join(BENCHMARK_HOME, "ais_bench", "benchmark", "configs", "models", "vllm_api")
|
|
DATASET_DIR = os.path.join(BENCHMARK_HOME, "ais_bench", "datasets")
|
|
|
|
|
|
class AisbenchRunner:
|
|
RESULT_MSG = {"performance": "Performance Result files located in ", "accuracy": "write csv to "}
|
|
DATASET_RENAME = {"aime2024": "aime", "gsm8k-lite": "gsm8k", "textvqa-lite": "textvqa"}
|
|
|
|
def _run_aisbench_task(self):
|
|
dataset_conf = self.dataset_conf.split("/")[-1]
|
|
if self.task_type == "accuracy":
|
|
aisbench_cmd = ["ais_bench", "--models", f"{self.request_conf}_custom", "--datasets", f"{dataset_conf}"]
|
|
if self.task_type == "performance":
|
|
aisbench_cmd = [
|
|
"ais_bench",
|
|
"--models",
|
|
f"{self.request_conf}_custom",
|
|
"--datasets",
|
|
f"{dataset_conf}_custom",
|
|
"--mode",
|
|
"perf",
|
|
]
|
|
if self.num_prompts:
|
|
aisbench_cmd.extend(["--num-prompts", str(self.num_prompts)])
|
|
self.stdout_file = f"output_{self.task_type}.txt"
|
|
aisbench_cmd = " ".join(aisbench_cmd) + f" --debug > {self.stdout_file} 2>&1 &"
|
|
print(f"running aisbench cmd: {aisbench_cmd}")
|
|
self.proc: subprocess.Popen = subprocess.Popen(aisbench_cmd, shell=True)
|
|
|
|
def __init__(self, model: str, port: int, aisbench_config: dict, host_ip: str = "localhost", verify=True):
|
|
self.model = model
|
|
self.dataset_path = aisbench_config.get("dataset_path_local")
|
|
if not self.dataset_path:
|
|
self.dataset_path = maybe_download_from_modelscope(aisbench_config["dataset_path"], repo_type="dataset")
|
|
self.model_path = aisbench_config.get("model_path")
|
|
if not self.model_path:
|
|
self.model_path = maybe_download_from_modelscope(model)
|
|
assert self.dataset_path is not None and self.model_path is not None, (
|
|
f"Failed to download dataset or model: dataset={self.dataset_path}, model={self.model_path}"
|
|
)
|
|
self.port = port
|
|
self.host_ip = host_ip
|
|
self.task_type = aisbench_config["case_type"]
|
|
self.request_conf = aisbench_config["request_conf"]
|
|
self.dataset_conf = aisbench_config.get("dataset_conf")
|
|
self.num_prompts = aisbench_config.get("num_prompts")
|
|
self.max_out_len = aisbench_config["max_out_len"]
|
|
self.batch_size = aisbench_config["batch_size"]
|
|
self.request_rate = aisbench_config.get("request_rate", 0)
|
|
self.trust_remote_code = aisbench_config.get("trust_remote_code", True)
|
|
self.temperature = aisbench_config.get("temperature")
|
|
self.top_k = aisbench_config.get("top_k")
|
|
self.top_p = aisbench_config.get("top_p")
|
|
self.seed = aisbench_config.get("seed")
|
|
self.min_p = aisbench_config.get("min_p")
|
|
self.presence_penalty = aisbench_config.get("presence_penalty")
|
|
self.repetition_penalty = aisbench_config.get("repetition_penalty")
|
|
self.no_pred = aisbench_config.get("no_pred")
|
|
self.thinking = aisbench_config.get("thinking")
|
|
self.exp_folder = None
|
|
self.result_line = None
|
|
self._init_dataset_conf()
|
|
self._init_request_conf()
|
|
self._run_aisbench_task()
|
|
self._wait_for_task()
|
|
if verify:
|
|
self.baseline = aisbench_config.get("baseline", 1)
|
|
if self.task_type == "accuracy":
|
|
self.threshold = aisbench_config.get("threshold", 1)
|
|
self._accuracy_verify()
|
|
if self.task_type == "performance":
|
|
self.threshold = aisbench_config.get("threshold", 0.97)
|
|
self._performance_verify()
|
|
|
|
def _init_dataset_conf(self):
|
|
if self.task_type == "accuracy":
|
|
dataset_name = os.path.basename(self.dataset_path)
|
|
dataset_rename = self.DATASET_RENAME.get(dataset_name, dataset_name)
|
|
dst_dir = os.path.join(DATASET_DIR, dataset_rename)
|
|
command = ["cp", "-r", self.dataset_path, dst_dir]
|
|
subprocess.call(command)
|
|
if self.task_type == "performance":
|
|
conf_path = os.path.join(DATASET_CONF_DIR, f"{self.dataset_conf}.py")
|
|
if self.dataset_conf.startswith("textvqa"):
|
|
self.dataset_path = os.path.join(self.dataset_path, "textvqa_val.jsonl")
|
|
with open(conf_path, encoding="utf-8") as f:
|
|
content = f.read()
|
|
content = re.sub(r"path=.*", f'path="{self.dataset_path}",', content)
|
|
conf_path_new = os.path.join(DATASET_CONF_DIR, f"{self.dataset_conf}_custom.py")
|
|
with open(conf_path_new, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
|
|
def _init_request_conf(self):
|
|
conf_path = os.path.join(REQUEST_CONF_DIR, f"{self.request_conf}.py")
|
|
with open(conf_path, encoding="utf-8") as f:
|
|
content = f.read()
|
|
content = re.sub(r"model=.*", f'model="{self.model}",', content)
|
|
content = re.sub(r"host_port.*", f"host_port={self.port},", content)
|
|
content = re.sub(r"host_ip.*", f'host_ip="{self.host_ip}",', content)
|
|
content = re.sub(r"max_out_len.*", f"max_out_len={self.max_out_len},", content)
|
|
content = re.sub(r"batch_size.*", f"batch_size={self.batch_size},", content)
|
|
content = re.sub(r"trust_remote_code=.*", f"trust_remote_code={self.trust_remote_code},", content)
|
|
if self.top_p:
|
|
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n top_p={self.top_p},", content)
|
|
if self.top_k:
|
|
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n top_k={self.top_k},", content)
|
|
if self.seed:
|
|
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n seed={self.seed},", content)
|
|
if self.min_p:
|
|
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n min_p={self.min_p},", content)
|
|
if self.presence_penalty:
|
|
content = re.sub(
|
|
r"ignore_eos.*", f"ignore_eos=False,\n presence_penalty={self.presence_penalty},", content
|
|
)
|
|
if self.repetition_penalty:
|
|
content = re.sub(
|
|
r"ignore_eos.*",
|
|
f"ignore_eos=False,\n repetition_penalty={self.repetition_penalty},",
|
|
content,
|
|
)
|
|
if self.thinking:
|
|
field_thinking = 'chat_template_kwargs={"thinking": True}'
|
|
content = re.sub(r"ignore_eos.*", f"ignore_eos=False,\n {field_thinking},", content)
|
|
if self.task_type == "performance":
|
|
content = re.sub(r"path=.*", f'path="{self.model_path}",', content)
|
|
content = re.sub(r"request_rate.*", f"request_rate={self.request_rate},", content)
|
|
content = re.sub(r"temperature.*", "temperature=0,", content)
|
|
content = re.sub(r"ignore_eos.*", "ignore_eos=True,", content)
|
|
if self.task_type == "accuracy":
|
|
content = re.sub(r"temperature.*", "temperature=0.6,", content)
|
|
if self.temperature is not None:
|
|
content = re.sub(r"temperature.*", f"temperature={self.temperature},", content)
|
|
if self.no_pred:
|
|
content = re.sub(r"pred_postprocessor.*", "#pred_postprocessor", content)
|
|
conf_path_new = os.path.join(REQUEST_CONF_DIR, f"{self.request_conf}_custom.py")
|
|
with open(conf_path_new, "w", encoding="utf-8") as f:
|
|
f.write(content)
|
|
print(f"The request config is\n {content}")
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self.proc.terminate()
|
|
try:
|
|
self.proc.wait(8)
|
|
except subprocess.TimeoutExpired:
|
|
# force kill if needed
|
|
self.proc.kill()
|
|
|
|
def _wait_for_exp_folder(self):
|
|
while True:
|
|
line = self._check_runtime_stdout()
|
|
if "Current exp folder: " in line:
|
|
self.exp_folder = re.search(r"Current exp folder: (.*)", line).group(1)
|
|
print(f"Current exp folder: {self.exp_folder}")
|
|
return
|
|
self._check_runtime_error(line)
|
|
|
|
def _wait_for_task(self):
|
|
self._wait_for_exp_folder()
|
|
result_msg = self.RESULT_MSG[self.task_type]
|
|
while True:
|
|
line = self._check_runtime_stdout()
|
|
if result_msg in line:
|
|
print(line)
|
|
self.result_line = line
|
|
return
|
|
self._check_runtime_error(line)
|
|
|
|
def _check_runtime_stdout(self):
|
|
time.sleep(5)
|
|
cmd = f"tail -100 {self.stdout_file}"
|
|
return subprocess.check_output(cmd, shell=True).decode()
|
|
|
|
@staticmethod
|
|
def _check_runtime_error(line):
|
|
if "ERROR" in line:
|
|
print(line)
|
|
error_msg = "Some errors happened to Aisbench runtime"
|
|
raise RuntimeError(error_msg) from None
|
|
|
|
def _get_result_performance(self):
|
|
result_dir = re.search(r"Performance Result files located in (.*)", self.result_line).group(1)[:-1]
|
|
dataset_type = self.dataset_conf.split("/")[0]
|
|
result_csv_file = os.path.join(result_dir, f"{dataset_type}.csv")
|
|
result_json_file = os.path.join(result_dir, f"{dataset_type}.json")
|
|
self.result_csv = pd.read_csv(result_csv_file, index_col=0)
|
|
print("Getting performance results from file: ", result_json_file)
|
|
with open(result_json_file, encoding="utf-8") as f:
|
|
self.result_json = json.load(f)
|
|
self.result = [self.result_csv, self.result_json]
|
|
|
|
def _get_result_accuracy(self):
|
|
acc_file = re.search(r"write csv to (.*)", self.result_line).group(1)
|
|
df = pd.read_csv(acc_file)
|
|
self.result = float(df.iloc[0, -1])
|
|
|
|
def _performance_verify(self):
|
|
self._get_result_performance()
|
|
output_throughput = self.result_json["Output Token Throughput"]["total"].replace("token/s", "")
|
|
assert float(output_throughput) >= self.threshold * self.baseline, (
|
|
"Performance verification failed. "
|
|
f"The current Output Token Throughput is {output_throughput} token/s, "
|
|
f"which is not greater than or equal to {self.threshold} * baseline {self.baseline}."
|
|
)
|
|
|
|
def _accuracy_verify(self):
|
|
self._get_result_accuracy()
|
|
acc_value = self.result
|
|
assert self.baseline - self.threshold <= acc_value <= self.baseline + self.threshold, (
|
|
"Accuracy verification failed. "
|
|
f"The accuracy of {self.dataset_path} is {acc_value}, "
|
|
f"which is not within {self.threshold} relative to baseline {self.baseline}."
|
|
)
|
|
|
|
|
|
def run_aisbench_cases(model, port, aisbench_cases, server_args="", host_ip="localhost"):
|
|
aisbench_results = []
|
|
aisbench_errors = []
|
|
total = sum(1 for c in aisbench_cases if c)
|
|
idx = 0
|
|
for aisbench_case in aisbench_cases:
|
|
if not aisbench_case:
|
|
continue
|
|
idx += 1
|
|
case_name = aisbench_case.get("case_name", "unknown")
|
|
case_type = aisbench_case.get("case_type", "unknown")
|
|
logging.info("=" * 60)
|
|
logging.info("[%d/%d] Starting benchmark: %s (type=%s)", idx, total, case_name, case_type)
|
|
logging.info("=" * 60)
|
|
try:
|
|
with AisbenchRunner(model=model, port=port, host_ip=host_ip, aisbench_config=aisbench_case) as aisbench:
|
|
aisbench_results.append(aisbench.result)
|
|
logging.info("[%d/%d] Finished benchmark: %s", idx, total, case_name)
|
|
except Exception as e:
|
|
aisbench_results.append("")
|
|
aisbench_errors.append([aisbench_case, e])
|
|
logging.error("[%d/%d] Benchmark failed: %s, reason: %s", idx, total, case_name, e)
|
|
print(e)
|
|
for failed_case, error_info in aisbench_errors:
|
|
error_msg = f"The following aisbench case failed: {failed_case}, reason is {error_info}"
|
|
if server_args:
|
|
error_msg += f"\nserver_args are {server_args}"
|
|
logging.error(error_msg)
|
|
assert not aisbench_errors, "some aisbench cases failed, info were shown above."
|
|
return aisbench_results
|
|
|
|
|
|
def get_TTFT(results):
|
|
TTFT = []
|
|
for i in range(len(results)):
|
|
TTFT.append(float(results[i][0].loc["TTFT", "Average"][:-3]))
|
|
return TTFT
|
|
|
|
|
|
temp_dir = tempfile.gettempdir()
|
|
|
|
|
|
def get_lock(model_name_or_path: str | Path, cache_dir: str | None = None):
|
|
lock_dir = cache_dir or temp_dir
|
|
model_name_or_path = str(model_name_or_path)
|
|
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
|
|
model_name = model_name_or_path.replace("/", "-")
|
|
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
|
|
# add hash to avoid conflict with old users' lock files
|
|
lock_file_name = hash_name + model_name + ".lock"
|
|
# mode 0o666 is required for the filelock to be shared across users
|
|
lock = filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666)
|
|
return lock
|
|
|
|
|
|
def maybe_download_from_modelscope(
|
|
model: str,
|
|
repo_type: str = "model",
|
|
revision: str | None = None,
|
|
download_dir: str | None = None,
|
|
ignore_patterns: str | list[str] | None = None,
|
|
allow_patterns: list[str] | str | None = None,
|
|
) -> str:
|
|
"""
|
|
Download model/dataset from ModelScope hub.
|
|
Returns the path to the downloaded model, or None if the model is not
|
|
downloaded from ModelScope.
|
|
"""
|
|
# Use file lock to prevent multiple processes from
|
|
# downloading the same model weights at the same time.
|
|
with get_lock(model, download_dir):
|
|
if not os.path.exists(model):
|
|
model_path = snapshot_download(
|
|
model_id=model,
|
|
repo_type=repo_type,
|
|
cache_dir=download_dir,
|
|
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
|
revision=revision,
|
|
ignore_file_pattern=ignore_patterns,
|
|
allow_patterns=allow_patterns,
|
|
)
|
|
else:
|
|
model_path = model
|
|
return model_path
|