recipe/pipeline/nex_harness.py

#!/usr/bin/env python3
"""Measurement harness for Nex-N2.5-mini on MAX-1 (stdlib only) - derived from the Agnes harness.
Nex's template has no thinking on/off flag: thinking is `reasoning_effort` = none (off) | high (on) | anything else
(adaptive). Tool calls use the upstream sampling (temp 0.7, top_p 0.95, top_k 40).
House protocol (Qwen3.8-27B card): ctx 65536, batch 1, greedy (temp 0, top_k 1), ignore_eos -> exactly
256 tokens, unique nonce + cache_prompt:false (cached tokens asserted 0), median of reps after 1 warm-up.
Subcommands:  bench | cachegate | tools | vision   (see argparse)"""
import argparse, base64, json, os, signal, statistics, subprocess, sys, time, urllib.request, uuid

W = "/mnt/models/nex-n2.5-mini"
CAL = "/mnt/models/agnes-3.0-flash/calib"      # shared calibration / prose corpora
THINK_OFF = {"reasoning_effort": "none"}
THINK_ON = {"reasoning_effort": "high"}
BIN = os.environ.get("AGNES_BIN", "/opt/llama-rocm/rocmfpx-724/build-hipvk/bin")
ENV = dict(os.environ, LD_LIBRARY_PATH=f"{BIN}:/opt/rocm-7.2.4/lib", HSA_OVERRIDE_GFX_VERSION="11.5.1",
           GGML_HIP_ENABLE_UNIFIED_MEMORY="1")

def post(port, path, body, timeout=1800):
    req = urllib.request.Request(f"http://127.0.0.1:{port}{path}", data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    return json.load(urllib.request.urlopen(req, timeout=timeout))

class Server:
    def __init__(self, a, port, mmproj=None, fa="on"):
        self.port = port
        cmd = [f"{BIN}/llama-server", "-m", a.model, "-dev", a.dev, "-ngl", "999", "-fa", fa, "-dio",
               "--jinja", "-fit", "off", "--parallel", "1", "-c", str(a.ctx), "-b", "2048", "-ub", "1024",
               "--host", "127.0.0.1", "--port", str(port), "--no-webui"]
        if a.draft or a.mtp_infile:
            cmd += ["--spec-type", "draft-mtp"]
            if a.draft:
                cmd += ["--model-draft", a.draft, "--spec-draft-ngl", "99", "--spec-draft-device", a.dev]
            cmd += ["--spec-draft-n-max", str(a.nmax), "--spec-draft-n-min", "0", "--spec-draft-p-min", str(a.pmin)]
            if a.strict:
                cmd += ["--spec-mtp-strict-qwen"]
        if mmproj:
            cmd += ["--mmproj", mmproj]
        self.cmd = cmd
        self.logf = open(a.serverlog, "w")
        self.t0 = time.time()
        self.p = subprocess.Popen(cmd, env=ENV, stdout=self.logf, stderr=subprocess.STDOUT, start_new_session=True)
        while True:
            if self.p.poll() is not None:
                raise SystemExit(f"SERVER DIED rc={self.p.returncode} see {a.serverlog}")
            try:
                if json.load(urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=3)).get("status") == "ok":
                    break
            except Exception:
                pass
            if time.time() - self.t0 > 900:
                self.stop(); raise SystemExit("SERVER LOAD TIMEOUT")
            time.sleep(2)
        self.load_s = time.time() - self.t0
    def stop(self):
        try: os.killpg(self.p.pid, signal.SIGTERM); self.p.wait(60)
        except Exception:
            try: os.killpg(self.p.pid, signal.SIGKILL)
            except Exception: pass
        self.logf.close()

def prompt_8k(kind):
    if kind == "code":
        src = open("/opt/llama-rocm/rocmfpx-724/convert_hf_to_gguf.py").read()[:30000]
        return src, "Above is part of a model converter. Write a new, complete Python function that validates a GGUF tensor-name map against a list of HF tensor names and reports unmapped names. Code only."
    txt = open(f"{CAL}/wikitext-2-raw/wiki.train.raw").read()[:34000]
    return txt, "Above are encyclopedia excerpts. Write a long, detailed new encyclopedia article in the same style about the history of lighthouses."

def one_request(port, ctx_text, instr, n=256, cache=False, nonce=True):
    tag = f"[req {uuid.uuid4()}]\n" if nonce else ""
    body = {"messages": [{"role": "user", "content": tag + ctx_text + "\n\n" + instr}],
            "max_tokens": n, "temperature": 0, "top_k": 1, "ignore_eos": True, "cache_prompt": cache,
            "chat_template_kwargs": THINK_OFF}
    t = time.time(); r = post(port, "/v1/chat/completions", body); el = time.time() - t
    tm = r.get("timings", {})
    return {"pred_n": tm.get("predicted_n"), "tg": tm.get("predicted_per_second"),
            "pp": tm.get("prompt_per_second"), "prompt_n": tm.get("prompt_n"), "cache_n": tm.get("cache_n"),
            "draft_n": tm.get("draft_n"), "draft_acc": tm.get("draft_n_accepted"), "wall": el,
            "usage": r.get("usage", {})}

def cmd_bench(a):
    s = Server(a, a.port)
    out = {"label": a.label, "model": os.path.basename(a.model),
           "draft": os.path.basename(a.draft) if a.draft else ("in-file" if a.mtp_infile else None),
           "nmax": a.nmax if (a.draft or a.mtp_infile) else None, "strict": bool(a.strict), "bin": BIN, "dev": a.dev, "ctx": a.ctx, "workload": a.workload,
           "load_s": round(s.load_s, 1), "cmd": " ".join(s.cmd), "runs": []}
    try:
        ctx_text, instr = prompt_8k(a.workload)
        one_request(a.port, ctx_text, instr)                      # warm-up, discarded
        for _ in range(a.reps):
            r = one_request(a.port, ctx_text, instr)
            cached = (r["usage"].get("prompt_tokens_details") or {}).get("cached_tokens", 0) or 0
            if r["pred_n"] != 256: raise SystemExit(f"GATE FAIL: predicted_n={r['pred_n']} != 256")
            if (r["cache_n"] or 0) != 0 or cached != 0: raise SystemExit(f"GATE FAIL: cache hit cache_n={r['cache_n']} cached={cached}")
            out["runs"].append(r)
    finally:
        s.stop()
    tg = [r["tg"] for r in out["runs"]]; pp = [r["pp"] for r in out["runs"]]
    dn = sum(r["draft_n"] or 0 for r in out["runs"]); da = sum(r["draft_acc"] or 0 for r in out["runs"])
    out.update(tg_median=round(statistics.median(tg), 2), tg_min=round(min(tg), 2), tg_max=round(max(tg), 2),
               pp_median=round(statistics.median(pp), 1), prompt_n=out["runs"][0]["prompt_n"],
               accept=(round(da / dn, 3) if dn else None))
    print(json.dumps({k: v for k, v in out.items() if k != "runs"}))
    with open(a.jsonl, "a") as f: f.write(json.dumps(out) + "\n")

def cmd_identity(a):
    """Fixed prompt, no nonce, no cache, greedy: return the exact generated text for cross-config diffing."""
    s = Server(a, a.port)
    texts = []
    try:
        ctx_text, instr = prompt_8k(a.workload)
        for _ in range(a.reps):
            body = {"messages": [{"role": "user", "content": ctx_text + "\n\n" + instr}], "max_tokens": 256,
                    "temperature": 0, "top_k": 1, "ignore_eos": True, "cache_prompt": False,
                    "chat_template_kwargs": THINK_OFF}
            r = post(a.port, "/v1/chat/completions", body)
            texts.append(r["choices"][0]["message"].get("content") or "")
    finally:
        s.stop()
    import hashlib
    res = {"label": a.label, "reps": a.reps, "sha256": [hashlib.sha256(t.encode()).hexdigest()[:16] for t in texts],
           "self_consistent": len(set(texts)) == 1, "text": texts[0]}
    print(json.dumps({k: v for k, v in res.items() if k != "text"}))
    with open(a.jsonl, "a") as f: f.write(json.dumps(res) + "\n")

def cmd_cachegate2(a):
    """Partial-prefix reuse WITH output identity: warm on A, run B warm (partial reuse), run B cold, diff outputs."""
    import hashlib
    s = Server(a, a.port)
    rows = []
    try:
        ctx_text, _ = prompt_8k("code")
        variants = ["Write a function that lists every tensor name in this file.",
                    "Write a function that counts the model classes registered in this file.",
                    "Write a function that finds the longest method in this file."]
        for i, vb in enumerate(variants[:a.reps]):
            def req(instr, n, cache):
                body = {"messages": [{"role": "user", "content": ctx_text + "\n\n" + instr}], "max_tokens": n,
                        "temperature": 0, "top_k": 1, "ignore_eos": True, "cache_prompt": cache,
                        "chat_template_kwargs": THINK_OFF}
                r = post(a.port, "/v1/chat/completions", body)
                tm = r.get("timings", {})
                return (r["choices"][0]["message"].get("content") or ""), tm.get("prompt_n"), tm.get("cache_n")
            req(f"Summarise this file. (warm-up {i})", 16, True)
            tw, pw, cw = req(vb, 192, True)
            tc, pc, cc = req(vb, 192, False)
            rows.append({"variant": i, "warm_prompt_n": pw, "warm_cache_n": cw, "cold_prompt_n": pc, "cold_cache_n": cc,
                         "identical": tw == tc, "warm_sha": hashlib.sha256(tw.encode()).hexdigest()[:12],
                         "cold_sha": hashlib.sha256(tc.encode()).hexdigest()[:12]})
            print("  ", json.dumps(rows[-1]), flush=True)
    finally:
        s.stop()
    reused = all((r["warm_cache_n"] or 0) > 0 for r in rows)
    ident = all(r["identical"] for r in rows)
    res = {"label": a.label, "rows": rows, "all_reused": reused, "all_identical": ident,
           "result": "PASS" if (reused and ident) else "FAIL"}
    print(json.dumps({k: v for k, v in res.items() if k != "rows"}))
    with open(a.jsonl, "a") as f: f.write(json.dumps(res) + "\n")
    sys.exit(0 if res["result"] == "PASS" else 1)

def cmd_cachegate(a):
    """Repeat-prompt test WITH the draft head loaded: turn 2 must reuse turn 1's prefix."""
    s = Server(a, a.port)
    try:
        ctx_text, instr = prompt_8k("code")
        r1 = one_request(a.port, ctx_text, instr, n=32, cache=True, nonce=False)
        r2 = one_request(a.port, ctx_text, instr + " Also add type hints.", n=32, cache=True, nonce=False)
    finally:
        s.stop()
    res = {"label": a.label, "turn1_prompt_n": r1["prompt_n"], "turn2_prompt_n": r2["prompt_n"],
           "turn2_cache_n": r2["cache_n"]}
    # hybrid recurrent models resume only from context checkpoints (~1024-token spacing), so the test is
    # "any prefix reuse" -- report the fraction rather than demand near-total reuse.
    ok = isinstance(r2["cache_n"], int) and r2["cache_n"] > 0
    res["reuse_fraction"] = round(r2["cache_n"] / r1["prompt_n"], 3) if ok and r1["prompt_n"] else 0.0
    res["result"] = "PASS" if ok else "FAIL"
    print(json.dumps(res))
    with open(a.jsonl, "a") as f: f.write(json.dumps(res) + "\n")
    sys.exit(0 if ok else 1)

TOOLS = [
 {"type": "function", "function": {"name": "get_weather", "description": "Current weather for a city",
  "parameters": {"type": "object", "properties": {"city": {"type": "string"},
   "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}}, "required": ["city", "unit"]}}},
 {"type": "function", "function": {"name": "create_event", "description": "Create a calendar event",
  "parameters": {"type": "object", "properties": {"title": {"type": "string"},
   "when": {"type": "object", "properties": {"date": {"type": "string"}, "time": {"type": "string"}},
            "required": ["date", "time"]},
   "attendees": {"type": "array", "items": {"type": "string"}}}, "required": ["title", "when", "attendees"]}}},
]
LEAK = ("<tool_call>", "<function=", "<parameter=", "<think>", "</think>")

def chat(port, msgs, think, stream=False, tools=TOOLS):
    body = {"messages": msgs, "tools": tools, "tool_choice": "auto", "temperature": 0.7, "top_p": 0.95,
            "top_k": 40, "max_tokens": 8192,
            "chat_template_kwargs": (THINK_ON if think else THINK_OFF)}
    if not stream:
        return post(port, "/v1/chat/completions", body)["choices"][0]["message"]
    body["stream"] = True
    req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/chat/completions", data=json.dumps(body).encode(),
                                 headers={"Content-Type": "application/json"})
    calls, content = {}, ""
    for line in urllib.request.urlopen(req, timeout=1800):
        line = line.decode().strip()
        if not line.startswith("data:") or line.endswith("[DONE]"): continue
        d = json.loads(line[5:])["choices"][0]["delta"]
        content += d.get("content") or ""
        for tc in d.get("tool_calls") or []:
            c = calls.setdefault(tc["index"], {"name": "", "arguments": ""})
            c["name"] += (tc.get("function") or {}).get("name") or ""
            c["arguments"] += (tc.get("function") or {}).get("arguments") or ""
    return {"content": content, "tool_calls": [{"function": v} for _, v in sorted(calls.items())]}

def args_of(m, i=0):
    return json.loads(m["tool_calls"][i]["function"]["arguments"])

def cmd_tools(a):
    s = Server(a, a.port)
    results = {}
    try:
        for think in (True, False):
            def check(name, fn):
                try: ok, why = fn()
                except Exception as e: ok, why = False, f"exception {e!r}"[:160]
                results[f"{name}|think={think}"] = (ok, why)
                print(f"  {'PASS' if ok else 'FAIL'} think={think!s:5} {name}: {why}", flush=True)
            def clean(m): return not any(x in (m.get("content") or "") for x in LEAK)
            def t1():
                m = chat(a.port, [{"role": "user", "content": "What's the weather in Paris in celsius?"}], think)
                ag = args_of(m); return (m["tool_calls"][0]["function"]["name"] == "get_weather" and ag.get("city", "").lower().startswith("paris")
                        and ag.get("unit") == "celsius" and clean(m)), f"args={ag}"
            def t2():
                m = chat(a.port, [{"role": "user", "content": "Book 'Design review' on 2026-10-02 at 14:00 with ana@x.io and bo@x.io."}], think)
                ag = args_of(m); return (isinstance(ag.get("when"), dict) and ag["when"].get("date") == "2026-10-02"
                        and sorted(ag.get("attendees", [])) == ["ana@x.io", "bo@x.io"] and clean(m)), f"args={ag}"
            def t3():
                m = chat(a.port, [{"role": "user", "content": "Weather in Denver, and give it to me in fahrenheit."}], think)
                return args_of(m).get("unit") == "fahrenheit" and clean(m), f"unit={args_of(m).get('unit')}"
            def t4():
                m = chat(a.port, [{"role": "user", "content": "What is 17 times 23? Answer directly."}], think)
                c = m.get("content") or ""
                return (not m.get("tool_calls")) and "391" in c and clean(m), f"content={c[:60]!r}"
            def t5():
                msgs = [{"role": "user", "content": "What's the weather in Tokyo in celsius?"}]
                m = chat(a.port, msgs, think)
                tc = m["tool_calls"][0]
                msgs += [{"role": "assistant", "content": m.get("content") or "", "tool_calls": [
                          {"id": "call_1", "type": "function", "function": tc["function"]}]},
                         {"role": "tool", "tool_call_id": "call_1", "content": json.dumps({"temp_c": 21, "sky": "clear"})}]
                m2 = chat(a.port, msgs, think)
                c = m2.get("content") or ""
                return ("21" in c and not m2.get("tool_calls") and clean(m2)), f"final={c[:70]!r}"
            def t6():
                m = chat(a.port, [{"role": "user", "content": "What's the weather in Rome in celsius?"}], think, stream=True)
                ag = args_of(m); return (m["tool_calls"][0]["function"]["name"] == "get_weather"
                        and ag.get("city", "").lower().startswith("rome") and clean(m)), f"stream args={ag}"
            def t7():
                m = chat(a.port, [{"role": "user", "content": "Get the weather in Oslo AND in Lima, both in celsius. Call the tool for each city."}], think)
                cities = sorted(args_of(m, i).get("city", "").lower() for i in range(len(m.get("tool_calls") or [])))
                return (len(cities) == 2 and cities[0].startswith("lima") and cities[1].startswith("oslo") and clean(m)), f"calls={cities}"
            for nm, fn in (("multi-arg", t1), ("nested-object", t2), ("enum", t3), ("correct-decline", t4),
                           ("multi-turn", t5), ("streaming", t6), ("parallel", t7)):
                check(nm, fn)
    finally:
        s.stop()
    n_ok = sum(v[0] for v in results.values())
    summary = {"label": a.label, "passed": n_ok, "total": len(results),
               "detail": {k: v[0] for k, v in results.items()}}
    print(json.dumps(summary))
    with open(a.jsonl, "a") as f: f.write(json.dumps(summary) + "\n")

def cmd_vision(a):
    """Image gate. A server that fails to load or dies on the image is a RESULT (FAIL row), not a harness crash."""
    fa = a.fa or "off"
    res = {"label": a.label, "fa": fa, "mtp": bool(a.draft or a.mtp_infile), "expected": a.expect,
           "answer": "", "hits": [], "error": None, "server_died": False, "server_log_errors": []}
    c = ""
    try:
        s = Server(a, a.port, mmproj=a.mmproj, fa=fa)
    except SystemExit as e:
        res.update(error=f"server did not start: {e}", server_died=True); s = None
    if s is not None:
        try:
            img = base64.b64encode(open(a.image, "rb").read()).decode()
            body = {"messages": [{"role": "user", "content": [
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img}"}},
                        {"type": "text", "text": a.question}]}],
                    "temperature": 0, "top_k": 1, "max_tokens": 300, "chat_template_kwargs": THINK_OFF}
            try:
                r = post(a.port, "/v1/chat/completions", body, timeout=900)
                c = r["choices"][0]["message"].get("content") or ""
            except Exception as e:
                res["error"] = f"{type(e).__name__}: {e}"[:300]
            time.sleep(1)
            res["server_died"] = s.p.poll() is not None
        finally:
            s.stop()
    try:
        res["server_log_errors"] = [l.strip()[-200:] for l in open(a.serverlog, errors="replace")
                                    if any(k in l for k in ("GGML_ABORT", "abort", "failed to process", " E "))][-5:]
    except OSError:
        pass
    res["answer"] = c[:300]
    res["hits"] = [w for w in a.expect.split(",") if w.lower() in c.lower()]
    ok = res["error"] is None and not res["server_died"] and len(res["hits"]) == len(a.expect.split(","))
    res["result"] = "PASS" if ok else "FAIL"
    print(json.dumps(res))
    with open(a.jsonl, "a") as f: f.write(json.dumps(res) + "\n")
    sys.exit(0 if ok else 1)

if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["bench", "cachegate", "cachegate2", "tools", "vision", "identity"])
    ap.add_argument("--model", required=True); ap.add_argument("--draft")
    ap.add_argument("--nmax", type=int, default=4); ap.add_argument("--pmin", type=float, default=0.0)
    ap.add_argument("--dev", default="ROCm0"); ap.add_argument("--ctx", type=int, default=65536)
    ap.add_argument("--port", type=int, default=18600); ap.add_argument("--reps", type=int, default=3)
    ap.add_argument("--workload", default="code", choices=["code", "prose"])
    ap.add_argument("--label", default=""); ap.add_argument("--jsonl", default=f"{W}/results/phase_b.jsonl")
    ap.add_argument("--serverlog", default=f"{W}/logs/server_last.log")
    ap.add_argument("--mmproj"); ap.add_argument("--image"); ap.add_argument("--question"); ap.add_argument("--expect")
    ap.add_argument("--mtp-infile", action="store_true"); ap.add_argument("--strict", action="store_true")
    ap.add_argument("--fa", choices=["on", "off", "auto"], help="vision only; default off")
    a = ap.parse_args()
    os.makedirs(os.path.dirname(a.jsonl), exist_ok=True)
    {"bench": cmd_bench, "cachegate": cmd_cachegate, "tools": cmd_tools, "vision": cmd_vision,
     "identity": cmd_identity, "cachegate2": cmd_cachegate2}[a.cmd](a)