Download src/tests/test_service.py from CIIL/Anuvani-Diarization: direct link, hf CLI and curl.
- Browser
- Download file 55.4 kB
-
https://huggingface.co/CIIL/Anuvani-Diarization/resolve/main/src/tests/test_service.py
- Command line
-
hf download hf://CIIL/Anuvani-Diarization/src/tests/test_service.py
-
curl -L -o test_service.py https://huggingface.co/CIIL/Anuvani-Diarization/resolve/main/src/tests/test_service.py
55.4 kB
| # Copyright (c) 2026 The Anuvani-Diarization authors | |
| # SPDX-License-Identifier: MIT OR Apache-2.0 | |
| """The FastAPI service (triton_inference/service) on CPU: the gateway's suite (tests/test_gateway.py), ported. | |
| The service runs the Rust gateway's engine and Triton backend in process (`anuvani_diar._core.TritonService`), so | |
| the same stand-in Triton (fake_triton.py) serves the stand-in networks of fake_nn.py, and every RTTM, error and | |
| front-end record must be byte-identical to the release's own `analyse` and `indic_mode` with the same networks | |
| (conftest.fake_reference). Around that, the gateway's behaviours: Triton going away and coming back, a model that | |
| stays away, model errors, a crashed model process, the service restarting under the client, admission limits, | |
| uploads, path rules, the token, cancellation, the synchronous call, the result cache, metrics, the spool, SIGTERM and | |
| 2,000 seeded random requests. Then what the service adds: the size and spool limits from Content-Length before the | |
| body, the decoded-audio budget, multipart uploads, Idempotency-Key, request ids and JSON logs, the drain on SIGTERM, | |
| OpenAPI, and the engine's calls releasing the GIL. | |
| Needs the extension (`make build`) and the service's packages (env/venv_extras.txt). | |
| """ | |
| import io | |
| import json | |
| import os | |
| import signal | |
| import subprocess | |
| import sys | |
| import threading | |
| import time | |
| import urllib.parse | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| import numpy as np | |
| import pytest | |
| from conftest import PROJ | |
| from fake_triton import FakeTriton | |
| from test_gateway import OOM, call, compare, free_port, longest, ok_items, raw_request, run_jobs, wait_for | |
| SERVICE = PROJ / "triton_inference/service" | |
| pytest.importorskip("fastapi") | |
| pytest.importorskip("uvicorn") | |
| pytest.importorskip("python_multipart") | |
| GATEWAY_METRICS = ["anuvani_diar_jobs_submitted_total", "anuvani_diar_jobs_done_total", "anuvani_diar_jobs_failed_total", | |
| "anuvani_diar_jobs_cancelled_total", "anuvani_diar_jobs_rejected_total", "anuvani_diar_cache_hits_total", | |
| "anuvani_diar_audio_seconds_done_total", "anuvani_diar_upload_bytes_total", "anuvani_diar_jobs_queued", | |
| "anuvani_diar_jobs_running", "anuvani_diar_uploads_in_progress", "anuvani_diar_spool_bytes", | |
| "anuvani_diar_triton_retries_total"] | |
| class Service: | |
| """`python -s -m anuvani_diar_service` with the gateway test's settings (two GPUs, two slots each).""" | |
| def __init__(self, triton, tmp, *extra, root="/", token=None, spool=None, wait=True, port=None): | |
| self.port = port or free_port() | |
| self.url = f"http://127.0.0.1:{self.port}" | |
| self.spool = spool or tmp / f"spool-{self.port}" | |
| self.records = tmp / f"records-{self.port}" | |
| self.log = tmp / f"service-{self.port}.log" | |
| cmd = [sys.executable, "-s", "-m", "anuvani_diar_service", "--listen", f"127.0.0.1:{self.port}", | |
| "--triton", triton.url, "--gpus", "0,1", "--slots-per-gpu", "2", "--spool", str(self.spool), | |
| "--dump-records", str(self.records), "--decode-threads", "3", "--post-threads", "4", | |
| "--drain-linger-s", "0", *extra] | |
| if root is not None: | |
| cmd += ["--allow-path-root", str(root)] | |
| if token is not None: | |
| cmd += ["--token", token] | |
| env = {k: v for k, v in os.environ.items() if k != "ANUVANI_DIAR_TOKEN"} | |
| env["PYTHONPATH"] = os.pathsep.join([str(SERVICE), *filter(None, [os.environ.get("PYTHONPATH")])]) | |
| self.proc = subprocess.Popen(cmd, stdout=open(self.log, "w"), stderr=subprocess.STDOUT, env=env) | |
| if wait: | |
| deadline = time.monotonic() + 60 | |
| while time.monotonic() < deadline and self.proc.poll() is None: | |
| try: | |
| import urllib.request | |
| urllib.request.urlopen(self.url + "/health", timeout=1) | |
| return | |
| except OSError: | |
| time.sleep(0.05) | |
| raise RuntimeError(f"the service did not start: {self.log.read_text()[-3000:]}") | |
| def stop(self): | |
| if self.proc.poll() is None: | |
| self.proc.send_signal(signal.SIGTERM) | |
| try: | |
| self.proc.wait(30) | |
| except subprocess.TimeoutExpired: | |
| self.proc.kill() | |
| self.proc.wait() | |
| def metrics(self, headers=None): | |
| _, body, _ = call(self.url + "/metrics", headers=headers) | |
| return {k: float(v) for k, v in (line.split() for line in body.decode().splitlines() if line and line[0] != "#")} | |
| def client(self, **kw): | |
| from anuvani_diar import Client | |
| return Client(self.url, timeout=120, **kw) | |
| def logs(self): | |
| return [json.loads(line) for line in self.log.read_text().splitlines() if line.startswith("{")] | |
| def stack(tmp_path): | |
| """Start fake Tritons and services on demand; stop them all afterwards.""" | |
| started = [] | |
| class Stack: | |
| def triton(self, **kw): | |
| t = FakeTriton(**kw).start() | |
| started.append(t) | |
| return t | |
| def service(self, triton, *extra, **kw): | |
| s = Service(triton, tmp_path, *extra, **kw) | |
| started.append(s) | |
| return s | |
| yield Stack() | |
| for s in reversed(started): | |
| s.stop() | |
| def read_reply(s): | |
| """Everything the server sends until it closes the connection.""" | |
| out = bytearray() | |
| while True: | |
| b = s.recv(65536) | |
| if not b: | |
| return bytes(out) | |
| out += b | |
| def post_audio(url, data, ctype="application/octet-stream", headers=None): | |
| return call(url, "POST", data, {"Content-Type": ctype, **(headers or {})}) | |
| def multipart(fields, boundary="anuvaniBoundary7MA4YWxk"): | |
| """A multipart/form-data body: fields are (name, filename or None, bytes).""" | |
| out = io.BytesIO() | |
| for name, filename, data in fields: | |
| out.write(f"--{boundary}\r\n".encode()) | |
| disp = f'form-data; name="{name}"' + (f'; filename="{filename}"' if filename is not None else "") | |
| out.write(f"Content-Disposition: {disp}\r\n".encode()) | |
| if filename is not None: | |
| out.write(b"Content-Type: application/octet-stream\r\n") | |
| out.write(b"\r\n" + data + b"\r\n") | |
| out.write(f"--{boundary}--\r\n".encode()) | |
| return out.getvalue(), f"multipart/form-data; boundary={boundary}" | |
| # ---------------------------------------------------------------- byte identity | |
| def test_jobs_identical_to_the_release(fake_reference, stack, schedule): | |
| items, ref = fake_reference | |
| svc = stack.service(stack.triton(), "--schedule", schedule) | |
| got = run_jobs(svc, items, ref, seed={"fair": 3, "fifo": 4, "share": 9}[schedule]) | |
| assert compare(items, got, ref, svc.records) == [] | |
| m = svc.metrics() | |
| log = json.loads((ref / "log.json").read_text()) | |
| assert m["anuvani_diar_jobs_done_total"] == sum(1 for it in items if "error" not in log[it["sid"]]) | |
| assert m["anuvani_diar_triton_retries_total"] == 0 and m["anuvani_diar_spool_bytes"] == 0 | |
| assert m["anuvani_diar_audio_bytes_admitted"] == 0 | |
| def test_one_triton_per_gpu(fake_reference, stack): | |
| """Topology B: each GPU's model on its own Triton (`--triton url0,url1` in the order of `--gpus`).""" | |
| items, ref = fake_reference | |
| t0, t1 = stack.triton(models=("diar_nn_g0", "resample_cpu")), stack.triton(models=("diar_nn_g1", "resample_cpu")) | |
| svc = stack.service(t0, "--triton", f"{t0.url},{t1.url}") | |
| got = run_jobs(svc, items, ref, seed=8) | |
| assert compare(items, got, ref, svc.records) == [] | |
| assert t0.calls["diar_nn_g0"] > 0 and t1.calls["diar_nn_g1"] > 0 | |
| bad = stack.service(t0, "--triton", f"{t0.url},{t1.url},{t1.url}", wait=False) | |
| assert bad.proc.wait(30) != 0 and "one URL, or one per GPU" in bad.log.read_text() | |
| def test_many_clients_at_once(fake_reference, stack): | |
| """64 client threads, every item three times (uploads repeat, so some answers come from the result cache); the | |
| event loop stays responsive meanwhile (the engine's work never holds the GIL).""" | |
| items, ref = fake_reference | |
| svc = stack.service(stack.triton()) | |
| worst, stop = [0.0], threading.Event() | |
| def probe(): | |
| while not stop.is_set(): | |
| t0 = time.monotonic() | |
| call(svc.url + "/health") | |
| worst[0] = max(worst[0], time.monotonic() - t0) | |
| time.sleep(0.02) | |
| th = threading.Thread(target=probe) | |
| th.start() | |
| try: | |
| for rep in range(3): | |
| got = run_jobs(svc, items, ref, uploads=12, seed=10 + rep, threads=64) | |
| assert compare(items, got, ref, svc.records) == [] | |
| finally: | |
| stop.set() | |
| th.join() | |
| assert svc.metrics()["anuvani_diar_cache_hits_total"] >= 12 | |
| assert worst[0] < 1.0, worst[0] | |
| def test_schedules_share_the_devices_as_documented(fake_reference, stack, schedule): | |
| """One device slot; a long file runs, then a shorter file (at least 4 blocks) arrives. Fair serves the newcomer's | |
| blocks first; share interleaves them.""" | |
| items, ref = fake_reference | |
| blocks = {it["sid"]: -(-np.load(ref / f"{it['sid']}.npz")["P"].shape[0] // 8) for it in ok_items(items, ref)} | |
| long_it = longest(ok_items(items, ref), ref, 1)[0] | |
| short_it = min((it for it in ok_items(items, ref) if blocks[it["sid"]] >= 4), key=lambda it: blocks[it["sid"]]) | |
| m = blocks[short_it["sid"]] | |
| assert blocks[long_it["sid"]] >= 2 * m + 4, (blocks[long_it["sid"]], m) | |
| tri = stack.triton(delay=0.05) | |
| svc = stack.service(tri, "--schedule", schedule, "--gpus", "0", "--slots-per-gpu", "1") | |
| c = svc.client() | |
| long_job = c.submit(long_it["path"], uri=long_it["sid"], server_path=True) | |
| wait_for(lambda: c.status(long_job)["blocks_done"] >= m + 2, timeout=120) | |
| short_job = c.submit(short_it["path"], uri=short_it["sid"], server_path=True) | |
| wait_for(lambda: c.status(short_job)["blocks_done"] >= 1, timeout=60) | |
| before = c.status(long_job)["blocks_done"] | |
| s = c.wait(short_job, poll=0.02) | |
| advanced = c.status(long_job)["blocks_done"] - before | |
| assert s["status"] == "done" and s["blocks"] == m | |
| if schedule == "fair": | |
| assert advanced <= 1, advanced | |
| else: | |
| assert advanced >= m - 3, (advanced, m) | |
| c.cancel(long_job) | |
| # ---------------------------------------------------------------- Triton going away, the service restarting | |
| def test_jobs_survive_a_triton_restart(fake_reference, stack): | |
| """Triton killed mid-run (open connections cut), down 2 s, models loading 2 s more: every job still finishes | |
| with identical bytes, the requests in flight having been sent again.""" | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.03) | |
| svc = stack.service(tri) | |
| restart = tri.restart_in_background(after_s=1.5, down_s=2.0, loading_s=2.0) | |
| got = run_jobs(svc, items, ref, seed=5) | |
| restart.join() | |
| assert compare(items, got, ref, svc.records) == [] | |
| assert svc.metrics()["anuvani_diar_triton_retries_total"] > 0 | |
| def test_jobs_survive_a_service_restart(fake_reference, stack): | |
| """The service killed (SIGKILL) mid-run and started again on the same port: the client backs off while it is | |
| down, finds its jobs gone, submits them again, and every result is still identical.""" | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.03) | |
| svc = stack.service(tri) | |
| spool, port = svc.spool, svc.port | |
| def crash_and_restart(): | |
| time.sleep(1.5) | |
| svc.proc.kill() | |
| svc.proc.wait() | |
| time.sleep(1.0) | |
| restarted.append(stack.service(tri, spool=spool, port=port)) | |
| restarted = [] | |
| th = threading.Thread(target=crash_and_restart) | |
| th.start() | |
| got = run_jobs(svc, items, ref, seed=6) | |
| th.join() | |
| new = restarted[0] | |
| for f in svc.records.glob("*"): | |
| if not (new.records / f.name).exists(): | |
| (new.records / f.name).write_bytes(f.read_bytes()) | |
| assert compare(items, got, ref, new.records) == [] | |
| assert new.metrics()["anuvani_diar_jobs_submitted_total"] > 0 | |
| def test_a_job_keeps_its_id_after_the_service_restarts(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.05) | |
| svc = stack.service(tri) | |
| c = svc.client(retries=200) | |
| it = longest(ok_items(items, ref), ref, 1)[0] | |
| job = c.submit(it["path"], uri=it["sid"], server_path=True) | |
| wait_for(lambda: c.status(job)["blocks_done"] > 0, timeout=60) | |
| svc.proc.kill() | |
| svc.proc.wait() | |
| t0 = time.monotonic() | |
| assert c.ready() is False and time.monotonic() - t0 < 5 | |
| new = stack.service(tri, spool=svc.spool, port=svc.port) | |
| assert c.wait(job, poll=0.05)["status"] == "done" | |
| for _ in range(2): | |
| r = c.result(job) | |
| assert r.ok and r.rttm.encode() == (ref / f"{it['sid']}.rttm").read_bytes() | |
| assert c.status(job)["status"] == "done" and c.cancel(job)["cancelling"] is False | |
| assert c.ready() is True and new.metrics()["anuvani_diar_jobs_submitted_total"] == 1 | |
| def test_ready_follows_triton(stack): | |
| tri = stack.triton() | |
| svc = stack.service(tri) | |
| assert call(svc.url + "/ready")[0] == 200 | |
| tri.stop() | |
| wait_for(lambda: call(svc.url + "/ready")[0] == 503, timeout=10) | |
| code, body, _ = call(svc.url + "/ready") | |
| d = json.loads(body) | |
| # names the model, not the address | |
| assert "diar_nn_g0" in d["error"] and b"127.0.0.1" not in body | |
| assert set(d["models"]) == {"diar_nn_g0", "diar_nn_g1", "resample_cpu"} | |
| tri.start() | |
| wait_for(lambda: call(svc.url + "/ready")[0] == 200, timeout=10) | |
| def test_jobs_fail_cleanly_when_triton_stays_away(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.05) | |
| svc = stack.service(tri, "--triton-retry-s", "2") | |
| c = svc.client() | |
| jobs = [c.submit(it["path"], uri=it["sid"], server_path=True) for it in longest(ok_items(items, ref), ref, 4)] | |
| wait_for(lambda: tri.calls["diar_nn_g0"] + tri.calls["diar_nn_g1"] > 0, timeout=30) | |
| tri.stop() | |
| t0 = time.monotonic() | |
| final = [c.wait(j, poll=0.1) for j in jobs] | |
| assert time.monotonic() - t0 < 60 | |
| for s in final: | |
| assert s["status"] == "failed" and "still away after 2 s" in s["error"], s | |
| assert svc.metrics()["anuvani_diar_jobs_failed_total"] == len(jobs) | |
| def test_model_errors_are_final(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton() | |
| for m in ("diar_nn_g0", "diar_nn_g1"): | |
| tri.fail(m, OOM) | |
| svc = stack.service(tri) | |
| c = svc.client() | |
| for it in ok_items(items, ref, 3): | |
| s = c.wait(c.submit(it["path"], uri=it["sid"], server_path=True), poll=0.05) | |
| assert s["status"] == "failed" and s["error"].endswith(f"(HTTP 500 Internal Server Error): {OOM}"), s | |
| assert svc.metrics()["anuvani_diar_triton_retries_total"] == 0 | |
| def test_a_dead_model_process_is_retried(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton() | |
| tri.fail("diar_nn_g0", "Failed to process the request(s) for model instance 'diar_nn_g0_0_0', message: " | |
| "Stub process 'diar_nn_g0_0_0' is not healthy.", times=3) | |
| svc = stack.service(tri, "--gpus", "0") | |
| sub = ok_items(items, ref, 4) | |
| got = run_jobs(svc, sub, ref, uploads=0) | |
| assert compare(sub, got, ref, svc.records) == [] | |
| assert svc.metrics()["anuvani_diar_triton_retries_total"] == 3 | |
| def test_engine_waits_release_the_gil(stack): | |
| """The engine's blocking calls (here: waiting for finished jobs, probing readiness) let other Python threads run.""" | |
| from anuvani_diar import _core | |
| tri = stack.triton() | |
| eng = _core.TritonService(tri.url, gpus=[0, 1], decode_threads=1, post_threads=2) | |
| count, stop = [0], threading.Event() | |
| def spin(): | |
| while not stop.is_set(): | |
| count[0] += 1 | |
| th = threading.Thread(target=spin) | |
| th.start() | |
| try: | |
| time.sleep(0.05) | |
| c0 = count[0] | |
| # one second waiting for nothing | |
| assert eng.finished(1.0) == [] | |
| c1 = count[0] | |
| tri.stop() | |
| assert all(m for _, m in eng.ready()) | |
| finally: | |
| stop.set() | |
| th.join() | |
| eng.shutdown() | |
| # the spinning thread ran while the main thread waited | |
| assert c1 - c0 > 100_000, c1 - c0 | |
| def test_engine_api(fake_reference, stack): | |
| """`anuvani_diar._core.TritonService` used directly: submit, progress, result (one job's end, also with a | |
| timeout), cancel, the feed, counters, the budget refusal (ServiceBusy) and a bounded shutdown.""" | |
| from anuvani_diar import _core | |
| items, ref = fake_reference | |
| # each under the 20 MiB budget, both together over it | |
| a, b = longest(ok_items(items, ref), ref, 2) | |
| tri = stack.triton(delay=0.05) | |
| eng = _core.TritonService([tri.url], gpus=[0, 1], decode_threads=2, post_threads=2, audio_budget_bytes=20 << 20) | |
| try: | |
| ia = eng.submit(a["path"], a["sid"], a["path"], "tag-a") | |
| assert eng.counters()["audio_bytes_admitted"] > 0 | |
| with pytest.raises(_core.ServiceBusy, match="decoded-audio budget"): | |
| eng.submit(b["path"], b["sid"], b["path"], "tag-b") | |
| assert eng.progress(ia)[0] in ("queued", "running") and eng.result(ia) is None | |
| r = eng.result(ia, 60.0) | |
| assert r["tag"] == "tag-a" and r["status"] == "done" and r["rttm"].encode() == (ref / f"{a['sid']}.rttm").read_bytes() | |
| assert r["blocks"][0] == r["blocks"][1] > 0 and set(r["timing"]) == {"decode_s", "frontend_s", "post_s", "wall_s"} | |
| assert eng.progress(ia) is None and eng.result(ia) is None and eng.finished(0.0) == [] and not eng.cancel(ia) | |
| ib = eng.submit(b["path"], b["sid"], b["path"], "tag-b") | |
| wait_for(lambda: eng.progress(ib) is not None and eng.progress(ib)[1] > 0, timeout=30) | |
| assert eng.cancel(ib) | |
| (r,) = eng.finished(30.0) | |
| assert r["index"] == ib and r["status"] == "cancelled" and r["error"] == "cancelled" | |
| c = eng.counters() | |
| assert (c["submitted"], c["done"], c["cancelled"], c["busy"], c["audio_bytes_admitted"]) == (2, 1, 1, 1, 0) | |
| finally: | |
| assert eng.shutdown(True, 10.0) | |
| # ---------------------------------------------------------------- admission and uploads | |
| def test_job_queue_limit(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.3) | |
| svc = stack.service(tri, "--max-jobs", "2") | |
| c = svc.client(retries=0) | |
| a, b = (c.submit(it["path"], uri=it["sid"], server_path=True) for it in longest(ok_items(items, ref), ref, 2)) | |
| body = json.dumps({"path": items[0]["path"]}).encode() | |
| code, out, headers = call(svc.url + "/v1/jobs", "POST", body, {"Content-Type": "application/json"}) | |
| assert code == 429 and headers.get("Retry-After") == "5" and b"queue is full" in out | |
| assert svc.metrics()["anuvani_diar_jobs_rejected_total"] == 1 | |
| c.cancel(a), c.cancel(b) | |
| assert c.wait(a)["status"] == "cancelled" and c.wait(b)["status"] == "cancelled" | |
| assert call(svc.url + "/v1/jobs", "POST", body, {"Content-Type": "application/json"})[0] == 202 | |
| def test_client_backs_off_when_the_queue_is_full(fake_reference, stack): | |
| items, ref = fake_reference | |
| svc = stack.service(stack.triton(), "--max-jobs", "2") | |
| sub = ok_items(items, ref, 6) | |
| got = run_jobs(svc, sub, ref, uploads=2, threads=6) | |
| assert compare(sub, got, ref, svc.records) == [] | |
| assert svc.metrics()["anuvani_diar_jobs_rejected_total"] > 0 | |
| def test_upload_size_limit(stack): | |
| svc = stack.service(stack.triton(), "--max-upload-mb", "1") | |
| code, out, _ = post_audio(svc.url + "/v1/jobs?name=big.wav", os.urandom(2 * 1024 * 1024)) | |
| assert code == 413 and b"larger than 1 MB" in out | |
| assert [p.name for p in svc.spool.iterdir()] == [] and svc.metrics()["anuvani_diar_spool_bytes"] == 0 | |
| # chunked (no Content-Length): enforced while the bytes arrive | |
| import http.client | |
| conn = http.client.HTTPConnection("127.0.0.1", svc.port, timeout=30) | |
| conn.request("POST", "/v1/jobs?name=big.wav", body=iter([b"\0" * 65536] * 40), | |
| headers={"Content-Type": "application/octet-stream"}, encode_chunked=True) | |
| r = conn.getresponse() | |
| assert r.status == 413 and b"larger than 1 MB" in r.read() | |
| conn.close() | |
| wait_for(lambda: svc.metrics()["anuvani_diar_uploads_in_progress"] == 0, timeout=10) | |
| assert [p.name for p in svc.spool.iterdir()] == [] and svc.metrics()["anuvani_diar_spool_bytes"] == 0 | |
| def test_limits_are_checked_from_content_length_before_the_body(stack): | |
| """A declared size over the upload limit (413) or over the spool's free space (503, Retry-After) is answered from | |
| the headers alone, before any byte of the body is read; a large unread body closes the connection, a small one | |
| is read and dropped after the answer so that the client sees it.""" | |
| # 537 KB of spool | |
| svc = stack.service(stack.triton(), "--max-upload-mb", "1", "--max-spool-gb", "0.0005") | |
| for length, status, text in ((10 * 1024**3, b"413", b"larger than 1 MB"), (100 << 20, b"413", b"larger than 1 MB")): | |
| s = raw_request(svc.port, f"POST /v1/jobs?name=a.wav HTTP/1.1\r\nHost: x\r\nContent-Type: application/octet-stream\r\n" | |
| f"Content-Length: {length}\r\n\r\n") | |
| # the server closes without reading the body | |
| reply = read_reply(s) | |
| assert reply.startswith(b"HTTP/1.1 " + status) and text in reply and b"connection: close" in reply.lower(), reply | |
| s.close() | |
| code, out, headers = post_audio(svc.url + "/v1/jobs?name=a.wav", os.urandom(600_000)) | |
| assert code == 503 and headers.get("Retry-After") == "5" and b"spool is full" in out | |
| # 10.7 MB of spool | |
| big = stack.service(stack.triton(), "--max-spool-gb", "0.01") | |
| s = raw_request(big.port, "POST /v1/diarize?name=a.wav HTTP/1.1\r\nHost: x\r\nContent-Type: application/octet-stream\r\n" | |
| f"Content-Length: {100 << 20}\r\n\r\n") | |
| reply = read_reply(s) | |
| assert reply.startswith(b"HTTP/1.1 503") and b"retry-after: 5" in reply.lower() and b"connection: close" in reply.lower() | |
| assert b"spool is full" in reply | |
| s.close() | |
| for x in (svc, big): | |
| m = x.metrics() | |
| assert m["anuvani_diar_upload_bytes_total"] == 0 and m["anuvani_diar_spool_bytes"] == 0 | |
| assert [p.name for p in x.spool.iterdir()] == [] | |
| assert svc.metrics()["anuvani_diar_jobs_rejected_total"] == 1 and big.metrics()["anuvani_diar_jobs_rejected_total"] == 1 | |
| def test_stalled_upload_and_upload_limit(stack): | |
| svc = stack.service(stack.triton(), "--max-uploads", "1", "--upload-idle-s", "1") | |
| s = raw_request(svc.port, "POST /v1/jobs?name=slow.wav HTTP/1.1\r\nHost: x\r\nContent-Type: application/octet-stream\r\n" | |
| "Content-Length: 1000000\r\n\r\n", b"RIFF" + b"\0" * 60) | |
| wait_for(lambda: svc.metrics()["anuvani_diar_uploads_in_progress"] == 1, timeout=10) | |
| code, out, _ = post_audio(svc.url + "/v1/jobs?name=b.wav", b"x" * 10) | |
| assert code == 429 and b"too many uploads" in out | |
| # 408, and the connection closed | |
| reply = read_reply(s) | |
| s.close() | |
| assert reply.startswith(b"HTTP/1.1 408") and b"upload stalled" in reply | |
| wait_for(lambda: svc.metrics()["anuvani_diar_uploads_in_progress"] == 0, timeout=10) | |
| assert [p.name for p in svc.spool.iterdir()] == [] | |
| def test_a_low_sample_rate_cannot_resample_past_the_limit(stack): | |
| """600,000 samples declared at 1 Hz: 2.4 MB uploaded, 38.4 GB at 16 kHz, refused before anything is resampled.""" | |
| import soundfile as sf | |
| tri = stack.triton() | |
| svc = stack.service(tri) | |
| buf = io.BytesIO() | |
| sf.write(buf, np.full(600000, 0.1, np.float32), 1, subtype="FLOAT", format="WAV") | |
| code, out, _ = post_audio(svc.url + "/v1/diarize?name=low.wav", buf.getvalue()) | |
| assert code == 422 and json.loads(out)["error"] == ("RefusedError: audio of 38400000000 bytes (resampled to 16 kHz) is " | |
| "larger than this engine's limit of 8589934592 bytes") | |
| assert tri.calls["resample_cpu"] == 0 | |
| def test_a_rate_that_does_not_reduce_cannot_build_a_huge_resampling_filter(stack): | |
| """A 500-byte WAV declaring 16,001 Hz would have the resampling model build a 2 GB filter (1,000,003 Hz: 128 GB); | |
| refused from the header before anything is resampled (docs/SECURITY.md, security review of 9 October 2026).""" | |
| import soundfile as sf | |
| tri = stack.triton() | |
| svc = stack.service(tri) | |
| buf = io.BytesIO() | |
| sf.write(buf, np.full(100, 0.1, np.float32), 16001, subtype="FLOAT", format="WAV") | |
| code, out, _ = post_audio(svc.url + "/v1/diarize?name=odd.wav", buf.getvalue()) | |
| assert code == 422 and json.loads(out)["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 tri.calls["resample_cpu"] == 0 | |
| def test_multipart_framing_is_bounded(stack): | |
| """Everything in a multipart body but the file's own bytes is bounded: at most 16 parts, and boundaries, part | |
| headers, fields and anything after the closing boundary (which the parser skips) together within 128 KiB above | |
| the upload limit, also while streaming, so empty parts with long names or an endless epilogue cannot grow memory | |
| or hold the request (security review of 9 October 2026). Nothing stays in the spool.""" | |
| svc = stack.service(stack.triton()) | |
| body, ctype = multipart([(f"f{i}", None, b"") for i in range(16)] + [("file", "a.wav", b"1")]) | |
| code, out, _ = post_audio(svc.url + "/v1/jobs", body, ctype) | |
| assert code == 413 and b"more than 16 parts" in out, out | |
| body, ctype = multipart([("file", "a.wav", b"1")]) | |
| code, out, _ = post_audio(svc.url + "/v1/jobs", body + b"e" * 200_000, ctype) | |
| assert code == 413 and b"framing and form fields too large" in out, out | |
| s = raw_request(svc.port, "POST /v1/jobs HTTP/1.1\r\nHost: x\r\nTransfer-Encoding: chunked\r\n" | |
| f"Content-Type: {ctype}\r\n\r\n", b"") | |
| # the body, then 3 x 64 KiB of epilogue, chunked | |
| try: | |
| s.sendall(f"{len(body):x}\r\n".encode() + body + b"\r\n") | |
| for _ in range(3): | |
| s.sendall(b"10000\r\n" + b"e" * 65536 + b"\r\n") | |
| # answered and closed before the last chunk | |
| except OSError: | |
| pass | |
| reply = read_reply(s) | |
| s.close() | |
| assert reply.startswith(b"HTTP/1.1 413") and b"framing and form fields too large" in reply, reply[:300] | |
| wait_for(lambda: svc.metrics()["anuvani_diar_uploads_in_progress"] == 0, timeout=10) | |
| assert [p.name for p in svc.spool.iterdir()] == [] | |
| def test_listening_beyond_the_loopback_needs_a_token(monkeypatch): | |
| """--listen on an address other than the loopback refuses to start without a token, unless --insecure-no-token | |
| (security review of 9 October 2026: the image listens on 0.0.0.0 inside its container).""" | |
| monkeypatch.syspath_prepend(str(SERVICE)) | |
| monkeypatch.delenv("ANUVANI_DIAR_TOKEN", raising=False) | |
| from anuvani_diar_service import settings | |
| with pytest.raises(SystemExit): | |
| settings.parse(["--listen", "0.0.0.0:8180"]) | |
| with pytest.raises(SystemExit): | |
| settings.parse(["--listen", "[::]:8180"]) | |
| assert settings.parse(["--listen", "0.0.0.0:8180", "--token", "t"]).token == "t" | |
| monkeypatch.setenv("ANUVANI_DIAR_TOKEN", "from-env") | |
| assert settings.parse(["--listen", "0.0.0.0:8180"]).token == "from-env" | |
| monkeypatch.delenv("ANUVANI_DIAR_TOKEN") | |
| assert settings.parse(["--listen", "0.0.0.0:8180", "--insecure-no-token"]).token is None | |
| for host in ("127.0.0.1", "127.0.0.9", "localhost", "[::1]"): | |
| assert settings.parse(["--listen", f"{host}:8180"]).token is None | |
| def test_decoded_audio_budget(fake_reference, stack): | |
| """A job whose decoded audio would take the admitted jobs past --audio-budget-gb is answered 503 with Retry-After | |
| before anything is decoded; one job alone always gets in, and room comes back when jobs end.""" | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.2) | |
| # 21.5 MB; one benchmark file is ~13 MB | |
| svc = stack.service(tri, "--audio-budget-gb", "0.02") | |
| a, b = longest(ok_items(items, ref), ref, 2) | |
| c = svc.client(retries=0) | |
| first = c.submit(a["path"], uri=a["sid"]) | |
| assert 0 < svc.metrics()["anuvani_diar_audio_bytes_admitted"] <= 0.02 * 2**30 | |
| q = urllib.parse.urlencode(dict(name=Path(b["path"]).name, uri=b["sid"])) | |
| code, out, headers = post_audio(svc.url + "/v1/jobs?" + q, Path(b["path"]).read_bytes()) | |
| assert code == 503 and headers.get("Retry-After") == "5" and b"decoded-audio budget" in out, out | |
| code, out, _ = call(svc.url + "/v1/jobs", "POST", json.dumps({"path": b["path"], "uri": b["sid"]}).encode(), | |
| {"Content-Type": "application/json"}) | |
| assert code == 503 and b"decoded-audio budget" in out | |
| m = svc.metrics() | |
| # only the first's file | |
| assert m["anuvani_diar_jobs_rejected_total"] == 2 and m["anuvani_diar_spool_bytes"] > 0 | |
| c.cancel(first) | |
| assert c.wait(first)["status"] == "cancelled" | |
| assert svc.metrics()["anuvani_diar_audio_bytes_admitted"] == 0 | |
| s = c.wait(c.submit(b["path"], uri=b["sid"]), poll=0.05) | |
| assert s["status"] == "done" and s["rttm"].encode() == (ref / f"{b['sid']}.rttm").read_bytes() | |
| def test_client_does_not_follow_redirects(): | |
| """The client is shared with the gateway; its redirect rule is tested there (test_gateway.py).""" | |
| from test_gateway import test_client_does_not_follow_redirects as t | |
| t() | |
| def test_upload_names_are_checked(stack): | |
| svc = stack.service(stack.triton()) | |
| for q, msg in ((dict(name="../a.wav"), b"plain file name"), (dict(name="a.wav", uri="a b"), b"uri:"), | |
| (dict(name="a.wav", uri="x" * 201), b"uri:"), (dict(name="a.wav", uri="a/b"), b"uri:"), | |
| (dict(name="a.wav", uri="a b"), b"uri:"), (dict(name="."), b"plain file name")): | |
| for path in ("/v1/jobs?", "/v1/diarize?"): | |
| code, out, _ = post_audio(svc.url + path + urllib.parse.urlencode(q), b"x") | |
| assert code == 400 and msg in out, (q, out) | |
| def test_multipart_uploads(fake_reference, stack): | |
| """A multipart/form-data body streams its file part to the spool: the same bytes, the same RTTM. The file | |
| part's filename is the default name, a "uri" field the uri; a missing, doubled or unterminated file part is 400.""" | |
| items, ref = fake_reference | |
| svc = stack.service(stack.triton()) | |
| it = ok_items(items, ref, 1)[0] | |
| data = Path(it["path"]).read_bytes() | |
| body, ctype = multipart([("note", None, b"ignored"), ("uri", None, it["sid"].encode()), ("file", "x.flac", data)]) | |
| code, out, _ = post_audio(svc.url + "/v1/jobs", body, ctype) | |
| assert code == 202, out | |
| s = svc.client().wait(json.loads(out)["job"], poll=0.05) | |
| assert s["status"] == "done" and s["uri"] == it["sid"] and s["rttm"].encode() == (ref / f"{it['sid']}.rttm").read_bytes() | |
| body, ctype = multipart([("file", f"{it['sid']}.flac", data)]) | |
| # uri: the filename's stem | |
| code, out, headers = post_audio(svc.url + "/v1/diarize", body, ctype) | |
| assert code == 200 and out == (ref / f"{it['sid']}.rttm").read_bytes() and svc.metrics()["anuvani_diar_cache_hits_total"] == 1 | |
| for fields, msg in (([("uri", None, b"a")], b"no file part"), ([("file", "a.wav", b"1"), ("f2", "b.wav", b"2")], b"more than one"), | |
| ([("file", "a.wav", b"1"), ("uri", None, b"a b")], b"uri:")): | |
| body, ctype = multipart(fields) | |
| code, out, _ = post_audio(svc.url + "/v1/jobs", body, ctype) | |
| assert code == 400 and msg in out, (fields, out) | |
| body, ctype = multipart([("file", "a.wav", data[:5000])]) | |
| code, out, _ = post_audio(svc.url + "/v1/jobs", body[:-30], ctype) | |
| assert code == 400 and b"closing boundary" in out | |
| code, out, _ = post_audio(svc.url + "/v1/jobs", b"x", "multipart/form-data") | |
| assert code == 400 and b"no boundary" in out | |
| assert [p.name for p in svc.spool.iterdir()] == [] and svc.metrics()["anuvani_diar_spool_bytes"] == 0 | |
| # ---------------------------------------------------------------- path jobs and the token | |
| def test_path_jobs_stay_under_the_allowed_root(fake_reference, stack, tmp_path): | |
| items, ref = fake_reference | |
| it = ok_items(items, ref, 1)[0] | |
| root = tmp_path / "allowed" | |
| root.mkdir() | |
| inside = root / "a.flac" | |
| inside.write_bytes(Path(it["path"]).read_bytes()) | |
| (root / "escape.flac").symlink_to(it["path"]) | |
| (root / "sub").mkdir() | |
| svc = stack.service(stack.triton(), root=root) | |
| def post(path, **kw): | |
| return call(svc.url + "/v1/jobs", "POST", json.dumps({"path": path, **kw}).encode(), {"Content-Type": "application/json"}) | |
| assert post(str(inside))[0] == 202 | |
| for p in (it["path"], str(root / "escape.flac"), str(root / "sub"), str(root / "missing.flac"), "a.flac", | |
| str(root / ".." / "allowed" / ".." / Path(it["path"]).name), str(root) + "/a.flac\0x", "/etc/passwd"): | |
| code, out, _ = post(p) | |
| assert code == 403 and b"not a readable regular file under the allowed root" in out, p | |
| assert post(str(inside), uri="a b")[0] == 400 | |
| for bad in (b"{not json", b'{"path": 3}', b"[]", b'{"path": "/x", "uri": 7}'): | |
| assert call(svc.url + "/v1/jobs", "POST", bad, {"Content-Type": "application/json"})[0] == 400, bad | |
| code, out, _ = call(svc.url + "/v1/jobs", "POST", b'{"path": "' + b"a" * 70000 + b'"}', {"Content-Type": "application/json"}) | |
| assert code == 413 | |
| closed = stack.service(stack.triton(), root=None) | |
| code, out, _ = call(closed.url + "/v1/jobs", "POST", json.dumps({"path": str(inside)}).encode(), | |
| {"Content-Type": "application/json"}) | |
| assert code == 403 and b"path jobs are disabled" in out | |
| def test_token(fake_reference, stack): | |
| items, ref = fake_reference | |
| svc = stack.service(stack.triton(), token="s3cret") | |
| assert call(svc.url + "/health")[0] == 200 | |
| for h in ({}, {"Authorization": "Bearer wrong"}, {"Authorization": "s3cret"}, {"Authorization": "Bearer s3cre"}, | |
| {"Authorization": "Bearer s3cret2"}, {"Authorization": "bearer s3cret"}): | |
| for path in ("/ready", "/metrics", "/v1/jobs/0", "/docs", "/openapi.json", "/nowhere"): | |
| code, _, headers = call(svc.url + path, headers=h) | |
| assert code == 401 and headers.get("WWW-Authenticate") == "Bearer", (h, path) | |
| assert post_audio(svc.url + "/v1/jobs?name=a.wav", b"x" * 1000, headers=h)[0] == 401 | |
| ok = {"Authorization": "Bearer s3cret"} | |
| assert call(svc.url + "/ready", headers=ok)[0] == 200 and call(svc.url + "/openapi.json", headers=ok)[0] == 200 | |
| sub = ok_items(items, ref, 2) | |
| from anuvani_diar import Client | |
| c = Client(svc.url, token="s3cret") | |
| got = {it["sid"]: c.wait(c.submit(it["path"], uri=it["sid"]), poll=0.05) for it in sub} | |
| assert compare(sub, got, ref, svc.records) == [] | |
| text = svc.log.read_text() | |
| assert "s3cret" not in text and str(Path(sub[0]["path"]).parent) not in text | |
| # ---------------------------------------------------------------- job life cycle | |
| def test_cancel(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.2) | |
| svc = stack.service(tri) | |
| c = svc.client() | |
| it = longest(ok_items(items, ref), ref, 1)[0] | |
| # an upload: its spool file goes too | |
| job = c.submit(it["path"], uri=it["sid"]) | |
| wait_for(lambda: c.status(job)["blocks_done"] > 0, timeout=30) | |
| assert c.cancel(job) == {"job": job, "cancelling": True} | |
| s = c.wait(job, poll=0.05) | |
| assert s["status"] == "cancelled" and s["error"] == "cancelled" | |
| # the rest was never sent | |
| assert tri.calls["diar_nn_g0"] + tri.calls["diar_nn_g1"] < 2 * s["blocks"] | |
| assert c.cancel(job) == {"job": job, "cancelling": False} | |
| m = svc.metrics() | |
| assert m["anuvani_diar_jobs_cancelled_total"] == 1 and m["anuvani_diar_spool_bytes"] == 0 | |
| assert [p.name for p in svc.spool.iterdir()] == [] | |
| def test_lookups(stack, fake_reference): | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.2) | |
| svc = stack.service(tri) | |
| assert call(svc.url + "/v1/jobs/feedbeef")[0] == 404 | |
| assert call(svc.url + "/v1/jobs/feedbeef", "DELETE")[0] == 404 | |
| assert call(svc.url + "/v1/jobs/feedbeef/rttm")[0] == 404 | |
| c = svc.client() | |
| it = longest(ok_items(items, ref), ref, 1)[0] | |
| job = c.submit(it["path"], uri=it["sid"], server_path=True) | |
| # 128 random bits | |
| assert len(job) == 32 and int(job, 16) >= 0 | |
| code, out, _ = call(svc.url + f"/v1/jobs/{job}/rttm") | |
| assert code == 409 and b"job is" in out | |
| s = c.wait(job, poll=0.05) | |
| # one model: no per-mode routes | |
| assert call(svc.url + f"/v1/jobs/{job}/rttm/public")[0] == 404 | |
| code, out, headers = call(svc.url + f"/v1/jobs/{job}/rttm") | |
| assert code == 200 and out == (ref / f"{it['sid']}.rttm").read_bytes() | |
| assert headers["content-type"].startswith("text/plain") | |
| assert set(s) == {"job", "uri", "status", "submitted", "blocks_done", "blocks", "seconds", "duration_s", "error", | |
| # the gateway's fields | |
| "timing", "rttm"} | |
| assert set(s["timing"]) == {"decode_s", "frontend_s", "post_s", "wall_s"} and s["blocks_done"] == s["blocks"] > 0 | |
| assert call(svc.url + "/v1/jobs", "PUT")[0] == 405 and call(svc.url + "/nowhere")[0] == 404 | |
| def test_synchronous_call(fake_reference, stack): | |
| items, ref = fake_reference | |
| svc = stack.service(stack.triton()) | |
| for it in ok_items(items, ref, 3): | |
| q = urllib.parse.urlencode(dict(name=Path(it["path"]).name, uri=it["sid"])) | |
| code, out, headers = post_audio(svc.url + "/v1/diarize?" + q, Path(it["path"]).read_bytes()) | |
| assert code == 200 and out == (ref / f"{it['sid']}.rttm").read_bytes() and len(headers["x-job"]) == 32 | |
| assert headers["content-type"] == "text/plain; charset=utf-8" | |
| code, out, _ = call(svc.url + "/v1/diarize?name=bad.wav", "POST", b"not audio at all") | |
| assert code == 422 and json.loads(out)["error"].startswith("LibsndfileError: Error opening 'bad.wav'") | |
| def test_synchronous_call_cancelled_when_the_client_leaves(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.2) | |
| svc = stack.service(tri) | |
| it = longest(ok_items(items, ref), ref, 1)[0] | |
| data = Path(it["path"]).read_bytes() | |
| s = raw_request(svc.port, f"POST /v1/diarize?name=a.flac HTTP/1.1\r\nHost: x\r\nContent-Type: application/octet-stream\r\n" | |
| f"Content-Length: {len(data)}\r\n\r\n", data) | |
| wait_for(lambda: svc.metrics()["anuvani_diar_jobs_running"] == 1, timeout=30) | |
| s.close() | |
| wait_for(lambda: svc.metrics()["anuvani_diar_jobs_cancelled_total"] == 1, timeout=60) | |
| wait_for(lambda: svc.metrics()["anuvani_diar_spool_bytes"] == 0, timeout=10) | |
| def test_result_cache(fake_reference, stack): | |
| items, ref = fake_reference | |
| tri = stack.triton() | |
| svc = stack.service(tri) | |
| c = svc.client() | |
| it = ok_items(items, ref, 1)[0] | |
| first = c.wait(c.submit(it["path"], uri=it["sid"]), poll=0.05) | |
| calls = sum(tri.calls.values()) | |
| again = c.wait(c.submit(it["path"], uri=it["sid"]), poll=0.05) | |
| assert svc.metrics()["anuvani_diar_cache_hits_total"] == 1 and sum(tri.calls.values()) == calls | |
| assert again["rttm"] == first["rttm"] | |
| # another uri is another result | |
| other = c.wait(c.submit(it["path"], uri="other_name"), poll=0.05) | |
| assert svc.metrics()["anuvani_diar_cache_hits_total"] == 1 and sum(tri.calls.values()) > calls | |
| assert other["rttm"] == first["rttm"].replace(f" {it['sid']} ", " other_name ") | |
| def test_idempotency_key(fake_reference, stack): | |
| """The same Idempotency-Key and request again answers with the same job (Idempotent-Replayed: true) without | |
| a new job; the key with another request is 422; while its first request is still uploading, 409.""" | |
| items, ref = fake_reference | |
| tri = stack.triton() | |
| svc = stack.service(tri, "--upload-idle-s", "2") | |
| it = ok_items(items, ref, 1)[0] | |
| data = Path(it["path"]).read_bytes() | |
| url = svc.url + "/v1/jobs?" + urllib.parse.urlencode(dict(name="a.flac", uri=it["sid"])) | |
| k = {"Idempotency-Key": "upload-1"} | |
| code, out, h1 = post_audio(url, data, headers=k) | |
| assert code == 202 and "Idempotent-Replayed" not in h1 | |
| code, out2, h2 = post_audio(url, data, headers=k) | |
| assert code == 202 and out2 == out and h2["Idempotent-Replayed"] == "true" | |
| # another uri, same key | |
| assert post_audio(url + "x", data, headers=k)[0] == 422 | |
| body = json.dumps({"path": it["path"], "uri": "p1"}).encode() | |
| pk = {"Content-Type": "application/json", "Idempotency-Key": "path-1"} | |
| j1, j2 = (json.loads(call(svc.url + "/v1/jobs", "POST", body, pk)[1])["job"] for _ in range(2)) | |
| assert j1 == j2 and call(svc.url + "/v1/jobs", "POST", body.replace(b"p1", b"p2"), pk)[0] == 422 | |
| assert post_audio(url, data, headers={"Idempotency-Key": "bad key"})[0] == 400 | |
| m = svc.metrics() | |
| assert m["anuvani_diar_jobs_submitted_total"] == 2 and m["anuvani_diar_upload_bytes_total"] == len(data) | |
| s = raw_request(svc.port, f"POST /v1/jobs?name=b.flac HTTP/1.1\r\nHost: x\r\nContent-Type: application/octet-stream\r\n" | |
| f"Idempotency-Key: slow\r\nContent-Length: {len(data)}\r\n\r\n", data[:1000]) | |
| wait_for(lambda: svc.metrics()["anuvani_diar_uploads_in_progress"] == 1, timeout=10) | |
| code, out, h = post_audio(svc.url + "/v1/jobs?name=b.flac", data, headers={"Idempotency-Key": "slow"}) | |
| assert code == 409 and h["Retry-After"] == "1" | |
| # the stalled first request | |
| assert read_reply(s).startswith(b"HTTP/1.1 408") | |
| s.close() | |
| # free again | |
| assert post_audio(svc.url + "/v1/jobs?name=b.flac", data, headers={"Idempotency-Key": "slow"})[0] == 202 | |
| assert svc.client().wait(json.loads(out2)["job"])["rttm"].encode() == (ref / f"{it['sid']}.rttm").read_bytes() | |
| def test_metrics_text(stack): | |
| svc = stack.service(stack.triton()) | |
| code, body, headers = call(svc.url + "/metrics") | |
| assert code == 200 and headers["content-type"].startswith("text/plain; version=0.0.4") | |
| lines = body.decode().splitlines() | |
| names = [line.split()[0] for line in lines if line[0] != "#"] | |
| assert all(f"# TYPE {n} " in body.decode() for n in names) and len(names) == len(set(names)) | |
| assert names[:13] == GATEWAY_METRICS and all(n.startswith("anuvani_diar_") for n in names) | |
| def test_request_ids_and_json_logs(stack): | |
| """Every answer carries X-Request-Id (the client's when it is plain, else a new one); every request is one JSON | |
| log line with that id, its method, path, status and time, never the query string.""" | |
| svc = stack.service(stack.triton()) | |
| _, _, h = call(svc.url + "/health", headers={"X-Request-Id": "abc-123"}) | |
| assert h["X-Request-Id"] == "abc-123" | |
| _, _, h = call(svc.url + "/health", headers={"X-Request-Id": "bad idé".encode("utf-8").decode("latin-1")}) | |
| rid = h["X-Request-Id"] | |
| assert len(rid) == 32 and int(rid, 16) >= 0 | |
| post_audio(svc.url + "/v1/jobs?name=secret.wav", b"not audio", headers={"X-Request-Id": "job-req"}) | |
| wait_for(lambda: any(r.get("request_id") == "job-req" for r in svc.logs()), timeout=10) | |
| recs = svc.logs() | |
| req = [r for r in recs if r.get("msg") == "request"] | |
| assert {"abc-123", rid, "job-req"} <= {r["request_id"] for r in req} | |
| r = next(r for r in req if r["request_id"] == "job-req") | |
| assert r["method"] == "POST" and r["path"] == "/v1/jobs" and r["status"] == 202 and len(r["job"]) == 32 and r["ms"] >= 0 | |
| assert "secret.wav" not in svc.log.read_text() | |
| assert all(set(x) >= {"ts", "level", "logger", "msg"} for x in recs) | |
| def test_openapi(stack): | |
| svc = stack.service(stack.triton()) | |
| code, body, _ = call(svc.url + "/openapi.json") | |
| spec = json.loads(body) | |
| assert code == 200 and {"/v1/jobs", "/v1/jobs/{job_id}", "/v1/jobs/{job_id}/rttm", "/v1/diarize", "/health", "/ready", | |
| "/metrics"} <= set(spec["paths"]) | |
| assert set(spec["paths"]["/v1/jobs"]["post"]["requestBody"]["content"]) == {"application/octet-stream", | |
| "multipart/form-data", "application/json"} | |
| assert spec["components"]["securitySchemes"]["bearer"]["scheme"] == "bearer" | |
| assert call(svc.url + "/docs")[0] == 200 | |
| # ---------------------------------------------------------------- spool and shutdown | |
| def test_spool_must_be_private(stack, tmp_path): | |
| tri = stack.triton() | |
| spool = tmp_path / "shared" | |
| spool.mkdir() | |
| os.chmod(spool, 0o775) | |
| svc = stack.service(tri, spool=spool, wait=False) | |
| assert svc.proc.wait(30) != 0 and "mode 0700" in svc.log.read_text() | |
| os.chmod(spool, 0o700) | |
| (spool / "spool-1-2-abc.wav").write_bytes(b"left over") | |
| (spool / "keep.txt").write_text("not ours") | |
| stack.service(tri, spool=spool) | |
| assert sorted(p.name for p in spool.iterdir()) == ["keep.txt"] | |
| def test_sigterm_drains_then_stops(fake_reference, stack): | |
| """SIGTERM: the job in progress finishes and can be read; meanwhile new jobs get 503 with Retry-After, /ready | |
| 503, /health 200. Then the service exits 0 with the spool empty.""" | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.1) | |
| svc = stack.service(tri, "--drain-linger-s", "2") | |
| c = svc.client() | |
| it = longest(ok_items(items, ref), ref, 1)[0] | |
| job = c.submit(it["path"], uri=it["sid"]) | |
| wait_for(lambda: c.status(job)["blocks_done"] > 0, timeout=30) | |
| svc.proc.send_signal(signal.SIGTERM) | |
| wait_for(lambda: svc.metrics()["anuvani_diar_draining"] == 1, timeout=10) | |
| code, out, h = post_audio(svc.url + "/v1/jobs?name=a.wav", b"x") | |
| assert code == 503 and h["Retry-After"] == "5" and b"shutting down" in out | |
| assert call(svc.url + "/ready")[0] == 503 and call(svc.url + "/health")[0] == 200 | |
| s = c.wait(job, poll=0.05) | |
| assert s["status"] == "done" and s["rttm"].encode() == (ref / f"{it['sid']}.rttm").read_bytes() | |
| assert svc.proc.wait(60) == 0 | |
| assert [p.name for p in svc.spool.iterdir()] == [] | |
| def test_sigterm_cancels_what_the_drain_cannot_finish(fake_reference, stack): | |
| """--drain-s 1 with a long job and a synchronous call waiting: both are cancelled after a second, the waiting | |
| caller is answered, the service exits 0 and the spool is empty.""" | |
| items, ref = fake_reference | |
| tri = stack.triton(delay=0.2) | |
| svc = stack.service(tri, "--drain-s", "1") | |
| c = svc.client() | |
| a, b = longest(ok_items(items, ref), ref, 2) | |
| job = c.submit(a["path"], uri=a["sid"]) | |
| answer = [] | |
| th = threading.Thread(target=lambda: answer.append(post_audio(svc.url + "/v1/diarize?name=b.flac", Path(b["path"]).read_bytes()))) | |
| th.start() | |
| wait_for(lambda: svc.metrics()["anuvani_diar_jobs_running"] == 2, timeout=30) | |
| t0 = time.monotonic() | |
| svc.proc.send_signal(signal.SIGTERM) | |
| assert svc.proc.wait(60) == 0 and time.monotonic() - t0 < 20 | |
| th.join() | |
| assert answer[0][0] == 422 and json.loads(answer[0][1])["error"] == "cancelled" | |
| assert [p.name for p in svc.spool.iterdir()] == [] | |
| assert any(r.get("msg") == "cancelling jobs left at shutdown" for r in svc.logs()) and job | |
| def test_seeded_random_requests_never_break_the_service(stack, tmp_path): | |
| """2,000 seeded random requests (methods, paths, query strings, headers, bodies: bytes, JSON, broken JSON, | |
| multipart, paths that do and do not exist): never a 5xx except 503/507, the service stays up, and the spool ends | |
| empty.""" | |
| import http.client | |
| import random as rnd | |
| tri = stack.triton() | |
| svc = stack.service(tri, "--max-upload-mb", "1", root=tmp_path) | |
| (tmp_path / "a.wav").write_bytes(b"RIFF" + os.urandom(100)) | |
| rng = rnd.Random(20261008) | |
| alphabet = "abc/._-%?&=+ \t\"'\\\x00\x7féक\U0001F600" | |
| def junk(n): | |
| return "".join(rng.choice(alphabet) for _ in range(rng.randint(0, n))) | |
| paths = ["/v1/jobs", "/v1/jobs/{}", "/v1/jobs/{}/rttm/{}", "/v1/diarize", "/health", "/ready", "/metrics", "/{}", | |
| "/v1/{}/{}", "/docs", "/openapi.json"] | |
| seen = set() | |
| t0 = time.monotonic() | |
| for _ in range(2000): | |
| method = rng.choice(["GET", "POST", "DELETE", "PUT", "PATCH", "HEAD"]) | |
| path = rng.choice(paths).format(*(urllib.parse.quote(junk(30), safe="") for _ in range(2))) | |
| q = urllib.parse.urlencode({k: junk(40) for k in rng.sample(["name", "uri", "y", "x"], rng.randint(0, 4))}) | |
| if rng.random() < 0.05: | |
| q += "&name=" + "x" * 300 | |
| kind = rng.randint(0, 5) | |
| if kind == 0: | |
| body, ctype = b"", None | |
| elif kind == 1: | |
| body, ctype = os.urandom(rng.randint(1, 65536)), "application/octet-stream" | |
| elif kind == 2: | |
| body, ctype = json.dumps({"path": rng.choice([str(tmp_path / "a.wav"), "/etc/passwd", junk(50), str(tmp_path)]), | |
| "uri": junk(20)}).encode(), "application/json" | |
| elif kind == 3: | |
| body, ctype = junk(200).encode("utf-8", "surrogatepass"), "application/json" | |
| elif kind == 4: | |
| parts = [(rng.choice(["file", "uri", junk(5)]), rng.choice([None, junk(12)]), os.urandom(rng.randint(0, 3000))) | |
| for _ in range(rng.randint(0, 3))] | |
| body, ctype = multipart(parts, boundary=rng.choice(["b0undary", "x" * 70, "q"])) | |
| if rng.random() < 0.3: | |
| body = body[: rng.randint(0, len(body))] | |
| if rng.random() < 0.2: | |
| ctype = "multipart/form-data" | |
| else: | |
| body, ctype = os.urandom(rng.randint(1, 4096)), rng.choice(["text/plain", junk(20), "application/json; x"]) | |
| headers = {"Content-Length": str(len(body))} | |
| if ctype is not None: | |
| headers["Content-Type"] = ctype.encode("latin-1", "replace").decode("latin-1") | |
| if rng.random() < 0.2: | |
| headers["Authorization"] = "Bearer " + junk(20).encode("latin-1", "replace").decode("latin-1") | |
| if rng.random() < 0.1: | |
| headers["Idempotency-Key"] = junk(10).encode("latin-1", "replace").decode("latin-1") | |
| c = http.client.HTTPConnection("127.0.0.1", svc.port, timeout=30) | |
| try: | |
| c.request(method, path + ("?" + q if q else ""), body=body, headers=headers) | |
| status = c.getresponse().status | |
| except (http.client.HTTPException, OSError): | |
| status = None | |
| finally: | |
| c.close() | |
| seen.add(status) | |
| assert status is None or status < 500 or status in (503, 507), (method, path, q, status) | |
| elapsed = time.monotonic() - t0 | |
| assert svc.proc.poll() is None and call(svc.url + "/health")[0] == 200 | |
| wait_for(lambda: svc.metrics()["anuvani_diar_jobs_queued"] + svc.metrics()["anuvani_diar_jobs_running"] == 0, timeout=120) | |
| assert svc.metrics()["anuvani_diar_spool_bytes"] == 0 and [p.name for p in svc.spool.iterdir()] == [] | |
| assert {200, 202, 400, 404}.issubset(seen) | |
| assert elapsed < 60, elapsed | |
| assert "Traceback" not in svc.log.read_text() | |
| # ---------------------------------------------------------------- the image (opt-in: ANUVANI_DIAR_SERVICE_IMAGE=tag) | |
| class Container: | |
| """The service image on the host network (it reaches the stand-in Triton on 127.0.0.1), as this user, with the | |
| test inputs mounted read-only at their own paths and the records directory writable.""" | |
| def __init__(self, image, triton, tmp, inputs, *extra): | |
| self.port = free_port() | |
| self.url = f"http://127.0.0.1:{self.port}" | |
| self.records = tmp / f"records-{self.port}" | |
| self.records.mkdir() | |
| self.name = f"anuvani-diar-service-test-{self.port}" | |
| mounts = sum((["-v", f"{d}:{d}:ro"] for d in sorted({str(Path(p).parent) for p in inputs})), []) | |
| cmd = ["docker", "run", "-d", "--name", self.name, "--network", "host", "--user", f"{os.getuid()}:{os.getgid()}", | |
| "-v", f"{self.records}:/records", *mounts, image, "--listen", f"127.0.0.1:{self.port}", "--triton", triton.url, | |
| "--gpus", "0,1", "--slots-per-gpu", "2", "--spool", "/tmp/spool", "--dump-records", "/records", | |
| "--decode-threads", "3", "--post-threads", "4", "--drain-linger-s", "0", "--allow-path-root", "/", *extra] | |
| subprocess.run(cmd, check=True, capture_output=True) | |
| deadline = time.monotonic() + 60 | |
| while time.monotonic() < deadline: | |
| if self._up(): | |
| return | |
| time.sleep(0.1) | |
| raise RuntimeError(subprocess.run(["docker", "logs", self.name], capture_output=True, text=True).stderr[-3000:]) | |
| def _up(self): | |
| try: | |
| import urllib.request | |
| urllib.request.urlopen(self.url + "/health", timeout=1) | |
| return True | |
| except OSError: | |
| return False | |
| def client(self, **kw): | |
| from anuvani_diar import Client | |
| return Client(self.url, timeout=120, **kw) | |
| def metrics(self): | |
| return Service.metrics(self) | |
| def stop(self): | |
| """docker stop: SIGTERM to tini, which forwards it; the service drains and exits 0.""" | |
| subprocess.run(["docker", "stop", "-t", "60", self.name], capture_output=True) | |
| def test_image_identical_to_the_release(fake_reference, tmp_path): | |
| """The built image (Ubuntu 22.04, the extension compiled inside it, the staged numeric binaries) gives the same | |
| bytes as the release with the stand-in networks: every RTTM, error and front-end record; then `docker stop` | |
| drains it and it exits 0.""" | |
| items, ref = fake_reference | |
| image = os.environ["ANUVANI_DIAR_SERVICE_IMAGE"] | |
| tri = FakeTriton().start() | |
| c = None | |
| try: | |
| c = Container(image, tri, tmp_path, [it["path"] for it in items]) | |
| got = run_jobs(c, items, ref, uploads=12, seed=21) | |
| assert compare(items, got, ref, c.records) == [] | |
| it = ok_items(items, ref, 1)[0] | |
| code, out, _ = post_audio(c.url + "/v1/diarize?" + urllib.parse.urlencode(dict(name="x.flac", uri=it["sid"])), | |
| Path(it["path"]).read_bytes()) | |
| assert code == 200 and out == (ref / f"{it['sid']}.rttm").read_bytes() | |
| assert call(c.url + "/ready")[0] == 200 and c.metrics()["anuvani_diar_spool_bytes"] == 0 | |
| user = subprocess.run(["docker", "exec", c.name, "python3", "-c", "import os, sys; print(os.getuid(), sys.version.split()[0])"], | |
| capture_output=True, text=True).stdout.split() | |
| assert user[0] == str(os.getuid()) | |
| c.stop() | |
| state = subprocess.run(["docker", "inspect", "-f", "{{.State.Status}} {{.State.ExitCode}}", c.name], | |
| capture_output=True, text=True).stdout.split() | |
| assert state == ["exited", "0"], state | |
| logs = subprocess.run(["docker", "logs", c.name], capture_output=True, text=True).stderr | |
| assert '"msg": "draining"' in logs and '"msg": "stopped"' in logs | |
| finally: | |
| if c is not None: | |
| subprocess.run(["docker", "rm", "-f", c.name], capture_output=True) | |
| tri.stop() | |