verae-mcp-timestamping-server/scripts/perf_stamp_load.py

#!/usr/bin/env python3
"""Latency/load stamps against api.veraetime.net with host CPU/RSS samples."""

from __future__ import annotations

import hashlib
import json
import os
import statistics
import subprocess
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
API = os.environ.get("VERAE_GLOBAL_URL", "https://api.veraetime.net")
SIZES = [20, 100, 1000, 10000, 100000]
WORKERS = int(os.environ.get("VERAE_PERF_WORKERS", "32"))


def utc() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ")


def load_token() -> str:
    path = Path.home() / ".verae" / "auth.properties"
    props: dict[str, str] = {}
    for line in path.read_text().splitlines():
        if "=" in line and not line.strip().startswith("#"):
            k, v = line.split("=", 1)
            props[k.strip()] = v.strip()
    return os.environ.get("VERAE_GLOBAL_TOKEN") or props.get("access_token") or props.get("auth_token") or ""


def pct(xs: list[float], p: float) -> float | None:
    if not xs:
        return None
    s = sorted(xs)
    i = min(len(s) - 1, max(0, int(round((p / 100.0) * (len(s) - 1)))))
    return s[i]


def stamp_one(token: str, digest: str, timeout: float = 60.0) -> dict:
    body = json.dumps(
        {
            "data": digest,
            "digest": digest,
            "hashAlg": "SHA256",
            "restamp": False,
            "wait": True,
        }
    ).encode()
    req = urllib.request.Request(
        API + "/api/timestamp",
        data=body,
        method="POST",
        headers={
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}",
        },
    )
    t0 = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read()
            ms = (time.perf_counter() - t0) * 1000
            parsed = json.loads(raw) if raw else {}
            return {
                "ok": True,
                "status": resp.status,
                "ms": ms,
                "jobId": parsed.get("jobId"),
                "certificate": parsed.get("certificate"),
            }
    except urllib.error.HTTPError as e:
        ms = (time.perf_counter() - t0) * 1000
        return {"ok": False, "status": e.code, "ms": ms, "error": e.read()[:200].decode("utf-8", "replace")}
    except Exception as e:
        ms = (time.perf_counter() - t0) * 1000
        return {"ok": False, "status": 0, "ms": ms, "error": f"{type(e).__name__}: {e}"}


def sample_host(stop: threading.Event, samples: list[dict]) -> None:
    while not stop.is_set():
        try:
            out = subprocess.check_output(
                [
                    "ssh",
                    "-o",
                    "BatchMode=yes",
                    "ns1",
                    "pid=$(pgrep -n -f 'timestamping-service/main -config' || true); "
                    "if [ -n \"$pid\" ]; then ps -p \"$pid\" -o %cpu=,rss=; "
                    "else echo NA NA; fi; "
                    "awk '/^cpu /{u=$2+$4; t=$2+$3+$4+$5; print u,t}' /proc/stat; "
                    "awk '/MemAvailable:/{print $2}' /proc/meminfo",
                ],
                text=True,
                timeout=8,
            )
            lines = [ln.strip() for ln in out.splitlines() if ln.strip()]
            cpu_p, rss = (None, None)
            host_u = host_t = None
            mem_avail = None
            if len(lines) >= 1 and lines[0] != "NA NA":
                parts = lines[0].split()
                if len(parts) >= 2:
                    cpu_p, rss = float(parts[0]), int(parts[1])
            if len(lines) >= 2:
                u, t = lines[1].split()[:2]
                host_u, host_t = int(u), int(t)
            if len(lines) >= 3:
                mem_avail = int(lines[2])
            samples.append(
                {
                    "at": utc(),
                    "proc_cpu_pct": cpu_p,
                    "proc_rss_kb": rss,
                    "host_cpu_jiffies_used": host_u,
                    "host_cpu_jiffies_total": host_t,
                    "mem_available_kb": mem_avail,
                }
            )
        except Exception as e:
            samples.append({"at": utc(), "error": str(e)})
        stop.wait(1.0)


def run_size(token: str, n: int, workers: int, run_id: str) -> dict:
    stop = threading.Event()
    samples: list[dict] = []
    mon = threading.Thread(target=sample_host, args=(stop, samples), daemon=True)
    mon.start()
    t0 = time.perf_counter()
    results = []
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futs = []
        for i in range(n):
            digest = hashlib.sha256(f"Verae-Perf-{run_id}-{i}".encode()).hexdigest()
            futs.append(pool.submit(stamp_one, token, digest))
        for fut in as_completed(futs):
            results.append(fut.result())
    wall_s = time.perf_counter() - t0
    stop.set()
    mon.join(timeout=3)
    lat = [r["ms"] for r in results]
    oks = [r for r in results if r.get("ok")]
    return {
        "n": n,
        "workers": workers,
        "wall_s": wall_s,
        "throughput_rps": n / wall_s if wall_s else None,
        "ok": len(oks),
        "fail": n - len(oks),
        "latency_ms": {
            "min": min(lat) if lat else None,
            "avg": statistics.fmean(lat) if lat else None,
            "p50": pct(lat, 50),
            "p75": pct(lat, 75),
            "p80": pct(lat, 80),
            "p90": pct(lat, 90),
            "p95": pct(lat, 95),
            "max": max(lat) if lat else None,
        },
        "host_samples": samples,
        "errors": [r for r in results if not r.get("ok")][:20],
    }


def host_summary(samples: list[dict]) -> dict:
    cpus = [s["proc_cpu_pct"] for s in samples if isinstance(s.get("proc_cpu_pct"), (int, float))]
    rss = [s["proc_rss_kb"] for s in samples if isinstance(s.get("proc_rss_kb"), int)]
    avail = [s["mem_available_kb"] for s in samples if isinstance(s.get("mem_available_kb"), int)]
    return {
        "samples": len(samples),
        "proc_cpu_pct": {
            "min": min(cpus) if cpus else None,
            "avg": statistics.fmean(cpus) if cpus else None,
            "p50": pct(cpus, 50),
            "p95": pct(cpus, 95),
            "max": max(cpus) if cpus else None,
        },
        "proc_rss_mb": {
            "min": (min(rss) / 1024) if rss else None,
            "avg": (statistics.fmean(rss) / 1024) if rss else None,
            "max": (max(rss) / 1024) if rss else None,
        },
        "host_mem_available_gb": {
            "min": (min(avail) / 1e6) if avail else None,
            "avg": (statistics.fmean(avail) / 1e6) if avail else None,
        },
    }


def html_report(data: dict) -> str:
    rows = []
    for b in data["batches"]:
        L = b["latency_ms"]
        H = b["host"]
        rows.append(
            "<tr>"
            + "".join(
                f"<td>{html_esc(x)}</td>"
                for x in [
                    b["n"],
                    b["workers"],
                    f"{b['wall_s']:.2f}",
                    f"{b['throughput_rps']:.1f}" if b["throughput_rps"] else "",
                    b["ok"],
                    b["fail"],
                    f"{L['min']:.1f}" if L["min"] is not None else "",
                    f"{L['avg']:.1f}" if L["avg"] is not None else "",
                    f"{L['p50']:.1f}" if L.get("p50") is not None else "",
                    f"{L['p75']:.1f}" if L.get("p75") is not None else "",
                    f"{L['p80']:.1f}" if L.get("p80") is not None else "",
                    f"{L['p90']:.1f}" if L.get("p90") is not None else "",
                    f"{L['p95']:.1f}" if L.get("p95") is not None else "",
                    f"{L['max']:.1f}" if L["max"] is not None else "",
                    f"{(H['proc_cpu_pct']['avg'] or 0):.1f}",
                    f"{(H['proc_rss_mb']['avg'] or 0):.1f}",
                ]
            )
            + "</tr>"
        )
    return f"""<!DOCTYPE html>
<html lang="en"><head><meta charset="utf-8"><title>stamp load {data['run_id']}</title>
<style>
body{{font:15px/1.4 ui-sans-serif,system-ui;background:#0f1419;color:#e7ecf1;margin:1.5rem}}
table{{border-collapse:collapse;width:100%;font-size:.85rem}}
th,td{{border-bottom:1px solid #2a3542;padding:.35rem .4rem;text-align:right}}
th:first-child,td:first-child{{text-align:left}}
code{{font-family:ui-monospace,Menlo,monospace}}
</style></head><body>
<h1>Stamp load test</h1>
<p>API <code>{html_esc(data['api'])}</code> · wait:true · unique hashes · workers {data['workers']}</p>
<p>Run <code>{html_esc(data['run_id'])}</code> · {html_esc(data['started'])} → {html_esc(data['ended'])}</p>
<p><strong>JetStream:</strong> not used (standing ban on job-complete messages).</p>
<table><thead><tr>
<th>n</th><th>workers</th><th>wall s</th><th>rps</th><th>ok</th><th>fail</th>
<th>min</th><th>avg</th><th>p50</th><th>p75</th><th>p80</th><th>p90</th><th>p95</th><th>max</th>
<th>CPU% avg</th><th>RSS MB avg</th>
</tr></thead><tbody>{''.join(rows)}</tbody></table>
<p>Latency columns are milliseconds.</p>
</body></html>
"""


def html_esc(s) -> str:
    return (
        str(s)
        .replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
    )


def main() -> int:
    token = load_token()
    if not token:
        raise SystemExit("no VERAE_GLOBAL_TOKEN / ~/.verae/auth.properties")
    run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    dest = ROOT / "debug-trace" / f"{run_id}-perf-load"
    dest.mkdir(parents=True, exist_ok=True)
    started = utc()
    batches = []
    sizes = [int(x) for x in os.environ.get("VERAE_PERF_SIZES", ",".join(map(str, SIZES))).split(",") if x]
    for n in sizes:
        workers = min(WORKERS, n)
        print(f"=== n={n} workers={workers} ===", flush=True)
        batch = run_size(token, n, workers, run_id)
        batch["host"] = host_summary(batch.pop("host_samples"))
        print(json.dumps({k: batch[k] for k in ("n", "ok", "fail", "wall_s", "throughput_rps", "latency_ms")}, indent=2), flush=True)
        batches.append(batch)
    ended = utc()
    data = {
        "run_id": run_id,
        "api": API,
        "workers": WORKERS,
        "started": started,
        "ended": ended,
        "jetstream": False,
        "batches": batches,
    }
    (dest / "perf-load.json").write_text(json.dumps(data, indent=2) + "\n")
    (dest / "perf-load.html").write_text(html_report(data))
    print("wrote", dest / "perf-load.html")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())