Initial inference engine build out

testing/test_inference_engine.py

@@ -0,0 +1,384 @@
+"""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()