| |
| |
| |
| """Fast checks of the wheel, CPU only (CUDA_VISIBLE_DEVICES is set empty for every step); about a minute. |
| |
| python3 tools/wheel_check.py [--wheel WHEEL] [--data-from DIR] [--serial] [--keep-venv] |
| |
| (a) This tree, in the project venv (.venv) and the canonical environment, as `make` runs the tests: the existing tests |
| that cover what the wheel bundles. OpenBLAS: the row-parallel VBx products against numpy's own (shapes up to |
| 5,000 rows; the 20,774-row ones take 16 s and are left to `make test-fast`) and the engine's bits with every BLAS |
| thread variable at 64. libsndfile: decoding against soundfile over every format, 2,750 damaged files, and the |
| release's load_audio. The assets and everything after the networks: 40 cached records' RTTMs against the |
| canonical oracle, through diarize_record and through the engine. |
| (b) The wheel in a fresh venv: Python 3.12, torch 2.11 (CPU wheels), numpy 2.x and the wheel; none of |
| env/canonical.sh's variables and no ANUVANI_DIAR_* variable set, so the package alone finds its libraries. |
| torch and numpy are imported first (the Diarizer loads torch before the engine), then anuvani_diar. The private |
| OpenBLAS must run 1 thread with core Zen, from the package, without starting a thread; two golden FLACs must |
| decode to the reference environment's samples (sha256 of the float32 data, computed there with soundfile); the |
| 1 h, 3 h and 6 h files' raw-audio oracle records must give RTTMs whose sha256 equal assets/long_golden.json. |
| |
| The data (cached records, oracle records, test audio) are read from this tree or, with --data-from, from another |
| checkout of this project; in a git worktree, which has none of the untracked data, the main checkout is used by |
| default. Part (a) then runs in a private mount namespace (unshare) where the other checkout's folders are bound |
| read-only over this tree's, so the tests run unchanged and nothing is written to either tree. |
| |
| (a) and (b) run at the same time unless --serial. The result is printed and written next to the wheel |
| (<wheel>.check.json). Exit status 0 only if every check passed. |
| """ |
| import argparse |
| import hashlib |
| import json |
| import os |
| import re |
| import shlex |
| import shutil |
| import subprocess |
| import sys |
| import tempfile |
| import time |
| from pathlib import Path |
|
|
| PROJ = Path(__file__).resolve().parents[1] |
| |
| REFERENCE_PYTHON = Path(os.environ.get("ANUVANI_DIAR_REFERENCE_PYTHON") or sys.executable) |
| TORCH = "torch==2.11.*" |
| TORCH_INDEX = "https://download.pytorch.org/whl/cpu" |
| NUMPY = "numpy>=2,<3" |
| LONG = ("long_1h", "long_3h", "long_6h") |
| FLACS = ("bengali_009", "telugu_020") |
| TESTS = [ |
| "tests/test_audio_decode.py::test_synthetic_matrix_is_broad", |
| "tests/test_audio_decode.py::test_sf_read_matches_soundfile", |
| "tests/test_audio_decode.py::test_load_audio_matches_release", |
| "tests/test_audio_decode.py::test_damaged_files_decode_as_soundfile", |
| "tests/test_audio_decode.py::test_damaged_files_load_as_the_release", |
| "tests/test_audio_decode.py::test_memory_error_text_matches_numpy", |
| "tests/test_npblas_par.py", |
| "tests/test_engine.py::test_engine_bits_do_not_depend_on_openblas_env", |
| "tests/test_engine.py::test_engine_records_sample", |
| "tests/test_engine.py::test_a_low_sample_rate_cannot_resample_past_the_limit", |
| "tests/test_identity_cached.py::test_identity_cached_sample", |
| ] |
| TESTS_K = "not 20774" |
|
|
| |
| VENV_CHECK = r''' |
| import hashlib, json, os, sys, sysconfig, time |
| spec = json.loads(sys.argv[1]) |
| out = dict(problems=[]) |
| def bad(msg): |
| out["problems"].append(msg) |
| def tasks(): |
| return len(os.listdir("/proc/self/task")) |
| set_vars = sorted(v for v in spec["must_be_unset"] if v in os.environ) |
| if set_vars: |
| bad(f"variables set: {set_vars}") |
| import numpy as np |
| import torch |
| out["python"] = sys.version.split()[0] |
| out["numpy"] = np.__version__ |
| out["torch"] = torch.__version__ |
| before = tasks() |
| t = time.perf_counter() |
| import anuvani_diar |
| from anuvani_diar import _core |
| blas = _core.openblas_info() |
| out["openblas"] = dict(blas, load_s=round(time.perf_counter() - t, 3), threads_started=tasks() - before) |
| site = sysconfig.get_paths()["platlib"] |
| out["package_dir"] = os.path.dirname(anuvani_diar.__file__) |
| if not out["package_dir"].startswith(site): |
| bad(f"anuvani_diar imported from {out['package_dir']}, not the venv") |
| if not blas["path"].startswith(out["package_dir"]): |
| bad(f"OpenBLAS loaded from {blas['path']}") |
| if blas["num_threads"] != 1 or blas["corename"] != "Zen" or out["openblas"]["threads_started"] != 0: |
| bad(f"OpenBLAS: {blas['num_threads']} threads, core {blas['corename']}, {out['openblas']['threads_started']} threads started") |
| snd = _core.sndfile_info() |
| out["sndfile"] = snd |
| if not (snd["checked"] and snd["version"] == "libsndfile-1.2.2" and snd["path"].startswith(out["package_dir"])): |
| bad(f"libsndfile: {snd}") |
| out["decoded"] = {} |
| for sid, f in spec["flacs"].items(): |
| data, sr = _core.sf_read_f32(f["path"]) |
| if data is None: |
| bad(f"{sid}: {sr}") |
| continue |
| got = dict(sha256=hashlib.sha256(np.ascontiguousarray(data).tobytes()).hexdigest(), samples=int(data.shape[0]), samplerate=int(sr)) |
| out["decoded"][sid] = dict(got, identical=got == f["reference"]) |
| if got != f["reference"]: |
| bad(f"{sid}: decoded {got}, reference {f['reference']}") |
| out["long"] = {} |
| for sid, r in spec["records"].items(): |
| d = dict(np.load(r["npz"])) |
| t = time.perf_counter() |
| text = _core.diarize_record(sid, np.ascontiguousarray(d["P"]).view(np.uint16), d["masks"], d["emb"], |
| list(map(float, d["chunk_sw"])), list(map(float, d["frames_sw"])), float(d["duration_s"]), |
| spec["clustering"], spec["decoder"]) |
| got = hashlib.sha256(text.encode()).hexdigest() |
| out["long"][sid] = dict(chunks=int(d["P"].shape[0]), wall_s=round(time.perf_counter() - t, 2), rttm_sha256=got, |
| identical=got == r["rttm"]) |
| if got != r["rttm"]: |
| bad(f"{sid}: RTTM sha256 {got}, golden {r['rttm']}") |
| after = _core.openblas_info() |
| if after["num_threads"] != 1: |
| bad(f"OpenBLAS runs {after['num_threads']} threads after the records") |
| print(json.dumps(out)) |
| ''' |
|
|
| REFERENCE_DECODE = r''' |
| import hashlib, json, sys |
| import numpy as np, soundfile as sf |
| out = {} |
| for sid, path in json.loads(sys.argv[1]).items(): |
| data, sr = sf.read(path, dtype="float32", always_2d=False) |
| out[sid] = dict(sha256=hashlib.sha256(np.ascontiguousarray(data).tobytes()).hexdigest(), samples=int(data.shape[0]), |
| samplerate=int(sr)) |
| print(json.dumps(dict(out, soundfile=sf.__version__, libsndfile=sf.__libsndfile_version__))) |
| ''' |
|
|
|
|
| def git(*args): |
| return subprocess.run(["git", "-C", str(PROJ), *args], capture_output=True, text=True, check=True).stdout.strip() |
|
|
|
|
| def main_checkout(): |
| """This project in the main checkout when this tree is a git worktree, else None.""" |
| try: |
| common = Path(git("rev-parse", "--path-format=absolute", "--git-common-dir")) |
| top = Path(git("rev-parse", "--show-toplevel")) |
| except (subprocess.CalledProcessError, FileNotFoundError): |
| return None |
| if common.parent == top: |
| return None |
| return common.parent / PROJ.relative_to(top) |
|
|
|
|
| def canonical_vars(): |
| return re.findall(r"^export (\w+)=", (PROJ / "env/canonical.sh").read_text(), re.M) |
|
|
|
|
| def run_tests(data, log): |
| """Part (a); returns (ok, summary line, seconds).""" |
| venv = PROJ / ".venv" |
| cmd = (f"cd {shlex.quote(str(PROJ))} && source env/canonical.sh && unset CONDA_PREFIX CONDA_DEFAULT_ENV && " |
| f"export VIRTUAL_ENV={shlex.quote(str(venv))} PATH={shlex.quote(str(venv / 'bin'))}:$PATH && " |
| f"python -s -m pytest -q {' '.join(map(shlex.quote, TESTS))} -k {shlex.quote(TESTS_K)} -p no:cacheprovider") |
| if data != PROJ: |
| binds = [(data.parents[1] / "diar", PROJ.parents[1] / "diar"), |
| (data.parent / "anuvani_diar_v1", PROJ.parent / "anuvani_diar_v1"), |
| (data / "artifacts", PROJ / "artifacts")] |
| mounts = " && ".join(f"mount --bind {shlex.quote(str(s))} {shlex.quote(str(d))} && " |
| f"mount -o remount,bind,ro {shlex.quote(str(d))}" for s, d in binds) |
| argv = ["unshare", "--user", "--map-root-user", "--mount", "bash", "-c", f"set -e; {mounts}; {cmd}"] |
| else: |
| argv = ["bash", "-c", cmd] |
| t = time.perf_counter() |
| with open(log, "w") as f: |
| p = subprocess.Popen(argv, stdout=f, stderr=subprocess.STDOUT) |
| return p, t |
|
|
|
|
| def finish_tests(p, t, log): |
| p.wait() |
| secs = time.perf_counter() - t |
| text = Path(log).read_text() |
| lines = [ln for ln in text.splitlines() if re.search(r"\d+ (passed|failed|error)", ln)] |
| summary = lines[-1].strip("= ") if lines else text[-2000:] |
| return dict(ok=p.returncode == 0 and "passed" in summary and "failed" not in summary, summary=summary, |
| seconds=round(secs, 1), log=str(log)) |
|
|
|
|
| def check_wheel(wheel, data, tmp, keep): |
| """Part (b).""" |
| t0 = time.perf_counter() |
| res = dict(problems=[]) |
| |
| |
| bundle = Path(os.environ.get("ANUVANI_DIAR_BUNDLE") or next( |
| b for b in (data / "artifacts/bundle", data.parent / "anuvani_diar_v1/bundle") if (b / "MANIFEST.json").is_file())) |
| res["bundle"] = str(bundle) |
| model = json.loads((bundle / "MANIFEST.json").read_text())["modes"]["indic"] |
| golden_long = json.loads((PROJ / "assets/long_golden.json").read_text())["files"] |
| inputs = {i["sid"]: i for i in json.loads((data / "artifacts/golden12_inputs.json").read_text())} |
| flacs = {} |
| for sid in FLACS: |
| path = Path(inputs[sid]["path"]) |
| if hashlib.sha256(path.read_bytes()).hexdigest() != inputs[sid]["sha256"]: |
| res["problems"].append(f"{path}: the file is not the golden input") |
| flacs[sid] = str(path) |
| ref = subprocess.run([str(REFERENCE_PYTHON), "-I", "-c", REFERENCE_DECODE, json.dumps(flacs)], capture_output=True, |
| text=True, check=True, env=dict(os.environ, OPENBLAS_NUM_THREADS="1")) |
| reference = json.loads(ref.stdout) |
| res["reference_decoder"] = dict(soundfile=reference.pop("soundfile"), libsndfile=reference.pop("libsndfile"), |
| python=str(REFERENCE_PYTHON)) |
| venv = tmp / "venv" |
| t = time.perf_counter() |
| subprocess.run(["uv", "venv", "-q", "--python", "3.12", str(venv)], check=True) |
| py = str(venv / "bin" / "python") |
| subprocess.run(["uv", "pip", "install", "-q", "--python", py, "--index-url", TORCH_INDEX, TORCH], check=True) |
| subprocess.run(["uv", "pip", "install", "-q", "--python", py, NUMPY], check=True) |
| subprocess.run(["uv", "pip", "install", "-q", "--python", py, str(wheel)], check=True) |
| res["venv_s"] = round(time.perf_counter() - t, 1) |
| spec = dict( |
| must_be_unset=sorted(set(canonical_vars()) | {"ANUVANI_DIAR_OPENBLAS_DIR", "ANUVANI_DIAR_SNDFILE", |
| "ANUVANI_DIAR_ASSETS", "ANUVANI_DIAR_SNDFILE_UNCHECKED"}), |
| flacs={sid: dict(path=flacs[sid], reference=reference[sid]) for sid in FLACS}, |
| records={sid: dict(npz=str(data / f"artifacts/oracle_audio/long/records/{sid}.npz"), rttm=golden_long[sid]["rttm"]) |
| for sid in LONG}, |
| clustering=model["clustering"], decoder=str(bundle / model["decoder"])) |
| env = dict(PATH="/usr/bin:/bin", HOME=os.environ.get("HOME", "/tmp"), CUDA_VISIBLE_DEVICES="", LANG="C.UTF-8") |
| t = time.perf_counter() |
| r = subprocess.run([py, "-I", "-c", VENV_CHECK, json.dumps(spec)], capture_output=True, text=True, env=env, cwd=tmp) |
| res["run_s"] = round(time.perf_counter() - t, 1) |
| if r.returncode != 0: |
| res["problems"].append(f"the check in the venv failed: {r.stderr[-3000:]}") |
| else: |
| out = json.loads(r.stdout.strip().splitlines()[-1]) |
| res["problems"] += out.pop("problems") |
| res.update(out) |
| res["stderr"] = r.stderr[-2000:] |
| if keep: |
| res["venv"] = str(venv) |
| res["seconds"] = round(time.perf_counter() - t0, 1) |
| res["ok"] = not res["problems"] |
| return res |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--wheel", type=Path, help="default: the newest artifacts/wheels/*.whl") |
| ap.add_argument("--data-from", type=Path, help="another checkout of this project (anuvani_diar_rs) holding the data") |
| ap.add_argument("--serial", action="store_true", help="run (a), then (b)") |
| ap.add_argument("--keep-venv", action="store_true") |
| a = ap.parse_args(argv) |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" |
| wheel = a.wheel or max((PROJ / "artifacts/wheels").glob("*.whl"), key=lambda p: p.stat().st_mtime) |
| data = (a.data_from or (None if (PROJ / "artifacts/dev_data").is_dir() else main_checkout()) |
| or PROJ).resolve() |
| print(f"wheel {wheel.name}\ndata {data}" + ("" if data == PROJ else " (bound read-only for the tests)"), flush=True) |
| t0 = time.perf_counter() |
| tmp = Path(tempfile.mkdtemp(prefix="anuvani_wheel_check_")) |
| log = tmp / "tests.log" |
| try: |
| (PROJ / "artifacts").mkdir(exist_ok=True) |
| p, t = run_tests(data, log) |
| if a.serial: |
| tests = finish_tests(p, t, log) |
| venv = check_wheel(wheel, data, tmp, a.keep_venv) |
| else: |
| venv = check_wheel(wheel, data, tmp, a.keep_venv) |
| tests = finish_tests(p, t, log) |
| if not tests["ok"]: |
| print(Path(log).read_text()[-4000:]) |
| finally: |
| if not a.keep_venv: |
| shutil.rmtree(tmp / "venv", ignore_errors=True) |
| total = round(time.perf_counter() - t0, 1) |
| result = dict(wheel=wheel.name, wheel_sha256=hashlib.sha256(wheel.read_bytes()).hexdigest(), data=str(data), |
| commit=git("rev-parse", "HEAD"), host=os.uname().nodename, cpus=os.cpu_count(), |
| mode="serial" if a.serial else "parallel", seconds=total, tests=tests, wheel_in_fresh_venv=venv, |
| ok=tests["ok"] and venv["ok"]) |
| (wheel.parent / f"{wheel.name}.check.json").write_text(json.dumps(result, indent=1) + "\n") |
| print(f"(a) tests in .venv: {'PASS' if tests['ok'] else 'FAIL'} {tests['summary']} ({tests['seconds']} s)") |
| print(f"(b) wheel in a fresh venv: {'PASS' if venv['ok'] else 'FAIL'} (venv {venv.get('venv_s')} s, checks " |
| f"{venv.get('run_s')} s; Python {venv.get('python')}, torch {venv.get('torch')}, numpy {venv.get('numpy')})") |
| if "openblas" in venv: |
| b = venv["openblas"] |
| print(f" OpenBLAS {b['num_threads']} thread, core {b['corename']}, forced {b['coretype_forced']}, " |
| f"{b['threads_started']} threads started, from {b['path']}") |
| for sid, d in venv.get("decoded", {}).items(): |
| print(f" {sid}: {d['samples']} samples at {d['samplerate']} Hz, sha256 {d['sha256'][:16]}..., " |
| f"{'identical to' if d['identical'] else 'DIFFERENT from'} the reference environment's") |
| for sid, d in venv.get("long", {}).items(): |
| print(f" {sid}: RTTM sha256 {d['rttm_sha256'][:16]}... {'identical' if d['identical'] else 'DIFFERENT'} " |
| f"({d['chunks']} chunks, {d['wall_s']} s)") |
| for msg in venv["problems"]: |
| print(f" PROBLEM: {msg}") |
| print(f"{'PASS' if result['ok'] else 'FAIL'} in {total} s; record: {wheel.parent / (wheel.name + '.check.json')}") |
| return 0 if result["ok"] else 1 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
