File size: 4,454 Bytes
63dd946 | 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 | # Copyright (c) 2026 The Anuvani-Diarization authors
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Identity check through the serving stack (Triton, the Rust gateway, the Python client) against the raw-audio oracle,
byte for byte: every RTTM, every error, and every front-end record the gateway dumps. Jobs arrive concurrently in a
shuffled order; most name a file on the server, the rest are uploaded through the gateway's spool.
Run through triton/identity_check.sh (it starts Triton and the gateway and sets GATEWAY_URL; its variants set
IDENTITY_SEED for another arrival order and RESTART_AT for a Triton restart in the middle of the run).
"""
import json
import os
import random
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import numpy as np
import pytest
from conftest import PROJ
from test_identity_audio import INPUTS, identity_sample, oracle_log
ORACLE = PROJ / "artifacts/oracle_audio"
RECORDS = PROJ / "artifacts/identity_serving/records"
pytestmark = pytest.mark.skipif("GATEWAY_URL" not in os.environ, reason="needs triton/identity_check.sh")
def run_through_gateway(keys, uploads=20, seed=7):
from anuvani_diar import Client
c = Client(os.environ["GATEWAY_URL"], timeout=3600)
order = list(keys)
random.Random(seed).shuffle(order)
# uploads only for files that succeed: a failing upload's message names the spool file, not the client's path
log = oracle_log()
upload = set([k for k in order if log[k]["ok"]][:uploads])
def one(k):
it = INPUTS[k]
job = c.submit(it["path"], uri=it["sid"], server_path=k not in upload)
return k, c.wait(job, poll=1.0)
with ThreadPoolExecutor(64) as ex:
return dict(ex.map(one, order))
def compare(keys, got):
log = oracle_log()
bad = []
for k in keys:
s, sid = k.split("/", 1)
want, r = log[k], got[k]
if not want["ok"]:
if r["error"] != want["error"]:
bad.append((k, "error", want["error"], r["error"]))
continue
# the release's own post-processing raised
want_err = want["indic"].get("error")
if r["error"] != want_err:
bad.append((k, "error", want_err, r["error"]))
continue
if want_err is None and (r["status"] != "done" or (r["rttm"] or "").encode() != (ORACLE / s / "indic" / f"{sid}.rttm").read_bytes()):
bad.append((k, r["status"], "rttm"))
d = np.load(ORACLE / s / "records" / f"{sid}.npz")
P = np.fromfile(RECORDS / f"{sid}.P", dtype="<u2")
masks = np.fromfile(RECORDS / f"{sid}.masks", dtype=np.uint8)
emb = np.fromfile(RECORDS / f"{sid}.emb", dtype="<f4")
if (P.tobytes() != np.ascontiguousarray(d["P"]).view(np.uint16).tobytes() or masks.tobytes() != d["masks"].tobytes()
or emb.tobytes() != d["emb"].tobytes()):
bad.append((k, "record"))
return bad
def gateway_metric(name):
with urllib.request.urlopen(os.environ["GATEWAY_URL"] + "/metrics") as r:
for line in r.read().decode().splitlines():
if line.startswith(name + " "):
return float(line.split()[1])
SEED = int(os.environ.get("IDENTITY_SEED", "7"))
def quick_sample():
"""check-exact's set (tools/check_exact.py): the golden 12, every robustness input, the 1 h file, 20 sampled files."""
import records
keys = [f"bench/{i['sid']}" for i in json.loads((PROJ / "artifacts/golden12_inputs.json").read_text())]
keys += [k for k, i in INPUTS.items() if i["set"] == "robust"] + ["long/long_1h"]
keys += [f"{s}/{p.stem}" for s, p in records.sample(20)]
return [k for k in dict.fromkeys(keys) if k in INPUTS]
@pytest.mark.gpu
def test_identity_serving_sample():
keys = quick_sample() if os.environ.get("IDENTITY_QUICK") else identity_sample()
got = run_through_gateway(keys, seed=SEED)
assert len(got) == len(keys)
assert compare(keys, got) == []
# the restart hit work in flight, which was sent again
if os.environ.get("RESTART_AT"):
assert gateway_metric("anuvani_diar_triton_retries_total") > 0
@pytest.mark.gpu
@pytest.mark.slow
def test_identity_serving_all():
keys = list(INPUTS)
got = run_through_gateway(keys, uploads=100, seed=SEED)
assert compare(keys, got) == []
if os.environ.get("RESTART_AT"):
assert gateway_metric("anuvani_diar_triton_retries_total") > 0
|