feat(activity): image tagging jobs for bulk Add tags and Auto-tag
POST /prompt_manager/jobs/image-tags starts one record-changing "image-tags" job over 1..500 listed images (202, or the same 409 busy envelope as /jobs/scan). Add tags hashes each current file and unions a fixed set into that byte record; Auto-tag captures each still image, asks one legacy engine for its top N suggestions on the capture (count from the options or autotag.defaults.suggestion_limit), re-hashes the file and unions only onto the same bytes. Outcomes are grouped as tagged, unchanged, skipped (video, too large, file changed, file not found, not a still image) and errors; cancel stops at a file boundary. Auto-tag holds the shared model gate for the whole job (409 model-busy at start when it is held), loads the model once and never downloads one: an unavailable model fails the job before the first file. The registry gains blocker() so a start answers busy before taking the gate.
This commit is contained in:
Vito Sansevero
@@ -7,6 +7,7 @@ Split into domain-specific mixins for maintainability:
|
||||
- LoggingRoutesMixin: log management endpoints
|
||||
- AutotagRoutesMixin: auto-tagging model management and tagging
|
||||
- JobRoutesMixin: the session job registry behind Activity (py/api/jobs.py)
|
||||
- ImageTagJobsMixin: bulk Add tags / Auto-tag as image-tags jobs
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -32,6 +33,7 @@ from .history import HistoryRoutesMixin
|
||||
from .images import ImageRoutesMixin
|
||||
from .admin import AdminRoutesMixin
|
||||
from .job_routes import JobRoutesMixin
|
||||
from .image_tag_jobs import ImageTagJobsMixin
|
||||
from .jobs import JobRegistry
|
||||
from .logging_routes import LoggingRoutesMixin
|
||||
from .autotag_routes import AutotagRoutesMixin
|
||||
@@ -260,6 +262,7 @@ class PromptManagerAPI(
|
||||
ImageRoutesMixin,
|
||||
AdminRoutesMixin,
|
||||
JobRoutesMixin,
|
||||
ImageTagJobsMixin,
|
||||
LoggingRoutesMixin,
|
||||
AutotagRoutesMixin,
|
||||
LoraIntegrationMixin,
|
||||
|
||||
@@ -0,0 +1,446 @@
|
||||
"""Bulk image tagging as an Activity job: Add tags and Auto-tag (#249 D1).
|
||||
|
||||
``POST /prompt_manager/jobs/image-tags`` starts one ``image-tags`` job over
|
||||
1..500 listed images, named by the listing identity ``{root_index,
|
||||
root_scope, path}``. The job is record-changing (it writes image tag sets and
|
||||
new vocabulary names), so it never runs beside the metadata scan.
|
||||
|
||||
- **Add tags** (``mode: "add"``) hashes each current file and unions a fixed
|
||||
tag set into the record of exactly those bytes.
|
||||
- **Auto-tag** (``mode: "auto"``) copies each still image into a private
|
||||
capture, asks one legacy engine for its top N tags on the capture, re-hashes
|
||||
the current file and unions the suggestions only when the bytes are still
|
||||
the captured ones. It holds the shared model gate for the whole job, loads
|
||||
the model once and never downloads one.
|
||||
|
||||
Every step reuses the Tag dialog's pieces: the capture store, the model gate,
|
||||
the suggestion worker and the storage call ``add_image_tags``. Files are
|
||||
handled one at a time and the job stops at a file boundary when cancelled,
|
||||
keeping what it already wrote. Rows and log lines carry root-relative paths
|
||||
and public root labels only.
|
||||
"""
|
||||
|
||||
import functools
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .. import model_gate
|
||||
from .admin import _plural
|
||||
from .annotation_suggest import (
|
||||
OptionError,
|
||||
unavailable_message,
|
||||
validate_options,
|
||||
)
|
||||
from .annotations import (
|
||||
_MESSAGES,
|
||||
_error,
|
||||
_pinned_suggestions,
|
||||
_settled,
|
||||
_submit,
|
||||
_target,
|
||||
)
|
||||
from .autotag_routes import model_busy_response
|
||||
from .images import _public_path
|
||||
from .job_routes import busy_response
|
||||
from .jobs import IMAGE_TAG_GROUP_NAMES, IMAGE_TAG_JOB_KIND, JobBusy
|
||||
from .prompts import bad_request
|
||||
|
||||
try:
|
||||
from ...database.annotations import ImageTagLimitError, normalize_image_tags
|
||||
from ...database.library_commands import LibraryCommandError
|
||||
from ...utils.media_capture import (
|
||||
CaptureCapacity,
|
||||
CaptureChanged,
|
||||
CaptureDenied,
|
||||
CaptureError,
|
||||
CaptureTooLarge,
|
||||
)
|
||||
from ...utils.video_metadata import is_video_path
|
||||
except ImportError:
|
||||
from database.annotations import ImageTagLimitError, normalize_image_tags
|
||||
from database.library_commands import LibraryCommandError
|
||||
from utils.media_capture import (
|
||||
CaptureCapacity,
|
||||
CaptureChanged,
|
||||
CaptureDenied,
|
||||
CaptureError,
|
||||
CaptureTooLarge,
|
||||
)
|
||||
from utils.video_metadata import is_video_path
|
||||
|
||||
MAX_JOB_TARGETS = 500
|
||||
# 500 targets of up to ~1.1 KiB each plus the options: the annotation routes'
|
||||
# 64 KiB bound is for one target, too small for a full selection.
|
||||
MAX_JOB_BODY_BYTES = 1024 * 1024
|
||||
|
||||
ADD_TITLE = "Add tags to images"
|
||||
AUTO_TITLE = "Auto-tag images"
|
||||
GATE_OWNER_KIND = "image-tag-job"
|
||||
|
||||
_FIELDS = {
|
||||
"add": frozenset({"mode", "targets", "tags"}),
|
||||
"auto": frozenset({"mode", "targets", "model_type", "options"}),
|
||||
}
|
||||
_VERBS = {"add": "Add tags", "auto": "Auto-tag"}
|
||||
_LOG_PREFIXES = {"tagged": "+", "unchanged": "=", "skipped": "-", "errors": "!"}
|
||||
|
||||
NO_NEW_TAGS = "No new tags"
|
||||
VIDEO_SKIPPED = "Video: manual tags only"
|
||||
TOO_LARGE = "Too large to verify"
|
||||
FILE_CHANGED = "File changed"
|
||||
FILE_NOT_FOUND = "File not found"
|
||||
NOT_A_STILL = "Not a supported still image"
|
||||
INFERENCE_FAILED = "Auto-tagging failed. Check server logs for details."
|
||||
SAVE_FAILED = "The tags could not be saved. Check server logs for details."
|
||||
READ_FAILED = "The file could not be read."
|
||||
SOURCE_CHANGED = "The source changed. Refresh the image collection."
|
||||
STORAGE_RESTORED = "Prompt storage was restored; start the job again."
|
||||
|
||||
|
||||
class _JobStopped(Exception):
|
||||
"""No later file can be written (prompt storage was restored)."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Spec:
|
||||
"""A validated start request."""
|
||||
|
||||
mode: str
|
||||
targets: tuple
|
||||
tags: tuple = ()
|
||||
model_type: str | None = None
|
||||
options: dict | None = None
|
||||
|
||||
@property
|
||||
def title(self) -> str:
|
||||
return ADD_TITLE if self.mode == "add" else AUTO_TITLE
|
||||
|
||||
def scope(self) -> dict:
|
||||
scope = {"mode": self.mode, "images": len(self.targets)}
|
||||
if self.mode == "auto":
|
||||
scope["model_type"] = self.model_type
|
||||
return scope
|
||||
|
||||
|
||||
def _parse_targets(raw):
|
||||
"""Unique validated targets in request order, or ``(None, error)``."""
|
||||
if not isinstance(raw, list) or not 1 <= len(raw) <= MAX_JOB_TARGETS:
|
||||
return None, bad_request(f"targets must list 1 to {MAX_JOB_TARGETS} images")
|
||||
seen, targets = set(), []
|
||||
for item in raw:
|
||||
target, error = _target(item)
|
||||
if error is not None:
|
||||
return None, error
|
||||
if target not in seen:
|
||||
seen.add(target)
|
||||
targets.append(target)
|
||||
return tuple(targets), None
|
||||
|
||||
|
||||
def _parse_add(targets, data):
|
||||
try:
|
||||
tags = normalize_image_tags(data["tags"])
|
||||
except ImageTagLimitError as exc:
|
||||
return None, _error(400, "tag-limit", str(exc))
|
||||
except ValueError as exc:
|
||||
return None, bad_request(str(exc))
|
||||
if not tags:
|
||||
return None, bad_request("tags must name at least one tag")
|
||||
return _Spec("add", targets, tags=tuple(tags)), None
|
||||
|
||||
|
||||
def _parse_auto(targets, data):
|
||||
from ..config import MAX_SUGGESTION_LIMIT, AutotagConfig
|
||||
|
||||
try:
|
||||
model_type, options = validate_options(
|
||||
data["model_type"],
|
||||
data["options"],
|
||||
default_limit=AutotagConfig.suggestion_limit(),
|
||||
max_limit=MAX_SUGGESTION_LIMIT,
|
||||
)
|
||||
except OptionError as exc:
|
||||
return None, bad_request(str(exc))
|
||||
return _Spec("auto", targets, model_type=model_type, options=options), None
|
||||
|
||||
|
||||
def parse_image_tag_request(data: dict):
|
||||
"""A :class:`_Spec` for a start body, or ``(None, error response)``.
|
||||
|
||||
``{mode:"add", targets, tags}`` or ``{mode:"auto", targets, model_type,
|
||||
options}`` exactly; options follow the Tag dialog's Generate rules.
|
||||
"""
|
||||
fields = _FIELDS.get(data.get("mode"))
|
||||
if fields is None:
|
||||
return None, bad_request('mode must be "add" or "auto"')
|
||||
if set(data) != fields:
|
||||
return None, bad_request(
|
||||
"Request must contain exactly " + ", ".join(sorted(fields))
|
||||
)
|
||||
targets, error = _parse_targets(data["targets"])
|
||||
if error is not None:
|
||||
return None, error
|
||||
if data["mode"] == "add":
|
||||
return _parse_add(targets, data)
|
||||
return _parse_auto(targets, data)
|
||||
|
||||
|
||||
def _capture_outcome(exc: CaptureError, root, relative: str):
|
||||
"""The (group, detail) of a file the capture store refused to read."""
|
||||
if isinstance(exc, CaptureTooLarge):
|
||||
return "skipped", TOO_LARGE
|
||||
if isinstance(exc, CaptureDenied):
|
||||
return "skipped", FILE_NOT_FOUND
|
||||
if isinstance(exc, CaptureChanged):
|
||||
path = os.path.join(root, *relative.split("/"))
|
||||
return "skipped", FILE_CHANGED if os.path.lexists(path) else FILE_NOT_FOUND
|
||||
if isinstance(exc, CaptureCapacity):
|
||||
return "errors", _MESSAGES["review-capacity"]
|
||||
return "errors", READ_FAILED
|
||||
|
||||
|
||||
def _load_model(service, model_type: str, use_gpu: bool) -> None:
|
||||
"""Worker thread: load (or swap to) *model_type*; never downloads."""
|
||||
if service.is_model_loaded() and service.get_loaded_model_type() == model_type:
|
||||
return
|
||||
if service.load_model(model_type, use_gpu) is False:
|
||||
raise RuntimeError(f"{model_type} could not be loaded")
|
||||
|
||||
|
||||
class ImageTagJobsMixin:
|
||||
"""The ``image-tags`` job kind: start route and runner."""
|
||||
|
||||
async def start_image_tag_job(self, request: web.Request) -> web.Response:
|
||||
"""Start bulk tagging (202), or 409 ``busy`` / ``model-busy``.
|
||||
|
||||
Every target's root scope is checked before the job starts
|
||||
(``source-changed``). Auto-tag takes the model gate here, so a start
|
||||
while anything else holds it is refused at once.
|
||||
"""
|
||||
data, error = await self._read_json_body(request, max_bytes=MAX_JOB_BODY_BYTES)
|
||||
if error is not None:
|
||||
return error
|
||||
spec, error = parse_image_tag_request(data)
|
||||
if error is not None:
|
||||
return error
|
||||
for index, scope in dict.fromkeys((t[0], t[1]) for t in spec.targets):
|
||||
_root, error = self._scoped_root(index, scope)
|
||||
if error is not None:
|
||||
return error
|
||||
epoch = await self._run_in_executor(self.db.get_execution_epoch)
|
||||
# A running job in the way answers first: it says which job to wait for.
|
||||
blocker = self._jobs.blocker(IMAGE_TAG_JOB_KIND, record_changing=True)
|
||||
if blocker is not None:
|
||||
return busy_response(JobBusy(blocker))
|
||||
owner = None
|
||||
if spec.mode == "auto":
|
||||
owner = model_gate.new_owner(GATE_OWNER_KIND)
|
||||
if not model_gate.try_acquire(owner):
|
||||
return model_busy_response()
|
||||
try:
|
||||
job = self._jobs.start(
|
||||
IMAGE_TAG_JOB_KIND,
|
||||
spec.title,
|
||||
record_changing=True,
|
||||
scope=spec.scope(),
|
||||
runner=functools.partial(self._run_image_tag_job, spec, epoch, owner),
|
||||
groups=IMAGE_TAG_GROUP_NAMES,
|
||||
)
|
||||
except JobBusy as busy:
|
||||
if owner is not None:
|
||||
model_gate.release(owner)
|
||||
return busy_response(busy)
|
||||
if owner is not None:
|
||||
# Also when the task is cancelled before its runner ever starts.
|
||||
job.task.add_done_callback(lambda _task: model_gate.release(owner))
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
# -- runner ----------------------------------------------------------------
|
||||
|
||||
async def _run_image_tag_job(self, spec: _Spec, epoch: str, owner, job) -> None:
|
||||
verb = _VERBS[spec.mode]
|
||||
total = len(spec.targets)
|
||||
job.update(total=total, status_line=f"Processing file 0 of {total:,}")
|
||||
started = f"{verb} started · {_plural(total, 'image')}"
|
||||
if spec.mode == "auto":
|
||||
started += f" · {spec.model_type}"
|
||||
job.log_line(started)
|
||||
try:
|
||||
service = None
|
||||
if spec.mode == "auto":
|
||||
service = await self._prepare_tagging_model(job, spec, owner)
|
||||
if service is None:
|
||||
job.finish("failed")
|
||||
return
|
||||
if not await self._tag_each_file(job, spec, epoch, owner, service):
|
||||
job.log_line(f"{verb} stopped · {STORAGE_RESTORED}")
|
||||
job.finish("failed")
|
||||
return
|
||||
state = "cancelled" if job.cancelled else "finished"
|
||||
job.log_line(f"{verb} {state} · {_plural(job.processed, 'file')} processed")
|
||||
job.publish({"type": state, "progress": job.progress})
|
||||
job.finish(state)
|
||||
finally:
|
||||
if owner is not None:
|
||||
model_gate.release(owner)
|
||||
|
||||
async def _tag_each_file(self, job, spec, epoch, owner, service) -> bool:
|
||||
"""Tag the targets in order; False when the job had to stop."""
|
||||
total = len(spec.targets)
|
||||
for number, target in enumerate(spec.targets, 1):
|
||||
if job.cancelled:
|
||||
break
|
||||
job.update(status_line=f"Processing file {number:,} of {total:,}")
|
||||
job.publish({"type": "progress", "progress": job.progress})
|
||||
stopped = False
|
||||
try:
|
||||
group, detail = await self._tag_one_file(
|
||||
spec, epoch, owner, service, target
|
||||
)
|
||||
except _JobStopped:
|
||||
stopped, group, detail = True, "errors", STORAGE_RESTORED
|
||||
except Exception:
|
||||
self.logger.exception("Image tagging failed for one file")
|
||||
group, detail = "errors", SAVE_FAILED
|
||||
self._record_tag_outcome(job, target, group, detail)
|
||||
job.update(processed=number, progress=number * 100 // total)
|
||||
if stopped:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _prepare_tagging_model(self, job, spec: _Spec, owner):
|
||||
"""The loaded service, or None (logged) when the job cannot run."""
|
||||
from ..autotag import MODELS, get_autotag_service
|
||||
|
||||
service = get_autotag_service()
|
||||
status = await self._run_in_executor(service.get_models_status)
|
||||
if not status.get(spec.model_type, {}).get("downloaded"):
|
||||
name = MODELS.get(spec.model_type, {}).get("name", spec.model_type)
|
||||
message = unavailable_message(name)
|
||||
job.update(status_line=message)
|
||||
job.log_line(f"! {message}")
|
||||
return None
|
||||
try:
|
||||
await model_gate.run(
|
||||
owner,
|
||||
_load_model,
|
||||
service,
|
||||
spec.model_type,
|
||||
spec.options["use_gpu"],
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("The auto-tag job could not load its model")
|
||||
job.update(status_line=INFERENCE_FAILED)
|
||||
job.log_line(f"! {INFERENCE_FAILED}")
|
||||
return None
|
||||
return service
|
||||
|
||||
async def _tag_one_file(self, spec, epoch, owner, service, target):
|
||||
"""The (group, detail) of one target; raises _JobStopped on a restore."""
|
||||
index, scope, relative = target
|
||||
root, error = self._scoped_root(index, scope)
|
||||
if error is not None:
|
||||
return "errors", SOURCE_CHANGED
|
||||
store = self._capture_store()
|
||||
if spec.mode == "add":
|
||||
try:
|
||||
facts = await self._run_in_executor(store.verify, root, relative)
|
||||
except CaptureError as exc:
|
||||
return _capture_outcome(exc, root, relative)
|
||||
return await self._union_image_tags(facts, list(spec.tags), epoch)
|
||||
if is_video_path(relative):
|
||||
return "skipped", VIDEO_SKIPPED
|
||||
try:
|
||||
capture = await _settled(
|
||||
_submit(
|
||||
store.capture,
|
||||
root,
|
||||
relative,
|
||||
epoch=epoch,
|
||||
scope=scope,
|
||||
root_index=index,
|
||||
),
|
||||
lambda made: made is not None and store.release(made.token),
|
||||
)
|
||||
except CaptureError as exc:
|
||||
return _capture_outcome(exc, root, relative)
|
||||
try:
|
||||
return await self._auto_tag_capture(
|
||||
store, capture, spec, epoch, owner, service, root, relative
|
||||
)
|
||||
finally:
|
||||
store.release(capture.token)
|
||||
|
||||
async def _auto_tag_capture(
|
||||
self, store, capture, spec, epoch, owner, service, root, relative
|
||||
):
|
||||
"""Suggest on the capture, then union only onto the same current bytes."""
|
||||
if not capture.inferable:
|
||||
return "skipped", NOT_A_STILL
|
||||
if store.acquire(capture.token) is None:
|
||||
raise _JobStopped # every capture was invalidated (restore)
|
||||
try:
|
||||
pairs = await model_gate.run(
|
||||
owner,
|
||||
_pinned_suggestions,
|
||||
store,
|
||||
capture.token,
|
||||
service,
|
||||
capture,
|
||||
spec.model_type,
|
||||
spec.options,
|
||||
on_skipped=functools.partial(store.release_use, capture.token),
|
||||
)
|
||||
except Exception:
|
||||
self.logger.exception("Auto-tagging failed for one file")
|
||||
return "errors", INFERENCE_FAILED
|
||||
try:
|
||||
facts = await self._run_in_executor(store.verify, root, relative)
|
||||
except CaptureError as exc:
|
||||
return _capture_outcome(exc, root, relative)
|
||||
if facts.locator != capture.locator or facts.sha256 != capture.sha256:
|
||||
return "skipped", FILE_CHANGED
|
||||
names = [str(name) for name, _score in pairs]
|
||||
if not names:
|
||||
return "unchanged", NO_NEW_TAGS
|
||||
return await self._union_image_tags(facts, names, epoch)
|
||||
|
||||
async def _union_image_tags(self, facts, tags: list, epoch: str):
|
||||
"""``add_image_tags`` on the verified bytes, as a (group, detail)."""
|
||||
try:
|
||||
receipt = await self._run_in_executor(
|
||||
self.db.add_image_tags,
|
||||
locator=facts.locator,
|
||||
sha256=facts.sha256,
|
||||
size=facts.size,
|
||||
mtime_ns=facts.mtime_ns,
|
||||
expected_epoch=epoch,
|
||||
tags=tags,
|
||||
)
|
||||
except LibraryCommandError as exc:
|
||||
if exc.code == "storage-changed":
|
||||
raise _JobStopped from None
|
||||
if exc.code == "tag-limit":
|
||||
return "errors", _MESSAGES["tag-limit"]
|
||||
self.logger.warning(f"Image tags not saved ({exc.code})")
|
||||
return "errors", SAVE_FAILED
|
||||
except ValueError:
|
||||
self.logger.exception("Image tags were refused")
|
||||
return "errors", SAVE_FAILED
|
||||
added = receipt.get("added", 0)
|
||||
if added:
|
||||
return "tagged", f"Added {_plural(added, 'tag')}"
|
||||
return "unchanged", NO_NEW_TAGS
|
||||
|
||||
def _record_tag_outcome(self, job, target, group: str, detail: str) -> None:
|
||||
index, _scope, relative = target
|
||||
roots = list(self._get_all_output_dirs())
|
||||
label = _public_path(roots[index]) if index < len(roots) else ""
|
||||
where = f"{label}/{relative}" if label else relative
|
||||
job.record(
|
||||
group,
|
||||
{"path": relative, "root": label, "root_index": index, "detail": detail},
|
||||
f"{_LOG_PREFIXES[group]} {where} · {detail}",
|
||||
)
|
||||
+23
-13
@@ -1,8 +1,9 @@
|
||||
"""Session job routes (Activity): list, inspect, follow and cancel jobs.
|
||||
|
||||
The registry itself is ``py/api/jobs.py``; the scan job's runner is in
|
||||
``py/api/admin.py``. Every response leaves out absolute server paths: job
|
||||
rows already carry root-relative paths and public root labels.
|
||||
The registry itself is ``py/api/jobs.py``. The scan job's runner is in
|
||||
``py/api/admin.py``, the image tagging job's in ``py/api/image_tag_jobs.py``.
|
||||
Every response leaves out absolute server paths: job rows already carry
|
||||
root-relative paths and public root labels.
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
@@ -42,6 +43,20 @@ async def _scan_scope(request):
|
||||
return scope if scope in SCAN_SCOPES else None
|
||||
|
||||
|
||||
def busy_response(busy: JobBusy):
|
||||
"""``409 busy``, naming the running job that holds the slot."""
|
||||
running = busy.running
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "busy",
|
||||
"error": str(busy),
|
||||
"running": {"id": running.id, "title": running.title},
|
||||
},
|
||||
status=409,
|
||||
)
|
||||
|
||||
|
||||
def _summary_event(job):
|
||||
return f"data: {json.dumps(job.summary())}\n\n".encode("utf-8")
|
||||
|
||||
@@ -58,6 +73,10 @@ class JobRoutesMixin:
|
||||
async def start_scan_job_route(request):
|
||||
return await self.start_scan_job(request)
|
||||
|
||||
@routes.post("/prompt_manager/jobs/image-tags")
|
||||
async def start_image_tag_job_route(request):
|
||||
return await self.start_image_tag_job(request)
|
||||
|
||||
@routes.get("/prompt_manager/jobs/{job_id}")
|
||||
async def job_detail_route(request):
|
||||
return await self.job_detail(request)
|
||||
@@ -121,16 +140,7 @@ class JobRoutesMixin:
|
||||
try:
|
||||
job = self._start_scan_job(files)
|
||||
except JobBusy as busy:
|
||||
running = busy.running
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"code": "busy",
|
||||
"error": str(busy),
|
||||
"running": {"id": running.id, "title": running.title},
|
||||
},
|
||||
status=409,
|
||||
)
|
||||
return busy_response(busy)
|
||||
return web.json_response({"success": True, "job": job.summary()}, status=202)
|
||||
|
||||
async def job_events(self, request):
|
||||
|
||||
+22
-9
@@ -9,8 +9,9 @@ as a snapshot and then every event published after it, until the job ends.
|
||||
|
||||
Jobs live only in server memory and are gone after a restart. The registry
|
||||
keeps the running jobs and this session's last ``MAX_FINISHED_JOBS`` finished
|
||||
ones. Record-changing jobs (the metadata scan, and later maintenance, cleanup
|
||||
and restore) never run together, and one job of a kind runs at a time.
|
||||
ones. Record-changing jobs (the metadata scan, bulk image tagging, and later
|
||||
maintenance, cleanup and restore) never run together, and one job of a kind
|
||||
runs at a time.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -54,6 +55,11 @@ SCAN_GROUPS = {
|
||||
"failed": "errors",
|
||||
}
|
||||
|
||||
# Bulk image tagging (#249 D1): Add tags and Auto-tag over selected images.
|
||||
# Record-changing: it writes image tag sets and new vocabulary names.
|
||||
IMAGE_TAG_JOB_KIND = "image-tags"
|
||||
IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
|
||||
|
||||
|
||||
def _relative_path(path) -> str:
|
||||
"""``path`` as a root-relative POSIX path; an absolute one keeps its name."""
|
||||
@@ -292,13 +298,9 @@ class JobRegistry:
|
||||
running or when ``record_changing`` and another record-changing job
|
||||
runs. Must be called from the event loop.
|
||||
"""
|
||||
same_kind = self.active(kind)
|
||||
if same_kind is not None:
|
||||
raise JobBusy(same_kind)
|
||||
if record_changing:
|
||||
for job in self._jobs.values():
|
||||
if job.running and job.record_changing:
|
||||
raise JobBusy(job)
|
||||
blocker = self.blocker(kind, record_changing=record_changing)
|
||||
if blocker is not None:
|
||||
raise JobBusy(blocker)
|
||||
self._prune()
|
||||
job = Job(
|
||||
kind, title, record_changing=record_changing, scope=scope, groups=groups
|
||||
@@ -307,6 +309,17 @@ class JobRegistry:
|
||||
job.task = asyncio.ensure_future(self._run(job, runner))
|
||||
return job
|
||||
|
||||
def blocker(self, kind, *, record_changing):
|
||||
"""The running job a new ``kind`` job would be refused for, or None."""
|
||||
same_kind = self.active(kind)
|
||||
if same_kind is not None:
|
||||
return same_kind
|
||||
if record_changing:
|
||||
for job in self._jobs.values():
|
||||
if job.running and job.record_changing:
|
||||
return job
|
||||
return None
|
||||
|
||||
async def _run(self, job, runner):
|
||||
try:
|
||||
await runner(job)
|
||||
|
||||
@@ -462,6 +462,15 @@ class TestUiRoutes(CoreRoutesTestCase):
|
||||
):
|
||||
self.assertIn(entry, table)
|
||||
|
||||
async def test_image_tag_job_route_is_registered_beside_the_scan_job(self):
|
||||
table = {
|
||||
(route.method, route.resource.canonical)
|
||||
for route in self.app.router.routes()
|
||||
if route.resource is not None
|
||||
}
|
||||
self.assertIn(("POST", "/prompt_manager/jobs/scan"), table)
|
||||
self.assertIn(("POST", "/prompt_manager/jobs/image-tags"), table)
|
||||
|
||||
|
||||
class TestStaticRoutes(CoreRoutesTestCase):
|
||||
|
||||
|
||||
@@ -0,0 +1,648 @@
|
||||
"""Bulk image tagging as an Activity job: Add tags and Auto-tag (#249 D1).
|
||||
|
||||
``POST /prompt_manager/jobs/image-tags`` starts one record-changing
|
||||
``image-tags`` job over 1..500 listed images. Add tags unions a fixed set
|
||||
into each file's byte record; Auto-tag runs one legacy engine on a private
|
||||
capture of each still image (under the shared model gate, never
|
||||
downloading) and unions its top N suggestions. No real model is loaded: the
|
||||
engine is B2's fake service.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from test_api_annotation_suggest import SCORED, WD14, FakeTagger, SuggestTestCase
|
||||
from test_api_images import make_png
|
||||
except ImportError:
|
||||
from tests.test_api_annotation_suggest import (
|
||||
SCORED,
|
||||
WD14,
|
||||
FakeTagger,
|
||||
SuggestTestCase,
|
||||
)
|
||||
from tests.test_api_images import make_png
|
||||
|
||||
from database.execution_history_schema import rotate_history_epoch
|
||||
from py import model_gate
|
||||
from py.api import jobs as jobs_module
|
||||
from py.config import AutotagConfig
|
||||
|
||||
START = "/prompt_manager/jobs/image-tags"
|
||||
WAIT = 30.0
|
||||
KIND = getattr(jobs_module, "IMAGE_TAG_JOB_KIND", "image-tags")
|
||||
GROUPS = ("tagged", "unchanged", "skipped", "errors")
|
||||
|
||||
|
||||
class HookedTagger(FakeTagger):
|
||||
"""The B2 fake plus a hook that runs inside each inference."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.on_generate = None
|
||||
|
||||
def generate_tag_suggestions(self, image_path, **kwargs):
|
||||
result = super().generate_tag_suggestions(image_path, **kwargs)
|
||||
if self.on_generate is not None:
|
||||
self.on_generate(image_path)
|
||||
return result
|
||||
|
||||
|
||||
class ImageTagJobTestCase(SuggestTestCase):
|
||||
async def get_application(self):
|
||||
app = await super().get_application()
|
||||
self.tagger = HookedTagger()
|
||||
patcher = patch("py.autotag.get_autotag_service", return_value=self.tagger)
|
||||
patcher.start()
|
||||
self.addCleanup(patcher.stop)
|
||||
return app
|
||||
|
||||
async def tearDownAsync(self):
|
||||
# Every job of the test has ended before the app goes away.
|
||||
for job in self.api._jobs.list():
|
||||
if job.task is not None and not job.task.done():
|
||||
job.task.cancel()
|
||||
await asyncio.gather(job.task, return_exceptions=True)
|
||||
await super().tearDownAsync()
|
||||
|
||||
# -- helpers ---------------------------------------------------------
|
||||
|
||||
def png(self, path, size=(8, 6), color=(200, 120, 40)):
|
||||
make_png(self.output_dir / path, size=size, color=color)
|
||||
return self.output_dir / path
|
||||
|
||||
def digest(self, path):
|
||||
return hashlib.sha256((self.output_dir / path).read_bytes()).hexdigest()
|
||||
|
||||
def stored(self, path):
|
||||
return self.api.db.read_image_tags(self.locator(path), self.digest(path))
|
||||
|
||||
def add_body(self, paths, tags, **changes):
|
||||
return {
|
||||
"mode": "add",
|
||||
"targets": [self.target(p) for p in paths],
|
||||
"tags": list(tags),
|
||||
**changes,
|
||||
}
|
||||
|
||||
def auto_body(self, paths, model_type="wd14-vit", options=None, **changes):
|
||||
return {
|
||||
"mode": "auto",
|
||||
"targets": [self.target(p) for p in paths],
|
||||
"model_type": model_type,
|
||||
"options": dict(WD14) if options is None else options,
|
||||
**changes,
|
||||
}
|
||||
|
||||
async def start(self, body):
|
||||
response = await self.client.post(START, json=body)
|
||||
return response.status, await self.json_of(response)
|
||||
|
||||
async def start_ok(self, body):
|
||||
status, data = await self.start(body)
|
||||
self.assertEqual(status, 202, data)
|
||||
return self.api._jobs.get(data["job"]["id"])
|
||||
|
||||
async def run_job(self, body):
|
||||
job = await self.start_ok(body)
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
return job
|
||||
|
||||
def rows(self, job, group):
|
||||
return [(row["path"], row["detail"]) for row in job.outcomes[group]]
|
||||
|
||||
async def wait_for(self, event):
|
||||
ok = await asyncio.get_running_loop().run_in_executor(None, event.wait, WAIT)
|
||||
self.assertTrue(ok, "the fake engine was never entered")
|
||||
|
||||
|
||||
class TestAddTags(ImageTagJobTestCase):
|
||||
async def test_add_job_unions_tags_on_each_target_and_records_outcomes(self):
|
||||
for name in ("a.png", "b.png", "c.png"):
|
||||
self.png(name, color=(len(name), ord(name[0]), 9))
|
||||
self.saved_set("a.png", (self.output_dir / "a.png").read_bytes(), ["keep_me"])
|
||||
self.saved_set(
|
||||
"c.png", (self.output_dir / "c.png").read_bytes(), ["Sandstone", "desert"]
|
||||
)
|
||||
|
||||
job = await self.run_job(
|
||||
self.add_body(
|
||||
["a.png", "b.png", "c.png", "gone.png"], ["sandstone", "desert"]
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(job.kind, KIND)
|
||||
self.assertEqual(job.title, "Add tags to images")
|
||||
self.assertEqual(job.state, "finished")
|
||||
self.assertTrue(job.record_changing)
|
||||
self.assertEqual(job.scope, {"mode": "add", "images": 4})
|
||||
self.assertEqual(
|
||||
job.counters, {"tagged": 2, "unchanged": 1, "skipped": 1, "errors": 0}
|
||||
)
|
||||
self.assertEqual(
|
||||
self.stored("a.png")["tags"], ["keep_me", "sandstone", "desert"]
|
||||
)
|
||||
self.assertEqual(self.stored("b.png")["tags"], ["sandstone", "desert"])
|
||||
# The stored spelling wins; nothing new means no write.
|
||||
self.assertEqual(self.stored("c.png")["tags"], ["Sandstone", "desert"])
|
||||
self.assertEqual(self.stored("c.png")["tag_revision"], 1)
|
||||
self.assertEqual(
|
||||
self.rows(job, "tagged"),
|
||||
[("a.png", "Added 2 tags"), ("b.png", "Added 2 tags")],
|
||||
)
|
||||
self.assertEqual(self.rows(job, "unchanged"), [("c.png", "No new tags")])
|
||||
self.assertEqual(self.rows(job, "skipped"), [("gone.png", "File not found")])
|
||||
self.assertEqual(job.outcomes["tagged"][0]["root_index"], 0)
|
||||
self.assertEqual((job.processed, job.total, job.progress), (4, 4, 100))
|
||||
self.assertTrue(any("+ " in line and "a.png" in line for line in job.log))
|
||||
self.assertTrue(job.log[-1].endswith("Add tags finished · 4 files processed"))
|
||||
# Add tags never needs the model.
|
||||
self.assertEqual(self.tagger.load_calls, [])
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_add_tags_vocabulary_names_become_global_tags(self):
|
||||
self.png("a.png")
|
||||
|
||||
await self.run_job(self.add_body(["a.png"], ["brand_new_tag"]))
|
||||
|
||||
names = set(self.api.db.get_all_tags())
|
||||
self.assertIn("brand_new_tag", names)
|
||||
|
||||
async def test_add_over_the_tag_cap_is_an_error_row_with_zero_writes(self):
|
||||
self.png("a.png")
|
||||
full = [f"t{i:03d}" for i in range(100)]
|
||||
self.saved_set("a.png", (self.output_dir / "a.png").read_bytes(), full)
|
||||
|
||||
job = await self.run_job(self.add_body(["a.png"], ["one_more"]))
|
||||
|
||||
self.assertEqual(job.counters["errors"], 1)
|
||||
self.assertEqual(
|
||||
self.rows(job, "errors"),
|
||||
[("a.png", "An image can carry at most 100 tags.")],
|
||||
)
|
||||
self.assertEqual(self.stored("a.png")["tags"], full)
|
||||
|
||||
async def test_video_skipped_in_auto_and_tagged_in_add(self):
|
||||
(self.output_dir / "clip.mp4").write_bytes(b"\x00\x00\x00\x18ftypmp42" * 8)
|
||||
|
||||
added = await self.run_job(self.add_body(["clip.mp4"], ["motion"]))
|
||||
self.assertEqual(self.rows(added, "tagged"), [("clip.mp4", "Added 1 tag")])
|
||||
self.assertEqual(self.stored("clip.mp4")["tags"], ["motion"])
|
||||
|
||||
auto = await self.run_job(self.auto_body(["clip.mp4"]))
|
||||
self.assertEqual(
|
||||
self.rows(auto, "skipped"), [("clip.mp4", "Video: manual tags only")]
|
||||
)
|
||||
self.assertEqual(self.tagger.calls, [])
|
||||
self.assertEqual(self.stored("clip.mp4")["tags"], ["motion"])
|
||||
|
||||
async def test_not_a_still_image_and_too_large_are_skipped(self):
|
||||
(self.output_dir / "broken.png").write_bytes(b"not an image")
|
||||
self.png("big.png")
|
||||
self.api._capture_store().max_file = 16
|
||||
|
||||
auto = await self.run_job(self.auto_body(["big.png"]))
|
||||
self.assertEqual(
|
||||
self.rows(auto, "skipped"), [("big.png", "Too large to verify")]
|
||||
)
|
||||
|
||||
self.api._capture_store().max_file = 2**20
|
||||
auto = await self.run_job(self.auto_body(["broken.png"]))
|
||||
self.assertEqual(
|
||||
self.rows(auto, "skipped"), [("broken.png", "Not a supported still image")]
|
||||
)
|
||||
|
||||
self.api._capture_store().max_file = 16
|
||||
added = await self.run_job(self.add_body(["big.png"], ["x"]))
|
||||
self.assertEqual(
|
||||
self.rows(added, "skipped"), [("big.png", "Too large to verify")]
|
||||
)
|
||||
self.assertEqual(self.tagger.calls, [])
|
||||
|
||||
|
||||
class TestAutoTag(ImageTagJobTestCase):
|
||||
async def test_auto_job_applies_the_requested_count_per_file_without_removing_existing( # noqa: E501
|
||||
self,
|
||||
):
|
||||
self.png("a.png")
|
||||
self.png("b.png", color=(1, 2, 3))
|
||||
self.saved_set("a.png", (self.output_dir / "a.png").read_bytes(), ["keep_me"])
|
||||
options = {**WD14, "suggestion_limit": 2}
|
||||
|
||||
job = await self.run_job(self.auto_body(["a.png", "b.png"], options=options))
|
||||
|
||||
self.assertEqual(job.title, "Auto-tag images")
|
||||
self.assertEqual(
|
||||
job.scope, {"mode": "auto", "images": 2, "model_type": "wd14-vit"}
|
||||
)
|
||||
self.assertEqual(job.state, "finished")
|
||||
self.assertEqual(
|
||||
self.api.db.read_image_tags(self.locator("a.png"), self.digest("a.png"))[
|
||||
"tags"
|
||||
],
|
||||
["keep_me", "sandstone", "desert"],
|
||||
)
|
||||
self.assertEqual(self.stored("b.png")["tags"], ["sandstone", "desert"])
|
||||
self.assertEqual(
|
||||
job.counters, {"tagged": 2, "unchanged": 0, "skipped": 0, "errors": 0}
|
||||
)
|
||||
# Inference read the private captures only, never the current files.
|
||||
for path in self.tagger.read_paths:
|
||||
self.assertNotIn(str(self.output_dir), path)
|
||||
self.assertEqual([call["limit"] for call in self.tagger.calls], [2, 2])
|
||||
|
||||
# limit 3: at most three names are added to a file.
|
||||
self.png("c.png", color=(9, 9, 9))
|
||||
options = {**WD14, "suggestion_limit": 3}
|
||||
job = await self.run_job(self.auto_body(["c.png"], options=options))
|
||||
self.assertEqual(self.stored("c.png")["tags"], [name for name, _ in SCORED[:3]])
|
||||
self.assertEqual(self.rows(job, "tagged"), [("c.png", "Added 3 tags")])
|
||||
|
||||
# Omitted: the configured default.
|
||||
AutotagConfig.SUGGESTION_LIMIT = 4
|
||||
self.png("d.png", color=(7, 7, 7))
|
||||
await self.run_job(self.auto_body(["d.png"]))
|
||||
self.assertEqual(self.tagger.calls[-1]["limit"], 4)
|
||||
self.assertEqual(len(self.stored("d.png")["tags"]), 4)
|
||||
|
||||
# Suggestions the file already carries add nothing.
|
||||
job = await self.run_job(self.auto_body(["d.png"]))
|
||||
self.assertEqual(self.rows(job, "unchanged"), [("d.png", "No new tags")])
|
||||
|
||||
async def test_auto_job_holds_gate_and_never_downloads(self):
|
||||
self.png("a.png")
|
||||
self.png("b.png", color=(1, 2, 3))
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
|
||||
job = await self.start_ok(self.auto_body(["a.png", "b.png"]))
|
||||
await self.wait_for(entered)
|
||||
holder = model_gate.holder()
|
||||
self.assertIsNotNone(holder)
|
||||
self.assertTrue(holder.startswith("image-tag-job:"), holder)
|
||||
# A legacy route cannot swap the model under the running job.
|
||||
response = await self.client.post("/prompt_manager/autotag/unload")
|
||||
self.assertEqual(response.status, 409)
|
||||
self.assertEqual((await response.json())["code"], "model-busy")
|
||||
proceed.set()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
|
||||
self.assertEqual(job.counters["tagged"], 2)
|
||||
self.assertEqual(self.tagger.load_calls, [("wd14-vit", True)])
|
||||
self.tagger.download_model.assert_not_called()
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_model_unavailable_fails_the_job_before_the_first_file(self):
|
||||
self.png("a.png")
|
||||
self.tagger.downloaded["wd14-vit"] = False
|
||||
|
||||
job = await self.run_job(self.auto_body(["a.png"]))
|
||||
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.processed, 0)
|
||||
self.assertEqual(job.counters, dict.fromkeys(GROUPS, 0))
|
||||
self.assertTrue(
|
||||
any(
|
||||
line.endswith(
|
||||
"is not downloaded. Download it from the legacy Auto-tag page."
|
||||
)
|
||||
for line in job.log
|
||||
),
|
||||
job.log,
|
||||
)
|
||||
self.assertEqual(self.tagger.load_calls, [])
|
||||
self.tagger.download_model.assert_not_called()
|
||||
self.assertIsNone(model_gate.holder())
|
||||
self.assertEqual(self.api._capture_store().stats()["captures"], 0)
|
||||
|
||||
async def test_model_load_failure_fails_the_job_and_frees_the_gate(self):
|
||||
self.png("a.png")
|
||||
self.tagger.load_error = RuntimeError("/secret/models/wd14 exploded")
|
||||
|
||||
job = await self.run_job(self.auto_body(["a.png"]))
|
||||
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.processed, 0)
|
||||
self.assertNotIn("/secret", json.dumps(job.detail()))
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_inference_failure_is_an_error_row_and_the_job_goes_on(self):
|
||||
self.png("a.png")
|
||||
self.png("b.png", color=(1, 2, 3))
|
||||
failing = {"first": True}
|
||||
|
||||
def explode(_path):
|
||||
if failing.pop("first", False):
|
||||
raise RuntimeError("/secret/path boom")
|
||||
|
||||
self.tagger.on_generate = explode
|
||||
|
||||
job = await self.run_job(self.auto_body(["a.png", "b.png"]))
|
||||
|
||||
self.assertEqual(job.counters["errors"], 1)
|
||||
self.assertEqual(job.counters["tagged"], 1)
|
||||
self.assertEqual(
|
||||
self.rows(job, "errors"),
|
||||
[("a.png", "Auto-tagging failed. Check server logs for details.")],
|
||||
)
|
||||
self.assertNotIn("/secret", json.dumps(job.report()))
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_changed_file_is_skipped_not_tagged(self):
|
||||
self.png("a.png")
|
||||
before = self.digest("a.png")
|
||||
|
||||
def rewrite(_capture_path):
|
||||
make_png(self.output_dir / "a.png", size=(8, 6), color=(1, 1, 1))
|
||||
|
||||
self.tagger.on_generate = rewrite
|
||||
|
||||
job = await self.run_job(self.auto_body(["a.png"]))
|
||||
|
||||
self.assertEqual(self.rows(job, "skipped"), [("a.png", "File changed")])
|
||||
self.assertEqual(job.counters["tagged"], 0)
|
||||
locator = self.locator("a.png")
|
||||
self.assertEqual(self.api.db.read_image_tags(locator, before)["tags"], [])
|
||||
self.assertEqual(self.stored("a.png")["tags"], [])
|
||||
|
||||
async def test_captures_released_and_handles_closed(self):
|
||||
try:
|
||||
from open_handles import assert_closed
|
||||
except ImportError:
|
||||
from tests.open_handles import assert_closed
|
||||
for index in range(3):
|
||||
self.png(f"f{index}.png", color=(index, 0, 0))
|
||||
|
||||
await self.run_job(self.auto_body([f"f{index}.png" for index in range(3)]))
|
||||
|
||||
store = self.api._capture_store()
|
||||
self.assertEqual(store.stats()["captures"], 0)
|
||||
assert_closed(self, store.directory)
|
||||
assert_closed(self, str(self.output_dir))
|
||||
|
||||
|
||||
class TestStartAndExclusion(ImageTagJobTestCase):
|
||||
async def test_start_returns_model_busy_or_job_busy(self):
|
||||
self.png("a.png")
|
||||
other = model_gate.new_owner("legacy-batch")
|
||||
self.assertTrue(model_gate.try_acquire(other))
|
||||
try:
|
||||
status, body = await self.start(self.auto_body(["a.png"]))
|
||||
self.assertEqual((status, body["code"]), (409, "model-busy"))
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
# Add tags needs no model: it starts while the gate is held.
|
||||
job = await self.run_job(self.add_body(["a.png"], ["x"]))
|
||||
self.assertEqual(job.state, "finished")
|
||||
finally:
|
||||
model_gate.release(other)
|
||||
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
job = await self.start_ok(self.auto_body(["a.png"]))
|
||||
await self.wait_for(entered)
|
||||
try:
|
||||
for body in (self.add_body(["a.png"], ["y"]), self.auto_body(["a.png"])):
|
||||
resp = await self.client.post(START, json=body)
|
||||
body = await resp.json()
|
||||
self.assertEqual((resp.status, body["code"]), (409, "busy"))
|
||||
self.assertEqual(
|
||||
body["running"], {"id": job.id, "title": "Auto-tag images"}
|
||||
)
|
||||
finally:
|
||||
proceed.set()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
# A refused start never leaves the gate taken.
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_record_changing_excludes_running_scan(self):
|
||||
self.png("a.png")
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def scan_runner(job):
|
||||
await gate.wait()
|
||||
|
||||
scan = self.api._jobs.start(
|
||||
"scan",
|
||||
"Scan for metadata",
|
||||
record_changing=True,
|
||||
scope={"files": "new"},
|
||||
runner=scan_runner,
|
||||
)
|
||||
try:
|
||||
status, body = await self.start(self.add_body(["a.png"], ["x"]))
|
||||
self.assertEqual((status, body["code"]), (409, "busy"))
|
||||
self.assertEqual(body["running"]["title"], "Scan for metadata")
|
||||
status, body = await self.start(self.auto_body(["a.png"]))
|
||||
self.assertEqual((status, body["code"]), (409, "busy"))
|
||||
self.assertIsNone(model_gate.holder())
|
||||
finally:
|
||||
gate.set()
|
||||
await scan.task
|
||||
|
||||
# The other way round: a running tagging job keeps the scan out.
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
job = await self.start_ok(self.auto_body(["a.png"]))
|
||||
await self.wait_for(entered)
|
||||
try:
|
||||
resp = await self.client.post("/prompt_manager/jobs/scan", json={})
|
||||
self.assertEqual(resp.status, 409)
|
||||
self.assertEqual((await resp.json())["running"]["id"], job.id)
|
||||
finally:
|
||||
proceed.set()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
|
||||
async def test_a_start_that_loses_the_slot_race_frees_the_gate(self):
|
||||
self.png("a.png")
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def scan_runner(job):
|
||||
await gate.wait()
|
||||
|
||||
scan = self.api._jobs.start(
|
||||
"scan",
|
||||
"Scan for metadata",
|
||||
record_changing=True,
|
||||
scope={},
|
||||
runner=scan_runner,
|
||||
)
|
||||
try:
|
||||
# The pre-check saw a free slot; registry.start still refuses.
|
||||
with patch.object(self.api._jobs, "blocker", side_effect=[None, scan]):
|
||||
status, body = await self.start(self.auto_body(["a.png"]))
|
||||
self.assertEqual((status, body["code"]), (409, "busy"))
|
||||
self.assertIsNone(model_gate.holder())
|
||||
finally:
|
||||
gate.set()
|
||||
await scan.task
|
||||
|
||||
async def test_cancel_stops_at_file_boundary_and_keeps_written_files(self):
|
||||
for index in range(3):
|
||||
self.png(f"f{index}.png", color=(index, 5, 5))
|
||||
entered, proceed = threading.Event(), threading.Event()
|
||||
self.tagger.block = (entered, proceed)
|
||||
|
||||
job = await self.start_ok(self.auto_body([f"f{i}.png" for i in range(3)]))
|
||||
await self.wait_for(entered)
|
||||
response = await self.client.post(f"/prompt_manager/jobs/{job.id}/cancel")
|
||||
self.assertTrue((await response.json())["cancellation_requested"])
|
||||
proceed.set()
|
||||
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
||||
|
||||
self.assertEqual(job.state, "cancelled")
|
||||
self.assertEqual(job.counters["tagged"], 1)
|
||||
self.assertEqual(job.processed, 1)
|
||||
self.assertEqual(len(self.stored("f0.png")["tags"]), 10)
|
||||
self.assertEqual(self.stored("f1.png")["tags"], [])
|
||||
self.assertEqual(len(self.tagger.calls), 1)
|
||||
self.assertTrue(job.log[-1].endswith("Auto-tag cancelled · 1 file processed"))
|
||||
self.assertIsNone(model_gate.holder())
|
||||
self.assertEqual(self.api._capture_store().stats()["captures"], 0)
|
||||
|
||||
async def test_stale_root_scope_refuses_the_start(self):
|
||||
self.png("a.png")
|
||||
body = self.add_body(["a.png"], ["x"])
|
||||
body["targets"][0]["root_scope"] = "0" * 64
|
||||
|
||||
status, data = await self.start(body)
|
||||
|
||||
self.assertEqual((status, data["code"]), (409, "source-changed"))
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
|
||||
async def test_restore_during_the_job_stops_it_without_writes(self):
|
||||
self.png("a.png")
|
||||
self.png("b.png", color=(1, 2, 3))
|
||||
|
||||
def rotate(_path):
|
||||
conn = sqlite3.connect(self._temp_db.name)
|
||||
try:
|
||||
rotate_history_epoch(conn)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
self.tagger.on_generate = rotate
|
||||
|
||||
job = await self.run_job(self.auto_body(["a.png", "b.png"]))
|
||||
|
||||
self.assertEqual(job.state, "failed")
|
||||
self.assertEqual(job.counters["tagged"], 0)
|
||||
self.assertEqual(
|
||||
self.rows(job, "errors"),
|
||||
[("a.png", "Prompt storage was restored; start the job again.")],
|
||||
)
|
||||
self.assertEqual(self.stored("a.png")["tags"], [])
|
||||
self.assertEqual(len(self.tagger.calls), 1)
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_invalid_bodies_rejected(self):
|
||||
self.png("a.png")
|
||||
target = self.target("a.png")
|
||||
bad_bodies = [
|
||||
{"mode": "add", "targets": [], "tags": ["x"]},
|
||||
{"mode": "add", "targets": [target] * 501, "tags": ["x"]},
|
||||
{"mode": "add", "targets": "a.png", "tags": ["x"]},
|
||||
{"mode": "add", "targets": [{**target, "path": "../a.png"}], "tags": ["x"]},
|
||||
{"mode": "add", "targets": [{**target, "root_index": True}], "tags": ["x"]},
|
||||
{"mode": "add", "targets": [{**target, "extra": 1}], "tags": ["x"]},
|
||||
{"mode": "add", "targets": [{**target, "path": "a.txt"}], "tags": ["x"]},
|
||||
{"mode": "add", "targets": [target], "tags": []},
|
||||
{"mode": "add", "targets": [target], "tags": "x"},
|
||||
{"mode": "add", "targets": [target], "tags": ["bad\nname"]},
|
||||
{"mode": "add", "targets": [target]},
|
||||
{"mode": "add", "targets": [target], "tags": ["x"], "extra": 1},
|
||||
{"mode": "add", "targets": [target], "tags": ["x"], "model_type": "hf"},
|
||||
{"mode": "remove", "targets": [target], "tags": ["x"]},
|
||||
{"targets": [target], "tags": ["x"]},
|
||||
{"mode": "auto", "targets": [target], "model_type": "wd14-vit"},
|
||||
{"mode": "auto", "targets": [target], "model_type": "x", "options": {}},
|
||||
{
|
||||
"mode": "auto",
|
||||
"targets": [target],
|
||||
"model_type": "wd14-vit",
|
||||
"options": {},
|
||||
},
|
||||
{
|
||||
"mode": "auto",
|
||||
"targets": [target],
|
||||
"model_type": "wd14-vit",
|
||||
"options": {**WD14, "general_threshold": 2},
|
||||
},
|
||||
{
|
||||
"mode": "auto",
|
||||
"targets": [target],
|
||||
"model_type": "wd14-vit",
|
||||
"options": {**WD14, "suggestion_limit": 0},
|
||||
},
|
||||
{
|
||||
"mode": "auto",
|
||||
"targets": [target],
|
||||
"model_type": "gguf",
|
||||
"options": dict(WD14),
|
||||
},
|
||||
{
|
||||
"mode": "auto",
|
||||
"targets": [target],
|
||||
"model_type": "wd14-vit",
|
||||
"options": dict(WD14),
|
||||
"tags": ["x"],
|
||||
},
|
||||
]
|
||||
for body in bad_bodies:
|
||||
with self.subTest(body=json.dumps(body)[:120]):
|
||||
status, data = await self.start(body)
|
||||
self.assertIn(status, (400, 403), data)
|
||||
self.assertFalse(data["success"])
|
||||
too_many = {
|
||||
"mode": "add",
|
||||
"targets": [target],
|
||||
"tags": [f"t{i}" for i in range(101)],
|
||||
}
|
||||
status, data = await self.start(too_many)
|
||||
self.assertEqual((status, data["code"]), (400, "tag-limit"))
|
||||
response = await self.client.post(START, data="not json")
|
||||
self.assertEqual(response.status, 400)
|
||||
self.assertEqual(self.api._jobs.list(), [])
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_duplicate_targets_are_tagged_once(self):
|
||||
self.png("a.png")
|
||||
|
||||
job = await self.run_job(self.add_body(["a.png", "a.png"], ["x"]))
|
||||
|
||||
self.assertEqual(job.scope["images"], 1)
|
||||
self.assertEqual(job.counters["tagged"], 1)
|
||||
|
||||
async def test_rows_and_report_carry_no_absolute_paths(self):
|
||||
self.png("sub/a.png")
|
||||
self.png("sub/b.png", color=(3, 3, 3))
|
||||
job = await self.run_job(
|
||||
self.auto_body(["sub/a.png", "sub/b.png", "sub/missing.png"])
|
||||
)
|
||||
|
||||
for url in (
|
||||
START.rsplit("/", 1)[0],
|
||||
f"/prompt_manager/jobs/{job.id}",
|
||||
f"/prompt_manager/jobs/{job.id}/report",
|
||||
):
|
||||
response = await self.client.get(url)
|
||||
self.assertEqual(response.status, 200)
|
||||
body = json.loads(await response.text())
|
||||
self.assert_no_paths(body)
|
||||
detail = job.detail()
|
||||
self.assertEqual(
|
||||
[row["path"] for row in detail["outcomes"]["tagged"]],
|
||||
["sub/a.png", "sub/b.png"],
|
||||
)
|
||||
self.assertTrue(
|
||||
all(row["root"] == "output" for row in detail["outcomes"]["tagged"])
|
||||
)
|
||||
self.assertTrue(any("output/sub/a.png" in line for line in detail["log"]))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -12,6 +12,7 @@ import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from py.api import jobs as jobs_module # noqa: E402
|
||||
from py.api.jobs import ( # noqa: E402
|
||||
DEFAULT_SCAN_SCOPE,
|
||||
LOG_CAP_LINE,
|
||||
@@ -415,6 +416,14 @@ class TestJobOutcomesAndLog(unittest.TestCase):
|
||||
SCAN_GROUP_NAMES, ("added", "known", "review", "no_metadata", "errors")
|
||||
)
|
||||
|
||||
def test_image_tag_job_kind_and_groups(self):
|
||||
# #249 D1: bulk Add tags / Auto-tag, counter-tile order.
|
||||
self.assertEqual(getattr(jobs_module, "IMAGE_TAG_JOB_KIND", None), "image-tags")
|
||||
self.assertEqual(
|
||||
getattr(jobs_module, "IMAGE_TAG_GROUP_NAMES", None),
|
||||
("tagged", "unchanged", "skipped", "errors"),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user