Spaces:
Sleeping
Sleeping
File size: 1,794 Bytes
f16b727 | 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 | """Bound CPU work without blocking the HTTP event loop or unbounded queues."""
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Callable, TypeVar
T = TypeVar("T")
class InferenceBusy(Exception):
pass
class InferenceTimedOut(Exception):
pass
class BoundedInferenceRunner:
def __init__(self, outstanding: int = 2, timeout_seconds: float = 90):
self._slots = asyncio.Semaphore(outstanding)
self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="inference")
self._timeout = timeout_seconds
self._closed = False
async def run(self, operation: Callable[[], T]) -> T:
if self._closed:
raise InferenceBusy()
try:
await asyncio.wait_for(self._slots.acquire(), timeout=0.05)
except asyncio.TimeoutError as exc:
raise InferenceBusy() from exc
try:
future = asyncio.get_running_loop().run_in_executor(self._executor, operation)
except BaseException:
self._slots.release()
raise
# A timed-out thread cannot be safely killed. Keep its slot occupied until
# the actual work finishes; otherwise retries accumulate hidden CPU work.
def completed(done):
self._slots.release()
if not done.cancelled():
done.exception() # Consume late failures without logging private inputs.
future.add_done_callback(completed)
try:
return await asyncio.wait_for(asyncio.shield(future), timeout=self._timeout)
except asyncio.TimeoutError as exc:
raise InferenceTimedOut() from exc
def close(self):
self._closed = True
self._executor.shutdown(wait=True, cancel_futures=True)
|