File size: 7,115 Bytes
df84e45 d992b32 df84e45 d992b32 df84e45 d992b32 df84e45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """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: # run from the package directory
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: # noqa: BLE001 - every failure is a result here
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())
|