"""TTFT / TPOT benchmark of the full DeepSeek-V4.1-Flash on one trn2.48xlarge, the vllm-neuron#71 method offline.

vllm-neuron#71 (``contrib/models/deepseek_v41/test/performance/run_bench.py``) serves the model and drives it with
``vllm bench serve --dataset-name random``: random token ids fix the input length exactly and share no prefix,
output 128 tokens, and per cell it reports the median TTFT, the median TPOT (decode only:
``(latency - ttft) / (output_tokens - 1)``) and the decode tokens per second. DeepInfer has no server, so this drives
the engine directly with the same requests: random prompt token ids of the cell's input length, ``ignore_eos`` and
``max_tokens`` 128, timed per request from submission.

Topology: one engine process per DP rank (EP=64, DP-attention), and every rank runs its own request of the cell at
the same time: batch 1 per rank, 64 requests in flight (the engine steps all DP ranks together). #71's bs1 cells are
one request for the whole machine; per request the two loads cost the same here, since the static MoE computes
every rank's row either way.

    export DEEPINFER_ACCELERATOR=neuron WORLD_SIZE=1 RANK=0 DEEPINFER_NEURON_PREFILL_BUCKET=128 \\
        DEEPINFER_MAX_BACKTRACKING_TOKENS=128
    python examples/neuron/bench_neuron.py --model_name /mnt/nvme/models/DeepSeek-V4.1-Flash-1017 --ep 64
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from pathlib import Path

EXAMPLES = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(EXAMPLES))

from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env, setup_cuda_visible_devices  # noqa: E402
from _launch import launch_dp  # noqa: E402

parser = argparse.ArgumentParser(description="V4.1 TTFT/TPOT benchmark on Neuron (vllm-neuron#71 method)")
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
parser.add_argument("--ep", default=64, type=int)
parser.add_argument("--input_lens", default="512,1024,2048,8192", help="comma-separated prompt lengths (tokens)")
parser.add_argument("--output_len", default=128, type=int)
parser.add_argument("--repeats", default=3, type=int, help="requests per cell (medians are reported)")
parser.add_argument("--warmup", default=1, type=int, help="untimed requests before the first cell")
parser.add_argument("--num_kv_blocks", default=400, type=int)
parser.add_argument("--max_model_len", default=16384, type=int)
parser.add_argument("--max_num_batched_tokens", default=128, type=int)
parser.add_argument("--max_num_seqs", default=1, type=int)
parser.add_argument("--out", default="", help="write the results as JSON here (rank 0)")
args = parser.parse_args()
dp_size = args.ep


def _fixed_kv_pool(num_blocks: int) -> None:
    os.environ["DEEPINFER_NEURON_NUM_KV_BLOCKS"] = str(num_blocks)  # CP followers read it too
    from deepinfer.worker.base import WorkerBase

    def profile_num_available_blocks(self):  # noqa: ANN001, ANN202
        self.model_runner.set_block_size(self.cache_config.block_size)
        self.cache_config.num_gpu_blocks = num_blocks
        return num_blocks

    WorkerBase.profile_num_available_blocks = profile_num_available_blocks


def _random_prompt(length: int, seed: int, vocab: int) -> list[int]:
    import numpy as np

    rng = np.random.default_rng(seed)
    return rng.integers(1000, vocab - 1000, size=length, dtype=np.int64).astype(np.int32).tolist()


def _run_one(llm, prompt: list[int] | None, params) -> tuple[float, float, int] | None:  # noqa: ANN001
    """Run the engine until it drains. With a prompt, returns (ttft_s, latency_s, output_tokens) for it."""
    engine = llm.llm_engine
    start = time.perf_counter()
    if prompt is not None:
        llm._add_request(None, params, prompt)
    first = last = None
    tokens = 0
    while True:
        unfinished = engine.has_unfinished_requests()
        outputs = engine.step() if unfinished else engine.flush_inflight_reqs()
        now = time.perf_counter()
        for out in outputs:
            tokens += 1
            first = first if first is not None else now
            last = now
        if not unfinished:
            break
    if prompt is None:
        return None
    return first - start, last - start, tokens


def main(dp_size: int, dp_rank: int) -> None:
    os.environ["DEEPINFER_USE_NEXTN_DECODE"] = "0"
    os.environ.setdefault("DEEPINFER_ACCELERATOR", "neuron")
    setup_common_env(model_name=args.model_name, sampler_force_topk=0, set_override_experts=False, enable_eplb=False)
    from deep_cluster.hardware import get_actual_num_gpus

    setup_cuda_visible_devices(dp_rank, 1, get_actual_num_gpus())
    _fixed_kv_pool(args.num_kv_blocks)
    llm = create_llm(
        args.model_name,
        dp_rank,
        ep_size=args.ep,
        cp_size=1,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": "100"},
        data_parallel_size=dp_size,
        max_num_batched_tokens=args.max_num_batched_tokens,
        max_model_len=args.max_model_len,
        max_num_seqs=args.max_num_seqs,
        gpu_memory_utilization=0.85,
        lookahead_tokens=0,
    )
    from deepinfer.sampling_params import SamplingParams

    params = SamplingParams(temperature=0.0, max_tokens=args.output_len, ignore_eos=True)
    vocab = llm.llm_engine.tokenizer.vocab_size if hasattr(llm.llm_engine.tokenizer, "vocab_size") else 128000
    busy = True
    lens = [int(x) for x in args.input_lens.split(",")]

    for w in range(args.warmup):
        _run_one(llm, _random_prompt(lens[0], 10_000 + w + 7919 * dp_rank, vocab) if busy else None, params)
    results = []
    for length in lens:
        # One untimed request a cell: a cell's first request pays for its own context bucket's lazy compile.
        _run_one(llm, _random_prompt(length, 20_000 + length + 7919 * dp_rank, vocab) if busy else None, params)
        cell = []
        for r in range(args.repeats):
            got = _run_one(llm, _random_prompt(length, 1000 * length + r + 7919 * dp_rank, vocab) if busy else None, params)
            if got is not None:
                ttft, latency, n = got
                tpot = (latency - ttft) / max(n - 1, 1)
                cell.append((ttft, tpot, n))
                if dp_rank == 0:
                    print(f"BENCH input={length} rep={r} ttft_ms={ttft * 1e3:.1f} tpot_ms={tpot * 1e3:.1f} tokens={n}", flush=True)
        if dp_rank == 0 and cell:
            ttft = statistics.median(c[0] for c in cell)
            tpot = statistics.median(c[1] for c in cell)
            results.append({"input": length, "output": args.output_len, "ttft_p50_ms": ttft * 1e3, "tpot_p50_ms": tpot * 1e3,
                            "decode_tok_s_per_request": 1.0 / tpot})
    if dp_rank == 0:
        print("BENCH_RESULTS " + json.dumps(results), flush=True)
        print(f"| input | output | TTFT p50 (ms) | TPOT p50 (ms) | decode tok/s per request |   (batch 1 per DP rank)")
        for r in results:
            print(f"| {r['input']} | {r['output']} | {r['ttft_p50_ms']:.0f} | {r['tpot_p50_ms']:.1f} | {r['decode_tok_s_per_request']:.2f} |")
        if args.out:
            Path(args.out).write_text(json.dumps(results, indent=2))
    llm.close()


if __name__ == "__main__":
    launch_dp(main, dp_size)