Video tag quality: cadence sampling + min-frame aggregation + ML thread cap (#747) #111
@@ -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