File size: 7,227 Bytes
0c85e62 | 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 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | #!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
#
# ComfyUI-Koolook -- live run/poll smoke check for the published-setup API.
"""Exercise the published-setup run -> poll flow against a live ComfyUI.
This is a *live* smoke check (not a unit test): it queues a real render through
the same routes the runner simulator uses, then polls the run endpoint exactly
like the simulator's poller, printing the status timeline. Use it to self-verify
end-to-end behaviour (e.g. "does a long render reach succeeded?") without driving
the browser UI.
python tools/check_setup_run.py # first valid setup, default inputs
python tools/check_setup_run.py --match upscale # first setup whose id contains "upscale"
python tools/check_setup_run.py --id pub_ltx-2.3_upscale_v09 --inputs '{"prompt":"a bear"}'
python tools/check_setup_run.py --list # just list setups and exit
Exits 0 on a succeeded run, 1 otherwise. Default base is http://127.0.0.1:8000.
"""
from __future__ import annotations
import argparse
import json
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
TERMINAL = {"succeeded", "failed", "lost"}
ACTIVE = {"queued", "running"}
ALLOWED_URL_SCHEMES = {"http", "https"}
def _checked_url(url: str) -> str:
scheme = urllib.parse.urlsplit(url).scheme
if scheme not in ALLOWED_URL_SCHEMES:
raise ValueError(f"unsupported URL scheme: {scheme or '<empty>'}")
return url
def _get(url: str, timeout: float = 30.0):
with urllib.request.urlopen(_checked_url(url), timeout=timeout) as r: # nosec B310
return json.load(r)
def _post(url: str, payload: dict, timeout: float = 60.0):
data = json.dumps(payload).encode()
req = urllib.request.Request(
_checked_url(url),
data=data,
headers={"Content-Type": "application/json"},
method="POST",
)
with urllib.request.urlopen(req, timeout=timeout) as r: # nosec B310
return json.load(r)
def _progress_signature(run: dict) -> str:
outputs = []
for output in run.get("outputs") or []:
if not isinstance(output, dict):
continue
items = output.get("items") if isinstance(output.get("items"), list) else []
outputs.append(
{
"key": output.get("key"),
"itemCount": len(items),
"values": [item.get("value", "") for item in items if isinstance(item, dict) and item.get("value")],
}
)
return json.dumps(
{
"status": str(run.get("status", "")).lower(),
"promptId": run.get("promptId", ""),
"queuePosition": run.get("queuePosition", run.get("position")),
"updatedAt": run.get("updatedAt") or run.get("timestamp") or run.get("lastUpdate") or "",
"outputs": outputs,
},
sort_keys=True,
)
def list_setups(base: str) -> list[dict]:
body = _get(f"{base}/koolook/api/setups")
rows = body.get("setups") if isinstance(body, dict) else body
return rows if isinstance(rows, list) else []
def pick_setup(rows: list[dict], setup_id: str | None, match: str | None) -> dict | None:
valid = [r for r in rows if (r.get("validation") or {}).get("status") == "valid"]
if setup_id:
return next((r for r in rows if r.get("id") == setup_id), None)
if match:
return next((r for r in valid if match.lower() in str(r.get("id", "")).lower()), None)
return valid[0] if valid else None
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--base", default="http://127.0.0.1:8000")
ap.add_argument("--id", default=None, help="exact setup id to run")
ap.add_argument("--match", default=None, help="substring to pick a setup id")
ap.add_argument("--inputs", default="{}", help="JSON object of run inputs")
ap.add_argument("--list", action="store_true", help="list setups and exit")
ap.add_argument("--interval", type=float, default=3.0, help="poll interval seconds")
ap.add_argument("--stall", type=float, default=180.0,
help="give up after this many seconds with no active/terminal status")
args = ap.parse_args()
base = args.base.rstrip("/")
try:
rows = list_setups(base)
except urllib.error.URLError as exc:
print(f"FAIL: cannot reach ComfyUI at {base}: {exc}")
return 1
print(f"{len(rows)} published setup(s):")
for r in rows:
print(f" {r.get('id'):<34} status={(r.get('validation') or {}).get('status')}")
if args.list:
return 0
setup = pick_setup(rows, args.id, args.match)
if setup is None:
print("FAIL: no matching/valid setup to run.")
return 1
setup_id = setup["id"]
try:
inputs = json.loads(args.inputs)
assert isinstance(inputs, dict)
except Exception as exc:
print(f"FAIL: --inputs must be a JSON object: {exc}")
return 1
print(f"\nRunning '{setup_id}' with inputs={inputs} ...")
started = time.monotonic()
try:
queued = _post(f"{base}/koolook/api/setups/{setup_id}/run", {"inputs": inputs})
except urllib.error.HTTPError as exc:
print(f"FAIL: run request HTTP {exc.code}: {exc.read().decode()[:300]}")
return 1
run = queued.get("run") or {}
run_id = run.get("runId")
print(f" queued runId={run_id} promptId={run.get('promptId')} status={run.get('status')}")
if not run_id:
print("FAIL: no runId returned.")
return 1
# Poll exactly like the simulator: treat --stall as a no-progress window.
# The deadline resets only when the run payload changes in a meaningful way;
# repeated queued/running responses with no changing marker are a stall.
deadline = time.monotonic() + args.stall
last_status = None
last_progress = _progress_signature(run)
while time.monotonic() <= deadline:
time.sleep(args.interval)
try:
body = _get(f"{base}/koolook/api/runs/{run_id}")
except urllib.error.HTTPError as exc:
print(f" poll HTTP {exc.code} -- retrying")
continue
run = body.get("run") or {}
status = str(run.get("status", "")).lower()
elapsed = int(time.monotonic() - started)
if status != last_status:
print(f" [t+{elapsed}s] status={status}")
last_status = status
if status in TERMINAL:
ok = status == "succeeded"
print(f"\n{'PASS' if ok else 'FAIL'}: run {status} after {elapsed}s")
for out in run.get("outputs") or []:
for item in out.get("items") or []:
if item.get("value"):
print(f" {out.get('label') or out.get('key')}: {item['value']}")
return 0 if ok else 1
if status in ACTIVE:
progress = _progress_signature(run)
if progress != last_progress:
last_progress = progress
deadline = time.monotonic() + args.stall
print(f"\nFAIL: no terminal status within the {args.stall}s no-progress window.")
return 1
if __name__ == "__main__":
sys.exit(main())
|