| |
| |
| """The gateway (crates/diar-gateway) as a service, on CPU. |
| |
| A stand-in Triton (fake_triton.py) serves the stand-in networks of fake_nn.py, and every RTTM, error and front-end |
| record the gateway returns must be byte-identical to the release's own `analyse` and `indic_mode` run with the same |
| networks (conftest.fake_reference). Around that: Triton going away and coming back mid-run, a model that stays away, model |
| errors, admission limits, uploads, path rules, the token, cancellation, the synchronous call, the result cache, |
| metrics, the spool and shutdown. |
| |
| Needs the gateway binary (`make build`, or $ANUVANI_DIAR_GATEWAY). |
| """ |
| import json |
| import os |
| import random |
| import signal |
| import socket |
| import subprocess |
| import time |
| import urllib.error |
| import urllib.parse |
| import urllib.request |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
|
|
| import numpy as np |
| import pytest |
|
|
| from conftest import PROJ |
| from fake_triton import FakeTriton |
|
|
| BIN = Path(os.environ.get("ANUVANI_DIAR_GATEWAY", PROJ / "target/release/anuvani-diar-gateway")) |
| pytestmark = pytest.mark.skipif(not BIN.exists(), reason="build the gateway first (make build)") |
| OOM = "RuntimeError: CUDA out of memory (simulated)" |
|
|
|
|
| def free_port(): |
| with socket.socket() as s: |
| s.bind(("127.0.0.1", 0)) |
| return s.getsockname()[1] |
|
|
|
|
| def call(url, method="GET", data=None, headers=None, timeout=60): |
| """(status, body, headers); header names are looked up case-insensitively.""" |
| req = urllib.request.Request(url, data=data, method=method, headers=headers or {}) |
| try: |
| with urllib.request.urlopen(req, timeout=timeout) as r: |
| return r.status, r.read(), r.headers |
| except urllib.error.HTTPError as e: |
| return e.code, e.read(), e.headers |
|
|
|
|
| class Gateway: |
| 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"gateway-{self.port}.log" |
| cmd = [str(BIN), "--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", *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"} |
| self.proc = subprocess.Popen(cmd, stdout=open(self.log, "w"), stderr=subprocess.STDOUT, env=env) |
| if wait: |
| deadline = time.monotonic() + 30 |
| while time.monotonic() < deadline and self.proc.poll() is None: |
| try: |
| urllib.request.urlopen(self.url + "/health", timeout=1) |
| return |
| except OSError: |
| time.sleep(0.05) |
| raise RuntimeError(f"the gateway did not start: {self.log.read_text()[-2000:]}") |
|
|
| def stop(self): |
| if self.proc.poll() is None: |
| self.proc.send_signal(signal.SIGTERM) |
| try: |
| self.proc.wait(20) |
| except subprocess.TimeoutExpired: |
| self.proc.kill() |
| self.proc.wait() |
|
|
| def metrics(self): |
| _, body, _ = call(self.url + "/metrics") |
| 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) |
|
|
|
|
| @pytest.fixture |
| def stack(tmp_path): |
| """Start fake Tritons and gateways on demand; stop them all afterwards.""" |
| started = [] |
|
|
| class Stack: |
| def triton(self, **kw): |
| t = FakeTriton(**kw).start() |
| started.append(t) |
| return t |
|
|
| def gateway(self, triton, *extra, **kw): |
| g = Gateway(triton, tmp_path, *extra, **kw) |
| started.append(g) |
| return g |
| yield Stack() |
| for s in reversed(started): |
| s.stop() |
|
|
|
|
| def ok_items(items, ref, n=None): |
| log = json.loads((ref / "log.json").read_text()) |
| ok = [it for it in items if not log[it["sid"]]] |
| return ok if n is None else ok[:n] |
|
|
|
|
| def longest(items, ref, n): |
| """The `n` items with the most 16 s windows (from the reference records).""" |
| return sorted(items, key=lambda it: -np.load(ref / f"{it['sid']}.npz")["P"].shape[0])[:n] |
|
|
|
|
| def run_jobs(gw, items, ref, uploads=8, seed=0, threads=16): |
| """Every item once, in a shuffled order, concurrently; the first `uploads` working files (in the items' own |
| order) are uploaded, the rest named by path. Returns {sid: final job status}.""" |
| c = gw.client(retries=200) |
| order = list(items) |
| random.Random(seed).shuffle(order) |
| upload = {it["sid"] for it in ok_items(items, ref, uploads)} |
|
|
| def one(it): |
| job = c.submit(it["path"], uri=it["sid"], server_path=it["sid"] not in upload) |
| return it["sid"], c.wait(job, poll=0.05) |
| with ThreadPoolExecutor(threads) as ex: |
| return dict(ex.map(one, order)) |
|
|
|
|
| def compare(items, got, ref, records): |
| """Differences between the gateway's answers and the release's: errors, the RTTM, the dumped records.""" |
| log = json.loads((ref / "log.json").read_text()) |
| bad = [] |
| for it in items: |
| sid = it["sid"] |
| want, r = log[sid], got[sid] |
| if "error" in want: |
| if r["error"] != want["error"]: |
| bad.append((sid, "error", want["error"], r["error"])) |
| continue |
| if r["status"] != "done": |
| bad.append((sid, r["status"], r["error"])) |
| continue |
| if (r["rttm"] or "").encode() != (ref / f"{sid}.rttm").read_bytes(): |
| bad.append((sid, "rttm")) |
| d = np.load(ref / f"{sid}.npz") |
| meta = json.loads((records / f"{sid}.json").read_text()) |
| same = (np.fromfile(records / f"{sid}.P", dtype="<u2").tobytes() == np.ascontiguousarray(d["P"]).view(np.uint16).tobytes() |
| and np.fromfile(records / f"{sid}.masks", dtype=np.uint8).tobytes() == d["masks"].tobytes() |
| and np.fromfile(records / f"{sid}.emb", dtype="<f4").tobytes() == d["emb"].tobytes() |
| and meta["chunks"] == d["P"].shape[0] |
| and np.float64(meta["duration_s"]).tobytes() == d["duration_s"].tobytes() |
| and np.array(meta["frames_sw"], np.float64).tobytes() == np.asarray(d["frames_sw"], np.float64).tobytes()) |
| if not same: |
| bad.append((sid, "record")) |
| return bad |
|
|
|
|
| def wait_for(cond, timeout=60.0, step=0.05): |
| deadline = time.monotonic() + timeout |
| while time.monotonic() < deadline: |
| v = cond() |
| if v: |
| return v |
| time.sleep(step) |
| raise AssertionError("timed out") |
|
|
|
|
| def raw_request(port, head, body=b""): |
| s = socket.create_connection(("127.0.0.1", port), timeout=30) |
| s.sendall(head.encode() + body) |
| return s |
|
|
|
|
| |
|
|
| @pytest.mark.parametrize("schedule", ["fair", "fifo", "share"]) |
| def test_jobs_identical_to_the_release(fake_reference, stack, schedule): |
| items, ref = fake_reference |
| gw = stack.gateway(stack.triton(), "--schedule", schedule) |
| got = run_jobs(gw, items, ref, seed={"fair": 3, "fifo": 4, "share": 9}[schedule]) |
| assert compare(items, got, ref, gw.records) == [] |
| m = gw.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 |
|
|
|
|
| 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")) |
| gw = stack.gateway(t0, "--triton", f"{t0.url},{t1.url}") |
| got = run_jobs(gw, items, ref, seed=8) |
| assert compare(items, got, ref, gw.records) == [] |
| assert t0.calls["diar_nn_g0"] > 0 and t1.calls["diar_nn_g1"] > 0 |
| bad = stack.gateway(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).""" |
| items, ref = fake_reference |
| gw = stack.gateway(stack.triton()) |
| for rep in range(3): |
| got = run_jobs(gw, items, ref, uploads=12, seed=10 + rep, threads=64) |
| assert compare(items, got, ref, gw.records) == [] |
| assert gw.metrics()["anuvani_diar_cache_hits_total"] >= 12 |
|
|
|
|
| @pytest.mark.parametrize("schedule", ["fair", "share"]) |
| 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 (the long file stands still until it is done); share interleaves them (the long file keeps |
| advancing).""" |
| 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) |
| gw = stack.gateway(tri, "--schedule", schedule, "--gpus", "0", "--slots-per-gpu", "1") |
| c = gw.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) |
|
|
|
|
| |
|
|
| 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) |
| gw = stack.gateway(tri) |
| restart = tri.restart_in_background(after_s=1.5, down_s=2.0, loading_s=2.0) |
| got = run_jobs(gw, items, ref, seed=5) |
| restart.join() |
| assert compare(items, got, ref, gw.records) == [] |
| assert gw.metrics()["anuvani_diar_triton_retries_total"] > 0 |
|
|
|
|
| def test_jobs_survive_a_gateway_restart(fake_reference, stack, tmp_path): |
| """The gateway killed 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.""" |
| import threading |
| items, ref = fake_reference |
| tri = stack.triton(delay=0.03) |
| gw = stack.gateway(tri) |
| spool, port = gw.spool, gw.port |
|
|
| def crash_and_restart(): |
| time.sleep(1.5) |
| gw.proc.kill() |
| gw.proc.wait() |
| time.sleep(1.0) |
| restarted.append(stack.gateway(tri, spool=spool, port=port)) |
| restarted = [] |
| th = threading.Thread(target=crash_and_restart) |
| th.start() |
| got = run_jobs(gw, items, ref, seed=6) |
| th.join() |
| new = restarted[0] |
| |
| for f in gw.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_gateway_restarts(fake_reference, stack): |
| """After a resubmission the caller's original job id keeps working: wait, result (twice), status, cancel. |
| Client.ready answers at once (no retries) while the gateway is down.""" |
| items, ref = fake_reference |
| tri = stack.triton(delay=0.05) |
| gw = stack.gateway(tri) |
| c = gw.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) |
| gw.proc.kill() |
| gw.proc.wait() |
| t0 = time.monotonic() |
| assert c.ready() is False and time.monotonic() - t0 < 5 |
| new = stack.gateway(tri, spool=gw.spool, port=gw.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() |
| gw = stack.gateway(tri) |
| assert call(gw.url + "/ready")[0] == 200 |
| tri.stop() |
| wait_for(lambda: call(gw.url + "/ready")[0] == 503, timeout=10) |
| code, body, _ = call(gw.url + "/ready") |
| |
| assert b"diar_nn_g0" in body and b"127.0.0.1" not in body |
| tri.start() |
| wait_for(lambda: call(gw.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) |
| gw = stack.gateway(tri, "--triton-retry-s", "2") |
| c = gw.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 gw.metrics()["anuvani_diar_jobs_failed_total"] == len(jobs) |
|
|
|
|
| def test_model_errors_are_final(fake_reference, stack): |
| """An error the model returns while it is ready (bad input, out of memory) fails the file at once, no retry.""" |
| items, ref = fake_reference |
| tri = stack.triton() |
| for m in ("diar_nn_g0", "diar_nn_g1"): |
| tri.fail(m, OOM) |
| gw = stack.gateway(tri) |
| c = gw.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 gw.metrics()["anuvani_diar_triton_retries_total"] == 0 |
|
|
|
|
| def test_a_dead_model_process_is_retried(fake_reference, stack): |
| """Triton's Python backend answers "Stub process ... is not healthy" for a model whose process died, then starts |
| a new one: such requests are sent again, and the bytes are unchanged.""" |
| 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) |
| gw = stack.gateway(tri, "--gpus", "0") |
| sub = ok_items(items, ref, 4) |
| got = run_jobs(gw, sub, ref, uploads=0) |
| assert compare(sub, got, ref, gw.records) == [] |
| assert gw.metrics()["anuvani_diar_triton_retries_total"] == 3 |
|
|
|
|
| |
|
|
| def test_job_queue_limit(fake_reference, stack): |
| items, ref = fake_reference |
| tri = stack.triton(delay=0.3) |
| gw = stack.gateway(tri, "--max-jobs", "2") |
| c = gw.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(gw.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 gw.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(gw.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 |
| gw = stack.gateway(stack.triton(), "--max-jobs", "2") |
| sub = ok_items(items, ref, 6) |
| got = run_jobs(gw, sub, ref, uploads=2, threads=6) |
| assert compare(sub, got, ref, gw.records) == [] |
| assert gw.metrics()["anuvani_diar_jobs_rejected_total"] > 0 |
|
|
|
|
| def test_upload_size_limit(stack, tmp_path): |
| gw = stack.gateway(stack.triton(), "--max-upload-mb", "1") |
| code, out, _ = call(gw.url + "/v1/jobs?name=big.wav", "POST", os.urandom(2 * 1024 * 1024), |
| {"Content-Type": "application/octet-stream"}) |
| assert code == 413 and b"larger than 1 MB" in out |
| assert [p.name for p in gw.spool.iterdir()] == [] and gw.metrics()["anuvani_diar_spool_bytes"] == 0 |
|
|
|
|
| def test_stalled_upload_and_upload_limit(stack): |
| gw = stack.gateway(stack.triton(), "--max-uploads", "1", "--upload-idle-s", "1") |
| s = raw_request(gw.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: gw.metrics()["anuvani_diar_uploads_in_progress"] == 1, timeout=10) |
| code, out, _ = call(gw.url + "/v1/jobs?name=b.wav", "POST", b"x" * 10, {"Content-Type": "application/octet-stream"}) |
| assert code == 429 and b"too many uploads" in out |
| reply = s.recv(4096) |
| s.close() |
| assert reply.startswith(b"HTTP/1.1 408") and b"upload stalled" in reply |
| wait_for(lambda: gw.metrics()["anuvani_diar_uploads_in_progress"] == 0, timeout=10) |
| assert [p.name for p in gw.spool.iterdir()] == [] |
|
|
|
|
| def test_a_low_sample_rate_cannot_resample_past_the_limit(stack, tmp_path): |
| """600,000 samples declared at 1 Hz: 2.4 MB uploaded, 38.4 GB at 16 kHz, refused before anything is resampled.""" |
| import io |
| import soundfile as sf |
| tri = stack.triton() |
| gw = stack.gateway(tri) |
| buf = io.BytesIO() |
| sf.write(buf, np.full(600000, 0.1, np.float32), 1, subtype="FLOAT", format="WAV") |
| code, out, _ = call(gw.url + "/v1/diarize?name=low.wav", "POST", buf.getvalue(), {"Content-Type": "application/octet-stream"}) |
| 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_client_does_not_follow_redirects(): |
| """A redirect would carry the token to another host; the client stops at it.""" |
| import threading |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| from anuvani_diar.client import Client, GatewayError |
| seen = [] |
|
|
| class Sink(BaseHTTPRequestHandler): |
| def do_GET(self): |
| seen.append(self.headers.get("Authorization")) |
| self.send_response(200) |
| self.send_header("Content-Length", "2") |
| self.end_headers() |
| self.wfile.write(b"{}") |
|
|
| def log_message(self, *a): |
| pass |
| sink = ThreadingHTTPServer(("127.0.0.1", 0), Sink) |
|
|
| class Redirect(BaseHTTPRequestHandler): |
| def do_GET(self): |
| self.send_response(302) |
| self.send_header("Location", f"http://127.0.0.1:{sink.server_address[1]}/steal") |
| self.send_header("Content-Length", "0") |
| self.end_headers() |
|
|
| def log_message(self, *a): |
| pass |
| red = ThreadingHTTPServer(("127.0.0.1", 0), Redirect) |
| for srv in (sink, red): |
| threading.Thread(target=srv.serve_forever, daemon=True).start() |
| try: |
| c = Client(f"http://127.0.0.1:{red.server_address[1]}", token="s3cret", retries=0) |
| with pytest.raises(GatewayError, match="HTTP 302"): |
| c.status("abc") |
| assert seen == [] |
| finally: |
| sink.shutdown() |
| red.shutdown() |
|
|
|
|
| def test_upload_names_are_checked(stack): |
| gw = stack.gateway(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:")): |
| code, out, _ = call(gw.url + "/v1/jobs?" + urllib.parse.urlencode(q), "POST", b"x", {"Content-Type": "application/octet-stream"}) |
| assert code == 400 and msg in out, (q, out) |
|
|
|
|
| |
|
|
| 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() |
| gw = stack.gateway(stack.triton(), root=root) |
|
|
| def post(path, **kw): |
| return call(gw.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)): |
| 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 |
| assert call(gw.url + "/v1/jobs", "POST", b"{not json", {"Content-Type": "application/json"})[0] == 400 |
| closed = stack.gateway(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 |
| gw = stack.gateway(stack.triton(), token="s3cret") |
| assert call(gw.url + "/health")[0] == 200 |
| for h in ({}, {"Authorization": "Bearer wrong"}, {"Authorization": "s3cret"}, {"Authorization": "Bearer s3cre"}, |
| {"Authorization": "Bearer s3cret2"}): |
| for path in ("/ready", "/metrics", "/v1/jobs/0"): |
| assert call(gw.url + path, headers=h)[0] == 401, (h, path) |
| assert call(gw.url + "/ready", headers={"Authorization": "Bearer s3cret"})[0] == 200 |
| sub = ok_items(items, ref, 2) |
| from anuvani_diar import Client |
| c = Client(gw.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, gw.records) == [] |
|
|
|
|
| |
|
|
| def test_cancel(fake_reference, stack): |
| items, ref = fake_reference |
| tri = stack.triton(delay=0.2) |
| gw = stack.gateway(tri) |
| c = gw.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) |
| assert c.cancel(job) == {"job": job, "cancelling": True} |
| s = c.wait(job, poll=0.05) |
| assert s["status"] == "cancelled" and s["error"] == "cancelled" |
| |
| assert tri.calls["diar_nn_g0"] + tri.calls["diar_nn_g1"] < 2 * s["blocks"] |
| assert c.cancel(job) == {"job": job, "cancelling": False} |
| m = gw.metrics() |
| assert m["anuvani_diar_jobs_cancelled_total"] == 1 and m["anuvani_diar_spool_bytes"] == 0 |
| assert [p.name for p in gw.spool.iterdir()] == [] |
|
|
|
|
| def test_lookups(stack, fake_reference): |
| items, ref = fake_reference |
| tri = stack.triton(delay=0.2) |
| gw = stack.gateway(tri) |
| assert call(gw.url + "/v1/jobs/feedbeef")[0] == 404 |
| assert call(gw.url + "/v1/jobs/feedbeef", "DELETE")[0] == 404 |
| assert call(gw.url + "/v1/jobs/feedbeef/rttm")[0] == 404 |
| c = gw.client() |
| it = longest(ok_items(items, ref), ref, 1)[0] |
| job = c.submit(it["path"], uri=it["sid"], server_path=True) |
| |
| assert len(job) == 32 and int(job, 16) >= 0 |
| code, out, _ = call(gw.url + f"/v1/jobs/{job}/rttm") |
| assert code == 409 and b"job is" in out |
| c.wait(job, poll=0.05) |
| |
| assert call(gw.url + f"/v1/jobs/{job}/rttm/public")[0] == 404 |
| code, out, headers = call(gw.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") |
|
|
|
|
| def test_synchronous_call(fake_reference, stack): |
| items, ref = fake_reference |
| gw = stack.gateway(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 = call(gw.url + "/v1/diarize?" + q, "POST", Path(it["path"]).read_bytes(), |
| {"Content-Type": "application/octet-stream"}) |
| assert code == 200 and out == (ref / f"{it['sid']}.rttm").read_bytes() and len(headers["x-job"]) == 32 |
| code, out, _ = call(gw.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) |
| gw = stack.gateway(tri) |
| it = longest(ok_items(items, ref), ref, 1)[0] |
| data = Path(it["path"]).read_bytes() |
| s = raw_request(gw.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: gw.metrics()["anuvani_diar_jobs_running"] == 1, timeout=30) |
| s.close() |
| wait_for(lambda: gw.metrics()["anuvani_diar_jobs_cancelled_total"] == 1, timeout=60) |
| wait_for(lambda: gw.metrics()["anuvani_diar_spool_bytes"] == 0, timeout=10) |
|
|
|
|
| def test_result_cache(fake_reference, stack): |
| items, ref = fake_reference |
| tri = stack.triton() |
| gw = stack.gateway(tri) |
| c = gw.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 gw.metrics()["anuvani_diar_cache_hits_total"] == 1 and sum(tri.calls.values()) == calls |
| assert again["rttm"] == first["rttm"] |
| |
| other = c.wait(c.submit(it["path"], uri="other_name"), poll=0.05) |
| assert gw.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_metrics_text(stack): |
| gw = stack.gateway(stack.triton()) |
| code, body, headers = call(gw.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)) == 13 |
| assert all(n.startswith("anuvani_diar_") for n in names) |
|
|
|
|
| |
|
|
| def test_spool_must_be_private(stack, tmp_path): |
| tri = stack.triton() |
| spool = tmp_path / "shared" |
| spool.mkdir() |
| os.chmod(spool, 0o775) |
| gw = stack.gateway(tri, spool=spool, wait=False) |
| assert gw.proc.wait(30) != 0 and "mode 0700" in gw.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.gateway(tri, spool=spool) |
| assert sorted(p.name for p in spool.iterdir()) == ["keep.txt"] |
|
|
|
|
| def test_sigterm_stops_and_cleans_the_spool(fake_reference, stack): |
| items, ref = fake_reference |
| tri = stack.triton(delay=0.2) |
| gw = stack.gateway(tri) |
| c = gw.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) |
| gw.proc.send_signal(signal.SIGTERM) |
| assert gw.proc.wait(60) == 0 |
| assert [p.name for p in gw.spool.iterdir()] == [] |
|
|
|
|
| def test_seeded_random_requests_never_break_the_gateway(stack, tmp_path): |
| """2,000 seeded random requests (methods, paths, query strings, headers, bodies: bytes, JSON, broken JSON, |
| paths that do and do not exist): never a 5xx except 503/507, the gateway stays up, and the spool ends empty.""" |
| import http.client |
| import random as rnd |
| tri = stack.triton() |
| gw = stack.gateway(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/{}/{}"] |
| seen = set() |
| 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, 4) |
| 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" |
| 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") |
| c = http.client.HTTPConnection("127.0.0.1", gw.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) |
| assert gw.proc.poll() is None and call(gw.url + "/health")[0] == 200 |
| wait_for(lambda: gw.metrics()["anuvani_diar_jobs_queued"] + gw.metrics()["anuvani_diar_jobs_running"] == 0, timeout=120) |
| assert gw.metrics()["anuvani_diar_spool_bytes"] == 0 and [p.name for p in gw.spool.iterdir()] == [] |
| assert {200, 202, 400, 404}.issubset(seen) |
|
|
