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
This commit is contained in:
@@ -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"])
|
||||
)
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user