File size: 2,460 Bytes
0018659 e5a8331 0018659 e5a8331 0018659 e5a8331 0018659 e5a8331 0018659 e5a8331 0018659 e5a8331 0018659 e5a8331 0018659 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | 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
|