Spaces:
Running
Running
File size: 1,415 Bytes
330f477 | 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 | 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()
|