602 lines
23 KiB
Python
602 lines
23 KiB
Python
"""Dictate service state machine and Unix-domain socket loop."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from concurrent.futures import Executor, Future, ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
import math
|
|
import os
|
|
from queue import Empty, SimpleQueue
|
|
import selectors
|
|
import socket
|
|
import time
|
|
from typing import Callable, Protocol
|
|
|
|
from .audio import AudioError, pcm_s16le_to_floats
|
|
from .model import NemoParakeetModel, TranscriptionModel
|
|
from .protocol import (
|
|
MAX_FRAME_BYTES,
|
|
MAX_TRANSCRIPT_BYTES,
|
|
PROTOCOL_VERSION,
|
|
FrameDecoder,
|
|
ProtocolError,
|
|
Request,
|
|
encode_frame,
|
|
error_event,
|
|
parse_request,
|
|
)
|
|
from .recorder import Capture, PwRecorder, RecorderError
|
|
|
|
DEFAULT_IDLE_SECONDS = 15 * 60
|
|
MAX_CLIENT_CONNECTIONS = 16
|
|
MAX_CONNECTION_OUTPUT_BYTES = MAX_FRAME_BYTES
|
|
LEVEL_EVENT_INTERVAL_SECONDS = 0.1
|
|
|
|
|
|
class Recorder(Protocol):
|
|
@property
|
|
def audio_level(self) -> float: ...
|
|
|
|
def start(self) -> None: ...
|
|
def stop(self) -> Capture: ...
|
|
def cancel(self) -> None: ...
|
|
def poll_unexpected(self) -> RecorderError | None: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Outbound:
|
|
client_id: str
|
|
frame: dict[str, object]
|
|
|
|
|
|
@dataclass
|
|
class _Generation:
|
|
owner: str
|
|
generation_id: str
|
|
request_id: str
|
|
state: str
|
|
recorder: Recorder | None = None
|
|
tombstoned: bool = False
|
|
last_level_event_at: float | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Completion:
|
|
kind: str
|
|
generation_id: str
|
|
future: Future[object]
|
|
|
|
|
|
class DictateServiceCore:
|
|
"""Single-generation protocol state machine with injected host boundaries."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: TranscriptionModel,
|
|
recorder_factory: Callable[[], Recorder] = PwRecorder,
|
|
executor: Executor | None = None,
|
|
clock: Callable[[], float] = time.monotonic,
|
|
idle_seconds: float = DEFAULT_IDLE_SECONDS,
|
|
) -> None:
|
|
self._model = model
|
|
self._recorder_factory = recorder_factory
|
|
self._executor = executor or ThreadPoolExecutor(max_workers=1, thread_name_prefix="dictate")
|
|
self._owns_executor = executor is None
|
|
self._clock = clock
|
|
self._idle_seconds = idle_seconds
|
|
self._last_activity = clock()
|
|
self._clients: dict[str, bool] = {}
|
|
self._active: _Generation | None = None
|
|
self._model_loaded = bool(getattr(model, "loaded", False))
|
|
self._worker_kind: str | None = None
|
|
self._completions: SimpleQueue[_Completion] = SimpleQueue()
|
|
self._events: list[Outbound] = []
|
|
|
|
@property
|
|
def active_generation_id(self) -> str | None:
|
|
return self._active.generation_id if self._active else None
|
|
|
|
def accepted(self, client_id: str) -> None:
|
|
"""Track a socket without treating acceptance as protocol activity."""
|
|
self._clients.setdefault(client_id, False)
|
|
|
|
def handle_request(self, client_id: str, request: Request) -> None:
|
|
"""Handle one already parsed request and queue its observable events."""
|
|
self._last_activity = self._clock()
|
|
self._clients.setdefault(client_id, False)
|
|
|
|
if request.type == "hello":
|
|
self._clients[client_id] = True
|
|
self._emit(
|
|
client_id,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "pong",
|
|
"requestId": request.request_id,
|
|
"service": "pi-dictate",
|
|
"ready": self._model_loaded,
|
|
},
|
|
)
|
|
return
|
|
|
|
if not self._clients[client_id]:
|
|
self._request_error(client_id, request, "hello_required", "hello is required first")
|
|
return
|
|
|
|
if request.type == "ping":
|
|
self._emit(
|
|
client_id,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "pong",
|
|
"requestId": request.request_id,
|
|
"service": "pi-dictate",
|
|
"ready": self._model_loaded,
|
|
"busy": self._active is not None,
|
|
},
|
|
)
|
|
elif request.type == "start":
|
|
self._start(client_id, request)
|
|
elif request.type == "stop":
|
|
self._stop(client_id, request)
|
|
elif request.type == "cancel":
|
|
self._cancel(client_id, request)
|
|
|
|
def poll(self) -> None:
|
|
"""Apply completed workers and recorder exits without blocking the socket loop."""
|
|
while True:
|
|
try:
|
|
completion = self._completions.get_nowait()
|
|
except Empty:
|
|
break
|
|
self._apply_completion(completion)
|
|
|
|
active = self._active
|
|
if active and active.state == "listening" and active.recorder and self._worker_kind is None:
|
|
try:
|
|
failure = active.recorder.poll_unexpected()
|
|
except RecorderError as error:
|
|
failure = error
|
|
if failure is not None:
|
|
self._async_error(active, failure.code, failure.detail)
|
|
self._finish_generation(active, notify_idle=True)
|
|
return
|
|
self._emit_audio_level(active)
|
|
|
|
def drain_events(self) -> list[Outbound]:
|
|
events, self._events = self._events, []
|
|
return events
|
|
|
|
def disconnect(self, client_id: str) -> None:
|
|
self._clients.pop(client_id, None)
|
|
active = self._active
|
|
if active is None or active.owner != client_id:
|
|
return
|
|
active.tombstoned = True
|
|
if active.state == "listening" and active.recorder is not None:
|
|
active.state = "stopping"
|
|
self._submit("cancel", active, active.recorder.cancel)
|
|
|
|
def should_exit(self) -> bool:
|
|
if self._active is not None or self._worker_kind is not None:
|
|
return False
|
|
return self._clock() - self._last_activity >= self._idle_seconds
|
|
|
|
def close(self) -> None:
|
|
active = self._active
|
|
if active and active.recorder and active.state == "listening":
|
|
try:
|
|
active.recorder.cancel()
|
|
except RecorderError:
|
|
pass
|
|
self._active = None
|
|
if self._owns_executor:
|
|
self._executor.shutdown(wait=True, cancel_futures=False)
|
|
|
|
def _start(self, client_id: str, request: Request) -> None:
|
|
assert request.generation_id is not None
|
|
if self._active is not None or self._worker_kind is not None:
|
|
self._request_error(client_id, request, "busy", "another generation is active")
|
|
return
|
|
active = _Generation(client_id, request.generation_id, request.request_id, "loading")
|
|
self._active = active
|
|
if not self._model_loaded:
|
|
self._state(active, "loading")
|
|
self._submit("load", active, self._model.load)
|
|
return
|
|
self._begin_recording(active)
|
|
|
|
def _begin_recording(self, active: _Generation) -> None:
|
|
if active.tombstoned:
|
|
self._finish_generation(active, notify_idle=False)
|
|
return
|
|
try:
|
|
recorder = self._recorder_factory()
|
|
recorder.start()
|
|
except RecorderError as error:
|
|
self._async_error(active, error.code, error.detail)
|
|
self._finish_generation(active, notify_idle=True)
|
|
return
|
|
except OSError as error:
|
|
self._async_error(active, "recorder_spawn", str(error))
|
|
self._finish_generation(active, notify_idle=True)
|
|
return
|
|
active.recorder = recorder
|
|
active.last_level_event_at = self._clock()
|
|
self._state(active, "listening")
|
|
|
|
def _emit_audio_level(self, active: _Generation) -> None:
|
|
now = self._clock()
|
|
if (
|
|
active.last_level_event_at is not None
|
|
and now - active.last_level_event_at < LEVEL_EVENT_INTERVAL_SECONDS
|
|
):
|
|
return
|
|
level = active.recorder.audio_level if active.recorder is not None else 0.0
|
|
if not math.isfinite(level):
|
|
level = 0.0
|
|
active.last_level_event_at = now
|
|
self._emit(
|
|
active.owner,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "level",
|
|
"requestId": active.request_id,
|
|
"generationId": active.generation_id,
|
|
"level": round(max(0.0, min(1.0, level)), 3),
|
|
},
|
|
)
|
|
|
|
def _stop(self, client_id: str, request: Request) -> None:
|
|
active = self._owned_generation(client_id, request)
|
|
if active is None:
|
|
return
|
|
if active.state != "listening" or active.recorder is None:
|
|
self._request_error(client_id, request, "invalid_state", f"cannot stop while {active.state}")
|
|
return
|
|
active.request_id = request.request_id
|
|
try:
|
|
failure = active.recorder.poll_unexpected()
|
|
except RecorderError as error:
|
|
failure = error
|
|
if failure is not None:
|
|
self._async_error(active, failure.code, failure.detail)
|
|
self._finish_generation(active, notify_idle=True)
|
|
return
|
|
active.state = "stopping"
|
|
self._state(active, "stopping", assign=False)
|
|
self._submit("stop", active, self._stop_and_convert, active.recorder)
|
|
|
|
def _cancel(self, client_id: str, request: Request) -> None:
|
|
active = self._owned_generation(client_id, request)
|
|
if active is None:
|
|
return
|
|
if active.tombstoned:
|
|
self._request_error(client_id, request, "invalid_state", "generation is already cancelled")
|
|
return
|
|
active.tombstoned = True
|
|
active.request_id = request.request_id
|
|
self._emit(
|
|
client_id,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "cancelled",
|
|
"requestId": active.request_id,
|
|
"generationId": active.generation_id,
|
|
},
|
|
)
|
|
if active.state == "listening" and active.recorder is not None:
|
|
active.state = "stopping"
|
|
self._submit("cancel", active, active.recorder.cancel)
|
|
|
|
def _owned_generation(self, client_id: str, request: Request) -> _Generation | None:
|
|
assert request.generation_id is not None
|
|
active = self._active
|
|
if active is None:
|
|
self._request_error(client_id, request, "invalid_state", "no generation is active")
|
|
return None
|
|
if active.owner != client_id:
|
|
self._request_error(client_id, request, "not_owner", "generation belongs to another client")
|
|
return None
|
|
if active.generation_id != request.generation_id:
|
|
self._request_error(
|
|
client_id,
|
|
request,
|
|
"generation_mismatch",
|
|
"generation ID does not match the active generation",
|
|
)
|
|
return None
|
|
return active
|
|
|
|
@staticmethod
|
|
def _stop_and_convert(recorder: Recorder) -> Sequence[float]:
|
|
capture = recorder.stop()
|
|
return pcm_s16le_to_floats(capture.pcm)
|
|
|
|
def _submit(
|
|
self,
|
|
kind: str,
|
|
active: _Generation,
|
|
function: Callable[..., object],
|
|
*args: object,
|
|
) -> None:
|
|
self._worker_kind = kind
|
|
future = self._executor.submit(function, *args)
|
|
future.add_done_callback(
|
|
lambda completed, kind=kind, generation_id=active.generation_id: self._completions.put(
|
|
_Completion(kind, generation_id, completed)
|
|
)
|
|
)
|
|
|
|
def _apply_completion(self, completion: _Completion) -> None:
|
|
if self._worker_kind == completion.kind:
|
|
self._worker_kind = None
|
|
active = self._active
|
|
if active is None or active.generation_id != completion.generation_id:
|
|
return
|
|
try:
|
|
result = completion.future.result()
|
|
except RecorderError as error:
|
|
if not active.tombstoned:
|
|
self._async_error(active, error.code, error.detail)
|
|
self._finish_generation(active, notify_idle=not active.tombstoned)
|
|
return
|
|
except AudioError as error:
|
|
if not active.tombstoned:
|
|
self._async_error(active, "invalid_audio", str(error))
|
|
self._finish_generation(active, notify_idle=not active.tombstoned)
|
|
return
|
|
except BaseException as error:
|
|
if not active.tombstoned:
|
|
code = "model_failure" if completion.kind in ("load", "infer") else "recorder_failure"
|
|
self._async_error(active, code, str(error))
|
|
self._finish_generation(active, notify_idle=not active.tombstoned)
|
|
return
|
|
|
|
if completion.kind == "load":
|
|
self._model_loaded = True
|
|
if active.tombstoned:
|
|
self._finish_generation(active, notify_idle=False)
|
|
else:
|
|
self._begin_recording(active)
|
|
elif completion.kind == "stop":
|
|
if active.tombstoned:
|
|
self._finish_generation(active, notify_idle=False)
|
|
return
|
|
samples = result
|
|
if not isinstance(samples, Sequence) or isinstance(samples, (str, bytes, bytearray)):
|
|
self._async_error(active, "invalid_audio", "audio conversion returned invalid samples")
|
|
self._finish_generation(active, notify_idle=True)
|
|
return
|
|
self._state(active, "transcribing")
|
|
self._submit("infer", active, self._model.transcribe, samples)
|
|
elif completion.kind == "infer":
|
|
if active.tombstoned:
|
|
self._finish_generation(active, notify_idle=False)
|
|
return
|
|
if not isinstance(result, str):
|
|
self._async_error(active, "model_failure", "model returned a non-string transcript")
|
|
elif len(result.encode("utf-8")) > MAX_TRANSCRIPT_BYTES:
|
|
self._async_error(active, "transcript_too_large", "transcript exceeds output limit")
|
|
else:
|
|
self._emit(
|
|
active.owner,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "result",
|
|
"requestId": active.request_id,
|
|
"generationId": active.generation_id,
|
|
"transcript": result,
|
|
},
|
|
)
|
|
self._finish_generation(active, notify_idle=True)
|
|
elif completion.kind == "cancel":
|
|
self._finish_generation(active, notify_idle=False)
|
|
|
|
def _finish_generation(self, active: _Generation, *, notify_idle: bool) -> None:
|
|
if self._active is not active:
|
|
return
|
|
self._active = None
|
|
self._last_activity = self._clock()
|
|
if notify_idle and active.owner in self._clients:
|
|
self._state(active, "idle", assign=False)
|
|
|
|
def _state(self, active: _Generation, state: str, *, assign: bool = True) -> None:
|
|
if assign:
|
|
active.state = state
|
|
self._emit(
|
|
active.owner,
|
|
{
|
|
"version": PROTOCOL_VERSION,
|
|
"type": "state",
|
|
"requestId": active.request_id,
|
|
"generationId": active.generation_id,
|
|
"state": state,
|
|
},
|
|
)
|
|
|
|
def _request_error(self, client_id: str, request: Request, code: str, detail: str) -> None:
|
|
error = ProtocolError(
|
|
code,
|
|
detail,
|
|
request_id=request.request_id,
|
|
generation_id=request.generation_id,
|
|
)
|
|
self._emit(client_id, error_event(error))
|
|
|
|
def _async_error(self, active: _Generation, code: str, detail: str) -> None:
|
|
error = ProtocolError(
|
|
code,
|
|
detail,
|
|
request_id=active.request_id,
|
|
generation_id=active.generation_id,
|
|
)
|
|
self._emit(active.owner, error_event(error))
|
|
|
|
def _emit(self, client_id: str, frame: dict[str, object]) -> None:
|
|
if client_id in self._clients:
|
|
self._events.append(Outbound(client_id, frame))
|
|
|
|
|
|
@dataclass
|
|
class _Connection:
|
|
socket: socket.socket
|
|
decoder: FrameDecoder
|
|
output: bytearray
|
|
close_after_write: bool = False
|
|
|
|
|
|
def inherited_listener(environ: dict[str, str] | None = None) -> socket.socket:
|
|
"""Return systemd's single inherited listener at fd 3."""
|
|
environment = os.environ if environ is None else environ
|
|
try:
|
|
listen_pid = int(environment.get("LISTEN_PID", "0"))
|
|
listen_fds = int(environment.get("LISTEN_FDS", "0"))
|
|
except ValueError as error:
|
|
raise RuntimeError("invalid systemd socket activation environment") from error
|
|
if listen_pid != os.getpid() or listen_fds != 1:
|
|
raise RuntimeError("expected exactly one systemd socket-activated listener")
|
|
listener = socket.socket(fileno=3)
|
|
listener.setblocking(False)
|
|
return listener
|
|
|
|
|
|
def serve(listener: socket.socket, core: DictateServiceCore) -> None:
|
|
"""Serve protocol v2 on an explicit listener until the core becomes idle."""
|
|
selector = selectors.DefaultSelector()
|
|
connections: dict[str, _Connection] = {}
|
|
sequence = 0
|
|
listener.setblocking(False)
|
|
selector.register(listener, selectors.EVENT_READ, None)
|
|
try:
|
|
while not core.should_exit():
|
|
core.poll()
|
|
for outbound in core.drain_events():
|
|
connection = connections.get(outbound.client_id)
|
|
if connection is None:
|
|
continue
|
|
try:
|
|
encoded = encode_frame(outbound.frame)
|
|
except ProtocolError:
|
|
_close_connection(selector, connections, core, outbound.client_id)
|
|
continue
|
|
if len(connection.output) + len(encoded) > MAX_CONNECTION_OUTPUT_BYTES:
|
|
_close_connection(selector, connections, core, outbound.client_id)
|
|
continue
|
|
connection.output.extend(encoded)
|
|
selector.modify(
|
|
connection.socket,
|
|
selectors.EVENT_READ | selectors.EVENT_WRITE,
|
|
outbound.client_id,
|
|
)
|
|
for key, mask in selector.select(timeout=0.1):
|
|
if key.data is None:
|
|
try:
|
|
client, _address = listener.accept()
|
|
except BlockingIOError:
|
|
continue
|
|
if len(connections) >= MAX_CLIENT_CONNECTIONS:
|
|
client.close()
|
|
continue
|
|
client.setblocking(False)
|
|
sequence += 1
|
|
client_id = f"client-{sequence}"
|
|
connections[client_id] = _Connection(client, FrameDecoder(), bytearray())
|
|
core.accepted(client_id)
|
|
selector.register(client, selectors.EVENT_READ, client_id)
|
|
continue
|
|
client_id = key.data
|
|
connection = connections.get(client_id)
|
|
if connection is None:
|
|
continue
|
|
if mask & selectors.EVENT_READ:
|
|
try:
|
|
data = connection.socket.recv(64 * 1024)
|
|
except BlockingIOError:
|
|
data = None
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
_close_connection(selector, connections, core, client_id)
|
|
continue
|
|
if data == b"":
|
|
_close_connection(selector, connections, core, client_id)
|
|
continue
|
|
if data:
|
|
try:
|
|
for frame in connection.decoder.feed(data):
|
|
core.handle_request(client_id, parse_request(frame))
|
|
except ProtocolError as error:
|
|
connection.output.extend(encode_frame(error_event(error)))
|
|
if error.close_connection:
|
|
connection.close_after_write = True
|
|
if mask & selectors.EVENT_WRITE and connection.output:
|
|
try:
|
|
sent = connection.socket.send(connection.output)
|
|
except BlockingIOError:
|
|
sent = 0
|
|
except (BrokenPipeError, ConnectionResetError):
|
|
_close_connection(selector, connections, core, client_id)
|
|
continue
|
|
del connection.output[:sent]
|
|
if connection.close_after_write and not connection.output:
|
|
_close_connection(selector, connections, core, client_id)
|
|
elif client_id in connections:
|
|
events = selectors.EVENT_READ
|
|
if connection.output:
|
|
events |= selectors.EVENT_WRITE
|
|
selector.modify(connection.socket, events, client_id)
|
|
finally:
|
|
for client_id in list(connections):
|
|
_close_connection(selector, connections, core, client_id)
|
|
selector.unregister(listener)
|
|
listener.close()
|
|
selector.close()
|
|
core.close()
|
|
|
|
|
|
def _close_connection(
|
|
selector: selectors.BaseSelector,
|
|
connections: dict[str, _Connection],
|
|
core: DictateServiceCore,
|
|
client_id: str,
|
|
) -> None:
|
|
connection = connections.pop(client_id, None)
|
|
if connection is None:
|
|
return
|
|
core.disconnect(client_id)
|
|
try:
|
|
selector.unregister(connection.socket)
|
|
except KeyError:
|
|
pass
|
|
connection.socket.close()
|
|
|
|
|
|
def main() -> None:
|
|
model_path = os.environ.get("PI_DICTATE_MODEL_PATH")
|
|
if not model_path:
|
|
raise SystemExit("PI_DICTATE_MODEL_PATH must name an explicit local .nemo file")
|
|
terms_path = os.environ.get("PI_DICTATE_TERMS_PATH")
|
|
if not terms_path:
|
|
raise SystemExit("PI_DICTATE_TERMS_PATH must name the host coding dictionary")
|
|
idle_value = os.environ.get("PI_DICTATE_IDLE_SECONDS", str(DEFAULT_IDLE_SECONDS))
|
|
boost_value = os.environ.get("PI_DICTATE_BOOST_ALPHA", "1.0")
|
|
try:
|
|
idle_seconds = float(idle_value)
|
|
boost_alpha = float(boost_value)
|
|
except ValueError as error:
|
|
raise SystemExit("PI_DICTATE_IDLE_SECONDS and PI_DICTATE_BOOST_ALPHA must be numeric") from error
|
|
if idle_seconds <= 0:
|
|
raise SystemExit("PI_DICTATE_IDLE_SECONDS must be positive")
|
|
if boost_alpha <= 0:
|
|
raise SystemExit("PI_DICTATE_BOOST_ALPHA must be positive")
|
|
|
|
model = NemoParakeetModel(model_path, terms_path=terms_path, boost_alpha=boost_alpha)
|
|
core = DictateServiceCore(model=model, idle_seconds=idle_seconds)
|
|
listener = inherited_listener()
|
|
serve(listener, core)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|