| """Load test a running design-canvas server: many concurrent episodes over the OpenEnv client, no model calls. |
| |
| DESIGN_CANVAS_CORPUS=<same as the server> python scripts/loadtest.py --url http://127.0.0.1:8007 \\ |
| --surface tools --episodes 64 --concurrency 32 |
| |
| Each episode is what training does to the server, minus the model: |
| - tools: replays the expert's tool calls, then submit (should score ~1.0); |
| - html: one write_html of the expert page, then submit; |
| - gui: `--steps` rounds of screenshot + click on the page + escape, then submit (built-in editor, or a real app with |
| --skin penpot|impress|inkscape-desktop|gimp). |
| |
| Prints throughput, reset/step/submit latency percentiles and every error, and writes them as JSON with --json. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import concurrent.futures as cf |
| import json |
| import random |
| import statistics as st |
| import sys |
| import threading |
| import time |
| from pathlib import Path |
|
|
| from openenv.core.env_server.mcp_types import CallToolAction |
| from openenv.core.mcp_client import MCPToolClient |
|
|
| try: |
| from design_canvas.core import expert |
| from design_canvas.core.pack import load_pack |
| except ImportError: |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) |
| from design_canvas.core import expert |
| from design_canvas.core.pack import load_pack |
|
|
|
|
| def authorize(token: str) -> None: |
| """Send an HF token with every request the OpenEnv client makes (a private Space needs it). The client has no |
| option for headers, so its WebSocket connect and its HTTP client are wrapped here, in this script only.""" |
| import functools |
|
|
| import httpx |
| import openenv.core.env_client as ec |
|
|
| auth = {"Authorization": f"Bearer {token}"} |
| ec.ws_connect = functools.partial(ec.ws_connect, additional_headers=auth) |
| base = httpx.AsyncClient.__init__ |
|
|
| def init(self, *a, **kw): |
| kw["headers"] = {**auth, **(kw.get("headers") or {})} |
| base(self, *a, **kw) |
|
|
| httpx.AsyncClient.__init__ = init |
|
|
|
|
| def pct(xs: list[float], q: float) -> float | None: |
| if not xs: |
| return None |
| xs = sorted(xs) |
| return round(xs[min(len(xs) - 1, int(q * len(xs)))], 3) |
|
|
|
|
| def episode(url: str, pack, task: dict, surface: str, skin: str | None, steps: int, seed: int) -> dict: |
| out = {"task_id": task["task_id"], "reset_s": None, "steps_s": [], "submit_s": None, "reward": None, "error": None, |
| "stage": "connect"} |
| rng = random.Random(seed) |
| try: |
| with MCPToolClient(base_url=url, message_timeout_s=600).sync() as env: |
| out["stage"] = "reset" |
| t = time.time() |
| env.reset(task_id=task["task_id"], surface=surface, mode="reference", |
| **({"skin": skin} if skin else {})) |
| out["reset_s"] = time.time() - t |
| gt = pack.gt_doc(task) |
| if surface == "tools": |
| calls = [c for c in expert.tool_calls(gt) if c["name"] != "submit"] |
| elif surface == "html": |
| calls = [{"name": "write_html", "arguments": {"html": expert.html_page(gt)}}] |
| else: |
| calls = [] |
| for _ in range(steps): |
| calls += [{"name": "screenshot", "arguments": {}}, |
| {"name": "left_click", "arguments": {"coordinate": [rng.randrange(300, 980), |
| rng.randrange(150, 650)]}}, |
| {"name": "key", "arguments": {"keys": "escape"}}] |
| for c in calls: |
| out["stage"] = c["name"] |
| t = time.time() |
| r = env.step(CallToolAction(tool_name=c["name"], arguments=c["arguments"])) |
| out["steps_s"].append(time.time() - t) |
| err = getattr(r.observation, "error", None) |
| if err is not None: |
| raise RuntimeError(f"{c['name']}: {getattr(err, 'message', err)}") |
| out["stage"] = "submit" |
| t = time.time() |
| r = env.step(CallToolAction(tool_name="submit", arguments={})) |
| out["submit_s"] = time.time() - t |
| out["reward"] = r.reward |
| out["stage"] = "close" |
| out["stage"] = "done" |
| except Exception as e: |
| out["error"] = f"{type(e).__name__} at {out['stage']}: {str(e)[:300]}" |
| return out |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--url", required=True) |
| ap.add_argument("--surface", choices=["tools", "html", "gui"], default="tools") |
| ap.add_argument("--skin", default=None) |
| ap.add_argument("--split", default="test") |
| ap.add_argument("--episodes", type=int, default=32) |
| ap.add_argument("--concurrency", type=int, default=16) |
| ap.add_argument("--steps", type=int, default=5, help="gui: screenshot+click+key rounds per episode") |
| ap.add_argument("--json", default=None, help="write the results here") |
| ap.add_argument("--hf-token-env", default=None, help="env var holding an HF token, for a private Space") |
| a = ap.parse_args(argv) |
| if a.hf_token_env: |
| import os |
|
|
| authorize(os.environ[a.hf_token_env]) |
|
|
| pack = load_pack() |
| tasks = [pack.at(a.split, i % pack.count(a.split)) for i in range(a.episodes)] |
| lock, done = threading.Lock(), [] |
| t0 = time.time() |
| with cf.ThreadPoolExecutor(a.concurrency) as pool: |
| futs = [pool.submit(episode, a.url, pack, t, a.surface, a.skin, a.steps, i) for i, t in enumerate(tasks)] |
| for f in cf.as_completed(futs): |
| r = f.result() |
| with lock: |
| done.append(r) |
| if r["error"]: |
| print(f" error {r['task_id']}: {r['error']}", flush=True) |
| wall = time.time() - t0 |
| ok = [r for r in done if not r["error"]] |
| steps = [s for r in ok for s in r["steps_s"]] |
| summary = { |
| "surface": a.surface, "skin": a.skin, "episodes": a.episodes, "concurrency": a.concurrency, |
| "ok": len(ok), "errors": len(done) - len(ok), "wall_s": round(wall, 1), |
| "episodes_per_min": round(60 * len(ok) / wall, 1), "steps_per_s": round(len(steps) / wall, 1), |
| "reset_p50": pct([r["reset_s"] for r in ok], .5), "reset_p95": pct([r["reset_s"] for r in ok], .95), |
| "step_p50": pct(steps, .5), "step_p95": pct(steps, .95), |
| "submit_p50": pct([r["submit_s"] for r in ok], .5), "submit_p95": pct([r["submit_s"] for r in ok], .95), |
| "reward_mean": round(st.fmean([r["reward"] or 0 for r in ok]), 4) if ok else None, |
| "error_kinds": sorted({r["error"].split(":")[0] for r in done if r["error"]}), |
| "results_lost": sum(1 for r in done if r["error"] and r["reward"] is not None), |
| } |
| print(json.dumps(summary)) |
| if a.json: |
| Path(a.json).write_text(json.dumps({"summary": summary, "episodes": done}, indent=1)) |
| return 0 if not summary["errors"] else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
