"""
=============================================================================
CS349E · 第 16 讲 · 02_deployment_planner.py —— 部署规划器：给定模型、负载和 SLO，选硬件、精度、切分和共置 / 解耦
=============================================================================

【这个脚本要回答的问题】
    一个推理服务上线前要回答：用哪种 GPU、BF16 还是 FP8、张量并行几路、prefill 和 decode 放在一起还是分开、要几块 GPU？
    planner.plan 枚举 3 种 GPU × 2 种精度 × 张量并行 {1, 2, 4, 8} × {共置（分块 prefill）, 解耦}，
    用解析模型估每个实例能承受的流量，按每百万输出 token 的成本排序。
      演示 1：四个场景（聊天、写代码的智能体、离线批处理、低延迟的小模型）各自最便宜的几个方案。
      演示 2：同一个场景里，每块 GPU 的效率 vs 按整数台算的成本：为什么流量小的时候"整数"很重要。
      演示 3：每个场景里，各种 GPU / 精度 / 模式的最好方案——差别有多大，瓶颈在哪里。

【关于数字】解析模型（planner.decode_capacity / prefill_capacity / colocated_capacity）+ 第 5 讲 StepModel 的每步时间（roofline 模型，
    不是实测）；GPU 价格是假设值（H100 2.0、H200 2.6、B200 4.0 美元 / GPU·小时）。这里只做**筛选**，03 脚本用模拟器检验。
    TTFT 的可行性按 P90 的提示长度检查（tail_aware=True，原因见 03 脚本演示 1）。
【运行】python3 02_deployment_planner.py      （约 2 秒）
【依赖】numpy；本目录的 planner；../lecture01 的 console。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from planner import CHIPS, SCENARIOS, plan  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402


def eff_cost(p, sc):
    """按"每块 GPU 的效率"算的成本：如果 GPU 可以按小数台租（流量足够大时接近这个值）。"""
    return p.gpus * p.price * (sc.rate / p.capacity()) / (sc.rate * sc.output * 3600) * 1e6


def describe(p):
    """方案的关键参数（批大小、块大小、单实例吞吐）。"""
    if p.mode == "colo":
        c = p.detail["colo"]
        return f"批 ≤ {c['batch']}、块 {c['chunk']}、每实例 {c['rps']:.1f} 请求/s、一步 {c['t_step'] * 1e3:.0f} ms"
    pc, dc = p.detail["prefill"], p.detail["decode"]
    return f"P：每实例 {pc['rps']:.1f} 请求/s（利用率 ≤ {pc['rho']:.0%}）；D：批 ≤ {dc['batch']}、每实例 {dc['rps']:.1f} 请求/s"


def demo1_scenarios():
    print_header("演示 1：四个场景各自最便宜的方案（解析模型筛选）")
    best = {}
    for key, sc in SCENARIOS.items():
        ps = plan(sc, tail_aware=True)
        best[key] = ps
        print(f"\n【{sc.name}】{sc.model.name}；{sc.note}；目标流量 {sc.rate:g} 请求/s；P90 提示长度 {sc.prompt_quantile(0.9):,}")
        rows = [[p.label(), str(p.gpus), f"{p.usd_per_mtok:.3f}", f"{p.capacity() / sc.rate:.2f}", describe(p)] for p in ps[:5]]
        print_table(["方案", "GPU", "美元 / 百万输出 token", "解析容量 / 目标流量", "关键参数"], rows)
    chat, agent, fast = best["chat"][0], best["agent"][0], best["fast"][0]
    print_observations([
        f"聊天：最便宜的是 {chat.label()}（{chat.gpus} 块 GPU）；几个方案成本一样——流量只有 20 请求/s 时，GPU 的\"整数\"决定了成本（演示 2）。",
        f"写代码的智能体（提示长 8 倍）：每百万输出 token 的成本约是聊天的 {agent.usd_per_mtok / chat.usd_per_mtok:.1f} 倍——"
        "提示越长，每个输出 token 要分摊的 prefill 越多，KV 也越大（批更小）。",
        f"低延迟的 8B：TPOT ≤ 15 ms 时一块 H100 就够（{fast.label()}），成本最低；"
        "小模型、严格的延迟，是\"单卡 + 小批\"的地盘。",
        "在这些场景和这组假设价格下，共置（分块 prefill）几乎总是比解耦便宜：分块 prefill 已经控制住了干扰，"
        "解耦还要多付 KV 传输和两边各自取整的代价。第 5 讲看到，解耦的优势在 SLO 很严、prefill 和 decode 的最优并行方式差别很大的时候才明显。",
    ])
    return best


def demo2_integer(best):
    print_header("演示 2：按整数台算的成本 vs 每块 GPU 的效率（聊天场景）")
    sc = SCENARIOS["chat"]
    ps = best["chat"]
    rows = []
    for p in sorted(ps, key=lambda p: eff_cost(p, sc))[:8]:
        rows.append([p.label(), str(p.gpus), f"{p.usd_per_mtok:.3f}", f"{eff_cost(p, sc):.3f}", f"{p.capacity() / sc.rate:.2f}"])
    print_table(["方案（按效率排序）", "GPU", "按整数台（美元/百万）", "按效率（美元/百万）", "容量 / 流量"], rows)
    e = sorted(ps, key=lambda p: eff_cost(p, sc))
    print_observations([
        f"按每块 GPU 的效率排，最好的是 {e[0].label()}（{eff_cost(e[0], sc):.3f} 美元/百万）；按整数台算，它要 {e[0].gpus} 块 GPU，"
        f"容量是目标流量的 {e[0].capacity() / sc.rate:.1f} 倍——多出来的都是闲置。",
        "流量小的时候，\"凑整\"的浪费可能比方案之间效率的差别还大：要么挑一个容量刚好合适的配置，要么把多个模型 / 多个租户放到同一组 GPU 上"
        "（多路复用，CS349D 第 12 讲），要么按负载弹性伸缩。流量很大时，两种算法趋于一致。",
    ])


def demo3_by_family(best):
    print_header("演示 3：每种 GPU / 精度 / 模式的最好方案（按效率算的成本，美元 / 百万输出 token）")
    rows = []
    for key, sc in SCENARIOS.items():
        cell = {}
        for p in best[key]:
            k = (p.chip.name, p.prec, p.mode)
            if k not in cell or eff_cost(p, sc) < eff_cost(cell[k], sc):
                cell[k] = p
        row = [sc.name]
        for chip in CHIPS:
            for prec in ("bf16", "fp8"):
                vals = [f"{eff_cost(cell[(chip.name, prec, m)], sc):.2f}" if (chip.name, prec, m) in cell else "—"
                        for m in ("colo", "disagg")]
                row.append(" / ".join(vals))
        rows.append(row)
    print_table(["场景"] + [f"{c.name.split('(')[0].split('-')[0]} {p.upper()}（共置 / 解耦）" for c in CHIPS for p in ("bf16", "fp8")], rows)
    h100 = [float(r[2].split(" / ")[0]) for r in rows]          # H100 FP8 共置
    h200 = [float(r[4].split(" / ")[0]) for r in rows]          # H200 FP8 共置
    dash = any("—" in c for r in rows for c in r[1:])
    obs = [
        "FP8 几乎在每一格都比 BF16 便宜：读权重和读 KV 的字节减半，decode 的批可以更大（第 3 讲）。前提是质量允许——"
        "Scenario.precisions 只放 \"bf16\" 就能模拟\"不许量化\"的要求（04 脚本）。",
        "按这组假设的价格，B200 的效率最高：它的带宽是 H100 的 2.4 倍、显存 2.4 倍、算力 2.5 倍，价格只贵 1 倍。"
        + (f"H200 和 H100 是同一颗计算芯片：H200 贵 30%，多出来的带宽和显存只帮了 decode（更大的批），prefill 的算力一样，"
           f"所以按效率算它在这四个场景里都比 H100 略贵（FP8 共置：{' / '.join(f'{b:.2f}' for b in h200)} vs {' / '.join(f'{a:.2f}' for a in h100)}）。"
           if all(b > a for a, b in zip(h100, h200)) else "H200 和 H100 是同一颗计算芯片，多出来的带宽和显存只帮了 decode。")
        + "价格一变，排序就变（04 脚本）。",
    ]
    if dash:
        obs.append("\"—\"表示这个组合满足不了 SLO（例如单卡放不下权重，或者 prefill 太慢）。")
    print_observations(obs)


if __name__ == "__main__":
    b = demo1_scenarios()
    demo2_integer(b)
    demo3_by_family(b)