5e2094f2ad
FastAPI service in infra/media-worker implementing every MediaInferenceProvider operation behind HTTP API v1, plus capabilities. Bearer auth, MEDIA_ROOT realpath enforcement, strict request schemas, cancellation on disconnect, GPU/CPU gates with 503 overflow, input size/duration limits, lazy SHA256-pinned model downloads, fail-closed redaction with atomic output, offline pytest suite with fake backends, CPU/CUDA/ROCm Dockerfile and a license review of every model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj
62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
"""Accelerator detection. ML libraries are imported only when installed."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import importlib.util
|
|
from typing import Any
|
|
|
|
|
|
def _installed(name: str) -> bool:
|
|
try:
|
|
return importlib.util.find_spec(name) is not None
|
|
except (ImportError, ValueError):
|
|
return False
|
|
|
|
|
|
@functools.cache
|
|
def torch_gpu_available() -> bool:
|
|
if not _installed("torch"):
|
|
return False
|
|
import torch
|
|
|
|
return bool(torch.cuda.is_available())
|
|
|
|
|
|
def effective_accelerator(settings: Any) -> str:
|
|
"""What the image was built for, downgraded to cpu when no GPU is visible.
|
|
|
|
ROCm builds of torch expose the GPU through the torch.cuda API.
|
|
"""
|
|
if settings.accelerator in ("cuda", "rocm") and torch_gpu_available():
|
|
return settings.accelerator
|
|
return "cpu"
|
|
|
|
|
|
def torch_device(settings: Any) -> str:
|
|
return "cuda" if effective_accelerator(settings) != "cpu" else "cpu"
|
|
|
|
|
|
def ctranslate2_device(settings: Any) -> str:
|
|
# CTranslate2 ships CUDA builds only; ROCm hosts transcribe on the CPU.
|
|
if settings.accelerator != "cuda":
|
|
return "cpu"
|
|
import ctranslate2
|
|
|
|
return "cuda" if ctranslate2.get_cuda_device_count() > 0 else "cpu"
|
|
|
|
|
|
def onnx_providers(settings: Any) -> list[str]:
|
|
import onnxruntime
|
|
|
|
available = set(onnxruntime.get_available_providers())
|
|
providers: list[str] = []
|
|
if settings.accelerator == "cuda" and "CUDAExecutionProvider" in available:
|
|
providers.append("CUDAExecutionProvider")
|
|
if settings.accelerator == "rocm":
|
|
for name in ("MIGraphXExecutionProvider", "ROCMExecutionProvider"):
|
|
if name in available:
|
|
providers.append(name)
|
|
providers.append("CPUExecutionProvider")
|
|
return providers
|