426 lines
16 KiB
Python
426 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
import threading
|
|
from collections.abc import Callable
|
|
from contextlib import AbstractContextManager, ExitStack, contextmanager
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from tests.dfx.conftest import (
|
|
create_paired_omni_benchmark_pytest_params,
|
|
create_test_parameter_mapping,
|
|
get_benchmark_params_for_server,
|
|
get_runtime_resource_label,
|
|
is_diffusion_perf_config,
|
|
load_benchmark_configs,
|
|
run_benchmark,
|
|
)
|
|
from tests.helpers.runtime import OmniServer
|
|
|
|
# Optional JSON field ``mark`` is applied as pytest marks via
|
|
# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks":
|
|
# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``).
|
|
|
|
|
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
|
|
|
|
|
def _get_config_file_from_argv() -> str | None:
|
|
"""Read ``--test-config-file`` from ``sys.argv`` at import time so parametrization can use it."""
|
|
import sys
|
|
|
|
for i, arg in enumerate(sys.argv):
|
|
if arg == "--test-config-file" and i + 1 < len(sys.argv):
|
|
return sys.argv[i + 1]
|
|
if arg.startswith("--test-config-file="):
|
|
return arg.split("=", 1)[1]
|
|
return None
|
|
|
|
|
|
_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"
|
|
|
|
CONFIG_FILE_PATH = _get_config_file_from_argv()
|
|
if CONFIG_FILE_PATH is None:
|
|
_all_configs = load_benchmark_configs(config_dir=_PERF_TESTS_DIR)
|
|
BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if not is_diffusion_perf_config(cfg)]
|
|
print(
|
|
f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts/generation case(s) from "
|
|
f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} diffusion; "
|
|
f"use -m to filter, e.g. -m tts)"
|
|
)
|
|
else:
|
|
_loaded = load_benchmark_configs(CONFIG_FILE_PATH)
|
|
BENCHMARK_CONFIGS = [cfg for cfg in _loaded if not is_diffusion_perf_config(cfg)]
|
|
skipped = len(_loaded) - len(BENCHMARK_CONFIGS)
|
|
if skipped:
|
|
print(
|
|
f"--test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts/generation case(s); "
|
|
f"skipped {skipped} remaining diffusion case(s) (chat completions / custom jsonl)"
|
|
)
|
|
|
|
DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy"
|
|
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS)
|
|
paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR)
|
|
|
|
_omni_server_lock = threading.Lock()
|
|
|
|
|
|
class _SingleActiveContext:
|
|
"""Reuse one active context while its configuration key is unchanged."""
|
|
|
|
def __init__(self) -> None:
|
|
self._key: Any = None
|
|
self._stack: ExitStack | None = None
|
|
self._value: Any = None
|
|
|
|
def acquire(self, key: Any, factory: Callable[[], AbstractContextManager[Any]]) -> Any:
|
|
if self._stack is not None and key == self._key:
|
|
return self._value
|
|
|
|
self.close()
|
|
stack = ExitStack()
|
|
value = stack.enter_context(factory())
|
|
self._key = key
|
|
self._stack = stack
|
|
self._value = value
|
|
return value
|
|
|
|
def close(self) -> None:
|
|
stack = self._stack
|
|
self._key = None
|
|
self._stack = None
|
|
self._value = None
|
|
if stack is not None:
|
|
stack.close()
|
|
|
|
|
|
# OmniServer defaults for flags not already present in JSON ``serve_args`` /
|
|
# ``extra_cli_args``. Add new (flag, value) pairs here rather than special-casing.
|
|
_OMNI_DEFAULT_SERVER_ARGS: tuple[tuple[str, str], ...] = (
|
|
("--stage-init-timeout", "600"),
|
|
("--init-timeout", "900"),
|
|
)
|
|
|
|
|
|
def _cli_flag_names(cli_args: tuple[str, ...] | list[str]) -> set[str]:
|
|
"""Return long-option names present in a flat CLI argv list."""
|
|
names: set[str] = set()
|
|
for item in cli_args:
|
|
token = str(item)
|
|
if not token.startswith("--"):
|
|
continue
|
|
names.add(token.split("=", 1)[0])
|
|
return names
|
|
|
|
|
|
def _merge_omni_default_server_args(
|
|
extra_cli_args: tuple[str, ...] | list[str],
|
|
*,
|
|
use_omni: bool,
|
|
defaults: tuple[tuple[str, str], ...] = _OMNI_DEFAULT_SERVER_ARGS,
|
|
) -> list[str]:
|
|
"""Fill Omni defaults for flags not already set in JSON-derived CLI args.
|
|
|
|
JSON ``serve_args`` / ``extra_cli_args`` win; only missing flags are appended.
|
|
"""
|
|
if not use_omni:
|
|
return []
|
|
present = _cli_flag_names(extra_cli_args)
|
|
args: list[str] = []
|
|
for flag, value in defaults:
|
|
if flag not in present:
|
|
args += [flag, value]
|
|
return args
|
|
|
|
|
|
def _omni_server_env() -> dict[str, str]:
|
|
"""Writable video/image storage for ``/v1/videos`` and related generation APIs."""
|
|
result_dir = Path(os.environ.get("BENCHMARK_DIR", "tests/dfx/perf/results"))
|
|
storage_path = Path(os.environ.get("VLLM_OMNI_STORAGE_PATH", str(result_dir / "storage")))
|
|
storage_path.mkdir(parents=True, exist_ok=True)
|
|
return {"VLLM_OMNI_STORAGE_PATH": str(storage_path)}
|
|
|
|
|
|
def _resolve_offline_model(model: str) -> str:
|
|
"""Resolve HF ids / MiniMax env overrides the same way as the diffusion runner."""
|
|
import huggingface_hub
|
|
|
|
from vllm_omni.transformers_utils.repo_utils import hf_api
|
|
|
|
if not model or os.path.isdir(model):
|
|
return model
|
|
|
|
model_env_overrides = {
|
|
"MiniMaxAI/MiniMax-H3": "VLLM_TEST_MINIMAX_H3_MODEL",
|
|
"MiniMaxAI/MiniMax-H3/FL2VA": "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL",
|
|
"MiniMaxAI/MiniMax-H3/Ref2VA": "VLLM_TEST_MINIMAX_H3_REF2VA_MODEL",
|
|
}
|
|
env_name = model_env_overrides.get(model)
|
|
if env_name:
|
|
env_model = os.environ.get(env_name)
|
|
if env_model:
|
|
return env_model
|
|
|
|
parts = model.split("/")
|
|
if len(parts) >= 3:
|
|
repo_id = "/".join(parts[:2])
|
|
subfolder = "/".join(parts[2:])
|
|
snapshot_root = hf_api().snapshot_download(
|
|
repo_id,
|
|
allow_patterns=[f"{subfolder}/**"],
|
|
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
|
)
|
|
return str(Path(snapshot_root) / subfolder)
|
|
|
|
if not huggingface_hub.constants.HF_HUB_OFFLINE:
|
|
return model
|
|
return hf_api().snapshot_download(model, local_files_only=True)
|
|
|
|
|
|
@contextmanager
|
|
def _start_omni_server(server_param):
|
|
test_name, model, stage_config_path, stage_overrides, extra_cli_args, use_omni = server_param
|
|
extra = tuple(extra_cli_args or ())
|
|
model = _resolve_offline_model(model)
|
|
|
|
print(f"Starting OmniServer with test: {test_name}, model: {model}")
|
|
|
|
server_args: list[str] = _merge_omni_default_server_args(extra, use_omni=use_omni)
|
|
# --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py):
|
|
# deploy-config sets the base; stage-overrides are applied on top. Both can be set.
|
|
if stage_config_path:
|
|
server_args = ["--deploy-config", stage_config_path] + server_args
|
|
if stage_overrides:
|
|
server_args = ["--stage-overrides", stage_overrides] + server_args
|
|
if extra:
|
|
server_args = list(extra) + server_args
|
|
with OmniServer(model, server_args, use_omni=use_omni, env_dict=_omni_server_env()) as server:
|
|
server.test_name = test_name
|
|
print("OmniServer started successfully")
|
|
yield server
|
|
print("OmniServer stopping...")
|
|
|
|
print("OmniServer stopped")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def omni_server_context():
|
|
"""Start vLLM-Omni server as a subprocess with actual model weights.
|
|
Reuse it for adjacent benchmark cases with the same server configuration.
|
|
Multi-stage initialization can take 10-20+ minutes.
|
|
"""
|
|
with _omni_server_lock:
|
|
active_context = _SingleActiveContext()
|
|
try:
|
|
yield active_context
|
|
finally:
|
|
active_context.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def omni_server(request, omni_server_context):
|
|
return omni_server_context.acquire(request.param, lambda: _start_omni_server(request.param))
|
|
|
|
|
|
@pytest.fixture
|
|
def benchmark_params(request):
|
|
"""Benchmark parameters fixture; paired with ``omni_server`` via parametrization."""
|
|
test_name, param_index = request.param
|
|
|
|
all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping)
|
|
|
|
if not all_params:
|
|
raise ValueError(f"No benchmark parameters found for test: {test_name}")
|
|
|
|
if param_index >= len(all_params):
|
|
raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}")
|
|
|
|
current = param_index + 1
|
|
total = len(all_params)
|
|
print(f"\n Running benchmark {current}/{total} for {test_name}")
|
|
|
|
return {
|
|
"test_name": test_name,
|
|
"params": all_params[param_index],
|
|
}
|
|
|
|
|
|
def _resolve_num_warmups(params: dict[str, Any], *, default: int) -> int:
|
|
value = params.get("num_warmups")
|
|
if value is None:
|
|
return default
|
|
if not isinstance(value, int) or isinstance(value, bool) or value < 0:
|
|
raise ValueError("num_warmups must be a non-negative integer")
|
|
return value
|
|
|
|
|
|
def assert_result(result, params, num_prompt) -> None:
|
|
assert result["completed"] == num_prompt, "Request failures exist"
|
|
if params.get("dataset_name") == "omniinteract":
|
|
summary = result.get("omniinteract")
|
|
assert isinstance(summary, dict), "OmniInteract summary is missing"
|
|
assert (summary.get("total"), summary.get("success"), summary.get("failed")) == (
|
|
num_prompt,
|
|
num_prompt,
|
|
0,
|
|
), "OmniInteract requests did not all succeed"
|
|
assert summary.get("artifacts_complete") is True, "OmniInteract artifacts are incomplete"
|
|
baseline = params.get("baseline")
|
|
hardware = result.get("Hardware")
|
|
hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
|
|
if isinstance(hardware_baseline, dict) and "mean_tpot_ms" in hardware_baseline:
|
|
num_tpot_samples = result.get("num_tpot_samples")
|
|
mean_tpot_ms = result.get("mean_tpot_ms")
|
|
assert isinstance(num_tpot_samples, int) and not isinstance(num_tpot_samples, bool) and num_tpot_samples > 0, (
|
|
"TPOT baseline is configured, but no measurable TPOT samples were produced"
|
|
)
|
|
assert (
|
|
isinstance(mean_tpot_ms, int | float) and not isinstance(mean_tpot_ms, bool) and math.isfinite(mean_tpot_ms)
|
|
), "TPOT baseline is configured, but mean_tpot_ms is not finite"
|
|
expected_audio_turns = params.get("expected_duplex_audio_turns_per_session")
|
|
if expected_audio_turns is not None:
|
|
session_metrics = result.get("duplex_session_metrics")
|
|
assert isinstance(session_metrics, list), "Duplex session metrics are missing"
|
|
assert len(session_metrics) == num_prompt, (
|
|
f"Expected {num_prompt} duplex session metric rows, got {len(session_metrics)}"
|
|
)
|
|
assert all(
|
|
isinstance(metric, dict) and metric.get("audio_turn_count") == expected_audio_turns
|
|
for metric in session_metrics
|
|
), f"Not every duplex session emitted {expected_audio_turns} audio turns"
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
@pytest.mark.parametrize(
|
|
"omni_server,benchmark_params",
|
|
paired_benchmark_params,
|
|
indirect=["omni_server", "benchmark_params"],
|
|
)
|
|
def test_performance_benchmark(omni_server, benchmark_params):
|
|
test_name = benchmark_params["test_name"]
|
|
params = benchmark_params["params"]
|
|
dataset_name = params.get("dataset_name", "")
|
|
|
|
host = omni_server.host
|
|
port = omni_server.port
|
|
model = omni_server.model
|
|
|
|
print(f"Running benchmark for model: {model}")
|
|
print(f"Benchmark parameters: {benchmark_params}")
|
|
|
|
resource_label = get_runtime_resource_label()
|
|
|
|
def to_list(value, default=None):
|
|
if value is None:
|
|
return [] if default is None else [default]
|
|
return [value] if not isinstance(value, (list, tuple)) else list(value)
|
|
|
|
qps_list = to_list(params.get("request_rate"))
|
|
num_prompt_list = to_list(params.get("num_prompts"))
|
|
max_concurrency_list = to_list(params.get("max_concurrency"))
|
|
|
|
max_len = max(len(qps_list), len(max_concurrency_list))
|
|
if len(num_prompt_list) == 1 and max_len > 1:
|
|
num_prompt_list = num_prompt_list * max_len
|
|
elif max_len == 1 and len(num_prompt_list) > 1:
|
|
if len(qps_list) == 1:
|
|
qps_list = qps_list * len(num_prompt_list)
|
|
if len(max_concurrency_list) == 1:
|
|
max_concurrency_list = max_concurrency_list * len(num_prompt_list)
|
|
max_len = max(len(qps_list), len(max_concurrency_list))
|
|
elif len(num_prompt_list) != max_len and max_len > 0:
|
|
raise ValueError("The number of prompts does not match the QPS or max_concurrency")
|
|
|
|
args = ["--host", host, "--port", str(port)]
|
|
exclude_keys = {
|
|
"request_rate",
|
|
"baseline",
|
|
"num_prompts",
|
|
"max_concurrency",
|
|
"num_warmups",
|
|
"task",
|
|
"enabled",
|
|
"eval_phase",
|
|
"trust_remote_code",
|
|
"expected_duplex_audio_turns_per_session",
|
|
"name",
|
|
"enable_negative_prompt",
|
|
"random_request_config",
|
|
"num_input_images",
|
|
"warmup_requests",
|
|
"warmup_concurrency",
|
|
"warmup_num_inference_steps",
|
|
}
|
|
|
|
param_keys = {str(key).replace("-", "_") for key in params}
|
|
if "model" not in param_keys:
|
|
args.extend(["--model", str(model)])
|
|
|
|
for key, value in params.items():
|
|
if key in exclude_keys or value is None:
|
|
continue
|
|
if key in {"extra_body", "extra-body"} and value == {}:
|
|
continue
|
|
|
|
arg_name = f"--{key.replace('_', '-')}"
|
|
|
|
if isinstance(value, bool) and value:
|
|
args.append(arg_name)
|
|
elif isinstance(value, dict):
|
|
json_str = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
|
|
args.extend([arg_name, json_str])
|
|
elif not isinstance(value, bool):
|
|
args.extend([arg_name, str(value)])
|
|
|
|
for config in BENCHMARK_CONFIGS:
|
|
if config.get("test_name") != test_name:
|
|
continue
|
|
server_params = config.get("server_params") or {}
|
|
if server_params.get("trust_remote_code") or params.get("trust_remote_code"):
|
|
args.append("--trust-remote-code")
|
|
break
|
|
|
|
# QPS / request-rate sweep
|
|
for sweep_index, (qps, num_prompt) in enumerate(zip(qps_list, num_prompt_list)):
|
|
args = args + ["--request-rate", str(qps), "--num-prompts", str(num_prompt)]
|
|
result = run_benchmark(
|
|
args=args,
|
|
test_name=test_name,
|
|
flow=qps,
|
|
dataset_name=dataset_name,
|
|
num_prompt=num_prompt,
|
|
baseline_config=params.get("baseline"),
|
|
sweep_index=sweep_index,
|
|
random_input_len=params.get("random_input_len"),
|
|
random_output_len=params.get("random_output_len"),
|
|
resource_label=resource_label,
|
|
num_warmups=_resolve_num_warmups(params, default=2),
|
|
benchmark_params_name=params.get("name") if isinstance(params.get("name"), str) else None,
|
|
)
|
|
assert_result(result, params, num_prompt)
|
|
|
|
# concurrency test
|
|
for sweep_index, (concurrency, num_prompt) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
|
args = args + ["--max-concurrency", str(concurrency), "--num-prompts", str(num_prompt), "--request-rate", "inf"]
|
|
result = run_benchmark(
|
|
args=args,
|
|
test_name=test_name,
|
|
flow=concurrency,
|
|
dataset_name=dataset_name,
|
|
num_prompt=num_prompt,
|
|
baseline_config=params.get("baseline"),
|
|
sweep_index=sweep_index,
|
|
random_input_len=params.get("random_input_len"),
|
|
random_output_len=params.get("random_output_len"),
|
|
resource_label=resource_label,
|
|
num_warmups=_resolve_num_warmups(params, default=max(2, int(concurrency))),
|
|
benchmark_params_name=params.get("name") if isinstance(params.get("name"), str) else None,
|
|
)
|
|
assert_result(result, params, num_prompt)
|