perf(ml): batch the auto-apply sweeps' image_tag inserts (#3072)

Item 1 of #3072. Both sweeps issued a single-row pg_insert(image_tag)
from inside their per-image loop. Steady state that is nothing; a first
sweep over a back-catalogue is one round-trip per applied tag, tens of
thousands of them. Each chunk now collects its rows and writes them in
one statement.

The ticket suggested one insert per chunk PER TAG. A single multi-row
VALUES carries every tag at once, so it is one statement per chunk full
stop — and the sweeps already accumulate across all heads before they
commit, so nothing had to be restructured to allow it.

Not a new helper: wip_title.apply_wip_image_tags was already doing the
chunked ON CONFLICT DO NOTHING insert, so that shape is extracted to
services/image_tag_apply.insert_image_tags and all three writers share
it. The extraction deliberately leaves wip_title's pre-SELECT behind
rather than pulling it into the shared function — the sweeps don't need
it (their `skip` sets already exclude applied and rejected images) and
it exists only to produce an accurate count, which the sweeps also
compute themselves. So the shared primitive returns nothing: psycopg
reports rowcount -1 for a multi-row ON CONFLICT DO NOTHING insert, and a
count taken from the statement would be a lie rather than an
approximation.

Ordering note for the system-tag sweep: tag rows are now written after
that chunk's PresentationReview rows rather than interleaved before
them. Safe — PresentationReview FKs to image_record and tag, not to
image_tag.

Chunk size stays 5000: 5000 rows x 3 bound params = 15000, inside
Postgres' 65535-parameter ceiling with room to spare.

tests/test_image_tag_apply.py covers the primitive directly, since it is
now the single place three writers can be wrong at once — most
importantly that a re-run never restamps a hand-applied tag's source,
which would silently poison head training (it excludes the auto
sources).

Left alone: _insert_presentation_review is still per-row, and the
retract path still deletes per-row. Both operate on sets that are small
by construction, unlike the apply path.

Refs #3072

backend/app/services/image_tag_apply.py

@@ -0,0 +1,51 @@
+"""Bulk, idempotent writes to the ``image_tag`` association table.
+
+Three writers attach tags to images in bulk: the WIP-title backfill
+(`wip_title.apply_wip_image_tags`), the concept-head auto-apply sweep and the
+system-tag auto-apply sweep (both in `ml/heads.py`). The two sweeps used to
+issue ONE INSERT PER ROW from inside their per-image loop — fine in steady
+state, but a first pass over a back-catalogue is tens of thousands of
+individual round-trips (#3072). All three share this one chunked multi-row
+insert now.
+
+Sync only: every caller runs on a sync ``Session`` (the Celery task path). No
+async service writes image_tag in bulk, so there is no async sibling to keep in
+step — unlike `db_helpers.get_or_create`, which does have one.
+"""
+
+from __future__ import annotations
+
+from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.orm import Session
+
+from ..models.tag import image_tag
+
+# 5000 rows x 3 bound params = 15000, comfortably inside Postgres' 65535-param
+# ceiling for a single statement. Raising this past ~21000 rows would exceed it.
+INSERT_CHUNK = 5000
+
+
+def insert_image_tags(
+    session: Session, rows: list[dict], *, chunk: int = INSERT_CHUNK
+) -> None:
+    """Attach ``rows`` to their images, skipping any tag already on one.
+
+    Each row is ``{"image_record_id": int, "tag_id": int, "source": str}``.
+    Does NOT commit — the caller owns the transaction.
+
+    ON CONFLICT DO NOTHING against the (image_record_id, tag_id) primary key,
+    so an existing tag keeps its ORIGINAL ``source``: re-running a sweep can
+    never re-stamp a tag the operator applied by hand as machine-applied.
+
+    Returns nothing on purpose. psycopg reports ``rowcount`` -1 for a multi-row
+    ON CONFLICT DO NOTHING insert (it runs via an executemany path), so a count
+    taken from the statement would be a lie rather than an approximation.
+    Callers that need an accurate count derive it themselves — see
+    `wip_title.apply_wip_image_tags`' pre-SELECT, and the sweeps' `skip` sets.
+    """
+    for start in range(0, len(rows), chunk):
+        session.execute(
+            pg_insert(image_tag)
+            .values(rows[start:start + chunk])
+            .on_conflict_do_nothing(index_elements=["image_record_id", "tag_id"])
+        )

backend/app/services/ml/heads.py

@@ -41,6 +41,7 @@ from ...models import (
     TagSuggestionRejection,
 )
 from ...models.tag import CHROME_SYSTEM_TAGS, PROCESS_SYSTEM_TAGS, image_tag
+from ..image_tag_apply import insert_image_tags
 from .training_data import (
     _AUTO_SOURCES,
     _applied_or_rejected,
@@ -757,6 +758,10 @@ def auto_apply_sweep(
         Xn = _l2norm(np.vstack([emb[i] for i in cids]).astype(np.float32), np)
         probs = _sigmoid(Xn @ W.T + b, np)              # (N, H)
         scanned += len(cids)
+        # Collected across every head, then written as ONE insert below. Was an
+        # insert per applied tag from inside this loop, which on a first sweep
+        # over a back-catalogue is tens of thousands of round-trips (#3072).
+        pending: list[dict] = []
         for h in range(len(rows)):
             tid = tag_ids[h]
             for idx in np.where(probs[:, h] >= thr[h])[0]:
@@ -766,12 +771,12 @@ def auto_apply_sweep(
                 skip[tid].add(iid)
                 applied[h] += 1
                 if not dry_run:
-                    session.execute(
-                        pg_insert(image_tag)
-                        .values(image_record_id=iid, tag_id=tid, source="head_auto")
-                        .on_conflict_do_nothing()
-                    )
+                    pending.append({
+                        "image_record_id": iid, "tag_id": tid,
+                        "source": "head_auto",
+                    })
         if not dry_run:
+            insert_image_tags(session, pending)
             session.commit()
         run.last_progress_at = datetime.now(UTC)
         session.commit()
@@ -913,6 +918,11 @@ def system_tag_auto_apply_sweep(
         if Wc is not None:
             max_c, arg_c = _conflict_scores(Xn, Wc, bc, np)   # (N,), (N,)
         scanned += len(cids)
+        # Same batching as auto_apply_sweep (#3072): collect the chunk's rows
+        # and write them once, below. The PresentationReview rows stay per-row —
+        # they FK to image_record/tag, not to image_tag, so writing the tags
+        # after them is safe, and a flagged conflict is rare by construction.
+        pending: list[dict] = []
         for p in range(len(pres)):
             tid = pres_tag_ids[p]
             for idx in np.where(probs[:, p] >= thr)[0]:
@@ -922,14 +932,10 @@ def system_tag_auto_apply_sweep(
                 skip[tid].add(iid)
                 applied[p] += 1
                 if not dry_run:
-                    session.execute(
-                        pg_insert(image_tag)
-                        .values(
-                            image_record_id=iid, tag_id=tid,
-                            source=source,
-                        )
-                        .on_conflict_do_nothing()
-                    )
+                    pending.append({
+                        "image_record_id": iid, "tag_id": tid,
+                        "source": source,
+                    })
                 # Guard 2: also looks like real content → still apply, but flag it
                 # for the review strip instead of silently marking (chrome hides,
                 # process stays visible — either way the operator gets a heads-up).
@@ -944,6 +950,7 @@ def system_tag_auto_apply_sweep(
                             mode=mode,
                         )
         if not dry_run:
+            insert_image_tags(session, pending)
             session.commit()
 
     concepts = [

backend/app/services/wip_title.py

@@ -20,10 +20,10 @@ family gains one member.
 import re
 
 from sqlalchemy import select
-from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.orm import Session
 
 from ..models.tag import WIP_SYSTEM_TAG, Tag, image_tag
+from .image_tag_apply import insert_image_tags
 
 # image_tag.source stamped on title-heuristic WIP tags — distinct from the other
 # apply sources so provenance stays legible and a future undo can target only these.
@@ -113,13 +113,9 @@ def apply_wip_image_tags(
         to_insert = [iid for iid in chunk if iid not in already]
         if not to_insert:
             continue
-        session.execute(
-            pg_insert(image_tag)
-            .values([
-                {"image_record_id": iid, "tag_id": tag_id, "source": source}
-                for iid in to_insert
-            ])
-            .on_conflict_do_nothing(index_elements=["image_record_id", "tag_id"])
-        )
+        insert_image_tags(session, [
+            {"image_record_id": iid, "tag_id": tag_id, "source": source}
+            for iid in to_insert
+        ])
         inserted += len(to_insert)
     return inserted

tests/test_image_tag_apply.py

@@ -0,0 +1,119 @@
+"""`insert_image_tags` — the shared bulk write behind the WIP-title backfill and
+both auto-apply sweeps (#3072).
+
+The sweeps previously issued one INSERT per applied tag from inside their
+per-image loop; they now hand this helper a chunk's worth of rows. That makes
+this function the single place three writers can be wrong at once, so it is
+tested directly rather than only through its callers.
+"""
+import pytest
+from sqlalchemy import select
+
+from backend.app.models import ImageRecord, Tag, TagKind
+from backend.app.models.tag import image_tag
+from backend.app.services.image_tag_apply import insert_image_tags
+
+pytestmark = pytest.mark.integration
+
+_N = 0
+
+
+def _img(db_sync):
+    global _N
+    _N += 1
+    rec = ImageRecord(
+        path=f"/images/ita/{_N}.jpg", sha256=f"a{_N:063d}",
+        size_bytes=1, mime="image/jpeg", width=1, height=1,
+        origin="imported_filesystem", integrity_status="unknown",
+    )
+    db_sync.add(rec)
+    db_sync.flush()
+    return rec
+
+
+def _tag(db_sync, name):
+    t = Tag(name=name, kind=TagKind.general)
+    db_sync.add(t)
+    db_sync.flush()
+    return t
+
+
+def _rows(db_sync, tag_id):
+    """(image_record_id, source) pairs currently carrying `tag_id`."""
+    return dict(db_sync.execute(
+        select(image_tag.c.image_record_id, image_tag.c.source)
+        .where(image_tag.c.tag_id == tag_id)
+    ).all())
+
+
+def _row(image_record_id, tag_id, source):
+    return {
+        "image_record_id": image_record_id, "tag_id": tag_id, "source": source,
+    }
+
+
+def test_inserts_every_row_in_one_call(db_sync):
+    t = _tag(db_sync, "ita-basic")
+    imgs = [_img(db_sync) for _ in range(3)]
+    insert_image_tags(
+        db_sync, [_row(i.id, t.id, "head_auto") for i in imgs]
+    )
+    assert _rows(db_sync, t.id) == {i.id: "head_auto" for i in imgs}
+
+
+def test_spans_several_tags_in_a_single_call(db_sync):
+    """The sweeps accumulate across ALL heads before flushing, so one call
+    carries rows for different tags. A per-tag implementation would drop all
+    but the first."""
+    t1, t2 = _tag(db_sync, "ita-multi-1"), _tag(db_sync, "ita-multi-2")
+    a, b = _img(db_sync), _img(db_sync)
+    insert_image_tags(db_sync, [
+        _row(a.id, t1.id, "head_auto"), _row(b.id, t1.id, "head_auto"),
+        _row(a.id, t2.id, "head_auto"),
+    ])
+    assert _rows(db_sync, t1.id) == {a.id: "head_auto", b.id: "head_auto"}
+    assert _rows(db_sync, t2.id) == {a.id: "head_auto"}
+
+
+def test_an_existing_tag_keeps_its_original_source(db_sync):
+    """THE assertion this helper exists for. A sweep re-running over an image
+    the operator tagged by hand must not restamp it as machine-applied — that
+    would silently poison the head's own training data, which excludes the
+    auto sources. ON CONFLICT DO NOTHING, never DO UPDATE."""
+    t = _tag(db_sync, "ita-manual")
+    rec = _img(db_sync)
+    insert_image_tags(db_sync, [_row(rec.id, t.id, "manual")])
+
+    insert_image_tags(db_sync, [_row(rec.id, t.id, "head_auto")])
+
+    assert _rows(db_sync, t.id) == {rec.id: "manual"}
+
+
+def test_a_repeat_within_one_call_does_not_raise(db_sync):
+    """Two heads can both fire on the same (image, tag) inside one chunk. The
+    conflict is resolved by the statement, not by the caller de-duplicating."""
+    t = _tag(db_sync, "ita-dupe")
+    rec = _img(db_sync)
+    insert_image_tags(db_sync, [
+        _row(rec.id, t.id, "head_auto"), _row(rec.id, t.id, "head_auto"),
+    ])
+    assert _rows(db_sync, t.id) == {rec.id: "head_auto"}
+
+
+def test_more_rows_than_the_chunk_size_all_land(db_sync):
+    """The chunk exists to stay under Postgres' 65535 bound-parameter ceiling.
+    Driven with a tiny chunk so the split is real rather than theoretical — at
+    the 5000 default no test would ever reach a second statement."""
+    t = _tag(db_sync, "ita-chunked")
+    imgs = [_img(db_sync) for _ in range(7)]
+    insert_image_tags(
+        db_sync, [_row(i.id, t.id, "head_auto") for i in imgs], chunk=2
+    )
+    assert _rows(db_sync, t.id) == {i.id: "head_auto" for i in imgs}
+
+
+def test_no_rows_is_a_no_op(db_sync):
+    """A dry-run chunk, or a chunk where every candidate was already skipped,
+    hands over an empty list. `.values([])` is a SQL error, so the empty case
+    must never reach the statement."""
+    insert_image_tags(db_sync, [])