core/inference_runner.py

"""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)