"""
=============================================================================
第 10 讲 · 代码 5：硬件故障与"有效训练时间"（goodput）——多久存一次检查点
=============================================================================

对应讲义：lectures/lecture10_training_in_practice.md 第 5 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
一张 GPU 平均几年才坏一次，可是一万张 GPU 放在一起，每隔几个小时就有一张坏掉。每次故障，从上一个检查点到故障
之间的训练都白做了，还要花时间重启。检查点存得越勤，丢的越少，但存检查点本身也要时间。

    演示 1：从三份训练报告反推"每张卡平均多久坏一次"（MTBF），再看集群越大、故障越频繁。
    演示 2：模拟一次 54 天的训练：故障随机发生（泊松过程），每隔 τ 存一次检查点。不同的 τ 下，
            有效训练时间占多少？和 Young–Daly 公式 τ* ≈ √(2δM) 给出的最优间隔对比。
    演示 3：把第 8 讲的硬件利用率（MFU）和这里的有效时间乘起来：一次训练实际用掉的算力里，有多少变成了模型。

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_failures_and_goodput.py      # 约 2 秒；只用 Python 标准库（random、math）

【记号】
    M : 整个集群的平均故障间隔时间（mean time between failures，小时）
    τ : 两次检查点之间的训练时间（小时）
    δ : 存一次检查点要停多久（小时）
    R : 一次故障之后，发现、换机器、重新加载检查点要多久（小时）

【Python 小知识】
    - random.Random(seed).expovariate(λ)：指数分布的随机数，均值 1/λ。故障之间的间隔服从指数分布 ⟺ 故障是泊松过程。
    - math.fsum(列表)：精确地求一组浮点数的和（不受加法顺序影响），保证不同 Python 版本输出相同。
"""

import math
import os
import random
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from console import pad, print_header  # noqa: E402


# =============================================================================
# 演示 1：MTBF
# =============================================================================
def demo1_mtbf():
    print_header("演示 1：从训练报告反推每张卡的平均故障间隔（MTBF），以及集群越大故障越频繁")
    reports = [
        ("Llama 3 405B（2024）", "16,384 张 H100，54 天里 419 次意外中断", 16384, 54 * 24 / 419),
        ("BLOOM 176B（2022）", "约 400 张 A100，平均每周 1–2 次 GPU 故障（取 1.5）", 400, 7 * 24 / 1.5),
        ("OPT-175B（2021）", "1,024 张 A100（128 台机器），节假日期间约每天 2 台机器故障", 1024, 24 / 2),
    ]
    print(f"{pad('训练', 22)}{pad('报告', 58)}{pad('集群 MTBF', 12, '>')}{pad('每张卡 MTBF', 14, '>')}")
    per_gpu = []
    for name, what, n, mtbf_cluster in reports:
        mtbf_gpu = mtbf_cluster * n
        per_gpu.append(mtbf_gpu)
        print(f"{pad(name, 22)}{pad(what, 58)}{pad(f'{mtbf_cluster:.1f} 小时', 12, '>')}{pad(f'{mtbf_gpu / 24 / 365:.1f} 年', 14, '>')}")
    print("\n如果每张卡的 MTBF 是 5 年（约 43,800 小时），不同规模的集群多久出一次故障：")
    for n in [8, 512, 4096, 16384, 100000]:
        m = 5 * 365 * 24 / n
        print(f"  {pad(f'{n:,} 张卡', 14)}平均每 {m:,.1f} 小时一次" + (f"（{m * 60:.0f} 分钟）" if m < 2 else ""))
    print("\n观察 Observations：")
    print("  1. Llama 3 和 BLOOM 的数字给出的每张卡 MTBF 都在 5–6 年；OPT 的数字（按机器算、而且是故障最频繁的节假日期间）")
    print("     低得多——2021 年的云上 A100 集群和后来专门建设的集群，可靠性差别很大。")
    print("  2. 单张卡很可靠，但故障率随卡数线性增加：8 张卡几乎从不出事，1.6 万张卡每 2.7 小时一次，")
    print("     10 万张卡每 26 分钟一次。规模大到一定程度，'训练会被打断'不是意外，而是常态。")
    return per_gpu


# =============================================================================
# 演示 2：模拟 goodput
# =============================================================================
def simulate(days, mtbf, tau, delta, restart, seed):
    """
    模拟一次训练：故障是速率 1/mtbf 的泊松过程；每训练 tau 小时停 delta 小时存检查点；
    故障发生时，从上一个检查点之后做的训练全部作废，再花 restart 小时恢复。

    【思路】按时间往前走：当前阶段要么是"训练 tau 小时 + 存检查点 delta 小时"（成功的话这 tau 小时算有效），
            要么在中途被一次故障打断（这一段的训练全部作废，再加上 restart）。用指数分布抽"下一次故障还有多久"。
    【参数】days：总的墙上时间（天）；mtbf、tau、delta、restart：小时；seed：随机种子
    【返回】dict(goodput=有效训练时间 / 总时间, failures=故障次数, lost=作废的训练小时数)
    【边界】假设恢复期间和存检查点期间不会再出故障（它们比 M 短得多时影响很小）；最后一段不完整的训练不算有效。
    """
    rng = random.Random(seed)
    total = days * 24.0
    t, useful, lost, failures = 0.0, [], [], 0
    next_fail = rng.expovariate(1 / mtbf)
    while t < total:
        segment = tau + delta
        if next_fail >= segment:                   # 这一段顺利完成：训练 tau 小时 + 存检查点
            useful.append(min(tau, max(0.0, total - t)))
            t += segment
            next_fail -= segment
        else:                                      # 中途出故障：这一段作废
            lost.append(min(next_fail, tau))
            t += next_fail + restart
            failures += 1
            next_fail = rng.expovariate(1 / mtbf)
    return dict(goodput=math.fsum(useful) / total, failures=failures, lost=math.fsum(lost))


def expected_goodput(mtbf, tau, delta, restart):
    """
    goodput 的近似公式：1 − δ/(τ + δ) − (τ/2 + R)/M（τ ≪ M 时）。

    【思路】三项损失：存检查点占的时间比例 δ/(τ + δ)；每次故障平均丢掉半个间隔的训练（τ/2）加上恢复时间 R，
            故障频率 1/M。对 τ 求导、令导数为 0，得到 Young (1974) / Daly (2006) 的最优间隔 τ* ≈ √(2δM)。
    """
    return 1 - delta / (tau + delta) - (tau / 2 + restart) / mtbf


def demo2_goodput():
    print_header("演示 2：54 天的训练，多久存一次检查点（模拟 5 次取平均）")
    settings = [("Llama 3 规模：M = 3.1 小时，δ = 3 分钟，R = 15 分钟", 54 * 24 / 419, 0.05, 0.25),
                ("BLOOM 规模：M = 112 小时，δ = 10 分钟，R = 1 小时", 7 * 24 / 1.5, 1 / 6, 1.0)]
    for title, mtbf, delta, restart in settings:
        tau_star = math.sqrt(2 * delta * mtbf)
        print(f"\n{title}；Young–Daly 最优间隔 τ* = √(2δM) = {tau_star:.2f} 小时")
        print(f"{pad('检查点间隔 τ', 16)}{pad('故障次数', 10, '>')}{pad('作废的训练', 14, '>')}{pad('模拟 goodput', 14, '>')}{pad('公式', 10, '>')}")
        for tau in sorted({0.1, 0.25, 0.5, 1.0, 3.0, 6.0, round(tau_star, 2)}):
            runs = [simulate(54, mtbf, tau, delta, restart, seed) for seed in range(5)]
            g = math.fsum(r["goodput"] for r in runs) / len(runs)
            f = math.fsum(r["failures"] for r in runs) / len(runs)
            lost = math.fsum(r["lost"] for r in runs) / len(runs)
            mark = "  ← τ*" if tau == round(tau_star, 2) else ""
            print(f"{pad(f'{tau:g} 小时', 16)}{pad(f'{f:.0f}', 10, '>')}{pad(f'{lost:.0f} 小时', 14, '>')}{pad(f'{g:.1%}', 14, '>')}"
                  f"{pad(f'{expected_goodput(mtbf, tau, delta, restart):.1%}', 10, '>')}{mark}")
    print("\n观察 Observations：")
    print("  1. 检查点存得太少（τ 大），每次故障丢掉的训练多；存得太勤（τ 小），存检查点本身占掉太多时间。")
    print("     最优间隔大约是 √(2δM)：Llama 3 规模的集群要每 30 分钟左右存一次，BLOOM 规模的每 6 小时左右。")
    print("     BLOOM 实际每 3 小时存一次，训练日志说每次故障平均丢 1.5 小时——正好是半个间隔。")
    print("  2. 即使用最优间隔，Llama 3 规模的集群也会损失约 23% 的时间（goodput 77%）；BLOOM 规模的只损失约 6%。")
    print("     损失由三部分组成：存检查点（δ/τ）、作废的训练（每次故障约半个间隔）、恢复（R × 故障次数，这里 381 次 × 15 分钟）。")
    print("     Llama 3 报告的有效训练时间'超过 90%'，说明他们的 δ 和 R 比这里假设的更小——")
    print("     把检查点写得更快（异步写、分片写）、把恢复做得更快（自动检测、备用机器），是大规模训练的核心工程。")
    print("  3. τ 远小于 M 时，公式和模拟基本一致；τ 接近或超过 M 时公式就不准了（它假设一个间隔里最多出一次故障，")
    print("     τ = 6 小时、M = 3.1 小时时甚至给出负数）。")


# =============================================================================
# 演示 3：把利用率和有效时间乘起来
# =============================================================================
def demo3_ledger():
    print_header("演示 3：一次训练实际用掉的算力里，有多少变成了模型")
    mfu, goodput = 0.462, 0.85
    print(f"假设 MFU = {mfu:.1%}（PaLM，第 8 讲 §7），有效训练时间 = {goodput:.0%}：")
    print(f"  有用的 FLOPs ÷ 买来的峰值 FLOPs × 时间 = {mfu:.3f} × {goodput:.2f} = {mfu * goodput:.1%}")
    print("第 8 讲 §7.3 的对照：PaLM 按实测吞吐不间断训练需要 37.9 天，实际折合 57.0 天（1.50 倍）——")
    print(f"  相当于有效时间约 {1 / 1.5:.0%}（包括 batch 预热阶段吞吐低、停机、尖峰回滚），总效率约 {mfu / 1.5:.0%}。")
    print("\n观察 Observations：")
    print("  1. 大模型训练里，买来的算力只有大约三分之一到五分之二变成了'模型必须做的运算'：")
    print("     一半多输在单卡效率上（第 8 讲），再有一部分输在故障、重启、回滚上（本讲）。")
    print("  2. 这也是为什么训练报告越来越重视'有效训练时间'这个指标：规模越大，它越难保持。")


if __name__ == "__main__":
    demo1_mtbf()
    demo2_goodput()
    demo3_ledger()