from __future__ import annotations
import json
import os
import sys
import time
from dataclasses import dataclass, field
from typing import Optional
import cloudpickle
import httpx
NODE1_KEY = os.environ.get("ZK0NODE01_ZAKURO_KEY") or os.environ.get("ZAKURO_API_KEY")
NODE2_KEY = os.environ.get("ZK0NODE02_ZAKURO_KEY")
NODE1_MASTER_KEY = os.environ.get("NODE1_MASTER_KEY", "node1-master-key")
if not NODE1_KEY:
sys.exit("ERROR: ZK0NODE01_ZAKURO_KEY not set. Source ~/.zakuro/env first.")
if not NODE2_KEY:
sys.exit("ERROR: ZK0NODE02_ZAKURO_KEY not set. Source ~/.zakuro/env first.")
import zakuro
from zakuro.compute import Compute
from zakuro.processors.base import ProcessorConfig
from zakuro.processors.broker import BrokerProcessor
NODE1_URI = "zc://localhost:9001"
NODE2_URI = "zc://localhost:9002"
NODE1_HTTP = "http://localhost:9001"
NODE2_HTTP = "http://localhost:9002"
@dataclass
class NodeClient:
label: str
uri: str http_url: str api_key: str
user_id: str = field(init=False)
def __post_init__(self) -> None:
os.environ["ZAKURO_AUTH"] = self.api_key
config = ProcessorConfig.from_uri(self.uri)
compute = Compute(uri=self.uri, cpus=1.0)
self._proc = BrokerProcessor(config, compute)
self._proc.connect()
if self.api_key.startswith("zk_"):
parts = self.api_key[3:]
idx = parts.rfind("_")
self.user_id = parts[:idx] if idx > 0 else "unknown"
else:
self.user_id = "master"
def whoami(self) -> dict:
return self._proc.whoami()
def balance(self) -> float:
return self.whoami()["balance"]
def workers(self) -> list[dict]:
return self._proc.list_workers()
def estimate_price(self, cpus: float = 1.0, memory_gib: float = 0.25,
duration_secs: float = 1.0) -> dict:
return self._proc.estimate_price(
cpus=cpus, memory_gib=memory_gib, gpus=0, duration_secs=duration_secs
)
def fund(self, amount: float, description: str = "test funding") -> float:
r = httpx.post(
f"{self.http_url}/credits/{self.user_id}/add",
json={"amount": amount, "description": description},
headers={"X-Api-Key": NODE1_MASTER_KEY},
timeout=10,
)
r.raise_for_status()
return r.json()["new_balance"]
def execute(
self,
fn,
*,
cpus: float = 1.0,
memory_mb: int = 256,
duration_secs: float = 1.0,
budget_credits: Optional[float] = None,
remote_only: bool = False,
strategy: str = "best_price",
worker_type: Optional[str] = None,
) -> "ExecResult":
payload = cloudpickle.dumps({"func": fn, "args": (), "kwargs": {}})
requirements: dict = {
"cpus": cpus,
"memory_bytes": memory_mb * 1024 * 1024,
"gpus": 0,
"estimated_duration_secs": duration_secs,
"strategy": strategy,
"remote_only": remote_only,
}
if budget_credits is not None:
requirements["budget_credits"] = budget_credits
if worker_type is not None:
requirements["worker_type"] = worker_type
r = httpx.post(
f"{self.http_url}/execute",
content=payload,
headers={
"Content-Type": "application/octet-stream",
"Authorization": f"Bearer {self.api_key}",
"X-Zakuro-Requirements": json.dumps(requirements),
},
timeout=httpx.Timeout(connect=10, read=120, write=30, pool=10),
)
cost = float(r.headers.get("X-Zakuro-Cost", "0") or "0")
remaining = float(r.headers.get("X-Zakuro-Credits-Remaining", "0") or "0")
worker_id = r.headers.get("X-Zakuro-Worker", "?")
return ExecResult(
status=r.status_code,
cost=cost,
remaining=remaining,
worker_id=worker_id,
error=_parse_error(r) if r.status_code >= 400 else None,
)
def health(self) -> dict:
r = httpx.get(f"{self.http_url}/health", timeout=5)
r.raise_for_status()
return r.json()
def close(self) -> None:
self._proc.disconnect()
@dataclass
class ExecResult:
status: int
cost: float
remaining: float
worker_id: str
error: Optional[str] = None
@property
def ok(self) -> bool:
return self.status == 200
@property
def rejected_budget(self) -> bool:
return self.status == 402
@property
def no_workers(self) -> bool:
return self.status == 503
def _parse_error(r: httpx.Response) -> str:
try:
return r.json().get("error", r.text)
except Exception:
return r.text
def workload_fast():
import time
time.sleep(0.05)
import socket
return {"ok": True, "host": socket.gethostname()}
def workload_cpu(n: int = 30):
def fib(k):
return k if k < 2 else fib(k - 1) + fib(k - 2)
return fib(n)
PASS = "\033[32m✓ PASS\033[0m"
FAIL = "\033[31m✗ FAIL\033[0m"
WARN = "\033[33m⚠ WARN\033[0m"
INFO = "\033[36m·\033[0m"
_results: list[tuple[str, bool, str]] = []
def check(name: str, condition: bool, detail: str = "") -> bool:
_results.append((name, condition, detail))
status = PASS if condition else FAIL
detail_str = f" {detail}" if detail else ""
print(f" {status} {name}{detail_str}")
return condition
def separator(title: str) -> None:
print(f"\n{'═' * 64}")
print(f" {title}")
print('═' * 64)
def info(msg: str) -> None:
print(f" {INFO} {msg}")
def balance_row(label: str, n1_bal: float, n2_bal: float) -> None:
print(f" {label:<30} node1={n1_bal:>12.6f} node2={n2_bal:>12.6f}")
def main() -> None:
separator("SETUP — Connect to brokers via zc:// URI")
n1 = NodeClient("Node 1", NODE1_URI, NODE1_HTTP, NODE1_KEY)
n2 = NodeClient("Node 2", NODE2_URI, NODE2_HTTP, NODE2_KEY)
h1 = n1.health()
h2 = n2.health()
info(f"node1 ({NODE1_URI}) name={h1['node_name']} ts_ip={h1.get('wireguard_ip','?')} ts={h1['wireguard_connected']}")
info(f"node2 ({NODE2_URI}) name={h2['node_name']} ts_ip={h2.get('wireguard_ip','?')} ts={h2['wireguard_connected']}")
check("node1 WireGuard connected", h1["wireguard_connected"])
check("node2 WireGuard connected", h2["wireguard_connected"])
separator("SETUP — Worker topology")
for node, label in [(n1, "node1"), (n2, "node2")]:
workers = node.workers()
local_w = [w for w in workers if "127.0.0.1" in w.get("uri", "")]
remote_w = [w for w in workers if "127.0.0.1" not in w.get("uri", "")]
info(f"{label}: {len(workers)} workers — {len(local_w)} local, {len(remote_w)} remote")
for w in workers:
loc = "LOCAL " if "127.0.0.1" in w.get("uri", "") else "REMOTE"
info(f" [{loc}] {w.get('name','?'):<22} uri={w.get('uri','?'):<30} min_charge={w.get('min_charge', 0):.3f}")
check(f"{label} has local workers", len(local_w) > 0)
check(f"{label} has remote workers", len(remote_w) > 0)
separator("SETUP — Fund node1 account (standalone ledger)")
fund_amount = 50.0
new_bal = n1.fund(fund_amount, "transaction test seed")
info(f"Funded node1 with {fund_amount} credits → new balance: {new_bal:.4f}")
check("node1 balance after funding > 0", new_bal > 0,
f"balance={new_bal:.4f}")
b1_base = n1.balance()
b2_base = n2.balance()
separator("BASELINE BALANCES")
balance_row("Baseline", b1_base, b2_base)
separator("TEST 1 — Local execution (cost must be 0)")
info("Submitting job to each node without remote_only → picks local worker")
for node, label, base in [(n1, "node1", b1_base), (n2, "node2", b2_base)]:
pre = node.balance()
r = node.execute(workload_fast, cpus=1.0, duration_secs=0.1,
remote_only=False, strategy="best_price")
post = node.balance()
delta = pre - post
check(f"{label} local exec succeeded (200)", r.ok,
f"status={r.status} error={r.error}")
check(f"{label} local exec cost = 0", r.cost == 0.0,
f"cost={r.cost:.6f}")
check(f"{label} local balance unchanged", abs(delta) < 1e-9,
f"delta={delta:+.6f}")
info(f" worker={r.worker_id} cost={r.cost:.6f} balance: {pre:.4f} → {post:.4f}")
separator("TEST 2 — Remote execution: node1 → node2 workers (MIN_CHARGE=1.0)")
info("node1 submits with remote_only=True → must route to node2's workers")
pre1 = n1.balance()
est = n1.estimate_price(cpus=1.0, memory_gib=0.25, duration_secs=1.0)
info(f"Price estimate: min={est['min_cost']:.4f} max={est['max_cost']:.4f} workers={est['matching_workers']}")
r = n1.execute(workload_fast, cpus=1.0, duration_secs=1.0,
budget_credits=10.0, remote_only=True, strategy="best_price")
post1 = n1.balance()
delta1 = pre1 - post1
check("node1 remote exec succeeded (200)", r.ok,
f"status={r.status} error={r.error}")
check("node1 remote exec cost > 0", r.cost > 0,
f"cost={r.cost:.6f}")
check("node1 remote exec cost ≥ min_charge (1.0)", r.cost >= 1.0,
f"cost={r.cost:.6f} (node2 workers MIN_CHARGE=1.0)")
check("node1 balance debited by exact cost", abs(delta1 - r.cost) < 1e-6,
f"delta={delta1:.6f} cost_header={r.cost:.6f} diff={abs(delta1 - r.cost):.2e}")
check("remaining header = post-balance", abs(r.remaining - post1) < 1e-4,
f"remaining_header={r.remaining:.6f} actual_balance={post1:.6f}")
info(f" worker={r.worker_id} cost={r.cost:.6f} balance: {pre1:.4f} → {post1:.4f}")
separator("TEST 3 — Remote execution: node2 → node1 workers (no min_charge)")
info("node2 submits with remote_only=True → routes to node1's workers")
pre2 = n2.balance()
r = n2.execute(workload_fast, cpus=1.0, duration_secs=1.0,
budget_credits=10.0, remote_only=True, strategy="best_price")
post2 = n2.balance()
delta2 = pre2 - post2
check("node2 remote exec succeeded (200)", r.ok,
f"status={r.status} error={r.error}")
check("node2 remote exec cost > 0", r.cost > 0,
f"cost={r.cost:.6f}")
check("node2 balance debited by exact cost", abs(delta2 - r.cost) < 1e-4,
f"delta={delta2:.6f} cost_header={r.cost:.6f}")
info(f" worker={r.worker_id} cost={r.cost:.6f} balance: {pre2:.4f} → {post2:.4f}")
separator("TEST 4 — Invalid budget_credits=0 rejected at intake (400)")
info("budget_credits must be positive; zero must be rejected before routing")
for node, label in [(n1, "node1"), (n2, "node2")]:
pre = node.balance()
r = node.execute(workload_fast, cpus=1.0, duration_secs=1.0,
budget_credits=0.0)
post = node.balance()
check(f"{label} budget=0 rejected (400)", r.status == 400,
f"status={r.status} error={r.error}")
check(f"{label} balance unchanged on invalid budget", abs(pre - post) < 1e-9,
f"delta={pre - post:+.6f}")
separator("TEST 5 — Budget enforcement: too low for hourly rate (503 + balance unchanged)")
info("node1 remote_only=True with budget=0.001 for a 3600s job at price_per_hour=3.6")
info("max_price_per_hour = 0.001/1.0hr = 0.001 < worker price_per_hour=3.6 → no worker selected")
pre1 = n1.balance()
r = n1.execute(workload_fast, cpus=1.0, duration_secs=3600.0,
budget_credits=0.001, remote_only=True)
post1 = n1.balance()
check("budget too low for 1hr job returns 503 (no worker within budget)", r.no_workers,
f"status={r.status} error={r.error}")
check("node1 balance unchanged when no worker within budget", abs(pre1 - post1) < 1e-9,
f"pre={pre1:.6f} post={post1:.6f}")
info(f" status={r.status} error={r.error}")
info("Corollary: min_charge is a billing floor, not a dispatch gate")
info(" budget=2.0 for 1s job dispatches to node2 workers, actual bill = max(0.001/hr×1s, 1.0) = 1.0")
pre1 = n1.balance()
r_floor = n1.execute(workload_fast, cpus=1.0, duration_secs=1.0,
budget_credits=2.0, remote_only=True)
post1 = n1.balance()
check("min_charge floor: billed 1.0 even though compute cost < 1.0", r_floor.ok and r_floor.cost >= 1.0,
f"cost={r_floor.cost:.6f}")
check("min_charge floor: balance debited by exact cost", abs((pre1 - post1) - r_floor.cost) < 1e-6,
f"delta={pre1-post1:.6f} cost={r_floor.cost:.6f}")
separator("TEST 6 — Price estimate accuracy")
info("Compare /price estimate with actual billed cost")
for node, label, opts in [
(n1, "node1→local", {"remote_only": False}),
(n2, "node2→local", {"remote_only": False}),
]:
est = node.estimate_price(cpus=1.0, memory_gib=0.25, duration_secs=1.0)
pre = node.balance()
r = node.execute(workload_fast, cpus=1.0, duration_secs=1.0, **opts)
post = node.balance()
check(f"{label} exec succeeded", r.ok, f"status={r.status}")
if r.ok and r.cost > 0:
within_range = est["min_cost"] <= r.cost <= est["max_cost"] * 2
check(f"{label} actual cost within estimate range", within_range,
f"cost={r.cost:.6f} estimate=[{est['min_cost']:.6f}, {est['max_cost']:.6f}]")
else:
info(f" {label}: cost=0 (local exec), skipping estimate check")
separator("TEST 7 — Cumulative balance tracking over N remote jobs")
N = 3
info(f"node1 runs {N} remote jobs (node2 workers, MIN_CHARGE=1.0 each)")
pre1 = n1.balance()
total_cost = 0.0
for i in range(N):
r = n1.execute(workload_fast, cpus=1.0, duration_secs=0.5,
budget_credits=5.0, remote_only=True, strategy="best_price")
if r.ok:
total_cost += r.cost
info(f" job {i+1}/{N}: cost={r.cost:.6f} worker={r.worker_id}")
else:
info(f" job {i+1}/{N}: FAILED status={r.status}")
post1 = n1.balance()
actual_debit = pre1 - post1
check(f"all {N} remote jobs succeeded", total_cost > 0,
f"total_cost={total_cost:.6f}")
check("cumulative cost matches balance deduction", abs(actual_debit - total_cost) < 1e-4,
f"sum_of_costs={total_cost:.6f} actual_debit={actual_debit:.6f} diff={abs(actual_debit-total_cost):.2e}")
check(f"each job charged ≥ min_charge (1.0)", total_cost >= N * 1.0,
f"total={total_cost:.6f} expected≥{N * 1.0:.1f}")
separator("FINAL BALANCE SUMMARY")
b1_final = n1.balance()
b2_final = n2.balance()
balance_row("Baseline", b1_base, b2_base)
balance_row("Final ", b1_final, b2_final)
balance_row("Δ spent ", b1_base - b1_final, b2_base - b2_final)
print()
separator("TEST RESULTS")
passed = sum(1 for _, ok, _ in _results if ok)
failed = sum(1 for _, ok, _ in _results if not ok)
total = len(_results)
print(f"\n {passed}/{total} passed", end="")
if failed:
print(f" ({failed} FAILED)")
print()
for name, ok, detail in _results:
if not ok:
print(f" {FAIL} {name} {detail}")
else:
print(" — all checks passed")
print()
n1.close()
n2.close()
sys.exit(0 if failed == 0 else 1)
if __name__ == "__main__":
main()