| |
| """K1 needle / cross-request corruption check for KV-layout changes (0057/0060) and long max-model-len. |
| |
| A haystack of public CourtListener material (the bench corpus) of an exact token length holds N "vault code" needles |
| at evenly spread depths; the model (T=0, thinking off, priority 10) must list every code. --concurrent C sends C |
| different haystacks (different seeds, so different codes) at once: a KV page shared or overwritten across requests shows |
| up as a missed code or, more tellingly, a FOREIGN code (one that belongs to another request in the same wave). |
| python3 offload/kvpack/needle.py --url http://SPARK_A_LAN_IP:8001 --ctx 250000 --needles 6 --out offload/results/K1/needle.jsonl |
| python3 offload/kvpack/needle.py --url ... --ctx 8000 30000 50000 100000 --concurrent 4 --out ... # corruption test |
| python3 offload/kvpack/needle.py --dry-run --ctx 8000 # offline: build + score |
| Exit status 1 when any request recalls fewer than --min-recall of its needles or returns a foreign code. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import random |
| import re |
| import sys |
| import threading |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
| NAMES = ["Abbeville", "Beaufort", "Calhoun", "Darlington", "Edgefield", "Fairfield", "Georgetown", "Horry", "Jasper", |
| "Kershaw", "Lancaster", "Marion", "Newberry", "Oconee", "Pickens", "Richland", "Saluda", "Union"] |
|
|
|
|
| def load_docs(data_dir: Path) -> list[dict]: |
| docs = [] |
| for name in ("opinions_cl.jsonl", "filings_recap.jsonl"): |
| with (data_dir / name).open() as f: |
| docs.extend(json.loads(line) for line in f) |
| return docs |
|
|
|
|
| class Builder: |
| def __init__(self, url: str | None, model: str | None, data_dir: Path): |
| self.url, self.model, self.docs = url, model, load_docs(data_dir) |
|
|
| def count(self, text: str) -> int: |
| if self.url is None: |
| return int(len(text) / 3.3) |
| req = urllib.request.Request(self.url + "/tokenize", data=json.dumps({"model": self.model, "prompt": text}).encode(), |
| headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(req, timeout=900) as r: |
| return int(json.load(r)["count"]) |
|
|
| def build(self, ctx: int, n: int, seed: int) -> tuple[str, dict[str, str]]: |
| rng = random.Random(seed) |
| names = rng.sample(NAMES, n) |
| codes = {name: f"{rng.randrange(10**5, 10**6)}" for name in names} |
| question = ("\n\nEND OF MATERIALS.\n\nSomewhere in the materials above are sentences giving the vault code for " |
| f"{n} vaults. List the vault code for each of these vaults: {', '.join(names)}. Answer only with " |
| "lines of the form `Name: code`, one per vault, nothing else.") |
| budget = ctx - self.count(question) - 64 * n |
| docs = self.docs[:] |
| rng.shuffle(docs) |
| parts, chars, ratio = [], 0, 3.3 |
| for d in docs * 4: |
| parts.append(f"=== {d.get('title', '')} ({d.get('court', '')}, {d.get('date', '')}) ===\n{d['text']}") |
| chars += len(parts[-1]) |
| if chars / ratio >= budget * 1.02: |
| got = self.count("\n\n".join(parts)) |
| ratio = chars / got |
| if got >= budget: |
| break |
| hay = "\n\n".join(parts) |
| while (got := self.count(hay)) > budget: |
| hay = hay[: int(len(hay) - (got - budget) * ratio) - 50] |
| |
| out, last = [], 0 |
| for i, name in enumerate(names): |
| pos = int(len(hay) * (0.05 + 0.9 * i / max(1, n - 1))) |
| cut = hay.find("\n", pos) |
| cut = len(hay) if cut < 0 else cut |
| out.append(hay[last:cut]) |
| out.append(f"\nThe vault code for {name} is {codes[name]}.\n") |
| last = cut |
| out.append(hay[last:]) |
| return "".join(out) + question, codes |
|
|
|
|
| def score(text: str, codes: dict[str, str], foreign: set[str]) -> dict: |
| found = {name: (m.group(1) if (m := re.search(rf"{name}[^\d\n]{{0,20}}(\d{{6}})", text)) else None) |
| for name in codes} |
| hits = sum(found[name] == code for name, code in codes.items()) |
| returned = set(re.findall(r"\b\d{6}\b", text)) |
| return {"recall": hits, "of": len(codes), "found": found, "foreign_codes": sorted(returned & foreign)} |
|
|
|
|
| def ask(url: str, model: str, prompt: str, max_tokens: int, salt: str) -> dict: |
| body = {"model": model, "messages": [{"role": "user", "content": prompt}], "temperature": 0, "top_p": 1, |
| "max_tokens": max_tokens, "priority": 10, "cache_salt": salt, |
| "chat_template_kwargs": {"thinking": False}} |
| req = urllib.request.Request(url + "/v1/chat/completions", data=json.dumps(body).encode(), |
| headers={"Content-Type": "application/json"}) |
| t0 = time.monotonic() |
| with urllib.request.urlopen(req, timeout=7200) as r: |
| res = json.load(r) |
| msg = res["choices"][0]["message"] |
| return {"text": msg.get("content") or "", "elapsed_s": round(time.monotonic() - t0, 2), "usage": res.get("usage")} |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--url") |
| ap.add_argument("--model") |
| ap.add_argument("--ctx", type=int, nargs="+", required=True) |
| ap.add_argument("--needles", type=int, default=6) |
| ap.add_argument("--concurrent", type=int, default=1) |
| ap.add_argument("--reps", type=int, default=1) |
| ap.add_argument("--seed", type=int, default=1022) |
| ap.add_argument("--max-tokens", type=int, default=300) |
| ap.add_argument("--min-recall", type=float, default=1.0) |
| ap.add_argument("--data-dir", type=Path, default=Path("/data/dsv41-offload/bench-data")) |
| ap.add_argument("--out", type=Path) |
| ap.add_argument("--dry-run", action="store_true") |
| args = ap.parse_args() |
| if not args.dry_run and not args.url: |
| ap.error("--url is required unless --dry-run") |
| url = None if args.dry_run else args.url.rstrip("/") |
| model = args.model |
| if url and not model: |
| with urllib.request.urlopen(url + "/v1/models", timeout=30) as r: |
| model = json.load(r)["data"][0]["id"] |
| b = Builder(url, model, args.data_dir) |
| rows, failed = [], False |
|
|
| def counters() -> dict[str, float]: |
| if url is None: |
| return {} |
| with urllib.request.urlopen(url + "/metrics", timeout=30) as r: |
| out: dict[str, float] = {} |
| for line in r.read().decode().splitlines(): |
| if line.startswith(("vllm:num_preemptions", "vllm:prefix_cache_queries", "vllm:prefix_cache_hits")): |
| name = line.split("{")[0].split(" ")[0] |
| out[name] = out.get(name, 0.0) + float(line.rsplit(" ", 1)[1]) |
| return out |
|
|
| m0 = counters() |
| for ctx in args.ctx: |
| for rep in range(args.reps): |
| seeds = [args.seed + 1000 * rep + 17 * i + ctx for i in range(args.concurrent)] |
| built = [b.build(ctx, args.needles, s) for s in seeds] |
| all_codes = [set(c.values()) for _, c in built] |
| results: list[dict | None] = [None] * len(built) |
|
|
| def run(i: int) -> None: |
| prompt, codes = built[i] |
| foreign = set().union(*(all_codes[:i] + all_codes[i + 1:])) if len(built) > 1 else set() |
| if url is None: |
| good = "\n".join(f"{k}: {v}" for k, v in codes.items()) |
| results[i] = {"dry_run": True, "prompt_chars": len(prompt), **score(good, codes, foreign)} |
| return |
| ans = ask(url, model, prompt, args.max_tokens, f"k1-needle-{seeds[i]}-{time.time_ns()}") |
| results[i] = {**ans, **score(ans["text"], codes, foreign)} |
|
|
| threads = [threading.Thread(target=run, args=(i,)) for i in range(len(built))] |
| for t in threads: |
| t.start() |
| for t in threads: |
| t.join() |
| for i, r in enumerate(results): |
| assert r is not None |
| ok = r["recall"] >= args.min_recall * r["of"] and not r["foreign_codes"] |
| failed |= not ok |
| row = {"ctx": ctx, "rep": rep, "slot": i, "concurrent": args.concurrent, "seed": seeds[i], "ok": ok, |
| **{k: v for k, v in r.items() if k != "text"}, "answer": r.get("text", "")[:400], |
| "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())} |
| rows.append(row) |
| usage = r.get("usage") or {} |
| print(f"ctx {ctx:>9,} rep {rep} slot {i}: {r['recall']}/{r['of']} " |
| f"{'OK' if ok else 'FAIL'} foreign={r['foreign_codes']} prompt_tokens={usage.get('prompt_tokens')} " |
| f"elapsed={r.get('elapsed_s')}", flush=True) |
| m1 = counters() |
| delta = {k: round(m1.get(k, 0.0) - m0.get(k, 0.0), 3) for k in sorted(set(m0) | set(m1))} |
| if delta: |
| print("metrics delta:", json.dumps(delta)) |
| if args.out: |
| args.out.parent.mkdir(parents=True, exist_ok=True) |
| with args.out.open("a") as f: |
| for row in rows: |
| f.write(json.dumps({**row, "metrics_delta": delta}) + "\n") |
| sys.exit(1 if failed else 0) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
