core/queue.py

"""A bounded, single-worker asynchronous inference queue."""

from __future__ import annotations

import asyncio
from contextvars import Context, copy_context
from dataclasses import dataclass
from typing import Any, Callable, Generic, TypeVar

from core.errors import GatewayError, QueueCapacityError


T = TypeVar("T")


@dataclass(slots=True)
class QueueJob(Generic[T]):
    request_id: str
    label: str
    operation: Callable[[], T]
    context: Context
    future: asyncio.Future[T]
    started: asyncio.Future[None]
    cancelled: bool = False


class InferenceQueue:
    """Serialize all accelerator work through one worker."""

    def __init__(self, capacity: int, timeout_seconds: int) -> None:
        self.capacity = capacity
        self.timeout_seconds = timeout_seconds
        self._queue: asyncio.Queue[QueueJob[Any] | None] = asyncio.Queue(maxsize=capacity)
        self._worker: asyncio.Task[None] | None = None
        self._active = False

    @property
    def depth(self) -> int:
        return self._queue.qsize()

    @property
    def active(self) -> bool:
        return self._active

    async def start(self) -> None:
        if self._worker is None or self._worker.done():
            self._worker = asyncio.create_task(self._work(), name="inference-worker")

    async def stop(self) -> None:
        if self._worker is None:
            return
        await self._queue.put(None)
        await self._worker
        self._worker = None

    async def submit(self, request_id: str, label: str, operation: Callable[[], T]) -> T:
        """Add a synchronous operation without blocking the event loop."""
        await self.start()
        loop = asyncio.get_running_loop()
        future: asyncio.Future[T] = loop.create_future()
        started: asyncio.Future[None] = loop.create_future()
        job = QueueJob(
            request_id=request_id,
            label=label,
            operation=operation,
            context=copy_context(),
            future=future,
            started=started,
        )
        try:
            self._queue.put_nowait(job)
        except asyncio.QueueFull as exc:
            raise QueueCapacityError() from exc
        try:
            await asyncio.wait_for(asyncio.shield(started), timeout=self.timeout_seconds)
        except asyncio.CancelledError:
            if started.done():
                future.add_done_callback(self._consume_late_result)
            else:
                job.cancelled = True
                future.cancel()
            raise
        except TimeoutError as exc:
            if started.done():
                return await asyncio.shield(future)
            job.cancelled = True
            future.cancel()
            raise GatewayError(
                "Inference timed out waiting in the queue",
                status_code=504,
                code="queue_timeout",
            ) from exc
        return await asyncio.shield(future)

    @staticmethod
    def _consume_late_result(future: asyncio.Future[Any]) -> None:
        if not future.cancelled():
            try:
                future.exception()
            except asyncio.CancelledError:
                pass

    async def _work(self) -> None:
        while True:
            job = await self._queue.get()
            if job is None:
                self._queue.task_done()
                break
            if job.cancelled:
                self._queue.task_done()
                continue
            self._active = True
            if not job.started.done():
                job.started.set_result(None)
            try:
                result = await asyncio.to_thread(job.context.run, job.operation)
            except Exception as exc:
                if not job.future.done():
                    job.future.set_exception(exc)
            else:
                if not job.future.done():
                    job.future.set_result(result)
            finally:
                self._active = False
                self._queue.task_done()