| |
| |
| """Client for anuvani-diar-gateway: the same one-line use as the local Diarizer, with the work on a server. |
| |
| from anuvani_diar import Client |
| c = Client("http://127.0.0.1:8180") |
| r = c("meeting.wav") # upload, wait, -> Diarization |
| r.save_rttm("meeting.rttm") |
| jobs = [c.submit(p) for p in paths] # many files: submit all, then collect |
| for j in jobs: print(c.result(j).uri) |
| |
| Standard library only (urllib); uploads stream from disk. The client rides out a busy or restarting server: it |
| backs off on 429/503 and on refused connections, and if the gateway restarts and no longer knows a job, `wait` |
| submits the same input again (the result is the same, byte for byte) and follows the new job. |
| """ |
| import json |
| import os |
| import threading |
| import time |
| from collections import OrderedDict |
| import urllib.error |
| import urllib.parse |
| import urllib.request |
| from pathlib import Path |
| from typing import Optional, Union |
|
|
| from .api import Diarization |
|
|
|
|
| class GatewayError(RuntimeError): |
| pass |
|
|
|
|
| class _NoRedirect(urllib.request.HTTPRedirectHandler): |
| """The gateway never redirects. urllib would follow a redirect with every header of the request, the token |
| included, to whatever host it names; a 3xx becomes an error instead.""" |
|
|
| def redirect_request(self, req, fp, code, msg, headers, newurl): |
| return None |
|
|
|
|
| _OPENER = urllib.request.build_opener(_NoRedirect) |
|
|
|
|
| class Client: |
| def __init__(self, url: str = "http://127.0.0.1:8180", timeout: float = 600.0, retries: int = 30, |
| token: Optional[str] = None): |
| self.url = url.rstrip("/") |
| self.timeout = timeout |
| self.retries = retries |
| self.token = token or os.environ.get("ANUVANI_DIAR_TOKEN") |
| |
| self._inputs = OrderedDict() |
| |
| self._alias = OrderedDict() |
| |
| self._lock = threading.Lock() |
| |
| self._resubmit = threading.Lock() |
|
|
| |
| _REMEMBER = 100_000 |
|
|
| def _remember(self, table, key, value): |
| with self._lock: |
| table[key] = value |
| while len(table) > self._REMEMBER: |
| table.popitem(last=False) |
|
|
| def _current(self, job: str) -> str: |
| """The id a job is known by now (it changes when the job is submitted again after a gateway restart).""" |
| with self._lock: |
| while job in self._alias: |
| job = self._alias[job] |
| return job |
|
|
| def _call(self, method, path, data=None, headers=None): |
| headers = dict(headers or {}) |
| if self.token: |
| headers["Authorization"] = f"Bearer {self.token}" |
| req = urllib.request.Request(self.url + path, data=data, method=method, headers=headers) |
| for attempt in range(self.retries + 1): |
| try: |
| with _OPENER.open(req, timeout=self.timeout) as r: |
| return r.status, r.read(), dict(r.headers) |
| except urllib.error.HTTPError as e: |
| body = e.read() |
| |
| if e.code in (429, 503) and attempt < self.retries: |
| time.sleep(float(e.headers.get("Retry-After", "2"))) |
| if hasattr(data, "seek"): |
| data.seek(0) |
| continue |
| try: |
| msg = json.loads(body).get("error", body.decode(errors="replace")) |
| except Exception: |
| msg = body.decode(errors="replace") |
| raise GatewayError(f"HTTP {e.code}: {msg}") from None |
| except (urllib.error.URLError, ConnectionError) as e: |
| |
| if attempt < self.retries: |
| time.sleep(min(2.0 * (attempt + 1), 10.0)) |
| if hasattr(data, "seek"): |
| data.seek(0) |
| continue |
| raise GatewayError(str(e)) from None |
|
|
| def submit(self, path: Union[str, Path], uri: Optional[str] = None, server_path: bool = False) -> str: |
| """Start a job; returns its id. `server_path=True` sends only the path (the file is on the server's disk).""" |
| path = Path(path) |
| if server_path: |
| body = json.dumps({"path": str(path), **({"uri": uri} if uri else {})}).encode() |
| _, out, _ = self._call("POST", "/v1/jobs", body, {"Content-Type": "application/json"}) |
| else: |
| q = urllib.parse.urlencode({"name": path.name, **({"uri": uri} if uri else {})}) |
| with open(path, "rb") as f: |
| _, out, _ = self._call("POST", f"/v1/jobs?{q}", f, {"Content-Type": "application/octet-stream", |
| "Content-Length": str(os.path.getsize(path))}) |
| job = json.loads(out)["job"] |
| self._remember(self._inputs, job, (path, uri, server_path)) |
| return job |
|
|
| def status(self, job: str) -> dict: |
| return json.loads(self._call("GET", f"/v1/jobs/{self._current(job)}")[1]) |
|
|
| def cancel(self, job: str) -> dict: |
| return json.loads(self._call("DELETE", f"/v1/jobs/{self._current(job)}")[1]) |
|
|
| def wait(self, job: str, poll: float = 0.5) -> dict: |
| """Poll until the job ends; its final status. A job the server no longer knows (it restarted) is submitted |
| again from the same input and followed under its new id; the original id keeps working in every call.""" |
| while True: |
| cur = self._current(job) |
| try: |
| s = self.status(cur) |
| except GatewayError as e: |
| with self._lock: |
| inp = self._inputs.get(cur) |
| if not str(e).startswith("HTTP 404") or inp is None: |
| raise |
| with self._resubmit: |
| |
| if self._current(job) == cur: |
| self._remember(self._alias, cur, self.submit(*inp)) |
| continue |
| if s["status"] not in ("queued", "running"): |
| return s |
| time.sleep(poll) |
|
|
| def result(self, job: str, path: str = "") -> Diarization: |
| s = self.wait(job) |
| err = s.get("error") or (None if s["status"] == "done" else s["status"]) |
| return Diarization(uri=s["uri"], path=path, rttm=s.get("rttm"), error=err, duration_s=s.get("duration_s"), |
| timing=s.get("timing") or {}) |
|
|
| def __call__(self, path: Union[str, Path], uri: Optional[str] = None) -> Diarization: |
| return self.result(self.submit(path, uri), str(path)) |
|
|
| def ready(self) -> bool: |
| """One attempt, no retries: is the gateway up with every network model ready?""" |
| req = urllib.request.Request(self.url + "/ready", headers={"Authorization": f"Bearer {self.token}"} if self.token else {}) |
| try: |
| with _OPENER.open(req, timeout=min(self.timeout, 10)) as r: |
| return r.status == 200 |
| except (urllib.error.URLError, OSError): |
| return False |
|
|
