from __future__ import annotations
import json
import os
import sys
import time
import statistics
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Optional
import cloudpickle
import httpx
NODE1_HTTP = "http://localhost:9001"
NODE2_HTTP = "http://localhost:9002"
NODE1_KEY = os.environ.get("ZK0NODE01_ZAKURO_KEY") or os.environ.get("ZAKURO_API_KEY")
NODE1_MASTER_KEY = os.environ.get("NODE1_MASTER_KEY", "node1-key")
if not NODE1_KEY:
sys.exit("ERROR: ZK0NODE01_ZAKURO_KEY not set. Source ~/.zakuro/env first.")
_parts = NODE1_KEY[3:] if NODE1_KEY.startswith("zk_") else NODE1_KEY
NODE1_USER_ID = _parts[:_parts.rfind("_")] if "_" in _parts else "unknown"
N_TASKS = 10_000
BUDGET_PER_TASK = 3.5 EXPECTED_MIN_CHARGE = 0.00035 CONCURRENCY = 64 TOLERANCE = 0.10
def workload_instant():
import socket
return {"ok": True, "host": socket.gethostname()}
_PAYLOAD = cloudpickle.dumps({"func": workload_instant, "args": (), "kwargs": {}})
def get_balance(http_url: str, api_key: str) -> float:
r = httpx.get(f"{http_url}/me",
headers={"Authorization": f"Bearer {api_key}"}, timeout=10)
r.raise_for_status()
return float(r.json()["balance"])
def fund(http_url: str, user_id: str, amount: float) -> float:
r = httpx.post(
f"{http_url}/credits/{user_id}/add",
json={"amount": amount, "description": "bench-prefetch seed"},
headers={"X-Api-Key": NODE1_MASTER_KEY},
timeout=10,
)
r.raise_for_status()
return float(r.json()["new_balance"])
@dataclass
class TaskResult:
idx: int
status: int
cost: float
remaining: float
worker: str
elapsed_ms: float
error: Optional[str]
@property
def ok(self) -> bool:
return self.status == 200
def execute_one(idx: int, client: httpx.Client) -> TaskResult:
requirements = {
"cpus": 1.0,
"memory_bytes": 64 * 1024 * 1024,
"gpus": 0,
"estimated_duration_secs": 0.1,
"strategy": "best_price",
"remote_only": True,
"budget_credits": BUDGET_PER_TASK,
}
t0 = time.perf_counter()
try:
r = client.post(
f"{NODE1_HTTP}/execute",
content=_PAYLOAD,
headers={
"Content-Type": "application/octet-stream",
"Authorization": f"Bearer {NODE1_KEY}",
"X-Zakuro-Requirements": json.dumps(requirements),
},
)
except Exception as exc:
elapsed = (time.perf_counter() - t0) * 1000
return TaskResult(idx=idx, status=0, cost=0.0, remaining=0.0,
worker="?", elapsed_ms=elapsed, error=str(exc))
elapsed = (time.perf_counter() - t0) * 1000
cost = float(r.headers.get("X-Zakuro-Cost", "0") or "0")
remaining = float(r.headers.get("X-Zakuro-Credits-Remaining", "0") or "0")
worker = r.headers.get("X-Zakuro-Worker", "?")
error = None
if r.status_code >= 400:
try: error = r.json().get("error", r.text)
except: error = r.text
return TaskResult(idx=idx, status=r.status_code, cost=cost,
remaining=remaining, worker=worker,
elapsed_ms=elapsed, error=error)
def pct(data: list[float], p: float) -> float:
if not data:
return 0.0
s = sorted(data)
i = (len(s) - 1) * p / 100
lo, hi = int(i), min(int(i) + 1, len(s) - 1)
return s[lo] + (s[hi] - s[lo]) * (i - lo)
PASS = "\033[32m✓ PASS\033[0m"
FAIL = "\033[31m✗ FAIL\033[0m"
INFO = "\033[36m·\033[0m"
_checks: list[tuple[str, bool, str]] = []
def check(name: str, cond: bool, detail: str = "") -> bool:
_checks.append((name, cond, detail))
print(f" {PASS if cond else FAIL} {name}" + (f" {detail}" if detail else ""))
return cond
def sep(title: str) -> None:
print(f"\n{'═' * 70}\n {title}\n{'═' * 70}")
def info(msg: str) -> None:
print(f" {INFO} {msg}")
def main() -> None:
sep("SETUP — Verify nodes and fund node01")
h1 = httpx.get(f"{NODE1_HTTP}/health", timeout=5).json()
h2 = httpx.get(f"{NODE2_HTTP}/health", timeout=5).json()
info(f"node1: {h1['node_name']} ts={h1['wireguard_connected']} ts_ip={h1.get('wireguard_ip','?')}")
info(f"node2: {h2['node_name']} ts={h2['wireguard_connected']} ts_ip={h2.get('wireguard_ip','?')}")
check("node1 healthy", h1["status"] == "healthy")
check("node2 healthy", h2["status"] == "healthy")
check("node1 WireGuard connected", h1["wireguard_connected"])
check("node2 WireGuard connected", h2["wireguard_connected"])
workers_resp = httpx.get(f"{NODE1_HTTP}/workers", timeout=10).json()["workers"]
remote_workers = [w for w in workers_resp if not w["uri"].startswith("http://127.")]
info(f"node01 sees {len(remote_workers)} remote (node02) workers:")
for w in remote_workers:
info(f" {w['name']:<24} min_charge={w['min_charge']:.5f} price/hr={w['price_per_hour']:.4f}")
check("node01 sees node02 workers", len(remote_workers) > 0,
f"found {len(remote_workers)}")
check("node02 workers min_charge = 0.00035",
all(abs(w["min_charge"] - 0.00035) < 1e-6 for w in remote_workers),
f"min_charges={[w['min_charge'] for w in remote_workers]}")
expected_total = N_TASKS * EXPECTED_MIN_CHARGE fund_target = expected_total + 1.0 current_bal = get_balance(NODE1_HTTP, NODE1_KEY)
info(f"Current node01 balance: {current_bal:.5f}")
info(f"Target balance for benchmark: {fund_target:.5f}")
if current_bal < fund_target:
top_up = fund_target - current_bal + 0.1
try:
new_bal = fund(NODE1_HTTP, NODE1_USER_ID, top_up)
info(f"Funded {top_up:.5f} → new balance: {new_bal:.5f}")
current_bal = new_bal
except httpx.HTTPStatusError as e:
if e.response.status_code == 501:
info("API mode: /credits/add not available — using existing dashboard balance")
else:
raise
if current_bal > fund_target + 0.5:
info(f"Balance {current_bal:.5f} exceeds target {fund_target:.5f} — draining excess for tight accounting")
info(" (cannot drain standalone ledger — will note initial balance precisely)")
balance_initial = get_balance(NODE1_HTTP, NODE1_KEY)
info(f"Baseline balance: {balance_initial:.6f}")
check("node01 balance ≥ expected total cost",
balance_initial >= expected_total,
f"balance={balance_initial:.5f} needed={expected_total:.5f}")
sep(f"BENCHMARK — {N_TASKS:,} tasks budget_per_task={BUDGET_PER_TASK} concurrency={CONCURRENCY}")
info(f"node01 → node02 workers (remote_only=True)")
info(f"Expected: min_charge={EXPECTED_MIN_CHARGE} total≈{expected_total:.4f} credits")
info(f"Workers: {len(remote_workers)} (each serves requests concurrently)")
print()
results: list[Optional[TaskResult]] = [None] * N_TASKS
done = 0
errors = 0
last_print = [time.perf_counter()]
t_start = time.perf_counter()
transport = httpx.HTTPTransport(limits=httpx.Limits(
max_connections=CONCURRENCY + 16,
max_keepalive_connections=CONCURRENCY,
))
client = httpx.Client(
transport=transport,
timeout=httpx.Timeout(connect=10, read=120, write=30, pool=30),
)
with ThreadPoolExecutor(max_workers=CONCURRENCY) as pool:
futures = {pool.submit(execute_one, i, client): i for i in range(N_TASKS)}
for fut in as_completed(futures):
r = fut.result()
results[futures[fut]] = r
done += 1
if not r.ok:
errors += 1
now = time.perf_counter()
if done % 500 == 0 or (now - last_print[0]) >= 5.0:
elapsed = now - t_start
rate = done / elapsed if elapsed > 0 else 0
running_cost = sum(x.cost for x in results if x and x.ok)
print(f" [{100*done/N_TASKS:5.1f}%] {done:>6}/{N_TASKS} "
f"{rate:>7.1f} tasks/s errors={errors} "
f"cost_so_far={running_cost:.4f}", flush=True)
last_print[0] = now
client.close()
t_total = time.perf_counter() - t_start
print()
time.sleep(1) balance_final = get_balance(NODE1_HTTP, NODE1_KEY)
sep("ANALYSIS")
ok_results = [r for r in results if r and r.ok]
fail_results = [r for r in results if r and not r.ok]
costs = [r.cost for r in ok_results]
latencies = [r.elapsed_ms for r in ok_results]
remainings = [r.remaining for r in ok_results]
sum_costs = sum(costs)
actual_debit = balance_initial - balance_final
throughput = N_TASKS / t_total
info(f"Tasks: {N_TASKS:>8,} succeeded={len(ok_results):,} failed={len(fail_results):,}")
info(f"Elapsed: {t_total:>8.2f} s ({throughput:.1f} tasks/s)")
print()
info(f"Balance initial: {balance_initial:>14.6f}")
info(f"Balance final: {balance_final:>14.6f}")
info(f"Actual debit: {actual_debit:>14.6f}")
info(f"Σ cost headers: {sum_costs:>14.6f}")
info(f"Expected total: {expected_total:>14.6f} (10k × 0.00035)")
info(f"Discrepancy: {abs(actual_debit - sum_costs):>14.2e}")
print()
if costs:
info(f"Cost/task — min={min(costs):.5f} max={max(costs):.5f} "
f"mean={statistics.mean(costs):.5f} stdev={statistics.stdev(costs) if len(costs)>1 else 0:.5f}")
if latencies:
info(f"Latency ms — p50={pct(latencies,50):.1f} p95={pct(latencies,95):.1f} "
f"p99={pct(latencies,99):.1f} max={max(latencies):.1f}")
sep("CHECKS — Balance pre-fetch and billing correctness")
check("all 10k tasks succeeded", len(fail_results) == 0,
f"{len(fail_results)} failed" if fail_results else "")
check("success rate ≥ 99.9%", len(ok_results) >= int(N_TASKS * 0.999),
f"{len(ok_results)}/{N_TASKS}")
check("benchmark completed in ≤ 120 s", t_total <= 120.0,
f"elapsed={t_total:.1f}s")
if costs:
check("every task cost ≥ min_charge (0.00035)",
all(c >= EXPECTED_MIN_CHARGE - 1e-7 for c in costs),
f"min_cost_seen={min(costs):.6f}")
check("every task cost ≤ budget_per_task (3.5)",
all(c <= BUDGET_PER_TASK + 1e-6 for c in costs),
f"max_cost_seen={max(costs):.6f}")
check("no overdraft (balance ≥ 0 after benchmark)", balance_final >= -1e-6,
f"balance_final={balance_final:.6f}")
debit_diff = abs(actual_debit - sum_costs)
check("Σ(cost headers) = actual balance deduction (diff < 1e-4)",
debit_diff < 1e-4,
f"sum={sum_costs:.6f} debit={actual_debit:.6f} diff={debit_diff:.2e}")
lo = expected_total * (1 - TOLERANCE)
hi = expected_total * (1 + TOLERANCE)
check(f"total billing ≈ 3.5 credits (within ±{int(TOLERANCE*100)}%)",
lo <= actual_debit <= hi,
f"debit={actual_debit:.5f} range=[{lo:.5f}, {hi:.5f}]")
stale_zero = sum(1 for r in ok_results if r.remaining == 0.0)
check("no task returned stale remaining=0 (pre-fetch bug absent)",
stale_zero == 0,
f"{stale_zero} tasks had remaining=0" if stale_zero else "")
if remainings:
max_rem = max(remainings)
check("remaining headers ≤ initial balance (no phantom credits)",
max_rem <= balance_initial + 0.01,
f"max_remaining={max_rem:.5f} initial={balance_initial:.5f}")
check("throughput ≥ 50 tasks/s", throughput >= 50.0,
f"{throughput:.1f} tasks/s")
if fail_results:
sep("FAILED TASKS (sample)")
error_counts: dict[str, int] = {}
for r in fail_results:
key = f"status={r.status} error={r.error or '?'}"
error_counts[key] = error_counts.get(key, 0) + 1
for msg, cnt in sorted(error_counts.items(), key=lambda x: -x[1]):
info(f" ×{cnt:>4} {msg}")
sep("RESULTS")
passed = sum(1 for _, ok, _ in _checks if ok)
failed = sum(1 for _, ok, _ in _checks if not ok)
total = len(_checks)
print(f"\n {passed}/{total} checks passed", end="")
if failed:
print(f" ({failed} FAILED)")
for name, ok, detail in _checks:
if not ok:
print(f" {FAIL} {name} {detail}")
else:
print(" — all checks passed")
print()
print(f" Balance: {balance_initial:.5f} → {balance_final:.5f} "
f"Δ={actual_debit:.5f} credits ({len(ok_results):,} tasks @ {throughput:.0f} tasks/s)")
print()
sys.exit(0 if failed == 0 else 1)
if __name__ == "__main__":
main()