"""How many GPUs a crowd needs: lesson 29 (how-many-gpus-for-1000-users).

Everything runs inside this one Python process. Nothing here talks to a network or to a real service.

The server is lesson 28's (settings_calc.py), which is lesson 24's continuous-batching server with a pool of cache blocks. The load generator is lesson
27's: random gaps between requests, a fixed seed. The plan has three steps.
  1. Demand. `peak_rate` turns users into new replies a second in the busiest hour.
  2. Limits of one GPU. Speed: the most chats at once for which the p99 time between words stays under a target (`speed_limit`, read off a table
     `step_table` made by running the server with ever more chats at once). Memory: the chats the cache pool holds (`memory_limit`).
  3. Count. One time W serves the whole plan: the seconds a chat holds a batch slot and its cache when a GPU is as full as its tighter limit allows, the
     design point (`stay`, read off the same table; with speed binding, every user sees the p99 target). Little's law gives the chats in progress, rate
     times W. `plan` divides them by what one GPU may carry of each kind, never more than HEADROOM of its limit, and takes the larger count.
     `load_test` then sends random arrivals at that rate through that many GPUs and reports what a user sees.
Each request is one reply of one chat: its prompt is the whole conversation so far, read in full, and lesson 28's `serve` frees the chat's cache blocks at
its last word. The toy keeps no cache between the turns of a conversation.
The widget (../widget.html) repeats `x_at`, `speed_limit`, `stay`, `peak_rate`, `gpus` and `plan` in JavaScript; test_widget.py runs both and compares them.
Python 3 standard library only.
"""
from __future__ import annotations

import importlib.util
import math
import random
import statistics
import sys
from pathlib import Path

sys.dont_write_bytecode = True                       # loading the sibling lessons' files below must not leave a __pycache__ in their folders
LESSONS = Path(__file__).resolve().parents[2]


def _load(name: str, folder: str, filename: str):
    """A sibling lesson folder is not a package, so its file is loaded by path."""
    spec = importlib.util.spec_from_file_location(name, LESSONS / folder / "code" / filename)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module                  # dataclasses looks the module up by name
    spec.loader.exec_module(module)
    return module


L26 = _load("gpu_split_lesson_26", "why-a-big-model-needs-several-gpus", "gpu_split.py")
L27 = _load("tail_sim_lesson_27", "why-average-speed-is-misleading", "tail_sim.py")
L28 = _load("settings_calc_lesson_28", "what-model-server-settings-change", "settings_calc.py")
Request = L28.Request

# --- the invented crowd and its targets -----------------------------------------------------------------------------------------------------
USERS = 1000
ACTIVE = 60                                          # percent of the users who send anything in the busiest hour
REPLIES = 120                                        # replies each of those users receives in that hour: one every 30 seconds
TARGET_MS = 50                                       # a user sees words arrive at most this far apart, for 99 replies in 100 (p99)
TARGETS = [30, 50, 75]                               # the targets the widget offers, in ms
SLIDERS = {"users": (100, 5000, 100), "active": (10, 100, 5), "replies": (10, 300, 10)}       # (lowest, highest, step) of the widget's three sliders
WAIT_TARGET = 2.0                                    # and the first word comes within this many seconds, for 99 replies in 100 (p99)
HEADROOM = 0.8                                       # plan each GPU at no more than this share of what it can carry: 20% in hand for random bursts
SHARE = 90                                           # lesson 28's opening memory share, in percent
UNLIMITED = 10**9                                    # blocks: a pool that never runs out, so the speed table shows speed alone
LOAD_SECONDS = 600                                   # the load test lasts 10 minutes of the busiest hour, shrunk to keep the tests quick

# --- three kinds of chat. Each has the ranges prompts and replies are drawn from, on the log scale of lesson 28, and "each": the replies one active user
# receives in the busiest hour in the lesson's example. A long answer takes over a minute at the design point, so 120 an hour would be a user who asks again
# before the last answer ends; the answers get 30 ---------------------------------------------------------------------------------------------------------
MIXES = {
    "usual": {"label": "Usual chats", "prompts": (300, 6000), "replies": (40, 400), "longest": 8000, "n": 600, "seed": 2026, "each": 120,
              "caps": [1, 2, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256]},
    "documents": {"label": "Long documents, short answers", "prompts": (24000, 31000), "replies": (40, 400), "longest": 32000, "n": 200, "seed": 1, "each": 120,
                  "caps": [1, 2, 3, 4, 6, 8, 12, 16, 24, 32, 48, 64]},
    "answers": {"label": "Long answers", "prompts": (1000, 4000), "replies": (1500, 4000), "longest": 8000, "n": 500, "seed": 1, "each": 30,
                "caps": [8, 16, 32, 64, 96, 128, 192, 256, 384]},
}
OPENING_MIX = "usual"
OPENING = {"users": USERS, "active": ACTIVE, "replies": REPLIES, "target": TARGET_MS, "mix": OPENING_MIX}
LOAD_SEED = 77                                       # gaps between arrivals
CHAT_SEED = 2029                                     # which chat each arrival is


def draw(seed: int, n: int, prompts: tuple, replies: tuple) -> dict:
    """Lesson 28's draw, with the ranges as arguments: with (300, 6000) and (40, 400) it gives lesson 28's 600 chats."""
    rng = random.Random(seed)
    out_p, out_r = [], []
    for _ in range(n):
        out_p.append(round(prompts[0] * (prompts[1] / prompts[0]) ** rng.random()))
        out_r.append(round(replies[0] * (replies[1] / replies[0]) ** rng.random()))
    return {"prompts": out_p, "replies": out_r}


def chats_of(mix: str) -> dict:
    m = MIXES[mix]
    return draw(m["seed"], m["n"], m["prompts"], m["replies"])


# --- step 1: demand -------------------------------------------------------------------------------------------------------------------------

def peak_rate(users: int, active: int, replies: int) -> float:
    """New replies a second in the busiest hour, taking the hour as steady."""
    return users * active * replies / 360000


# --- step 2: what one GPU can carry ---------------------------------------------------------------------------------------------------------

def per_chat_gaps(chats: dict, run: dict) -> list[float]:
    """Seconds between words for each reply: the time from its first word to its last, over the words after the first."""
    return [(done - first) / (reply - 1) for done, first, reply in zip(run["finished"], run["first_token"], chats["replies"])]


_TABLES: dict = {}


def step_table(mix: str) -> list[list[float]]:
    """One row per cap: [chats in progress on average, replies a second, tokens a second, p99 ms between words], from a backlog served with a pool that never
    runs out. A row whose chats in progress did not rise above the last one's (the backlog was too short to fill the cap) is dropped."""
    if mix not in _TABLES:
        m, chats = MIXES[mix], chats_of(mix)
        backlog = L28.backlog(chats)
        rows = []
        for cap in m["caps"]:
            run = L28.serve(backlog, cap, UNLIMITED, m["longest"])
            row = [run["running"], run["served"] / run["seconds"], run["tokens_per_second"], L27.percentile(per_chat_gaps(chats, run), 99) * 1000]
            if not rows or row[0] > rows[-1][0] + 0.5:
                rows.append(row)
        _TABLES[mix] = rows
    return _TABLES[mix]


def x_at(rows: list, chats: float) -> float:
    """Replies a second one GPU finishes with `chats` in progress. Below the first row the line runs from zero; past the last it stays flat."""
    if chats <= rows[0][0]:
        return rows[0][1] * chats / rows[0][0]
    for a, b in zip(rows, rows[1:]):
        if chats <= b[0]:
            return a[1] + (b[1] - a[1]) * (chats - a[0]) / (b[0] - a[0])
    return rows[-1][1]


def speed_limit(rows: list, target_ms: float) -> float:
    """The most chats in progress at which the p99 time between words is still at most `target_ms`, between two rows of the table."""
    for a, b in zip(rows, rows[1:]):
        if b[3] > target_ms:
            return a[0] + (b[0] - a[0]) * (target_ms - a[3]) / (b[3] - a[3])
    raise ValueError("the table never reaches the target")


def stay(rows: list, chats: float) -> float:
    """Seconds one chat holds its batch slot and its cache with `chats` in progress on a GPU: Little's law read off the table, chats divided by replies a second.
    A chat is in a slot from the step it joins to its last word, so the wait for the first word after joining is in it."""
    return chats / x_at(rows, chats)


def memory_limit(mix: str, share: int = SHARE) -> int:
    """Chats of the mix's average length (prompt plus reply) that lesson 28's cache pool holds."""
    m, chats = MIXES[mix], chats_of(mix)
    average = round(statistics.fmean(chats["prompts"]) + statistics.fmean(chats["replies"]))
    return L28.room(share, m["longest"])["blocks"] // L28.blocks_for(average)


# --- step 3: how many GPUs ------------------------------------------------------------------------------------------------------------------

def gpus(chats: float, limit: float) -> int:
    """GPUs so that none carries more than HEADROOM of `limit` chats. The 1e-9 keeps a division that is whole in exact arithmetic from rounding up."""
    return max(1, math.ceil(chats / (HEADROOM * limit) - 1e-9))


def plan(users: int, active: int, replies: int, mix: str, target_ms: float) -> dict:
    """The count for a crowd. `full` is the design point, the chats one GPU holds when it is as full as its tighter limit allows; `stay` is W there; `chats` is
    rate times W, the chats in progress by Little's law. Each limit divides the same `chats`."""
    rows, rate = step_table(mix), peak_rate(users, active, replies)
    memory, speed = memory_limit(mix), speed_limit(rows, target_ms)
    full = min(speed, memory)
    seconds = stay(rows, full)
    chats = rate * seconds
    n_speed, n_memory = gpus(chats, speed), gpus(chats, memory)
    return {"rate": rate, "memory": memory, "speed": speed, "binding": "speed" if speed < memory else "memory", "full": full, "stay": seconds, "chats": chats,
            "n_speed": n_speed, "n_memory": n_memory, "n": max(n_speed, n_memory),
            "without": max(1, math.ceil(chats / full - 1e-9))}                       # no headroom: every GPU full to its tighter limit


def commas(n: int) -> str:
    return f"{n:,}"


def one_decimal(x: float) -> str:
    return f"{math.floor(x * 10 + 0.5) / 10:.1f}"


def rate_text(rate: float) -> str:
    """Replies a second to one decimal, or to two below 1 so that a small crowd does not read 0.0."""
    return one_decimal(rate) if rate >= 1 else f"{math.floor(rate * 100 + 0.5) / 100:.2f}"


def gpu_word(n: int) -> str:
    return f"{n} GPU" if n == 1 else f"{commas(n)} GPUs"


def chats_words(chats: float) -> str:
    return "less than 1 chat" if chats < 0.5 else f"about {commas(math.floor(chats + 0.5))} chats"


def describe(users: int, active: int, replies: int, mix: str, target_ms: float) -> dict:
    """Everything the widget shows for one setting of its controls: the two bars and four lines of words."""
    p = plan(users, active, replies, mix, target_ms)
    people = users * active // 100
    m = MIXES[mix]
    crowded = " At the plan's speed each user would need a new reply before the last one ends: lower the replies per user." if p["chats"] > people else ""
    return {"bars": [p["n_memory"], p["n_speed"]],
            "line": f"{commas(users)} users: {chats_words(p['chats'])} at once in the busiest hour. Speed needs {gpu_word(p['n_speed'])}, memory {gpu_word(p['n_memory'])}. "
                    f"Plan {gpu_word(p['n'] + 1)}, with a spare.",
            "rate": f"The busiest hour brings {rate_text(p['rate'])} new replies a second from {commas(people)} active users.{crowded}",
            "note": f"Prompts of {commas(m['prompts'][0])} to {commas(m['prompts'][1])} tokens, replies of {commas(m['replies'][0])} to {commas(m['replies'][1])}.",
            "limits": f"One GPU holds {commas(p['memory'])} chats in memory, and keeps the p99 under {target_ms} ms up to {commas(math.floor(p['speed']))}: {p['binding']} sets the limit."}


# --- the load test --------------------------------------------------------------------------------------------------------------------------

def arrivals(rate: float, seconds: float, mix: str) -> list:
    """Random arrivals at `rate` a second for `seconds`: lesson 27's generator, the gaps from a fixed seed, and a fixed draw of which chat arrives."""
    m, count = MIXES[mix], int(rate * seconds)
    rng = random.Random(LOAD_SEED)
    times = L27.intended_times([round(rng.expovariate(1.0), 4) for _ in range(count)], rate)
    chats = draw(CHAT_SEED, count, m["prompts"], m["replies"])
    return [Request(t, p, r) for t, p, r in zip(times, chats["prompts"], chats["replies"])]


def load_test(count: int, requests: list, mix: str, cap: int, share: int = SHARE) -> dict:
    """Spread the requests over `count` GPUs, request k to GPU k mod count, each a lesson-28 server with `cap` chats at most and its real pool.
    Reports over all requests the p99 time between words (ms), the p99 wait before the first word (s), the mean time per reply, and the chats in progress on
    average: the sum of every reply's time (arrival to last word) over the length of the run, which is Little's law written as a sum."""
    m = MIXES[mix]
    blocks = L28.room(share, m["longest"])["blocks"]
    gaps, waits, spans = [], [], []
    for g in range(count):
        mine = requests[g::count]
        run = L28.serve(mine, cap, blocks, m["longest"])
        gaps += [(done - first) / (r.reply - 1) for done, first, r in zip(run["finished"], run["first_token"], mine)]
        waits += [first - r.arrival for first, r in zip(run["first_token"], mine)]
        spans += [(r.arrival, done) for r, done in zip(mine, run["finished"])]
    length = max(done for _, done in spans)
    times = [done - arrival for arrival, done in spans]
    gap, wait = L27.percentile(gaps, 99) * 1000, L27.percentile(waits, 99)
    return {"gpus": count, "gap": gap, "wait": wait, "passes": gap <= TARGET_MS and wait <= WAIT_TARGET, "mean_time": statistics.fmean(times), "rate": len(spans) / length,
            "in_progress": sum(times) / length, "spans": spans}


_LOADS: dict = {}


def opening_load_test() -> dict:
    """The lesson's load test, run once per process: the opening crowd's random arrivals through the GPUs the arithmetic alone gives, the plan, and one more."""
    if not _LOADS:
        p = plan(USERS, ACTIVE, REPLIES, OPENING_MIX, TARGET_MS)
        cap, requests = cap_for(OPENING_MIX), arrivals(p["rate"], LOAD_SECONDS, OPENING_MIX)
        _LOADS.update({"plan": p, "cap": cap, "requests": requests, "busiest": busiest_minute(requests),
                       "runs": {count: load_test(count, requests, OPENING_MIX, cap) for count in (p["without"], p["n"], p["n"] + 1)}})
    return _LOADS


def in_progress_curve(spans: list, seconds: int) -> list[int]:
    """Chats in progress at the start of each second, counted from the arrival and last-word times by sweeping through them in order."""
    events = sorted([(a, 1) for a, _ in spans] + [(d, -1) for _, d in spans])
    out, now, k = [], 0, 0
    for t in range(seconds):
        while k < len(events) and events[k][0] <= t:
            now += events[k][1]
            k += 1
        out.append(now)
    return out


def busiest_minute(requests: list) -> dict:
    """Arrivals in each whole minute: the busiest against the average."""
    minutes = int(max(r.arrival for r in requests) // 60)
    counts = [0] * minutes
    for r in requests:
        if r.arrival // 60 < minutes:
            counts[int(r.arrival // 60)] += 1
    return {"busiest": max(counts), "average": statistics.fmean(counts), "ratio": max(counts) / statistics.fmean(counts)}


def cap_for(mix: str, target_ms: float = TARGET_MS) -> int:
    """The cap on chats at once per GPU for a load test: the smaller limit, rounded down."""
    return math.floor(min(speed_limit(step_table(mix), target_ms), memory_limit(mix)))


def gpus_to_hold(model: str, replicas: int) -> int:
    """Lesson 26: a model too big for one GPU takes `least_n` of them for every copy."""
    return replicas * L26.least_n(model)