105 lines
4.1 KiB
Python
105 lines
4.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Run the benchmark matrix and aggregate results.
|
|
|
|
Each cell shells out to agent_swarm.py (cache OFF baseline vs cache ON) for a
|
|
given (model, context, #agents, kv_type) on the currently-running stack, then
|
|
collects the JSON into a single matrix file consumed by plot.py.
|
|
|
|
The store-backend axis (RDMA / TCP / local-file) is swept by restarting the
|
|
store layer between groups (re-run proxy_start.sh with a different protocol).
|
|
"""
|
|
import argparse, json, os, subprocess, sys, time
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
def run_cell(py, cell, bridge, out_dir):
|
|
out = os.path.join(out_dir, f"cell_{cell['name']}.json")
|
|
# model_path may be a bare filename; resolve it against $MODELS_DIR so the
|
|
# checked-in config is portable across machines/checkouts.
|
|
model_path = cell["model_path"]
|
|
if not os.path.isabs(model_path) and not os.path.exists(model_path):
|
|
models_dir = os.environ.get("MODELS_DIR", "")
|
|
if models_dir:
|
|
cand = os.path.join(models_dir, os.path.basename(model_path))
|
|
if os.path.exists(cand):
|
|
model_path = cand
|
|
cmd = [py, os.path.join(HERE, "agent_swarm.py"),
|
|
"--bridge", bridge,
|
|
"--llamas", cell["llamas"],
|
|
"--model-path", model_path,
|
|
"--kv-type", cell.get("kv_type", "f16"),
|
|
"--agents", str(cell["agents"]),
|
|
"--ctx-tokens", str(cell["ctx_tokens"]),
|
|
"--share-mode", cell.get("share_mode", "shared_prefix"),
|
|
"--slots-per-server", str(cell.get("slots_per_server", 4)),
|
|
"--concurrency", str(cell.get("concurrency", 4)),
|
|
"--n-predict", str(cell.get("n_predict", 8)),
|
|
"--namespace", cell["name"],
|
|
"--out", out]
|
|
print(f"\n===== cell {cell['name']} =====")
|
|
print(" ", " ".join(cmd))
|
|
t0 = time.time()
|
|
r = subprocess.run(cmd, capture_output=True, text=True)
|
|
print(r.stdout[-1200:])
|
|
if r.returncode != 0:
|
|
print("STDERR:", r.stderr[-1500:])
|
|
return None
|
|
with open(out) as f:
|
|
data = json.load(f)
|
|
data["wall_total_s"] = round(time.time() - t0, 1)
|
|
data["store_backend"] = cell.get("store_backend", "rdma")
|
|
return data
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--config", required=True, help="JSON file with a list of cells")
|
|
ap.add_argument("--bridge", default="http://127.0.0.1:52052")
|
|
ap.add_argument("--out", default="run/matrix_results.json")
|
|
args = ap.parse_args()
|
|
|
|
with open(args.config) as f:
|
|
cfg = json.load(f)
|
|
out_dir = os.path.dirname(os.path.abspath(args.out))
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
results = []
|
|
for cell in cfg["cells"]:
|
|
d = run_cell(sys.executable, cell, args.bridge, out_dir)
|
|
if d:
|
|
results.append({"name": cell["name"], **summ(d)})
|
|
matrix = {"cells": results, "raw": cfg}
|
|
with open(args.out, "w") as f:
|
|
json.dump(matrix, f, indent=2)
|
|
|
|
print("\n================ MATRIX SUMMARY ================")
|
|
hdr = f"{'cell':<22}{'backend':<8}{'agents':>7}{'TTFT off':>10}{'TTFT on':>10}{'red%':>7}{'thru x':>8}{'recompute%':>12}"
|
|
print(hdr); print("-" * len(hdr))
|
|
for r in results:
|
|
print(f"{r['name']:<22}{r['backend']:<8}{r['agents']:>7}{r['ttft_off']:>10.0f}{r['ttft_on']:>10.0f}"
|
|
f"{r['ttft_red_pct']:>7.0f}{r['thru_x']:>8.2f}{r['recompute_on_pct']:>12.0f}")
|
|
print(f"\nsaved {args.out}")
|
|
|
|
|
|
def summ(d):
|
|
b, c = d["baseline"], d["cached"]
|
|
reuse = d["per_agent_cached"][1:] if len(d["per_agent_cached"]) > 1 else d["per_agent_cached"]
|
|
reuse_ttft = sum(r["ttft_ms"] for r in reuse) / max(len(reuse), 1)
|
|
return {
|
|
"backend": d.get("store_backend", "rdma"),
|
|
"agents": c["agents"],
|
|
"ttft_off": b["ttft_ms_mean"],
|
|
"ttft_on": c["ttft_ms_mean"],
|
|
"ttft_red_pct": d["ttft_reduction_pct"],
|
|
"reuser_ttft_on": round(reuse_ttft, 1),
|
|
"thru_x": d["throughput_speedup_x"],
|
|
"recompute_off_pct": round(b["recompute_ratio"] * 100, 1),
|
|
"recompute_on_pct": round(c["recompute_ratio"] * 100, 1),
|
|
"hit_agents": c["hit_agents"],
|
|
}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|