diff --git a/backend/app/services/ml/device.py b/backend/app/services/ml/device.py new file mode 100644 index 0000000..aa6c252 --- /dev/null +++ b/backend/app/services/ml/device.py @@ -0,0 +1,31 @@ +"""ML device selection (#872 — GPU enablement for the ml-worker). + +The ml-worker is GPU-capable but must run unchanged on CPU (CI, non-GPU hosts). +Selection is a per-worker-HOST bootstrap concern (the GPU host runs CUDA, others +CPU), so it's an env var, not a DB setting — different workers need different +values. Each framework still ANDs this intent with its OWN runtime availability +(onnxruntime providers / torch.cuda), so "want GPU but none present" falls back +to CPU cleanly. + +Env: + FC_ML_DEVICE auto (default) | cuda | gpu -> try GPU; cpu -> force CPU + FC_ML_ONNX_GPU_MEM_GB ONNX CUDA arena cap, GB (default 3) — the P4 is 8GB + total and torch shares it, so keep headroom. + FC_ML_TORCH_MEM_FRACTION fraction of total VRAM torch may use (default 0.6). +""" + +import os + + +def gpu_requested() -> bool: + return os.environ.get("FC_ML_DEVICE", "auto").strip().lower() in ( + "auto", "cuda", "gpu", + ) + + +def onnx_gpu_mem_bytes() -> int: + return int(float(os.environ.get("FC_ML_ONNX_GPU_MEM_GB", "3")) * 1024 ** 3) + + +def torch_mem_fraction() -> float: + return float(os.environ.get("FC_ML_TORCH_MEM_FRACTION", "0.6")) diff --git a/backend/app/services/ml/embedder.py b/backend/app/services/ml/embedder.py index 5da36c5..3ab7a4f 100644 --- a/backend/app/services/ml/embedder.py +++ b/backend/app/services/ml/embedder.py @@ -1,8 +1,11 @@ -"""SigLIP SO400M image-embedding wrapper (PyTorch CPU). +"""SigLIP SO400M image-embedding wrapper (PyTorch). -Direct port of ImageRepo's siglip.py. torch/transformers are imported -lazily inside load() so this module can be imported in the web container -(which never runs inference) without paying the torch import cost. +Runs on CPU by default; moves to CUDA when requested (FC_ML_DEVICE) and a GPU is +available (#872), else stays on CPU. fp32 is kept on GPU too so GPU-computed +embeddings stay in the same numeric space as the existing CPU ones (cosine +comparisons). torch/transformers are imported lazily inside load() so this +module can be imported in the web container (which never runs inference) without +paying the torch import cost. """ import os @@ -29,6 +32,7 @@ class Embedder: self._model = None self._processor = None self._torch = None + self._device = "cpu" def load(self) -> None: if self._model is not None: @@ -36,7 +40,17 @@ class Embedder: import torch from transformers import AutoModel, SiglipImageProcessor + from .device import gpu_requested, torch_mem_fraction + self._torch = torch + # GPU (#872) when requested AND a CUDA device is present; else CPU. Cap + # torch's share of the 8GB P4 (the ONNX tagger shares the card). + if gpu_requested() and torch.cuda.is_available(): + self._device = "cuda" + try: + torch.cuda.set_per_process_memory_fraction(torch_mem_fraction()) + except Exception: # noqa: BLE001 — best-effort cap; never block load + pass # FC's embedder only does IMAGE inference — never text. AutoProcessor # loads the full processor including SiglipTokenizer, which requires # the sentencepiece library at import time even if we never call it. @@ -51,6 +65,8 @@ class Embedder: ) self._model = AutoModel.from_pretrained(str(self._model_dir)) self._model.eval() + if self._device == "cuda": + self._model = self._model.to("cuda") def infer(self, image_path: Path) -> np.ndarray: """Return a 1152-dim float32 embedding (SigLIP MAP-pooled output).""" @@ -58,9 +74,12 @@ class Embedder: img = Image.open(image_path).convert("RGB") with self._torch.no_grad(): inputs = self._processor(images=img, return_tensors="pt") + if self._device == "cuda": + inputs = {k: v.to("cuda") for k, v in inputs.items()} out = self._model.get_image_features(**inputs) pooled = out.pooler_output if hasattr(out, "pooler_output") else out - return pooled[0].numpy().astype(np.float32) + # .detach().cpu() so a CUDA tensor converts to numpy (no-op on CPU). + return pooled[0].detach().cpu().numpy().astype(np.float32) _default_embedder: Embedder | None = None diff --git a/backend/app/services/ml/tagger.py b/backend/app/services/ml/tagger.py index 6348595..178387d 100644 --- a/backend/app/services/ml/tagger.py +++ b/backend/app/services/ml/tagger.py @@ -1,8 +1,10 @@ """Camie-tagger-v2 ONNX wrapper. -CPU-only, single-image at a time. Loaded lazily inside the ml-worker -process; NOT thread-safe — the ml queue worker must run --concurrency=1 -(set by the FC-1 entrypoint). +Single-image at a time. Runs on CPU by default; uses the CUDA execution +provider when requested (FC_ML_DEVICE) and onnxruntime-gpu + a GPU are present +(#872), else falls back to CPU. Loaded lazily inside the ml-worker process; NOT +thread-safe — the ml queue worker must run --concurrency=1 (set by the FC-1 +entrypoint). v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB) @@ -12,6 +14,7 @@ ImageNet normalize, NCHW layout, sigmoid on refined logits (output[1]). """ import json +import logging import os from dataclasses import dataclass from pathlib import Path @@ -19,6 +22,8 @@ from pathlib import Path import numpy as np from PIL import Image, ImageFile +log = logging.getLogger(__name__) + # onnxruntime lives in requirements-ml.txt only — it is NOT installed in the # lean web image or in CI. Imported lazily inside Tagger.load() so this module # imports fine without it (the suggestion service imports SURFACED_CATEGORIES @@ -117,9 +122,20 @@ class Tagger: # without onnxruntime (CI / lean web image). import onnxruntime as ort - session = ort.InferenceSession( - str(model_path), providers=["CPUExecutionProvider"] - ) + from .device import gpu_requested, onnx_gpu_mem_bytes + + # GPU (#872) when requested AND the CUDA provider is actually present + # (onnxruntime-gpu in the ml image); otherwise CPU. gpu_mem_limit caps + # the CUDA arena so the tagger + the torch embedder co-exist on the 8GB + # P4. Falls back to CPU automatically on the CPU onnxruntime package. + providers: list = ["CPUExecutionProvider"] + if gpu_requested() and "CUDAExecutionProvider" in ort.get_available_providers(): + providers = [ + ("CUDAExecutionProvider", {"gpu_mem_limit": onnx_gpu_mem_bytes()}), + "CPUExecutionProvider", + ] + session = ort.InferenceSession(str(model_path), providers=providers) + log.info("tagger ONNX providers: %s", session.get_providers()) self._input_name = session.get_inputs()[0].name # Assign sentinels last so a partial load isn't observable. self._tag_names = names diff --git a/tests/test_ml_device.py b/tests/test_ml_device.py new file mode 100644 index 0000000..2aed435 --- /dev/null +++ b/tests/test_ml_device.py @@ -0,0 +1,31 @@ +"""ML device-selection env parsing (#872). Pure logic — no models/GPU/DB.""" + +from backend.app.services.ml import device + + +def test_gpu_requested_default_is_auto(monkeypatch): + monkeypatch.delenv("FC_ML_DEVICE", raising=False) + assert device.gpu_requested() is True + + +def test_gpu_requested_modes(monkeypatch): + for v in ("auto", "cuda", "gpu", "CUDA", " Auto "): + monkeypatch.setenv("FC_ML_DEVICE", v) + assert device.gpu_requested() is True + for v in ("cpu", "CPU", "none", "0"): + monkeypatch.setenv("FC_ML_DEVICE", v) + assert device.gpu_requested() is False + + +def test_onnx_gpu_mem_bytes(monkeypatch): + monkeypatch.delenv("FC_ML_ONNX_GPU_MEM_GB", raising=False) + assert device.onnx_gpu_mem_bytes() == 3 * 1024 ** 3 + monkeypatch.setenv("FC_ML_ONNX_GPU_MEM_GB", "2") + assert device.onnx_gpu_mem_bytes() == 2 * 1024 ** 3 + + +def test_torch_mem_fraction(monkeypatch): + monkeypatch.delenv("FC_ML_TORCH_MEM_FRACTION", raising=False) + assert device.torch_mem_fraction() == 0.6 + monkeypatch.setenv("FC_ML_TORCH_MEM_FRACTION", "0.5") + assert device.torch_mem_fraction() == 0.5