TerslerThe little junk rig that could

Source files / scripts · raw

concurrency-test.py

#!/usr/bin/env python3
"""Multi-session test against a running llama-server (agent-style use).

  concurrency-test.py [--port 8080] [--levels 1,2,4] [--replay] [--label NAME]

1. Concurrency: N parallel chat requests (distinct ~1.5k-token prompts, thinking
   off, MAX_TOKENS each). Reports per-request and total decode tok/s and draft
   acceptance.
2. --replay: agent-style prompt reuse, with a ~20k-token shared "system" prefix:
   a) cold prompt, b) same + one appended turn, c) a tool output edited in the
   middle of the history, d) the same system prefix on another slot (subagent).
   Reports how many prompt tokens the server had to (re)process each time.
3. --swarm K: like an agent harness: the root (no slot pinned) finishes a turn,
   then K subagents with the root's system prompt start at once, then the root
   continues. Reports the prompt tokens each one had to process.
Results are appended as JSON lines to --out.
"""
import argparse, glob, json, os, random, threading, time, urllib.request

ap = argparse.ArgumentParser()
ap.add_argument("--port", type=int, default=8080)
ap.add_argument("--levels", default="1,2,4")
ap.add_argument("--max-tokens", type=int, default=768)
ap.add_argument("--replay", action="store_true")
ap.add_argument("--swarm", type=int, default=0)
ap.add_argument("--label", default="")
ap.add_argument("--src", default="/opt/tersler/src/llama.cpp/src")
ap.add_argument("--out", default="/opt/tersler/bench/concurrency.jsonl")
a = ap.parse_args()
URL = f"http://127.0.0.1:{a.port}"


def post(path, body, timeout=3600):
    req = urllib.request.Request(URL + path, json.dumps(body).encode(), {"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def chat(messages, max_tokens, id_slot=None, think=False):
    body = {"messages": messages, "max_tokens": max_tokens, "stream": False,
            "chat_template_kwargs": {"enable_thinking": think}}
    if id_slot is not None:
        body["id_slot"] = id_slot
    t0 = time.time()
    r = post("/v1/chat/completions", body)
    wall = time.time() - t0
    t = r.get("timings")
    if not t:   # not llama.cpp (Strata): OpenAI usage only, timed on the client
        u = r.get("usage", {})
        cached = (u.get("prompt_tokens_details") or {}).get("cached_tokens") or 0
        return {"wall": wall, "prompt_n": u.get("prompt_tokens", 0) - cached, "cache_n": cached, "pp": None,
                "n": u.get("completion_tokens"), "tg": u.get("completion_tokens", 0) / wall, "accept": None}
    acc = t["draft_n_accepted"] / t["draft_n"] if t.get("draft_n") else None
    return {"wall": wall, "prompt_n": t.get("prompt_n"), "cache_n": t.get("cache_n"),
            "pp": t.get("prompt_per_second"), "n": t.get("predicted_n"), "tg": t.get("predicted_per_second"),
            "accept": acc}


def log(rec):
    rec.update(label=a.label, date=time.strftime("%Y-%m-%dT%H:%M:%S"))
    with open(a.out, "a") as f:
        f.write(json.dumps(rec) + "\n")


files = sorted(glob.glob(os.path.join(a.src, "*.cpp")))
random.seed(1)


def excerpt(i, chars):
    s = open(files[i % len(files)], errors="replace").read()
    return s[:chars]


TASK = ("Explain in detail what this source file does, section by section, for a new contributor. "
        "Write at least 800 words of continuous prose, no code blocks.")

for n in [int(x) for x in a.levels.split(",") if x]:
    res = [None] * n

    def run(i):
        nonce = f"[session {n}-{i}-{random.random():.6f}]\n"
        res[i] = chat([{"role": "user", "content": nonce + excerpt(3 * i + n, 5000) + "\n\n" + TASK}], a.max_tokens)

    t0 = time.time()
    th = [threading.Thread(target=run, args=(i,)) for i in range(n)]
    [x.start() for x in th]
    [x.join() for x in th]
    wall = time.time() - t0
    total = sum(r["n"] for r in res)
    accs = [r["accept"] for r in res if r["accept"] is not None]
    rec = {"test": "concurrency", "n": n, "wall_s": round(wall, 1), "tokens": total,
           "total_tok_s": round(total / wall, 1),
           "per_req_tg": [round(r["tg"], 1) for r in res],
           "accept": round(sum(accs) / len(accs), 3) if accs else None}
    print(json.dumps(rec), flush=True)
    log(rec)

if a.replay:
    # ~20k-token shared prefix: system prompt + "tool outputs" in the history
    system = "You are a coding agent. Reference material follows.\n\n" + "\n\n".join(excerpt(i, 12000) for i in range(6))
    hist = [{"role": "system", "content": system}]
    for k in range(4):
        hist += [{"role": "user", "content": f"Step {k}: read the next file."},
                 {"role": "assistant", "content": f"Reading file {k}."},
                 {"role": "user", "content": "Tool output:\n" + excerpt(10 + k, 8000)},
                 {"role": "assistant", "content": f"Noted file {k}."}]
    q = {"role": "user", "content": "In one sentence: which file defines the vocabulary?"}
    steps = []
    steps.append(("a_cold", chat(hist + [q], 32, id_slot=0)))
    hist2 = hist + [q, {"role": "assistant", "content": "llama-vocab.cpp."},
                    {"role": "user", "content": "Tool output:\n" + excerpt(20, 4000)}]
    steps.append(("b_append_turn", chat(hist2, 32, id_slot=0)))
    # prune a tool output in the middle (what agent harnesses do on compaction)
    edited = [dict(m) for m in hist2]
    edited[7]["content"] = "Tool output: [pruned]"
    steps.append(("c_edit_middle", chat(edited, 32, id_slot=0)))
    # subagent: same system prompt, own task, another slot
    try:
        slots = json.load(urllib.request.urlopen(URL + "/props")).get("total_slots", 1)
    except Exception:   # no /props outside llama.cpp
        slots = 1
    steps.append(("d_other_slot_same_prefix", chat([hist[0], {"role": "user", "content": "List three functions in the reference material."}], 32, id_slot=1 if slots > 1 else 0)))
    for name, r in steps:
        rec = {"test": "replay", "step": name, **{k: r[k] for k in ("prompt_n", "cache_n", "pp", "wall")}}
        rec["wall"] = round(rec["wall"], 1)
        print(json.dumps(rec), flush=True)
        log(rec)

if a.swarm:
    system = "You are a coding agent. Reference material follows.\n\n" + "\n\n".join(excerpt(30 + i, 12000) for i in range(6))
    root = [{"role": "system", "content": system},
            {"role": "user", "content": "Plan the work. Tool output:\n" + excerpt(40, 8000)}]
    r = chat(root, 32)
    out = [("root_turn1", r)]
    res = [None] * a.swarm

    def sub(i):
        res[i] = chat([root[0], {"role": "user", "content": f"Subtask {i}: summarise file {i} of the reference material in two sentences."}], 32)

    th = [threading.Thread(target=sub, args=(i,)) for i in range(a.swarm)]
    [x.start() for x in th]
    [x.join() for x in th]
    out += [(f"sub{i}", r) for i, r in enumerate(res)]
    root += [{"role": "assistant", "content": "Delegated."},
             {"role": "user", "content": "Subagent results: done. Continue."}]
    out.append(("root_turn2", chat(root, 32)))
    for name, r in out:
        rec = {"test": "swarm", "k": a.swarm, "step": name, **{k: r[k] for k in ("prompt_n", "cache_n", "pp", "wall")}}
        rec["wall"] = round(rec["wall"], 1)
        print(json.dumps(rec), flush=True)
        log(rec)