diff --git a/opt_prime/examples/pd/http_matrix/.gitignore b/opt_prime/examples/pd/http_matrix/.gitignore new file mode 100644 index 0000000..05f8f03 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/.gitignore @@ -0,0 +1,6 @@ +# Run outputs stay local; only the frozen held-out inputs are tracked. +evidence*/ +*.json +*.jsonl +*.log +!fixtures/fixture-manifest.json diff --git a/opt_prime/examples/pd/http_matrix/README.md b/opt_prime/examples/pd/http_matrix/README.md new file mode 100644 index 0000000..d6b047b --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/README.md @@ -0,0 +1,67 @@ +# HTTP validation and benchmark harness for PD serving + +Scripts that drive the OpenAI-style HTTP stage server (`opt_prime.accel.serve.llmd.serve_stage`) +as built by `examples/pd/serve_batched.py`, for aggregated and prefill/decode (PD) serving. +They write a JSON or JSONL record per run; outputs are git-ignored (see `.gitignore`). +Models are loaded with `local_files_only=True`, so pass a local model directory or a cached name. + +| Script | Purpose | +| --- | --- | +| `validate.py` | HTTP lifecycle check for one configuration in one process: readiness, 128/2K/8K-token completions, concurrency 4 and 8 against a serial baseline, stream cancel and client disconnect with drain checks, invalid-request rejection, prefix-cache and host-tier metrics. | +| `benchmark.py` | Warm HTTP comparison of `native` (one GPU), `dp2` (two replicas) and `pd` (prefill `cuda:0`, decode `cuda:1`). Starts isolated server processes, checks token parity against warmup references, then times streamed batches (first SSE event, e2e, estimated TPOT). | +| `benchmark_server.py` | Server launcher used by `benchmark.py`: runs `serve_batched.py` with the given flags, optionally loading an HTTP module from a file first. | +| `atom_server.py` | Single-device Rebellions ATOM decode server for `benchmark.py --server-config`, optionally fed by a subprocess prefill worker (`--worker-python` + `--engine-factory`). Requires `optimum-rbln` and exactly one visible `/dev/rbln*` device. | +| `heldout_quality.py` | Held-out structured-answer check: `fixtures` freezes 36 synthetic cases per model (addition, passkey, Korean table lookup at short/2K/8K tokens); `run` sends the frozen token IDs to an endpoint and scores exact JSON answers; `compare` (Python API) pairs two runs case by case. | +| `fixtures/` | The frozen inputs for four models (`*-fixtures.json.gz`) and `fixture-manifest.json` with the SHA-256 of each uncompressed file. | + +Tests: `opt_prime/tests/test_benchmark.py` and `opt_prime/tests/test_heldout_quality.py` (no GPU needed). + +## Running + +Paths below are relative to this directory. + +```bash +# HTTP lifecycle validation (torch backend; add --pd-prefill-device cuda:0 for PD) +python validate.py --model MODEL_DIR --device cuda:1 --output evidence/validate.json +python validate.py --model MODEL_DIR --device cuda:1 --prefix-cache --host-kv-gib 4 --output evidence/validate-host.json + +# Benchmark: native first, then reuse its prompts and token references for the other modes +python benchmark.py --model MODEL_DIR --mode native --output evidence/native.json +python benchmark.py --model MODEL_DIR --mode pd --reference evidence/native.json --output evidence/pd.json +python benchmark.py --model MODEL_DIR --mode dp2 --reference evidence/native.json --output evidence/dp2.json +``` + +`benchmark.py` defaults: input lengths 128 and 8192, 32 output tokens, concurrency 1 and 4, 3 repeats, +4 slots, bfloat16. It fails (and keeps the raw record) when output parity breaks. + +External backends: `--server-config FILE` takes JSON with `servers`, a list of argv arrays each +containing exactly one `"{port}"` argument, and `metadata` with resource facts (`physical_gpus`, +`physical_atom`, `physical_rngd`, `slots_per_decoder`, `producer_cache`, `decoder_cache`, +`producer_chunk`, `local_chunk`, `context`). For example: + +```json +{"servers": [["python", "atom_server.py", "--model", "MODEL_DIR", "--cache-dir", "ATOM_CACHE", "--port", "{port}"]], + "metadata": {"physical_gpus": 0, "physical_atom": 1, "physical_rngd": 0}} +``` + +`--http-module` (validate, benchmark) and the module options of `atom_server.py` load a trusted +local file in place of the in-tree module, for testing a candidate change in isolation. + +## Held-out structured answers + +```bash +# Use the frozen inputs as they are (gzip is read directly): +python heldout_quality.py run --fixtures fixtures/quality263-qwen4b-fixtures.json.gz \ + --endpoint http://localhost:8000 --route native-gpu --model SERVED_MODEL_NAME --output results.jsonl +# Or freeze new inputs for another model, before any inference: +python heldout_quality.py fixtures --model MODEL_DIR --output new-fixtures.json +``` + +The fixtures record the model path used when they were frozen (the Qwen3-0.6B file records the +Hub name `Qwen/Qwen3-0.6B` instead of a local cache path); the server rejects a request +whose `model` differs from its own, so pass `--model` with the served name (input IDs and hashes +are unchanged). Serve with greedy sampling and normal EOS; the runner sends `max_tokens` 128, +refuses to overwrite an existing output, and stops at the first HTTP failure. +`answer_correct` needs a parsed JSON object whose `answer` string matches exactly; `format_correct` +needs exactly one `answer` key. This measures structured-answer adherence on 36 synthetic cases, +not general model quality. diff --git a/opt_prime/examples/pd/http_matrix/atom_server.py b/opt_prime/examples/pd/http_matrix/atom_server.py new file mode 100644 index 0000000..7687f69 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/atom_server.py @@ -0,0 +1,76 @@ +"""Single assigned ATOM HTTP backend, optionally fed by the established RNGD subprocess adapter.""" +import argparse +import importlib.util +import signal +import sys +from pathlib import Path + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--model', required=True) + p.add_argument('--cache-dir', required=True) + p.add_argument('--port', required=True, type=int) + p.add_argument('--batching', help='optional trusted ATOM batching module file (default: the in-tree module)') + p.add_argument('--http-module', help='optional trusted HTTP module file (default: the in-tree module)') + p.add_argument('--pd-module', help='optional trusted subprocess worker module file (default: the in-tree module)') + p.add_argument('--worker-python') + p.add_argument('--engine-factory') + p.add_argument('--width', type=int, default=8192) + p.add_argument('--kv-blocks', type=int, default=3) + p.add_argument('--max-output', type=int, default=128) + a = p.parse_args() + if bool(a.worker_python) != bool(a.engine_factory): + p.error('worker-python and engine-factory must be supplied together') + for name, path in [('opt_prime.accel.pd.pd_batching', a.pd_module), + ('opt_prime.accel.platforms.atom.batching', a.batching), + ('opt_prime.accel.serve.llmd', a.http_module)]: + if not path: + continue + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + from opt_prime.accel.pd.pd_batching import SubprocessPrefillWorker + from opt_prime.accel.platforms.atom.engine import AtomRblnEngine + from opt_prime.accel.serve.batched import BatchedOwner, BatchedService + from opt_prime.accel.serve.llmd import serve_stage + from transformers import AutoTokenizer + + nodes = list(Path('/dev').glob('rbln[0-9]*')) + if len(nodes) != 1: + raise RuntimeError('benchmark requires exactly one assigned ATOM device') + engine = AtomRblnEngine(devices=[int(nodes[0].name[4:])], max_seq_len=9216, + batch_size=1, kv_blocks=a.kv_blocks, cache_dir=a.cache_dir) + worker = owner = server = http = None + try: + engine.start(a.model) + c = engine.model.config.get_text_config() + if a.worker_python: + worker = SubprocessPrefillWorker(a.worker_python, + ['--model', a.model, '--width', str(a.width), '--engine-factory', a.engine_factory], + slots=1, layers=c.num_hidden_layers, kv_heads=c.num_key_value_heads, + head_dim=getattr(c, 'head_dim', None) or c.hidden_size // c.num_attention_heads, + max_context=a.width) + tok = AutoTokenizer.from_pretrained(a.model, local_files_only=True) + batcher = engine.batcher(prefix_cache=False, prefill_worker=worker, local_prefill=worker is None) + service = BatchedService(batcher, model=a.model, vocab_size=c.vocab_size, + max_output=a.max_output, tokenizer=tok, + render=lambda ids: tok.decode(ids, skip_special_tokens=True)) + owner = BatchedOwner(service) + # The benchmark uses fixed-length ignore_eos requests; no EOS-dependent timing comparisons. + signal.signal(signal.SIGTERM, lambda *_: owner.submit('stop')) + server, http = serve_stage(owner, 'decode', a.port) + owner.run() + finally: + if server: + server.shutdown() + server.server_close() + http.join(5) + if worker: + worker.close() + engine.close() + + +if __name__ == '__main__': + main() diff --git a/opt_prime/examples/pd/http_matrix/benchmark.py b/opt_prime/examples/pd/http_matrix/benchmark.py new file mode 100644 index 0000000..8c273ba --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/benchmark.py @@ -0,0 +1,264 @@ +"""Warm HTTP native/DP/PD comparison; isolated server processes, exact controls and SSE timings.""" +import argparse +import hashlib +import json +import os +import signal +import socket +import statistics +import subprocess +import sys +import threading +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + + +def timing_summary(rows, wall_s): + return { + "requests": len(rows), "output_tokens": sum(r["output_tokens"] for r in rows), + "wall_s": wall_s, + "output_tok_s": sum(r["output_tokens"] for r in rows) / wall_s, + "median_first_sse_s": statistics.median(r["first_sse_s"] for r in rows), + "median_tpot_est_s": statistics.median(r["tpot_est_s"] for r in rows), + "median_e2e_s": statistics.median(r["e2e_s"] for r in rows), + "all_text_equal": all(r["text_equal"] for r in rows), + } + + +RESOURCE_FIELDS = {"physical_gpus", "physical_atom", "physical_rngd", "slots_per_decoder", + "producer_cache", "decoder_cache", "producer_chunk", "local_chunk", "context"} + + +def load_server_config(path): + """Trusted executable argv arrays; no shell expansion. Resource facts are supplied by the experiment.""" + config = json.loads(Path(path).read_text()) + if not isinstance(config, dict): + raise TypeError("server config must be an object") + servers, metadata = config.get("servers"), config.get("metadata") + if not isinstance(servers, list) or not servers or not isinstance(metadata, dict): + raise ValueError("server config needs nonempty servers and metadata object") + if set(metadata) - RESOURCE_FIELDS: + raise ValueError("unsupported resource metadata key") + for argv in servers: + if not isinstance(argv, list) or not argv or not all(isinstance(x, str) and x for x in argv): + raise ValueError("each server must be a nonempty argv array") + if argv.count("{port}") != 1: + raise ValueError("each server needs exactly one standalone {port} argument") + return config + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model", required=True) + p.add_argument("--mode", choices=["native", "dp2", "pd"], required=True) + p.add_argument("--output", required=True) + p.add_argument("--startup-timeout", type=int, default=180, help="server readiness deadline in seconds; excluded from timing") + p.add_argument("--server-config", help="trusted JSON argv templates and resource metadata for external backends") + p.add_argument("--reference", help="native benchmark JSON with identical token-ID prompts") + p.add_argument("--http-module", help="trusted local HTTP candidate") + p.add_argument("--lengths", nargs="+", type=int, default=[128, 8192]) + p.add_argument("--outputs", nargs="+", type=int, default=[32]) + p.add_argument("--concurrency", nargs="+", type=int, default=[1, 4]) + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--slots", type=int, default=4) + p.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16") + a = p.parse_args() + if a.startup_timeout <= 0: + p.error("startup timeout must be positive") + external = load_server_config(a.server_config) if a.server_config else None + if min(*a.lengths, *a.outputs, *a.concurrency, a.repeats, a.slots) < 1: + p.error("sizes must be positive") + if max(a.lengths) + max(a.outputs) > 9216: + p.error("input plus output exceeds context9216") + import torch + import transformers + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(a.model, local_files_only=True) + head = [tok.bos_token_id] if tok.bos_token_id is not None else [] + fill = tok.encode("This is background context. The red box contains apples. The blue box contains books.\n", + add_special_tokens=False) + tail = tok.encode("\nQuestion: What is the capital of France?\nAnswer:", add_special_tokens=False) + prompts = {} + for n in a.lengths: + left = n - len(head) - len(tail) + if left < 0 or not fill: + p.error("prompt length too short") + prompts[str(n)] = head + (fill * ((left + len(fill) - 1) // len(fill)))[:left] + tail + assert len(prompts[str(n)]) == n + reference = json.loads(Path(a.reference).read_text()) if a.reference else None + if reference: + assert reference["prompts"] == prompts and reference["dtype"] == a.dtype + refs = {} if reference is None else reference["references"].copy() + output = Path(a.output) + output.parent.mkdir(parents=True, exist_ok=True) + procs, logs, ports = [], [], [] + result = {"passed": False, "model": Path(a.model).name, "mode": a.mode, "dtype": a.dtype, + "torch": torch.__version__, "transformers": transformers.__version__, + "replicas": 2 if a.mode == "dp2" else 1, "slots_per_decoder": a.slots, + "physical_gpus": 1 if a.mode == "native" else 2, "cpu_threads_per_server": 4, + "producer_cache": False, "decoder_cache": False, "producer_chunk": 256, + "local_chunk": 256, "context": 9216, "prompts": prompts, + "warmups": [], "correctness_batches": [], "batches": [], + "harness_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "startup_timeout_s": a.startup_timeout, + "timing_scope": "Client first SSE event, full response e2e and estimated TPOT=(e2e-first_event)/(output_tokens-1). SSE can coalesce tokens; this is not per-token kernel timing. Loading/warmup/correctness checks excluded."} + devices = ["cuda:0", "cuda:1"] if a.mode == "dp2" else ["cuda:1"] + if external: + devices = [None] * len(external["servers"]) + for key in RESOURCE_FIELDS: + result.pop(key, None) + result.update(external["metadata"]) + result["replicas"] = len(devices) + result["server_config"] = external + started = time.perf_counter() + try: + for endpoint, device in enumerate(devices): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + log = output.with_suffix(f".server{endpoint}.log").open("w") + logs.append(log) + cmd = [sys.executable, str(Path(__file__).with_name("benchmark_server.py")), + "--model", a.model, "--device", device, "--dtype", a.dtype, + "--slots", str(a.slots), "--max-context", "9216", "--max-output", str(max(a.outputs)), + "--kv", "paged", "--page-size", "64", "--chunk", "256", "--port", str(port)] + if a.http_module: + cmd += ["--http-module", a.http_module] + if a.mode == "pd": + cmd += ["--pd-prefill-device", "cuda:0", "--pd-prefill-no-prefix-cache", + "--pd-prefill-chunk", "256"] + if external: + cmd = [str(port) if arg == "{port}" else arg for arg in external["servers"][endpoint]] + proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, start_new_session=True, + env=dict(os.environ, OMP_NUM_THREADS="4", MKL_NUM_THREADS="4")) + procs.append(proc) + ports.append(port) + end = time.monotonic() + a.startup_timeout + while True: + if proc.poll() is not None: + raise RuntimeError(f"server{endpoint} exited{proc.returncode}; see server log") + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/ready", timeout=1) as response: + if response.status == 200 and json.load(response)["ready"]: + break + except (OSError, urllib.error.HTTPError): + pass + if time.monotonic() >= end: + raise TimeoutError("server startup") + time.sleep(.05) + result["load_s"] = time.perf_counter() - started + + def request(n, count, rid, endpoint, stream=False, gate=None): + body = {"model": a.model, "prompt": prompts[str(n)], "max_tokens": count, + "temperature": 0, "ignore_eos": True, "stream": stream} + req = urllib.request.Request(f"http://127.0.0.1:{ports[endpoint]}/v1/completions", + data=json.dumps(body).encode(), headers={"x-request-id": rid}) + if gate is not None: + gate.wait(timeout=30) + start = time.perf_counter() + key = f"{n}:{count}" + with urllib.request.urlopen(req, timeout=180) as response: + assert response.status == 200 + if not stream: + out = json.load(response) + choice = out["choices"][0] + assert choice["finish_reason"] == "length" and len(choice["token_ids"]) == count + return {"id": rid, "endpoint": endpoint, "tokens": choice["token_ids"], + "text": choice["text"], "tokens_equal": key not in refs or choice["token_ids"] == refs[key]["tokens"]} + first_event = None + usage = finish = None + done = False + text = [] + for line in response: + if not line.startswith(b"data: "): + continue + if first_event is None: + first_event = time.perf_counter() - start + if line[6:].strip() == b"[DONE]": + done = True + break + event = json.loads(line[6:]) + if "usage" in event: + usage = event["usage"] + for choice in event.get("choices", []): + text.append(choice.get("text", "")) + if choice.get("finish_reason"): + finish = choice["finish_reason"] + elapsed = time.perf_counter() - start + assert done and finish == "length" and usage["completion_tokens"] == count + text = "".join(text) + return {"id": rid, "endpoint": endpoint, "input_tokens": n, "output_tokens": count, + "text": text, "text_equal": text == refs[key]["text"], "first_sse_s": first_event, + "tpot_est_s": (elapsed - first_event) / (count - 1) if count > 1 else 0, + "e2e_s": elapsed} + + def batch(n, count, concurrency, rid, repeat, stream): + gate = threading.Barrier(concurrency + 1) + with ThreadPoolExecutor(max_workers=concurrency) as clients: + futures = [clients.submit(request, n, count, f"{rid}-{i}", + (repeat + i) % len(ports), stream, gate) for i in range(concurrency)] + start = time.perf_counter() + gate.wait(timeout=30) + rows = [future.result(timeout=240) for future in futures] + wall = time.perf_counter() - start + return rows, wall + + warm_start = time.perf_counter() + for endpoint in range(len(ports)): + for n in a.lengths: + for count in a.outputs: + row = request(n, count, f"warm-{endpoint}-{n}-{count}", endpoint) + refs.setdefault(f"{n}:{count}", {"tokens": row["tokens"], "text": row["text"]}) + result["warmups"].append(row) + result["references"] = refs + for n in a.lengths: + for count in a.outputs: + for concurrency in a.concurrency: + rows, _ = batch(n, count, concurrency, f"check-{n}-{count}-{concurrency}", 0, False) + result["correctness_batches"].append({"input_tokens": n, "output_tokens": count, + "concurrency": concurrency, "rows": rows}) + result["warmup_and_checks_s"] = time.perf_counter() - warm_start + for repeat in range(a.repeats): + for n in (a.lengths if repeat % 2 == 0 else list(reversed(a.lengths))): + for count in a.outputs: + for concurrency in a.concurrency: + rows, wall = batch(n, count, concurrency, f"r{repeat}-n{n}-o{count}-c{concurrency}", repeat, True) + result["batches"].append({"repeat": repeat, "input_tokens": n, + "output_tokens_per_request": count, "concurrency": concurrency, + "summary": timing_summary(rows, wall), "rows": rows}) + result["metrics"] = [] + for port in ports: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/metrics", timeout=10) as response: + result["metrics"].append(response.read().decode().replace(a.model, result["model"])) + result["parity_passed"] = all(r["tokens_equal"] for r in result["warmups"]) and all( + r["tokens_equal"] for b in result["correctness_batches"] for r in b["rows"]) and all( + b["summary"]["all_text_equal"] for b in result["batches"]) + result["passed"] = True + except Exception as error: + result["error"] = repr(error) + raise + finally: + for proc in procs: + try: + os.killpg(proc.pid, signal.SIGTERM) + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(proc.pid, signal.SIGKILL) + proc.wait(timeout=10) + except ProcessLookupError: + proc.wait(timeout=10) + for log in logs: + log.close() + output.write_text(json.dumps(result, indent=2) + "\n") + if not result["parity_passed"]: + raise AssertionError("output parity failed; raw data retained, no equivalent-output speedup claim") + print(json.dumps({"passed": result["passed"], "parity_passed": result["parity_passed"], + "batches": len(result["batches"]), "mode": a.mode}), flush=True) + + +if __name__ == "__main__": + main() diff --git a/opt_prime/examples/pd/http_matrix/benchmark_server.py b/opt_prime/examples/pd/http_matrix/benchmark_server.py new file mode 100644 index 0000000..e37fa5f --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/benchmark_server.py @@ -0,0 +1,18 @@ +"""Launch the production serving CLI with an optional trusted HTTP candidate.""" +import argparse +import importlib.util +import runpy +import sys +from pathlib import Path + +p = argparse.ArgumentParser(description=__doc__, add_help=False) +p.add_argument("--http-module") +a, rest = p.parse_known_args() +if a.http_module: + spec = importlib.util.spec_from_file_location("opt_prime.accel.serve.llmd", a.http_module) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) +path = Path(__file__).resolve().parents[1] / "serve_batched.py" +sys.argv = [str(path), *rest] +runpy.run_path(str(path), run_name="__main__") diff --git a/opt_prime/examples/pd/http_matrix/fixtures/fixture-manifest.json b/opt_prime/examples/pd/http_matrix/fixtures/fixture-manifest.json new file mode 100644 index 0000000..4bdc540 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/fixtures/fixture-manifest.json @@ -0,0 +1,74 @@ +{ + "quality263-gemma-fixtures.json": { + "sha256": "0c24b1e759148dd7de68e15fa435e03c77a85b947fd6e923eee46af8c87edf44", + "cases": 36, + "length_ranges": { + "0": [ + 38, + 68 + ], + "2048": [ + 2034, + 2048 + ], + "8192": [ + 8183, + 8187 + ] + } + }, + "quality263-koni-fixtures.json": { + "sha256": "22f6e5b117913530ef6d58ed5352bb47684d55409535157b75f0462d64797eaf", + "cases": 36, + "length_ranges": { + "0": [ + 150, + 182 + ], + "2048": [ + 2040, + 2045 + ], + "8192": [ + 8186, + 8191 + ] + } + }, + "quality263-qwen06b-fixtures.json": { + "sha256": "5fd8133ab6026cf757ae005d671e7778c69e141660d1771627cf1df7070e6a83", + "cases": 36, + "length_ranges": { + "0": [ + 40, + 77 + ], + "2048": [ + 2037, + 2046 + ], + "8192": [ + 8183, + 8192 + ] + } + }, + "quality263-qwen4b-fixtures.json": { + "sha256": "701be76546e6fa3c7410475bc026a50c21c677abb7725565fb650b514fc71c98", + "cases": 36, + "length_ranges": { + "0": [ + 40, + 77 + ], + "2048": [ + 2037, + 2046 + ], + "8192": [ + 8183, + 8192 + ] + } + } +} diff --git a/opt_prime/examples/pd/http_matrix/fixtures/quality263-gemma-fixtures.json.gz b/opt_prime/examples/pd/http_matrix/fixtures/quality263-gemma-fixtures.json.gz new file mode 100644 index 0000000..4cb6f6c Binary files /dev/null and b/opt_prime/examples/pd/http_matrix/fixtures/quality263-gemma-fixtures.json.gz differ diff --git a/opt_prime/examples/pd/http_matrix/fixtures/quality263-koni-fixtures.json.gz b/opt_prime/examples/pd/http_matrix/fixtures/quality263-koni-fixtures.json.gz new file mode 100644 index 0000000..2b2316a Binary files /dev/null and b/opt_prime/examples/pd/http_matrix/fixtures/quality263-koni-fixtures.json.gz differ diff --git a/opt_prime/examples/pd/http_matrix/fixtures/quality263-qwen06b-fixtures.json.gz b/opt_prime/examples/pd/http_matrix/fixtures/quality263-qwen06b-fixtures.json.gz new file mode 100644 index 0000000..a3e39ec Binary files /dev/null and b/opt_prime/examples/pd/http_matrix/fixtures/quality263-qwen06b-fixtures.json.gz differ diff --git a/opt_prime/examples/pd/http_matrix/fixtures/quality263-qwen4b-fixtures.json.gz b/opt_prime/examples/pd/http_matrix/fixtures/quality263-qwen4b-fixtures.json.gz new file mode 100644 index 0000000..87fe0d0 Binary files /dev/null and b/opt_prime/examples/pd/http_matrix/fixtures/quality263-qwen4b-fixtures.json.gz differ diff --git a/opt_prime/examples/pd/http_matrix/heldout_quality.py b/opt_prime/examples/pd/http_matrix/heldout_quality.py new file mode 100644 index 0000000..f9fd246 --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/heldout_quality.py @@ -0,0 +1,139 @@ +"""Preregistered synthetic structured-answer checks; not a general quality benchmark.""" +import argparse +import gzip +import hashlib +import json +import random +import time +import urllib.request +from pathlib import Path + + +def score(text, expected): + def unique(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError('duplicate JSON key') + result[key] = value + return result + try: + value = json.loads(text, object_pairs_hook=unique) + except (ValueError, TypeError): + value = None + valid = isinstance(value, dict) and isinstance(value.get('answer'), str) + return dict(answer_correct=bool(valid and value['answer'] == expected), + format_correct=bool(valid and set(value) == {'answer'})) + + +def compare(baseline, candidate): + left = {x['case']: x for x in baseline} + right = {x['case']: x for x in candidate} + if len(left) != len(baseline) or len(right) != len(candidate) or left.keys() != right.keys(): + raise ValueError('duplicate or missing paired cases') + result = dict(cases=len(left), correct_to_wrong=0, wrong_to_correct=0, token_equal=0) + for key, a in left.items(): + b = right[key] + if a['input_sha256'] != b['input_sha256']: + raise ValueError('paired input IDs differ') + result['correct_to_wrong'] += bool(a['answer_correct'] and not b['answer_correct']) + result['wrong_to_correct'] += bool(not a['answer_correct'] and b['answer_correct']) + result['token_equal'] += a['tokens'] == b['tokens'] + return result + + +def fixtures(tokenizer): + rng = random.Random(26320261024) + tasks = [] + for variant in range(4): + a, b = rng.randrange(21, 80), rng.randrange(11, 60) + tasks.append(('arithmetic', variant, '', f'Compute {a} + {b}.', str(a + b))) + secret = ''.join(rng.choice('ABCDEFGHJKLMNPQRSTUVWXYZ23456789') for _ in range(8)) + tasks.append(('passkey', variant, f'The access key for archive {variant} is {secret}.', + f'What is the access key for archive {variant}?', secret)) + city = ['춘천', '목포', '전주', '속초'][variant] + tasks.append(('korean_table', variant, + f'배송표: 주문 가{variant}=대전; 주문 나{variant}={city}; 주문 다{variant}=부산.', + f'배송표에서 주문 나{variant}의 도착 도시를 찾으세요.', city)) + filler = 'Background: shelves hold ordinary books; this note contains no task answer.\n' + result = [] + for target in [0, 2048, 8192]: + for task, variant, fact, question, expected in tasks: + def render(repeats): + prompt = (filler * (repeats // 2) + fact + '\n' + filler * (repeats - repeats // 2) + + '\n' + question + '\nReturn only a JSON object with one key "answer" and a string value. No explanation.') + ids = tokenizer.apply_chat_template([{'role': 'user', 'content': prompt}], + tokenize=True, return_dict=False, add_generation_prompt=True, enable_thinking=False) + return prompt, [int(x) for x in ids] + low, high = 0, target + if target and len(render(0)[1]) > target: + raise ValueError('base prompt exceeds context target') + while low < high: + mid = (low + high + 1) // 2 + if len(render(mid)[1]) <= target: + low = mid + else: + high = mid - 1 + prompt, ids = render(low) + digest = hashlib.sha256(json.dumps(ids, separators=(',', ':')).encode()).hexdigest() + result.append(dict(case=f'{task}-{variant}-{target}', task=task, target_tokens=target, + expected=expected, prompt=prompt, ids=ids, input_sha256=digest)) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest='command', required=True) + export = sub.add_parser('fixtures') + export.add_argument('--model', required=True) + export.add_argument('--output', type=Path, required=True) + run = sub.add_parser('run') + run.add_argument('--fixtures', type=Path, required=True) + run.add_argument('--endpoint', required=True) + run.add_argument('--route', required=True) + run.add_argument('--output', type=Path, required=True) + run.add_argument('--model', help='request model name (default: the name frozen in the fixture); ' + 'input IDs and hashes are unchanged') + args = parser.parse_args() + if args.command == 'fixtures': + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True) + data = dict(model=args.model, seed=26320261024, max_tokens=128, + thinking=False, cases=fixtures(tokenizer)) + args.output.write_text(json.dumps(data, ensure_ascii=False, indent=2) + '\n') + return + opener = gzip.open if args.fixtures.suffix == '.gz' else open + with opener(args.fixtures, 'rt', encoding='utf-8') as source: + data = json.load(source) + model = args.model or data['model'] + # Exclusive creation prevents silently replacing a prior run after seeing answers. + with args.output.open('x') as output: + for case in data['cases']: + body = dict(model=model, prompt=case['ids'], temperature=0, + max_tokens=data['max_tokens'], ignore_eos=False, stream=False) + request = urllib.request.Request(args.endpoint.rstrip('/') + '/v1/completions', + data=json.dumps(body).encode(), headers={'Content-Type': 'application/json', + 'x-request-id': 'quality263-' + case['case']}) + start = time.perf_counter() + row = dict(case=case['case'], route=args.route, input_sha256=case['input_sha256'], + input_tokens=len(case['ids']), expected=case['expected']) + try: + with urllib.request.urlopen(request, timeout=600) as response: + payload = json.load(response) + choice = payload['choices'][0] + row.update(text=choice['text'], tokens=choice['token_ids'], + finish_reason=choice['finish_reason'], response=payload, + **score(choice['text'], case['expected'])) + except Exception as error: + row.update(error=repr(error), answer_correct=False, format_correct=False) + output.write(json.dumps(row, ensure_ascii=False) + '\n') + output.flush() + raise + row['elapsed_seconds'] = time.perf_counter() - start + output.write(json.dumps(row, ensure_ascii=False) + '\n') + output.flush() + print(case['case'], row['answer_correct'], row['format_correct'], flush=True) + + +if __name__ == '__main__': + main() diff --git a/opt_prime/examples/pd/http_matrix/validate.py b/opt_prime/examples/pd/http_matrix/validate.py new file mode 100644 index 0000000..ab19bcc --- /dev/null +++ b/opt_prime/examples/pd/http_matrix/validate.py @@ -0,0 +1,310 @@ +"""Actual HTTP validation through serve_batched.build; one model/configuration per process.""" + +import argparse +import importlib.util +import itertools +import json +import sys +import threading +import time +import urllib.error +import urllib.request +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from types import SimpleNamespace + +import torch +from transformers import AutoTokenizer + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--model", required=True) + p.add_argument("--output", required=True) + p.add_argument("--http-module", help="trusted local HTTP module candidate for isolated validation") + p.add_argument("--device", default="cuda:0") + p.add_argument("--pd-prefill-device", default=None) + p.add_argument("--pd-prefill-no-prefix-cache", action="store_true") + p.add_argument("--pd-prefill-chunk", type=int, default=2048) + p.add_argument("--prefix-cache", action="store_true") + p.add_argument("--backend", choices=["torch", "atom"], default="torch") + p.add_argument("--atom-kv-blocks", type=int, default=3) + p.add_argument("--host-kv-gib", type=float, default=0) + p.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16") + a = p.parse_args() + if a.http_module: + spec = importlib.util.spec_from_file_location("opt_prime.accel.serve.llmd", a.http_module) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + # Load consumers after the candidate so StageError identity is shared across the server and owner. + from opt_prime.accel.serve.llmd import serve_stage + from serve_batched import build + if a.backend == "atom" and a.dtype != "bfloat16": + p.error("ATOM precision is determined by the compiled artifact") + if a.host_kv_gib < 0 or (a.host_kv_gib and not a.prefix_cache): + p.error("host KV requires --prefix-cache and a nonnegative capacity") + devices = [] + if a.backend == "atom": + nodes = list(Path("/dev").glob("rbln[0-9]*")) + assert len(nodes) == 1, "requires one isolated allocated ATOM" + devices = [int(nodes[0].name[4:])] + torch.set_num_threads(4) + args = SimpleNamespace( + model=a.model, + backend=a.backend, + device=a.device, + dtype=a.dtype, + slots=4 if a.backend == "torch" else 1, + devices=devices, + atom_kv_blocks=a.atom_kv_blocks, + pd_remote_only=False, + max_context=9216, + max_output=256, + kv="paged", + page_size=64, + chunk=256, + prefix_cache=a.prefix_cache, + kv_events=None, + kv_host_memory_gb=a.host_kv_gib, + reload_bandwidth=None, + prefill_rate=None, + remote_kv=None, + pd_prefill_device=a.pd_prefill_device, + pd_prefill_no_prefix_cache=a.pd_prefill_no_prefix_cache, + pd_prefill_chunk=a.pd_prefill_chunk, + pd_prefill_python=None, + ) + result = { + "model_name": Path(a.model).name, + "prefix_cache": a.prefix_cache, + "host_kv_gib": a.host_kv_gib, + "passed": False, + "scope": "HTTP lifecycle and completion counts; numerical equality reported separately", + "torch": torch.__version__, + "dtype": a.dtype if a.backend == "torch" else "compiled artifact", + "backend": a.backend, + "decode_device": a.device, + "pd_prefill_device": a.pd_prefill_device, + "producer_prefix_cache": None, + "slots": args.slots, + "atom_kv_blocks": a.atom_kv_blocks if a.backend == "atom" else None, + "device": torch.cuda.get_device_name(torch.device(a.device)) + if a.backend == "torch" + else "allocated ATOM", + "cases": [], + } + output = Path(a.output) + output.parent.mkdir(parents=True, exist_ok=True) + try: + owner = build(args) + if a.pd_prefill_device: + producer = owner.service.batcher.worker.engine + result["producer_prefix_cache"] = producer.kv.prefix_cache + result["producer_chunk"] = producer.chunk + except Exception as exc: + result["startup_error"] = repr(exc) + output.write_text(json.dumps(result, indent=2) + "\n") + raise + tok = AutoTokenizer.from_pretrained(a.model, local_files_only=True) + head = [tok.bos_token_id] if tok.bos_token_id is not None else [] + fill = tok.encode( + "This is background context. The red box contains apples. The blue box contains books.\n", + add_special_tokens=False, + ) + tail = tok.encode( + "\nQuestion: What is the capital of France?\nAnswer:", add_special_tokens=False + ) + prompts = {} + for n in (128, 2048, 8192): + left = n - len(head) - len(tail) + assert left >= 0 and fill + prompts[n] = head + (fill * ((left + len(fill) - 1) // len(fill)))[:left] + tail + assert len(prompts[n]) == n + server, server_thread = serve_stage(owner, "decode", 0) + url = f"http://127.0.0.1:{server.server_port}" + serial = itertools.count() + + def request(path, body=None, rid=None): + req = urllib.request.Request( + url + path, + data=None if body is None else json.dumps(body).encode(), + headers={ + "Content-Type": "application/json", + "x-request-id": rid or f"http211-{next(serial)}", + }, + ) + return urllib.request.urlopen(req, timeout=180) + + def body(n, count=32, stream=False): + return { + "model": a.model, + "prompt": prompts[n], + "max_tokens": count, + "temperature": 0, + "ignore_eos": True, + "stream": stream, + } + + def complete(n): + start = time.perf_counter() + with request("/v1/completions", body(n)) as response: + assert response.status == 200 + data = json.load(response) + tokens = data["choices"][0]["token_ids"] + assert len(tokens) == 32 and data["usage"]["completion_tokens"] == 32 + return { + "input_tokens": n, + "tokens": tokens, + "e2e_s": time.perf_counter() - start, + } + + def drained(rid, timeout=30): + end = time.monotonic() + timeout + while rid in owner.service.batcher.by_id: + if time.monotonic() > end: + raise TimeoutError(f"undrained request {rid}") + time.sleep(0.01) + + def client(): + try: + end = time.monotonic() + 30 + while not owner.started: + if time.monotonic() > end: + raise TimeoutError("owner startup") + time.sleep(0.01) + with request("/ready") as response: + assert response.status == 200 + baseline = {n: complete(n) for n in prompts} + result["baseline"] = baseline + for concurrency in (4, 8): + lengths = [(128, 2048, 8192)[i % 3] for i in range(concurrency)] + start = time.perf_counter() + with ThreadPoolExecutor(max_workers=concurrency) as pool: + rows = list(pool.map(complete, lengths)) + for row in rows: + row["baseline_equal"] = ( + row["tokens"] == baseline[row["input_tokens"]]["tokens"] + ) + result["cases"].append( + { + "concurrency": concurrency, + "wall_s": time.perf_counter() - start, + "rows": rows, + } + ) + # Explicit cancellation and socket disconnect after the first streamed token. + for mode in ("cancel", "disconnect"): + before_cancel = complete(128) + rid = f"http211-{mode}" + with request("/v1/completions", body(128, 256, True), rid) as response: + for raw in response: + if raw.startswith(b"data: ") and b"[DONE]" not in raw: + data = json.loads(raw[6:]) + if data.get("choices") and data["choices"][0].get("text"): + break + else: + raise AssertionError("stream had no text") + if mode == "cancel": + with request("/v1/cancel", {}, rid) as cancelled: + assert cancelled.status == 202 + drained(rid) + after = complete(128) + result[mode] = { + "drained": True, + "subsequent_tokens_equal": after["tokens"] + == before_cancel["tokens"], + "cold_baseline_equal": after["tokens"] == baseline[128]["tokens"], + "before_tokens": before_cancel["tokens"], + "after_tokens": after["tokens"], + } + assert result[mode]["subsequent_tokens_equal"], ( + "post-cancellation output changed" + ) + + try: + with request( + "/v1/completions", + {"model": a.model, "prompt": [], "max_tokens": 32}, + ): + raise AssertionError("invalid request accepted") + except urllib.error.HTTPError as exc: + assert exc.code == 400 + result["invalid_status"] = exc.code + with request("/ready") as response: + assert response.status == 200 + result["remaining_requests"] = len(owner.service.batcher.by_id) + assert result["remaining_requests"] == 0 + with request("/metrics") as response: + result["metrics"] = ( + response.read().decode().replace(a.model, result["model_name"]) + ) + hits = [ + float(line.rsplit(" ", 1)[1]) + for line in result["metrics"].splitlines() + if line.startswith("vllm:prefix_cache_hits_total{") + ] + assert hits, "prefix hit metric missing" + result["prefix_hit_tokens"] = sum(hits) + if a.prefix_cache: + assert result["prefix_hit_tokens"] > 0, ( + "cache enabled without observed hits" + ) + if a.host_kv_gib: + host_hits = [ + float(line.rsplit(" ", 1)[1]) + for line in result["metrics"].splitlines() + if line.startswith("airm_kv_host_hit_tokens_total{") + ] + result["host_hit_tokens"] = sum(host_hits) + assert host_hits and result["host_hit_tokens"] > 0, ( + "host tier not exercised" + ) + result["batcher_stats"] = dict(owner.service.batcher.stats) + if a.pd_prefill_device: + # Nineteen completed normal requests; cancellations can add two more prefills. + assert result["batcher_stats"]["prefills"] >= 19 + assert result["batcher_stats"]["transferred_tokens"] > 0 + result["passed"] = True + except Exception as exc: # noqa: BLE001 - preserve failed hardware evidence + result["error"] = repr(exc) + finally: + for rid in list(owner.service.batcher.by_id): + owner.service.batcher.cancel(rid) + owner.submit("stop") + + thread = threading.Thread(target=client, daemon=True) + thread.start() + try: + owner.run() + except Exception as exc: # noqa: BLE001 - preserve failure and stop this process + result["engine_error"] = repr(exc) + finally: + thread.join(190) + server.shutdown() + server.server_close() + server_thread.join(5) + worker = getattr(owner.service.batcher, "worker", None) + if worker is not None: + worker.close() + output.write_text(json.dumps(result, indent=2) + "\n") + if not result["passed"]: + raise RuntimeError( + result.get("error", result.get("engine_error", "validation failed")) + ) + print( + json.dumps( + { + k: v + for k, v in result.items() + if k not in ("cases", "baseline", "metrics") + } + ), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/opt_prime/tests/test_benchmark.py b/opt_prime/tests/test_benchmark.py new file mode 100644 index 0000000..a40564f --- /dev/null +++ b/opt_prime/tests/test_benchmark.py @@ -0,0 +1,52 @@ +"""Benchmark accounting uses all output tokens and whole-batch wall time.""" +import importlib.util +from pathlib import Path + +import pytest + +BENCHMARK = Path(__file__).parents[1] / 'examples/pd/http_matrix/benchmark.py' +spec = importlib.util.spec_from_file_location('http_matrix_benchmark', BENCHMARK) +benchmark = importlib.util.module_from_spec(spec) +spec.loader.exec_module(benchmark) +timing_summary = benchmark.timing_summary + + +def test_batch_accounting_and_parity(): + rows = [ + {"output_tokens": 2, "first_sse_s": .2, "tpot_est_s": .1, "e2e_s": .4, "text_equal": True}, + {"output_tokens": 4, "first_sse_s": .4, "tpot_est_s": .2, "e2e_s": .8, "text_equal": False}, + ] + out = timing_summary(rows, 2) + assert out["requests"] == 2 and out["output_tokens"] == 6 and out["output_tok_s"] == 3 + assert out["median_first_sse_s"] == pytest.approx(.3) + assert out["median_tpot_est_s"] == pytest.approx(.15) + assert out["median_e2e_s"] == pytest.approx(.6) + assert out["all_text_equal"] is False + + +def test_external_server_config_requires_port_and_metadata(tmp_path): + import json + + path = tmp_path / 'server.json' + good = {'servers': [['python', 'atom_server.py', '--port', '{port}', '--model', 'a model']], + 'metadata': {'physical_gpus': 0, 'physical_atom': 1, 'physical_rngd': 1, + 'producer_chunk': None}} + path.write_text(json.dumps(good)) + assert benchmark.load_server_config(path) == good + for bad in [[], {'servers': [['python', '{port}']], 'metadata': {'passed': True}}, {}, {'servers': []}, {'servers': [['python', 'server.py']], 'metadata': {}}, + {'servers': ['python server.py {port}'], 'metadata': {}}, + {'servers': [['python', '{port}']], 'metadata': []}]: + path.write_text(json.dumps(bad)) + with pytest.raises((ValueError, TypeError)): + benchmark.load_server_config(path) + + +def test_startup_timeout_rejects_nonpositive_before_loading_model(tmp_path): + import subprocess + import sys + + run = subprocess.run([sys.executable, str(BENCHMARK), + '--model', 'unused', '--mode', 'native', '--output', str(tmp_path / 'unused.json'), + '--startup-timeout', '0'], check=False, capture_output=True, text=True) + assert run.returncode == 2 + assert 'startup timeout must be positive' in run.stderr diff --git a/opt_prime/tests/test_heldout_quality.py b/opt_prime/tests/test_heldout_quality.py new file mode 100644 index 0000000..0d9c587 --- /dev/null +++ b/opt_prime/tests/test_heldout_quality.py @@ -0,0 +1,134 @@ +"""Prevent answer substrings, duplicate keys, or missing pairs from inflating scores.""" +import importlib.util +from pathlib import Path +import pytest + +path = Path(__file__).parents[1] / 'examples/pd/http_matrix/heldout_quality.py' + +def module(): + assert path.exists(), 'held-out evaluator is not implemented' + spec = importlib.util.spec_from_file_location('heldout_quality', path) + value = importlib.util.module_from_spec(spec) + spec.loader.exec_module(value) + return value + +@pytest.mark.parametrize('text,correct,formatted', [ + ('{"answer":"73"}', True, True), + ('{"answer":"73","note":"ok"}', True, False), + ('{"answer":73}', False, False), + ('The answer is 73', False, False), + ('{"answer":"173"}', False, True), + ('{"answer":"wrong","answer":"73"}', False, False), + ('{"answer":"73"} trailing', False, False), +]) +def test_score(text, correct, formatted): + got = module().score(text, '73') + assert (got['answer_correct'], got['format_correct']) == (correct, formatted) + +def test_paired_regressions_are_not_hidden_by_equal_aggregate(): + m = module() + base = [{'case': 'a', 'input_sha256': 'x', 'answer_correct': True, 'tokens': [1]}, + {'case': 'b', 'input_sha256': 'y', 'answer_correct': False, 'tokens': [2]}] + other = [{'case': 'a', 'input_sha256': 'x', 'answer_correct': False, 'tokens': [3]}, + {'case': 'b', 'input_sha256': 'y', 'answer_correct': True, 'tokens': [2]}] + got = m.compare(base, other) + assert got == {'cases': 2, 'correct_to_wrong': 1, 'wrong_to_correct': 1, 'token_equal': 1} + with pytest.raises(ValueError): + m.compare(base, other[:1]) + other[0]['input_sha256'] = 'changed' + with pytest.raises(ValueError): + m.compare(base, other) + +def test_context_builder_keeps_retrieval_fact_inside_budget(): + class CharacterTokenizer: + def apply_chat_template(self, messages, **kwargs): + return list(messages[0]['content'].encode()) + m = module() + cases = m.fixtures(CharacterTokenizer()) + assert len(cases) == len({x['case'] for x in cases}) == 36 + for case in cases: + assert len(case['ids']) <= (case['target_tokens'] or 2048) + if case['task'] != 'arithmetic': + assert case['expected'] in case['prompt'] + if case['target_tokens']: + assert len(case['ids']) > case['target_tokens'] - 150 + +def test_http_runner_supplies_required_request_id_and_frozen_tokens(tmp_path, monkeypatch): + import json + import sys + import threading + from http.server import BaseHTTPRequestHandler, HTTPServer + seen = [] + class Handler(BaseHTTPRequestHandler): + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers['Content-Length']))) + seen.append((self.headers.get('x-request-id'), body)) + self.send_response(200) + self.end_headers() + self.wfile.write(json.dumps({'choices': [{'text': '{"answer":"73"}', + 'token_ids': [9], 'finish_reason': 'stop'}]}).encode()) + def log_message(self, *args): + pass + server = HTTPServer(('127.0.0.1', 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + fixture = tmp_path / 'input.json' + output = tmp_path / 'output.jsonl' + fixture.write_text(json.dumps({'model':'model', 'max_tokens':128, + 'cases':[{'case':'a','ids':[4,5], 'input_sha256':'hash','expected':'73'}]})) + monkeypatch.setattr(sys, 'argv', ['heldout_quality', 'run', '--fixtures',str(fixture), + '--endpoint',f'http://127.0.0.1:{server.server_port}', '--route','test', '--output',str(output)]) + try: + module().main() + finally: + server.shutdown() + thread.join() + server.server_close() + assert seen[0][0], 'stage HTTP API requires a nonempty request ID' + assert seen[0][1]['prompt'] == [4,5] + assert seen[0][1]['ignore_eos'] is False + assert json.loads(output.read_text())['answer_correct'] is True + with pytest.raises(FileExistsError): + module().main() + +def test_frozen_fixtures_match_manifest(): + import gzip + import hashlib + import json + folder = path.parent / 'fixtures' + manifest = json.loads((folder / 'fixture-manifest.json').read_text()) + assert len(manifest) == 4 + for name, entry in manifest.items(): + raw = gzip.decompress((folder / (name + '.gz')).read_bytes()) + assert hashlib.sha256(raw).hexdigest() == entry['sha256'] + cases = json.loads(raw)['cases'] + assert len(cases) == len({x['case'] for x in cases}) == entry['cases'] + for case in cases: + ids = json.dumps(case['ids'], separators=(',', ':')).encode() + assert hashlib.sha256(ids).hexdigest() == case['input_sha256'] + +def test_runner_reads_gzip_fixture_and_overrides_model(tmp_path, monkeypatch): + import gzip + import json + import sys + m = module() + seen = [] + class Response: + def __enter__(self): + return self + def __exit__(self, *args): + return False + def read(self): + return json.dumps({'choices': [{'text': '{"answer":"73"}', 'token_ids': [9], + 'finish_reason': 'stop'}]}) + def urlopen(request, timeout): + seen.append(json.loads(request.data)) + return Response() + monkeypatch.setattr(m.urllib.request, 'urlopen', urlopen) + fixture = tmp_path / 'input.json.gz' + fixture.write_bytes(gzip.compress(json.dumps({'model': 'frozen', 'max_tokens': 128, + 'cases': [{'case': 'a', 'ids': [4, 5], 'input_sha256': 'hash', 'expected': '73'}]}).encode())) + monkeypatch.setattr(sys, 'argv', ['heldout_quality', 'run', '--fixtures', str(fixture), '--endpoint', 'http://x', + '--route', 'test', '--output', str(tmp_path / 'out.jsonl'), '--model', 'served-name']) + m.main() + assert seen[0]['model'] == 'served-name' and seen[0]['prompt'] == [4, 5]