Ollama on CPU only: tokens per second on 4 cores

#!/usr/bin/env python3
"""Runs in the VM. bench.py <ollama_url> <outdir> <model> [<model> ...]
Per model: unload, one cold request (load time, time to first token), then 3 prompts x 3 repetitions with temperature 0, seed 42 and
num_predict 200, streamed. Tokens per second = eval_count / eval_duration (Ollama's own counters). Time to first token is the client-side time
until the first non-empty chunk. A failed load (HTTP error) is recorded with the message and the container's OOM flag."""
import sys, os, json, time, subprocess, urllib.request, urllib.error
url, out = sys.argv[1], sys.argv[2]; models = sys.argv[3:]; os.makedirs(out, exist_ok=True)
TEXT = ("A home server backup plan usually has three parts. First, a copy of the data is taken on a schedule, for example every night, by a tool that can detect which files changed. "
        "Second, the copy is stored somewhere that does not share a failure with the original: a second disk in another machine, a drive that is unplugged afterwards, or a remote object store. "
        "Third, the copy is restored on a spare machine at least once, because a backup that has never been restored is only a hope. Many people stop after the first part and discover the gaps "
        "during their first real failure: a database that was copied while it was being written, a volume that was left out of the job, or an encryption key that was kept only on the machine that died.")
PROMPTS = {"short-answer": "Explain in two sentences what a reverse proxy does and why someone would run one at home.",
           "summary": TEXT + "\n\nSummarize the passage above in three bullet points.",
           "code": "Write a Python function that parses one nginx access log line in the combined format and returns the HTTP status code and the request path. Include a short docstring."}
def post(path, body, stream=False, timeout=900):
    r = urllib.request.Request(url + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    return urllib.request.urlopen(r, timeout=timeout)
def sh(c): return subprocess.run(c, shell=True, capture_output=True, text=True).stdout
oname = sh("docker ps --format '{{.Names}}' | grep ollama").strip().splitlines()[0]
def gen(model, prompt):
    body = {"model": model, "prompt": prompt, "stream": True, "think": False, "keep_alive": "10m", "options": {"temperature": 0, "seed": 42, "num_predict": 200}}
    t0 = time.perf_counter(); first = None; final = None; text = ""
    with post("/api/generate", body) as r:
        for line in r:
            d = json.loads(line)
            if d.get("response") and first is None: first = time.perf_counter() - t0
            text += d.get("response", "")
            if d.get("done"): final = d
    return {"ttft_s": first, "wall_s": time.perf_counter() - t0, "eval_count": final["eval_count"], "eval_duration_s": final["eval_duration"] / 1e9, "tps": final["eval_count"] / (final["eval_duration"] / 1e9),
            "prompt_eval_count": final["prompt_eval_count"], "prompt_eval_s": final["prompt_eval_duration"] / 1e9, "load_s": final["load_duration"] / 1e9, "done_reason": final.get("done_reason"), "text_head": text[:120]}
res = {"ollama_container": oname, "models": {}}
for m in models:
    r = {"model": m}; res["models"][m] = r
    try: post("/api/generate", {"model": m, "keep_alive": 0}).read()
    except Exception as e: r["unload_note"] = str(e)[:100]
    time.sleep(3)
    csv = f"{out}/mem-{m.replace(':','_')}.csv"; sp = subprocess.Popen(f"while true; do echo $(date +%s.%N),$(docker stats --no-stream --format '{{{{.MemUsage}}}},{{{{.CPUPerc}}}}' {oname}); sleep 0.5; done > {csv}", shell=True, preexec_fn=os.setsid)
    r["t_start"] = time.time(); time.sleep(2)
    try:
        r["cold"] = gen(m, PROMPTS["short-answer"]); print(m, "cold", {k: v for k, v in r["cold"].items() if k != "text_head"}, flush=True)
        r["ps"] = sh(f"docker exec {oname} ollama ps")
        r["warm"] = []
        for name, p in PROMPTS.items():
            for rep in range(1, 4):
                x = gen(m, p); x.update(prompt=name, rep=rep); r["warm"].append(x); print(m, name, rep, round(x["tps"], 2), "tok/s", round(x["ttft_s"], 2), "s ttft", x["eval_count"], "tokens", flush=True)
    except urllib.error.HTTPError as e:
        r["error"] = {"http": e.code, "body": e.read().decode()[:500]}; print(m, "FAILED", r["error"], flush=True)
    except Exception as e:
        r["error"] = {"exception": repr(e)[:300]}; print(m, "FAILED", r["error"], flush=True)
    r["t_end"] = time.time(); time.sleep(2); os.killpg(os.getpgid(sp.pid), 15)
    r["oom_killed"] = sh(f"docker inspect -f '{{{{.State.OOMKilled}}}} restarts={{{{.RestartCount}}}}' {oname}").strip()
    r["dmesg_oom"] = sh("sudo dmesg | grep -i 'killed process\\|out of memory' | tail -3")
    r["logs_tail"] = sh(f"docker logs --tail 15 {oname} 2>&1 | cut -c1-260")
    json.dump(res, open(f"{out}/bench.json", "w"), indent=1)
    try: post("/api/generate", {"model": m, "keep_alive": 0}).read()
    except Exception: pass
res["free_m"] = sh("free -m"); json.dump(res, open(f"{out}/bench.json", "w"), indent=1)