File size: 3,355 Bytes
b83a9c1 41a10e6 b83a9c1 41a10e6 b83a9c1 | 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 | #!/usr/bin/env python3
"""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()
|