refactor(ml): drop GPU code, cap inference threads by default (#747/#872)
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user