File size: 2,645 Bytes
96d64c6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | import argparse
import json
import subprocess
import sys
import time
from pathlib import Path
import psutil
TEST_PREFIX = "test_benchmark_"
def _total_rss(procs):
total = 0
for p in procs:
try:
total += p.memory_info().rss
except (psutil.NoSuchProcess, psutil.AccessDenied):
pass
return total
def _percentile(sorted_data, pct):
if not sorted_data:
return 0.0
idx = min(int(len(sorted_data) * pct), len(sorted_data) - 1)
return sorted_data[idx]
def _short_name(full_name: str) -> str:
return full_name[len(TEST_PREFIX):] if full_name.startswith(TEST_PREFIX) else full_name
def run_once(raw_json_path: Path) -> dict:
cmd = [
sys.executable,
"-m",
"pytest",
"tests/",
"--benchmark-only",
f"--benchmark-json={raw_json_path}",
]
psutil.cpu_percent(interval=None)
proc = subprocess.Popen(cmd)
parent = psutil.Process(proc.pid)
peak_rss = 0
cpu_samples = []
while proc.poll() is None:
try:
procs = [parent] + parent.children(recursive=True)
except psutil.NoSuchProcess:
procs = []
peak_rss = max(peak_rss, _total_rss(procs))
cpu_samples.append(psutil.cpu_percent(interval=None))
time.sleep(0.05)
proc.wait()
with raw_json_path.open() as f:
benchmarks = json.load(f)["benchmarks"]
row = {
"memory_peak": peak_rss / (1024 * 1024),
"cpu_usage": sum(cpu_samples) / len(cpu_samples) if cpu_samples else 0.0,
}
for b in sorted(benchmarks, key=lambda x: x["name"]):
name = _short_name(b["name"])
stats = b["stats"]
data_sorted = sorted(stats["data"])
row[f"{name}_throughput"] = stats["ops"]
row[f"{name}_latency_p50"] = stats["median"] * 1_000_000
row[f"{name}_latency_p99"] = _percentile(data_sorted, 0.99) * 1_000_000
return row
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--runs", type=int, default=1)
parser.add_argument("--output", type=Path, default=Path("artemis_results.csv"))
parser.add_argument("--raw-json", type=Path, default=Path("artemis_raw.json"))
args = parser.parse_args()
rows = [run_once(args.raw_json) for _ in range(args.runs)]
columns = list(rows[0].keys())
with args.output.open("w") as f:
f.write(",".join(columns) + "\n")
for r in rows:
f.write(",".join(f"{r[c]}" for c in columns) + "\n")
try:
args.raw_json.unlink()
except FileNotFoundError:
pass
if __name__ == "__main__":
main()
|