patdev/k3-a40-bootstrap

1"""Vitesse de PREREMPLISSAGE : jetons d'entree par seconde, une requete, max_tokens=1.
2
3Le debit de decodage ne dit rien de ce que coute un gros prompt systeme ou un
4depot entier colle dans le contexte : c'est le prefill qui fixe le temps avant
5le premier jeton. On mesure a plusieurs tailles, a froid puis a chaud (meme
6prompt, cache de prefixe), en streaming pour passer le proxy Runpod (~125 s).
7
8    python prefill_bench.py --base https://<pod>-8080.proxy.runpod.net --model claude-qwen38 --tailles 8000,32000,128000
9"""
10from __future__ import annotations
11import argparse, json, time, urllib.request, uuid
12
13LIGNE = ("Le registre {} consigne les operations de maintenance courante du parc, "
14         "sans incident notable ni intervention requise.\n")
15
16def compte(base, model, texte):
17    r = urllib.request.urlopen(urllib.request.Request(base + "/v1/messages/count_tokens",
18        data=json.dumps({"model": model, "messages": [{"role": "user", "content": texte}]}).encode(),
19        headers={"Content-Type": "application/json", "anthropic-version": "2023-06-01", "User-Agent": "prefill-bench"}), timeout=300)
20    return json.load(r)["input_tokens"]
21
22def une(base, model, texte, budget):
23    body = {"model": model, "max_tokens": 1, "temperature": 0, "stream": True,
24            "thinking": {"type": "disabled"},
25            "messages": [{"role": "user", "content": texte + "\nReponds par un seul mot."}]}
26    req = urllib.request.Request(base + "/v1/messages", data=json.dumps(body).encode(),
27        headers={"Content-Type": "application/json", "anthropic-version": "2023-06-01",
28                 "Accept": "text/event-stream", "User-Agent": "prefill-bench"})
29    t0 = time.time(); entree = 0; premier = None
30    with urllib.request.urlopen(req, timeout=budget) as r:
31        for raw in r:
32            l = raw.decode("utf-8", "ignore").strip()
33            if not l.startswith("data:"): continue
34            try: ev = json.loads(l[5:])
35            except json.JSONDecodeError: continue
36            t = ev.get("type")
37            if t == "content_block_delta" and premier is None: premier = time.time() - t0
38            if t in ("message_start", "message_delta"):
39                u = ev.get("usage") or (ev.get("message") or {}).get("usage") or {}
40                entree = u.get("input_tokens") or entree
41            if t == "error": raise RuntimeError(str(ev)[:200])
42    return entree, (premier if premier is not None else time.time() - t0)
43
44def main():
45    ap = argparse.ArgumentParser()
46    ap.add_argument("--base", required=True); ap.add_argument("--model", default="claude-ornith")
47    ap.add_argument("--tailles", default="8000,32000,128000"); ap.add_argument("--budget", type=float, default=280)
48    a = ap.parse_args(); base = a.base.rstrip("/")
49    par_ligne = compte(base, a.model, LIGNE.format(0) * 500) / 500
50    for cible in (int(x) for x in a.tailles.split(",")):
51        sel = f"[{uuid.uuid4().hex}] "
52        n = max(10, int(cible / par_ligne)); texte = sel + "".join(LIGNE.format(i) for i in range(n))
53        res = []
54        for passe in ("froid", "chaud"):
55            try:
56                e, ttft = une(base, a.model, texte, a.budget)
57                res.append(f"{passe}: {e:,} jetons en {ttft:.1f}s = {e/ttft:,.0f} j/s")
58            except Exception as ex:
59                res.append(f"{passe}: ECHEC {type(ex).__name__} {str(ex)[:60]}")
60        print(f"  ~{cible:,} : " + " | ".join(res), flush=True)
61
62if __name__ == "__main__":
63    main()
64