| |
| import argparse, json, os, re, resource, subprocess, sys, threading, time |
|
|
| PAGE = resource.getpagesize() |
|
|
| def evict(path): |
| fd = os.open(path, os.O_RDONLY) |
| try: |
| sz = os.fstat(fd).st_size |
| step = 1 << 30 |
| off = 0 |
| while off < sz: |
| n = min(step, sz - off) |
| os.posix_fadvise(fd, off, n, os.POSIX_FADV_DONTNEED) |
| off += n |
| finally: |
| os.close(fd) |
|
|
| def read_status(pid): |
| d = {} |
| try: |
| with open(f"/proc/{pid}/status") as f: |
| for line in f: |
| k, _, v = line.partition(":") |
| d[k.strip()] = v.strip() |
| except OSError: |
| return None |
| return d |
|
|
| def kb(s): |
| if s is None: |
| return 0 |
| return int(s.split()[0]) * 1024 |
|
|
| def swap_used(): |
| with open("/proc/meminfo") as f: |
| for line in f: |
| if line.startswith("SwapTotal:"): |
| tot = int(line.split()[1]) |
| elif line.startswith("SwapFree:"): |
| free = int(line.split()[1]) |
| return (tot - free) * 1024 |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--bin", required=True) |
| ap.add_argument("--levels-mb", type=int, nargs="+", required=True) |
| ap.add_argument("--threads", type=int, default=8) |
| ap.add_argument("--ctx", type=int, default=2048) |
| ap.add_argument("--ubatch", type=int, default=128) |
| ap.add_argument("--n-predict", type=int, default=48) |
| ap.add_argument("--prompt", default="Write a detailed paragraph about mixture-of-experts routing.") |
| ap.add_argument("--out-json", required=True) |
| ap.add_argument("--evidence-dir", required=True) |
| args = ap.parse_args() |
|
|
| os.makedirs(args.evidence_dir, exist_ok=True) |
| records = [] |
| for mb in args.levels_mb: |
| lim = mb * 1024 * 1024 |
| evict(args.model) |
| swap0 = swap_used() |
| cmd = ["prlimit", f"--data={lim}", "--nofile=65536", args.bin, |
| "-m", args.model, "-c", str(args.ctx), "-b", str(args.ubatch), |
| "-t", str(args.threads), "-n", str(args.n_predict), "--temp", "0", |
| "-no-cnv", "--no-jinja", "-p", args.prompt] |
| t0 = time.time() |
| p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) |
| hwm_anon = 0 |
| hwm_rss = 0 |
| hwm_data = 0 |
| peak_read = 0 |
| infinity = 0 |
| while p.poll() is None: |
| st = read_status(p.pid) |
| if st: |
| hwm_anon = max(hwm_anon, kb(st.get("RssAnon"))) |
| hwm_rss = max(hwm_rss, kb(st.get("VmRSS"))) |
| hwm_data = max(hwm_data, kb(st.get("VmData"))) |
| try: |
| with open(f"/proc/{p.pid}/io") as f: |
| for line in f: |
| if line.startswith("read_bytes:"): |
| peak_read = max(peak_read, int(line.split()[1])) |
| except OSError: |
| pass |
| time.sleep(0.2) |
| out, err = p.communicate() |
| dt = time.time() - t0 |
| text = (out + b"\n" + err).decode("utf-8", "replace") |
| log = os.path.join(args.evidence_dir, f"upstream_{mb}m.log") |
| open(log, "w").write(text) |
|
|
| decode = prefill = None |
| for ln in text.replace("\r", "\n").splitlines(): |
| m = re.search(r"(prompt eval|eval) time =.*?([\d.]+)\s*tokens per second", ln) |
| if m: |
| if m.group(1) == "prompt eval": |
| prefill = float(m.group(2)) |
| else: |
| decode = float(m.group(2)) |
| fail = p.returncode != 0 or decode is None |
| rec = { |
| "level_mb": mb, |
| "limit_bytes": lim, |
| "returncode": p.returncode, |
| "seconds": round(dt, 2), |
| "anon_hwm_gib": round(hwm_anon / (1 << 30), 3), |
| "rss_hwm_gib": round(hwm_rss / (1 << 30), 3), |
| "vmdata_hwm_gib": round(hwm_data / (1 << 30), 3), |
| "read_bytes_gib": round(peak_read / (1 << 30), 3), |
| "swap_delta_gib": round(max(0, swap_used() - swap0) / (1 << 30), 3), |
| "decode_tps": decode, |
| "prefill_tps": prefill, |
| "result": "FAIL" if fail else "PASS", |
| "log": log, |
| } |
| records.append(rec) |
| print(json.dumps(rec), flush=True) |
|
|
| json.dump(records, open(args.out_json, "w"), indent=2) |
| print("wrote", args.out_json) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
