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,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"))
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user