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
|
torch/transformers are imported lazily inside load() so this module can be
|
||||||
available (#872), else stays on CPU. fp32 is kept on GPU too so GPU-computed
|
imported in the web container (which never runs inference) without paying the
|
||||||
embeddings stay in the same numeric space as the existing CPU ones (cosine
|
torch import cost.
|
||||||
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.
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
@@ -16,6 +13,11 @@ from PIL import Image, ImageFile
|
|||||||
|
|
||||||
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
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(
|
MODEL_NAME = os.environ.get(
|
||||||
"SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
|
"SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
|
||||||
)
|
)
|
||||||
@@ -32,7 +34,6 @@ class Embedder:
|
|||||||
self._model = None
|
self._model = None
|
||||||
self._processor = None
|
self._processor = None
|
||||||
self._torch = None
|
self._torch = None
|
||||||
self._device = "cpu"
|
|
||||||
|
|
||||||
def load(self) -> None:
|
def load(self) -> None:
|
||||||
if self._model is not None:
|
if self._model is not None:
|
||||||
@@ -40,17 +41,10 @@ class Embedder:
|
|||||||
import torch
|
import torch
|
||||||
from transformers import AutoModel, SiglipImageProcessor
|
from transformers import AutoModel, SiglipImageProcessor
|
||||||
|
|
||||||
from .device import gpu_requested, torch_mem_fraction
|
|
||||||
|
|
||||||
self._torch = torch
|
self._torch = torch
|
||||||
# GPU (#872) when requested AND a CUDA device is present; else CPU. Cap
|
# Bound torch's CPU thread pool (see _INTRA_OP_THREADS) so each replica
|
||||||
# torch's share of the 8GB P4 (the ONNX tagger shares the card).
|
# stays a predictable core consumer on a shared node.
|
||||||
if gpu_requested() and torch.cuda.is_available():
|
torch.set_num_threads(_INTRA_OP_THREADS)
|
||||||
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
|
|
||||||
# FC's embedder only does IMAGE inference — never text. AutoProcessor
|
# FC's embedder only does IMAGE inference — never text. AutoProcessor
|
||||||
# loads the full processor including SiglipTokenizer, which requires
|
# loads the full processor including SiglipTokenizer, which requires
|
||||||
# the sentencepiece library at import time even if we never call it.
|
# 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 = AutoModel.from_pretrained(str(self._model_dir))
|
||||||
self._model.eval()
|
self._model.eval()
|
||||||
if self._device == "cuda":
|
|
||||||
self._model = self._model.to("cuda")
|
|
||||||
|
|
||||||
def infer(self, image_path: Path) -> np.ndarray:
|
def infer(self, image_path: Path) -> np.ndarray:
|
||||||
"""Return a 1152-dim float32 embedding (SigLIP MAP-pooled output)."""
|
"""Return a 1152-dim float32 embedding (SigLIP MAP-pooled output)."""
|
||||||
@@ -74,12 +66,9 @@ class Embedder:
|
|||||||
img = Image.open(image_path).convert("RGB")
|
img = Image.open(image_path).convert("RGB")
|
||||||
with self._torch.no_grad():
|
with self._torch.no_grad():
|
||||||
inputs = self._processor(images=img, return_tensors="pt")
|
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)
|
out = self._model.get_image_features(**inputs)
|
||||||
pooled = out.pooler_output if hasattr(out, "pooler_output") else out
|
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].numpy().astype(np.float32)
|
||||||
return pooled[0].detach().cpu().numpy().astype(np.float32)
|
|
||||||
|
|
||||||
|
|
||||||
_default_embedder: Embedder | None = None
|
_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
|
Single-image at a time. Loaded lazily inside the ml-worker process; NOT
|
||||||
provider when requested (FC_ML_DEVICE) and onnxruntime-gpu + a GPU are present
|
thread-safe — the ml queue worker runs --concurrency=1 per process (scale ML by
|
||||||
(#872), else falls back to CPU. Loaded lazily inside the ml-worker process; NOT
|
running multiple worker replicas, not threads).
|
||||||
thread-safe — the ml queue worker must run --concurrency=1 (set by the FC-1
|
|
||||||
entrypoint).
|
|
||||||
|
|
||||||
v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has
|
v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has
|
||||||
camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB)
|
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 json
|
||||||
import logging
|
|
||||||
import os
|
import os
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -22,7 +19,10 @@ from pathlib import Path
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
from PIL import Image, ImageFile
|
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
|
# 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
|
# 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).
|
# without onnxruntime (CI / lean web image).
|
||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
|
|
||||||
from .device import gpu_requested, onnx_gpu_mem_bytes
|
# 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
|
||||||
# GPU (#872) when requested AND the CUDA provider is actually present
|
# core and oversubscribe (and starve the co-located DB/web). Bounding it
|
||||||
# (onnxruntime-gpu in the ml image); otherwise CPU. gpu_mem_limit caps
|
# makes each replica a predictable core consumer — run N replicas where
|
||||||
# the CUDA arena so the tagger + the torch embedder co-exist on the 8GB
|
# N × _INTRA_OP_THREADS stays within the cores you allot to ML.
|
||||||
# P4. Falls back to CPU automatically on the CPU onnxruntime package.
|
opts = ort.SessionOptions()
|
||||||
providers: list = ["CPUExecutionProvider"]
|
opts.intra_op_num_threads = _INTRA_OP_THREADS
|
||||||
if gpu_requested() and "CUDAExecutionProvider" in ort.get_available_providers():
|
session = ort.InferenceSession(
|
||||||
providers = [
|
str(model_path), sess_options=opts, providers=["CPUExecutionProvider"],
|
||||||
("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())
|
|
||||||
self._input_name = session.get_inputs()[0].name
|
self._input_name = session.get_inputs()[0].name
|
||||||
# Assign sentinels last so a partial load isn't observable.
|
# Assign sentinels last so a partial load isn't observable.
|
||||||
self._tag_names = names
|
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