feat(dictate): support configurable ASR backends
This commit is contained in:
@@ -1,11 +1,22 @@
|
||||
from concurrent.futures import Future
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
import errno
|
||||
import selectors
|
||||
import subprocess
|
||||
import sys
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from pi_dictate_service.protocol import MAX_TRANSCRIPT_BYTES, PROTOCOL_VERSION, parse_request
|
||||
from pi_dictate_service.protocol import (
|
||||
MAX_FRAME_BYTES,
|
||||
MAX_TRANSCRIPT_BYTES,
|
||||
PROTOCOL_VERSION,
|
||||
encode_frame,
|
||||
parse_request,
|
||||
)
|
||||
from pi_dictate_service.recorder import Capture, RecorderError
|
||||
from pi_dictate_service.server import DictateServiceCore, Outbound
|
||||
from pi_dictate_service.server import DictateServiceCore, Outbound, serve
|
||||
|
||||
|
||||
class Clock:
|
||||
@@ -93,6 +104,154 @@ class FakeRecorder:
|
||||
return failure
|
||||
|
||||
|
||||
class SpawnFailRecorder(FakeRecorder):
|
||||
cleanup_complete = True
|
||||
|
||||
def start(self) -> None:
|
||||
raise RecorderError("recorder_spawn", "could not start pw-record")
|
||||
|
||||
|
||||
class SpawnFailRecorderFactory:
|
||||
def __call__(self) -> SpawnFailRecorder:
|
||||
return SpawnFailRecorder()
|
||||
|
||||
|
||||
class DelayedCleanupRecorder(FakeRecorder):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.cleanup_complete = False
|
||||
self.cleanup_attempts = 0
|
||||
self.cleanup_started = False
|
||||
|
||||
def stop(self) -> Capture:
|
||||
self.cleanup_started = True
|
||||
raise RecorderError("recorder_timeout", "pw-record cleanup did not complete")
|
||||
|
||||
def poll_unexpected(self) -> RecorderError | None:
|
||||
if not self.cleanup_started:
|
||||
return None
|
||||
self.cleanup_attempts += 1
|
||||
if self.cleanup_attempts >= 3:
|
||||
self.cleanup_complete = True
|
||||
return None
|
||||
|
||||
|
||||
class DelayedCleanupRecorderFactory:
|
||||
def __init__(self) -> None:
|
||||
self.instance = DelayedCleanupRecorder()
|
||||
|
||||
def __call__(self) -> DelayedCleanupRecorder:
|
||||
return self.instance
|
||||
|
||||
|
||||
class DelayedCloseTurn:
|
||||
sample_rate_hz = 16_000
|
||||
retain_pcm = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.close_calls = 0
|
||||
|
||||
@property
|
||||
def cleanup_complete(self) -> bool:
|
||||
return self.close_calls >= 3
|
||||
|
||||
def enqueue_pcm(self, chunk: bytes) -> None:
|
||||
return None
|
||||
|
||||
def finish(self, capture: Capture | None = None) -> str:
|
||||
return "cloud result"
|
||||
|
||||
def poll_error(self) -> None:
|
||||
return None
|
||||
|
||||
def cancel(self) -> None:
|
||||
return None
|
||||
|
||||
def close(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
|
||||
class StreamingTurn:
|
||||
sample_rate_hz = 24_000
|
||||
retain_pcm = False
|
||||
|
||||
def __init__(self, events: list[str]) -> None:
|
||||
self.events = events
|
||||
self.closed = False
|
||||
|
||||
@property
|
||||
def cleanup_complete(self) -> bool:
|
||||
return self.closed
|
||||
|
||||
def enqueue_pcm(self, chunk: bytes) -> None:
|
||||
self.events.append(f"audio-{chunk.decode('ascii')}")
|
||||
|
||||
def finish(self, capture: Capture | None = None) -> str:
|
||||
self.events.append("finish")
|
||||
self.assert_capture = capture
|
||||
return "cloud result"
|
||||
|
||||
def poll_error(self) -> None:
|
||||
return None
|
||||
|
||||
def cancel(self) -> None:
|
||||
self.events.append("cancel")
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class StreamingRecorder(FakeRecorder):
|
||||
def __init__(self, sink: StreamingTurn, events: list[str]) -> None:
|
||||
super().__init__()
|
||||
self.sink = sink
|
||||
self.events = events
|
||||
self.capture = Capture(b"", -15, 15, b"")
|
||||
|
||||
def start(self) -> None:
|
||||
super().start()
|
||||
self.sink.enqueue_pcm(b"a")
|
||||
self.sink.enqueue_pcm(b"b")
|
||||
|
||||
def stop(self) -> Capture:
|
||||
self.events.append("recorder-stop")
|
||||
return super().stop()
|
||||
|
||||
|
||||
class StreamingRecorderFactory:
|
||||
def __init__(self, events: list[str]) -> None:
|
||||
self.events = events
|
||||
self.sample_rate_hz: int | None = None
|
||||
self.instance: StreamingRecorder | None = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
sample_rate_hz: int,
|
||||
pcm_sink: StreamingTurn,
|
||||
retain_pcm: bool,
|
||||
) -> StreamingRecorder:
|
||||
self.sample_rate_hz = sample_rate_hz
|
||||
if retain_pcm:
|
||||
raise AssertionError("cloud capture must not retain PCM")
|
||||
self.instance = StreamingRecorder(pcm_sink, self.events)
|
||||
return self.instance
|
||||
|
||||
|
||||
class FakeBackend:
|
||||
sample_rate_hz = 16_000
|
||||
descriptor = type("Descriptor", (), {"backend": "fake", "model": "fake-model"})()
|
||||
|
||||
def __init__(self, turn: DelayedCloseTurn | StreamingTurn) -> None:
|
||||
self.turn = turn
|
||||
|
||||
def start_turn(self) -> DelayedCloseTurn:
|
||||
return self.turn
|
||||
|
||||
def cancel_prepare(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class RecorderFactory:
|
||||
def __init__(self) -> None:
|
||||
self.instances: list[FakeRecorder] = []
|
||||
@@ -271,6 +430,121 @@ class StopAndInferenceTests(CoreTestCase):
|
||||
self.assertEqual(completed[0]["transcript"], "dictated text")
|
||||
self.assertEqual(completed[1]["state"], "idle")
|
||||
|
||||
def test_realtime_backend_streams_ordered_pcm_and_finalizes_after_recorder_stop(self) -> None:
|
||||
events: list[str] = []
|
||||
turn = StreamingTurn(events)
|
||||
recorders = StreamingRecorderFactory(events)
|
||||
core = DictateServiceCore(
|
||||
backend=FakeBackend(turn),
|
||||
recorder_factory=recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
core.accepted("a")
|
||||
core.handle_request("a", parse_request({"version": 2, "type": "hello", "requestId": "hello"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "start",
|
||||
"requestId": "start",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
self.assertEqual(recorders.sample_rate_hz, 24_000)
|
||||
self.assertEqual(events, ["audio-a", "audio-b"])
|
||||
core.drain_events()
|
||||
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "stop",
|
||||
"requestId": "stop",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
|
||||
self.assertEqual(events, ["audio-a", "audio-b", "recorder-stop", "finish"])
|
||||
result = [event.frame for event in core.drain_events() if event.frame["type"] == "result"]
|
||||
self.assertEqual(result[0]["transcript"], "cloud result")
|
||||
self.assertIsNone(core.active_generation_id)
|
||||
|
||||
def test_cancel_after_prepare_completion_tracks_turn_cleanup_before_release(self) -> None:
|
||||
events: list[str] = []
|
||||
turn = StreamingTurn(events)
|
||||
core = DictateServiceCore(
|
||||
backend=FakeBackend(turn),
|
||||
recorder_factory=StreamingRecorderFactory(events),
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
core.accepted("a")
|
||||
core.handle_request("a", parse_request({"version": 2, "type": "hello", "requestId": "hello"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "start",
|
||||
"requestId": "start",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
# The future has completed, but the socket loop has not yet applied it.
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "cancel",
|
||||
"requestId": "cancel",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
|
||||
core.poll()
|
||||
|
||||
self.assertIn("cancel", events)
|
||||
self.assertTrue(turn.cleanup_complete)
|
||||
self.assertIsNone(core.active_generation_id)
|
||||
|
||||
def test_completed_turn_retries_bounded_cleanup_before_releasing_generation(self) -> None:
|
||||
turn = DelayedCloseTurn()
|
||||
core = DictateServiceCore(
|
||||
backend=FakeBackend(turn),
|
||||
recorder_factory=self.recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
core.accepted("a")
|
||||
core.handle_request("a", parse_request({"version": 2, "type": "hello", "requestId": "hello"}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "start",
|
||||
"requestId": "start",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "stop",
|
||||
"requestId": "stop",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
self.executor.run_next()
|
||||
|
||||
core.poll()
|
||||
self.assertEqual(core.active_generation_id, "g1")
|
||||
self.assertEqual(turn.close_calls, 2)
|
||||
|
||||
core.poll()
|
||||
self.assertIsNone(core.active_generation_id)
|
||||
self.assertEqual(turn.close_calls, 3)
|
||||
|
||||
def test_empty_transcript_is_a_result(self) -> None:
|
||||
self.model.transcript = ""
|
||||
self.load_to_listening()
|
||||
@@ -321,6 +595,8 @@ class StopAndInferenceTests(CoreTestCase):
|
||||
recorder = self.load_to_listening()
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 3")
|
||||
|
||||
self.core.poll()
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
failed = self.frames(self.core.drain_events(), "a")
|
||||
|
||||
@@ -332,12 +608,286 @@ class StopAndInferenceTests(CoreTestCase):
|
||||
recorder.unexpected = RecorderError("recorder_exit", "pw-record exited with code 0")
|
||||
|
||||
failed = self.request("a", "stop", "stop", "g1")
|
||||
self.executor.run_next()
|
||||
self.core.poll()
|
||||
failed.extend(self.frames(self.core.drain_events(), "a"))
|
||||
|
||||
self.assertEqual(failed[0]["code"], "recorder_exit")
|
||||
self.assertEqual(failed[-1]["state"], "idle")
|
||||
self.assertEqual(self.executor.pending, [])
|
||||
|
||||
|
||||
def test_recorder_spawn_failure_releases_generation_without_cleanup_loop(self) -> None:
|
||||
core = DictateServiceCore(
|
||||
model=self.model,
|
||||
recorder_factory=SpawnFailRecorderFactory(),
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
core.accepted("a")
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2, "type": "hello", "requestId": "hello"
|
||||
}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "start",
|
||||
"requestId": "start",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
|
||||
core.poll()
|
||||
|
||||
events = self.frames(core.drain_events(), "a")
|
||||
self.assertEqual(events[-2]["code"], "recorder_spawn")
|
||||
self.assertEqual(events[-1]["state"], "idle")
|
||||
self.assertIsNone(core.active_generation_id)
|
||||
self.assertEqual(self.executor.pending, [])
|
||||
|
||||
def test_unreaped_recorder_keeps_generation_busy_until_cleanup_completes(self) -> None:
|
||||
recorders = DelayedCleanupRecorderFactory()
|
||||
core = DictateServiceCore(
|
||||
model=self.model,
|
||||
recorder_factory=recorders,
|
||||
executor=self.executor,
|
||||
clock=self.clock,
|
||||
idle_seconds=900,
|
||||
)
|
||||
core.accepted("a")
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2, "type": "hello", "requestId": "hello"
|
||||
}))
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "start",
|
||||
"requestId": "start",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
core.drain_events()
|
||||
core.handle_request("a", parse_request({
|
||||
"version": 2,
|
||||
"type": "stop",
|
||||
"requestId": "stop",
|
||||
"generationId": "g1",
|
||||
}))
|
||||
self.executor.run_next()
|
||||
|
||||
core.poll()
|
||||
self.assertEqual(core.active_generation_id, "g1")
|
||||
# First reap the cancellation worker scheduled by the terminal stop.
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
for attempt in range(1, 4):
|
||||
self.executor.run_next()
|
||||
core.poll()
|
||||
if attempt < 3:
|
||||
self.assertEqual(core.active_generation_id, "g1")
|
||||
self.assertIsNone(core.active_generation_id)
|
||||
self.assertEqual(recorders.instance.cleanup_attempts, 3)
|
||||
|
||||
|
||||
class SocketLoopRegressionTests(unittest.TestCase):
|
||||
def test_output_overflow_does_not_use_the_closed_socket_from_same_selector_event(self) -> None:
|
||||
frame = {"padding": "x" * (MAX_FRAME_BYTES - len(encode_frame({"padding": ""})))}
|
||||
self.assertEqual(len(encode_frame(frame)), MAX_FRAME_BYTES)
|
||||
client = _FakeClientSocket()
|
||||
listener = _FakeListener(client)
|
||||
selector = _FakeSelector(client)
|
||||
core = _FakeLoopCore(
|
||||
Outbound("client-1", frame),
|
||||
)
|
||||
|
||||
with patch("pi_dictate_service.server.selectors.DefaultSelector", return_value=selector):
|
||||
serve(listener, core)
|
||||
|
||||
self.assertEqual(client.send_calls, 0)
|
||||
self.assertTrue(client.closed)
|
||||
self.assertTrue(listener.closed)
|
||||
|
||||
def test_close_does_not_wait_for_owned_executor_workers(self) -> None:
|
||||
executor = _ShutdownRecordingExecutor()
|
||||
with patch("pi_dictate_service.server._DaemonExecutor", return_value=executor):
|
||||
core = DictateServiceCore(model=FakeModel())
|
||||
core.accepted("a")
|
||||
core.handle_request(
|
||||
"a", parse_request({"version": PROTOCOL_VERSION, "type": "hello", "requestId": "h"})
|
||||
)
|
||||
core.drain_events()
|
||||
core.handle_request(
|
||||
"a",
|
||||
parse_request(
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "start",
|
||||
"requestId": "s",
|
||||
"generationId": "g",
|
||||
}
|
||||
),
|
||||
)
|
||||
core.close()
|
||||
|
||||
self.assertEqual(executor.shutdown_args, (False, True))
|
||||
|
||||
def test_real_owned_executor_does_not_block_interpreter_exit(self) -> None:
|
||||
program = """
|
||||
import threading
|
||||
import time
|
||||
from pi_dictate_service.protocol import parse_request
|
||||
from pi_dictate_service.server import DictateServiceCore
|
||||
|
||||
class StuckModel:
|
||||
loaded = False
|
||||
def load(self):
|
||||
threading.Event().wait()
|
||||
def transcribe(self, _samples):
|
||||
return ''
|
||||
|
||||
core = DictateServiceCore(model=StuckModel())
|
||||
core.accepted('a')
|
||||
core.handle_request('a', parse_request({'version': 2, 'type': 'hello', 'requestId': 'h'}))
|
||||
core.drain_events()
|
||||
core.handle_request('a', parse_request({
|
||||
'version': 2, 'type': 'start', 'requestId': 's', 'generationId': 'g'
|
||||
}))
|
||||
time.sleep(0.05)
|
||||
core.close()
|
||||
print('closed')
|
||||
"""
|
||||
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", program],
|
||||
text=True,
|
||||
capture_output=True,
|
||||
timeout=2,
|
||||
check=False,
|
||||
)
|
||||
|
||||
self.assertEqual(result.returncode, 0, result.stderr)
|
||||
self.assertEqual(result.stdout.strip(), "closed")
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeSelectorKey:
|
||||
data: str | None
|
||||
|
||||
|
||||
class _FakeSelector:
|
||||
def __init__(self, client: "_FakeClientSocket") -> None:
|
||||
self._events = [
|
||||
[(_FakeSelectorKey(None), selectors.EVENT_READ)],
|
||||
[(_FakeSelectorKey("client-1"), selectors.EVENT_READ | selectors.EVENT_WRITE)],
|
||||
]
|
||||
self.client = client
|
||||
self.unregistered: list[object] = []
|
||||
self.closed = False
|
||||
|
||||
def register(self, fileobj: object, _events: int, _data: object = None) -> None:
|
||||
return None
|
||||
|
||||
def unregister(self, fileobj: object) -> None:
|
||||
self.unregistered.append(fileobj)
|
||||
|
||||
def modify(self, fileobj: object, _events: int, _data: object = None) -> None:
|
||||
if fileobj is self.client and self.client.closed:
|
||||
raise OSError(errno.EBADF, "bad file descriptor")
|
||||
|
||||
def select(self, timeout: float | None = None) -> list[tuple[_FakeSelectorKey, int]]:
|
||||
return self._events.pop(0)
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeClientSocket:
|
||||
def __init__(self) -> None:
|
||||
self.closed = False
|
||||
self.send_calls = 0
|
||||
self._read = True
|
||||
|
||||
def setblocking(self, _value: bool) -> None:
|
||||
return None
|
||||
|
||||
def recv(self, _size: int) -> bytes:
|
||||
if self._read:
|
||||
self._read = False
|
||||
return b"{broken}\n"
|
||||
return b""
|
||||
|
||||
def send(self, _data: bytearray) -> int:
|
||||
self.send_calls += 1
|
||||
raise OSError(errno.EBADF, "bad file descriptor")
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeListener:
|
||||
def __init__(self, client: _FakeClientSocket) -> None:
|
||||
self.client = client
|
||||
self.closed = False
|
||||
self._accepted = False
|
||||
|
||||
def setblocking(self, _value: bool) -> None:
|
||||
return None
|
||||
|
||||
def accept(self) -> tuple[_FakeClientSocket, object]:
|
||||
if self._accepted:
|
||||
raise BlockingIOError
|
||||
self._accepted = True
|
||||
return self.client, None
|
||||
|
||||
def close(self) -> None:
|
||||
self.closed = True
|
||||
|
||||
|
||||
class _FakeLoopCore:
|
||||
def __init__(self, outbound: Outbound) -> None:
|
||||
self.outbound = outbound
|
||||
self.accepted_client: str | None = None
|
||||
self._outbound_sent = False
|
||||
self._exit_checks = 0
|
||||
|
||||
def should_exit(self) -> bool:
|
||||
self._exit_checks += 1
|
||||
return self._exit_checks >= 3
|
||||
|
||||
def poll(self) -> None:
|
||||
return None
|
||||
|
||||
def drain_events(self) -> list[Outbound]:
|
||||
if self.accepted_client is not None and not self._outbound_sent:
|
||||
self._outbound_sent = True
|
||||
return [self.outbound]
|
||||
return []
|
||||
|
||||
def accepted(self, client_id: str) -> None:
|
||||
self.accepted_client = client_id
|
||||
|
||||
def disconnect(self, _client_id: str) -> None:
|
||||
return None
|
||||
|
||||
def close(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
class _ShutdownRecordingExecutor:
|
||||
def __init__(self) -> None:
|
||||
self.shutdown_args: tuple[bool, bool] | None = None
|
||||
|
||||
def submit(self, _function: object, *_args: object) -> Future[object]:
|
||||
return Future()
|
||||
|
||||
def shutdown(self, *, wait: bool = True, cancel_futures: bool = False) -> None:
|
||||
self.shutdown_args = (wait, cancel_futures)
|
||||
if wait:
|
||||
raise AssertionError("shutdown waited for a hung worker")
|
||||
|
||||
|
||||
class IdleTests(CoreTestCase):
|
||||
def test_only_valid_requests_reset_idle_and_workers_suspend_then_restart_timer(self) -> None:
|
||||
self.core.accepted("a")
|
||||
|
||||
Reference in New Issue
Block a user