#!/usr/bin/env python3
"""Profile vLLM 0.20 TP all-reduce and assert FlashInfer TRTLLM dispatch."""

from __future__ import annotations

import argparse
import json
import os
import statistics
import subprocess
from pathlib import Path
from typing import Any

import torch
import torch.distributed as dist
import vllm


VLLM_VERSION = "0.20.0"
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vllm-source", type=Path, required=True)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--num-tokens", type=int, nargs="+", default=[8])
    parser.add_argument("--hidden-dim", type=int, default=2048)
    parser.add_argument("--warmup-iters", type=int, default=3)
    parser.add_argument("--repeats", type=int, default=10)
    return parser.parse_args()


def git_head(repo: Path) -> str:
    return subprocess.check_output(
        ["git", "-C", str(repo), "rev-parse", "HEAD"], text=True
    ).strip()


def stats_ms(samples: list[float]) -> dict[str, float]:
    return {
        "min": min(samples),
        "max": max(samples),
        "mean": statistics.fmean(samples),
        "median": statistics.median(samples),
        "std": statistics.pstdev(samples),
    }


def main() -> None:
    args = parse_args()
    if vllm.__version__ != VLLM_VERSION:
        raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}")
    source_head = git_head(args.vllm_source)
    if source_head != VLLM_COMMIT:
        raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
    if os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER") != "1":
        raise SystemExit("VLLM_ALLREDUCE_USE_FLASHINFER must equal 1")
    if os.getenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND") != "trtllm":
        raise SystemExit("VLLM_FLASHINFER_ALLREDUCE_BACKEND must equal trtllm")
    if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
        raise SystemExit("launch with torchrun")

    from vllm.distributed import (
        destroy_distributed_environment,
        destroy_model_parallel,
        init_distributed_environment,
        initialize_model_parallel,
        tensor_model_parallel_all_reduce,
    )
    from vllm.distributed.parallel_state import get_tp_group
    from vllm.config import (
        ModelConfig,
        ParallelConfig,
        VllmConfig,
        set_current_vllm_config,
    )

    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    if world_size not in (2, 4):
        raise SystemExit(f"expected TP world size 2 or 4, got {world_size}")
    device = torch.device(f"cuda:{local_rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)
    init_distributed_environment()
    model_config = ModelConfig(
        model=str(args.model),
        dtype="bfloat16",
        max_model_len=8192,
        skip_tokenizer_init=True,
        generation_config="vllm",
    )
    vllm_config = VllmConfig(
        model_config=model_config,
        parallel_config=ParallelConfig(tensor_parallel_size=world_size)
    )
    with set_current_vllm_config(vllm_config):
        initialize_model_parallel(tensor_model_parallel_size=world_size)

    rows: list[dict[str, Any]] = []
    expected_sum = world_size * (world_size + 1) / 2
    try:
        for num_tokens in args.num_tokens:
            input_tensor = torch.full(
                (num_tokens, args.hidden_dim),
                float(rank + 1),
                dtype=torch.bfloat16,
                device=device,
            )
            for _ in range(args.warmup_iters):
                output = tensor_model_parallel_all_reduce(input_tensor)
            torch.accelerator.synchronize()
            torch.testing.assert_close(
                output,
                torch.full_like(output, expected_sum),
                atol=0.0,
                rtol=0.0,
            )

            communicator = get_tp_group().device_communicator
            fi_comm = communicator.fi_ar_comm
            if fi_comm is None or fi_comm.disabled:
                raise SystemExit(
                    f"FlashInfer all-reduce was not selected at TP={world_size}, "
                    f"tokens={num_tokens}"
                )
            uses_flashinfer = fi_comm.should_use_fi_ar(input_tensor)

            samples: list[float] = []
            for _ in range(args.repeats):
                dist.barrier()
                start = torch.cuda.Event(enable_timing=True)
                end = torch.cuda.Event(enable_timing=True)
                start.record()
                output = tensor_model_parallel_all_reduce(input_tensor)
                end.record()
                torch.accelerator.synchronize()
                samples.append(float(start.elapsed_time(end)))

            gathered: list[list[float] | None] = [None] * world_size
            dist.all_gather_object(gathered, samples)
            if rank == 0:
                per_rank = [stats_ms(item) for item in gathered if item is not None]
                row = {
                    "tensor_parallel_size": world_size,
                    "num_tokens": num_tokens,
                    "hidden_dim": args.hidden_dim,
                    "payload_bytes": num_tokens
                    * args.hidden_dim
                    * torch.tensor([], dtype=torch.bfloat16).element_size(),
                    "dtype": "bfloat16",
                    "communicator": "vllm.tensor_model_parallel_all_reduce",
                    "selected_backend": (
                        "flashinfer_trtllm" if uses_flashinfer else "nccl_fallback"
                    ),
                    "per_rank_time_ms": per_rank,
                    "critical_path_median_ms": max(
                        rank_stats["median"] for rank_stats in per_rank
                    ),
                }
                rows.append(row)
                print(json.dumps(row, sort_keys=True), flush=True)
    finally:
        destroy_model_parallel()
        destroy_distributed_environment()

    if rank == 0:
        payload = {
            "schema_version": "qwen30_vllm020_allreduce_raw.v1",
            "environment": {
                "vllm_version": vllm.__version__,
                "vllm_source_commit": source_head,
                "torch_version": torch.__version__,
                "torch_cuda": torch.version.cuda,
                "gpu": torch.cuda.get_device_name(device),
                "model": str(args.model),
                "backend_env": {
                    "VLLM_ALLREDUCE_USE_FLASHINFER": "1",
                    "VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm",
                },
            },
            "rows": rows,
        }
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


if __name__ == "__main__":
    main()