From 369e3de68401ef4f55d785acfb9fc52b1d9cd445 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Tue, 16 Jun 2026 11:07:00 -0400 Subject: [PATCH 1/3] feat(ml): cadence-based video frame sampling + min-frame tag aggregation (#747) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Video tag noise root cause: frames were a FIXED count (6) max-pooled — a tag firing on one frame survived at peak confidence, and a fixed count under-samples long multi-scene videos so real scene-local tags looked like noise. Redesign (operator-steered): - Sample at a fixed CADENCE — one frame every `video_frame_interval_seconds` (default 4) across the 5–95% window — so a tag's frame-presence reflects real screen time independent of video length. Capped at `video_max_frames` (default 64): a long video stretches the spacing instead of exploding into hundreds of inferences, bounding per-video cost on the single ml-worker (per-frame ffmpeg timeout also cut 60s→30s). - Aggregate with `_aggregate_video_predictions`: keep a tag only if it appears in >= `video_min_tag_frames` sampled frames (≈ that many × interval seconds on screen — duration-independent noise rejection), with confidence = MEAN over the frames it appears in (not max). Clamps the threshold to the sample count so a 1–2-frame short video still tags. - All three knobs are DB-backed ml_settings (migration 0053), patchable via /api/ml/settings + sliders in the ML settings card — replaces the VIDEO_ML_FRAMES env var (product-not-project). Tests: aggregation drops one-frame noise + means corroborated tags + clamps on short videos; settings round-trip + min>max validation. Replaced the _maxpool_predictions unit test. NOTE: this is the QUALITY half of #747. The perf half — the ml-worker runs CPU-only — is GPU enablement, tracked separately in #872. Co-Authored-By: Claude Opus 4.8 --- .../0053_ml_settings_video_tagging.py | 49 +++++++++++ backend/app/api/ml_admin.py | 15 ++++ backend/app/models/ml_settings.py | 15 ++++ backend/app/tasks/ml.py | 88 ++++++++++++++----- .../settings/MLThresholdSliders.vue | 38 ++++++++ tests/test_api_ml_admin.py | 30 +++++++ tests/test_tasks_ml.py | 41 ++++++--- 7 files changed, 241 insertions(+), 35 deletions(-) create mode 100644 alembic/versions/0053_ml_settings_video_tagging.py diff --git a/alembic/versions/0053_ml_settings_video_tagging.py b/alembic/versions/0053_ml_settings_video_tagging.py new file mode 100644 index 0000000..1f192a4 --- /dev/null +++ b/alembic/versions/0053_ml_settings_video_tagging.py @@ -0,0 +1,49 @@ +"""ml_settings: video tagging knobs (cadence sampling + noise floor) + +#747. Video tag quality/perf: sample frames at a fixed cadence (interval) so a +tag's frame-presence reflects real screen time, cap total frames so long videos +stay bounded, and keep a tag only if it appears in >= min_tag_frames sampled +frames. Operator-tunable via Settings → ML (replaces the VIDEO_ML_FRAMES env var). + +Revision ID: 0053 +Revises: 0052 +Create Date: 2026-06-16 +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +revision: str = "0053" +down_revision: Union[str, None] = "0052" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column( + "ml_settings", + sa.Column( + "video_frame_interval_seconds", sa.Float(), nullable=False, + server_default="4.0", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "video_max_frames", sa.Integer(), nullable=False, server_default="64", + ), + ) + op.add_column( + "ml_settings", + sa.Column( + "video_min_tag_frames", sa.Integer(), nullable=False, + server_default="3", + ), + ) + + +def downgrade() -> None: + op.drop_column("ml_settings", "video_min_tag_frames") + op.drop_column("ml_settings", "video_max_frames") + op.drop_column("ml_settings", "video_frame_interval_seconds") diff --git a/backend/app/api/ml_admin.py b/backend/app/api/ml_admin.py index ef3d2fa..b3e8d94 100644 --- a/backend/app/api/ml_admin.py +++ b/backend/app/api/ml_admin.py @@ -14,6 +14,9 @@ _EDITABLE = ( "centroid_similarity_threshold", "min_reference_images", "tagger_store_floor", + "video_frame_interval_seconds", + "video_max_frames", + "video_min_tag_frames", ) @@ -32,6 +35,9 @@ async def get_settings(): "centroid_similarity_threshold": s.centroid_similarity_threshold, "min_reference_images": s.min_reference_images, "tagger_store_floor": s.tagger_store_floor, + "video_frame_interval_seconds": s.video_frame_interval_seconds, + "video_max_frames": s.video_max_frames, + "video_min_tag_frames": s.video_min_tag_frames, "tagger_model_version": s.tagger_model_version, "embedder_model_version": s.embedder_model_version, } @@ -85,6 +91,15 @@ def _validate(p: dict) -> str | None: f"suggestion_threshold_{cat} cannot be below tagger_store_floor " f"({floor}) — predictions below the floor are not stored" ) + # Video tagging (#747). + if p["video_frame_interval_seconds"] <= 0: + return "video_frame_interval_seconds must be > 0" + if p["video_max_frames"] < 1: + return "video_max_frames must be >= 1" + if p["video_min_tag_frames"] < 1: + return "video_min_tag_frames must be >= 1" + if p["video_min_tag_frames"] > p["video_max_frames"]: + return "video_min_tag_frames cannot exceed video_max_frames" return None diff --git a/backend/app/models/ml_settings.py b/backend/app/models/ml_settings.py index e28686d..3ae5a1b 100644 --- a/backend/app/models/ml_settings.py +++ b/backend/app/models/ml_settings.py @@ -40,6 +40,21 @@ class MLSettings(Base): min_reference_images: Mapped[int] = mapped_column( Integer, nullable=False, default=5 ) + # Video tagging (#747). Sample one frame every N seconds (fixed CADENCE, not a + # fixed count) so a tag's frame-presence reflects real screen time regardless + # of video length; cap the total so a long video can't explode into hundreds + # of inferences (the cadence stretches past the cap). A tag is kept only if it + # appears in >= video_min_tag_frames sampled frames (≈ that many × interval + # seconds on screen) — duration-independent noise rejection. Operator-tunable. + video_frame_interval_seconds: Mapped[float] = mapped_column( + Float, nullable=False, default=4.0 + ) + video_max_frames: Mapped[int] = mapped_column( + Integer, nullable=False, default=64 + ) + video_min_tag_frames: Mapped[int] = mapped_column( + Integer, nullable=False, default=3 + ) tagger_model_version: Mapped[str] = mapped_column( String(128), nullable=False, default="camie-tagger-v2" ) diff --git a/backend/app/tasks/ml.py b/backend/app/tasks/ml.py index 05d1581..7017e9c 100644 --- a/backend/app/tasks/ml.py +++ b/backend/app/tasks/ml.py @@ -49,11 +49,12 @@ def tag_and_embed(self, image_id: int) -> dict: """Run Camie + SigLIP on one image; store predictions + embedding; then enqueue per-image allowlist application. - Video: sample frames between 10% and 90% of duration (VIDEO_ML_FRAMES, - default 6). Max-pool tagger confidences across frames, mean-pool the + Video (#747): sample frames at a fixed cadence (ml_settings + video_frame_interval_seconds, capped at video_max_frames), keep a tag only if + it appears in >= video_min_tag_frames frames and average its confidence over + those frames (mean-pool, not max — kills one-frame noise); mean-pool the SigLIP embeddings. On no-frames returns status='no_frames' (not an error). """ - import os import time from ..services.ml.embedder import get_embedder @@ -116,7 +117,9 @@ def tag_and_embed(self, image_id: int) -> dict: phase = "video_sample_frames" t0 = time.monotonic() frames = _sample_video_frames( - src, int(os.environ.get("VIDEO_ML_FRAMES", "6")) + src, + interval=settings.video_frame_interval_seconds, + max_frames=settings.video_max_frames, ) log.info( "tag_and_embed sampled %d frame(s) in %.1fs: %s", @@ -127,13 +130,19 @@ def tag_and_embed(self, image_id: int) -> dict: phase = "video_infer" import numpy as np - preds = _maxpool_predictions( + preds = _aggregate_video_predictions( [tagger.infer(f, store_floor=settings.tagger_store_floor) - for f in frames] + for f in frames], + min_frames=settings.video_min_tag_frames, ) embedding = np.mean( [embedder.infer(f) for f in frames], axis=0 ).astype("float32") + log.info( + "tag_and_embed video aggregated %d tag(s) from %d frame(s) " + "(min_frames=%d): %s", + len(preds), len(frames), settings.video_min_tag_frames, ctx, + ) for f in frames: f.unlink(missing_ok=True) else: @@ -208,9 +217,17 @@ def tag_and_embed(self, image_id: int) -> dict: return {"status": "ok", "image_id": image_id, "tags": len(preds)} -def _sample_video_frames(src: Path, n: int) -> list[Path]: - """Extract n frames evenly between 10% and 90% of duration via ffmpeg. - Returns temp file paths (caller deletes). Empty list on failure.""" +def _sample_video_frames( + src: Path, *, interval: float, max_frames: int, +) -> list[Path]: + """Sample frames at a fixed CADENCE — one every `interval` seconds — so a + tag's frame-presence reflects real screen time regardless of video length + (#747). The count is capped at `max_frames`: a video longer than + interval×max_frames stretches the spacing instead of exploding the frame + count (keeps cost bounded so a long video can't hog the single ml-worker). + Frames are taken across the 5%–95% window (skip intro/outro black/cards) via + per-frame fast-seek. Returns temp file paths (caller deletes); [] on failure. + """ import json import subprocess import tempfile @@ -229,20 +246,25 @@ def _sample_video_frames(src: Path, n: int) -> list[Path]: if duration <= 0: return [] - start, end = duration * 0.10, duration * 0.90 - step = (end - start) / max(n - 1, 1) + start, end = duration * 0.05, duration * 0.95 + span = max(end - start, 0.0) + # Cadence count, clamped to [1, max_frames]. int(span/interval)+1 ≈ one frame + # per `interval` seconds across the window; the cap stretches spacing on very + # long videos. + n = max(1, min(int(span / interval) + 1, max(1, max_frames))) + step = span / max(n - 1, 1) out: list[Path] = [] tmpdir = Path(tempfile.mkdtemp(prefix="fc_vid_")) for i in range(n): ts = start + i * step - dest = tmpdir / f"frame_{i:02d}.jpg" + dest = tmpdir / f"frame_{i:04d}.jpg" try: subprocess.run( [ "ffmpeg", "-ss", f"{ts:.2f}", "-i", str(src), "-frames:v", "1", "-q:v", "3", "-y", str(dest), ], - check=True, capture_output=True, timeout=60, + check=True, capture_output=True, timeout=30, ) if dest.is_file(): out.append(dest) @@ -251,18 +273,38 @@ def _sample_video_frames(src: Path, n: int) -> list[Path]: return out -def _maxpool_predictions(per_frame: list[dict]) -> dict: - """Aggregate per-frame {name: TagPrediction} dicts by max confidence.""" - merged: dict[str, dict] = {} +def _aggregate_video_predictions(per_frame: list[dict], *, min_frames: int) -> dict: + """Aggregate per-frame {name: TagPrediction} into one prediction set (#747). + + A tag is kept only if it appears (≥ the tagger store floor, already applied) + in at least `min_frames` of the sampled frames — because sampling is at a + fixed cadence, that means it was on screen for roughly min_frames×interval + seconds, so a single-frame flicker / scene-transition artifact is dropped + while a genuine scene-local tag in a long video survives. Confidence is the + MEAN over the frames where the tag appears (not max — max re-inflated the + one-frame noise this whole change exists to remove). + + `min_frames` is clamped to the number of frames actually sampled so a very + short video (1–2 frames) still tags instead of dropping everything. + """ + n = len(per_frame) + if n == 0: + return {} + threshold = max(1, min(min_frames, n)) + agg: dict[str, dict] = {} for frame_preds in per_frame: for name, p in frame_preds.items(): - cur = merged.get(name) - if cur is None or p.confidence > cur["confidence"]: - merged[name] = { - "category": p.category, - "confidence": p.confidence, - } - return merged + cur = agg.get(name) + if cur is None: + agg[name] = {"category": p.category, "sum": p.confidence, "count": 1} + else: + cur["sum"] += p.confidence + cur["count"] += 1 + return { + name: {"category": v["category"], "confidence": v["sum"] / v["count"]} + for name, v in agg.items() + if v["count"] >= threshold + } @celery.task(name="backend.app.tasks.ml.backfill", bind=True) diff --git a/frontend/src/components/settings/MLThresholdSliders.vue b/frontend/src/components/settings/MLThresholdSliders.vue index 11bc76b..f0028ca 100644 --- a/frontend/src/components/settings/MLThresholdSliders.vue +++ b/frontend/src/components/settings/MLThresholdSliders.vue @@ -30,6 +30,39 @@ + + + +
Video tagging
+
+ Videos are tagged by sampling frames at a fixed cadence. A tag is kept + only if it shows up in enough frames (≈ that many × the interval in + seconds of screen time), which filters one-frame noise without losing + tags that only appear in part of a longer video. +
+ + + + + + + + + + + @@ -59,9 +92,14 @@ async function save() { const floor = local.tagger_store_floor local.suggestion_threshold_character = Math.max(local.suggestion_threshold_character, floor) local.suggestion_threshold_general = Math.max(local.suggestion_threshold_general, floor) + // Mirror the server invariant: a tag can't require more frames than are sampled. + local.video_min_tag_frames = Math.min(local.video_min_tag_frames, local.video_max_frames) const patch = {} for (const f of fields) patch[f.key] = local[f.key] patch.tagger_store_floor = local.tagger_store_floor + patch.video_frame_interval_seconds = local.video_frame_interval_seconds + patch.video_max_frames = local.video_max_frames + patch.video_min_tag_frames = local.video_min_tag_frames try { await store.patchSettings(patch) } catch (e) { toast({ text: e.message, type: 'error' }) } } diff --git a/tests/test_api_ml_admin.py b/tests/test_api_ml_admin.py index 76b1002..e46be48 100644 --- a/tests/test_api_ml_admin.py +++ b/tests/test_api_ml_admin.py @@ -56,6 +56,36 @@ async def test_suggestion_threshold_below_store_floor_rejected(client): assert "tagger_store_floor" in (await resp.get_json())["error"] +@pytest.mark.asyncio +async def test_video_tagging_settings_default_and_patch(client): + """#747: video cadence/noise knobs are exposed + patchable.""" + body = await (await client.get("/api/ml/settings")).get_json() + assert body["video_frame_interval_seconds"] == pytest.approx(4.0) + assert body["video_max_frames"] == 64 + assert body["video_min_tag_frames"] == 3 + + resp = await client.patch( + "/api/ml/settings", + json={"video_frame_interval_seconds": 5, "video_max_frames": 40, + "video_min_tag_frames": 4}, + ) + assert resp.status_code == 200 + out = await resp.get_json() + assert out["video_frame_interval_seconds"] == pytest.approx(5.0) + assert out["video_max_frames"] == 40 + assert out["video_min_tag_frames"] == 4 + + +@pytest.mark.asyncio +async def test_video_min_tag_frames_above_max_rejected(client): + resp = await client.patch( + "/api/ml/settings", + json={"video_max_frames": 10, "video_min_tag_frames": 20}, + ) + assert resp.status_code == 400 + assert "video_min_tag_frames" in (await resp.get_json())["error"] + + @pytest.mark.asyncio async def test_backfill_and_recompute_trigger(client): r1 = await client.post("/api/ml/backfill") diff --git a/tests/test_tasks_ml.py b/tests/test_tasks_ml.py index 3577360..62e5c5b 100644 --- a/tests/test_tasks_ml.py +++ b/tests/test_tasks_ml.py @@ -1,6 +1,6 @@ """tag_and_embed / backfill task tests. Models aren't in CI, so we test -the pure helpers (_maxpool_predictions, _is_video) as unit tests, and the -DB-touching backfill query as an integration test with monkeypatched +the pure helpers (_aggregate_video_predictions, _is_video) as unit tests, and +the DB-touching backfill query as an integration test with monkeypatched inference. """ @@ -9,7 +9,7 @@ from pathlib import Path import pytest from backend.app.services.ml.tagger import TagPrediction -from backend.app.tasks.ml import _is_video, _maxpool_predictions +from backend.app.tasks.ml import _aggregate_video_predictions, _is_video def test_is_video(): @@ -18,15 +18,32 @@ def test_is_video(): assert _is_video(Path("a.jpg")) is False -def test_maxpool_predictions(): - f1 = {"smile": TagPrediction("smile", "general", 0.6)} - f2 = { - "smile": TagPrediction("smile", "general", 0.9), - "sword": TagPrediction("sword", "general", 0.7), - } - merged = _maxpool_predictions([f1, f2]) - assert merged["smile"]["confidence"] == 0.9 - assert merged["sword"]["confidence"] == 0.7 +def _pred(name, conf, cat="general"): + return {name: TagPrediction(name, cat, conf)} + + +def test_aggregate_video_keeps_corroborated_and_means(): + # #747: 4 frames; "smile" in 3, "sword" in 1 (noise). min_frames=2. + per_frame = [ + {"smile": TagPrediction("smile", "general", 0.6), + "sword": TagPrediction("sword", "general", 0.9)}, + _pred("smile", 0.8), + _pred("smile", 0.7), + {}, + ] + out = _aggregate_video_predictions(per_frame, min_frames=2) + assert "sword" not in out # one-frame flicker dropped + assert abs(out["smile"]["confidence"] - (0.6 + 0.8 + 0.7) / 3) < 1e-9 # mean, not max + + +def test_aggregate_video_clamps_min_frames_to_sample_count(): + # Short video: 1 frame but min_frames=3 — clamp so it still tags. + out = _aggregate_video_predictions([_pred("solo", 0.8)], min_frames=3) + assert out["solo"]["confidence"] == 0.8 + + +def test_aggregate_video_empty(): + assert _aggregate_video_predictions([], min_frames=3) == {} @pytest.mark.integration From db7e1f2b599fde1ba90d4a4b288d1acc660c251e Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Tue, 16 Jun 2026 12:49:24 -0400 Subject: [PATCH 2/3] feat(ml): GPU-capable tagger + embedder with CPU fallback (#872) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Step 1 of GPU enablement (code only — CPU-safe, CI-green; the CUDA image is a separate step pending the host driver version). - New services/ml/device.py: FC_ML_DEVICE (auto|cuda|cpu) intent + VRAM knobs (FC_ML_ONNX_GPU_MEM_GB, FC_ML_TORCH_MEM_FRACTION). Per-worker-host bootstrap → env, not a DB setting (the GPU host runs CUDA, others CPU). - tagger: use CUDAExecutionProvider (with gpu_mem_limit) when requested AND the provider is actually present (onnxruntime-gpu), else CPUExecutionProvider. Logs the active providers. - embedder: move model + inputs to cuda when requested AND torch.cuda is available; cap torch's VRAM share; .detach().cpu() before numpy. fp32 kept so GPU embeddings stay in the same space as existing CPU ones. Both AND the env intent with the framework's real availability, so on CPU (CI / CPU onnxruntime / no GPU) they fall back cleanly — behavior unchanged. The 8GB P4 is shared by both frameworks, hence the conservative default caps. Tests: device env parsing. (tagger/embedder GPU paths are operator-verified on the GPU host — models aren't in CI.) Co-Authored-By: Claude Opus 4.8 --- backend/app/services/ml/device.py | 31 +++++++++++++++++++++++++++++ backend/app/services/ml/embedder.py | 29 ++++++++++++++++++++++----- backend/app/services/ml/tagger.py | 28 ++++++++++++++++++++------ tests/test_ml_device.py | 31 +++++++++++++++++++++++++++++ 4 files changed, 108 insertions(+), 11 deletions(-) create mode 100644 backend/app/services/ml/device.py create mode 100644 tests/test_ml_device.py diff --git a/backend/app/services/ml/device.py b/backend/app/services/ml/device.py new file mode 100644 index 0000000..aa6c252 --- /dev/null +++ b/backend/app/services/ml/device.py @@ -0,0 +1,31 @@ +"""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")) diff --git a/backend/app/services/ml/embedder.py b/backend/app/services/ml/embedder.py index 5da36c5..3ab7a4f 100644 --- a/backend/app/services/ml/embedder.py +++ b/backend/app/services/ml/embedder.py @@ -1,8 +1,11 @@ -"""SigLIP SO400M image-embedding wrapper (PyTorch CPU). +"""SigLIP SO400M image-embedding wrapper (PyTorch). -Direct port of ImageRepo's siglip.py. 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. +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. """ import os @@ -29,6 +32,7 @@ class Embedder: self._model = None self._processor = None self._torch = None + self._device = "cpu" def load(self) -> None: if self._model is not None: @@ -36,7 +40,17 @@ 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 # 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. @@ -51,6 +65,8 @@ 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).""" @@ -58,9 +74,12 @@ 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 - return pooled[0].numpy().astype(np.float32) + # .detach().cpu() so a CUDA tensor converts to numpy (no-op on CPU). + return pooled[0].detach().cpu().numpy().astype(np.float32) _default_embedder: Embedder | None = None diff --git a/backend/app/services/ml/tagger.py b/backend/app/services/ml/tagger.py index 6348595..178387d 100644 --- a/backend/app/services/ml/tagger.py +++ b/backend/app/services/ml/tagger.py @@ -1,8 +1,10 @@ """Camie-tagger-v2 ONNX wrapper. -CPU-only, single-image at a time. 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. 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). v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB) @@ -12,6 +14,7 @@ ImageNet normalize, NCHW layout, sigmoid on refined logits (output[1]). """ import json +import logging import os from dataclasses import dataclass from pathlib import Path @@ -19,6 +22,8 @@ from pathlib import Path import numpy as np from PIL import Image, ImageFile +log = logging.getLogger(__name__) + # 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 # imports fine without it (the suggestion service imports SURFACED_CATEGORIES @@ -117,9 +122,20 @@ class Tagger: # without onnxruntime (CI / lean web image). import onnxruntime as ort - session = ort.InferenceSession( - str(model_path), providers=["CPUExecutionProvider"] - ) + 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()) self._input_name = session.get_inputs()[0].name # Assign sentinels last so a partial load isn't observable. self._tag_names = names diff --git a/tests/test_ml_device.py b/tests/test_ml_device.py new file mode 100644 index 0000000..2aed435 --- /dev/null +++ b/tests/test_ml_device.py @@ -0,0 +1,31 @@ +"""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 From 60a9c9e6ef074ce3ae84598c8b1b76b0edb64c23 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Tue, 16 Jun 2026 13:39:55 -0400 Subject: [PATCH 3/3] refactor(ml): drop GPU code, cap inference threads by default (#747/#872) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- backend/app/services/ml/device.py | 31 ---------------------- backend/app/services/ml/embedder.py | 37 ++++++++++---------------- backend/app/services/ml/tagger.py | 40 +++++++++++++---------------- tests/test_ml_device.py | 31 ---------------------- 4 files changed, 31 insertions(+), 108 deletions(-) delete mode 100644 backend/app/services/ml/device.py delete mode 100644 tests/test_ml_device.py diff --git a/backend/app/services/ml/device.py b/backend/app/services/ml/device.py deleted file mode 100644 index aa6c252..0000000 --- a/backend/app/services/ml/device.py +++ /dev/null @@ -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")) diff --git a/backend/app/services/ml/embedder.py b/backend/app/services/ml/embedder.py index 3ab7a4f..d94f71e 100644 --- a/backend/app/services/ml/embedder.py +++ b/backend/app/services/ml/embedder.py @@ -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 diff --git a/backend/app/services/ml/tagger.py b/backend/app/services/ml/tagger.py index 178387d..0f15c4a 100644 --- a/backend/app/services/ml/tagger.py +++ b/backend/app/services/ml/tagger.py @@ -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 diff --git a/tests/test_ml_device.py b/tests/test_ml_device.py deleted file mode 100644 index 2aed435..0000000 --- a/tests/test_ml_device.py +++ /dev/null @@ -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