Spaces:
Running on Zero
Running on Zero
| """maya3.Engine latency/throughput bench (same protocol as baseline_fastbatch.py, through the public submit() path). | |
| B=1 : one request at a time; TTFA = submit -> first PCM chunk delivered to the callback (prompt build, prefix-KV copy, | |
| inline prefill, first frame, Mimi, D2H all included); RTF = wall / audio. | |
| N>1 : closed loop, N requests in flight; a finished request immediately submits the next; reps*N requests total. | |
| TTFA = submit -> first chunk (queue wait included). Tick p50/p99 from the engine. | |
| Usage: bench_engine.py <model.pt> <out.json> [--seats 1,8,32,64,100] [--depth fp16] [--no-guards] | |
| """ | |
| import argparse, json, os, sys, threading, time | |
| import numpy as np, torch | |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) | |
| from maya3.engine import Engine, EngineConfig, SR | |
| from maya3.bench import common as C | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("ckpt"); ap.add_argument("out") | |
| ap.add_argument("--seats", default="1,8,32,64,100"); ap.add_argument("--depth", default="fp16") | |
| ap.add_argument("--no-guards", action="store_true"); ap.add_argument("--b1-n", type=int, default=40) | |
| ap.add_argument("--reps", type=int, default=3); ap.add_argument("--capacity", type=int, default=2304) | |
| ap.add_argument("--max-inline-rows", type=int, default=384); ap.add_argument("--max-admit", type=int, default=8) | |
| ap.add_argument("--label", default="maya3"); ap.add_argument("--cb0-rep", type=float, default=None); ap.add_argument("--ramp-s", type=float, default=2.0) | |
| A = ap.parse_args() | |
| q = lambda x, p: round(float(np.quantile(x, p)), 4) if len(x) else None | |
| rows = C.rows() | |
| pack = C.load_pack() | |
| res = {"ckpt": A.ckpt, "label": A.label, "depth": A.depth, "guards": not A.no_guards, "gpu": torch.cuda.get_device_name(), "runs": {}} | |
| E = None | |
| for N in [int(x) for x in A.seats.split(",")]: | |
| cfg = EngineConfig(max_seats=N, capacity=A.capacity, depth_dtype=A.depth, max_inline_rows=A.max_inline_rows, | |
| max_admit_per_tick=A.max_admit, free_reference_depth=False, | |
| **({} if A.cb0_rep is None else {"cb0_rep": A.cb0_rep})) # rebuilt per N below | |
| cfg.guards.enabled = not A.no_guards | |
| if E is None: | |
| E = Engine(A.ckpt, C.BACKBONE, C.MIMI, cfg) | |
| model, mimi, registry = E.model, E.mimi, E.registry | |
| else: # rebuild the fast components only (same weights) | |
| E.cfg = cfg; E.backbone = E.depth = E.codec = None; torch.cuda.empty_cache(); E._build_fast() | |
| vmap = C.register_pack_voices(E, pack, rows) | |
| t0 = time.time(); E.warmup(); warm = time.time() - t0 | |
| for v in set(vmap.values()): | |
| E.prepare_voice(E.registry.get(v)) | |
| E.start() | |
| reqs = [(r["pieces"][0], r["lang"], vmap[r["id"]], r["mode"]) for r in rows] | |
| lock = threading.Lock() | |
| rec = {"warmup_s": round(warm, 1)} | |
| # extra warm-up through the live loop | |
| ev = threading.Event(); E.submit(reqs[0][0], reqs[0][1], reqs[0][2], on_event=lambda k, p: ev.set() if k != "audio" else None); ev.wait(60) | |
| torch.cuda.reset_peak_memory_stats(); E.tick_ms.clear() | |
| if N == 1: | |
| ttfa, rtf, audio = {"plain": [], "icl_spk": []}, {"plain": [], "icl_spk": []}, 0.0 | |
| tw = time.time() | |
| for i, (text, lang, vid, mode) in enumerate(reqs[:A.b1_n]): | |
| done = threading.Event(); info = {} | |
| t_sub = time.perf_counter() | |
| first = [] | |
| def cb(k, p, first=first, info=info, done=done): | |
| if k == "audio": | |
| if not first: first.append(time.perf_counter()) | |
| else: | |
| info.update(p if isinstance(p, dict) else {}); info["kind"] = k; done.set() | |
| E.submit(text, lang, vid, on_event=cb, split=False) | |
| done.wait(120) | |
| if info.get("kind") != "end" or not first: continue | |
| ttfa[mode].append(first[0] - t_sub); rtf[mode].append(info["wall_s"] / max(info["audio_s"], 1e-3)); audio += info["audio_s"] | |
| wall = time.time() - tw | |
| rec.update({"n_req": A.b1_n, "audio_s_per_s": round(audio / wall, 2)}) | |
| for m in ttfa: | |
| rec[m] = {"n": len(ttfa[m]), "ttfa_ms_p50": round(q(ttfa[m], .5) * 1e3, 1), "ttfa_ms_p90": round(q(ttfa[m], .9) * 1e3, 1), | |
| "ttfa_ms_p99": round(q(ttfa[m], .99) * 1e3, 1), "rtf_p50": q(rtf[m], .5), "rtf_p90": q(rtf[m], .9)} | |
| else: | |
| total = A.reps * N | |
| state = {"sent": 0, "done": 0, "audio": 0.0} | |
| tt, rt, reasons, steady = [], [], {}, [] | |
| alldone = threading.Event() | |
| def send(): | |
| with lock: | |
| if state["sent"] >= total: return | |
| i = state["sent"]; state["sent"] += 1 | |
| text, lang, vid, mode = reqs[i % len(reqs)] | |
| t_sub = time.perf_counter(); first = [] | |
| def cb(k, p): | |
| if k == "audio": | |
| if not first: first.append(time.perf_counter()) | |
| return | |
| with lock: | |
| state["done"] += 1 | |
| if k == "end": | |
| tt.append(first[0] - t_sub if first else None); rt.append(p["wall_s"] / max(p["audio_s"], 1e-3)) | |
| if t_sub - tw0[0] > A.ramp_s + 1.0 and first: steady.append((first[0] - t_sub, rt[-1])) | |
| state["audio"] += p["audio_s"]; reasons[p["reason"]] = reasons.get(p["reason"], 0) + 1 | |
| else: | |
| reasons[p.get("code")] = reasons.get(p.get("code"), 0) + 1 | |
| fin = state["done"] >= total | |
| if fin: alldone.set() | |
| else: send() | |
| E.submit(text, lang, vid, on_event=cb, split=False) | |
| tw = time.time(); tw0 = [time.perf_counter()] | |
| for i in range(N): threading.Timer(A.ramp_s * i / N, send).start() | |
| alldone.wait(3600) | |
| wall = time.time() - tw | |
| tt = [x for x in tt if x is not None] | |
| tk = sorted(E.tick_ms) | |
| rec.update({"n_req": total, "wall_s": round(wall, 1), "audio_s": round(state["audio"], 1), | |
| "audio_s_per_s": round(state["audio"] / wall, 2), "ttfa_ms_p50": round(q(tt, .5) * 1e3, 1), | |
| "ttfa_ms_p99": round(q(tt, .99) * 1e3, 1), "rtf_p50": q(rt, .5), "rtf_p99": q(rt, .99), "reasons": reasons, "ramp_s": A.ramp_s, | |
| "steady": {"n": len(steady), "ttfa_ms_p50": round(q([a for a, _ in steady], .5) * 1e3, 1) if steady else None, | |
| "ttfa_ms_p99": round(q([a for a, _ in steady], .99) * 1e3, 1) if steady else None, | |
| "rtf_p99": q([b for _, b in steady], .99) if steady else None}, | |
| "tick_ms_p50": round(tk[len(tk) // 2], 2), "tick_ms_p99": round(tk[int(len(tk) * .99)], 2)}) | |
| rec["peak_alloc_gb"] = round(torch.cuda.max_memory_allocated() / 1e9, 2) | |
| rec["phase_ms_last"] = E.phase_ms | |
| rec["stats"] = dict(E.stats) | |
| E.close(); E._thread = None; E._stop = False | |
| res["runs"][f"N{N}"] = rec | |
| print(f"N={N}", json.dumps(rec), flush=True) | |
| json.dump(res, open(A.out, "w"), indent=1) | |
| print("BENCH_DONE") | |