| """JSONL protocol between the benchmark orchestrator and model workers. |
| |
| Each model family runs in its own virtual environment (dependency pins conflict, see |
| docs/benchmarks/B1-B2-tasarim.md §1.3), so the orchestrator and the workers only share |
| files. This module is the single definition of those files and must stay stdlib-only: |
| it is imported inside every worker environment. |
| |
| Layout for one worker invocation:: |
| |
| requests.jsonl one ``WorkRequest`` per line (written by the orchestrator) |
| results.jsonl one ``WorkResult`` per line (appended by the worker) |
| results.summary.json one ``WorkerSummary`` (written by the worker at the end) |
| |
| Results are appended one line at a time and flushed + fsynced, so an interrupted worker |
| leaves a valid prefix plus at most one partial trailing line. On restart the worker |
| calls :func:`load_results` with ``repair=True``, which drops the partial line, and |
| skips every request whose ``clip_id`` already has a result. Failed items are recorded |
| with an ``error`` and are never retried silently: callers decide via ``retry_failed``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from collections.abc import Iterable, Iterator, Mapping |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import IO |
|
|
| __all__ = [ |
| "ProtocolError", |
| "WorkRequest", |
| "WorkResult", |
| "WorkerSummary", |
| "append_result", |
| "drop_results", |
| "load_requests", |
| "load_results", |
| "pending_requests", |
| "read_summary", |
| "write_requests", |
| "write_summary", |
| ] |
|
|
|
|
| class ProtocolError(ValueError): |
| """A protocol file is malformed or inconsistent.""" |
|
|
|
|
| @dataclass(frozen=True) |
| class WorkRequest: |
| """One unit of work for a worker. |
| |
| Attributes: |
| clip_id: Content-addressed clip id; unique within a requests file. |
| audio_path: Absolute path of the input audio file. |
| params: Worker-specific string parameters (e.g. ``language``, ``context``, |
| ``out_path``). Values are strings so the format stays trivially portable. |
| """ |
|
|
| clip_id: str |
| audio_path: str |
| params: Mapping[str, str] = field(default_factory=dict) |
|
|
| def to_json(self) -> str: |
| return json.dumps( |
| {"clip_id": self.clip_id, "audio_path": self.audio_path, "params": dict(self.params)}, |
| ensure_ascii=False, |
| sort_keys=True, |
| ) |
|
|
| @staticmethod |
| def from_obj(obj: object, where: str) -> WorkRequest: |
| if not isinstance(obj, dict): |
| raise ProtocolError(f"{where}: request must be a JSON object") |
| clip_id = obj.get("clip_id") |
| audio_path = obj.get("audio_path") |
| params = obj.get("params", {}) |
| if not isinstance(clip_id, str) or not clip_id: |
| raise ProtocolError(f"{where}: clip_id must be a non-empty string") |
| if not isinstance(audio_path, str) or not audio_path: |
| raise ProtocolError(f"{where}: audio_path must be a non-empty string") |
| if not isinstance(params, dict) or not all( |
| isinstance(k, str) and isinstance(v, str) for k, v in params.items() |
| ): |
| raise ProtocolError(f"{where}: params must map strings to strings") |
| return WorkRequest(clip_id=clip_id, audio_path=audio_path, params=dict(params)) |
|
|
|
|
| @dataclass(frozen=True) |
| class WorkResult: |
| """Outcome for one request: exactly one of ``output`` or ``error`` is set. |
| |
| Attributes: |
| clip_id: Id of the request this answers. |
| output: Worker-specific JSON-compatible payload (e.g. ``{"text": ...}``). |
| error: ``"<ExceptionType>: <message>"`` when the item failed. |
| infer_seconds: Wall-clock inference time attributed to this item (batch time |
| divided by batch size), excluding model loading. |
| """ |
|
|
| clip_id: str |
| output: Mapping[str, object] | None |
| error: str | None |
| infer_seconds: float |
|
|
| def __post_init__(self) -> None: |
| if (self.output is None) == (self.error is None): |
| raise ProtocolError(f"{self.clip_id}: exactly one of output/error must be set") |
| if self.infer_seconds < 0: |
| raise ProtocolError(f"{self.clip_id}: infer_seconds must be >= 0") |
|
|
| @property |
| def ok(self) -> bool: |
| return self.error is None |
|
|
| def to_json(self) -> str: |
| return json.dumps( |
| { |
| "clip_id": self.clip_id, |
| "output": None if self.output is None else dict(self.output), |
| "error": self.error, |
| "infer_seconds": self.infer_seconds, |
| }, |
| ensure_ascii=False, |
| sort_keys=True, |
| ) |
|
|
| @staticmethod |
| def from_obj(obj: object, where: str) -> WorkResult: |
| if not isinstance(obj, dict): |
| raise ProtocolError(f"{where}: result must be a JSON object") |
| clip_id = obj.get("clip_id") |
| output = obj.get("output") |
| error = obj.get("error") |
| seconds = obj.get("infer_seconds") |
| if not isinstance(clip_id, str) or not clip_id: |
| raise ProtocolError(f"{where}: clip_id must be a non-empty string") |
| if output is not None and not isinstance(output, dict): |
| raise ProtocolError(f"{where}: output must be an object or null") |
| if error is not None and not isinstance(error, str): |
| raise ProtocolError(f"{where}: error must be a string or null") |
| if isinstance(seconds, bool) or not isinstance(seconds, int | float): |
| raise ProtocolError(f"{where}: infer_seconds must be a number") |
| return WorkResult(clip_id=clip_id, output=output, error=error, infer_seconds=float(seconds)) |
|
|
|
|
| @dataclass(frozen=True) |
| class WorkerSummary: |
| """Facts about one worker invocation, written once at the end.""" |
|
|
| worker: str |
| model_id: str |
| revision: str |
| load_seconds: float |
| peak_vram_gb: float | None |
| n_ok: int |
| n_error: int |
| started_at: str |
| ended_at: str |
| packages: Mapping[str, str] = field(default_factory=dict) |
| extra: Mapping[str, str] = field(default_factory=dict) |
|
|
|
|
| def _fsync_line(handle: IO[str], line: str) -> None: |
| handle.write(line + "\n") |
| handle.flush() |
| os.fsync(handle.fileno()) |
|
|
|
|
| def write_requests(path: Path, requests: Iterable[WorkRequest]) -> int: |
| """Write a requests file atomically and return the number of requests. |
| |
| Raises: |
| ProtocolError: on duplicate clip ids. |
| """ |
| from nida.utils.io import atomic_write_text |
|
|
| seen: set[str] = set() |
| lines: list[str] = [] |
| for request in requests: |
| if request.clip_id in seen: |
| raise ProtocolError(f"duplicate clip_id in requests: {request.clip_id}") |
| seen.add(request.clip_id) |
| lines.append(request.to_json()) |
| atomic_write_text(path, "".join(line + "\n" for line in lines)) |
| return len(lines) |
|
|
|
|
| def _iter_json_lines(path: Path) -> Iterator[tuple[int, str]]: |
| with path.open("r", encoding="utf-8") as handle: |
| yield from enumerate(handle, start=1) |
|
|
|
|
| def load_requests(path: Path) -> list[WorkRequest]: |
| """Read and validate a requests file (no partial lines tolerated).""" |
| requests: list[WorkRequest] = [] |
| seen: set[str] = set() |
| for number, line in _iter_json_lines(path): |
| where = f"{path}:{number}" |
| if not line.endswith("\n"): |
| raise ProtocolError(f"{where}: truncated request line") |
| try: |
| obj = json.loads(line) |
| except json.JSONDecodeError as err: |
| raise ProtocolError(f"{where}: invalid JSON ({err.msg})") from err |
| request = WorkRequest.from_obj(obj, where) |
| if request.clip_id in seen: |
| raise ProtocolError(f"{where}: duplicate clip_id {request.clip_id}") |
| seen.add(request.clip_id) |
| requests.append(request) |
| return requests |
|
|
|
|
| def load_results(path: Path, *, repair: bool = False) -> dict[str, WorkResult]: |
| """Read a results file into ``{clip_id: result}``. |
| |
| A missing file yields an empty mapping. Only the *last* line may be incomplete (an |
| interrupted append): with ``repair=True`` it is removed from the file, otherwise a |
| ``ProtocolError`` is raised. Any malformed line before the last one, and any |
| duplicate clip id, is always an error. |
| """ |
| if not path.exists(): |
| return {} |
| raw = path.read_bytes() |
| |
| |
| keep = raw.rfind(b"\n") + 1 |
| body, trailing = raw[:keep], raw[keep:] |
| try: |
| text = body.decode("utf-8") |
| except UnicodeDecodeError as err: |
| raise ProtocolError(f"{path}: complete lines are not valid UTF-8 ({err})") from err |
| lines = text.split("\n")[:-1] if text else [] |
| results: dict[str, WorkResult] = {} |
| for number, line in enumerate(lines, start=1): |
| where = f"{path}:{number}" |
| try: |
| obj = json.loads(line) |
| except json.JSONDecodeError as err: |
| raise ProtocolError(f"{where}: invalid JSON ({err.msg})") from err |
| result = WorkResult.from_obj(obj, where) |
| if result.clip_id in results: |
| raise ProtocolError(f"{where}: duplicate result for {result.clip_id}") |
| results[result.clip_id] = result |
| if trailing: |
| if not repair: |
| raise ProtocolError(f"{path}: incomplete trailing line ({len(trailing)} bytes)") |
| with path.open("r+b") as handle: |
| handle.truncate(keep) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| return results |
|
|
|
|
| def pending_requests( |
| requests: Iterable[WorkRequest], |
| results: Mapping[str, WorkResult], |
| *, |
| retry_failed: bool = False, |
| ) -> list[WorkRequest]: |
| """Requests that still need work, in input order. |
| |
| Raises: |
| ProtocolError: if ``results`` contains clip ids that are not in ``requests`` |
| (the results file belongs to a different request set). |
| """ |
| request_list = list(requests) |
| known = {r.clip_id for r in request_list} |
| stray = sorted(set(results) - known) |
| if stray: |
| raise ProtocolError(f"results contain unknown clip ids, e.g. {stray[:3]}") |
| pending: list[WorkRequest] = [] |
| for request in request_list: |
| done = results.get(request.clip_id) |
| if done is None or (retry_failed and not done.ok): |
| pending.append(request) |
| return pending |
|
|
|
|
| def append_result(handle: IO[str], result: WorkResult) -> None: |
| """Append one result line and fsync, so an interruption loses at most this line.""" |
| _fsync_line(handle, result.to_json()) |
|
|
|
|
| def write_summary(path: Path, summary: WorkerSummary) -> None: |
| """Write the summary JSON next to the results (stdlib-only atomic replace).""" |
| tmp = path.with_name(f".{path.name}.tmp") |
| with tmp.open("w", encoding="utf-8") as handle: |
| json.dump(asdict(summary), handle, ensure_ascii=False, indent=2, sort_keys=True) |
| handle.write("\n") |
| handle.flush() |
| os.fsync(handle.fileno()) |
| tmp.replace(path) |
|
|
|
|
| def read_summary(path: Path) -> WorkerSummary: |
| """Read and validate a summary file.""" |
| obj = json.loads(path.read_text(encoding="utf-8")) |
| if not isinstance(obj, dict): |
| raise ProtocolError(f"{path}: summary must be a JSON object") |
| try: |
| return WorkerSummary(**obj) |
| except TypeError as err: |
| raise ProtocolError(f"{path}: invalid summary fields ({err})") from err |
|
|
|
|
| def drop_results(path: Path, clip_ids: Iterable[str]) -> int: |
| """Remove the given clip ids from a results file (atomic rewrite); returns how many. |
| |
| Used when a result's side effect is gone (e.g. its output file was lost), so the |
| item becomes pending again instead of being trusted. |
| """ |
| doomed = set(clip_ids) |
| results = load_results(path, repair=True) |
| kept = [r.to_json() + "\n" for cid, r in results.items() if cid not in doomed] |
| removed = len(results) - len(kept) |
| if removed: |
| tmp = path.with_name(f".{path.name}.tmp") |
| with tmp.open("w", encoding="utf-8") as handle: |
| handle.writelines(kept) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| tmp.replace(path) |
| return removed |
|
|
