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
+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