"""5.7 本地模型测速（MLX 进程内版，只能在 Apple Silicon Mac 上跑）。

不经过 HTTP 服务，直接在 Python 里调 mlx-lm：单流和一次并行 N 路（batch_generate），
输出和 bench.py 同样的表，方便与 llama.cpp、vLLM 的结果并排比较。

用法：
    python bench_mlx.py --model mlx-community/Qwen3.5-2B-4bit --label "MLX 4bit"
    python bench_mlx.py --model mlx-community/Qwen3.5-2B-8bit --label "MLX 8bit"
"""
import argparse
import json
import time
from datetime import datetime
from pathlib import Path

from mlx_lm import batch_generate, load

FILLER = ('The quick brown fox jumps over the lazy dog while the model server keeps answering requests. ' * 40)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', default='mlx-community/Qwen3.5-2B-4bit')
    ap.add_argument('--label', default='')
    ap.add_argument('--max-tokens', type=int, default=256)
    ap.add_argument('--concurrency', type=int, nargs='+', default=[1, 4, 8, 16])
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    model, tok = load(args.model)
    msgs = [{'role': 'system', 'content': FILLER},
            {'role': 'user', 'content': 'Count from 1 upward, separated by commas. Do not stop.'}]
    ids = tok.apply_chat_template(msgs, add_generation_prompt=True, enable_thinking=False)
    print(f'{args.label or args.model}（MLX 进程内），每个请求 {len(ids)} token 输入、{args.max_tokens} token 输出，{datetime.now():%Y-%m-%d %H:%M}')
    batch_generate(model, tok, [ids], max_tokens=16)                     # 预热

    rows = []
    print('\n| 并发路数 | 首 token 延迟（整批 prefill） | 每一路的生成速度 | 所有路加起来 |')
    print('|---|---|---|---|')
    for n in args.concurrency:
        t0 = time.perf_counter()
        r = batch_generate(model, tok, [ids] * n, max_tokens=args.max_tokens, return_token_ids=True)
        wall = time.perf_counter() - t0
        tokens = [len(t) for t in r.token_ids]
        prefill = r.stats.prompt_time
        per_stream = (sum(tokens) / n) / max(wall - prefill, 1e-9)
        rows.append(dict(concurrency=n, ttft=prefill, per_stream_tok_s=per_stream, total_tok_s=sum(tokens) / wall,
                         tokens=tokens, peak_memory_gb=r.stats.peak_memory))
        print(f'| {n} | {prefill:.2f} 秒 | {per_stream:.0f} token/秒 | {sum(tokens) / wall:.0f} token/秒 |')
    if any(min(r['tokens']) < args.max_tokens * 0.9 for r in rows):
        print('\n注意：有的请求没生成满 max_tokens 就停了，速度会偏低或不可比。')
    Path(args.out).mkdir(exist_ok=True)
    stem = (args.label or args.model).replace('/', '_').replace(' ', '_')
    path = Path(args.out) / f'{stem}_{datetime.now():%Y%m%d_%H%M}.json'
    path.write_text(json.dumps(dict(label=args.label, model=args.model, max_tokens=args.max_tokens, rows=rows), indent=1), encoding='utf-8')
    print(f'\n已存到 {path}')


if __name__ == '__main__':
    main()