patdev/k3-a40-bootstrap

1"""Enchaine les variantes sur le pod sans jamais le recreer.2 3Le pod recharge le bootstrap a chaud des qu'il change sur le Hub. On modifie4donc les DEFAUTS du bloc EXPERIENCE (VL_KV / VL_SPEC), on republie, on attend5que vLLM redemarre avec les nouveaux flags (verifie dans son journal, pas6suppose), puis on mesure avec pod_bench.py. Les poids restent sur le disque :7un redemarrage coute ~2-3 min, pas un retelechargement.8 9    python drive_pod.py --variante kv=turboquant_k3v4_nc,spec=off10"""11from __future__ import annotations12 13import argparse14import json15import os16import re17import subprocess18import sys19import time20import urllib.request21 22ICI = os.path.dirname(os.path.abspath(__file__))23BOOT = r"g:\Environements\HuggingFace\vllm_bootstrap.sh"24BENCH = r"g:\Environements\HuggingFace\pod_bench.py"25POD = open(os.path.join(ICI, "podid")).read().strip()26BASE = f"https://{POD}-8080.proxy.runpod.net"27SSH_HOST, SSH_PORT = open(os.path.join(ICI, "podssh")).read().split()  # "ip port", ecrit par le moniteur28KEY = os.path.join(ICI, "ssh", "pod_key")29 30 31def ssh(cmd: str, timeout: int = 60) -> str:32    r = subprocess.run(["ssh", "-i", KEY, "-o", "StrictHostKeyChecking=no",33                        "-o", "UserKnownHostsFile=/dev/null", "-o", "ConnectTimeout=20",34                        "-o", "LogLevel=ERROR", "-p", SSH_PORT, f"root@{SSH_HOST}", cmd],35                       capture_output=True, text=True, timeout=timeout)36    return r.stdout37 38 39def publier(kv: str, spec: str, n_dspark: int) -> None:40    s = open(BOOT, encoding="utf-8", newline="").read()41    s = re.sub(r'^: "\$\{VL_KV:=[^}]*\}"', f': "${{VL_KV:={kv}}}"', s, flags=re.M)42    s = re.sub(r'^: "\$\{VL_SPEC:=[^}]*\}"', f': "${{VL_SPEC:={spec}}}"', s, flags=re.M)43    s = re.sub(r'^: "\$\{VL_DSPARK_N:=[^}]*\}"', f': "${{VL_DSPARK_N:={n_dspark}}}"', s, flags=re.M)44    open(BOOT, "w", encoding="utf-8", newline="").write(s)45    # (pas de `bash -n` ici : depuis Python, `bash` resout vers le bash WSL,46    #  qui ne lit pas les chemins Windows ; la syntaxe est verifiee a la main)47    subprocess.run(["hf", "upload", "patdev/k3-a40-bootstrap", BOOT, "vllm_bootstrap.sh",48                    "--commit-message", f"experience kv={kv or 'bf16'} spec={spec}"],49                   check=True, capture_output=True)50    print(f"  publie : VL_KV={kv or '(bf16)'} VL_SPEC={spec}", flush=True)51 52 53def attendre(kv: str, spec: str, budget: int = 900) -> bool:54    """Vrai quand vLLM a redemarre AVEC les flags voulus et repond."""55    t0 = time.time()56    motif_kv = f"'kv_cache_dtype': '{kv}'" if kv else None57    motif_spec = {"dspark": "dspark", "mtp": "'method': 'mtp'", "on": "ngram"}.get(spec)58    while time.time() - t0 < budget:59        log = ssh("grep -h 'non-default args' /tmp/vllm.log 2>/dev/null | tail -1 | cut -c1-6000; "60                  "echo ---; grep -c 'Application startup complete' /tmp/vllm.log 2>/dev/null")61        args, _, pret = log.partition("---")62        ok_kv = (motif_kv in args) if motif_kv else ("kv_cache_dtype" not in args)63        ok_spec = (motif_spec in args) if motif_spec else ("speculative_config" not in args)64        if ok_kv and ok_spec and pret.strip() not in ("", "0"):65            try:66                urllib.request.urlopen(BASE + "/v1/models", timeout=20)67                print(f"  pret en {time.time() - t0:.0f}s", flush=True)68                return True69            except Exception:                                    # noqa: BLE00170                pass71        time.sleep(20)72    print("  DELAI depasse ; journal :", ssh("tail -n 5 /tmp/vllm.log | cut -c1-200"))73    return False74 75 76def mesurer(etiquette: str, seqs: str, max_tokens: int) -> dict:77    r = subprocess.run([sys.executable, BENCH, "--base", BASE, "--seqs", seqs,78                        "--max-tokens", str(max_tokens)], capture_output=True, text=True, timeout=1800)79    print(r.stdout, flush=True)80    lignes = [l for l in r.stdout.splitlines() if l.startswith("[")]81    res = json.loads(lignes[-1]) if lignes else []82    kv_tok = ssh("grep -h 'GPU KV cache size' /tmp/vllm.log | tail -1 | grep -oE '[0-9,]+ tokens'")83    out = {"variante": etiquette, "kv_jetons": kv_tok.strip(), "debits": res}84    chemin = os.path.join(ICI, "mesures_pod.jsonl")85    with open(chemin, "a", encoding="utf-8") as f:86        f.write(json.dumps(out) + "\n")87    return out88 89 90def main() -> None:91    ap = argparse.ArgumentParser()92    ap.add_argument("--variante", required=True, help="kv=<dtype|''>,spec=<off|dspark|mtp|on>[,n=8]")93    ap.add_argument("--seqs", default="1,4,8,16,32")94    ap.add_argument("--max-tokens", type=int, default=256)95    ap.add_argument("--sans-publier", action="store_true", help="mesurer l'etat courant")96    a = ap.parse_args()97    kv, spec, n = "", "off", 898    for part in a.variante.split(","):99        k, _, v = part.partition("=")100        if k == "kv": kv = v101        elif k == "spec": spec = v102        elif k == "n": n = int(v)103    if not a.sans_publier:104        publier(kv, spec, n)105        if not attendre(kv, spec):106            sys.exit(2)107    print(json.dumps(mesurer(a.variante, a.seqs, a.max_tokens), ensure_ascii=False))108 109 110if __name__ == "__main__":111    main()112