scripts/benchmark-bundles-wan.py
Rama main · réplica de lectura del repositorio original.
#!/usr/bin/env python3
"""WAN comparison using the existing SSH runner shape and original manifests.
Control traffic and file data use the private croc relay, with SSH used only
for orchestration. Both builds use the same relay and fixed transport settings.
It captures croc output through pipes and retains timings, checksums, and logs.
"""
from bundle_bench_common import (
manifest,
sha256,
transfer_secret,
stop_process,
drain_output,
)
import argparse
import json
import os
from pathlib import Path
import secrets
import selectors
import shlex
import shutil
import subprocess
import time
from dataclasses import dataclass
@dataclass(frozen=True)
class BenchmarkConfig:
options: argparse.Namespace
root: Path
remote: str
ssh_command: list[str]
manifests: dict
settings: dict
def ssh(config, command, **kwargs):
return subprocess.run(
config.ssh_command + [command],
check=True,
text=True,
capture_output=True,
**kwargs,
)
def run(config, build, dataset, direction, profile, trial=0, timeout=900):
options, root, remote_root = config.options, config.root, config.remote
manifests, settings, ssh_command = (
config.manifests,
config.settings,
config.ssh_command,
)
run_id = (
"bundle-"
+ time.strftime("%H%M%S")
+ "-"
+ secrets.token_hex(3)
+ "-"
+ "-".join([build, dataset, direction, profile])
)
secret = transfer_secret()
work = root / "work" / run_id
work.mkdir(parents=True)
dest = work / "received"
if direction == "down":
dest.mkdir()
rdest = remote_root + "/receive/" + run_id
if direction == "up":
ssh(config, "mkdir -p " + shlex.quote(rdest))
common = [
"--relay6",
"",
"--debug",
"--yes",
"--ignore-stdin",
"--disable-clipboard",
"--overwrite",
"--relay",
options.relay_host + ":" + str(settings["port"]),
"--pass",
settings["password"],
]
if dataset == "large" or profile.endswith("nocompress"):
common += ["--no-compress"]
sender_extra = ["send"] + [
"--no-local",
"--transport",
"relay",
"--transfers",
"8",
"--port",
str(settings["port"]),
"--hash",
"xxhash",
]
if "zip" in profile:
sender_extra += ["--zip"]
local_source = str(options.video) if dataset == "large" else str(options.source)
remote_source = (
options.remote_video if dataset == "large" else options.remote_source
)
if direction == "up":
local_args = common + sender_extra + [local_source]
remote_args = common + ["--out", rdest]
else:
remote_args = common + sender_extra + [remote_source]
local_args = common + ["--out", str(dest)]
env = os.environ.copy()
for key in list(env):
if key.startswith("CROC_") or key.upper() in (
"HTTP_PROXY",
"HTTPS_PROXY",
"ALL_PROXY",
"SOCKS5_PROXY",
):
env.pop(key)
config_directory = root / "config"
config_directory.mkdir(exist_ok=True)
env["CROC_CONFIG_DIR"] = str(config_directory)
env["CROC_SECRET"] = secret
(config_directory / "version-check.json").write_text(
json.dumps(
{
"checked_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
"latest_version": "11.5.5",
}
)
)
record = {
"id": run_id,
"harness_version": 4,
"trial": trial,
"build": build,
"dataset": dataset,
"direction": direction,
"profile": profile,
"bytes": sum(x["size"] for x in manifests[dataset].values()),
"files": len(manifests[dataset]),
"local_args": local_args,
"remote_args": remote_args,
"destination": rdest if direction == "up" else str(dest),
}
# Passwords are redacted in saved metadata; runnable commands use settings.json.
for args in (record["local_args"], record["remote_args"]):
args[args.index("--pass") + 1] = "<private-relay-password>"
local_args = [
settings["password"] if a == "<private-relay-password>" else a
for a in local_args
]
remote_args = [
settings["password"] if a == "<private-relay-password>" else a
for a in remote_args
]
(root / "active.json").write_text(json.dumps(record, indent=2))
print("START " + run_id, flush=True)
started = time.monotonic()
remote = subprocess.Popen(
ssh_command
+ ["python3 " + shlex.quote(remote_root + "/bundle_remote_runner.py")],
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
)
remote.stdin.write(
json.dumps(
{
"id": run_id,
"build": build,
"secret": secret,
"args": remote_args,
"timeout": timeout,
"binary": (
options.remote_base if build == "base" else options.remote_new
),
}
)
)
remote.stdin.close()
ready = remote.stdout.readline()
if not ready:
raise RuntimeError("Remote did not start: " + remote.stderr.read())
event = json.loads(ready)
assert event["event"] == "started", event
with (root / "logs" / (run_id + "-local.log")).open("wb") as logfile:
local = subprocess.Popen(
[str(options.local_base if build == "base" else options.local_new)]
+ local_args,
cwd=work,
env=env,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
start_new_session=True,
)
drain = drain_output(local.stdout, logfile)
selector = selectors.DefaultSelector()
selector.register(remote.stdout, selectors.EVENT_READ)
remote_result = None
next_notice = time.monotonic() + 20
while True:
for key, _ in selector.select(timeout=0.2):
line = key.fileobj.readline()
if line:
candidate = json.loads(line)
if candidate["event"] == "finished":
remote_result = candidate
else:
selector.unregister(key.fileobj)
if local.poll() is not None and remote_result is not None:
break
elapsed = time.monotonic() - started
failed = local.poll() not in (None, 0) or (
remote_result is not None and remote_result.get("returncode") != 0
)
if elapsed > timeout + 15 or (failed and elapsed > 5):
stop_process(local, group=True)
# Kill only the croc process created for this particular run.
ssh(config, "kill -TERM " + str(event["pid"]) + " 2>/dev/null || true")
if remote_result is None:
line = remote.stdout.readline()
if line:
remote_result = json.loads(line)
break
if time.monotonic() > next_notice:
print(
"RUNNING " + run_id + " " + str(round(elapsed, 1)) + "s", flush=True
)
next_notice = time.monotonic() + 20
record["seconds"] = time.monotonic() - started
record["local_returncode"] = local.wait()
drain.join()
record["remote_returncode"] = (remote_result or {}).get("returncode")
record["remote_seconds"] = (remote_result or {}).get("seconds")
(root / "logs" / (run_id + "-remote.log")).write_text(
(remote_result or {}).get("log", "")
)
remote.wait(timeout=10)
if record["local_returncode"] == 0 and record["remote_returncode"] == 0:
if direction == "up":
code = (
"import sys;sys.path.insert(0,"
+ repr(remote_root)
+ ");from bundle_remote_runner import manifest;from pathlib import Path;import json;print(json.dumps(manifest(Path("
+ repr(rdest)
+ "))))"
)
got = json.loads(ssh(config, "python3 -c " + shlex.quote(code)).stdout)
else:
got = manifest(dest)
record["verified"] = got == manifests[dataset]
record["received_files"] = len(got)
if not record["verified"]:
record["missing"] = sorted(set(manifests[dataset]) - set(got))
record["extra"] = sorted(set(got) - set(manifests[dataset]))
record["different"] = [
key
for key in got
if key in manifests[dataset] and got[key] != manifests[dataset][key]
]
else:
record["verified"] = False
if direction == "up":
code = (
"from pathlib import Path;import shutil;p=Path("
+ repr(rdest)
+ ");assert p.parent==Path("
+ repr(remote_root + "/receive")
+ ') and p.name.startswith("bundle-");shutil.rmtree(p)'
)
ssh(config, "python3 -c " + shlex.quote(code))
elif record["verified"]:
shutil.rmtree(dest)
record["MBps"] = (
record["bytes"] / record["seconds"] / 1e6 if record["verified"] else None
)
with (root / "results.jsonl").open("a") as stream:
stream.write(json.dumps(record) + "\n")
(root / "active.json").unlink(missing_ok=True)
print(
"RESULT "
+ json.dumps(
{
k: record.get(k)
for k in (
"id",
"seconds",
"MBps",
"verified",
"local_returncode",
"remote_returncode",
"missing",
"extra",
"different",
)
}
),
flush=True,
)
return record
def main():
parser = argparse.ArgumentParser(
description="Paired WAN measurements against a real frozen client; inputs must be frozen outside the checkout."
)
parser.add_argument("--artifacts", type=Path, required=True)
parser.add_argument("--host", required=True)
parser.add_argument("--relay-host", required=True)
parser.add_argument("--remote-root", required=True)
parser.add_argument("--relay-settings", type=Path, required=True)
parser.add_argument("--manifest", type=Path, required=True)
parser.add_argument("--source", type=Path, required=True)
parser.add_argument("--video", type=Path, required=True)
parser.add_argument("--remote-source", required=True)
parser.add_argument("--remote-video", required=True)
parser.add_argument("--local-base", type=Path, required=True)
parser.add_argument("--local-new", type=Path, required=True)
parser.add_argument("--remote-base", required=True)
parser.add_argument("--remote-new", required=True)
parser.add_argument("--repeats", type=int, default=3)
parser.add_argument(
"--datasets", nargs="+", choices=["src", "large"], default=["src", "large"]
)
parser.add_argument(
"--directions", nargs="+", choices=["up", "down"], default=["up", "down"]
)
parser.add_argument(
"--profiles",
nargs="+",
choices=["raw", "zip"],
help="Restrict repeat measurements to selected existing modes",
)
parser.add_argument(
"--resume",
action="store_true",
help="Reuse verified trials after checking binary and input digests",
)
options = parser.parse_args()
root = options.artifacts.resolve()
root.mkdir(parents=True, exist_ok=True)
(root / "logs").mkdir(exist_ok=True)
remote_root = options.remote_root
ssh_command = [
"ssh",
"-o",
"BatchMode=yes",
"-o",
"ControlMaster=auto",
"-o",
"ControlPersist=1800",
"-o",
"ControlPath=" + str(root / "ssh-control"),
options.host,
]
# Establish the persistent master with detached stdio. Otherwise its
# background process can retain capture_output pipes after the command ends.
subprocess.run(
ssh_command + ["true"],
check=True,
stdin=subprocess.DEVNULL,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
manifests = json.loads(options.manifest.read_text())
settings = json.loads(options.relay_settings.read_text())
config = BenchmarkConfig(
options, root, remote_root, ssh_command, manifests, settings
)
remote_paths = [options.remote_base, options.remote_new]
code = (
"from pathlib import Path;import sys,json;sys.path.insert(0,"
+ repr(remote_root)
+ ");from bundle_bench_common import sha256;print(json.dumps({p:sha256(Path(p)) for p in "
+ repr(remote_paths)
+ "}))"
)
checks = {
str(path): sha256(path) for path in [options.local_base, options.local_new]
}
checks.update(json.loads(ssh(config, "python3 -c " + shlex.quote(code)).stdout))
frozen = {"binary_sha256": checks, "input_manifests": manifests}
record_path = root / "frozen-inputs-and-builds.json"
if record_path.exists() and json.loads(record_path.read_text()) != frozen:
raise SystemExit(
"Benchmark binaries or frozen manifests changed; use a fresh artifact directory"
)
record_path.write_text(json.dumps(frozen, indent=2) + "\n")
results_path = root / "results.jsonl"
completed = set()
if results_path.exists():
if not options.resume:
raise SystemExit(
"Results already exist; use --resume or a fresh artifact directory"
)
completed = {
(r["dataset"], r["direction"], r["profile"], r["trial"], r["build"])
for line in results_path.read_text().splitlines()
if (r := json.loads(line)).get("verified")
}
for dataset in options.datasets:
for direction in options.directions:
for profile in (["raw", "zip"] if dataset == "src" else ["raw"]):
if options.profiles and profile not in options.profiles:
continue
for trial in range(options.repeats):
builds = ["base", "new"] if trial % 2 == 0 else ["new", "base"]
for build in builds:
if (dataset, direction, profile, trial, build) in completed:
continue
result = run(config, build, dataset, direction, profile, trial)
if not result["verified"]:
raise SystemExit(
"Independent verification failed: " + result["id"]
)
if __name__ == "__main__":
main()