| from __future__ import annotations | |
| import asyncio | |
| from concurrent.futures import ThreadPoolExecutor | |
| from typing import Any, AsyncIterator | |
| from audio.pipeline import AudioPipeline | |
| _SENTINEL = object() | |
| class GpuSerialQueue: | |
| """Bounded replica pool. Each replica runs batched Whisper on one file.""" | |
| def __init__(self, pipeline: AudioPipeline, max_queue: int, gpu_concurrency: int) -> None: | |
| self.pipeline = pipeline | |
| workers = max(1, gpu_concurrency) | |
| self._gpu = asyncio.Semaphore(workers) | |
| self._slots = asyncio.Semaphore(max_queue) | |
| self._pool = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="asr") | |
| async def run(self, **kwargs: Any) -> dict: | |
| try: | |
| await asyncio.wait_for(self._slots.acquire(), timeout=0.05) | |
| except TimeoutError as exc: | |
| raise QueueFullError("audio GPU queue is full") from exc | |
| try: | |
| async with self._gpu: | |
| loop = asyncio.get_running_loop() | |
| return await loop.run_in_executor( | |
| self._pool, lambda: self.pipeline.run(**kwargs) | |
| ) | |
| finally: | |
| self._slots.release() | |
| async def stream(self, **kwargs: Any) -> AsyncIterator[dict[str, Any]]: | |
| try: | |
| await asyncio.wait_for(self._slots.acquire(), timeout=0.05) | |
| except TimeoutError as exc: | |
| raise QueueFullError("audio GPU queue is full") from exc | |
| await self._gpu.acquire() | |
| loop = asyncio.get_running_loop() | |
| out: asyncio.Queue[Any] = asyncio.Queue() | |
| def produce() -> None: | |
| try: | |
| for event in self.pipeline.run_events(**kwargs): | |
| asyncio.run_coroutine_threadsafe(out.put(event), loop).result() | |
| asyncio.run_coroutine_threadsafe(out.put(_SENTINEL), loop).result() | |
| except Exception as exc: | |
| asyncio.run_coroutine_threadsafe(out.put(exc), loop).result() | |
| try: | |
| self._pool.submit(produce) | |
| while True: | |
| item = await out.get() | |
| if item is _SENTINEL: | |
| break | |
| if isinstance(item, Exception): | |
| raise item | |
| yield item | |
| finally: | |
| self._gpu.release() | |
| self._slots.release() | |
| class QueueFullError(RuntimeError): | |
| pass | |