592 lines
24 KiB
Python
592 lines
24 KiB
Python
"""Concurrent, crash-recoverable recording and transcription sessions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import queue
|
|
import shutil
|
|
import threading
|
|
import time
|
|
import uuid
|
|
import wave
|
|
from collections import deque
|
|
from dataclasses import asdict, dataclass, is_dataclass, replace
|
|
from pathlib import Path
|
|
from types import MappingProxyType
|
|
from typing import Any, Callable, Mapping
|
|
|
|
import numpy as np
|
|
|
|
from blurt_audio import StreamingAudio
|
|
|
|
|
|
_ACTIVE_PATHS: set[Path] = set()
|
|
_ACTIVE_LOCK = threading.Lock()
|
|
|
|
|
|
def _atomic_json(path: Path, value: dict[str, Any]) -> None:
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
temporary.write_text(json.dumps(value, indent=2), encoding="utf-8")
|
|
temporary.replace(path)
|
|
|
|
|
|
def _write_wav(path: Path, samples: np.ndarray) -> None:
|
|
temporary = path.with_suffix(path.suffix + ".tmp")
|
|
with wave.open(str(temporary), "wb") as output:
|
|
output.setnchannels(1)
|
|
output.setsampwidth(2)
|
|
output.setframerate(16000)
|
|
output.writeframes(np.asarray(samples, dtype="<i2").tobytes())
|
|
temporary.replace(path)
|
|
|
|
|
|
def _read_wav(path: Path) -> np.ndarray:
|
|
with wave.open(str(path), "rb") as source:
|
|
if source.getnchannels() != 1 or source.getsampwidth() != 2:
|
|
raise ValueError(f"unsupported recovery WAV format: {path}")
|
|
data = source.readframes(source.getnframes())
|
|
rate = source.getframerate()
|
|
samples = np.frombuffer(data, dtype="<i2").copy()
|
|
if rate != 16000:
|
|
raise ValueError(f"recovery WAV must be 16 kHz: {path}")
|
|
return samples
|
|
|
|
|
|
def _config_dict(config: Any) -> dict[str, Any]:
|
|
if is_dataclass(config):
|
|
return asdict(config)
|
|
if isinstance(config, Mapping):
|
|
return dict(config)
|
|
return dict(vars(config))
|
|
|
|
|
|
def _public_config(config: Any) -> dict[str, Any]:
|
|
blocked = ("key", "token", "secret", "password", "auth", "credential")
|
|
return {
|
|
str(k): v
|
|
for k, v in _config_dict(config).items()
|
|
if not any(word in str(k).lower() for word in blocked)
|
|
and isinstance(v, (str, int, float, bool, type(None)))
|
|
}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SessionResult:
|
|
session_id: str
|
|
text: str
|
|
complete: bool
|
|
errors: tuple[str, ...] = ()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RecoveryRecord:
|
|
session_id: str
|
|
path: Path
|
|
state: str
|
|
config: Mapping[str, Any]
|
|
done: int
|
|
total: int
|
|
created_at: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class _Job:
|
|
session: "RecordingSession"
|
|
index: int
|
|
path: Path
|
|
config: Any
|
|
|
|
|
|
class TranscriptionScheduler:
|
|
"""A shared worker pool whose pending payloads live as WAV files on disk."""
|
|
|
|
def __init__(self, request: Callable[[np.ndarray, Any], str], workers: int = 3,
|
|
queue_size: int = 96):
|
|
if workers < 1:
|
|
raise ValueError("workers must be positive")
|
|
self.request = request
|
|
self._ready: queue.Queue[_Job | None] = queue.Queue(maxsize=max(workers, queue_size))
|
|
self._overflow: deque[_Job] = deque()
|
|
self._lock = threading.Lock()
|
|
self._closed = False
|
|
self._threads = [threading.Thread(target=self._worker, name=f"blurt-upload-{i}", daemon=True)
|
|
for i in range(workers)]
|
|
for thread in self._threads:
|
|
thread.start()
|
|
|
|
def submit(self, session: "RecordingSession", index: int, path: Path,
|
|
config: Any) -> None:
|
|
job = _Job(session, index, path, config)
|
|
with self._lock:
|
|
if self._closed:
|
|
raise RuntimeError("transcription scheduler is shut down")
|
|
try:
|
|
self._ready.put_nowait(job)
|
|
except queue.Full:
|
|
# Only tiny disk references are retained here; audio remains on disk.
|
|
self._overflow.append(job)
|
|
|
|
def _refill(self) -> None:
|
|
with self._lock:
|
|
while self._overflow:
|
|
try:
|
|
self._ready.put_nowait(self._overflow[0])
|
|
except queue.Full:
|
|
break
|
|
self._overflow.popleft()
|
|
|
|
def _worker(self) -> None:
|
|
while True:
|
|
try:
|
|
job = self._ready.get(timeout=0.1)
|
|
except queue.Empty:
|
|
with self._lock:
|
|
if self._closed and not self._overflow:
|
|
return
|
|
continue
|
|
text = None
|
|
error = None
|
|
try:
|
|
samples = _read_wav(job.path)
|
|
text = self.request(samples, job.config)
|
|
if not isinstance(text, str) or not text.strip():
|
|
raise ValueError("transcription returned empty text")
|
|
text = text.strip()
|
|
except Exception as exception:
|
|
error = f"{type(exception).__name__}: {exception}"
|
|
try:
|
|
job.session._job_finished(job.index, text, error)
|
|
except Exception:
|
|
# Session persistence/callback failures must not terminate the pool.
|
|
logging.exception("Could not record transcription result for %s", job.path)
|
|
finally:
|
|
self._ready.task_done()
|
|
self._refill()
|
|
|
|
def shutdown(self, wait: bool = True) -> None:
|
|
with self._lock:
|
|
self._closed = True
|
|
if wait:
|
|
for thread in self._threads:
|
|
thread.join()
|
|
|
|
|
|
class RecordingSession:
|
|
"""Owns capture, chunk persistence, ordered results, and one recording lifecycle."""
|
|
|
|
def __init__(self, rate: int, channels: int, config: Any, spool_dir: str | Path,
|
|
scheduler: TranscriptionScheduler, on_update: Callable[["RecordingSession"], None] | None = None,
|
|
on_complete: Callable[["RecordingSession"], None] | None = None,
|
|
input_queue_size: int = 64):
|
|
self.session_id = uuid.uuid4().hex
|
|
self.path = Path(spool_dir) / self.session_id
|
|
self.path.mkdir(parents=True, exist_ok=False)
|
|
self.config = config
|
|
self.scheduler = scheduler
|
|
self.on_update = on_update
|
|
self.on_complete = on_complete
|
|
self.state = "recording"
|
|
self.done = 0
|
|
self.total = 0
|
|
self.result: SessionResult | None = None
|
|
self.created_at = time.time()
|
|
self._entries: list[dict[str, Any]] = []
|
|
self._submitted: set[int] = set()
|
|
self._capture_persisted_ok = False
|
|
self._lock = threading.RLock()
|
|
self._condition = threading.Condition(self._lock)
|
|
self._input: queue.Queue[np.ndarray] = queue.Queue(maxsize=input_queue_size)
|
|
self._stop = threading.Event()
|
|
self._producer_done = threading.Event()
|
|
self._overflow_error = threading.Event()
|
|
self._input_statuses: queue.SimpleQueue[str] = queue.SimpleQueue()
|
|
self._audio = StreamingAudio(rate, channels)
|
|
self._last_checkpoint = time.monotonic()
|
|
self._tail_start = 0
|
|
self._manifest_path = self.path / "session.json"
|
|
with _ACTIVE_LOCK:
|
|
_ACTIVE_PATHS.add(self.path.resolve())
|
|
try:
|
|
self._save()
|
|
except Exception:
|
|
self.release_for_retry()
|
|
raise
|
|
self._producer = threading.Thread(target=self._produce, name=f"blurt-capture-{self.session_id[:6]}", daemon=True)
|
|
self._producer.start()
|
|
|
|
def feed(self, indata: np.ndarray) -> bool:
|
|
if self.state != "recording":
|
|
return False
|
|
try:
|
|
self._input.put_nowait(np.array(indata, copy=True))
|
|
return True
|
|
except queue.Full:
|
|
# The sound callback must never wait for a lock or filesystem.
|
|
self.state = "capture_failed"
|
|
self._overflow_error.set()
|
|
self._stop.set()
|
|
return False
|
|
|
|
def stop(self) -> None:
|
|
with self._lock:
|
|
if self.state == "recording":
|
|
self.state = "finishing"
|
|
self._stop.set()
|
|
self._save()
|
|
|
|
def report_input_status(self, status: object) -> None:
|
|
"""Record a PortAudio status without blocking its callback thread."""
|
|
if status:
|
|
self._input_statuses.put(str(status))
|
|
|
|
def producer_join(self, timeout: float | None = None) -> bool:
|
|
self._producer.join(timeout)
|
|
return not self._producer.is_alive()
|
|
|
|
def join(self, timeout: float | None = None) -> SessionResult | None:
|
|
deadline = None if timeout is None else time.monotonic() + timeout
|
|
with self._condition:
|
|
while self.result is None:
|
|
remaining = None if deadline is None else deadline - time.monotonic()
|
|
if remaining is not None and remaining <= 0:
|
|
return None
|
|
self._condition.wait(remaining)
|
|
return self.result
|
|
|
|
def _produce(self) -> None:
|
|
try:
|
|
while True:
|
|
try:
|
|
block = self._input.get(timeout=0.05)
|
|
except queue.Empty:
|
|
if self._stop.is_set():
|
|
break
|
|
continue
|
|
try:
|
|
for chunk in self._audio.feed(block):
|
|
self._seal(chunk)
|
|
if time.monotonic() - self._last_checkpoint >= 1.0:
|
|
self._checkpoint()
|
|
finally:
|
|
self._input.task_done()
|
|
for chunk in self._audio.finish():
|
|
self._seal(chunk)
|
|
self._remove_tail()
|
|
if self._overflow_error.is_set():
|
|
with self._lock:
|
|
self._entries.append({"index": len(self._entries), "speech": False,
|
|
"status": "failed", "error": "capture input queue overflowed; audio was lost"})
|
|
self._save()
|
|
self._persist_input_statuses()
|
|
self._capture_persisted_ok = True
|
|
except Exception as error:
|
|
with self._lock:
|
|
self.state = "capture_failed"
|
|
for entry in self._entries:
|
|
if (entry.get("speech") and entry.get("status") == "pending"
|
|
and int(entry["index"]) not in self._submitted):
|
|
entry["status"] = "failed"
|
|
entry["error"] = f"capture: {type(error).__name__}: {error}"
|
|
self.done += 1
|
|
self._entries.append({"index": len(self._entries), "speech": False,
|
|
"status": "failed",
|
|
"error": f"capture: {type(error).__name__}: {error}"})
|
|
try:
|
|
self._save()
|
|
except OSError:
|
|
pass
|
|
finally:
|
|
self._producer_done.set()
|
|
self._maybe_complete()
|
|
|
|
def _persist_input_statuses(self) -> None:
|
|
statuses = []
|
|
while not self._input_statuses.empty():
|
|
statuses.append(self._input_statuses.get())
|
|
if statuses:
|
|
with self._lock:
|
|
self._entries.append({"index": len(self._entries), "speech": False,
|
|
"status": "failed", "error": "input stream status: " + ", ".join(statuses)})
|
|
self._save()
|
|
|
|
@staticmethod
|
|
def _chunk_value(chunk: Any, name: str, default: Any = None) -> Any:
|
|
return getattr(chunk, name, default)
|
|
|
|
def _seal(self, chunk: Any) -> None:
|
|
samples = np.asarray(self._chunk_value(chunk, "samples"), dtype=np.int16)
|
|
speech = bool(self._chunk_value(chunk, "speech", True))
|
|
with self._lock:
|
|
index = len(self._entries)
|
|
filename = f"chunk-{index:06d}.wav"
|
|
entry = {"index": index, "file": filename, "speech": speech,
|
|
"start": self._chunk_value(chunk, "start", None),
|
|
"end": self._chunk_value(chunk, "end", None), "status": "pending"}
|
|
if not speech:
|
|
entry["status"] = "skipped"
|
|
entry["text"] = ""
|
|
self._entries.append(entry)
|
|
self.total += int(speech)
|
|
_write_wav(self.path / filename, samples)
|
|
# Commit the sealed file before removing the previous recovery tail.
|
|
self._save()
|
|
self._remove_tail_locked()
|
|
self._tail_start = int(entry.get("end") or self._tail_start)
|
|
self._save()
|
|
if speech:
|
|
self.scheduler.submit(self, index, self.path / filename, self.config)
|
|
with self._lock:
|
|
self._submitted.add(index)
|
|
self._notify()
|
|
|
|
def _checkpoint(self) -> None:
|
|
tail = np.asarray(self._audio.checkpoint(), dtype=np.int16)
|
|
with self._lock:
|
|
if tail.size:
|
|
_write_wav(self.path / "tail.wav", tail)
|
|
ends = [int(e["end"]) for e in self._entries if e.get("end") is not None]
|
|
self._tail_start = max(ends, default=0)
|
|
else:
|
|
self._remove_tail_locked()
|
|
self._last_checkpoint = time.monotonic()
|
|
self._save()
|
|
|
|
def _remove_tail_locked(self) -> None:
|
|
(self.path / "tail.wav").unlink(missing_ok=True)
|
|
|
|
def _remove_tail(self) -> None:
|
|
with self._lock:
|
|
self._remove_tail_locked()
|
|
self._save()
|
|
|
|
def _job_finished(self, index: int, text: str | None, error: str | None) -> None:
|
|
with self._lock:
|
|
entry = self._entries[index]
|
|
entry["status"] = "done" if error is None else "failed"
|
|
if text is not None:
|
|
entry["text"] = text
|
|
if error is not None:
|
|
entry["error"] = error
|
|
self.done += 1
|
|
try:
|
|
self._save()
|
|
except OSError as persistence_error:
|
|
entry.pop("text", None)
|
|
entry["status"] = "failed"
|
|
entry["error"] = (
|
|
f"could not persist transcription: {type(persistence_error).__name__}: "
|
|
f"{persistence_error}"
|
|
)
|
|
self._notify()
|
|
self._maybe_complete()
|
|
|
|
def _maybe_complete(self) -> None:
|
|
callback = None
|
|
with self._condition:
|
|
if not self._producer_done.is_set() or self.done < self.total or self.result is not None:
|
|
return
|
|
errors = tuple(str(e["error"]) for e in self._entries if e.get("error"))
|
|
has_text = any(entry.get("text") for entry in self._entries)
|
|
parts = []
|
|
gap = False
|
|
for entry in self._entries:
|
|
if entry.get("text"):
|
|
if gap and has_text:
|
|
parts.append("[…]")
|
|
gap = False
|
|
parts.append(str(entry["text"]))
|
|
elif entry.get("error"):
|
|
gap = True
|
|
if gap and has_text:
|
|
parts.append("[…]")
|
|
text = " ".join(parts)
|
|
complete = not errors
|
|
self.state = "complete" if complete else "needs_retry"
|
|
self.result = SessionResult(self.session_id, text, complete, errors)
|
|
try:
|
|
self._save()
|
|
except OSError:
|
|
complete = False
|
|
self.state = "needs_retry"
|
|
errors = errors + ("could not persist final session state",)
|
|
self.result = SessionResult(self.session_id, text, False, errors)
|
|
self._condition.notify_all()
|
|
callback = self.on_complete
|
|
if callback:
|
|
try:
|
|
callback(self)
|
|
except Exception:
|
|
logging.exception("Session completion callback failed: %s", self.session_id)
|
|
self.release_for_retry()
|
|
|
|
def _notify(self) -> None:
|
|
if self.on_update:
|
|
try:
|
|
self.on_update(self)
|
|
except Exception:
|
|
logging.exception("Session progress callback failed: %s", self.session_id)
|
|
|
|
def _save(self) -> None:
|
|
_atomic_json(self._manifest_path, {"version": 1, "session_id": self.session_id,
|
|
"state": self.state, "config": _public_config(self.config), "done": self.done,
|
|
"total": self.total, "created_at": self.created_at, "entries": self._entries,
|
|
"has_tail": (self.path / "tail.wav").exists(), "tail_start": self._tail_start})
|
|
|
|
def acknowledge_delivery(self) -> None:
|
|
with self._lock:
|
|
if self.result is None:
|
|
raise RuntimeError("session is not finished")
|
|
if self.result.complete:
|
|
shutil.rmtree(self.path)
|
|
else:
|
|
self.state = "needs_retry"
|
|
self._save()
|
|
with _ACTIVE_LOCK:
|
|
_ACTIVE_PATHS.discard(self.path.resolve())
|
|
|
|
def release_for_retry(self) -> None:
|
|
"""Make retained data discoverable after delivery is deferred or fails."""
|
|
with _ACTIVE_LOCK:
|
|
_ACTIVE_PATHS.discard(self.path.resolve())
|
|
|
|
@classmethod
|
|
def recover(cls, path: str | Path, config: Any, scheduler: TranscriptionScheduler,
|
|
on_update=None, on_complete=None) -> "RecordingSession":
|
|
return _recover_session(Path(path), config, scheduler, on_update, on_complete)
|
|
|
|
@classmethod
|
|
def from_wav(cls, wav_path: str | Path, config: Any, spool_dir: str | Path,
|
|
scheduler: TranscriptionScheduler, on_update=None, on_complete=None) -> "RecordingSession":
|
|
samples = _read_wav(Path(wav_path))
|
|
session = cls(16000, 1, config, spool_dir, scheduler, on_update, on_complete)
|
|
session.feed(samples.reshape(-1, 1).astype(np.float32) / 32768.0)
|
|
session.stop()
|
|
session.producer_join()
|
|
if session._capture_persisted_ok:
|
|
Path(wav_path).unlink(missing_ok=True)
|
|
return session
|
|
|
|
|
|
def _recover_session(path: Path, config: Any, scheduler: TranscriptionScheduler,
|
|
on_update=None, on_complete=None) -> RecordingSession:
|
|
data = json.loads((path / "session.json").read_text(encoding="utf-8"))
|
|
session = object.__new__(RecordingSession)
|
|
session.session_id, session.path = data["session_id"], path
|
|
saved = data.get("config", {})
|
|
if is_dataclass(config):
|
|
fields = getattr(config, "__dataclass_fields__", {})
|
|
session.config = replace(config, **{k: v for k, v in saved.items() if k in fields})
|
|
else:
|
|
merged = dict(saved); merged.update(_config_dict(config))
|
|
session.config = MappingProxyType(merged)
|
|
session.scheduler, session.on_update, session.on_complete = scheduler, on_update, on_complete
|
|
session.state, session.done, session.total = "finishing", 0, 0
|
|
session.result = None
|
|
session.created_at = float(data.get("created_at", path.stat().st_mtime))
|
|
session._entries = data.get("entries", [])
|
|
session._submitted = set()
|
|
session._capture_persisted_ok = True
|
|
session._lock = threading.RLock(); session._condition = threading.Condition(session._lock)
|
|
session._input = queue.Queue(); session._stop = threading.Event(); session._stop.set()
|
|
session._producer_done = threading.Event(); session._producer_done.set()
|
|
session._overflow_error = threading.Event()
|
|
session._input_statuses = queue.SimpleQueue()
|
|
session._audio = None; session._last_checkpoint = time.monotonic()
|
|
session._tail_start = int(data.get("tail_start", 0))
|
|
session._manifest_path = path / "session.json"
|
|
referenced = {entry.get("file") for entry in session._entries}
|
|
orphans = [candidate for candidate in sorted(path.glob("chunk-*.wav"))
|
|
if candidate.name not in referenced]
|
|
sealed_end = max((int(e.get("end") or 0) for e in session._entries), default=0)
|
|
for orphan in orphans:
|
|
samples = _read_wav(orphan)
|
|
index = len(session._entries)
|
|
speech = bool(samples.size and np.max(np.abs(samples.astype(np.int32))) > 1)
|
|
session._entries.append({"index": index, "file": orphan.name, "speech": speech,
|
|
"status": "pending" if speech else "skipped",
|
|
"start": sealed_end, "end": sealed_end + len(samples),
|
|
"recovered_orphan": True})
|
|
sealed_end += len(samples)
|
|
tail = path / "tail.wav"
|
|
if tail.exists():
|
|
tail_samples = _read_wav(tail)
|
|
overlap = max(0, sealed_end - session._tail_start)
|
|
if overlap:
|
|
tail_samples = tail_samples[min(overlap, len(tail_samples)):]
|
|
index = len(session._entries)
|
|
target = path / f"chunk-{index:06d}.wav"
|
|
if len(tail_samples):
|
|
_write_wav(target, tail_samples)
|
|
speech = bool(np.max(np.abs(tail_samples.astype(np.int32))) > 1)
|
|
session._entries.append({"index": index, "file": target.name, "speech": speech,
|
|
"status": "pending" if speech else "skipped",
|
|
"start": sealed_end, "end": sealed_end + len(tail_samples),
|
|
"recovered_tail": True})
|
|
pending = []
|
|
for entry in session._entries:
|
|
if entry.get("speech") and entry.get("status") != "done" and entry.get("file") and (path / entry["file"]).exists():
|
|
entry["status"] = "pending"; entry.pop("error", None)
|
|
pending.append(entry)
|
|
elif entry.get("speech") and entry.get("status") != "done":
|
|
entry["status"] = "failed"
|
|
entry["error"] = "recovery audio file is missing"
|
|
session.done += 1
|
|
elif entry.get("speech"):
|
|
session.done += 1
|
|
session.total = sum(bool(e.get("speech")) for e in session._entries)
|
|
with _ACTIVE_LOCK:
|
|
_ACTIVE_PATHS.add(path.resolve())
|
|
# Commit adopted orphan/suffix entries before deleting the original tail.
|
|
try:
|
|
session._save()
|
|
tail.unlink(missing_ok=True)
|
|
session._save()
|
|
except Exception:
|
|
session.release_for_retry()
|
|
raise
|
|
|
|
def start_recovery() -> None:
|
|
for entry in pending:
|
|
try:
|
|
scheduler.submit(session, entry["index"], path / entry["file"], session.config)
|
|
session._submitted.add(int(entry["index"]))
|
|
except Exception as error:
|
|
session._job_finished(entry["index"], None, f"Could not schedule retry: {error}")
|
|
session._maybe_complete()
|
|
|
|
session._producer = threading.Thread(target=start_recovery, name="BlurtRecovery", daemon=True)
|
|
session._producer.start()
|
|
return session
|
|
|
|
|
|
class RecoveryStore:
|
|
def __init__(self, spool_dir: str | Path):
|
|
self.spool_dir = Path(spool_dir)
|
|
|
|
def discover(self) -> list[RecoveryRecord]:
|
|
records = []
|
|
if not self.spool_dir.exists():
|
|
return records
|
|
for manifest in sorted(self.spool_dir.glob("*/session.json")):
|
|
try:
|
|
with _ACTIVE_LOCK:
|
|
if manifest.parent.resolve() in _ACTIVE_PATHS:
|
|
continue
|
|
data = json.loads(manifest.read_text(encoding="utf-8"))
|
|
records.append(RecoveryRecord(data["session_id"], manifest.parent,
|
|
data.get("state", "interrupted"), MappingProxyType(data.get("config", {})),
|
|
int(data.get("done", 0)), int(data.get("total", 0)),
|
|
float(data.get("created_at", manifest.stat().st_mtime))))
|
|
except (OSError, ValueError, KeyError):
|
|
continue
|
|
return sorted(records, key=lambda record: record.created_at)
|
|
|
|
def resume(self, record: RecoveryRecord | str, config: Any,
|
|
scheduler: TranscriptionScheduler, on_update=None, on_complete=None) -> RecordingSession:
|
|
if isinstance(record, str):
|
|
path = self.spool_dir / record
|
|
else:
|
|
path = record.path
|
|
return RecordingSession.recover(path, config, scheduler, on_update, on_complete)
|