1"""E6/E2: genuine two-node cache transfer over the LAN (laptop -> M3U96b worker) vs text transfer vs local recompute.2Usage: python scripts/e6_network.py --host 192.168.2.85 --port 9430 --model qwen3b --docs 6 --lengths 2000 8000 16000 --repeats 33Requires the worker: on the remote node `python -m wp30 worker --port 9430 --model <repo>` (same checkpoint revision).4Loopback mode (--host 127.0.0.1) is labelled 'loopback' and is NOT a network measurement."""5import argparse, json, time, statistics, subprocess, sys, socket6import yaml, numpy as np, mlx.core as mx7from wp30.util import ROOT, read_json, utc_now, append_jsonl, sha256_text8from wp30.data.prepare import load_doc, load_key9from wp30.backends.mlx_backend import MLXModel10from wp30.cache.sparse import quantize_cache, cache_tensor_bytes, to_sparse_all11from wp30.cache.serialize import build_manifest, serialize_cache, codec_compress12from wp30.transport import Client13from wp30.prompts import SYSTEM_RECEIVER, receiver_user_text14from wp30.scoring import score_synthetic1516ap = argparse.ArgumentParser()17ap.add_argument("--host", required=True); ap.add_argument("--port", type=int, default=9430); ap.add_argument("--model", default="qwen3b")18ap.add_argument("--docs", type=int, default=6); ap.add_argument("--lengths", type=int, nargs="+", default=[2000, 8000, 16000]); ap.add_argument("--repeats", type=int, default=3)19ap.add_argument("--label", default=None); ap.add_argument("--warmup", type=int, default=1); ap.add_argument("--split", default="dev"); ap.add_argument("--out", default="runs/raw/E6_network.jsonl")20a = ap.parse_args()21label = a.label or ("loopback" if a.host.startswith("127.") else "lan")22mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]23m = MLXModel(mcfg[a.model]["repo"])24cl = Client(a.host, a.port)25h, _, meta = cl.call({"cmd": "ping"})26assert h["config_hash"] == m.info.config_hash, "remote model differs"27print("remote ok", h["model_revision"], "rtt_ms", round(meta["roundtrip_ms"], 2))28# raw link bandwidth probe with an incompressible 64 MiB payload (frame echo is not implemented; use run_text no-op? -> use a dedicated ping with payload)29man = read_json(ROOT / "data/manifests/splits.json")30items = [it for it in man["items"] if it["family"] == "synthetic" and it["split"] == a.split and it.get("target_tokens") in a.lengths]31by_len = {L: [it for it in items if it["target_tokens"] == L][: a.docs] for L in a.lengths}32out = ROOT / a.out33# warm-up: compile kernels on both ends (not recorded)34if a.warmup:35 _w = m.chat_prompt_tokens([{"role": "user", "content": "warm up " * 200}])36 for _ in range(a.warmup):37 _c, _ = m.prefill(_w); m.generate_from_cache(_c, _w[-3:], max_tokens=4); cl.call({"cmd": "run_text", "token_ids": _w, "max_tokens": 4})38 _man = build_manifest(m.info, _w, to_sparse_all(_c), "warm", "B1", {}); cl.call({"cmd": "run_cache", "codec": "none", "prefix_hash": _man["prefix_token_hash"], "query_ids": _w[-3:], "max_tokens": 4}, serialize_cache(to_sparse_all(_c), _man))39for L, its in by_len.items():40 for it in its:41 doc = load_doc(it["document_id"], "synthetic"); key = load_key(it["document_id"], "synthetic")42 q = doc["questions"][0]; qk = key["questions"][0]43 full = m.chat_prompt_tokens([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": receiver_user_text("Document", doc["text"], q["question"], "synthetic")}])44 pre_txt = m.tokenizer.apply_chat_template([{"role": "system", "content": SYSTEM_RECEIVER}, {"role": "user", "content": "Document:\n" + doc["text"] + "\n\n"}], add_generation_prompt=False, tokenize=False)45 pre = m.encode(pre_txt); n = 046 while n < min(len(pre), len(full)) and pre[n] == full[n]: n += 147 pre, query = full[:n], full[n:]48 for rep in range(a.repeats):49 # (1) local sender prefill + serialize + transfer cache + remote load + remote query [warm remote cache scenario: cache moved]50 for variant in [("B1", 8 if False else None, "none"), ("B1", None, "zstd"), ("B2", 8, "none"), ("B2", 4, "none")]:51 meth, bits, codec = variant52 t0 = time.perf_counter(); c, tp = m.prefill(pre)53 c2 = to_sparse_all(c) if bits is None else quantize_cache(c, bits)54 mx.eval([x for cc in c2 for x in (cc.keys if isinstance(cc.keys, tuple) else (cc.keys,))]); mx.synchronize()55 t1 = time.perf_counter()56 manf = build_manifest(m.info, pre, c2, it["document_id"], meth, {"bits": bits})57 payload = serialize_cache(c2, manf); t2 = time.perf_counter()58 comp, cms = codec_compress(payload, codec) if codec != "none" else (payload, 0.0)59 hdr = {"cmd": "run_cache", "codec": codec, "prefix_hash": manf["prefix_token_hash"], "query_ids": query, "max_tokens": 32}60 rh, _, tm = cl.call(hdr, comp)61 rec = {"experiment_id": "E6_network", "label": label, "timestamp_utc": utc_now(), "document_id": it["document_id"], "target_tokens": L, "prefix_tokens": len(pre), "query_tokens": len(query), "repeat": rep,62 "method": meth, "bits": bits, "codec": codec, "status": rh.get("status"), "answer": rh.get("text"), "correct": score_synthetic(rh.get("text", ""), qk)["correct"] if rh.get("text") else None,63 "cache_tensor_bytes": cache_tensor_bytes(c2), "serialized_bytes": len(payload), "transport_bytes": tm["transport_bytes"], "sender_prefill_ms": tp["prefill_ms"], "compress_ms": (t1 - t0) * 1000 - tp["prefill_ms"],64 "serialize_ms": (t2 - t1) * 1000, "codec_ms": cms, "send_ms": tm["send_ms"], "roundtrip_ms": tm["roundtrip_ms"], "remote_recv_ms": rh.get("recv_ms"), "remote_decompress_ms": rh.get("decompress_ms"),65 "remote_load_ms": rh.get("load_ms"), "remote_ttft_ms": rh.get("ttft_ms"), "remote_decode_ms": rh.get("decode_ms"), "remote_server_ms": rh.get("server_ms"), "remote_peak_bytes": rh.get("peak_bytes"),66 "effective_gbps": (tm["transport_bytes"] * 8 / 1e9) / max(1e-9, (tm["roundtrip_ms"] - (rh.get("server_ms") or 0)) / 1000)}67 append_jsonl(out, rec)68 print(f"{label} L={L} {meth} bits={bits} codec={codec} rep={rep} bytes={len(comp):>10d} send={tm['send_ms']:7.1f}ms rt={tm['roundtrip_ms']:7.1f}ms remote_load={rh.get('load_ms',0):6.1f} ttft={rh.get('ttft_ms',0):6.1f} eff={rec['effective_gbps']:.2f}Gb/s corr={rec['correct']}")69 # (2) text transfer: send token ids, remote does full prefill [cold remote scenario]70 rh, _, tm = cl.call({"cmd": "run_text", "token_ids": full, "max_tokens": 32})71 rec = {"experiment_id": "E6_network", "label": label, "timestamp_utc": utc_now(), "document_id": it["document_id"], "target_tokens": L, "prefix_tokens": len(pre), "query_tokens": len(query), "repeat": rep,72 "method": "text_remote_prefill", "bits": None, "codec": "none", "status": "success", "answer": rh.get("text"), "correct": score_synthetic(rh.get("text", ""), qk)["correct"],73 "serialized_bytes": len(doc["text"].encode()), "transport_bytes": tm["transport_bytes"], "send_ms": tm["send_ms"], "roundtrip_ms": tm["roundtrip_ms"], "remote_prefill_ms": rh.get("prefill_ms"),74 "remote_ttft_ms": rh.get("ttft_ms"), "remote_decode_ms": rh.get("decode_ms"), "remote_server_ms": rh.get("server_ms"), "remote_peak_bytes": rh.get("peak_bytes")}75 append_jsonl(out, rec)76 print(f"{label} L={L} text_remote rep={rep} bytes={tm['transport_bytes']:>10d} rt={tm['roundtrip_ms']:7.1f}ms remote_prefill={rh.get('prefill_ms',0):7.1f} ttft={rh.get('ttft_ms',0):6.1f}")77 # (3) local processing (no transfer): run where the data is78 c, tp = m.prefill(pre); g = m.generate_from_cache(c, query, max_tokens=32)79 append_jsonl(out, {"experiment_id": "E6_network", "label": "local", "timestamp_utc": utc_now(), "document_id": it["document_id"], "target_tokens": L, "prefix_tokens": len(pre), "query_tokens": len(query), "repeat": rep,80 "method": "local_prefill", "status": "success", "answer": g["text"], "correct": score_synthetic(g["text"], qk)["correct"], "local_prefill_ms": tp["prefill_ms"], "local_ttft_ms": g["ttft_ms"], "local_decode_ms": g["decode_ms"]})81cl.close()82print("done ->", out)83