# Balayage capacite / cache / prefill sur le pod de production.
#
# Socle = la meilleure config connue : 0.93 + --enable-flashinfer-autotune.
# UNE variable change par configuration ; la reference est mesuree d'abord.
#
# TROIS NOUVEAUTES par rapport aux balayages precedents :
#
# 1. --kv-cache-memory-bytes. On est a 12,26 Gio de KV avec 0.93. vLLM annonce
#    19 243 520 512 (17,92 Gio) comme "fully utilize" -- et cette valeur-la a
#    tue le serveur. L'intervalle entre les deux n'a JAMAIS ete essaye.
#    C'est le seul levier qui attaque le goulot demontre : la capacite.
#
# 2. Le prefill FROID APRES EVICTION, pas seulement le froid initial. C'est le
#    seul cas que le dechargement KV peut ameliorer. Les prompts sont SALES en
#    TETE : sans ca le cache de prefixe fausse le "froid" d'un facteur 20.
#
# 3. Une sonde de FUITE CJK. Le 26/08 le modele a ecrit un ideogramme au
#    milieu d'un mot francais -- un glissement de jeton vers le chinois.
#    C'est la signature d'une marge de logit ecrasee. On la mesure a
#    temperature 1,0 (le regime de Claude Code) et a 0,7 (celui des bancs).
#    Un debit ne vaut rien si le texte deraille.
#
# Le script rend TOUJOURS un serveur : la meilleure config qui a demarre, ou
# la reference en repli. Ce pod sert les sessions de l'utilisateur.
import glob
import json
import os
import re
import subprocess
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

BASE = "http://127.0.0.1:8000"
MODELE = "RadixArk/Qwen3.8-Flash-Next-NVFP4"
LOG = "/travail/vllm.log"
JOURNAL = "/travail/banc.log"
DEBUT = time.time()
PLAFOND = 4200          # 70 min : au-dela on assemble et on rend le serveur.

SOCLE = {
    "--served-model-name": "flashnext",
    "--host": "127.0.0.1",
    "--port": "8000",
    "--max-model-len": "262144",
    "--max-num-seqs": "16",
    "--gpu-memory-utilization": "0.93",
    "--distributed-executor-backend": "mp",
    "--reasoning-parser": "qwen3",
    "--tool-call-parser": "qwen3_coder",
    "--limit-mm-per-prompt": '{"image":0,"video":0}',
}
FANIONS = ["--trust-remote-code", "--enable-prefix-caching",
           "--enable-auto-tool-choice", "--enable-flashinfer-autotune",
           "--enable-prompt-tokens-details"]

GIO = 1024 ** 3


def dire(msg):
    ligne = "[%5.1f min] %s" % ((time.time() - DEBUT) / 60, msg)
    print(ligne, flush=True)
    with open(JOURNAL, "a") as f:
        f.write(ligne + "\n")


# ---------------------------------------------------------------- indexer
def cherche_config():
    """Le seul levier qui reduise le TRAVAIL de prefill au lieu de
    l'ordonnancer : le budget de l'attention eparse QSA. On ne devine pas sa
    valeur, on la lit -- et on note ou elle se trouve dans l'arbre, parce que
    --hf-overrides doit reproduire exactement ce chemin."""
    motifs = [
        "/root/.cache/huggingface/hub/models--RadixArk--*Flash-Next*/snapshots/*/config.json",
        "/travail/**/models--RadixArk--*Flash-Next*/snapshots/*/config.json",
        os.path.expanduser(
            "~/.cache/huggingface/hub/models--RadixArk--*Flash-Next*/snapshots/*/config.json"),
    ]
    for m in motifs:
        for p in glob.glob(m, recursive=True):
            try:
                return json.load(open(p)), p
            except Exception:
                pass
    return None, None


def surcharge_indexer(facteur=0.5):
    cfg, chemin = cherche_config()
    if not cfg:
        return None, "config.json introuvable"
    if "indexer_budget" in cfg:
        v = cfg["indexer_budget"]
        return (json.dumps({"indexer_budget": int(v * facteur)}),
                "racine %s -> %d" % (v, int(v * facteur)))
    tc = cfg.get("text_config") or {}
    if "indexer_budget" in tc:
        v = tc["indexer_budget"]
        return (json.dumps({"text_config": {"indexer_budget": int(v * facteur)}}),
                "text_config %s -> %d (%s)" % (v, int(v * facteur), chemin))
    return None, "pas d'indexer_budget dans %s" % chemin


# ------------------------------------------------------------- lancement
def arreter():
    out = subprocess.run(["ps", "-eo", "pid,rss,comm", "--no-headers"],
                         capture_output=True, text=True).stdout
    for l in out.splitlines():
        p = l.split(None, 2)
        if len(p) != 3:
            continue
        pid, rss, comm = int(p[0]), int(p[1]), p[2].strip()
        # Filtrer sur le NOM, jamais sur la ligne de commande : ce script
        # contient lui-meme la commande de lancement.
        if re.match(r"^(vllm|VLLM|Ple|EngineCore)", comm) or \
           (rss > 10_000_000 and re.match(r"^python3?$", comm)):
            try:
                os.kill(pid, 9)
            except Exception:
                pass
    time.sleep(12)


FATAL = re.compile(r"Engine core initialization failed|EngineDeadError|"
                   r"Traceback \(most recent call last\)")


def raison(txt):
    m = re.findall(r"^.*(?:Error|Exception|not supported|NotImplemented).*$",
                   txt, re.M)
    return " || ".join(x.strip()[:200] for x in m[-3:]) or "cause non identifiee"


def lancer(nom, remplace, ajout, retire):
    if os.path.exists(LOG):
        os.replace(LOG, "/travail/vllm-%s-precedent.log" % nom)
    opts = dict(SOCLE)
    opts.update(remplace)
    cmd = ["vllm", "serve", MODELE]
    for k, v in opts.items():
        cmd += [k, str(v)]
    cmd += [f for f in FANIONS if f not in retire] + list(ajout)
    env = dict(os.environ, VLLM_PLE_CPU_OFFLOAD="1")
    with open(LOG, "w") as sortie:
        subprocess.Popen(cmd, stdout=sortie, stderr=subprocess.STDOUT,
                         stdin=subprocess.DEVNULL, start_new_session=True, env=env)
    for i in range(45):
        time.sleep(20)
        try:
            urllib.request.urlopen(BASE + "/v1/models", timeout=5)
            return True, ""
        except Exception:
            pass
        try:
            txt = open(LOG, errors="replace").read()
        except Exception:
            txt = ""
        if FATAL.search(txt):
            return False, raison(txt)
    return False, "jamais pret en 15 min"


def kv_annonce():
    try:
        txt = open(LOG, errors="replace").read()
        k = re.findall(r"GPU KV cache size: ([\d,]+) tokens", txt)
        c = re.findall(r"concurrency for [\d,]+ tokens per request: ([\d.]+)x", txt)
        return (int(k[-1].replace(",", "")) if k else None,
                float(c[-1]) if c else None)
    except Exception:
        return None, None


# ------------------------------------------------------------------ banc
def metrique(nom):
    try:
        t = urllib.request.urlopen(BASE + "/metrics", timeout=5).read().decode()
        m = re.search(r"^vllm:%s\S*\s+([0-9.eE+-]+)$" % nom, t, re.M)
        return float(m.group(1)) if m else None
    except Exception:
        return None


def attendre_vide(limite=90):
    """Trois mesures de debit solo ont deja ete polluees par le trafic de
    l'utilisateur. On refuse de mesurer plutot que de publier un faux."""
    t0 = time.time()
    while time.time() - t0 < limite:
        n = metrique("num_requests_running")
        if n is not None and n < 0.5:
            return True
        time.sleep(5)
    return False


def appel(prompt, n, ignore=True, temp=0.7):
    corps = json.dumps({"model": "flashnext", "prompt": prompt, "max_tokens": n,
                        "ignore_eos": ignore, "temperature": temp}).encode()
    t = time.time()
    r = json.loads(urllib.request.urlopen(urllib.request.Request(
        BASE + "/v1/completions", data=corps,
        headers={"content-type": "application/json"}), timeout=600).read().decode())
    u = r["usage"]
    return (u["completion_tokens"], u["prompt_tokens"], time.time() - t,
            r["choices"][0]["text"])


def prompt_sale(graine, mots=24000):
    """Le sel est en TETE : rien ne peut correspondre dans le cache."""
    return ("sel%08d " % graine) + " ".join("mot%d" % (i + graine)
                                            for i in range(mots))


CJK = re.compile(r"[぀-ヿ㐀-䶿一-鿿가-힯]")

Q_PROMPT = ("Redige en francais, en trois paragraphes techniques et precis, "
            "comment on compile un projet Zig avec build.zig, ce que fait "
            "l'etape de verification des types, et pourquoi il faut relancer "
            "la compilation apres avoir modifie un module partage.\n\n")


def div4(texte):
    mots = texte.split()
    if len(mots) < 40:
        return None
    q = [tuple(mots[i:i + 4]) for i in range(len(mots) - 3)]
    return round(len(set(q)) / len(q), 3)


def qualite(r):
    """La sonde du bug du 26/08 : des ideogrammes au milieu du francais.
    Trois tirages, dont deux a temperature 1,0 -- c'est ce que Claude Code
    envoie, et c'est la que le glissement a ete observe."""
    total, echantillons, pire = 0, [], ""
    for temp in (1.0, 1.0, 0.7):
        try:
            _, _, _, txt = appel(Q_PROMPT, 420, ignore=False, temp=temp)
        except Exception:
            continue
        n = len(CJK.findall(txt))
        total += n
        echantillons.append({"temp": temp, "cjk": n, "div4": div4(txt)})
        if n and not pire:
            i = CJK.search(txt).start()
            pire = txt[max(0, i - 45):i + 25].replace("\n", " ")
    r["cjk"] = total
    r["qualite"] = echantillons
    r["div4"] = echantillons[0]["div4"] if echantillons else None
    if pire:
        r["fuite_cjk"] = pire


def banc(nom, eviction):
    r = {"nom": nom}
    r["kv"], r["conc"] = kv_annonce()
    try:
        appel("Bonjour.", 64)                        # rodage des noyaux
        if not attendre_vide():
            r["erreur"] = "serveur non oisif : mesure refusee"
            return r
        n, _, d, _ = appel("Ecris un texte long et detaille sur la mer.", 300)
        r["solo"] = round(n / d, 1)
        r["ms_pas"] = round(d * 1000 / n, 2)

        t = time.time()
        with ThreadPoolExecutor(max_workers=8) as ex:
            res = list(ex.map(lambda i: appel(
                "Ecris un texte long sur la mer numero %d." % i, 200), range(8)))
        r["agg8"] = round(sum(x[0] for x in res) / (time.time() - t), 1)

        A = prompt_sale(1)
        _, pA, d1, _ = appel(A, 1)
        r["prefill_froid"] = round(pA / d1)
        r["jetons"] = pA
        _, _, d2, _ = appel(A, 1)
        r["prefill_chaud"] = round(pA / d2)

        if eviction:
            # Chasser A du cache GPU : le KV tient ~500 k jetons, on pousse
            # 5 x 133 k = 665 k de prompts distincts.
            for i in range(2, 7):
                appel(prompt_sale(i), 1)
            _, _, d3, _ = appel(A, 1)
            r["prefill_evince"] = round(pA / d3)
            # 1,0 = rien recupere, > 1 = le dechargement sert vraiment.
            r["gain_eviction"] = round(d1 / d3, 2)
        r["hit"] = metrique("gpu_prefix_cache_hit_rate")
        qualite(r)
    except Exception as e:
        r["erreur"] = repr(e)[:200]
    return r


# ------------------------------------------------------- configurations
ov_indexer, note_indexer = surcharge_indexer(0.5)

CONFIGS = [
    # --- LE GOULOT DEMONTRE : la capacite. 0.93 rend 12,26 Gio ; le "fully
    # utilize" a 17,92 Gio avait tue le serveur. On sonde l'intervalle.
    ("kvmem15", {"--kv-cache-memory-bytes": 15 * GIO}, [], [], False),
    ("kvmem17", {"--kv-cache-memory-bytes": 17 * GIO}, [], [], False),
    # --- le dechargement : spectaculaire ou nul, rien entre les deux
    ("kvoff8", {"--kv-offloading-size": "8"}, [], [], True),
    # --- 166 blocs Pickle+SHA256 par prompt de 133 k, dans le chemin chaud
    ("xxhash", {"--prefix-caching-hash-algo": "xxhash"}, [], [], False),
    # --- le seul levier qui reduise le TRAVAIL de prefill
    ("indexer", {"--hf-overrides": ov_indexer} if ov_indexer else None,
     [], [], False),
    # --- granularite des succes de prefixe (bloc force a 800 pour Mamba)
    ("matchunit", {"--prefix-match-unit": "64"}, [], [], False),
    # --- partage de KV entre couches : le nom promet exactement notre regime
    ("kvshare", {}, ["--kv-sharing-fast-prefill"], [], False),
]

resultats = []
dire("=== indexer_budget : %s" % note_indexer)
dire("=== REFERENCE (socle 0.93 + autotune) -- relance propre")
arreter()
ok, motif = lancer("base", {}, [], [])
if ok:
    ref = banc("base", eviction=True)
    resultats.append(ref)
    dire("  base : %s" % json.dumps(ref, ensure_ascii=False))
else:
    ref = {}
    dire("  base NE DEMARRE PAS -- %s" % motif)
    resultats.append({"nom": "base", "echec": motif})

for nom, remplace, ajout, retire, evic in CONFIGS:
    if remplace is None:
        dire("=== %s : IGNOREE (%s)" % (nom, note_indexer))
        resultats.append({"nom": nom,
                          "echec": "non applicable : " + str(note_indexer)})
        continue
    if time.time() - DEBUT > PLAFOND:
        dire("!!! plafond de temps -- on assemble et on rend le serveur")
        break
    dire("=== %s : %s %s" % (nom, remplace, ajout))
    arreter()
    ok, motif = lancer(nom, remplace, ajout, retire)
    if not ok:
        dire("  %s : NE DEMARRE PAS -- %s" % (nom, motif))
        resultats.append({"nom": nom, "echec": motif})
        continue
    r = banc(nom, eviction=evic)
    resultats.append(r)
    dire("  %s : %s" % (nom, json.dumps(r, ensure_ascii=False)))

with open("/travail/balayage_cache.json", "w") as f:
    json.dump(resultats, f, indent=1, ensure_ascii=False)


# ------------------------------------------------- assemblage et remise
def trouve(nom):
    for r in resultats:
        if r["nom"] == nom and "echec" not in r and "erreur" not in r:
            return r
    return None


final_opts, final_fanions, raisons = {}, [], []
agg_ref = (ref or {}).get("agg8") or 0
kv_ref = (ref or {}).get("kv") or 0

# La capacite d'abord : le plus gros KV qui a demarre ET qui ne perd pas plus
# de 5 % d'agrege. Un cache plus grand ne vaut rien s'il ralentit.
for nom, octets in (("kvmem17", 17 * GIO), ("kvmem15", 15 * GIO)):
    r = trouve(nom)
    if r and r.get("kv") and (r.get("agg8") or 0) >= 0.95 * agg_ref \
            and r["kv"] > kv_ref:
        final_opts["--kv-cache-memory-bytes"] = octets
        raisons.append("%s (kv %d vs %d, agg8 %s vs %s)"
                       % (nom, r["kv"], kv_ref, r["agg8"], agg_ref))
        break

# Puis ce qui rend du debit sans rien couter.
for nom, cle, val in (("xxhash", "--prefix-caching-hash-algo", "xxhash"),
                      ("matchunit", "--prefix-match-unit", "64")):
    r = trouve(nom)
    if r and r.get("agg8") and r["agg8"] > agg_ref * 1.01:
        final_opts[cle] = val
        raisons.append("%s (agg8 %s > %s)" % (nom, r["agg8"], agg_ref))

# Le dechargement : seulement s'il recupere vraiment apres eviction.
r = trouve("kvoff8")
if r and (r.get("gain_eviction") or 0) > 1.15:
    final_opts["--kv-offloading-size"] = "8"
    raisons.append("kvoff8 (gain apres eviction %s)" % r["gain_eviction"])

r = trouve("kvshare")
if r and r.get("prefill_froid") and (ref or {}).get("prefill_froid") \
        and r["prefill_froid"] > 1.05 * ref["prefill_froid"]:
    final_fanions.append("--kv-sharing-fast-prefill")
    raisons.append("kvshare (prefill froid %s > %s)"
                   % (r["prefill_froid"], ref["prefill_froid"]))

dire("=== CONFIG FINALE : %s" % (", ".join(raisons) or "reference telle quelle"))
arreter()
ok, motif = lancer("final", final_opts, final_fanions, [])
if not ok:
    dire("  la config finale ne demarre pas (%s) -- retour a la reference" % motif)
    arreter()
    ok, motif = lancer("repli", {}, [], [])
    dire("  repli : %s %s" % ("OK" if ok else "ECHEC", motif))
    final_opts, final_fanions = {}, []

kv, conc = kv_annonce()
dire("=== SERVEUR RENDU : kv=%s jetons, concurrence %sx, options=%s %s"
     % (kv, conc, final_opts, final_fanions))

print("\n\n########## TABLEAU ##########")
ENT = "%-10s %9s %5s %6s %7s %8s %8s %6s %4s %6s"
print(ENT % ("config", "kv", "conc", "solo", "agg8", "pref_fr", "pref_ev",
             "gain", "cjk", "div4"))
for r in resultats:
    if "echec" in r:
        print("%-10s  NE DEMARRE PAS : %s" % (r["nom"], r["echec"][:95]))
    elif "erreur" in r:
        print("%-10s  ERREUR : %s" % (r["nom"], r["erreur"][:95]))
    else:
        print(ENT % (r["nom"], r.get("kv"), r.get("conc"), r.get("solo"),
                     r.get("agg8"), r.get("prefill_froid"),
                     r.get("prefill_evince", "-"), r.get("gain_eviction", "-"),
                     r.get("cjk"), r.get("div4")))
print("\n  gain = prefill apres eviction / prefill froid. 1,0 = rien recupere.")
print("  cjk  = ideogrammes dans 3 generations FRANCAISES (2 a temperature 1,0).")
print("         Zero est le seul resultat acceptable.")
print("\n########## FUITES CJK OBSERVEES ##########")
vu = False
for r in resultats:
    if r.get("fuite_cjk"):
        vu = True
        print("  %-10s ...%s..." % (r["nom"], r["fuite_cjk"]))
if not vu:
    print("  aucune")