openmusic_analysis/runtime.py

from __future__ import annotations

import asyncio
from collections.abc import Callable
from typing import Any, TypeVar


T = TypeVar("T")


class DeviceManager:
    @staticmethod
    def select(configured: str = "auto") -> str:
        if configured != "auto":
            return configured
        try:
            import torch

            if torch.cuda.is_available():
                return "cuda"
            if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                return "mps"
        except ImportError:
            pass
        return "cpu"


class InferenceGate:
    """One process-wide bound on model loading and inference concurrency."""

    def __init__(self, concurrency: int) -> None:
        self._semaphore = asyncio.Semaphore(max(1, concurrency))

    async def run(self, function: Callable[..., T], *args: Any) -> T:
        task = asyncio.create_task(self._guarded(function, *args))
        try:
            return await asyncio.shield(task)
        except asyncio.CancelledError:
            task.add_done_callback(_consume_background_exception)
            raise

    async def _guarded(self, function: Callable[..., T], *args: Any) -> T:
        async with self._semaphore:
            return await asyncio.to_thread(function, *args)


def _consume_background_exception(task: asyncio.Task[Any]) -> None:
    if not task.cancelled():
        task.exception()