Extension channels: dev and main each carry their own signed extension #237

Merged
bvandeusen merged 16 commits from dev into main 2026-08-27 12:49:40 -04:00
4 changed files with 195 additions and 22 deletions
Showing only changes of commit 5a0e1bbd03 - Show all commits
+51
View File
@@ -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"])
)
+20 -13
View File
@@ -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 = [
+5 -9
View File
@@ -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
+119
View File
@@ -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, [])