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,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, [])
|
||||
Reference in New Issue
Block a user