feat(ml): GPU-capable tagger + embedder with CPU fallback (#872)

Step 1 of GPU enablement (code only — CPU-safe, CI-green; the CUDA image is a
separate step pending the host driver version).

- New services/ml/device.py: FC_ML_DEVICE (auto|cuda|cpu) intent + VRAM knobs
  (FC_ML_ONNX_GPU_MEM_GB, FC_ML_TORCH_MEM_FRACTION). Per-worker-host bootstrap →
  env, not a DB setting (the GPU host runs CUDA, others CPU).
- tagger: use CUDAExecutionProvider (with gpu_mem_limit) when requested AND the
  provider is actually present (onnxruntime-gpu), else CPUExecutionProvider. Logs
  the active providers.
- embedder: move model + inputs to cuda when requested AND torch.cuda is
  available; cap torch's VRAM share; .detach().cpu() before numpy. fp32 kept so
  GPU embeddings stay in the same space as existing CPU ones.

Both AND the env intent with the framework's real availability, so on CPU
(CI / CPU onnxruntime / no GPU) they fall back cleanly — behavior unchanged.
The 8GB P4 is shared by both frameworks, hence the conservative default caps.

Tests: device env parsing. (tagger/embedder GPU paths are operator-verified on
the GPU host — models aren't in CI.)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

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"))

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

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

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