| |
| |
| """The engine (crate diar-engine) against the release, byte for byte, without a GPU: |
| |
| 1. cached front-end records through the engine's post-processing pool vs the canonical cache oracle; |
| 2. audio -> RTTM with the stand-in networks of fake_nn.py in both the engine and the release's own `analyse` and |
| `indic_mode` (run in a subprocess): records, the RTTM and every error must be identical, with and without the |
| repacked voiceprint trunk, and every network call must be one of the release's batch shapes. |
| """ |
| import json |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import pytest |
|
|
| import fake_nn |
| import records |
| from conftest import PROJ, REL |
|
|
| MODEL = json.loads((REL / "bundle/MANIFEST.json").read_text())["modes"]["indic"] |
| DECODER = str(REL / "bundle" / MODEL["decoder"]) |
| ORACLE = json.loads((PROJ / "oracle/cache_canonical.json").read_text())["files"] |
| CACHE_RTTM = PROJ / "artifacts/oracle_cache/canonical" |
| WIN, STEP = 256000, 25600 |
|
|
|
|
| def make_engine(core, nn, **kw): |
| return core.Engine(nn, clustering=MODEL["clustering"], decoder_path=DECODER, **kw) |
|
|
|
|
| def record_jobs(items): |
| jobs = [] |
| for set_name, path in items: |
| d = records.load(path) |
| jobs.append((f"{set_name}/{path.stem}", np.ascontiguousarray(d["P"]).view(np.uint16), d["masks"], d["emb"], |
| list(map(float, d["frames_sw"])), float(d["duration_s"]))) |
| return jobs |
|
|
|
|
| def check_records_against_cache(core, items, post_threads): |
| eng = make_engine(core, fake_nn.FakeNN(), post_threads=post_threads) |
| got = {r["uri"]: r for r in eng.run_records(record_jobs(items))} |
| assert len(got) == len(items) |
| bad = [] |
| for set_name, path in items: |
| key = f"{set_name}/{path.stem}" |
| r = got[key] |
| |
| want = (CACHE_RTTM / set_name / "indic" / f"{path.stem}.rttm").read_bytes() |
| if r["error"] is not None or r["rttm"].replace(f"SPEAKER {key} ", f"SPEAKER {path.stem} ").encode() != want: |
| bad.append(key) |
| return bad |
|
|
|
|
| def test_engine_records_sample(core): |
| assert check_records_against_cache(core, records.sample(40), post_threads=8) == [] |
|
|
|
|
| @pytest.mark.slow |
| def test_engine_records_all(core): |
| assert check_records_against_cache(core, records.all_records(), post_threads=48) == [] |
|
|
|
|
| |
|
|
| @pytest.mark.parametrize("repack", [True, False]) |
| def test_engine_audio_with_stand_in_networks(core, fake_reference, repack): |
| items, ref = fake_reference |
| log = json.loads((ref / "log.json").read_text()) |
| nn = fake_nn.FakeNN(devices=(0, 0, 1, 1)) |
| eng = make_engine(core, nn, repack=repack, keep_records=True, decode_threads=3, post_threads=6, max_ready=3) |
| jobs = [(it["sid"], str(Path(it["path"])), it["path"]) for it in items] |
| got = {r["uri"]: r for r in eng.run(jobs)} |
| assert len(got) == len(items) |
| bad = [] |
| for it in items: |
| sid, r, want = it["sid"], got[it["sid"]], log[it["sid"]] |
| if "error" in want or r["error"]: |
| if r["error"] != want.get("error"): |
| bad.append((sid, "error", want.get("error"), r["error"])) |
| continue |
| d = np.load(ref / f"{sid}.npz") |
| rec = r["record"] |
| same = (rec["P_bits"].tobytes() == np.ascontiguousarray(d["P"]).view(np.uint16).tobytes() |
| and rec["P_bits"].shape == d["P"].shape |
| and rec["masks"].tobytes() == d["masks"].tobytes() and rec["masks"].shape == d["masks"].shape |
| and rec["emb"].tobytes() == d["emb"].tobytes() and rec["emb"].shape == d["emb"].shape |
| and rec["chunk_sw"].tobytes() == d["chunk_sw"].tobytes() |
| and rec["frames_sw"].tobytes() == d["frames_sw"].tobytes() |
| and np.float64(rec["duration_s"]).tobytes() == d["duration_s"].tobytes()) |
| if not same: |
| bad.append((sid, "record")) |
| if r["rttm"] is None or r["rttm"].encode() != (ref / f"{sid}.rttm").read_bytes(): |
| bad.append((sid, "rttm")) |
| assert not bad, bad |
| |
| for call in nn.calls: |
| if call[0] == "segment": |
| assert call[1][1:] == (1, WIN) and 1 <= call[1][0] <= 8 |
| elif call[0] == "embed": |
| assert call[1][1:] == (1, WIN) and call[1][0] in ((4,) if repack else (4, 8)) and call[2] == (call[1][0], 799) |
| else: |
| assert repack and call[1] == (8, 1, WIN) and call[2][1] == 8 and call[3][1:] == (8, 799) |
|
|
|
|
| |
|
|
| class SlowFake(fake_nn.FakeNN): |
| def segment(self, slot, x): |
| import time |
| time.sleep(0.05) |
| return super().segment(slot, x) |
|
|
|
|
| class FailingFake(fake_nn.FakeNN): |
| def segment(self, slot, x): |
| |
| if x[0, 0, :1000].sum() > 0: |
| raise RuntimeError("CUDA out of memory (simulated)") |
| return super().segment(slot, x) |
|
|
|
|
| def test_engine_stops_when_the_iterator_is_dropped(core, fake_reference): |
| import time |
| items, _ = fake_reference |
| eng = make_engine(core, SlowFake(devices=(0,)), decode_threads=2, post_threads=2) |
| it = eng.run([(i["sid"], i["path"], i["path"]) for i in items] * 3) |
| first = next(it) |
| t = time.perf_counter() |
| |
| del it |
| assert time.perf_counter() - t < 30 and first["uri"] |
|
|
|
|
| def test_engine_reports_backend_errors_per_file(core, fake_reference): |
| items, ref = fake_reference |
| log = json.loads((ref / "log.json").read_text()) |
| eng = make_engine(core, FailingFake(), decode_threads=2, post_threads=2) |
| got = list(eng.run([(i["sid"], i["path"], i["path"]) for i in items])) |
| assert len(got) == len(items) |
| for r in got: |
| want = log[r["uri"]].get("error") |
| assert r["error"] == want or r["error"] == "RuntimeError: CUDA out of memory (simulated)", (r["uri"], r["error"]) |
|
|
|
|
| def test_engine_bits_do_not_depend_on_openblas_env(tmp_path): |
| """The engine pins its private OpenBLAS to one thread itself: run with OPENBLAS_NUM_THREADS=64 in a fresh process.""" |
| import os |
| import subprocess |
| code = f""" |
| import json, sys, numpy as np |
| sys.path.insert(0, {str(PROJ / 'tests')!r}) |
| import records, fake_nn |
| from anuvani_diar import _core |
| m = json.loads(open({str(REL / 'bundle/MANIFEST.json')!r}).read())["modes"]["indic"] |
| eng = _core.Engine(fake_nn.FakeNN(), clustering=m["clustering"], decoder_path={DECODER!r}, post_threads=8) |
| jobs = [] |
| for s, p in records.sample(40): |
| d = records.load(p) |
| jobs.append((s + "/" + p.stem, np.ascontiguousarray(d["P"]).view(np.uint16), d["masks"], d["emb"], |
| list(map(float, d["frames_sw"])), float(d["duration_s"]))) |
| bad = 0 |
| for r in eng.run_records(jobs): |
| s, sid = r["uri"].split("/", 1) |
| want = open({str(CACHE_RTTM)!r} + f"/{{s}}/indic/{{sid}}.rttm", "rb").read() |
| bad += r["rttm"].replace("SPEAKER " + r["uri"] + " ", "SPEAKER " + sid + " ").encode() != want |
| print("BAD", bad) |
| """ |
| env = dict(os.environ, OPENBLAS_NUM_THREADS="64", OMP_NUM_THREADS="64", MKL_NUM_THREADS="64") |
| out = subprocess.run([sys.executable, "-s", "-c", code], env=env, capture_output=True, text=True, check=True) |
| assert "BAD 0" in out.stdout, out.stdout + out.stderr |
|
|
|
|
| def test_a_low_sample_rate_cannot_resample_past_the_limit(core, tmp_path): |
| """The size limit applies to the audio after resampling too, computed from the header before anything is |
| allocated: 1,000 samples declared at 1 Hz are 4 KB decoded but 64 MB at 16 kHz.""" |
| import soundfile as sf |
| low, ok = tmp_path / "one_hertz.wav", tmp_path / "eight_khz.wav" |
| sf.write(low, np.full(1000, 0.1, np.float32), 1, subtype="FLOAT") |
| sf.write(ok, np.full(16000, 0.1, np.float32), 8000, subtype="FLOAT") |
| nn = fake_nn.FakeNN() |
| eng = make_engine(core, nn, max_audio_bytes=1 << 20, decode_threads=1, post_threads=2) |
| got = {r["uri"]: r for r in eng.run([("low", str(low), str(low)), ("ok", str(ok), str(ok))])} |
| assert got["low"]["error"] == ("RefusedError: audio of 64000000 bytes (resampled to 16 kHz) is larger than this " |
| "engine's limit of 1048576 bytes") |
| assert got["ok"]["error"] is None |
|
|
|
|
| def test_a_rate_that_does_not_reduce_cannot_build_a_huge_resampling_filter(core, tmp_path): |
| """The resampler's filter grows with the declared rate, not the file's length: 100 samples declared at 16,001 Hz |
| would build a 2 GB table. Refused from the header before anything is resampled (docs/SECURITY.md); 44.1 kHz runs.""" |
| import soundfile as sf |
| odd, ok = tmp_path / "odd.wav", tmp_path / "cd.wav" |
| sf.write(odd, np.full(100, 0.1, np.float32), 16001, subtype="FLOAT") |
| sf.write(ok, np.full(44100, 0.1, np.float32), 44100, subtype="FLOAT") |
|
|
| class Counting(fake_nn.FakeNN): |
| rates = [] |
|
|
| def resample(self, audio, sr): |
| self.rates.append(sr) |
| return super().resample(audio, sr) |
| nn = Counting() |
| eng = make_engine(core, nn, decode_threads=1, post_threads=2) |
| got = {r["uri"]: r for r in eng.run([("odd", str(odd), str(odd)), ("ok", str(ok), str(ok))])} |
| assert got["odd"]["error"] == ("RefusedError: resampling 16001 Hz audio to 16 kHz needs a filter of 2049920000 bytes, " |
| "larger than this engine's limit of 67108864 bytes") |
| assert got["ok"]["error"] is None and nn.rates == [44100] |
|
|
