| from __future__ import annotations |
|
|
| import asyncio |
| import json |
| import logging |
| import struct |
| import time |
| import uuid |
| from dataclasses import dataclass, field |
| from typing import Any, Optional |
|
|
| import numpy as np |
| from aiortc import ( |
| RTCConfiguration, |
| RTCIceServer, |
| RTCPeerConnection, |
| RTCSessionDescription, |
| ) |
| from starlette.websockets import WebSocket, WebSocketState |
|
|
| from app.config import Settings |
| from app.engine import ( |
| FRAME_DIMENSION, |
| MOTION_LAYOUT, |
| GenerationRequest, |
| GenerationResult, |
| SignMotionEngine, |
| split_sentences, |
| ) |
|
|
| logger = logging.getLogger("realtime") |
|
|
| PROTOCOL_VERSION = 1 |
| MOTION_BLOCK_HEADER = struct.Struct("<III") |
|
|
|
|
| def pack_motion_block(request_index: int, frame_offset: int, frames: np.ndarray) -> bytes: |
| contiguous_frames = np.ascontiguousarray(frames, dtype="<f4") |
| header = MOTION_BLOCK_HEADER.pack(request_index, frame_offset, int(contiguous_frames.shape[0])) |
| return header + contiguous_frames.tobytes() |
|
|
|
|
| class Transport: |
| async def send_text(self, text: str) -> None: |
| raise NotImplementedError |
|
|
| async def send_bytes(self, data: bytes) -> None: |
| raise NotImplementedError |
|
|
| async def wait_until_writable(self) -> None: |
| return None |
|
|
| def is_open(self) -> bool: |
| raise NotImplementedError |
|
|
|
|
| class WebSocketTransport(Transport): |
| def __init__(self, websocket: WebSocket) -> None: |
| self.websocket = websocket |
|
|
| async def send_text(self, text: str) -> None: |
| await self.websocket.send_text(text) |
|
|
| async def send_bytes(self, data: bytes) -> None: |
| await self.websocket.send_bytes(data) |
|
|
| def is_open(self) -> bool: |
| client_connected = self.websocket.client_state == WebSocketState.CONNECTED |
| application_connected = self.websocket.application_state == WebSocketState.CONNECTED |
| return client_connected and application_connected |
|
|
|
|
| class DataChannelTransport(Transport): |
| def __init__(self, channel: Any, high_water_mark: int) -> None: |
| self.channel = channel |
| self.high_water_mark = high_water_mark |
|
|
| async def send_text(self, text: str) -> None: |
| self.channel.send(text) |
|
|
| async def send_bytes(self, data: bytes) -> None: |
| self.channel.send(data) |
|
|
| async def wait_until_writable(self) -> None: |
| while self.channel.readyState == "open" and self.channel.bufferedAmount > self.high_water_mark: |
| await asyncio.sleep(0.005) |
|
|
| def is_open(self) -> bool: |
| return self.channel.readyState == "open" |
|
|
|
|
| class ConnectionLimiter: |
| def __init__(self, maximum: int) -> None: |
| self.maximum = maximum |
| self.active = 0 |
|
|
| def acquire(self) -> bool: |
| if self.active >= self.maximum: |
| return False |
| self.active += 1 |
| return True |
|
|
| def release(self) -> None: |
| if self.active > 0: |
| self.active -= 1 |
|
|
|
|
| @dataclass |
| class RequestState: |
| request_id: str |
| request_index: int |
| sentence_count: int |
| completed_sentences: int = 0 |
| frame_offset: int = 0 |
| started_at: float = field(default_factory=time.perf_counter) |
| cancelled: bool = False |
|
|
|
|
| @dataclass |
| class Job: |
| session: "RealtimeSession" |
| request_index: int |
| sentence_index: int |
| generation_request: GenerationRequest |
|
|
|
|
| @dataclass |
| class OutgoingItem: |
| kind: str |
| payload: Optional[dict] = None |
| job: Optional[Job] = None |
| result: Optional[GenerationResult] = None |
| message: str = "" |
|
|
|
|
| class Scheduler: |
| def __init__(self, engine: SignMotionEngine, settings: Settings) -> None: |
| self.engine = engine |
| self.settings = settings |
| self.queue: asyncio.Queue = asyncio.Queue(maxsize=settings.queue_size) |
| self.worker_task: Optional[asyncio.Task] = None |
|
|
| async def start(self) -> None: |
| self.worker_task = asyncio.create_task(self.run_worker()) |
|
|
| async def stop(self) -> None: |
| if self.worker_task is None: |
| return |
| self.worker_task.cancel() |
| try: |
| await self.worker_task |
| except asyncio.CancelledError: |
| pass |
| self.worker_task = None |
|
|
| def free_slots(self) -> int: |
| return self.settings.queue_size - self.queue.qsize() |
|
|
| def submit_all(self, jobs: list) -> bool: |
| if len(jobs) > self.free_slots(): |
| return False |
| for job in jobs: |
| self.queue.put_nowait(job) |
| return True |
|
|
| async def run_worker(self) -> None: |
| while True: |
| job = await self.queue.get() |
| try: |
| if job.session.should_skip(job.request_index): |
| continue |
| result = await asyncio.to_thread(self.engine.generate, job.generation_request) |
| job.session.enqueue_result(job, result) |
| except asyncio.CancelledError: |
| raise |
| except Exception as error: |
| logger.exception("generation failed") |
| job.session.enqueue_failure(job, str(error)) |
| finally: |
| self.queue.task_done() |
|
|
|
|
| class RealtimeSession: |
| def __init__( |
| self, |
| transport: Transport, |
| engine: SignMotionEngine, |
| scheduler: Scheduler, |
| settings: Settings, |
| label: str, |
| ) -> None: |
| self.transport = transport |
| self.engine = engine |
| self.scheduler = scheduler |
| self.settings = settings |
| self.label = label |
| self.session_id = uuid.uuid4().hex[:12] |
| self.outbox: asyncio.Queue = asyncio.Queue() |
| self.requests: dict = {} |
| self.next_request_index = 1 |
| self.closed = False |
| self.sender_task: Optional[asyncio.Task] = None |
|
|
| def start(self) -> None: |
| self.sender_task = asyncio.create_task(self.run_sender()) |
| self.enqueue_json(self.build_hello()) |
|
|
| def build_hello(self) -> dict: |
| return { |
| "type": "hello", |
| "protocol_version": PROTOCOL_VERSION, |
| "session_id": self.session_id, |
| "transport": self.label, |
| "ready": self.engine.ready, |
| "fps": self.settings.fps, |
| "frame_dimension": FRAME_DIMENSION, |
| "layout": MOTION_LAYOUT, |
| "binary_block_header": "uint32 request_index, uint32 frame_offset, uint32 frame_count, then float32 frames, all little endian", |
| "limits": { |
| "maximum_text_characters": self.settings.maximum_text_characters, |
| "maximum_sentences_per_request": self.settings.maximum_sentences_per_request, |
| "maximum_pending_requests": self.settings.maximum_pending_requests_per_connection, |
| "frames_per_message": self.settings.frames_per_message, |
| }, |
| } |
|
|
| def enqueue_json(self, payload: dict) -> None: |
| self.outbox.put_nowait(OutgoingItem(kind="json", payload=payload)) |
|
|
| def enqueue_error(self, request_id: Optional[str], code: str, message: str) -> None: |
| self.enqueue_json( |
| {"type": "error", "id": request_id, "code": code, "message": message} |
| ) |
|
|
| def enqueue_result(self, job: Job, result: GenerationResult) -> None: |
| if self.closed: |
| return |
| self.outbox.put_nowait(OutgoingItem(kind="sentence", job=job, result=result)) |
|
|
| def enqueue_failure(self, job: Job, message: str) -> None: |
| if self.closed: |
| return |
| self.outbox.put_nowait(OutgoingItem(kind="failure", job=job, message=message)) |
|
|
| def should_skip(self, request_index: int) -> bool: |
| if self.closed: |
| return True |
| state = self.requests.get(request_index) |
| if state is None: |
| return True |
| return state.cancelled |
|
|
| def handle_message(self, message: Any) -> None: |
| if self.closed: |
| return |
|
|
| if isinstance(message, (bytes, bytearray)): |
| self.enqueue_error(None, "invalid_message", "Binary messages are not accepted") |
| return |
|
|
| try: |
| data = json.loads(message) |
| except (TypeError, json.JSONDecodeError): |
| self.enqueue_error(None, "invalid_json", "Message is not valid JSON") |
| return |
|
|
| if not isinstance(data, dict): |
| self.enqueue_error(None, "invalid_message", "Message must be a JSON object") |
| return |
|
|
| message_type = data.get("type") |
|
|
| if message_type == "generate": |
| self.handle_generate(data) |
| elif message_type == "cancel": |
| self.handle_cancel(data) |
| elif message_type == "ping": |
| self.enqueue_json({"type": "pong", "time": data.get("time")}) |
| else: |
| self.enqueue_error(None, "unknown_type", "Unknown message type") |
|
|
| def handle_cancel(self, data: dict) -> None: |
| request_id = str(data.get("id")) |
| cancelled_indexes = [] |
|
|
| for request_index, state in self.requests.items(): |
| if state.request_id == request_id: |
| state.cancelled = True |
| cancelled_indexes.append(request_index) |
|
|
| for request_index in cancelled_indexes: |
| del self.requests[request_index] |
|
|
| self.enqueue_json({"type": "cancelled", "id": request_id, "count": len(cancelled_indexes)}) |
|
|
| def handle_generate(self, data: dict) -> None: |
| raw_request_id = data.get("id") |
| if raw_request_id is None: |
| request_id = uuid.uuid4().hex[:8] |
| else: |
| request_id = str(raw_request_id) |
|
|
| if not self.engine.ready: |
| self.enqueue_error(request_id, "loading", "The model is still loading") |
| return |
|
|
| text = data.get("text") |
| if not isinstance(text, str) or text.strip() == "": |
| self.enqueue_error(request_id, "invalid_text", "The text field must be a non empty string") |
| return |
|
|
| if len(text) > self.settings.maximum_text_characters: |
| self.enqueue_error( |
| request_id, |
| "text_too_long", |
| "The text is longer than " + str(self.settings.maximum_text_characters) + " characters", |
| ) |
| return |
|
|
| pace = data.get("pace", 1.0) |
| pace_is_number = isinstance(pace, (int, float)) and not isinstance(pace, bool) |
| if not pace_is_number or pace < 0.5 or pace > 2.0: |
| self.enqueue_error(request_id, "invalid_pace", "The pace must be a number from 0.5 to 2.0") |
| return |
|
|
| use_retrieval = data.get("retrieval", self.settings.use_retrieval) |
| if not isinstance(use_retrieval, bool): |
| self.enqueue_error(request_id, "invalid_retrieval", "The retrieval field must be true or false") |
| return |
|
|
| do_sample = data.get("sample", self.settings.do_sample) |
| if not isinstance(do_sample, bool): |
| self.enqueue_error(request_id, "invalid_sample", "The sample field must be true or false") |
| return |
|
|
| seed = data.get("seed", self.settings.seed) |
| if not isinstance(seed, int) or isinstance(seed, bool): |
| self.enqueue_error(request_id, "invalid_seed", "The seed field must be an integer") |
| return |
|
|
| if len(self.requests) >= self.settings.maximum_pending_requests_per_connection: |
| self.enqueue_error(request_id, "too_many_requests", "Too many unfinished requests on this connection") |
| return |
|
|
| sentences = split_sentences(text, self.settings.maximum_words_per_sentence) |
| if len(sentences) == 0: |
| self.enqueue_error(request_id, "invalid_text", "No sentences were found in the text") |
| return |
|
|
| if len(sentences) > self.settings.maximum_sentences_per_request: |
| self.enqueue_error(request_id, "too_many_sentences", "The text has too many sentences") |
| return |
|
|
| request_index = self.next_request_index |
| self.next_request_index += 1 |
|
|
| state = RequestState( |
| request_id=request_id, |
| request_index=request_index, |
| sentence_count=len(sentences), |
| ) |
| self.requests[request_index] = state |
|
|
| jobs = [] |
| for sentence_index, sentence in enumerate(sentences): |
| token_count = self.engine.plan_token_count(sentence, float(pace)) |
| generation_request = GenerationRequest( |
| text=sentence, |
| token_count=token_count, |
| use_retrieval=use_retrieval, |
| do_sample=do_sample, |
| temperature=self.settings.temperature, |
| top_k=self.settings.top_k, |
| seed=seed, |
| ) |
| jobs.append( |
| Job( |
| session=self, |
| request_index=request_index, |
| sentence_index=sentence_index, |
| generation_request=generation_request, |
| ) |
| ) |
|
|
| accepted = self.scheduler.submit_all(jobs) |
| if not accepted: |
| del self.requests[request_index] |
| self.enqueue_error(request_id, "busy", "The server queue is full, try again shortly") |
| return |
|
|
| self.enqueue_json( |
| { |
| "type": "accepted", |
| "id": request_id, |
| "request_index": request_index, |
| "sentences": sentences, |
| "fps": self.settings.fps, |
| } |
| ) |
|
|
| async def run_sender(self) -> None: |
| try: |
| while True: |
| item = await self.outbox.get() |
| await self.dispatch(item) |
| except asyncio.CancelledError: |
| raise |
| except Exception as error: |
| logger.info("sender stopped: %s", error) |
| self.closed = True |
|
|
| async def dispatch(self, item: OutgoingItem) -> None: |
| if not self.transport.is_open(): |
| raise ConnectionError("The transport is closed") |
|
|
| if item.kind == "json": |
| await self.transport.send_text(json.dumps(item.payload)) |
| return |
|
|
| if item.kind == "failure": |
| failed_job = item.job |
| failed_state = self.requests.get(failed_job.request_index) |
| if failed_state is None: |
| return |
| failed_state.cancelled = True |
| del self.requests[failed_job.request_index] |
| await self.transport.send_text( |
| json.dumps( |
| { |
| "type": "error", |
| "id": failed_state.request_id, |
| "code": "generation_failed", |
| "message": item.message, |
| } |
| ) |
| ) |
| return |
|
|
| if item.kind == "sentence": |
| await self.dispatch_sentence(item) |
|
|
| async def dispatch_sentence(self, item: OutgoingItem) -> None: |
| job = item.job |
| result = item.result |
| state = self.requests.get(job.request_index) |
|
|
| if state is None or state.cancelled: |
| return |
|
|
| frames = result.motion |
| frame_count = int(frames.shape[0]) |
|
|
| await self.transport.send_text( |
| json.dumps( |
| { |
| "type": "sentence", |
| "id": state.request_id, |
| "request_index": job.request_index, |
| "sentence_index": job.sentence_index, |
| "sentence_count": state.sentence_count, |
| "text": job.generation_request.text, |
| "frame_offset": state.frame_offset, |
| "frame_count": frame_count, |
| "fps": result.fps, |
| "seconds": frame_count / result.fps, |
| "generation_ms": round(result.elapsed_seconds * 1000.0, 1), |
| "cache_hit": result.from_cache, |
| "keywords": result.keywords, |
| } |
| ) |
| ) |
|
|
| frames_per_message = self.settings.frames_per_message |
| for start_index in range(0, frame_count, frames_per_message): |
| if state.cancelled: |
| return |
| block = frames[start_index:start_index + frames_per_message] |
| await self.transport.wait_until_writable() |
| await self.transport.send_bytes( |
| pack_motion_block(job.request_index, state.frame_offset + start_index, block) |
| ) |
|
|
| state.frame_offset += frame_count |
| state.completed_sentences += 1 |
|
|
| if state.completed_sentences >= state.sentence_count: |
| elapsed_milliseconds = (time.perf_counter() - state.started_at) * 1000.0 |
| await self.transport.send_text( |
| json.dumps( |
| { |
| "type": "done", |
| "id": state.request_id, |
| "request_index": state.request_index, |
| "total_frames": state.frame_offset, |
| "total_seconds": state.frame_offset / result.fps, |
| "elapsed_ms": round(elapsed_milliseconds, 1), |
| } |
| ) |
| ) |
| self.requests.pop(state.request_index, None) |
|
|
| async def close(self) -> None: |
| self.closed = True |
|
|
| for state in self.requests.values(): |
| state.cancelled = True |
| self.requests.clear() |
|
|
| if self.sender_task is not None: |
| self.sender_task.cancel() |
| try: |
| await self.sender_task |
| except (asyncio.CancelledError, Exception): |
| pass |
| self.sender_task = None |
|
|
|
|
| class WebRTCManager: |
| def __init__( |
| self, |
| settings: Settings, |
| engine: SignMotionEngine, |
| scheduler: Scheduler, |
| limiter: ConnectionLimiter, |
| ) -> None: |
| self.settings = settings |
| self.engine = engine |
| self.scheduler = scheduler |
| self.limiter = limiter |
| self.peers: dict = {} |
|
|
| def build_configuration(self) -> RTCConfiguration: |
| ice_servers = [] |
| for entry in self.settings.ice_servers: |
| ice_servers.append( |
| RTCIceServer( |
| urls=entry["urls"], |
| username=entry.get("username"), |
| credential=entry.get("credential"), |
| ) |
| ) |
| return RTCConfiguration(iceServers=ice_servers) |
|
|
| async def handle_offer(self, sdp: str, offer_type: str) -> dict: |
| if not self.limiter.acquire(): |
| raise ConnectionRefusedError("Too many active connections") |
|
|
| peer = RTCPeerConnection(configuration=self.build_configuration()) |
| holder: dict = {"session": None, "watchdog": None} |
| self.peers[peer] = holder |
|
|
| @peer.on("datachannel") |
| def on_datachannel(channel: Any) -> None: |
| if holder["session"] is not None: |
| channel.close() |
| return |
|
|
| transport = DataChannelTransport(channel, self.settings.data_channel_high_water_mark) |
| session = RealtimeSession(transport, self.engine, self.scheduler, self.settings, "webrtc") |
| holder["session"] = session |
|
|
| @channel.on("message") |
| def on_message(message: Any) -> None: |
| session.handle_message(message) |
|
|
| @channel.on("close") |
| def on_close() -> None: |
| asyncio.ensure_future(session.close()) |
|
|
| if channel.readyState == "open": |
| session.start() |
| else: |
|
|
| @channel.on("open") |
| def on_open() -> None: |
| session.start() |
|
|
| @peer.on("connectionstatechange") |
| async def on_connectionstatechange() -> None: |
| if peer.connectionState in ("failed", "closed"): |
| await self.release_peer(peer) |
|
|
| holder["watchdog"] = asyncio.create_task(self.expire_if_idle(peer, holder)) |
|
|
| try: |
| await peer.setRemoteDescription(RTCSessionDescription(sdp=sdp, type=offer_type)) |
| answer = await peer.createAnswer() |
| await peer.setLocalDescription(answer) |
| except Exception: |
| await self.release_peer(peer) |
| raise |
|
|
| return { |
| "sdp": peer.localDescription.sdp, |
| "type": peer.localDescription.type, |
| } |
|
|
| async def expire_if_idle(self, peer: RTCPeerConnection, holder: dict) -> None: |
| await asyncio.sleep(self.settings.idle_peer_timeout_seconds) |
| if peer in self.peers and holder["session"] is None: |
| await self.release_peer(peer) |
|
|
| async def release_peer(self, peer: RTCPeerConnection) -> None: |
| holder = self.peers.pop(peer, None) |
| if holder is None: |
| return |
|
|
| session = holder["session"] |
| if session is not None: |
| await session.close() |
|
|
| try: |
| await peer.close() |
| except Exception: |
| logger.exception("closing peer failed") |
|
|
| self.limiter.release() |
|
|
| async def shutdown(self) -> None: |
| for peer in list(self.peers.keys()): |
| await self.release_peer(peer) |
|
|
