diff --git a/backend/app/services/image_tag_apply.py b/backend/app/services/image_tag_apply.py new file mode 100644 index 0000000..65257f4 --- /dev/null +++ b/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"]) + ) diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py index dbaa8cf..c6d2a02 100644 --- a/backend/app/services/ml/heads.py +++ b/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 = [ diff --git a/backend/app/services/wip_title.py b/backend/app/services/wip_title.py index 49a18a4..8608cab 100644 --- a/backend/app/services/wip_title.py +++ b/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 diff --git a/tests/test_image_tag_apply.py b/tests/test_image_tag_apply.py new file mode 100644 index 0000000..797414d --- /dev/null +++ b/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, [])