scripts/benchmark.py

#!/usr/bin/env python3
"""Measure real, authenticated Space outputs. Never imports the app locally."""

import argparse
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import random
import shutil
import sys
import time

from gradio_client import Client, handle_file
from huggingface_hub import HfApi, get_token
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from settings import DEFAULT_QUALITY, EXAMPLE_PROMPTS, SIZES, STEPS, duration_cap

# Held out from the four UI studies; do not change these after seeing outputs.
CHALLENGE_PROMPTS = [
    "A straight-on studio photograph of exactly three objects on one white shelf: "
    "a red wooden cube on the left, a green ceramic sphere in the middle, and a blue "
    "metal pyramid on the right. All objects are separate, equally spaced and fully "
    "visible. Soft neutral lighting, no other objects, no text.",
    "A documentary portrait of an adult female ceramic artist at a wooden workbench. "
    "Both of her hands are fully visible, one on either side of a single small white "
    "ceramic bowl. Five natural fingers on each hand. Linen shirt, daylight from a "
    "large window, realistic skin and clay texture, calm expression, eye-level camera.",
]


def sha256(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def file_path(value):
    if isinstance(value, str):
        return value
    if isinstance(value, dict):
        return value.get("path") or value.get("name")
    raise ValueError(f"Unexpected file response type: {type(value).__name__}")


def trials(suite, repeats):
    if suite == "smoke":
        return [{"prompt_index": 0, "quality": "Draft · 20 steps", "resolution": "1K",
                 "aspect": "1:1", "seed": 42, "repeat": 0}]
    if suite == "repeatability":
        return [{"prompt_index": 0, "quality": DEFAULT_QUALITY, "resolution": "1K",
                 "aspect": "1:1", "seed": 42, "repeat": r} for r in range(repeats)]
    if suite == "shapes":
        return [{"prompt_index": 0, "quality": "Draft · 20 steps", "resolution": "1K",
                 "aspect": aspect, "seed": 42, "repeat": 0} for aspect in SIZES["1K"]]
    if suite == "resolution":
        return [{"prompt_index": 0, "quality": DEFAULT_QUALITY, "resolution": resolution,
                 "aspect": "1:1", "seed": 42, "repeat": r}
                for r in range(repeats) for resolution in ["1K", "2K"]]
    if suite == "steps":
        rows = [{"prompt_index": p, "quality": quality, "resolution": "1K",
                 "aspect": "1:1", "seed": seed, "repeat": r}
                for r in range(repeats) for p in range(4) for seed in [42, 12345]
                for quality in STEPS]
        random.Random(20260920).shuffle(rows)
        return rows
    if suite == "challenge":
        rows = [{"prompt_index": p+4, "prompt": prompt, "quality": quality, "resolution": "1K",
                 "aspect": "1:1", "seed": seed, "repeat": r}
                for r in range(repeats) for p, prompt in enumerate(CHALLENGE_PROMPTS)
                for seed in [42, 12345] for quality in STEPS]
        random.Random(20260921).shuffle(rows)
        return rows
    raise ValueError(suite)


def main():
    sys.stdout.reconfigure(line_buffering=True)
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("space")
    parser.add_argument("--suite", choices=["smoke", "repeatability", "shapes", "resolution", "steps", "challenge"], default="smoke")
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--max-calls", type=int, default=1,
                        help="Explicit call budget. Increase only after checking remaining ZeroGPU quota.")
    parser.add_argument("--max-gpu-seconds", type=float, default=180,
                        help="Stop before another call if observed usage plus the last requested cap exceeds this budget.")
    parser.add_argument("--cold-after-restart", action="store_true",
                        help="Assert a separately observed restart; only first app request can qualify.")
    parser.add_argument("--reference", type=Path)
    parser.add_argument("--prompt", help="Override the curated prompt, useful for edit verification.")
    parser.add_argument("--transparent", action="store_true")
    parser.add_argument("--resolution", choices=list(SIZES))
    parser.add_argument("--quality", choices=list(STEPS))
    parser.add_argument("--aspect", choices=list(SIZES["1K"]))
    parser.add_argument("--seed", type=int)
    args = parser.parse_args()
    if args.repeats < 1 or args.max_calls < 1 or args.max_gpu_seconds <= 0:
        parser.error("Repeats, call count, and GPU budget must be positive.")
    token = get_token()
    if not token:
        sys.exit("BLOCKED: Existing Hugging Face authentication is required; no anonymous private-Space calls.")
    api = HfApi(token=token)
    info = api.space_info(args.space)
    runtime = api.get_space_runtime(args.space)
    if str(runtime.stage) != "RUNNING" or str(runtime.hardware) != "zero-a10g":
        sys.exit(f"Space must be RUNNING on zero-a10g: stage={runtime.stage}, hardware={runtime.hardware}")
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
    out = ROOT / "artifacts" / "benchmarks" / f"{stamp}-{args.suite}"
    out.mkdir(parents=True, exist_ok=False)
    client = Client(args.space, token=token, httpx_kwargs={"timeout": 600},
                    download_files=str(out / "downloads"), verbose=False, analytics_enabled=False)
    discovered = client.view_api(return_format="dict")
    (out / "api.json").write_text(json.dumps(discovered, indent=2, default=str))
    named = discovered.get("named_endpoints", {})
    if "/generate" not in named:
        sys.exit("/generate is absent from discovered API. Inspect api.json before making a request.")
    variables = api.get_space_variables(args.space)
    manifest = {"space": args.space, "hub_commit": info.sha, "private": info.private,
                "runtime": runtime.raw, "suite": args.suite, "started_at": stamp,
                "cold_assertion": args.cold_after_restart,
                "max_calls": args.max_calls, "gpu_budget_seconds": args.max_gpu_seconds,
                "reference_sha256": sha256(args.reference) if args.reference else None,
                "experimental_variables": {key: variables[key].value for key in (
                    "STILL_GPU_SIZE", "STILL_GPU_SECONDS", "STILL_KV_CACHE", "STILL_VAE_TILING")
                    if key in variables}}
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2, default=str))
    configured_cap = int(variables["STILL_GPU_SECONDS"].value) if "STILL_GPU_SECONDS" in variables else None
    configured_size = variables["STILL_GPU_SIZE"].value if "STILL_GPU_SIZE" in variables else "xlarge"
    spent = 0.0
    records = []
    for index, trial in enumerate(trials(args.suite, args.repeats)[:args.max_calls]):
        trial.update({key: getattr(args, key) for key in ["resolution", "quality", "aspect", "seed"]
                      if getattr(args, key) is not None})
        cap = configured_cap * (1 if configured_size == "xlarge" else 1.5) if configured_cap else duration_cap(
            STEPS[trial["quality"]], int(trial["resolution"][0])*1024, bool(args.reference), configured_size)
        next_cap = cap * (2 if configured_size == "xlarge" else 1)
        if spent + next_cap > args.max_gpu_seconds:
            print("Stopped at the experiment's GPU budget. This is not a measurement of platform quota.")
            break
        prompt = args.prompt or trial.get("prompt") or EXAMPLE_PROMPTS[trial["prompt_index"]]
        params = [prompt, handle_file(str(args.reference)) if args.reference else None,
                  trial["quality"], trial["resolution"], trial["aspect"],
                  args.transparent or (trial["prompt_index"] == 3), trial["seed"], False]
        row = {"trial_index": index, "trial": trial, "status": "started",
               "started_at": datetime.now(timezone.utc).isoformat()}
        started = time.perf_counter()
        try:
            response = client.predict(*params, api_name="/generate")
            row["client_wall_seconds"] = time.perf_counter() - started
            if not isinstance(response, (list, tuple)) or len(response) != 4:
                raise ValueError("Expected image, status, recipe and seed from /generate")
            image_src, status, recipe_src, returned_seed = response
            image_dest = out / f"{index:03d}.png"
            recipe_dest = out / f"{index:03d}.json"
            shutil.copyfile(file_path(image_src), image_dest)
            shutil.copyfile(file_path(recipe_src), recipe_dest)
            recipe = json.loads(recipe_dest.read_text())
            with Image.open(image_dest) as img:
                img.load()
                expected = SIZES[trial["resolution"]][trial["aspect"]]
                if img.format != "PNG" or img.size != expected:
                    raise ValueError("Output format or dimensions do not match the requested image")
                pixel_hash = hashlib.sha256(img.mode.encode() + str(img.size).encode() + img.tobytes()).hexdigest()
                alpha_extrema = img.getchannel("A").getextrema() if "A" in img.getbands() else None
            if returned_seed != trial["seed"] or recipe["parameters"]["seed"] != trial["seed"]:
                raise ValueError("Returned seed does not match the fixed request")
            if sha256(image_dest) != recipe["output"]["png_sha256"]:
                raise ValueError("Downloaded PNG hash does not match the server recipe")
            is_first = recipe["request_index_in_boot"] == 1
            if args.cold_after_restart and index == 0 and not is_first:
                raise ValueError("Cold-start assertion failed: this was not the app's first GPU request")
            weight = 2 if recipe["metrics"]["gpu"]["size"] == "xlarge" else 1
            spent += recipe["metrics"]["gpu_call_wall_seconds"] * weight
            row.update(status="ok", recipe=recipe, image_file=image_dest.name,
                       recipe_file=recipe_dest.name, png_sha256=sha256(image_dest),
                       pixel_sha256=pixel_hash, alpha_extrema=alpha_extrema,
                       classification="first_after_observed_restart" if args.cold_after_restart and is_first
                       else "first_request_in_boot" if is_first
                       else "reused_gpu_worker" if recipe["metrics"].get("gpu_worker_call_index", 1) > 1
                       else "new_gpu_worker" if "gpu_worker_call_index" in recipe["metrics"]
                       else "subsequent_request_worker_warmth_unknown")
            print(f"{index:03d} {trial['quality']} {trial['resolution']} "
                  f"inference={recipe['metrics']['inference_seconds']:.2f}s "
                  f"client={row['client_wall_seconds']:.2f}s")
        except Exception as exc:
            row.update(status="failed", client_wall_seconds=time.perf_counter() - started,
                       error_type=type(exc).__name__)
            # Keep sensitive HTTP headers and credential-bearing traces out of artifacts.
            print(f"Trial failed: {type(exc).__name__}. Inspect Space logs; no automatic retry.")
        records.append(row)
        with (out / "results.jsonl").open("a") as file:
            file.write(json.dumps(row, ensure_ascii=False) + "\n")
        if row["status"] == "failed":
            break
    print(f"Actual completed records: {len(records)}. Evidence: {out}")
    if not records or any(row["status"] == "failed" for row in records):
        sys.exit(1)


if __name__ == "__main__":
    main()