refactor(ml): drop GPU code, cap inference threads by default (#747/#872)
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 18s
CI / backend-lint-and-test (push) Successful in 27s
CI / integration (push) Successful in 3m18s

GPU enablement (#872) cancelled — not worth the Pascal-specific build for a
modest CPU→GPU win on an old P4. Remove the dead GPU code (device.py, the CUDA
provider branch in tagger, the .to('cuda') path in embedder) so nothing carries
it forward.

Instead, bound CPU inference threads by default so the ml-worker is a predictable
core consumer on a SHARED node — the intended scaling model is multiple worker
replicas (each --concurrency=1, each its own cgroup limit), not one big
container. ONNX Runtime and torch otherwise size their thread pools to ALL host
cores, so each replica would grab every core and oversubscribe / starve the
co-located DB+web. Cap both to _INTRA_OP_THREADS=4 (matches the prior per-worker
cpus:4 unit): run N replicas where N×4 stays within the cores allotted to ML.

- tagger: ort.SessionOptions().intra_op_num_threads = 4 (CPUExecutionProvider).
- embedder: torch.set_num_threads(4).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-16 13:39:55 -04:00
parent db7e1f2b59
commit 60a9c9e6ef
4 changed files with 31 additions and 108 deletions
-31
View File
@@ -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"))
+13 -24
View File
@@ -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
+18 -22
View File
@@ -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
-31
View File
@@ -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