zc2 0.0.26

P2P compute broker with credit-based billing, WAL, and broker mesh support
#!/usr/bin/env python3
"""
Throughput benchmark for Zakuro broker /execute endpoint.

Sends POST /execute with JSON {"n": 4} (fibonacci(4)); expects worker response
with result, and optionally pid + ip in dictionary format, e.g.:
  {"result": 3, "n": 4, "pid": 12345, "ip": "192.168.1.1"}

Use scripts/fibonacci_worker.py as the worker to get pid and ip in responses.

  - Phase 1 (optional): send N requests to a GATEWAY broker (no local workers);
    traffic is relayed via QUIC to the executor → measures QUIC relay throughput.
  - Phase 2: send N requests to the EXECUTOR broker (local execution) → baseline.

Requirements:
  - ZAKURO_API_URL and ZAKURO_API_KEY (to resolve user and form Bearer token),
    OR pass --api-url and --api-key, OR --bearer-token directly.
  - Broker(s) and worker(s) already running (this script only sends HTTP requests).

Usage:
  # Via API (WireGuard): point BROKER_URL at your broker's WireGuard hostname
  export ZAKURO_API_URL=https://my.zakuro-ai.com
  export ZAKURO_API_KEY=zk_<user_id>_<secret>
  export BROKER_URL=http://mybroker:9000   # or http://100.x.x.x:9000
  python scripts/benchmark_quic_throughput.py
  # Output includes credit distribution and per-worker statistics (cost, % requests).

  # Two-phase (QUIC relay vs local)
  export GATEWAY_BROKER_URL=http://gateway:9000
  export EXECUTOR_BROKER_URL=http://executor:9000
  python scripts/benchmark_quic_throughput.py --two-phase
"""

from __future__ import annotations

import argparse
import os
import sys
import threading
import time
from collections import Counter
from typing import Any, Dict, List, Tuple

try:
    import requests
except ImportError:
    print("Requires: pip install requests", file=sys.stderr)
    sys.exit(1)

# fibonacci(4) = 3
FIB4_EXPECTED = 3


def get_bearer_token(api_url: str, api_key: str) -> Tuple[str, str, float]:
    """Resolve zakuro_user_id, build Bearer token, and return current credits balance from dashboard."""
    url = f"{api_url.rstrip('/')}/api/auth/me/api-key"
    r = requests.get(url, headers={"Authorization": f"Bearer {api_key}"}, timeout=10)
    r.raise_for_status()
    data = r.json()
    uid = data.get("zakuro_user_id") or data.get("user_id")
    if not uid:
        raise SystemExit("Dashboard did not return zakuro_user_id")
    balance = float(data.get("credits_balance", 0) or 0)
    return f"zk_{uid}_0001", uid, balance


def run_phase(
    base_url: str,
    bearer: str,
    n: int,
    threads: int,
    timeout: float,
) -> Tuple[int, int, List[float], List[str], List[Dict[str, Any]], List[Dict[str, float]]]:
    """Send n requests to base_url/execute (fibonacci(4)); return (ok, total, latencies_ms, transports, worker_responses, credits_data)."""
    latencies: List[float] = [0.0] * n
    transports: List[str] = ["?"] * n
    worker_responses: List[Dict[str, Any]] = [{} for _ in range(n)]
    credits_data: List[Dict[str, float]] = [{"cost": 0.0, "credits_remaining": 0.0} for _ in range(n)]
    ok_count = 0
    lock = threading.Lock()
    idx = [0]

    def worker():
        nonlocal ok_count
        session = requests.Session()
        session.headers.update({
            "Authorization": f"Bearer {bearer}",
            "Content-Type": "application/json",
            "X-Zakuro-Requirements": '{"strategy":"round_robin","estimated_duration_secs":0.01}',
        })
        url = f"{base_url.rstrip('/')}/execute"
        while True:
            with lock:
                i = idx[0]
                if i >= n:
                    break
                idx[0] = i + 1
            t0 = time.perf_counter()
            body: Dict[str, Any] = {}
            cost, credits_remaining = 0.0, 0.0
            try:
                resp = session.post(url, json={"n": 4}, timeout=timeout)
                ms = (time.perf_counter() - t0) * 1000
                transport = resp.headers.get("X-Zakuro-Transport", "?")
                try:
                    cost = float(resp.headers.get("X-Zakuro-Cost", 0) or 0)
                except (TypeError, ValueError):
                    pass
                try:
                    credits_remaining = float(resp.headers.get("X-Zakuro-Credits-Remaining", 0) or 0)
                except (TypeError, ValueError):
                    pass
                success = resp.status_code == 200
                if success:
                    body = resp.json() if resp.content else {}
                fib_ok = body.get("result") == FIB4_EXPECTED
            except Exception:
                ms = (time.perf_counter() - t0) * 1000
                transport = "?"
                success = False
                fib_ok = False
            latencies[i] = ms
            transports[i] = transport
            worker_responses[i] = {k: body.get(k) for k in ("result", "n", "pid", "ip") if body.get(k) is not None}
            credits_data[i] = {"cost": cost, "credits_remaining": credits_remaining}
            if success and fib_ok:
                with lock:
                    ok_count += 1

    ths = [threading.Thread(target=worker) for _ in range(threads)]
    for t in ths:
        t.start()
    for t in ths:
        t.join()
    return ok_count, n, latencies, transports, worker_responses, credits_data


def format_worker_summary(worker_responses: List[Dict[str, Any]]) -> str:
    """Summarise unique worker (pid, ip) from response dicts."""
    seen = set()
    for w in worker_responses:
        pid, ip = w.get("pid"), w.get("ip")
        if pid is not None and ip is not None:
            seen.add((pid, ip))
    if not seen:
        return ""
    parts = [f"pid={p} ip={i}" for p, i in sorted(seen)]
    return "  Worker(s): " + "  ".join(parts)


def task_distribution(
    worker_responses: List[Dict[str, Any]], total: int
) -> List[Tuple[str, str, int, float]]:
    """Return per-worker rows (label_pid, label_ip, count, pct). Sorted by count descending."""
    counts: Dict[Tuple[Any, Any], int] = {}
    no_id = 0
    for w in worker_responses:
        pid, ip = w.get("pid"), w.get("ip")
        if pid is not None and ip is not None:
            key = (pid, ip)
            counts[key] = counts.get(key, 0) + 1
        else:
            no_id += 1
    if not total:
        return []
    rows = [
        (str(pid), str(ip), c, 100.0 * c / total)
        for (pid, ip), c in sorted(counts.items(), key=lambda x: -x[1])
    ]
    if no_id:
        rows.append(("(no pid)", "(no ip)", no_id, 100.0 * no_id / total))
    return rows


def credit_distribution_and_per_worker(
    worker_responses: List[Dict[str, Any]],
    credits_data: List[Dict[str, float]],
    total: int,
) -> Tuple[float, float, float, List[Tuple[str, str, int, float, float, float]]]:
    """Return (total_cost, balance_first, balance_last), and per-worker (pid, ip, count, pct_req, cost, pct_cost)."""
    total_cost = sum(c["cost"] for c in credits_data)
    balance_first = 0.0
    balance_last = 0.0
    for c in credits_data:
        if c["credits_remaining"] > 0 or c["cost"] > 0:
            balance_first = c["credits_remaining"] + c["cost"]
            break
    for i in range(len(credits_data) - 1, -1, -1):
        c = credits_data[i]
        if c["credits_remaining"] > 0 or c["cost"] > 0:
            balance_last = c["credits_remaining"]
            break
    # Per-worker: (pid, ip) -> (count, cost)
    worker_cost: Dict[Tuple[Any, Any], Tuple[int, float]] = {}
    no_id_count, no_id_cost = 0, 0.0
    for w, c in zip(worker_responses, credits_data):
        pid, ip = w.get("pid"), w.get("ip")
        cost = c["cost"]
        if pid is not None and ip is not None:
            key = (pid, ip)
            cnt, tot = worker_cost.get(key, (0, 0.0))
            worker_cost[key] = (cnt + 1, tot + cost)
        else:
            no_id_count += 1
            no_id_cost += cost
    rows = []
    for (pid, ip), (cnt, cost) in sorted(worker_cost.items(), key=lambda x: -x[1][1]):
        pct_req = 100.0 * cnt / total if total else 0
        pct_cost = 100.0 * cost / total_cost if total_cost else 0
        rows.append((str(pid), str(ip), cnt, pct_req, cost, pct_cost))
    if no_id_count:
        pct_req = 100.0 * no_id_count / total if total else 0
        pct_cost = 100.0 * no_id_cost / total_cost if total_cost else 0
        rows.append(("(no pid)", "(no ip)", no_id_count, pct_req, no_id_cost, pct_cost))
    return total_cost, balance_first, balance_last, rows


def stats(latencies: List[float]) -> dict:
    s = sorted(latencies)
    nn = len(s)
    if nn == 0:
        return {"min": 0, "avg": 0, "p50": 0, "p90": 0, "p99": 0, "max": 0}
    return {
        "min": s[0],
        "avg": sum(s) / nn,
        "p50": s[nn * 50 // 100],
        "p90": s[nn * 90 // 100],
        "p99": s[nn * 99 // 100],
        "max": s[-1],
    }


def main() -> None:
    ap = argparse.ArgumentParser(description="Zakuro broker /execute throughput benchmark")
    ap.add_argument("--broker-url", default=os.environ.get("BROKER_URL"), help="Single broker URL")
    ap.add_argument("--gateway-url", default=os.environ.get("GATEWAY_BROKER_URL"), help="Gateway broker (no workers)")
    ap.add_argument("--executor-url", default=os.environ.get("EXECUTOR_BROKER_URL"), help="Executor broker (has worker)")
    ap.add_argument("--api-url", default=os.environ.get("ZAKURO_API_URL"), help="Dashboard API URL")
    ap.add_argument("--api-key", default=os.environ.get("ZAKURO_API_KEY"), help="Dashboard API key")
    ap.add_argument("--bearer-token", default=os.environ.get("ZAKURO_BEARER_TOKEN"), help="Bearer token (skip dashboard)")
    ap.add_argument("--requests", "-n", type=int, default=1000, help="Requests per phase (default 1000 for distribution stats)")
    ap.add_argument("--threads", "-j", type=int, default=10, help="Concurrent threads")
    ap.add_argument("--timeout", type=float, default=30.0, help="Request timeout (s)")
    ap.add_argument("--two-phase", action="store_true", help="Run gateway (QUIC) + executor (local) phases")
    args = ap.parse_args()

    if args.bearer_token:
        bearer = args.bearer_token
        uid = "?"
        balance_start_api = None
    elif args.api_url and args.api_key:
        bearer, uid, balance_start_api = get_bearer_token(args.api_url, args.api_key)
        print("Resolved user: zakuro_user_id={}  balance (API)={:.6f}".format(uid, balance_start_api))
    else:
        print("Set ZAKURO_API_URL + ZAKURO_API_KEY, or --bearer-token, or --api-url/--api-key", file=sys.stderr)
        sys.exit(1)

    n, threads = args.requests, args.threads
    timeout = args.timeout

    if args.two_phase:
        if not args.gateway_url or not args.executor_url:
            print("--two-phase requires --gateway-url and --executor-url (or env GATEWAY_BROKER_URL, EXECUTOR_BROKER_URL)", file=sys.stderr)
            sys.exit(1)
        # Phase 1: gateway (QUIC relay)
        t0 = time.perf_counter()
        ok1, _, lat1, trans1, workers1, credits1 = run_phase(args.gateway_url, bearer, n, threads, timeout)
        wall1 = time.perf_counter() - t0
        quic_count = sum(1 for t in trans1 if t == "quic")
        # Phase 2: executor (local)
        t0 = time.perf_counter()
        ok2, _, lat2, trans2, workers2, credits2 = run_phase(args.executor_url, bearer, n, threads, timeout)
        wall2 = time.perf_counter() - t0
        local_count = sum(1 for t in trans2 if t == "local")

        print("\n" + "=" * 72)
        print("  QUIC Throughput Benchmark (Python) — API/WireGuard, fibonacci(4), {} requests × 2 phases".format(n))
        print("=" * 72)
        if balance_start_api is not None:
            print("  Balance (dashboard before run): {:.6f} credits".format(balance_start_api))
        print("\n  Phase 1: QUIC Relay (requests → Gateway → QUIC → Executor)")
        print("  " + "-" * 68)
        print("    URL:        {}".format(args.gateway_url))
        print("    Transport:  {} QUIC / {} other".format(quic_count, n - quic_count))
        print("    Success:    {}/{} ({:.1f}%)  fibonacci(4)={}".format(ok1, n, 100.0 * ok1 / n if n else 0, FIB4_EXPECTED))
        print("    Wall:       {:.2f}s   Throughput: {:.1f} req/s".format(wall1, n / wall1 if wall1 else 0))
        s1 = stats(lat1)
        print("    Latency:    min={:.2f}  avg={:.2f}  p50={:.2f}  p90={:.2f}  p99={:.2f}  max={:.2f} ms".format(
            s1["min"], s1["avg"], s1["p50"], s1["p90"], s1["p99"], s1["max"]))
        tot1, bal1_first, bal1_last, per_worker1 = credit_distribution_and_per_worker(workers1, credits1, n)
        print("    Credit distribution:  total cost={:.6f}  balance after={:.6f}".format(tot1, bal1_last))
        print("    Per-worker statistics:")
        for pid_s, ip_s, cnt, pct_r, cost, pct_c in per_worker1:
            print("      pid={:<10}  ip={:<15}  {:>4} req  {:5.1f}%  cost={:.6f}  {:5.1f}%".format(pid_s, ip_s, cnt, pct_r, cost, pct_c))
        print("\n  Phase 2: Local Baseline (requests → Executor directly)")
        print("  " + "-" * 68)
        print("    URL:        {}".format(args.executor_url))
        print("    Transport:  {} local / {} other".format(local_count, n - local_count))
        print("    Success:    {}/{} ({:.1f}%)  fibonacci(4)={}".format(ok2, n, 100.0 * ok2 / n if n else 0, FIB4_EXPECTED))
        print("    Wall:       {:.2f}s   Throughput: {:.1f} req/s".format(wall2, n / wall2 if wall2 else 0))
        s2 = stats(lat2)
        print("    Latency:    min={:.2f}  avg={:.2f}  p50={:.2f}  p90={:.2f}  p99={:.2f}  max={:.2f} ms".format(
            s2["min"], s2["avg"], s2["p50"], s2["p90"], s2["p99"], s2["max"]))
        tot2, bal2_first, bal2_last, per_worker2 = credit_distribution_and_per_worker(workers2, credits2, n)
        print("    Credit distribution:  total cost={:.6f}  balance after={:.6f}".format(tot2, bal2_last))
        print("    Per-worker statistics:")
        for pid_s, ip_s, cnt, pct_r, cost, pct_c in per_worker2:
            print("      pid={:<10}  ip={:<15}  {:>4} req  {:5.1f}%  cost={:.6f}  {:5.1f}%".format(pid_s, ip_s, cnt, pct_r, cost, pct_c))
        if s2["avg"] > 0:
            overhead = (s1["avg"] - s2["avg"]) / s2["avg"] * 100
            print("\n  Comparison:  QUIC relay overhead {:.2f}ms avg ({:+.1f}% vs local)".format(
                s1["avg"] - s2["avg"], overhead))
        print("\n" + "=" * 72 + "\n")
        sys.exit(0 if (ok1 == n and ok2 == n) else 1)

    # Single broker (e.g. WireGuard hostname)
    url = args.broker_url
    if not url:
        print("Set BROKER_URL or --broker-url for single-broker mode", file=sys.stderr)
        sys.exit(1)
    t0 = time.perf_counter()
    ok, _, latencies, transports, worker_responses, credits_data = run_phase(url, bearer, n, threads, timeout)
    wall = time.perf_counter() - t0
    s = stats(latencies)
    print("\n  Benchmark:   API/WireGuard, fibonacci(4), {} requests".format(n))
    if balance_start_api is not None:
        print("  Balance (dashboard before): {:.6f} credits".format(balance_start_api))
    print("  Broker:      {}".format(url))
    print("  Success:     {}/{}  Throughput: {:.1f} req/s".format(ok, n, n / wall if wall else 0))
    print("  Latency (ms): min={:.2f}  avg={:.2f}  p50={:.2f}  p90={:.2f}  max={:.2f}".format(
        s["min"], s["avg"], s["p50"], s["p90"], s["max"]))
    if any(t != "?" for t in transports):
        for t, c in Counter(transports).most_common():
            print("  Transport {}: {} requests".format(t, c))
    tot_cost, bal_first, bal_last, per_worker = credit_distribution_and_per_worker(worker_responses, credits_data, n)
    print("  Credit distribution:  total cost={:.6f}  balance after={:.6f}".format(tot_cost, bal_last))
    if balance_start_api is not None and tot_cost > 0:
        print("  Credits deducted (this run): {:.6f}".format(tot_cost))
    print("  Per-worker statistics:")
    for pid_s, ip_s, cnt, pct_r, cost, pct_c in per_worker:
        print("    pid={:<10}  ip={:<15}  {:>4} req  {:5.1f}%  cost={:.6f}  {:5.1f}%".format(pid_s, ip_s, cnt, pct_r, cost, pct_c))
    print()
    sys.exit(0 if ok == n else 1)


if __name__ == "__main__":
    main()