#!/usr/bin/env python3
"""Profile the vLLM 0.20 TP all-reduce path under a frozen runtime contract."""

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(
        "--collective-contract",
        choices=("flashinfer-trtllm", "qwen235-serving-projected"),
        default="flashinfer-trtllm",
    )
    parser.add_argument("--warmup-iters", type=int, default=3)
    parser.add_argument("--repeats", type=int, default=10)
    parser.add_argument("--trials", type=int, default=1)
    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 args.repeats <= 0 or args.trials <= 0 or args.warmup_iters < 0:
        raise SystemExit("warmup/repeats/trials must be non-negative/positive")
    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")
    expected_symm_mem = (
        "1" if args.collective_contract == "qwen235-serving-projected" else "0"
    )
    if os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM") != expected_symm_mem:
        raise SystemExit(
            f"VLLM_ALLREDUCE_USE_SYMM_MEM must equal {expected_symm_mem}"
        )
    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,
        set_custom_all_reduce,
        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, 8):
        raise SystemExit(f"expected TP world size 2, 4, or 8, got {world_size}")
    device = torch.device(f"cuda:{local_rank}")
    torch.accelerator.set_device_index(device)
    torch.set_default_device(device)
    disable_custom_all_reduce = (
        args.collective_contract == "qwen235-serving-projected"
    )
    set_custom_all_reduce(not disable_custom_all_reduce)
    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,
            disable_custom_all_reduce=disable_custom_all_reduce,
        )
    )
    with set_current_vllm_config(vllm_config):
        initialize_model_parallel(tensor_model_parallel_size=world_size)

    from vllm.distributed.device_communicators.all_reduce_utils import (
        should_nccl_symm_mem_allreduce,
    )

    def resolve_profile_operation(input_tensor: torch.Tensor):
        communicator = get_tp_group().device_communicator
        if args.collective_contract == "qwen235-serving-projected":
            from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
                FI_ALLREDUCE_FUSION_MAX_SIZE_MB,
            )

            major, minor = torch.cuda.get_device_capability(device)
            capability = major * 10 + minor
            try:
                fusion_limit_bytes = int(
                    FI_ALLREDUCE_FUSION_MAX_SIZE_MB[capability][world_size]
                    * 1024
                    * 1024
                )
            except KeyError as error:
                raise SystemExit(
                    f"no FlashInfer fusion limit for SM{capability}, TP{world_size}"
                ) from error
            if input_tensor.nbytes <= fusion_limit_bytes:
                from vllm.distributed.device_communicators.flashinfer_all_reduce import (
                    get_fi_ar_workspace,
                )

                fi_comm = communicator.fi_ar_comm
                if fi_comm is None or fi_comm.disabled:
                    raise SystemExit("FlashInfer all-reduce communicator is unavailable")
                max_token_num = fusion_limit_bytes // (
                    args.hidden_dim * input_tensor.element_size()
                )
                workspace = get_fi_ar_workspace(
                    world_size=world_size,
                    rank=rank,
                    max_token_num=max_token_num,
                    hidden_dim=args.hidden_dim,
                    dtype=input_tensor.dtype,
                    group=get_tp_group().device_group,
                )
                if workspace is None:
                    raise SystemExit(
                        "FlashInfer rejected the serving-matched integer workspace"
                    )
                if not fi_comm.should_use_fi_ar(input_tensor):
                    raise SystemExit("FlashInfer rejected a fusion-eligible payload")
                return (
                    "flashinfer_trtllm_fused_projection",
                    fi_comm.all_reduce,
                    fusion_limit_bytes,
                )
            if (
                communicator.pynccl_comm is not None
                and should_nccl_symm_mem_allreduce(world_size, input_tensor)
            ):
                return (
                    "pynccl_symmetric_with_copy",
                    torch.ops.vllm.all_reduce_symmetric_with_copy,
                    fusion_limit_bytes,
                )
            if (
                communicator.symm_mem_comm is not None
                and communicator.symm_mem_comm.should_use_symm_mem(input_tensor)
            ):
                return (
                    "torch_symmetric_memory",
                    communicator.symm_mem_comm.all_reduce,
                    fusion_limit_bytes,
                )
            if communicator.pynccl_comm is not None:
                return "pynccl", communicator.pynccl_comm.all_reduce, fusion_limit_bytes
            return "torch_distributed", tensor_model_parallel_all_reduce, fusion_limit_bytes

        fi_comm = communicator.fi_ar_comm
        selected_backend = (
            "flashinfer_trtllm"
            if fi_comm is not None
            and not fi_comm.disabled
            and fi_comm.should_use_fi_ar(input_tensor)
            else "non_flashinfer_fallback"
        )
        return selected_backend, tensor_model_parallel_all_reduce, None

    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,
            )
            selected_backend, profile_operation, fusion_limit_bytes = (
                resolve_profile_operation(input_tensor)
            )
            for _ in range(args.warmup_iters):
                output = profile_operation(input_tensor)
            torch.accelerator.synchronize()
            torch.testing.assert_close(
                output,
                torch.full_like(output, expected_sum),
                atol=0.0,
                rtol=0.0,
            )

            if args.collective_contract == "flashinfer-trtllm":
                if selected_backend != "flashinfer_trtllm":
                    raise SystemExit(
                        f"expected FlashInfer TRTLLM, got {selected_backend} at "
                        f"TP={world_size}, tokens={num_tokens}"
                    )
            elif selected_backend not in {
                "flashinfer_trtllm_fused_projection",
                "pynccl_symmetric_with_copy",
                "torch_symmetric_memory",
                "pynccl",
            }:
                raise SystemExit(
                    f"unexpected Qwen235 serving backend: {selected_backend}"
                )

            all_trial_rank_samples: list[list[list[float]]] = []
            critical_path_samples: list[float] = []
            for _ in range(args.trials):
                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 = profile_operation(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:
                    rank_samples = [item for item in gathered if item is not None]
                    all_trial_rank_samples.append(rank_samples)
                    critical_path_samples.extend(
                        max(per_rank[index] for per_rank in rank_samples)
                        for index in range(args.repeats)
                    )
            if rank == 0:
                flattened_by_rank = [
                    [
                        sample
                        for trial in all_trial_rank_samples
                        for sample in trial[rank_index]
                    ]
                    for rank_index in range(world_size)
                ]
                per_rank = [stats_ms(samples) for samples in flattened_by_rank]
                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",
                    "collective_contract": args.collective_contract,
                    "selected_backend": selected_backend,
                    "real_fusion_limit_bytes": fusion_limit_bytes,
                    "trials": args.trials,
                    "repeats_per_trial": args.repeats,
                    "per_trial_rank_samples_ms": all_trial_rank_samples,
                    "per_rank_time_ms": per_rank,
                    "critical_path_time_ms": stats_ms(critical_path_samples),
                    "critical_path_median_ms": statistics.median(
                        critical_path_samples
                    ),
                }
                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"
                if args.collective_contract == "flashinfer-trtllm"
                else "vllm020_allreduce_raw.v2"
            ),
            "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": os.getenv(
                        "VLLM_ALLREDUCE_USE_FLASHINFER", "0"
                    ),
                    "VLLM_FLASHINFER_ALLREDUCE_BACKEND": os.getenv(
                        "VLLM_FLASHINFER_ALLREDUCE_BACKEND"
                    ),
                    "VLLM_ALLREDUCE_USE_SYMM_MEM": os.getenv(
                        "VLLM_ALLREDUCE_USE_SYMM_MEM", "1"
                    ),
                },
                "collective_contract": args.collective_contract,
                "disable_custom_all_reduce": disable_custom_all_reduce,
            },
            "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()
