From bbb8e4850cd8fb2fdf75d565827dda29025b4302 Mon Sep 17 00:00:00 2001 From: dtoro Date: Sun, 12 Apr 2026 22:12:21 +0200 Subject: [PATCH] feat: GPU acceleration support for ONNX Runtime inference Centralize execution provider selection in providers.py with auto-detection and graceful fallback. All ONNX sessions (embedder, detector, face processor, recognizer) now use the configured providers. - New VISION_EXECUTION_PROVIDERS env var: "auto" for GPU auto-detect, or explicit "CUDAExecutionProvider,CPUExecutionProvider" - Provider priority: CUDA > ROCm > OpenVINO > CPU (when set to "auto") - docker-compose.yml includes commented-out NVIDIA GPU deploy section - Supports onnxruntime-gpu as a drop-in replacement for onnxruntime Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/app/config.py | 20 ++++- backend/app/services/vision/detect.py | 6 +- backend/app/services/vision/embed.py | 20 ++--- backend/app/services/vision/faces.py | 6 +- .../services/vision/insightface_processor.py | 5 +- backend/app/services/vision/providers.py | 86 +++++++++++++++++++ docker-compose.yml | 15 ++++ 7 files changed, 138 insertions(+), 20 deletions(-) create mode 100644 backend/app/services/vision/providers.py diff --git a/backend/app/config.py b/backend/app/config.py index 7665a26..0438e51 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -66,6 +66,11 @@ class VisionSettings(BaseModel): enabled: bool = True backend: str = "onnx" # "onnx" | "rocm" (future) models_dir: str = "/data/models" + # ONNX Runtime execution providers in priority order. + # Auto-detected at startup; falls back to CPU if GPU is unavailable. + # Options: "CUDAExecutionProvider", "ROCMExecutionProvider", + # "OpenVINOExecutionProvider", "CPUExecutionProvider" + execution_providers: list[str] = ["CPUExecutionProvider"] embedder: EmbedderSettings = EmbedderSettings() ocr: OCRSettings = OCRSettings() detector: DetectorSettings = DetectorSettings() @@ -183,9 +188,22 @@ class Settings(BaseSettings): def performance(self) -> PerformanceSettings: return self.config.performance + # ONNX Runtime execution providers, overridable via env var. + # Comma-separated: "CUDAExecutionProvider,CPUExecutionProvider" + # or "auto" for GPU auto-detection. + vision_execution_providers: str = Field( + default="CPUExecutionProvider", + env="VISION_EXECUTION_PROVIDERS", + ) + @property def vision(self) -> VisionSettings: - return self.config.vision + v = self.config.vision + # Override execution_providers from env if set. + providers = [p.strip() for p in self.vision_execution_providers.split(",") if p.strip()] + if providers: + v.execution_providers = providers + return v class Config: env_file = ".env" diff --git a/backend/app/services/vision/detect.py b/backend/app/services/vision/detect.py index 9884570..f0d0ebb 100644 --- a/backend/app/services/vision/detect.py +++ b/backend/app/services/vision/detect.py @@ -120,12 +120,10 @@ class YOLOv8Detector(ObjectDetector): def __init__(self, settings: VisionSettings): model_path = Path(settings.models_dir) / "detect" / "yolov8n.onnx" - opts = ort.SessionOptions() - opts.inter_op_num_threads = 2 - opts.intra_op_num_threads = 2 + from app.services.vision.providers import create_session logger.info("Loading YOLOv8n from %s", model_path) - self._session = ort.InferenceSession(str(model_path), opts, providers=["CPUExecutionProvider"]) + self._session = create_session(str(model_path), configured_providers=settings.execution_providers) self._conf_threshold = settings.detector.min_confidence self._max_detections = settings.detector.max_detections diff --git a/backend/app/services/vision/embed.py b/backend/app/services/vision/embed.py index 8681bd4..ae3cd03 100644 --- a/backend/app/services/vision/embed.py +++ b/backend/app/services/vision/embed.py @@ -67,15 +67,15 @@ class OpenCLIPEmbedder(Embedder): visual_path = model_dir / "visual.onnx" textual_path = model_dir / "textual.onnx" - opts = ort.SessionOptions() - opts.inter_op_num_threads = 2 - opts.intra_op_num_threads = 2 + from app.services.vision.providers import create_session + from app.config import settings as app_settings + providers = app_settings.vision.execution_providers logger.info("Loading OpenCLIP visual encoder from %s", visual_path) - self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"]) + self._visual = create_session(str(visual_path), configured_providers=providers) logger.info("Loading OpenCLIP textual encoder from %s", textual_path) - self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"]) + self._textual = create_session(str(textual_path), configured_providers=providers) def embed_image(self, image: np.ndarray) -> np.ndarray: inp = _preprocess_image(image, _OPENCLIP_SIZE, _OPENCLIP_MEAN, _OPENCLIP_STD) @@ -110,15 +110,15 @@ class SigLIP2Embedder(Embedder): visual_path = model_dir / "visual.onnx" textual_path = model_dir / "textual.onnx" - opts = ort.SessionOptions() - opts.inter_op_num_threads = 2 - opts.intra_op_num_threads = 2 + from app.services.vision.providers import create_session + from app.config import settings as app_settings + providers = app_settings.vision.execution_providers logger.info("Loading SigLIP2 visual encoder from %s", visual_path) - self._visual = ort.InferenceSession(str(visual_path), opts, providers=["CPUExecutionProvider"]) + self._visual = create_session(str(visual_path), configured_providers=providers) logger.info("Loading SigLIP2 textual encoder from %s", textual_path) - self._textual = ort.InferenceSession(str(textual_path), opts, providers=["CPUExecutionProvider"]) + self._textual = create_session(str(textual_path), configured_providers=providers) def embed_image(self, image: np.ndarray) -> np.ndarray: inp = _preprocess_image(image, _SIGLIP2_SIZE, _SIGLIP2_MEAN, _SIGLIP2_STD) diff --git a/backend/app/services/vision/faces.py b/backend/app/services/vision/faces.py index d79309c..4c3e860 100644 --- a/backend/app/services/vision/faces.py +++ b/backend/app/services/vision/faces.py @@ -71,11 +71,9 @@ class YuNetSFaceProcessor(FaceProcessor): logger.info("YuNet face detector loaded via OpenCV") # SFace via ONNX Runtime - opts = ort.SessionOptions() - opts.inter_op_num_threads = 2 - opts.intra_op_num_threads = 2 + from app.services.vision.providers import create_session ort.set_default_logger_severity(3) - self._recognizer = ort.InferenceSession(sface_path, opts, providers=["CPUExecutionProvider"]) + self._recognizer = create_session(sface_path, configured_providers=settings.execution_providers) logger.info("SFace recognizer loaded via ONNX Runtime") self._min_face_size = settings.faces.min_face_size diff --git a/backend/app/services/vision/insightface_processor.py b/backend/app/services/vision/insightface_processor.py index aea8bb0..14683c0 100644 --- a/backend/app/services/vision/insightface_processor.py +++ b/backend/app/services/vision/insightface_processor.py @@ -23,10 +23,13 @@ class InsightFaceProcessor(FaceProcessor): model_root = str(Path(settings.models_dir) / "face" / "insightface") logger.info("Loading InsightFace buffalo_l from %s", model_root) + from app.services.vision.providers import get_providers + providers = get_providers(settings.execution_providers) + self._app = FaceAnalysis( name="buffalo_l", root=model_root, - providers=["CPUExecutionProvider"], + providers=providers, ) self._app.prepare(ctx_id=-1, det_size=(640, 640)) self._min_det_score = settings.faces.recognition_threshold diff --git a/backend/app/services/vision/providers.py b/backend/app/services/vision/providers.py new file mode 100644 index 0000000..0a225d9 --- /dev/null +++ b/backend/app/services/vision/providers.py @@ -0,0 +1,86 @@ +""" +ONNX Runtime execution provider resolution with GPU auto-detection. + +Resolves configured execution providers against what's actually available +in the current ONNX Runtime build. Falls back to CPU if no GPU provider +is available. Logs the selected provider so users can confirm GPU is active. +""" +import logging + +import onnxruntime as ort + +logger = logging.getLogger(__name__) + +_resolved: list[str] | None = None + + +def get_providers(configured: list[str] | None = None) -> list[str]: + """Return the best available execution providers. + + 1. If `configured` is provided, filter to only those that are + actually available in the current ORT build. + 2. If none of the configured providers are available, fall back + to CPUExecutionProvider. + 3. Auto-detect: if configured is ["auto"], probe for GPU providers. + + Results are cached after first call. + """ + global _resolved + if _resolved is not None: + return _resolved + + available = set(ort.get_available_providers()) + logger.info("ONNX Runtime available providers: %s", sorted(available)) + + if configured is None or configured == ["CPUExecutionProvider"]: + _resolved = ["CPUExecutionProvider"] + return _resolved + + if configured == ["auto"]: + # Auto-detect: prefer CUDA > ROCm > OpenVINO > CPU + priority = [ + "CUDAExecutionProvider", + "ROCMExecutionProvider", + "OpenVINOExecutionProvider", + ] + for p in priority: + if p in available: + _resolved = [p, "CPUExecutionProvider"] + logger.info("Auto-detected GPU provider: %s", p) + return _resolved + _resolved = ["CPUExecutionProvider"] + logger.info("No GPU provider detected, using CPU") + return _resolved + + # Filter configured list to available providers. + resolved = [p for p in configured if p in available] + if not resolved: + logger.warning( + "None of the configured providers %s are available. " + "Falling back to CPU. Available: %s", + configured, + sorted(available), + ) + resolved = ["CPUExecutionProvider"] + else: + # Always include CPU as fallback. + if "CPUExecutionProvider" not in resolved: + resolved.append("CPUExecutionProvider") + + _resolved = resolved + logger.info("Using ONNX Runtime providers: %s", _resolved) + return _resolved + + +def create_session( + model_path: str, + opts: ort.SessionOptions | None = None, + configured_providers: list[str] | None = None, +) -> ort.InferenceSession: + """Create an ONNX InferenceSession with the best available providers.""" + providers = get_providers(configured_providers) + if opts is None: + opts = ort.SessionOptions() + opts.inter_op_num_threads = 2 + opts.intra_op_num_threads = 2 + return ort.InferenceSession(model_path, opts, providers=providers) diff --git a/docker-compose.yml b/docker-compose.yml index af347a8..04c5036 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -143,6 +143,13 @@ services: - LOG_LEVEL=${LOG_LEVEL:-INFO} - TZ=${TZ:-UTC} - MULITA_CELERY_WORKER=1 + # ONNX Runtime execution providers. Set to "auto" to auto-detect + # GPU (CUDA > ROCm > OpenVINO > CPU), or explicitly: + # "CUDAExecutionProvider,CPUExecutionProvider" + # "ROCMExecutionProvider,CPUExecutionProvider" + # Default: CPU only. To enable GPU, also uncomment the deploy + # section below and install nvidia-container-toolkit on the host. + - VISION_EXECUTION_PROVIDERS=${VISION_EXECUTION_PROVIDERS:-CPUExecutionProvider} # Pin each ONNX session to one intra-op thread so N prefork children # × default-all-cores doesn't oversubscribe the box. With # concurrency=5 and OMP=1, vision peaks at 5 busy cores, leaving @@ -151,6 +158,14 @@ services: - OMP_NUM_THREADS=1 - OPENBLAS_NUM_THREADS=1 - MKL_NUM_THREADS=1 + # Uncomment for NVIDIA GPU passthrough: + # deploy: + # resources: + # reservations: + # devices: + # - driver: nvidia + # count: all + # capabilities: [gpu] depends_on: redis: condition: service_started