tab_queue.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""
Ingested from inspection-backend/tab_queue.py.
FastAPI HTTPException replaced with RuntimeError — CLI has no HTTP layer.
"""
from __future__ import annotations
 
import asyncio
import os
 
 
class TabQueue:
    def __init__(self):
        self._inference_slots = self._calculate_inference_slots()
        self.inference = asyncio.Semaphore(self._inference_slots)
        self.training  = asyncio.Semaphore(int(os.getenv("TRAINING_SLOTS", "1")))
        self._inference_waiters = 0
        self._training_waiters  = 0
        self._training_active = False
 
    def _calculate_inference_slots(self) -> int:
        override = os.getenv("INFERENCE_SLOTS")
        if override:
            return int(override)
        return 1
 
    async def inference_request(self, fn, timeout: float = 30.0):
        if self._training_active:
            raise RuntimeError("Model is busy with a training run — try again when training completes")
        self._inference_waiters += 1
        try:
            await asyncio.wait_for(self.inference.acquire(), timeout=timeout)
        except asyncio.TimeoutError:
            raise RuntimeError("GPU busy, try again shortly")
        finally:
            self._inference_waiters -= 1
        try:
            return await fn()
        finally:
            self.inference.release()
 
    async def training_request(self, fn):
        self._training_waiters += 1
        await self.training.acquire()
        self._training_waiters -= 1
 
        await self.inference.acquire()
        self._training_active = True
        try:
            return await fn()
        finally:
            self._training_active = False
            self.inference.release()
            self.training.release()
 
    def status(self) -> dict:
        inference_active = self._inference_slots - self.inference._value
        return {
            "inference_slots":   self._inference_slots,
            "inference_active":  inference_active,
            "inference_queued":  self._inference_waiters,
            "training_slots":    int(os.getenv("TRAINING_SLOTS", "1")),
            "training_active":   int(self._training_active),
            "training_queued":   self._training_waiters,
        }
 
 
tab_queue: TabQueue | None = None
 
 
def init_tab_queue() -> TabQueue:
    global tab_queue
    tab_queue = TabQueue()
    return tab_queue