"""
=============================================================================
CS349E · 第 5 讲共享模块：serving_sim.py —— LLM 推理服务的离散事件模拟器：共置 vs 解耦
=============================================================================

【这个模块要解决什么问题】
    前四讲回答的都是"一步 prefill / 一步 decode 要多久"。可真实的服务里，请求随时到达，长短不一，
    prefill 和 decode 要么挤在同一组 GPU 上（共置，colocation），要么分开放（解耦，disaggregation）。
    用户感受到的是：
        TTFT（time to first token）：从请求到达到第一个 token 出来；
        TPOT（time per output token）：之后每个 token 的平均间隔。
    服务方关心的是：在"90% 的请求同时满足 TTFT 和 TPOT 的要求（SLO）"的前提下，每块 GPU 每秒能服务多少请求
    ——DistServe 把它叫做 **per-GPU goodput**。这些量没有简单的公式（排队、批处理、干扰交织在一起），
    DistServe 的做法是写一个模拟器去算。本模块就是这样一个模拟器。

【模拟器里有什么】
    1. 每步时间的模型 StepModel：用第 1 讲的 roofline（FLOPs / 峰值、字节 / 带宽）加效率折扣、每步固定开销、
       张量并行的 all-reduce（前置讲 P5 的 α–β 公式）。
    2. 三种"实例"（instance）：
         ColocatedInstance：prefill 和 decode 共用 GPU，迭代级调度（第 6 讲的连续批处理）。两种策略：
             "prefill_first"：有新请求就先做 prefill，正在 decode 的请求等着（vLLM 早期的默认行为）；
             "chunked"：每一步 = 所有 decode + 一块不超过预算的 prefill（Sarathi-Serve 的分块 prefill）。
         PrefillInstance：只做 prefill。可以用流水线并行（pp 段）：新的一批每 D/pp 秒就能开始，一批要 D 秒做完。
         DecodeInstance：只做 decode，连续批处理；KV cache 有容量上限；pp 段时相当于 pp 条独立的"车道"。
    3. 请求的流转：到达 → 选一个 prefill（或共置）实例排队 → prefill 完成（第一个 token）→（解耦时）通过网络把
       KV cache 搬到 decode 实例 → decode 实例在迭代边界接纳它 → 生成完剩下的 token。
    4. 指标：TTFT、TPOT、SLO 达成率；goodput 搜索（对请求速率做二分查找）。

【模型的简化（务必知道）】
    - 每步时间是 roofline 的估计（带效率折扣），不是测量值；用来比较方案，不预测真实系统的绝对数字。
    - KV cache 按"提示 + 全部输出"预留（不做抢占 / 换出，第 7 讲再讲）。
    - 不模拟 KV 传输的网络争用（每个请求单独计 α + 字节 / 带宽）。
    - 调度都是先到先服务（FCFS）。

【运行】python3 serving_sim.py   # 自检
【依赖】numpy；从 ../lecture01 导入 hw_specs、llm_math；从 ../prereq5 导入 comm_sim（α–β 公式与链路参数）。
"""

import heapq
import math
import os
import sys
from dataclasses import dataclass

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
for sub in ("lecture01", "prereq5"):
    sys.path.insert(0, os.path.join(HERE, "..", sub))

from comm_sim import IB_NDR, NVLINK4, t_formula  # noqa: E402


# ---------------------------------------------------------------------------
# 1. 每步时间的模型
# ---------------------------------------------------------------------------
@dataclass
class StepModel:
    """
    一个实例（tp 块 GPU 做张量并行、pp 段流水线）跑一步要多久。

    【字段】
        model, chip：llm_math.ModelConfig、hw_specs.Chip
        tp, pp     ：张量并行度、流水线段数（一个实例一共 tp × pp 块 GPU）
        w_bytes    ：权重每个数的字节；kv_bytes：KV 每个数的字节
        mfu, mbu   ：能达到的峰值算力比例、峰值带宽比例（假设值，真实的大矩阵乘法约 0.6–0.8）
        overhead   ：每一步的固定开销（秒）：调度、内核启动、采样……（假设值）
        link       ：张量并行用的链路（默认 NVLink 4）
        mem_reserve：显存里预留给激活、工作区的比例
    【思路】一步的时间 = max(计算时间, 读写时间) + 张量并行的通信 + 固定开销。
        流水线并行时，一批数据要依次走过 pp 段：**一批的延迟**是整个模型的时间（每段 1/pp 的层、
        tp 块 GPU 的带宽），而每段只占 1/pp 的时间——所以新的一批每 1/pp 的时间就能进来（吞吐 ×pp）。
    """
    model: object
    chip: object
    tp: int = 1
    pp: int = 1
    w_bytes: float = 2.0
    kv_bytes: float = 2.0
    mfu: float = 0.6
    mbu: float = 0.8
    overhead: float = 1e-3
    link: object = NVLINK4
    mem_reserve: float = 0.1

    @property
    def n_gpus(self):
        return self.tp * self.pp

    def _peak(self):
        return self.chip.peak_flops("bf16" if self.w_bytes >= 2 else "fp8") * self.mfu * self.tp

    def _bw(self):
        return self.chip.mem_bw * self.mbu * self.tp

    def _comm(self, tokens):
        """张量并行每层两次 all-reduce（消息 = tokens × d × 2 字节），用 ring 和递归倍增里更快的那个。"""
        if self.tp == 1:
            return 0.0
        msg = tokens * self.model.d_model * 2.0
        one = min(t_formula("ring", self.tp, msg, self.link.alpha, self.link.bw),
                  t_formula("rd", self.tp, msg, self.link.alpha, self.link.bw))
        return 2 * self.model.n_layers * one

    def weight_bytes(self):
        return self.model.matmul_params() * self.w_bytes

    def prefill_time(self, prompt_lens):
        """
        一批提示一起 prefill 的时间（整个模型走一遍）。

        【参数】prompt_lens：这一批里每个提示的长度（列表）
        【返回】秒（一批的延迟；流水线时每段占 1/pp）
        【公式】FLOPs = Σ L × 每 token 的矩阵乘法 FLOPs + 注意力 attn(1)·Σ L(L+1)/2；字节 = 权重 + 写 KV。
        """
        m = self.model
        tokens = float(sum(prompt_lens))
        flops = tokens * m.linear_flops_per_token() + m.attn_flops_per_token(1) * sum(L * (L + 1) / 2 for L in prompt_lens)
        by = self.weight_bytes() + tokens * m.kv_bytes_per_token(self.kv_bytes)
        return max(flops / self._peak(), by / self._bw()) + self._comm(tokens) + self.overhead

    def decode_time(self, batch, ctx_sum):
        """
        一步 decode：batch 个序列各生成 1 个 token，它们的上下文长度之和是 ctx_sum。

        【返回】秒
        【公式】FLOPs = batch × 矩阵乘法 FLOPs + attn(1) × ctx_sum；字节 = 权重 + ctx_sum × 每 token KV。
        """
        if batch == 0:
            return 0.0
        m = self.model
        flops = batch * m.linear_flops_per_token() + m.attn_flops_per_token(1) * ctx_sum
        by = self.weight_bytes() + ctx_sum * m.kv_bytes_per_token(self.kv_bytes)
        return max(flops / self._peak(), by / self._bw()) + self._comm(batch) + self.overhead

    def mixed_time(self, batch, ctx_sum, chunk_tokens, chunk_attn):
        """
        分块 prefill 的一步：batch 个 decode + 一块 chunk_tokens 个 prefill token（注意力运算 chunk_attn FLOPs）。
        """
        m = self.model
        tokens = batch + chunk_tokens
        flops = tokens * m.linear_flops_per_token() + m.attn_flops_per_token(1) * ctx_sum + chunk_attn
        by = self.weight_bytes() + (ctx_sum + chunk_tokens) * m.kv_bytes_per_token(self.kv_bytes)
        return max(flops / self._peak(), by / self._bw()) + self._comm(tokens) + self.overhead

    def kv_capacity_tokens(self):
        """
        这个实例的显存（tp × pp 块 GPU）放下权重之后还能放多少个 token 的 KV。

        【边界】放不下权重时返回 0。张量并行度超过 KV 组数时 KV 会被复制（第 4 讲 §4.4），这里按复制后的大小算。
        """
        free = self.n_gpus * self.chip.mem_cap * (1 - self.mem_reserve) - self.model.n_params() * self.w_bytes
        if free <= 0:
            return 0
        rep = max(1.0, self.tp / self.model.n_kv_heads)
        return int(free / (self.model.kv_bytes_per_token(self.kv_bytes) * rep))


# ---------------------------------------------------------------------------
# 2. 请求与工作负载
# ---------------------------------------------------------------------------
@dataclass
class Request:
    """一个请求。prompt：提示长度；output：要生成的 token 数（≥ 1，第一个由 prefill 产生）。"""
    rid: int
    arrival: float
    prompt: int
    output: int
    first_token: float = math.nan     # prefill 完成、第一个 token 出来的时间
    finish: float = math.nan          # 最后一个 token 出来的时间
    kv_ready: float = math.nan        # （解耦时）KV 搬到 decode 实例的时间

    @property
    def ttft(self):
        return self.first_token - self.arrival

    @property
    def tpot(self):
        """除第一个 token 以外，每个 token 的平均间隔（output = 1 时定义为 0）。"""
        return 0.0 if self.output <= 1 else (self.finish - self.first_token) / (self.output - 1)


def make_workload(n, rate, prompt, output, seed=0, arrivals="poisson", cv=1.0):
    """
    生成 n 个请求。默认到达过程是泊松过程（相邻到达间隔服从指数分布，平均 1/rate 秒）。

    【参数】prompt、output：整数（所有请求一样长），或者函数 f(rng) → 整数（每个请求单独抽）
        arrivals："poisson"（指数间隔）、"uniform"（等间隔，完全不随机）、"bursty"（伽马分布间隔，变异系数 cv > 1 时
            比泊松更"扎堆"）
    【返回】Request 列表（按到达时间排序）
    【边界】同样的 seed 下，到达间隔是"单位速率的间隔 / rate"——改变 rate 时请求的长度和相对顺序不变，
        只是整体被压缩或拉伸。这让"SLO 达成率随速率变化"的曲线更平滑（goodput 二分查找需要它）。
    【Python 小知识】rng.exponential(1.0, n) 抽 n 个均值为 1 的指数分布随机数；rng.gamma(k, θ, n) 抽伽马分布
        （均值 kθ、变异系数 1/√k）；np.cumsum 求累加和。
    """
    rng = np.random.default_rng(seed)
    if arrivals == "poisson":
        gaps = rng.exponential(1.0, n) / rate
    elif arrivals == "uniform":
        gaps = np.full(n, 1.0 / rate)
    elif arrivals == "bursty":
        k = 1.0 / cv ** 2
        gaps = rng.gamma(k, 1.0 / k, n) / rate
    else:
        raise ValueError(arrivals)
    t = np.cumsum(gaps)
    reqs = []
    for i in range(n):
        p = prompt(rng) if callable(prompt) else int(prompt)
        o = output(rng) if callable(output) else int(output)
        reqs.append(Request(i, float(t[i]), max(1, p), max(1, o)))
    return reqs


def lognormal_len(mean, sigma, lo, hi):
    """返回一个函数：按对数正态分布抽长度（均值约为 mean），截断到 [lo, hi]。用来模拟"长短不一"的真实请求。"""
    mu = math.log(mean) - sigma ** 2 / 2

    def f(rng):
        return int(min(hi, max(lo, round(rng.lognormal(mu, sigma)))))
    return f


# ---------------------------------------------------------------------------
# 3. 事件驱动的模拟
# ---------------------------------------------------------------------------
class Sim:
    """
    一个最小的离散事件模拟器：事件按时间排序放在堆里，依次取出处理。

    【Python 小知识】heapq 是"最小堆"：heappush 放入、heappop 取出最小的元素。
        事件是 (时间, 序号, 回调函数)：序号保证时间相同时按放入的顺序处理（也避免比较函数对象）。
    """

    def __init__(self):
        self.now = 0.0
        self._q = []
        self._seq = 0

    def at(self, t, fn):
        heapq.heappush(self._q, (t, self._seq, fn))
        self._seq += 1

    def run(self):
        while self._q:
            t, _, fn = heapq.heappop(self._q)
            self.now = t
            fn()


class PrefillInstance:
    """
    只做 prefill 的实例。先到先服务；一批最多 max_tokens 个提示 token（至少一个请求）。

    【流水线】pp > 1 时，一批要 D 秒做完，但第一段只占 D/pp 秒：第一段空出来就能开始下一批。
    【回调】on_done(request)：prefill 完成（第一个 token 已产生）时调用。
    """

    def __init__(self, sim, step, on_done, max_tokens=4096):
        self.sim, self.step, self.on_done, self.max_tokens = sim, step, on_done, max_tokens
        self.queue = []
        self.busy = False

    def load(self):
        return sum(r.prompt for r in self.queue)

    def submit(self, r):
        self.queue.append(r)
        if not self.busy:
            self._start()

    def _start(self):
        if not self.queue:
            self.busy = False
            return
        batch, tok = [], 0
        while self.queue and (not batch or tok + self.queue[0].prompt <= self.max_tokens):
            r = self.queue.pop(0)
            batch.append(r)
            tok += r.prompt
        D = self.step.prefill_time([r.prompt for r in batch])
        self.busy = True
        done_t = self.sim.now + D
        free_t = self.sim.now + D / self.step.pp
        for r in batch:
            self.sim.at(done_t, lambda r=r, t=done_t: self._finish(r, t))
        self.sim.at(free_t, self._start)

    def _finish(self, r, t):
        r.first_token = t
        self.on_done(r)


class DecodeLane:
    """
    decode 实例里的一条"车道"：连续批处理，每一步给车道里所有请求各生成 1 个 token。

    【思路】请求在迭代边界加入（KV 容量够时），生成完 output − 1 个 token 后离开。
        用一个小顶堆记录"每个请求在第几步结束"，这样每一步的工作量和批大小无关（只更新几个累加量）。
    【字段】ctx_sum：车道里所有请求当前上下文长度之和（决定读 KV 的字节）；reserved：预留的 KV token 数。
    """

    def __init__(self, sim, step, capacity):
        self.sim, self.step, self.capacity = sim, step, capacity
        self.waiting = []
        self.running = {}                     # rid → (request, 结束的步号)
        self.finish_heap = []
        self.iter = 0
        self.ctx_sum = 0
        self.reserved = 0
        self.busy = False

    def load(self):
        return len(self.running) + len(self.waiting)

    def submit(self, r):
        self.waiting.append(r)
        if not self.busy:
            self._start()

    def _admit(self):
        while self.waiting:
            r = self.waiting[0]
            need = r.prompt + r.output
            if self.reserved + need > self.capacity and self.running:
                break
            self.waiting.pop(0)
            if r.output <= 1:
                r.finish = r.first_token
                continue
            self.reserved += need
            self.ctx_sum += r.prompt + 1
            end = self.iter + (r.output - 1)
            self.running[r.rid] = (r, end)
            heapq.heappush(self.finish_heap, (end, r.rid))

    def _start(self):
        self._admit()
        if not self.running:
            self.busy = False
            return
        self.busy = True
        B = len(self.running)
        D = self.step.decode_time(B, self.ctx_sum)
        self.sim.at(self.sim.now + D, self._finish_step)

    def _finish_step(self):
        self.iter += 1
        self.ctx_sum += len(self.running)
        while self.finish_heap and self.finish_heap[0][0] <= self.iter:
            _, rid = heapq.heappop(self.finish_heap)
            r, _ = self.running.pop(rid)
            r.finish = self.sim.now
            self.ctx_sum -= r.prompt + r.output
            self.reserved -= r.prompt + r.output
        self._start()


class DecodeInstance:
    """
    只做 decode 的实例：pp 条车道（流水线的每个微批一条），请求分给负载最小的车道。

    【为什么一条车道的一步 = tp 块 GPU 上整个模型的一步】一个微批依次走过 pp 段，每段读 1/pp 的权重和自己那部分 KV，
        用 tp 块 GPU 的带宽：加起来正好是"tp 块 GPU 读整个模型"的时间。pp 个微批交替占用各段，所有段一直在忙。
        和 pp 个独立副本相比，区别只在显存：每块 GPU 只存 1/pp 的权重，KV 的空间更大。
    """

    def __init__(self, sim, step):
        cap = step.kv_capacity_tokens() // step.pp
        self.lanes = [DecodeLane(sim, step, cap) for _ in range(step.pp)]

    def load(self):
        return sum(l.load() for l in self.lanes)

    def submit(self, r):
        min(self.lanes, key=lambda l: l.load()).submit(r)


class ColocatedInstance:
    """
    prefill 和 decode 共用 GPU 的实例（连续批处理，迭代级调度）。

    【策略】
        "prefill_first"：只要有等待的新请求（且 KV 放得下），这一步就只做它们的 prefill（最多 max_tokens 个 token），
            正在 decode 的请求这一步停住——decode 的 TPOT 被 prefill 打断；
        "chunked"：每一步 = 全部 decode + 一块 prefill，总 token 数不超过 chunk_budget（Sarathi-Serve）。
            一个长提示要分几步做完；decode 不会被整段 prefill 卡住，但每一步都变长了一些。
    """

    def __init__(self, sim, step, policy="prefill_first", max_tokens=4096, chunk_budget=512):
        self.sim, self.step, self.policy = sim, step, policy
        self.max_tokens, self.chunk_budget = max_tokens, chunk_budget
        self.capacity = step.kv_capacity_tokens()
        self.queue = []                       # 等待 prefill 的请求
        self.partial = None                   # chunked：正在分块 prefill 的请求和已完成的 token 数
        self.lane = DecodeLane(sim, step, self.capacity)
        self.busy = False

    def load(self):
        return len(self.queue) + self.lane.load()

    def submit(self, r):
        self.queue.append(r)
        if not self.busy:
            self._start()

    def _kv_fits(self, r):
        """KV 还放不放得下 r：已接纳的 + prefill 完、等着接纳的 + r 自己。空的时候总是放得下（至少服务一个）。"""
        pending = sum(x.prompt + x.output for x in self.lane.waiting)
        if not self.lane.running and not self.lane.waiting:
            return True
        return self.lane.reserved + pending + r.prompt + r.output <= self.capacity

    def _start(self):
        self.lane._admit()
        if self.policy == "prefill_first" and self.queue and self._kv_fits(self.queue[0]):
            batch, tok = [], 0
            while self.queue and (not batch or tok + self.queue[0].prompt <= self.max_tokens) and self._kv_fits(self.queue[0]):
                r = self.queue.pop(0)
                batch.append(r)
                tok += r.prompt
            D = self.step.prefill_time([r.prompt for r in batch])
            self.busy = True
            self.sim.at(self.sim.now + D, lambda b=batch: self._prefill_done(b))
            return
        if self.policy == "chunked":
            self._start_chunked()
            return
        if self.lane.running:
            self.busy = True
            D = self.step.decode_time(len(self.lane.running), self.lane.ctx_sum)
            self.sim.at(self.sim.now + D, self._decode_done)
            return
        self.busy = False

    def _prefill_done(self, batch):
        for r in batch:
            r.first_token = self.sim.now
            self.lane.waiting.append(r)
        self._start()

    def _decode_done(self):
        self.lane.busy = False
        self.lane._finish_step_no_restart()
        self._start()

    def _start_chunked(self):
        B = len(self.lane.running)
        budget = self.chunk_budget - B
        chunk, attn, done_reqs = 0, 0.0, []
        m = self.step.model
        while budget > 0 and (self.partial is not None or (self.queue and self._kv_fits(self.queue[0]))):
            if self.partial is None:
                self.partial = [self.queue.pop(0), 0]
            r, done = self.partial
            take = min(budget, r.prompt - done)
            # 这一块的注意力：第 done+1 … done+take 个 token 各自看前面所有位置
            attn += m.attn_flops_per_token(1) * (take * done + take * (take + 1) / 2)
            chunk += take
            budget -= take
            done += take
            if done >= r.prompt:
                done_reqs.append(r)
                self.partial = None
            else:
                self.partial = [r, done]
        if B == 0 and chunk == 0:
            self.busy = False
            return
        D = self.step.mixed_time(B, self.lane.ctx_sum, chunk, attn)
        self.busy = True
        self.sim.at(self.sim.now + D, lambda d=done_reqs: self._chunk_done(d))

    def _chunk_done(self, done_reqs):
        if self.lane.running:
            self.lane._finish_step_no_restart()
        for r in done_reqs:
            r.first_token = self.sim.now
            self.lane.waiting.append(r)
        self._start()


def _finish_step_no_restart(self):
    """DecodeLane 的一步结束，但不自动开始下一步（共置实例自己决定下一步做什么）。"""
    self.iter += 1
    self.ctx_sum += len(self.running)
    while self.finish_heap and self.finish_heap[0][0] <= self.iter:
        _, rid = heapq.heappop(self.finish_heap)
        r, _ = self.running.pop(rid)
        r.finish = self.sim.now
        self.ctx_sum -= r.prompt + r.output
        self.reserved -= r.prompt + r.output


DecodeLane._finish_step_no_restart = _finish_step_no_restart


# ---------------------------------------------------------------------------
# 4. 搭一个系统、跑一遍、算指标
# ---------------------------------------------------------------------------
def kv_transfer_time(model, prompt, link=IB_NDR, kv_bytes=2.0):
    """把一个请求的 KV cache（prompt 个 token）从 prefill 实例搬到 decode 实例：α + 字节 / 带宽。"""
    return link.alpha + prompt * model.kv_bytes_per_token(kv_bytes) / link.bw


def run_colocated(reqs, step, n_instances=1, policy="prefill_first", chunk_budget=512, max_tokens=4096):
    """
    n 个共置实例（每个 step.n_gpus 块 GPU），请求发给负载最小的实例。

    【返回】reqs（字段已填好）
    """
    sim = Sim()
    insts = [ColocatedInstance(sim, step, policy, max_tokens, chunk_budget) for _ in range(n_instances)]
    for r in reqs:
        sim.at(r.arrival, lambda r=r: min(insts, key=lambda i: i.load()).submit(r))
    sim.run()
    return reqs


def run_disaggregated(reqs, p_step, d_step, n_prefill=1, n_decode=1, link=IB_NDR, max_tokens=4096):
    """
    解耦：n_prefill 个 prefill 实例 + n_decode 个 decode 实例。prefill 完成后，KV 通过 link 搬到负载最小的 decode 实例。

    【返回】reqs（字段已填好）
    """
    sim = Sim()
    decs = [DecodeInstance(sim, d_step) for _ in range(n_decode)]

    def to_decode(r):
        t = kv_transfer_time(p_step.model, r.prompt, link, p_step.kv_bytes)
        r.kv_ready = sim.now + t
        sim.at(sim.now + t, lambda r=r: min(decs, key=lambda d: d.load()).submit(r))

    pres = [PrefillInstance(sim, p_step, to_decode, max_tokens) for _ in range(n_prefill)]
    for r in reqs:
        sim.at(r.arrival, lambda r=r: min(pres, key=lambda p: p.load()).submit(r))
    sim.run()
    return reqs


def run_prefill_only(reqs, step, n_instances=1, max_tokens=4096):
    """
    只模拟 prefill（DistServe 算法 1 的 simu_prefill）：n 个 prefill 实例，只关心 TTFT。

    【返回】reqs（first_token 已填好；finish 不填）
    """
    sim = Sim()
    pres = [PrefillInstance(sim, step, lambda r: None, max_tokens) for _ in range(n_instances)]
    for r in reqs:
        sim.at(r.arrival, lambda r=r: min(pres, key=lambda p: p.load()).submit(r))
    sim.run()
    return reqs


def run_decode_only(reqs, step, n_instances=1):
    """
    只模拟 decode（DistServe 算法 1 的 simu_decode）：请求到达时就已经有了 KV 和第一个 token
    （first_token = 到达时间），直接进 decode 实例。只关心 TPOT。
    """
    sim = Sim()
    decs = [DecodeInstance(sim, step) for _ in range(n_instances)]
    for r in reqs:
        r.first_token = r.arrival
        sim.at(r.arrival, lambda r=r: min(decs, key=lambda d: d.load()).submit(r))
    sim.run()
    return reqs


def attainment(reqs, ttft_slo=None, tpot_slo=None, warmup=0.2, cooldown=0.2):
    """
    SLO 达成率：同时满足 TTFT ≤ ttft_slo 和 TPOT ≤ tpot_slo 的请求比例（None 表示不检查这一项）。

    【边界】只统计中间那部分请求：丢掉前 warmup 比例（系统刚启动时是空的，偏乐观）和后 cooldown 比例
        （请求不再到达、系统逐渐排空，也偏乐观）。decode 的请求要活很久（几百个 token × 几十毫秒），
        不去掉两头，模拟出的 goodput 会随请求总数明显变化。
    """
    k0 = int(len(reqs) * warmup)
    k1 = len(reqs) - int(len(reqs) * cooldown)
    rs = reqs[k0:k1]
    ok = 0
    for r in rs:
        good = True
        if ttft_slo is not None and not r.ttft <= ttft_slo:
            good = False
        if tpot_slo is not None and not r.tpot <= tpot_slo:
            good = False
        ok += good
    return ok / len(rs)


def percentile(reqs, attr, q, warmup=0.2, cooldown=0.2):
    """某个指标（"ttft" 或 "tpot"）的 q 分位数（同样只统计中间那部分请求）。"""
    k0 = int(len(reqs) * warmup)
    k1 = len(reqs) - int(len(reqs) * cooldown)
    return float(np.percentile([getattr(r, attr) for r in reqs[k0:k1]], q))


def max_rate(run_fn, target=0.9, lo=0.05, hi=64.0, iters=14):
    """
    二分查找：SLO 达成率 ≥ target 的最大请求速率（请求 / 秒）。

    【参数】run_fn(rate) → 达成率（调用方负责生成工作负载、跑模拟、算达成率）
    【返回】速率；连最小速率都达不到时返回 0
    【边界】假设达成率随速率单调下降（make_workload 的"同一组间隔整体缩放"让这近似成立）。
    """
    if run_fn(lo) < target:
        return 0.0
    if run_fn(hi) >= target:
        return hi
    for _ in range(iters):
        mid = math.sqrt(lo * hi)
        if run_fn(mid) >= target:
            lo = mid
        else:
            hi = mid
    return lo


# ---------------------------------------------------------------------------
# 自检
# ---------------------------------------------------------------------------
def _selftest():
    from hw_specs import H100
    from llm_math import LLAMA31_8B
    step = StepModel(LLAMA31_8B, H100)
    # 1. 单个请求、空闲系统：TTFT = 一次 prefill，TPOT = 一步 decode 的时间
    r = make_workload(1, 1.0, 512, 8)
    run_colocated(r, step)
    exp_ttft = step.prefill_time([512])
    assert abs(r[0].ttft - exp_ttft) < 1e-12, (r[0].ttft, exp_ttft)
    steps = [step.decode_time(1, 513 + i) for i in range(7)]
    assert abs(r[0].tpot - sum(steps) / 7) < 1e-12
    # 2. 解耦：TTFT 相同，decode 开始前多了 KV 传输
    r = make_workload(1, 1.0, 512, 8)
    run_disaggregated(r, step, step)
    assert abs(r[0].ttft - exp_ttft) < 1e-12
    assert abs(r[0].kv_ready - r[0].first_token - kv_transfer_time(LLAMA31_8B, 512)) < 1e-12
    # 3. 所有请求都完成
    reqs = make_workload(300, 5.0, lognormal_len(400, 0.8, 16, 4000), lognormal_len(150, 0.8, 2, 1000), seed=1)
    for fn in (lambda q: run_colocated(q, step), lambda q: run_colocated(q, step, policy="chunked"),
               lambda q: run_disaggregated(q, step, step)):
        out = fn([Request(x.rid, x.arrival, x.prompt, x.output) for x in reqs])
        assert all(not math.isnan(x.finish) for x in out)
    print(f"serving_sim 自检通过：空闲时 TTFT = 一次 prefill（{exp_ttft * 1e3:.2f} ms），TPOT = 平均一步 decode；"
          f"三种系统都处理完了全部请求")


if __name__ == "__main__":
    _selftest()