| |
| """Exercise a deployed RunPod serverless endpoint. |
| |
| python tests/endpoint_test.py --endpoint-id XXXX main.jpg [ref.png] \ |
| [--prompt "..."] [--mode turbo-8] [--seed 1] [--sync] |
| |
| Reads RUNPOD_API_KEY from the environment. Saves returned images next to this |
| script and prints per-phase timings (delay/queue vs execution). |
| """ |
| import argparse |
| import base64 |
| import json |
| import os |
| import pathlib |
| import sys |
| import time |
| import urllib.request |
|
|
|
|
| def api(method, url, payload=None, timeout=120): |
| req = urllib.request.Request( |
| url, method=method, |
| data=json.dumps(payload).encode() if payload is not None else None, |
| headers={"Content-Type": "application/json", |
| "Authorization": f"Bearer {os.environ['RUNPOD_API_KEY']}"}) |
| with urllib.request.urlopen(req, timeout=timeout) as r: |
| return json.load(r) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("images", nargs="+") |
| ap.add_argument("--endpoint-id", required=True) |
| ap.add_argument("--prompt", default="replace the outfit of the subject in the first " |
| "image with the outfit in the second image, don't change anything else.") |
| ap.add_argument("--mode", default=None, help="only for workflows that support it") |
| ap.add_argument("--seed", type=int) |
| ap.add_argument("--sync", action="store_true", help="use /runsync instead of /run+poll") |
| ap.add_argument("--timeout", type=float, default=900) |
| args = ap.parse_args() |
|
|
| imgs = [{"name": os.path.basename(p), |
| "image": base64.b64encode(open(p, "rb").read()).decode()} |
| for p in args.images] |
| params = {"prompt": args.prompt} |
| if args.mode: |
| params["mode"] = args.mode |
| if args.seed is not None: |
| params["seed"] = args.seed |
| payload = {"input": {"images": imgs, "params": params}} |
| base = f"https://api.runpod.ai/v2/{args.endpoint_id}" |
|
|
| t0 = time.monotonic() |
| if args.sync: |
| result = api("POST", f"{base}/runsync", payload, timeout=args.timeout) |
| else: |
| job = api("POST", f"{base}/run", payload) |
| print(f"job {job['id']} -> {job['status']}", file=sys.stderr) |
| while True: |
| result = api("GET", f"{base}/status/{job['id']}") |
| if result["status"] in ("COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"): |
| break |
| if time.monotonic() - t0 > args.timeout: |
| sys.exit(f"timed out after {args.timeout}s (status {result['status']})") |
| print(f" {result['status']} ... {time.monotonic()-t0:.0f}s", file=sys.stderr) |
| time.sleep(2) |
| wall = time.monotonic() - t0 |
|
|
| if result.get("status") != "COMPLETED": |
| sys.exit(f"FAILED: {json.dumps(result, indent=2)[:3000]}") |
| out = result["output"] |
| if "error" in out: |
| sys.exit(f"handler error: {out['error']}") |
| for i, img in enumerate(out["images"]): |
| dest = pathlib.Path(__file__).parent / f"endpoint-out-{i}-{img['filename']}" |
| dest.write_bytes(base64.b64decode(img["data"])) |
| print(f"saved {dest}") |
| print(f"status=COMPLETED wall={wall:.1f}s delayTime={result.get('delayTime')}ms " |
| f"executionTime={result.get('executionTime')}ms\n" |
| f"handler timings={out.get('timings')} seed={out.get('seed')}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
