Buckets:
| """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() | |
Xet Storage Details
- Size:
- 16.1 kB
- Xet hash:
- 2a5f5766fce5437a01f9b366ca2e82dbe91a8c5916fd9eded344153a87692413
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.