"""
Performance benchmark CI runner for diffusion models.

This runner separates two concepts:

1. ``server_type``: how the serving process is started.
   Currently only ``vllm-omni`` is supported here.
2. ``benchmark_endpoint``: which serving API the benchmark client calls.
   Examples: ``/v1/chat/completions`` and ``/v1/videos``.

A config JSON file may be passed via --test-config-file. If omitted, every ``*.json`` under
``tests/dfx/perf/tests/`` is loaded and pytest ``-m`` filters by each case's ``mark``:
  pytest run_diffusion_benchmark.py -m "diffusion"
  pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json

Optional JSON field ``mark`` is applied as pytest marks on that case via
``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``).

All benchmark results are written under BENCHMARK_RESULT_DIR (override via the
DIFFUSION_BENCHMARK_DIR environment variable). Each source JSON file gets one
aggregated ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` (JSON array
of all runs from cases in that file). Bulk load without ``--test-config-file`` uses
the same per-file aggregation; ``-m`` only selects which cases run.
"""

import json
import os
import socket
import subprocess
import sys
import tempfile
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any, cast

import psutil
import pytest

from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint
from tests.dfx.conftest import (
    create_paired_benchmark_pytest_params,
    get_runtime_resource_label,
    hardware_json_value,
    is_diffusion_perf_config,
    resolve_baseline_for_sweep,
    resolve_pytest_marks,
    resource_label_for_filename,
)
from tests.helpers.runtime import get_open_port

os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")


# ---------------------------------------------------------------------------
# Inline field processing
# ---------------------------------------------------------------------------
def _process_inline_fields(obj: Any, parent_key: str = "") -> None:
    """Recursively process '*-inline' fields into temp files."""

    if isinstance(obj, list):
        for item in obj:
            _process_inline_fields(item, parent_key)
        return

    if not isinstance(obj, dict):
        return

    import atexit

    import yaml

    for key in list(obj.keys()):
        value = obj[key]

        if not key.endswith("-inline"):
            _process_inline_fields(value, key)
            continue

        base_key = key[:-7]
        full_key = f"{parent_key}.{key}" if parent_key else key

        try:
            if not isinstance(value, dict):
                raise ValueError("must be a dict")

            file_type = value.get("type")
            content = value.get("content")

            if file_type not in {"yaml", "jsonl"}:
                raise ValueError(f"invalid type: {file_type}")

            fd, path = tempfile.mkstemp(
                suffix=f".{file_type}",
                prefix=f"{base_key}_",
            )

            atexit.register(Path(path).unlink, missing_ok=True)

            with os.fdopen(fd, "w", encoding="utf-8") as f:
                if file_type == "jsonl":
                    items = content if isinstance(content, list) else [content]
                    f.writelines(json.dumps(x, ensure_ascii=False) + "\n" for x in items)
                else:
                    yaml.dump(
                        content,
                        f,
                        allow_unicode=True,
                        sort_keys=False,
                        indent=2,
                    )

            obj[base_key] = path
            del obj[key]

        except Exception as e:
            print(f"Warning: failed processing '{full_key}': {e}")


# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

_DEFAULT_RESULT_DIR = Path(__file__).parent.parent / "results"
BENCHMARK_RESULT_DIR = Path(os.environ.get("DIFFUSION_BENCHMARK_DIR", str(_DEFAULT_RESULT_DIR)))

BENCHMARK_SCRIPT = str(
    Path(__file__).parent.parent.parent.parent.parent / "benchmarks" / "diffusion" / "diffusion_benchmark_serving.py"
)

# Single aggregated result file for the entire benchmark session.
# Populated lazily after CONFIG_FILE_PATH is resolved.
_SESSION_TIMESTAMP = datetime.now().strftime("%Y%m%d-%H%M%S")
_RESULT_LOCK = threading.Lock()
_BRANCHPOINT_COMMIT_SHA: str | None = None
DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json"


_DIFFUSION_SOURCE_CONFIG_KEY = "_source_config_file"
_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests"


def _get_config_file_from_argv() -> str | None:
    """Read --test-config-file from sys.argv at import time so pytest parametrize can use it.

    pytest_addoption (below) registers the same flag so pytest does not reject it.
    Supports both ``--test-config-file path`` and ``--test-config-file=path`` forms.
    Returns None if the flag is not present; callers must handle the missing case.
    """
    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


CONFIG_FILE_PATH = _get_config_file_from_argv()

# ---------------------------------------------------------------------------
# Config loading
# ---------------------------------------------------------------------------


def _resolve_refs(configs: list[dict[str, Any]], config_dir: Path) -> list[dict[str, Any]]:
    """Resolve {"$ref": "filename.json"} in benchmark_params fields."""
    for cfg in configs:
        bp = cfg.get("benchmark_params")
        if isinstance(bp, dict) and "$ref" in bp:
            ref_path = config_dir / bp["$ref"]
            try:
                with open(ref_path, encoding="utf-8") as f:
                    cfg["benchmark_params"] = json.load(f)
            except FileNotFoundError:
                raise ValueError(f"benchmark_params $ref not found: {ref_path}")
            except json.JSONDecodeError as e:
                raise ValueError(f"JSON parsing error in {ref_path}: {e}")
    return configs


def load_configs(config_path: str) -> list[dict[str, Any]]:
    """Load benchmark configs from JSON file and process inline fields."""
    try:
        abs_path = Path(config_path).resolve()
        with open(abs_path, encoding="utf-8") as f:
            configs = json.load(f)
        configs = _resolve_refs(configs, abs_path.parent)
        _process_inline_fields(configs)
        return configs
    except json.JSONDecodeError as e:
        raise ValueError(f"JSON parsing error: {str(e)}")
    except FileNotFoundError:
        raise ValueError(f"Configuration file not found: {config_path}")
    except Exception as e:
        raise RuntimeError(f"Failed to load configuration file: {str(e)}")


def load_diffusion_benchmark_configs(
    config_path: str | None = None,
    *,
    config_dir: Path | None = None,
) -> list[dict[str, Any]]:
    """Load one diffusion benchmark JSON, or merge all ``*.json`` under *config_dir*."""
    if config_path is not None:
        configs = load_configs(config_path)
        source = str(Path(config_path).resolve())
        for cfg in configs:
            cfg.setdefault(_DIFFUSION_SOURCE_CONFIG_KEY, source)
        return configs
    if config_dir is None:
        raise ValueError("load_diffusion_benchmark_configs requires config_path or config_dir")
    configs: list[dict[str, Any]] = []
    for path in sorted(config_dir.glob("*.json")):
        source = str(path.resolve())
        for cfg in load_configs(str(path)):
            cfg[_DIFFUSION_SOURCE_CONFIG_KEY] = source
            configs.append(cfg)
    if not configs:
        raise ValueError(f"No benchmark JSON files found under {config_dir}")
    return configs


if CONFIG_FILE_PATH is None:
    _all_configs = load_diffusion_benchmark_configs(config_dir=_PERF_TESTS_DIR)
    BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if is_diffusion_perf_config(cfg)]
    print(
        f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} diffusion case(s) from "
        f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} omni/tts; "
        f"use -m to filter, e.g. -m diffusion)"
    )
else:
    BENCHMARK_CONFIGS = load_diffusion_benchmark_configs(CONFIG_FILE_PATH)

_AGGREGATED_RESULT_FILES_BY_SOURCE: dict[str, Path] = {}


def _normalized_source_path(source_file: str) -> str:
    return str(Path(source_file).resolve())


def _aggregated_result_file_for_source(source_file: str) -> Path:
    """One session aggregate per source JSON (same naming as single ``--test-config-file``)."""
    key = _normalized_source_path(source_file)
    if key not in _AGGREGATED_RESULT_FILES_BY_SOURCE:
        stem = Path(key).stem
        resource = resource_label_for_filename(get_runtime_resource_label())
        if resource:
            result_name = f"diffusion_result_{stem}_{resource}_{_SESSION_TIMESTAMP}.json"
        else:
            result_name = f"diffusion_result_{stem}_{_SESSION_TIMESTAMP}.json"
        _AGGREGATED_RESULT_FILES_BY_SOURCE[key] = BENCHMARK_RESULT_DIR / result_name
    return _AGGREGATED_RESULT_FILES_BY_SOURCE[key]


def _write_result_record(record: dict[str, Any]) -> Path:
    """Append one benchmark record to the aggregate file for its source JSON."""
    source_file = record.get("source_file") or CONFIG_FILE_PATH
    if not source_file:
        raise ValueError("benchmark record missing source_file")
    target = _aggregated_result_file_for_source(str(source_file))
    with _RESULT_LOCK:
        BENCHMARK_RESULT_DIR.mkdir(parents=True, exist_ok=True)
        if target.exists():
            with open(target, encoding="utf-8") as f:
                records: list[dict] = json.load(f)
        else:
            records = []
        records.append(record)
        with open(target, "w", encoding="utf-8") as f:
            json.dump(records, f, indent=2, ensure_ascii=False)
    return target


_server_lock = threading.Lock()

# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------


def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
    """Block until the given host:port accepts connections or timeout expires.

    If *proc* is provided, also monitors the process; raises RuntimeError
    immediately if the server process exits before the port becomes available.
    """
    start = time.time()
    while time.time() - start < timeout:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(1)
                if s.connect_ex((host, port)) == 0:
                    return
        except Exception:
            pass
        if proc is not None:
            ret = proc.poll()
            if ret is not None:
                raise RuntimeError(f"Server process exited with code {ret} before port {host}:{port} became ready")
        time.sleep(2)
    raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s")


def _kill_process_tree(pid: int) -> None:
    try:
        parent = psutil.Process(pid)
        children = parent.children(recursive=True)
        all_pids = [pid] + [c.pid for c in children]

        for child in children:
            try:
                child.terminate()
            except psutil.NoSuchProcess:
                pass

        gone, alive = psutil.wait_procs(children, timeout=10)
        for child in alive:
            try:
                child.kill()
            except psutil.NoSuchProcess:
                pass

        try:
            parent.terminate()
            parent.wait(timeout=10)
        except (psutil.NoSuchProcess, psutil.TimeoutExpired):
            try:
                parent.kill()
            except psutil.NoSuchProcess:
                pass

        time.sleep(1)
        still_alive = [p for p in all_pids if psutil.pid_exists(p)]
        if still_alive:
            print(f"Warning: processes still alive after shutdown: {still_alive}")
            for p in still_alive:
                try:
                    subprocess.run(["kill", "-9", str(p)], timeout=2)
                except Exception:
                    pass
    except psutil.NoSuchProcess:
        pass


# ---------------------------------------------------------------------------
# Server classes
# ---------------------------------------------------------------------------


def _resolve_offline_model(model: str) -> str:
    """Under HF_HUB_OFFLINE, resolve a HF repo id to its local snapshot dir."""
    import huggingface_hub

    if not model or os.path.isdir(model) or not huggingface_hub.constants.HF_HUB_OFFLINE:
        return model
    from huggingface_hub import snapshot_download

    return snapshot_download(model, local_files_only=True)


class DiffusionServer:
    """Start a vLLM-Omni diffusion model server as a subprocess.

    Launched via vllm_omni.entrypoints.cli.main with the diffusion-specific
    parallelism flags (--usp, --ring, --cfg-parallel-size, etc.) passed directly
    on the CLI.  Minimum hardware: 4× NVIDIA H100 80 GB.
    """

    server_type = "vllm-omni"

    def __init__(
        self,
        server_cfg: dict[str, Any],
        *,
        port: int | None = None,
    ) -> None:
        self.server_cfg: dict[str, Any] = server_cfg
        self.model = _resolve_offline_model(server_cfg["model"])
        self.serve_args = server_cfg["serve_args"]
        self.host = "127.0.0.1"
        self.port = port if port is not None else get_open_port(self.host)
        self.proc: subprocess.Popen | None = None
        self.test_name: str = ""

    def _start_server(self) -> None:
        env = os.environ.copy()
        env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"

        cmd = [
            sys.executable,
            "-m",
            "vllm_omni.entrypoints.cli.main",
            "serve",
            self.model,
            "--omni",
            "--host",
            self.host,
            "--port",
            str(self.port),
        ] + self.serve_args

        print(f"Launching DiffusionServer: {' '.join(cmd)}")
        self.proc = subprocess.Popen(
            cmd,
            env=env,
            cwd=str(Path(__file__).parent.parent.parent.parent),
        )
        _wait_for_port(self.host, self.port, proc=self.proc)
        print(f"DiffusionServer ready on {self.host}:{self.port}")

    def __enter__(self):
        self._start_server()
        return self

    def __exit__(self, *_):
        if self.proc:
            _kill_process_tree(self.proc.pid)


# ---------------------------------------------------------------------------
# Config helpers
# ---------------------------------------------------------------------------


def _build_serve_args(serve_args_dict: dict[str, Any]) -> list[str]:
    """Convert a serve_args dict from test.json into a flat CLI argument list."""
    args: list[str] = []
    for key, value in serve_args_dict.items():
        flag = f"--{key}"
        if isinstance(value, bool):
            if value:
                args.append(flag)
        elif isinstance(value, dict):
            args.extend([flag, json.dumps(value, separators=(",", ":"))])
        else:
            args.extend([flag, str(value)])
    return args


def _get_branchpoint_commit_sha() -> str:
    """Return the branch-point commit SHA against main.

    Uses git command: ``git merge-base HEAD origin/main``.
    """
    global _BRANCHPOINT_COMMIT_SHA
    if _BRANCHPOINT_COMMIT_SHA is not None:
        return _BRANCHPOINT_COMMIT_SHA

    repo_root = Path(__file__).parent.parent.parent.parent
    try:
        sha = (
            subprocess.check_output(
                ["git", "merge-base", "HEAD", "origin/main"],
                cwd=str(repo_root),
                stderr=subprocess.STDOUT,
                text=True,
            )
            .strip()
            .splitlines()[0]
        )
        _BRANCHPOINT_COMMIT_SHA = sha
    except Exception as e:
        print(f"Warning: failed to get branch-point commit SHA: {e}")
        _BRANCHPOINT_COMMIT_SHA = ""
    return _BRANCHPOINT_COMMIT_SHA


def _to_resolution_string(params: dict[str, Any]) -> str:
    width = params.get("width", "unknown width")
    height = params.get("height", "unknown height")
    return f"{width}x{height}"


def _to_parallelism_string(framework: str, serve_args_dict: dict[str, Any]) -> str:
    parts: list[str] = []
    if framework == "vllm-omni":
        keys = [
            "num-gpus",
            "usp",
            "ulysses-degree",
            "ring",
            "ring-degree",
            "cfg-parallel-size",
            "vae-patch-parallel-size",
            "vae-use-tiling",
            "tensor-parallel-size",
        ]
        for key in keys:
            if key in serve_args_dict:
                parts.append(f"{key}={serve_args_dict[key]}")
    return ",".join(parts) if parts else "none"


def _to_cache_string(framework: str, serve_args_dict: dict[str, Any]) -> str:
    if framework == "vllm-omni":
        if "cache-backend" in serve_args_dict:
            return str(serve_args_dict["cache-backend"])
    return "disabled"


def _to_offload_string(framework: str, serve_args_dict: dict[str, Any]) -> str:
    selected: list[str] = []
    if framework == "vllm-omni":
        offload_keys = [
            "enable-cpu-offload",
            "enable-layerwise-offload",
        ]
        for key in offload_keys:
            if key in serve_args_dict:
                selected.append(key)
    return f"enabled({';'.join(selected)})" if selected else "disabled"


def _to_compile_value(framework: str, serve_args_dict: dict[str, Any]) -> str:
    if framework == "vllm-omni":
        if "enforce-eager" in serve_args_dict:
            return "disabled"
        return "enabled"
    return "disabled"


def _to_quantization_value(framework: str, serve_args_dict: dict[str, Any]) -> str:
    if framework == "vllm-omni":
        quant = serve_args_dict.get("quantization")
        return str(quant) if quant else "disabled"
    return "disabled"


def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Return one server-config dict per unique test_name."""
    seen: set[str] = set()
    result: list[dict[str, Any]] = []
    for cfg in configs:
        test_name = cfg["test_name"]
        if test_name in seen:
            continue
        seen.add(test_name)
        server_type = cfg.get("server_type", "vllm-omni")
        if server_type != "vllm-omni":
            raise ValueError(f"Unsupported server_type in config: {server_type}")
        serve_args_dict = cfg["server_params"].get("serve_args", {})
        result.append(
            {
                "test_name": test_name,
                "server_type": server_type,
                "model": cfg["server_params"]["model"],
                "serve_args_dict": serve_args_dict,
                "serve_args": _build_serve_args(serve_args_dict),
                "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")),
                "server_params": cfg["server_params"],
                "mark": cfg.get("mark"),
            }
        )
    return result


def _test_param_mapping(configs: list[dict[str, Any]]) -> dict[str, list[dict]]:
    mapping: dict[str, list[dict]] = {}
    for cfg in configs:
        name = cfg["test_name"]
        mapping.setdefault(name, [])
        mapping[name].extend(cfg["benchmark_params"])
    return mapping


def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]:
    return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs}


def _paired_diffusion_benchmark_pytest_params(configs: list[dict[str, Any]]) -> list[Any]:
    """Paired params for ``run_diffusion_benchmark.py``; same shape as omni runner."""
    test_param_map = _test_param_mapping(configs)
    server_entries = [(cfg, cfg["test_name"]) for cfg in _unique_server_params(configs)]
    return create_paired_benchmark_pytest_params(server_entries, test_param_map, _marks_by_test_name(configs))


# ---------------------------------------------------------------------------
# Parametrize data
# ---------------------------------------------------------------------------

test_param_map = _test_param_mapping(BENCHMARK_CONFIGS)
paired_benchmark_params = _paired_diffusion_benchmark_pytest_params(BENCHMARK_CONFIGS)


def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer:
    """Factory: return a vLLM-Omni diffusion server instance for the config."""
    return DiffusionServer(server_cfg=server_cfg)


# ---------------------------------------------------------------------------
# Pytest fixtures
# ---------------------------------------------------------------------------


@pytest.fixture(scope="module")
def diffusion_server(request):
    """Start one vLLM-Omni server per unique test configuration."""
    with _server_lock:
        server_cfg: dict[str, Any] = request.param
        test_name = server_cfg["test_name"]
        server_type = server_cfg["server_type"]

        print(f"\nStarting {server_type} server for test: {test_name}")
        with _make_server(server_cfg) as server:
            server.test_name = test_name
            print(f"{server_type} server started successfully")
            yield server
            print(f"{server_type} server stopping…")

    print(f"{server_type} server stopped")


@pytest.fixture
def benchmark_params(request):
    """Benchmark params for the paired server/index parametrization."""
    test_name, param_index = request.param

    params_list = test_param_map.get(test_name, [])
    if not params_list:
        raise ValueError(f"No benchmark params for test: {test_name}")

    current = param_index + 1
    total = len(params_list)
    print(f"\n  Running benchmark {current}/{total} for {test_name}")
    return {"test_name": test_name, "params": params_list[param_index]}


# ---------------------------------------------------------------------------
# Benchmark runner
# ---------------------------------------------------------------------------


_STAGE_METRICS_ENDPOINTS = {"/v1/chat/completions"}
_DIFFUSION_PIPELINE_PROFILER_ARG = "enable-diffusion-pipeline-profiler"


def run_benchmark(
    host: str,
    port: int,
    model: str,
    params: dict[str, Any],
    test_name: str,
    endpoint: str = "/v1/chat/completions",
    server_cfg: dict[str, Any] | None = None,
    source_file: str = "",
) -> dict[str, Any]:
    """Run diffusion_benchmark_serving.py as a subprocess and return parsed metrics.

    The raw metrics are written to a temporary file by the subprocess.  After
    the run completes the metrics are merged with full metadata (test_name,
    endpoint, benchmark_params, timestamp, flat reporting fields) and appended
    to ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` for the
    source JSON file. The temporary file is removed afterwards.  Subprocess stdout/stderr are tee'd
    to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in
    the record.
    """
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    endpoint = normalize_endpoint(endpoint)

    log_dir = BENCHMARK_RESULT_DIR / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    endpoint_label = endpoint_filename_token(endpoint)
    resource_label = get_runtime_resource_label()
    hw_for_filename = resource_label_for_filename(resource_label)
    if hw_for_filename:
        log_file = log_dir / f"{test_name}_{hw_for_filename}_{endpoint_label}_{timestamp}.log"
    else:
        log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log"

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
        tmp_result_file = Path(tmp.name)

    exclude_keys = {"baseline", "dataset", "task", "name"}

    cmd = [
        sys.executable,
        BENCHMARK_SCRIPT,
        "--host",
        host,
        "--port",
        str(port),
        "--model",
        model,
        "--endpoint",
        endpoint,
        "--dataset",
        params.get("dataset", "random"),
        "--task",
        params.get("task", "t2i"),
        "--output-file",
        str(tmp_result_file),
    ]

    serve_args_dict = (server_cfg or {}).get("serve_args_dict")
    profiler_enabled = isinstance(serve_args_dict, dict) and bool(serve_args_dict.get(_DIFFUSION_PIPELINE_PROFILER_ARG))
    if endpoint in _STAGE_METRICS_ENDPOINTS and profiler_enabled:
        cmd.append("--return-stage-metrics")

    for key, value in params.items():
        if key in exclude_keys or value is None:
            continue
        flag = f"--{key}"
        if isinstance(value, bool):
            if value:
                cmd.append(flag)
        elif isinstance(value, (dict, list)):
            cmd.extend([flag, json.dumps(value, separators=(",", ":"))])
        else:
            cmd.extend([flag, str(value)])

    # Insert -u so the subprocess runs with unbuffered stdout/stderr, ensuring
    # all print() output is flushed to the pipe immediately instead of being
    # held in Python's internal block-buffer until process exit (which can
    # cause truncated or out-of-order log output when stdout is piped).
    cmd = [cmd[0], "-u"] + cmd[1:]

    print(f"\nRunning benchmark (endpoint={endpoint}): {' '.join(cmd)}")
    print(f"  Log file: {log_file}")

    # Redirect stdout + stderr directly to the log file at the OS level
    # (equivalent to `cmd > log 2>&1`), so no output is ever lost regardless
    # of how the subprocess exits.  The log is echoed to the terminal afterwards.
    with open(log_file, "w", encoding="utf-8") as log_fh:
        log_fh.write(f"cmd: {' '.join(cmd)}\n\n")
        log_fh.flush()

        process = subprocess.Popen(
            cmd,
            stdout=log_fh,
            stderr=log_fh,
            cwd=str(Path(__file__).parent.parent.parent.parent),
        )
        process.wait()

    with open(log_file, encoding="utf-8") as log_fh:
        print(log_fh.read(), end="")

    if process.returncode != 0:
        tmp_result_file.unlink(missing_ok=True)
        print(f"ERROR:Benchmark script exited with code {process.returncode}")

    if not tmp_result_file.exists():
        with open(DIFFUSION_RESULT_TEMPLATE_PATH, encoding="utf-8") as f:
            template_payload = json.load(f)
        # Template schema is fixed and owned by this repo:
        # ``diffusion_result_template.json`` is a one-item list and metrics live at [0]["result"].
        template_metrics: dict[str, Any] = template_payload[0]["result"]
        with open(tmp_result_file, "w", encoding="utf-8") as f:
            json.dump(template_metrics, f, ensure_ascii=False, indent=2)
        print(f"Benchmark result file not generated, fallback to template: {tmp_result_file}")

    try:
        with open(tmp_result_file, encoding="utf-8") as f:
            metrics: dict[str, Any] = json.load(f)
    finally:
        tmp_result_file.unlink(missing_ok=True)

    server_cfg = server_cfg or {}
    server_type = cast(str, server_cfg.get("server_type", "vllm-omni"))
    serve_args_dict = server_cfg.get("serve_args_dict", {})
    if not isinstance(serve_args_dict, dict):
        serve_args_dict = {}

    completed = metrics.get("completed_requests", metrics.get("completed", 0))
    failed = metrics.get("failed_requests", metrics.get("failed", 0))

    # Persist sweep-resolved baseline from params (already narrowed in _build_run_params).
    baseline = params.get("baseline") or {}
    metrics["baseline"] = baseline

    record: dict[str, Any] = {
        "test_name": test_name,
        "endpoint": endpoint,
        "timestamp": timestamp,
        "server_params": server_cfg.get("server_params"),
        "benchmark_params": params,
        "baseline": baseline,
        "result": metrics,
        "log_file": str(log_file),
        "Model": model,
        "Framework": server_type,
        "API Endpoint": endpoint,
        "Hardware": hardware_json_value(resource_label),
        "Deployment": "",
        "Task": params.get("task", "t2i"),
        "Dataset": params.get("dataset", "random"),
        "resolution": _to_resolution_string(params),
        "Parallelism": _to_parallelism_string(server_type, serve_args_dict),
        "max_concurrency": params.get("max-concurrency", ""),
        "Cache": _to_cache_string(server_type, serve_args_dict),
        "Quantization": _to_quantization_value(server_type, serve_args_dict),
        "offload": _to_offload_string(server_type, serve_args_dict),
        "compile": _to_compile_value(server_type, serve_args_dict),
        "Attn_backend": os.environ.get("DIFFUSION_ATTENTION_BACKEND", ""),
        "num_inference_steps": params.get("num-inference-steps", ""),
        "completed": completed,
        "failed": failed,
        "throughput_qps": metrics.get("throughput_qps"),
        "latency_mean": metrics.get("latency_mean"),
        "latency_median": metrics.get("latency_median"),
        "latency_p99": metrics.get("latency_p99"),
        "latency_p95": metrics.get("latency_p95"),
        "latency_p50": metrics.get("latency_p50"),
        "peak_memory_mb_max": metrics.get("peak_memory_mb_max"),
        "peak_memory_mb_mean": metrics.get("peak_memory_mb_mean"),
        "peak_memory_mb_median": metrics.get("peak_memory_mb_median"),
        "stage_durations_mean": metrics.get("stage_durations_mean"),
        "stage_durations_p50": metrics.get("stage_durations_p50"),
        "stage_durations_p99": metrics.get("stage_durations_p99"),
        "commit_sha": _get_branchpoint_commit_sha(),
        "build_id": os.environ.get("BUILDKITE_BUILD_ID", ""),
        "build_url": os.environ.get("BUILDKITE_BUILD_URL", ""),
        "source_file": source_file,
    }
    result_path = _write_result_record(record)
    print(f"\n  Result saved to: {result_path}")
    print(f"  Log saved to:       {log_file}")

    return metrics


# ---------------------------------------------------------------------------
# Assertions
# ---------------------------------------------------------------------------


def assert_result(
    result: dict[str, Any],
    num_prompts: int,
) -> None:
    """Assert that the expected number of requests completed."""
    completed = result.get("completed_requests", result.get("completed", 0))
    assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"


def _default_benchmark_endpoint_for_task(task: str) -> str:
    """Return the default client-side benchmark endpoint for a diffusion task."""
    if task in {"t2v", "i2v", "ti2v", "v2v"}:
        return "/v1/videos"
    if task in {"t2i", "i2i", "ti2i"}:
        return "/v1/chat/completions"
    raise ValueError(f"Unsupported task for benchmark endpoint resolution: {task}")


def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
    """Resolve which serving API the benchmark client should call."""
    configured = server_cfg.get("benchmark_endpoint")
    if configured:
        return normalize_endpoint(cast(str, configured))
    return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))


def _to_list(value: Any) -> list[Any]:
    if value is None:
        return []
    return [value] if not isinstance(value, (list, tuple)) else list(value)


def _build_run_params(
    params: dict[str, Any],
    *,
    num_prompts: int,
    sweep_index: int | None = None,
    request_rate: Any | None = None,
    max_concurrency: Any | None = None,
) -> dict[str, Any]:
    run_params = {
        key: value for key, value in params.items() if key not in {"request-rate", "max-concurrency", "num-prompts"}
    }
    run_params["num-prompts"] = num_prompts
    if request_rate is not None:
        run_params["request-rate"] = request_rate
    if max_concurrency is not None:
        run_params["max-concurrency"] = max_concurrency
    if "baseline" in params:
        # Keep all hardware buckets; pick the metric value for this sweep step only.
        run_params["baseline"] = resolve_baseline_for_sweep(
            params.get("baseline"),
            sweep_index=sweep_index,
        )
    return run_params


def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
    request_rate_list = _to_list(params.get("request-rate"))
    num_prompt_list = _to_list(params.get("num-prompts", 10))
    max_concurrency_list = _to_list(params.get("max-concurrency"))

    max_len = max(len(request_rate_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(request_rate_list) == 1:
            request_rate_list = request_rate_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(request_rate_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 request-rate or max-concurrency")

    sweep_runs: list[dict[str, Any]] = []

    for sweep_index, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
        sweep_runs.append(
            {
                "params": _build_run_params(
                    params,
                    request_rate=request_rate,
                    num_prompts=num_prompts,
                    sweep_index=sweep_index,
                ),
                "num_prompts": num_prompts,
            }
        )

    for sweep_index, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
        sweep_runs.append(
            {
                "params": _build_run_params(
                    params,
                    max_concurrency=max_concurrency,
                    num_prompts=num_prompts,
                    request_rate="inf",
                    sweep_index=sweep_index,
                ),
                "num_prompts": num_prompts,
            }
        )

    if not sweep_runs:
        default_num_prompts = num_prompt_list[0]
        sweep_runs.append(
            {
                "params": _build_run_params(
                    params,
                    num_prompts=default_num_prompts,
                    sweep_index=0,
                ),
                "num_prompts": default_num_prompts,
            }
        )

    return sweep_runs


# ---------------------------------------------------------------------------
# Test entry point
# ---------------------------------------------------------------------------
@pytest.mark.benchmark
@pytest.mark.parametrize(
    "diffusion_server,benchmark_params",
    paired_benchmark_params,
    indirect=["diffusion_server", "benchmark_params"],
)
def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
    """Run the diffusion performance benchmark and verify request completion.

    One server is started per unique parallel configuration (module scope).
    For each server, all benchmark parameter sets defined in the config JSON
    are executed sequentially; metrics are recorded to the aggregated result file.
    """
    test_name = benchmark_params["test_name"]
    params = benchmark_params["params"]
    server_cfg = getattr(diffusion_server, "server_cfg", {})
    sweep_runs = _iter_sweep_runs(params)

    for sweep_run in sweep_runs:
        endpoint = _resolve_benchmark_endpoint(server_cfg, sweep_run["params"])
        result = run_benchmark(
            host=diffusion_server.host,
            port=diffusion_server.port,
            model=diffusion_server.model,
            params=sweep_run["params"],
            test_name=test_name,
            endpoint=endpoint,
            server_cfg=server_cfg,
            source_file=server_cfg.get(
                _DIFFUSION_SOURCE_CONFIG_KEY,
                CONFIG_FILE_PATH or f"{_PERF_TESTS_DIR}/*.json",
            ),
        )

        print(f"\n{'=' * 60}")
        print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
        for key in (
            "throughput_qps",
            "latency_mean",
            "latency_median",
            "latency_p50",
            "latency_p99",
            "peak_memory_mb_max",
            "peak_memory_mb_mean",
            "peak_memory_mb_median",
        ):
            if key in result:
                print(f"  {key}: {result[key]:.4f}")

        source = server_cfg.get(_DIFFUSION_SOURCE_CONFIG_KEY) or CONFIG_FILE_PATH
        if source:
            print(f"\n  Aggregated results: {_aggregated_result_file_for_source(str(source))}")
        print("=" * 60)

        assert_result(result, sweep_run["num_prompts"])