feat(dictate): use local socket-activated Parakeet
This commit is contained in:
@@ -0,0 +1,570 @@
|
||||
"""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 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
|
||||
|
||||
|
||||
class Recorder(Protocol):
|
||||
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
|
||||
|
||||
|
||||
@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)
|
||||
|
||||
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
|
||||
self._state(active, "listening")
|
||||
|
||||
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 v1 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()
|
||||
Reference in New Issue
Block a user