385 lines
16 KiB
Python
385 lines
16 KiB
Python
"""Inference engine end-to-end test (extensions_built_in/inference_engine).
|
|
|
|
Starts the engine in-process on an ephemeral loopback port, sends real
|
|
generate requests over HTTP, decodes the streamed frames and asserts on
|
|
them. Model/sample settings come from testing/test_model_loading.py's
|
|
MODEL_TESTS (tiny sizes: this proves the path, not quality).
|
|
|
|
Usage:
|
|
python testing/test_inference_engine.py --arch zimage # one arch
|
|
python testing/test_inference_engine.py --swap zimage,wan21 # sequential
|
|
# model switches
|
|
python testing/test_inference_engine.py --arch zimage --cancel # cancel mid-run
|
|
python testing/test_inference_engine.py --all # every arch,
|
|
# subprocess each
|
|
--allow-download permit hub downloads (default HF_HUB_OFFLINE=1 -> SKIP)
|
|
--device cuda:0
|
|
"""
|
|
|
|
import argparse
|
|
import http.client
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
TOOLKIT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, TOOLKIT_ROOT)
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv(os.path.join(TOOLKIT_ROOT, ".env"))
|
|
|
|
OUTPUT_ROOT = os.path.join(TOOLKIT_ROOT, "testing", ".engine_test_outputs")
|
|
|
|
from testing.test_model_loading import MODEL_TESTS, SKIP_MARKERS, classify_error # noqa: E402
|
|
|
|
|
|
class Client:
|
|
def __init__(self, host, port, token):
|
|
self.host, self.port, self.token = host, port, token
|
|
|
|
def _conn(self):
|
|
return http.client.HTTPConnection(self.host, self.port, timeout=3600)
|
|
|
|
def _headers(self):
|
|
return {"Content-Type": "application/json", "X-Engine-Token": self.token}
|
|
|
|
def get(self, path):
|
|
c = self._conn()
|
|
c.request("GET", path, headers=self._headers())
|
|
r = c.getresponse()
|
|
return r.status, json.loads(r.read())
|
|
|
|
def post_json(self, path, body=None):
|
|
c = self._conn()
|
|
c.request("POST", path, body=json.dumps(body or {}), headers=self._headers())
|
|
r = c.getresponse()
|
|
return r.status, json.loads(r.read())
|
|
|
|
def upload(self, path, name):
|
|
c = self._conn()
|
|
with open(path, "rb") as f:
|
|
data = f.read()
|
|
c.request("POST", f"/assets?name={name}", body=data, headers={"X-Engine-Token": self.token})
|
|
r = c.getresponse()
|
|
return r.status, json.loads(r.read())
|
|
|
|
def generate_stream(self, body):
|
|
"""Yields (header, payload) frames as they arrive."""
|
|
from extensions_built_in.inference_engine.protocol import FrameReader
|
|
|
|
c = self._conn()
|
|
c.request("POST", "/generate", body=json.dumps(body), headers=self._headers())
|
|
r = c.getresponse()
|
|
if r.status != 200:
|
|
raise RuntimeError(f"/generate -> {r.status}: {r.read()[:400]}")
|
|
|
|
def chunks():
|
|
while True:
|
|
data = r.read1(65536) # read1 returns the available chunk instead of blocking for 64KB
|
|
if not data:
|
|
return
|
|
yield data
|
|
|
|
reader = FrameReader(chunks())
|
|
for frame in reader:
|
|
yield frame
|
|
c.close()
|
|
|
|
|
|
def start_engine(device: str, out_root: str):
|
|
from extensions_built_in.inference_engine.engine import Engine
|
|
from extensions_built_in.inference_engine.server import create_app, serve_in_thread
|
|
|
|
import threading
|
|
|
|
engine = Engine(device=device, output_folder=os.path.join(out_root, "outputs"), assets_folder=os.path.join(out_root, "assets"))
|
|
token = "test-token"
|
|
server, thread, host, port = serve_in_thread(create_app(engine, token))
|
|
loop = threading.Thread(target=engine.run_forever, daemon=True, name="engine-loop")
|
|
loop.start()
|
|
return engine, server, Client(host, port, token)
|
|
|
|
|
|
def stop_engine(engine, server):
|
|
engine.stop()
|
|
server.should_exit = True
|
|
engine.unload()
|
|
|
|
|
|
QUANTIZE_ALL = False # --quantize: force convrot8 on transformer + TE (small GPUs)
|
|
|
|
|
|
def request_body(arch: str, entry: dict, client: Client, out_dir: str, tiny=True, stream=None):
|
|
model = dict(entry["model"], arch=arch)
|
|
if QUANTIZE_ALL:
|
|
model["quantize"] = True
|
|
model["quantize_te"] = True
|
|
if model.get("quantize"):
|
|
model.setdefault("qtype", "convrot8")
|
|
if model.get("quantize_te"):
|
|
model.setdefault("qtype_te", "convrot8")
|
|
sample = dict(entry["sample"])
|
|
sample["prompt"] = "a photo of a cat sitting on a wooden table"
|
|
if tiny and not entry.get("size_locked"):
|
|
# 384 divides every bucket size (16/32/64); video runs also keep the
|
|
# frame count small so the swap tests fit next to a resident TE
|
|
cap = 384 if tiny == "tiny" else 512
|
|
sample["width"] = min(sample["width"], cap)
|
|
sample["height"] = min(sample["height"], cap)
|
|
if tiny == "tiny" and "num_frames" in sample:
|
|
sample["num_frames"] = min(sample["num_frames"], 9)
|
|
if entry.get("needs_control_image"):
|
|
from PIL import Image
|
|
|
|
ctrl_path = os.path.join(out_dir, "ctrl.png")
|
|
Image.new("RGB", (sample["width"], sample["height"]), (128, 128, 128)).save(ctrl_path)
|
|
status, asset = client.upload(ctrl_path, "ctrl.png")
|
|
assert status == 200, asset
|
|
sample["ctrl_img"] = asset["path"]
|
|
return {"model": model, "sample": sample, "stream": stream or {"latents": "raw", "every_n_steps": 1, "max_frames": 2}}
|
|
|
|
|
|
def run_request(client: Client, body: dict, cancel_after: int = 0) -> dict:
|
|
"""Consume one streamed generation; returns a summary of the frames."""
|
|
import numpy as np
|
|
|
|
from extensions_built_in.inference_engine.protocol import payload_to_array
|
|
|
|
seen = {"start": 0, "status": 0, "progress": 0, "latent": 0, "result": 0, "error": 0, "end": 0}
|
|
summary = {"frames": seen, "results": [], "latent_shapes": [], "error": None, "cancelled": False, "request_id": None, "preview": None}
|
|
for header, payload in client.generate_stream(body):
|
|
t = header["type"]
|
|
seen[t] = seen.get(t, 0) + 1
|
|
summary["request_id"] = header.get("request_id")
|
|
if t == "start":
|
|
summary["preview"] = header.get("preview")
|
|
summary["modality"] = header.get("modality")
|
|
elif t == "status":
|
|
print(f" [status] {header.get('message')}")
|
|
elif t == "progress":
|
|
print(f" [progress] {header['step']}/{header['total']} {header['elapsed']}s", flush=True)
|
|
if cancel_after and header["step"] >= cancel_after:
|
|
st, res = client.post_json(f"/cancel/{summary['request_id']}")
|
|
assert st == 200 and res["ok"], res
|
|
cancel_after = 0
|
|
elif t == "latent":
|
|
arr = payload_to_array(header, payload)
|
|
assert np.isfinite(arr.astype(np.float32)).all(), "non-finite latent"
|
|
summary["latent_shapes"].append((header["layout"], list(arr.shape)))
|
|
elif t == "result":
|
|
summary["results"].append(header)
|
|
print(f" [result] {header['kind']} {header['path']} ({header['bytes']} bytes)")
|
|
elif t == "error":
|
|
summary["error"] = header.get("message")
|
|
summary["cancelled"] = bool(header.get("cancelled"))
|
|
if header.get("traceback"):
|
|
print(header["traceback"])
|
|
elif t == "end":
|
|
summary["status"] = header.get("status")
|
|
return summary
|
|
|
|
|
|
def assert_success(summary: dict, modality: str):
|
|
assert summary["frames"]["start"] == 1, summary
|
|
assert summary["frames"]["end"] == 1, summary
|
|
assert summary["error"] is None, summary["error"]
|
|
assert summary["frames"]["progress"] >= 1, "no progress frames"
|
|
assert summary["frames"]["latent"] >= 1, "no latent frames"
|
|
assert summary["results"], "no result frames"
|
|
for r in summary["results"]:
|
|
assert os.path.isfile(r["path"]) and r["bytes"] > 1024, r
|
|
kinds = {r["kind"] for r in summary["results"]}
|
|
assert modality in kinds or (modality == "video" and "image" in kinds), (modality, kinds)
|
|
|
|
|
|
def run_one(arch: str, device: str, cancel: bool = False) -> dict:
|
|
from toolkit.models.registry import get_arch_entry
|
|
|
|
entry = MODEL_TESTS[arch]
|
|
out_root = os.path.join(OUTPUT_ROOT, arch.replace("/", "_"))
|
|
shutil.rmtree(out_root, ignore_errors=True)
|
|
os.makedirs(out_root, exist_ok=True)
|
|
engine, server, client = start_engine(device, out_root)
|
|
try:
|
|
st, health = client.get("/health")
|
|
assert st == 200 and health["ok"], health
|
|
st, models = client.get("/models")
|
|
assert st == 200 and any(m["arch"] == arch for m in models["archs"]), f"{arch} missing from /models"
|
|
modality = get_arch_entry(arch)["modality"]
|
|
body = request_body(arch, entry, client, out_root)
|
|
|
|
t0 = time.perf_counter()
|
|
if cancel:
|
|
s = run_request(client, body, cancel_after=2)
|
|
assert s["cancelled"], f"expected a cancelled stream, got {s}"
|
|
assert s["frames"]["end"] == 1
|
|
print(" cancel: ok, re-running the same request")
|
|
s = run_request(client, body)
|
|
gen_seconds = time.perf_counter() - t0
|
|
assert_success(s, modality)
|
|
st, health = client.get("/health")
|
|
assert health["active"]["arch"] == arch, health
|
|
return {
|
|
"arch": arch,
|
|
"status": "PASS",
|
|
"seconds": round(gen_seconds, 2),
|
|
"frames": s["frames"],
|
|
"latent_shapes": s["latent_shapes"][:2],
|
|
"results": [r["path"] for r in s["results"]],
|
|
"preview_format": (s["preview"] or {}).get("format"),
|
|
"stats": engine.stats,
|
|
}
|
|
finally:
|
|
stop_engine(engine, server)
|
|
|
|
|
|
def run_swap(archs, device: str) -> dict:
|
|
out_root = os.path.join(OUTPUT_ROOT, "swap_" + "_".join(a.replace("/", "_") for a in archs))
|
|
shutil.rmtree(out_root, ignore_errors=True)
|
|
os.makedirs(out_root, exist_ok=True)
|
|
engine, server, client = start_engine(device, out_root)
|
|
from toolkit.models.registry import get_arch_entry
|
|
|
|
timeline = []
|
|
try:
|
|
for i, arch in enumerate(archs):
|
|
body = request_body(arch, MODEL_TESTS[arch], client, out_root, tiny="tiny")
|
|
t0 = time.perf_counter()
|
|
s = run_request(client, body)
|
|
assert_success(s, get_arch_entry(arch)["modality"])
|
|
pool = engine.pool.stats()
|
|
resident = [(e["cls"], round(e["bytes"] / 1e9, 2), e["device"]) for e in pool["resident"]]
|
|
entry = {
|
|
"arch": arch,
|
|
"seconds": round(time.perf_counter() - t0, 2),
|
|
"holder_loads": engine.stats["holder_loads"],
|
|
"oom_retries": engine.stats["oom_retries"],
|
|
"pool_hits": pool["hits"],
|
|
"pool_misses": pool["misses"],
|
|
"pool_evictions": pool["evictions"],
|
|
"resident_gb": round(pool["resident_bytes"] / 1e9, 2),
|
|
"resident": resident,
|
|
}
|
|
timeline.append(entry)
|
|
print(f" {arch}: {json.dumps(entry)}")
|
|
# after a request only this holder's components may stay resident
|
|
assert pool["untouched"] == 0, f"stale components left resident: {pool}"
|
|
if i > 0 and archs[i - 1] == arch:
|
|
assert entry["holder_loads"] == timeline[i - 1]["holder_loads"], "same model should reuse the holder"
|
|
if len(set(archs)) > 1:
|
|
assert engine.pool.hits > 0, "expected shared components between the swapped archs"
|
|
return {"arch": ",".join(archs), "status": "PASS", "timeline": timeline}
|
|
finally:
|
|
stop_engine(engine, server)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--arch", type=str, default=None)
|
|
parser.add_argument("--swap", type=str, default=None, help="comma-separated archs to run back to back")
|
|
parser.add_argument("--cancel", action="store_true")
|
|
parser.add_argument("--all", action="store_true")
|
|
parser.add_argument("--list", action="store_true")
|
|
parser.add_argument("--device", type=str, default="cuda:0")
|
|
parser.add_argument("--allow-download", action="store_true")
|
|
parser.add_argument("--json-result", type=str, default=None)
|
|
parser.add_argument("--quantize", action="store_true", help="force convrot8 quantization of transformer + TE")
|
|
args = parser.parse_args()
|
|
global QUANTIZE_ALL
|
|
QUANTIZE_ALL = args.quantize
|
|
|
|
if args.list:
|
|
for arch in MODEL_TESTS:
|
|
print(arch)
|
|
return
|
|
|
|
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
|
|
if not args.allow_download:
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
|
|
def _record(result):
|
|
if args.json_result:
|
|
with open(args.json_result, "w") as f:
|
|
json.dump(result, f)
|
|
print(f"[{result['status']}] {result['arch']}" + (f" - {result.get('error', '')}" if result["status"] != "PASS" else f" {json.dumps({k: v for k, v in result.items() if k not in ('arch', 'status')})}"))
|
|
if result["status"] == "FAIL":
|
|
sys.exit(1)
|
|
|
|
if args.swap:
|
|
archs = [a.strip() for a in args.swap.split(",") if a.strip()]
|
|
try:
|
|
result = run_swap(archs, args.device)
|
|
except BaseException as err:
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
result = {"arch": args.swap, "status": classify_error(err), "error": f"{type(err).__name__}: {err}"}
|
|
_record(result)
|
|
return
|
|
|
|
if args.arch:
|
|
if args.arch not in MODEL_TESTS:
|
|
raise SystemExit(f"arch {args.arch!r} is not registered; --list shows options")
|
|
try:
|
|
result = run_one(args.arch, args.device, cancel=args.cancel)
|
|
except BaseException as err:
|
|
status = classify_error(err)
|
|
if status == "FAIL":
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
result = {"arch": args.arch, "status": status, "error": f"{type(err).__name__}: {err}"}
|
|
_record(result)
|
|
return
|
|
|
|
if not args.all:
|
|
parser.print_help()
|
|
return
|
|
|
|
results = []
|
|
total = len(MODEL_TESTS)
|
|
t0 = time.perf_counter()
|
|
for i, arch in enumerate(MODEL_TESTS, start=1):
|
|
print(f"\n===== {arch} ({i}/{total}) =====", flush=True)
|
|
result_path = os.path.join(OUTPUT_ROOT, f".{arch.replace('/', '_')}.result.json")
|
|
os.makedirs(OUTPUT_ROOT, exist_ok=True)
|
|
cmd = [sys.executable, os.path.abspath(__file__), "--arch", arch, "--device", args.device, "--json-result", result_path]
|
|
if args.allow_download:
|
|
cmd.append("--allow-download")
|
|
if args.quantize:
|
|
cmd.append("--quantize")
|
|
proc = subprocess.run(cmd, cwd=TOOLKIT_ROOT)
|
|
if os.path.exists(result_path):
|
|
with open(result_path) as f:
|
|
results.append(json.load(f))
|
|
os.remove(result_path)
|
|
else:
|
|
results.append({"arch": arch, "status": "FAIL", "error": f"subprocess died (exit {proc.returncode})"})
|
|
tally = {"PASS": 0, "FAIL": 0, "SKIP": 0}
|
|
for r in results:
|
|
tally[r["status"]] = tally.get(r["status"], 0) + 1
|
|
elapsed = time.perf_counter() - t0
|
|
print(f">>> progress: {i}/{total} - {tally['PASS']} pass, {tally['FAIL']} fail, {tally['SKIP']} skip - elapsed {elapsed / 60:.0f}m", flush=True)
|
|
|
|
print("\n===== summary =====")
|
|
counts = {"PASS": 0, "FAIL": 0, "SKIP": 0}
|
|
for r in results:
|
|
counts[r["status"]] = counts.get(r["status"], 0) + 1
|
|
line = f"[{r['status']}] {r['arch']}"
|
|
if r["status"] != "PASS":
|
|
line += f" - {r.get('error', '')[:160]}"
|
|
print(line)
|
|
print(f"\n{counts['PASS']} passed, {counts['FAIL']} failed, {counts['SKIP']} skipped")
|
|
with open(os.path.join(OUTPUT_ROOT, "report.json"), "w") as f:
|
|
json.dump(results, f, indent=2)
|
|
if counts["FAIL"]:
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|