feat(dictate): use local socket-activated Parakeet
This commit is contained in:
@@ -0,0 +1,225 @@
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import socket
|
||||
import struct
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from pi_dictate_service.server import (
|
||||
MAX_CLIENT_CONNECTIONS,
|
||||
DictateServiceCore,
|
||||
inherited_listener,
|
||||
main,
|
||||
serve,
|
||||
)
|
||||
|
||||
|
||||
class LoadedModel:
|
||||
loaded = True
|
||||
|
||||
def load(self) -> None:
|
||||
return None
|
||||
|
||||
def transcribe(self, samples: list[float]) -> str:
|
||||
return ""
|
||||
|
||||
|
||||
class UnixLoopTests(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.temporary = tempfile.TemporaryDirectory()
|
||||
self.socket_path = Path(self.temporary.name) / "dictate.sock"
|
||||
self.listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
self.listener.bind(str(self.socket_path))
|
||||
self.listener.listen()
|
||||
self.core = DictateServiceCore(model=LoadedModel(), idle_seconds=0.4)
|
||||
self.failure: BaseException | None = None
|
||||
|
||||
def run() -> None:
|
||||
try:
|
||||
serve(self.listener, self.core)
|
||||
except BaseException as error:
|
||||
self.failure = error
|
||||
|
||||
self.thread = threading.Thread(target=run)
|
||||
self.thread.start()
|
||||
|
||||
def tearDown(self) -> None:
|
||||
self.thread.join(timeout=2)
|
||||
if self.thread.is_alive():
|
||||
self.fail("server loop did not exit after its idle timeout")
|
||||
if self.failure:
|
||||
raise self.failure
|
||||
self.temporary.cleanup()
|
||||
|
||||
def connect(self) -> socket.socket:
|
||||
client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
client.settimeout(1)
|
||||
client.connect(str(self.socket_path))
|
||||
return client
|
||||
|
||||
@staticmethod
|
||||
def receive_frame(client: socket.socket) -> dict[str, object]:
|
||||
data = bytearray()
|
||||
while not data.endswith(b"\n"):
|
||||
chunk = client.recv(4096)
|
||||
if not chunk:
|
||||
break
|
||||
data.extend(chunk)
|
||||
return json.loads(data)
|
||||
|
||||
def test_accepts_fragmented_requests_and_writes_persistent_responses(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello",')
|
||||
client.sendall(b'"requestId":"hello"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
client.sendall(b'{"version":1,"type":"ping","requestId":"ping"}\n')
|
||||
pong = self.receive_frame(client)
|
||||
self.assertEqual(pong["requestId"], "ping")
|
||||
self.assertEqual(pong["service"], "pi-dictate")
|
||||
|
||||
def test_valid_state_error_keeps_open_but_malformed_frame_errors_once_then_closes(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.receive_frame(client)
|
||||
client.sendall(
|
||||
b'{"version":1,"type":"stop","requestId":"stop","generationId":"g"}\n'
|
||||
)
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_state")
|
||||
client.sendall(b'{"version":1,"type":"ping","requestId":"ping"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
client.sendall(b'{broken}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "malformed_json")
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
|
||||
self.assertTrue(self.socket_path.exists(), "the service must not unlink a listener-owned path")
|
||||
|
||||
def test_malicious_json_values_do_not_crash_the_service(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":[],"requestId":"bad-type"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
|
||||
client.sendall(b'{"version":1,"type":"ping","requestId":"\\ud800"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["code"], "invalid_request")
|
||||
|
||||
with self.connect() as client:
|
||||
huge_integer = b"9" * 5000
|
||||
client.sendall(
|
||||
b'{"version":' + huge_integer + b',"type":"hello","requestId":"huge"}\n'
|
||||
)
|
||||
self.assertEqual(self.receive_frame(client)["code"], "malformed_json")
|
||||
self.assertEqual(client.recv(1), b"")
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
def test_closes_non_reading_client_before_output_queue_can_grow_unbounded(self) -> None:
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.receive_frame(client)
|
||||
ping = b'{"version":1,"type":"ping","requestId":"p"}\n'
|
||||
try:
|
||||
client.sendall(ping * 1200)
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
received = bytearray()
|
||||
while True:
|
||||
try:
|
||||
chunk = client.recv(64 * 1024)
|
||||
except ConnectionResetError:
|
||||
break
|
||||
if not chunk:
|
||||
break
|
||||
received.extend(chunk)
|
||||
self.assertLessEqual(len(received), 64 * 1024)
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"healthy"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
def test_rejects_connections_above_the_fixed_limit(self) -> None:
|
||||
clients = [self.connect() for _ in range(MAX_CLIENT_CONNECTIONS)]
|
||||
try:
|
||||
time.sleep(0.1)
|
||||
overflow = self.connect()
|
||||
try:
|
||||
overflow.sendall(b'{"version":1,"type":"hello","requestId":"overflow"}\n')
|
||||
with self.assertRaises((ConnectionResetError, BrokenPipeError)):
|
||||
if overflow.recv(1) == b"":
|
||||
raise ConnectionResetError
|
||||
finally:
|
||||
overflow.close()
|
||||
finally:
|
||||
for client in clients:
|
||||
client.close()
|
||||
|
||||
def test_client_reset_does_not_crash_the_service(self) -> None:
|
||||
reset_client = self.connect()
|
||||
reset_client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
|
||||
reset_client.sendall(b'{"version":1,"type":"hello","requestId":"abandoned"}\n')
|
||||
reset_client.close()
|
||||
time.sleep(0.15)
|
||||
|
||||
with self.connect() as client:
|
||||
client.sendall(b'{"version":1,"type":"hello","requestId":"hello"}\n')
|
||||
self.assertEqual(self.receive_frame(client)["type"], "pong")
|
||||
|
||||
|
||||
class ActivationTests(unittest.TestCase):
|
||||
def test_inherited_listener_requires_one_fd_for_this_process_and_wraps_fd_three(self) -> None:
|
||||
listener = Mock()
|
||||
with patch("pi_dictate_service.server.os.getpid", return_value=42), patch(
|
||||
"pi_dictate_service.server.socket.socket", return_value=listener
|
||||
) as socket_factory:
|
||||
self.assertIs(
|
||||
inherited_listener({"LISTEN_PID": "42", "LISTEN_FDS": "1"}),
|
||||
listener,
|
||||
)
|
||||
|
||||
socket_factory.assert_called_once_with(fileno=3)
|
||||
listener.setblocking.assert_called_once_with(False)
|
||||
|
||||
with patch("pi_dictate_service.server.os.getpid", return_value=42):
|
||||
with self.assertRaisesRegex(RuntimeError, "exactly one"):
|
||||
inherited_listener({"LISTEN_PID": "41", "LISTEN_FDS": "1"})
|
||||
|
||||
def test_main_requires_local_model_path_and_reads_idle_override(self) -> None:
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with self.assertRaisesRegex(SystemExit, "PI_DICTATE_MODEL_PATH"):
|
||||
main()
|
||||
|
||||
model = object()
|
||||
listener = object()
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"PI_DICTATE_MODEL_PATH": "/models/parakeet.nemo",
|
||||
"PI_DICTATE_TERMS_PATH": "/config/coding-terms.txt",
|
||||
"PI_DICTATE_BOOST_ALPHA": "0.5",
|
||||
"PI_DICTATE_IDLE_SECONDS": "12.5",
|
||||
},
|
||||
clear=True,
|
||||
), patch("pi_dictate_service.server.NemoParakeetModel", return_value=model) as adapter, patch(
|
||||
"pi_dictate_service.server.DictateServiceCore"
|
||||
) as core_factory, patch(
|
||||
"pi_dictate_service.server.inherited_listener", return_value=listener
|
||||
), patch("pi_dictate_service.server.serve") as serve_loop:
|
||||
core = core_factory.return_value
|
||||
main()
|
||||
|
||||
adapter.assert_called_once_with(
|
||||
"/models/parakeet.nemo",
|
||||
terms_path="/config/coding-terms.txt",
|
||||
boost_alpha=0.5,
|
||||
)
|
||||
core_factory.assert_called_once_with(model=model, idle_seconds=12.5)
|
||||
serve_loop.assert_called_once_with(listener, core)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user