diff --git a/backend/app/services/ml/device.py b/backend/app/services/ml/device.py deleted file mode 100644 index aa6c252..0000000 --- a/backend/app/services/ml/device.py +++ /dev/null @@ -1,31 +0,0 @@ -"""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 3ab7a4f..d94f71e 100644 --- a/backend/app/services/ml/embedder.py +++ b/backend/app/services/ml/embedder.py @@ -1,11 +1,8 @@ -"""SigLIP SO400M image-embedding wrapper (PyTorch). +"""SigLIP SO400M image-embedding wrapper (PyTorch CPU). -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. +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 @@ -16,6 +13,11 @@ from PIL import Image, ImageFile ImageFile.LOAD_TRUNCATED_IMAGES = True +# Cap torch's intra-op threads so each ml-worker replica is a bounded core +# consumer on a shared node (torch otherwise uses all cores). Keep +# N_replicas × this within the cores allotted to ML to avoid oversubscription. +_INTRA_OP_THREADS = 4 + MODEL_NAME = os.environ.get( "SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384" ) @@ -32,7 +34,6 @@ class Embedder: self._model = None self._processor = None self._torch = None - self._device = "cpu" def load(self) -> None: if self._model is not None: @@ -40,17 +41,10 @@ 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 + # Bound torch's CPU thread pool (see _INTRA_OP_THREADS) so each replica + # stays a predictable core consumer on a shared node. + torch.set_num_threads(_INTRA_OP_THREADS) # 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. @@ -65,8 +59,6 @@ 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).""" @@ -74,12 +66,9 @@ 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 - # .detach().cpu() so a CUDA tensor converts to numpy (no-op on CPU). - return pooled[0].detach().cpu().numpy().astype(np.float32) + return pooled[0].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 178387d..0f15c4a 100644 --- a/backend/app/services/ml/tagger.py +++ b/backend/app/services/ml/tagger.py @@ -1,10 +1,8 @@ -"""Camie-tagger-v2 ONNX wrapper. +"""Camie-tagger-v2 ONNX wrapper (CPU). -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). +Single-image at a time. Loaded lazily inside the ml-worker process; NOT +thread-safe — the ml queue worker runs --concurrency=1 per process (scale ML by +running multiple worker replicas, not threads). v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB) @@ -14,7 +12,6 @@ ImageNet normalize, NCHW layout, sigmoid on refined logits (output[1]). """ import json -import logging import os from dataclasses import dataclass from pathlib import Path @@ -22,7 +19,10 @@ from pathlib import Path import numpy as np from PIL import Image, ImageFile -log = logging.getLogger(__name__) +# Cap inference threads (see Tagger.load) so each ml-worker replica is a bounded +# core consumer on a shared node — keep N_replicas × this within the cores +# allotted to ML so replicas don't oversubscribe the box / starve the DB. +_INTRA_OP_THREADS = 4 # 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 @@ -122,20 +122,16 @@ class Tagger: # without onnxruntime (CI / lean web image). import onnxruntime as ort - 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()) + # Cap the intra-op thread pool. ONNX Runtime otherwise sizes it to ALL + # host cores, so on a shared node each ml-worker replica would grab every + # core and oversubscribe (and starve the co-located DB/web). Bounding it + # makes each replica a predictable core consumer — run N replicas where + # N × _INTRA_OP_THREADS stays within the cores you allot to ML. + opts = ort.SessionOptions() + opts.intra_op_num_threads = _INTRA_OP_THREADS + session = ort.InferenceSession( + str(model_path), sess_options=opts, providers=["CPUExecutionProvider"], + ) 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 deleted file mode 100644 index 2aed435..0000000 --- a/tests/test_ml_device.py +++ /dev/null @@ -1,31 +0,0 @@ -"""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