| """Measure cold load and end-to-end warm latency for an exported bundle.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import platform |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnxruntime as ort |
| from transformers import AutoTokenizer |
|
|
| from vons.data import read_jsonl |
|
|
|
|
| def percentile(values: list[float], fraction: float) -> float: |
| ordered = sorted(values) |
| index = (len(ordered) - 1) * fraction |
| lower, upper = int(index), min(int(index) + 1, len(ordered) - 1) |
| return ordered[lower] + (ordered[upper] - ordered[lower]) * (index - lower) |
|
|
|
|
| def benchmark(bundle_manifest: Path, input_path: Path, limit: int) -> dict[str, object]: |
| manifest = json.loads(bundle_manifest.read_text(encoding="utf-8")) |
| tokenizer = AutoTokenizer.from_pretrained(manifest["tokenizer"]) |
| rows = read_jsonl(input_path)[:limit] |
| load_started = time.perf_counter() |
| session = ort.InferenceSession(manifest["onnx"], providers=["CPUExecutionProvider"]) |
| cold_load_ms = (time.perf_counter() - load_started) * 1000 |
| backend = manifest["backend"] |
| option_count = int(manifest["option_count"]) |
| sequence_length = int(manifest["sequence_length"]) |
| rng = np.random.default_rng(7) |
| latencies: list[float] = [] |
| for row in rows: |
| texts = [f"{row.state}\nQuestion: {row.question}\nCandidate: {option}" for option in row.options] |
| started = time.perf_counter() |
| encoded = tokenizer(texts, padding="max_length", truncation=True, max_length=sequence_length, return_tensors="np") |
| if backend == "diffusion": |
| pad = option_count - len(row.options) |
| if pad < 0: |
| raise ValueError("row has more candidates than the fixed diffusion bundle") |
| feed = { |
| "input_ids": np.pad(encoded["input_ids"][None, :, :], ((0, 0), (0, pad), (0, 0))), |
| "attention_mask": np.pad(encoded["attention_mask"][None, :, :], ((0, 0), (0, pad), (0, 0))), |
| "token_type_ids": np.pad(encoded["token_type_ids"][None, :, :], ((0, 0), (0, pad), (0, 0))), |
| "option_mask": np.asarray([[True] * len(row.options) + [False] * pad], dtype=np.bool_), |
| "initial_noise": rng.standard_normal((1, option_count), dtype=np.float32), |
| } |
| else: |
| feed = { |
| "input_ids": encoded["input_ids"][None, :, :], |
| "attention_mask": encoded["attention_mask"][None, :, :], |
| "token_type_ids": encoded["token_type_ids"][None, :, :], |
| "option_mask": np.ones((1, len(row.options)), dtype=np.bool_), |
| } |
| session.run(None, feed) |
| latencies.append((time.perf_counter() - started) * 1000) |
| return { |
| "backend": backend, |
| "bundle": manifest["onnx"], |
| "input": str(input_path), |
| "rows": len(latencies), |
| "provider": "CPUExecutionProvider", |
| "onnxruntime_version": ort.__version__, |
| "platform": platform.platform(), |
| "candidate_slots": option_count if backend == "diffusion" else "row-defined", |
| "sequence_length": sequence_length, |
| "padding": "max_length", |
| "inference_steps": manifest.get("inference_steps") if backend == "diffusion" else None, |
| "cold_load_ms": cold_load_ms, |
| "warm_p50_ms": percentile(latencies, 0.5) if latencies else None, |
| "warm_p95_ms": percentile(latencies, 0.95) if latencies else None, |
| "latency_scope": "tokenization plus ONNX Runtime inference", |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--manifest", required=True, type=Path) |
| parser.add_argument("--input", required=True, type=Path) |
| parser.add_argument("--limit", type=int, default=20) |
| parser.add_argument("--output", required=True, type=Path) |
| args = parser.parse_args() |
| result = benchmark(args.manifest, args.input, args.limit) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8") |
| print(json.dumps(result, indent=2, sort_keys=True)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
