fix(jobs): refuse legacy record-changing operations while a job runs
The job registry gains admission classes. A records job is refused while any records or guarded job runs, a guarded job while a records job runs, and read jobs share only the one-job-per-kind rule. record_changing=True stays as the spelling of records. Job summaries add cancellable and admission, and a non-cancellable job ignores cancel. JobRegistry.hold() registers a running, non-cancellable job for a legacy-page operation and ends it when the request does. The legacy restore, maintenance, prompt cleanup, duplicate image delete, LoRA import and batch auto-tag routes hold as records; thumbnail generation (both routes) and thumbnail clear hold as guarded. A maintenance run of only read checks holds as read. So a restore during a scan, or a scan during a legacy restore, is refused with the running job named instead of interleaving; this closes the gap noted on #269. Refusals keep each route's shape: the JSON routes return the legacy 409 "<title> already running", and the thumbnail and LoRA streams send their existing error event.
This commit is contained in:
+20
-1
@@ -34,7 +34,7 @@ from .images import ImageRoutesMixin
|
||||
from .admin import AdminRoutesMixin
|
||||
from .job_routes import JobRoutesMixin
|
||||
from .image_tag_jobs import ImageTagJobsMixin
|
||||
from .jobs import JobRegistry
|
||||
from .jobs import JobBusy, JobRegistry
|
||||
from .logging_routes import LoggingRoutesMixin
|
||||
from .autotag_routes import AutotagRoutesMixin
|
||||
from .lora_integration import LoraIntegrationMixin
|
||||
@@ -361,6 +361,25 @@ class PromptManagerAPI(
|
||||
"""409 envelope for a long job that is already running."""
|
||||
return _json_error(f"{label} already running", 409)
|
||||
|
||||
async def _run_held(self, kind, title, admission, operation):
|
||||
"""Run a legacy-page operation under a registry hold (D80).
|
||||
|
||||
While ``operation()`` runs, the registry lists it as a running,
|
||||
non-cancellable job, so workspace jobs its admission class conflicts
|
||||
with are refused, naming it. When the registry refuses the hold, the
|
||||
legacy 409 envelope names the running job and nothing is done. A
|
||||
response of 400 or above ends the hold ``failed``.
|
||||
"""
|
||||
try:
|
||||
held = self._jobs.hold(kind, title, admission=admission)
|
||||
except JobBusy as busy:
|
||||
return self._long_job_busy_response(busy.running.title)
|
||||
async with held as job:
|
||||
response = await operation()
|
||||
if getattr(response, "status", 200) >= 400:
|
||||
job.finish("failed")
|
||||
return response
|
||||
|
||||
async def _run_in_executor(self, func, *args, **kwargs):
|
||||
"""Run a blocking function in the default thread pool executor.
|
||||
|
||||
|
||||
+67
-7
@@ -22,6 +22,9 @@ from .jobs import (
|
||||
SCAN_GROUP_NAMES,
|
||||
SCAN_GROUPS,
|
||||
SCAN_SCOPES,
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
MAINTENANCE_JOB_KIND,
|
||||
RESTORE_JOB_KIND,
|
||||
Job,
|
||||
JobBusy,
|
||||
)
|
||||
@@ -696,7 +699,15 @@ class AdminRoutesMixin:
|
||||
job.finish()
|
||||
|
||||
async def cleanup_duplicates_endpoint(self, request):
|
||||
"""Cleanup duplicate prompts endpoint."""
|
||||
"""Cleanup duplicate prompts endpoint (holds the registry: records)."""
|
||||
return await self._run_held(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
"Prompt cleanup (admin page)",
|
||||
"records",
|
||||
self._cleanup_duplicates,
|
||||
)
|
||||
|
||||
async def _cleanup_duplicates(self):
|
||||
try:
|
||||
removed_count = await self._run_in_executor(self.db.cleanup_duplicates)
|
||||
|
||||
@@ -811,12 +822,26 @@ class AdminRoutesMixin:
|
||||
return hash_sha256.hexdigest()
|
||||
|
||||
async def delete_duplicate_images_endpoint(self, request):
|
||||
"""Delete duplicate image files from disk (single flight with the scan)."""
|
||||
job = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
|
||||
if job is None:
|
||||
"""Delete duplicate image files from disk.
|
||||
|
||||
Single flight with the duplicate scan, and a records hold on the
|
||||
registry, so no scan, restore or tagging job runs meanwhile.
|
||||
"""
|
||||
guard = self._acquire_long_job(self.LONG_JOB_DUPLICATES)
|
||||
if guard is None:
|
||||
return self._long_job_busy_response("Duplicate scan or delete")
|
||||
async with job:
|
||||
return await self._delete_duplicate_images(request)
|
||||
|
||||
async def delete():
|
||||
# No await since the check above: the lock is still free.
|
||||
async with guard:
|
||||
return await self._delete_duplicate_images(request)
|
||||
|
||||
return await self._run_held(
|
||||
DUPLICATE_CLEANUP_JOB_KIND,
|
||||
"Duplicate image delete (gallery page)",
|
||||
"records",
|
||||
delete,
|
||||
)
|
||||
|
||||
async def _delete_duplicate_images(self, request):
|
||||
try:
|
||||
@@ -1268,6 +1293,11 @@ class AdminRoutesMixin:
|
||||
"vacuum",
|
||||
"cleanup_orphaned_images",
|
||||
)
|
||||
# Operations that only read; a run of only these holds the registry as
|
||||
# "read", every other run as "records".
|
||||
READ_ONLY_MAINTENANCE_OPERATIONS = frozenset(
|
||||
{"check_hash_duplicates", "statistics", "check_consistency"}
|
||||
)
|
||||
|
||||
async def run_maintenance(self, request):
|
||||
"""Perform database maintenance operations and report each outcome."""
|
||||
@@ -1280,7 +1310,27 @@ class AdminRoutesMixin:
|
||||
operations = data.get(
|
||||
"operations", list(self.DEFAULT_MAINTENANCE_OPERATIONS)
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Maintenance error: {e}", exc_info=True)
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Maintenance failed"}, status=500
|
||||
)
|
||||
writes = not (
|
||||
isinstance(operations, list)
|
||||
and all(
|
||||
isinstance(name, str) and name in self.READ_ONLY_MAINTENANCE_OPERATIONS
|
||||
for name in operations
|
||||
)
|
||||
)
|
||||
return await self._run_held(
|
||||
MAINTENANCE_JOB_KIND,
|
||||
"Maintenance (admin page)",
|
||||
"records" if writes else "read",
|
||||
lambda: self._run_maintenance(operations),
|
||||
)
|
||||
|
||||
async def _run_maintenance(self, operations):
|
||||
try:
|
||||
results = await self._run_in_executor(
|
||||
self._run_maintenance_operations, operations
|
||||
)
|
||||
@@ -1434,8 +1484,18 @@ class AdminRoutesMixin:
|
||||
|
||||
The upload is streamed to a temp file (capped at restore_max_bytes),
|
||||
verified with PromptModel.verify_database_file and only then copied
|
||||
over the live database, which is backed up first.
|
||||
over the live database, which is backed up first. The whole request
|
||||
holds the registry (records): it is refused while a scan or another
|
||||
record-changing job runs, and those are refused while it runs.
|
||||
"""
|
||||
return await self._run_held(
|
||||
RESTORE_JOB_KIND,
|
||||
"Database restore (admin page)",
|
||||
"records",
|
||||
lambda: self._restore_database(request),
|
||||
)
|
||||
|
||||
async def _restore_database(self, request):
|
||||
max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024)
|
||||
temp_path = None
|
||||
try:
|
||||
|
||||
@@ -8,6 +8,7 @@ from pathlib import Path
|
||||
from typing import Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from .jobs import AUTOTAG_BATCH_JOB_KIND, JobBusy
|
||||
from .prompts import thumbnail_url_for
|
||||
from .. import model_gate
|
||||
|
||||
@@ -440,6 +441,27 @@ class AutotagRoutesMixin:
|
||||
"message": _INTERNAL_ERROR,
|
||||
}
|
||||
|
||||
# The batch writes prompt tags: a records hold on the registry, so it
|
||||
# never runs beside a scan, restore or the workspace tagging job.
|
||||
try:
|
||||
held = self._jobs.hold(
|
||||
AUTOTAG_BATCH_JOB_KIND,
|
||||
"Auto-tag images (legacy page)",
|
||||
admission="records",
|
||||
)
|
||||
except JobBusy as busy:
|
||||
if busy.running.kind == AUTOTAG_BATCH_JOB_KIND: # unchanged payload
|
||||
return web.json_response(
|
||||
{"success": False, "error": _BUSY_ERROR}, status=409
|
||||
)
|
||||
return self._long_job_busy_response(busy.running.title)
|
||||
async with held as job:
|
||||
response = await self._start_autotag_batch(request, opts, events)
|
||||
if response.status >= 400:
|
||||
job.finish("failed")
|
||||
return response
|
||||
|
||||
async def _start_autotag_batch(self, request, opts, events):
|
||||
if not _autotag_batch_lock.acquire(blocking=False):
|
||||
return web.json_response(
|
||||
{"success": False, "error": _BUSY_ERROR}, status=409
|
||||
|
||||
+47
-5
@@ -19,6 +19,7 @@ from typing import Callable, NamedTuple
|
||||
from aiohttp import web
|
||||
from PIL import Image, UnidentifiedImageError
|
||||
|
||||
from .jobs import THUMBNAILS_JOB_KIND, JobBusy
|
||||
from .prompts import bad_request, parse_page_params, publish_image_paths
|
||||
from .prompts import thumbnail_url_for
|
||||
from .prompts import safe_error_message as _safe_error
|
||||
@@ -104,6 +105,10 @@ def _card_header(path: Path) -> tuple:
|
||||
return inspection.width, inspection.height, records
|
||||
|
||||
|
||||
# The legacy gallery page's thumbnail generation, as Activity lists it (D80).
|
||||
THUMBNAIL_GENERATION_TITLE = "Thumbnail generation (gallery page)"
|
||||
|
||||
|
||||
class _CardFacts(NamedTuple):
|
||||
"""One image file's card facts, derived from its header alone.
|
||||
|
||||
@@ -1736,7 +1741,19 @@ class ImageRoutesMixin:
|
||||
return cls.THUMBNAIL_SIZES.get(quality, cls.THUMBNAIL_SIZES["medium"])
|
||||
|
||||
async def generate_thumbnails(self, request):
|
||||
"""Generate every thumbnail in one executor job (POST, blocking)."""
|
||||
"""Generate every thumbnail in one executor job (POST, blocking).
|
||||
|
||||
Holds the registry as guarded work: refused while a record-changing
|
||||
job runs, and such jobs are refused while it runs.
|
||||
"""
|
||||
return await self._run_held(
|
||||
THUMBNAILS_JOB_KIND,
|
||||
THUMBNAIL_GENERATION_TITLE,
|
||||
"guarded",
|
||||
lambda: self._generate_thumbnails(request),
|
||||
)
|
||||
|
||||
async def _generate_thumbnails(self, request):
|
||||
try:
|
||||
data, error = await _json_body_or_error(request)
|
||||
if error is not None:
|
||||
@@ -1935,14 +1952,31 @@ class ImageRoutesMixin:
|
||||
"Connection": "keep-alive",
|
||||
},
|
||||
)
|
||||
await response.prepare(request)
|
||||
try:
|
||||
held = self._jobs.hold(
|
||||
THUMBNAILS_JOB_KIND, THUMBNAIL_GENERATION_TITLE, admission="guarded"
|
||||
)
|
||||
except JobBusy as busy:
|
||||
# The gallery reads this stream: refuse with its error event.
|
||||
await response.prepare(request)
|
||||
message = f"{busy.running.title} already running"
|
||||
await self._emit_sse(
|
||||
response, "error", {"error": message, "message": message}
|
||||
)
|
||||
return response
|
||||
async with held:
|
||||
await response.prepare(request)
|
||||
await self._thumbnail_progress_stream(response, thumbnail_size)
|
||||
return response
|
||||
|
||||
async def _thumbnail_progress_stream(self, response, thumbnail_size):
|
||||
try:
|
||||
output_dirs = self._get_all_output_dirs()
|
||||
if not output_dirs:
|
||||
await self._emit_sse(
|
||||
response, "error", {"error": "ComfyUI output directory not found"}
|
||||
)
|
||||
return response
|
||||
return
|
||||
|
||||
await self._emit_sse(
|
||||
response,
|
||||
@@ -1970,7 +2004,6 @@ class ImageRoutesMixin:
|
||||
"message": "Thumbnail generation failed. Check server logs.",
|
||||
},
|
||||
)
|
||||
return response
|
||||
|
||||
async def _stream_thumbnails(
|
||||
self, response, targets, thumbnail_size, skipped=0, remaining=0
|
||||
@@ -2179,8 +2212,17 @@ class ImageRoutesMixin:
|
||||
"""Clear generated thumbnails across gallery roots, preserving originals.
|
||||
|
||||
Invalid thumbnail directories return 400; unavailable gallery roots
|
||||
return 404, and unexpected filesystem failures return 500.
|
||||
return 404, and unexpected filesystem failures return 500. Holds the
|
||||
registry as guarded work, like generation.
|
||||
"""
|
||||
return await self._run_held(
|
||||
THUMBNAILS_JOB_KIND,
|
||||
"Thumbnail clear (gallery page)",
|
||||
"guarded",
|
||||
self._clear_thumbnails,
|
||||
)
|
||||
|
||||
async def _clear_thumbnails(self) -> web.Response:
|
||||
try:
|
||||
output_dirs = self._get_all_output_dirs()
|
||||
if not output_dirs:
|
||||
|
||||
+147
-25
@@ -9,9 +9,21 @@ 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, bulk image tagging, and later
|
||||
maintenance, cleanup and restore) never run together, and one job of a kind
|
||||
runs at a time.
|
||||
ones. One job of a kind runs at a time, and each job has an admission class
|
||||
(``ADMISSIONS``) that decides what may run beside it:
|
||||
|
||||
- ``records`` changes library records or files (the metadata scan, bulk image
|
||||
tagging, and the legacy restore, maintenance, cleanup, duplicate delete,
|
||||
LoRA import and batch auto-tag routes). It is refused while any
|
||||
``records`` or ``guarded`` job runs.
|
||||
- ``guarded`` must not overlap record changes but may overlap other guarded
|
||||
work (thumbnail generation and clearing, later backup). It is refused while
|
||||
a ``records`` job runs.
|
||||
- ``read`` only reads; it is refused only by a running job of its own kind.
|
||||
|
||||
A short legacy-page operation takes a ``hold``: a running, non-cancellable
|
||||
job that lasts as long as the request, so the Activity page lists it and a
|
||||
workspace job is refused while it runs (and the other way round).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
@@ -60,6 +72,41 @@ SCAN_GROUPS = {
|
||||
IMAGE_TAG_JOB_KIND = "image-tags"
|
||||
IMAGE_TAG_GROUP_NAMES = ("tagged", "unchanged", "skipped", "errors")
|
||||
|
||||
ADMISSIONS = ("read", "guarded", "records")
|
||||
# What each admission class is refused beside (besides its own kind).
|
||||
_REFUSED_BESIDE = {
|
||||
"read": frozenset(),
|
||||
"guarded": frozenset({"records"}),
|
||||
"records": frozenset({"records", "guarded"}),
|
||||
}
|
||||
|
||||
# Kinds of the legacy-page operations that hold the registry (D80; titles in
|
||||
# their handlers end in "(admin page)" or "(gallery page)").
|
||||
RESTORE_JOB_KIND = "restore"
|
||||
MAINTENANCE_JOB_KIND = "maintenance"
|
||||
DUPLICATE_CLEANUP_JOB_KIND = "duplicate_cleanup"
|
||||
THUMBNAILS_JOB_KIND = "thumbnails"
|
||||
LORA_IMPORT_JOB_KIND = "lora_import"
|
||||
AUTOTAG_BATCH_JOB_KIND = "autotag"
|
||||
|
||||
|
||||
def resolve_admission(admission=None, record_changing=None) -> str:
|
||||
"""The admission class for a start call; ``record_changing`` is an alias.
|
||||
|
||||
``record_changing=True`` means ``records`` and ``False`` means ``read``;
|
||||
with neither given the class is ``read``. Raises ValueError for an
|
||||
unknown class or one that contradicts ``record_changing``.
|
||||
"""
|
||||
if admission is None:
|
||||
return "records" if record_changing else "read"
|
||||
if admission not in ADMISSIONS:
|
||||
raise ValueError(f"Unknown job admission {admission!r}")
|
||||
if record_changing is not None and record_changing != (admission == "records"):
|
||||
raise ValueError(
|
||||
f"record_changing={record_changing!r} contradicts admission {admission!r}"
|
||||
)
|
||||
return admission
|
||||
|
||||
|
||||
def _relative_path(path) -> str:
|
||||
"""``path`` as a root-relative POSIX path; an absolute one keeps its name."""
|
||||
@@ -85,13 +132,23 @@ class Job:
|
||||
"""One tracked operation: state, counters, outcomes, log and subscribers."""
|
||||
|
||||
def __init__(
|
||||
self, kind, title, *, record_changing=False, scope=None, groups=(), clock=None
|
||||
self,
|
||||
kind,
|
||||
title,
|
||||
*,
|
||||
admission=None,
|
||||
record_changing=None,
|
||||
cancellable=True,
|
||||
scope=None,
|
||||
groups=(),
|
||||
clock=None,
|
||||
):
|
||||
self._clock = clock or time.time
|
||||
self.id = uuid.uuid4().hex
|
||||
self.kind = kind
|
||||
self.title = title
|
||||
self.record_changing = record_changing
|
||||
self.admission = resolve_admission(admission, record_changing)
|
||||
self.cancellable = bool(cancellable)
|
||||
self.scope = dict(scope or {})
|
||||
self.state = "running"
|
||||
self.started_at = self._clock()
|
||||
@@ -110,6 +167,10 @@ class Job:
|
||||
self.task = None
|
||||
self._subscribers = set()
|
||||
|
||||
@property
|
||||
def record_changing(self):
|
||||
return self.admission == "records"
|
||||
|
||||
@property
|
||||
def running(self):
|
||||
return self.state == "running"
|
||||
@@ -240,6 +301,9 @@ class Job:
|
||||
# What detail() lists, so a view re-reads it only when it changed.
|
||||
"listed": {group: len(rows) for group, rows in self.outcomes.items()},
|
||||
"log_length": len(self.log),
|
||||
# A view offers Cancel only when the server can stop the job.
|
||||
"cancellable": self.cancellable,
|
||||
"admission": self.admission,
|
||||
}
|
||||
|
||||
def detail(self):
|
||||
@@ -291,33 +355,76 @@ class JobRegistry:
|
||||
self._logger = logger_ or logger
|
||||
self._jobs = {}
|
||||
|
||||
def start(self, kind, title, *, record_changing, scope, runner, groups=()) -> Job:
|
||||
def start(
|
||||
self,
|
||||
kind,
|
||||
title,
|
||||
*,
|
||||
admission=None,
|
||||
record_changing=None,
|
||||
cancellable=True,
|
||||
scope=None,
|
||||
runner,
|
||||
groups=(),
|
||||
) -> Job:
|
||||
"""Create a job and run ``runner(job)`` as a task.
|
||||
|
||||
Raises JobBusy, naming the job in the way, when a job of ``kind`` is
|
||||
running or when ``record_changing`` and another record-changing job
|
||||
runs. Must be called from the event loop.
|
||||
``admission`` is one of ``ADMISSIONS`` (default ``read``);
|
||||
``record_changing=True`` is the older spelling of ``records``. Raises
|
||||
JobBusy, naming the job in the way, when the admission class refuses
|
||||
it (see the module docstring). Must be called from the event loop.
|
||||
"""
|
||||
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
|
||||
job = self._admit(
|
||||
kind,
|
||||
title,
|
||||
admission=admission,
|
||||
record_changing=record_changing,
|
||||
cancellable=cancellable,
|
||||
scope=scope,
|
||||
groups=groups,
|
||||
)
|
||||
self._jobs[job.id] = job
|
||||
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."""
|
||||
def hold(self, kind, title, *, admission, scope=None):
|
||||
"""Register a running, non-cancellable job for a legacy-page operation.
|
||||
|
||||
Returns an async context manager yielding the job; the job ends when
|
||||
the block exits: ``finished``, or ``failed`` when it raised (unless
|
||||
the body already ended it). Raises JobBusy at once, without waiting,
|
||||
when the admission class refuses it. Use it directly in ``async with``
|
||||
so the job cannot outlive the operation.
|
||||
"""
|
||||
job = self._admit(
|
||||
kind, title, admission=admission, cancellable=False, scope=scope
|
||||
)
|
||||
return _Hold(job)
|
||||
|
||||
def _admit(self, kind, title, *, admission, record_changing=None, **job_args):
|
||||
resolved = resolve_admission(admission, record_changing)
|
||||
blocker = self.blocker(kind, admission=resolved)
|
||||
if blocker is not None:
|
||||
raise JobBusy(blocker)
|
||||
self._prune()
|
||||
job = Job(kind, title, admission=resolved, **job_args)
|
||||
self._jobs[job.id] = job
|
||||
return job
|
||||
|
||||
def blocker(self, kind, *, admission=None, record_changing=None):
|
||||
"""The running job a new ``kind`` job would be refused for, or None.
|
||||
|
||||
The same kind refuses first; then the oldest running job whose class
|
||||
the new job's admission class may not run beside.
|
||||
"""
|
||||
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
|
||||
refused = _REFUSED_BESIDE[resolve_admission(admission, record_changing)]
|
||||
running = [job for job in self._jobs.values() if job.running]
|
||||
running.sort(key=lambda job: job._started_seq)
|
||||
for job in running:
|
||||
if job.admission in refused:
|
||||
return job
|
||||
return None
|
||||
|
||||
async def _run(self, job, runner):
|
||||
@@ -362,11 +469,26 @@ class JobRegistry:
|
||||
def cancel(self, job_id) -> bool:
|
||||
"""Ask a running job to stop at its next safe boundary.
|
||||
|
||||
Returns whether a cancellation was requested: False for a finished or
|
||||
unknown job, so a cancel after the end is a no-op.
|
||||
Returns whether a cancellation was requested: False for a finished,
|
||||
unknown or non-cancellable job, so a cancel after the end is a no-op.
|
||||
"""
|
||||
job = self._jobs.get(job_id)
|
||||
if job is None or job.finished:
|
||||
if job is None or job.finished or not job.cancellable:
|
||||
return False
|
||||
job.cancelled = True
|
||||
return True
|
||||
|
||||
|
||||
class _Hold:
|
||||
"""The async context manager ``JobRegistry.hold`` returns."""
|
||||
|
||||
def __init__(self, job):
|
||||
self._job = job
|
||||
|
||||
async def __aenter__(self):
|
||||
return self._job
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if self._job.running:
|
||||
self._job.finish("finished" if exc_type is None else "failed")
|
||||
return False
|
||||
|
||||
@@ -6,6 +6,7 @@ from pathlib import Path
|
||||
from aiohttp import web
|
||||
|
||||
from ..config_store import ConfigUnreadable, RevisionConflict, get_config_store
|
||||
from .jobs import LORA_IMPORT_JOB_KIND, JobBusy
|
||||
|
||||
# Flags of POST /lora/enable; each must be a JSON boolean when present.
|
||||
_LORA_FLAGS = ("enabled", "trigger_words_enabled", "clear_civitai_api_key")
|
||||
@@ -273,19 +274,38 @@ class LoraIntegrationMixin:
|
||||
reason="OK",
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
)
|
||||
await response.prepare(request)
|
||||
|
||||
async def send_progress(data):
|
||||
line = f"data: {json.dumps(data)}\n\n"
|
||||
await response.write(line.encode("utf-8"))
|
||||
|
||||
# The import replaces prompts and images: a records hold, so it
|
||||
# never runs beside a scan, restore or tagging job.
|
||||
try:
|
||||
await self._lora_scan_stream(allowed, send_progress)
|
||||
except Exception as e:
|
||||
self.logger.exception("LoRA scan failed")
|
||||
await send_progress({"type": "error", "message": self._public_error(e)})
|
||||
held = self._jobs.hold(
|
||||
LORA_IMPORT_JOB_KIND,
|
||||
"LoRA import (admin page)",
|
||||
admission="records",
|
||||
)
|
||||
except JobBusy as busy:
|
||||
await response.prepare(request)
|
||||
message = f"{busy.running.title} already running"
|
||||
await send_progress({"type": "error", "message": message})
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
await response.write_eof()
|
||||
async with held as job:
|
||||
await response.prepare(request)
|
||||
try:
|
||||
await self._lora_scan_stream(allowed, send_progress)
|
||||
except Exception as e:
|
||||
self.logger.exception("LoRA scan failed")
|
||||
job.finish("failed")
|
||||
await send_progress(
|
||||
{"type": "error", "message": self._public_error(e)}
|
||||
)
|
||||
|
||||
await response.write_eof()
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
|
||||
@@ -5,12 +5,16 @@ stream the upload to disk, reject anything that is not a healthy prompts
|
||||
database and leave the live database untouched when it does.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import glob
|
||||
import hashlib
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
_mock_server = MagicMock()
|
||||
@@ -142,6 +146,85 @@ class TestBackupRoute(BackupRestoreTestCase):
|
||||
self.assertIn("backup", data["error"].lower())
|
||||
|
||||
|
||||
class TestRestoreHoldsTheRegistry(BackupRestoreTestCase):
|
||||
"""A legacy restore and a workspace job refuse each other (#269 gap)."""
|
||||
|
||||
def _digest(self):
|
||||
return {
|
||||
suffix: hashlib.sha256(Path(self.db_path + suffix).read_bytes()).hexdigest()
|
||||
for suffix in ("", "-wal")
|
||||
if os.path.exists(self.db_path + suffix)
|
||||
}
|
||||
|
||||
async def test_legacy_restore_refused_while_scan_job_runs(self):
|
||||
self._save("keep me")
|
||||
upload = self._make_db_file(["restored"])
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def runner(job):
|
||||
await gate.wait()
|
||||
|
||||
scan = self.api._jobs.start(
|
||||
"scan", "Scan for metadata", record_changing=True, scope={}, runner=runner
|
||||
)
|
||||
before = self._digest()
|
||||
try:
|
||||
resp = await self._restore(upload)
|
||||
self.assertEqual(resp.status, 409)
|
||||
self.assertEqual(
|
||||
await resp.json(),
|
||||
{"success": False, "error": "Scan for metadata already running"},
|
||||
)
|
||||
self.assertEqual(self._digest(), before)
|
||||
self.assertEqual(
|
||||
[p["text"] for p in self.api.db.search_prompts()], ["keep me"]
|
||||
)
|
||||
self.assertEqual(glob.glob(self.db_path + ".backup_*"), [])
|
||||
finally:
|
||||
gate.set()
|
||||
await asyncio.wait_for(scan.task, 5)
|
||||
|
||||
async def test_scan_refused_while_legacy_restore_holds(self):
|
||||
upload = self._make_db_file(["restored"])
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
model = self.api.db.model
|
||||
real_verify = model.verify_database_file
|
||||
|
||||
def gated(path, **kwargs):
|
||||
entered.set()
|
||||
release.wait(10)
|
||||
return real_verify(path, **kwargs)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
with patch.object(model, "verify_database_file", gated):
|
||||
restore = asyncio.ensure_future(self._restore(upload))
|
||||
self.assertTrue(await loop.run_in_executor(None, entered.wait, 10))
|
||||
resp = await self.client.post("/prompt_manager/jobs/scan")
|
||||
body = await resp.json()
|
||||
listed = (await (await self.client.get("/prompt_manager/jobs")).json())[
|
||||
"jobs"
|
||||
]
|
||||
release.set()
|
||||
restored = await restore
|
||||
|
||||
self.assertEqual((resp.status, body["code"]), (409, "busy"))
|
||||
self.assertEqual(body["running"]["title"], "Database restore (admin page)")
|
||||
self.assertEqual(
|
||||
[(j["kind"], j["state"], j["cancellable"]) for j in listed],
|
||||
[("restore", "running", False)],
|
||||
)
|
||||
self.assertEqual(restored.status, 200)
|
||||
self.assertEqual(self.api._jobs.list()[0].state, "finished")
|
||||
self.assertIsNone(self.api._jobs.active("restore"))
|
||||
|
||||
async def test_rejected_upload_ends_its_hold_failed(self):
|
||||
resp = await self._restore(self._write("notes.txt", b"not a database"))
|
||||
self.assertEqual(resp.status, 400)
|
||||
job = self.api._jobs.list()[0]
|
||||
self.assertEqual((job.kind, job.state), ("restore", "failed"))
|
||||
|
||||
|
||||
class TestRestoreRoute(BackupRestoreTestCase):
|
||||
async def test_restore_replaces_database_and_reports_backup(self):
|
||||
self._save("old prompt")
|
||||
|
||||
@@ -5,11 +5,13 @@ Uses aiohttp's test client against a real PromptManagerAPI wired to a
|
||||
temporary SQLite database and a temporary output directory.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from contextlib import ExitStack
|
||||
from pathlib import Path
|
||||
@@ -638,6 +640,93 @@ class TestGenerateThumbnailsPost(ThumbnailTestCase):
|
||||
self.assertEqual(data["total_images"], 0)
|
||||
|
||||
|
||||
class TestThumbnailRoutesHoldTheRegistry(ThumbnailTestCase):
|
||||
"""Thumbnail work is ``guarded``: never beside a record-changing job."""
|
||||
|
||||
def _run_scan(self):
|
||||
gate = asyncio.Event()
|
||||
|
||||
async def runner(job):
|
||||
await gate.wait()
|
||||
|
||||
job = self.api._jobs.start(
|
||||
"scan", "Scan for metadata", admission="records", scope={}, runner=runner
|
||||
)
|
||||
|
||||
async def release():
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 5)
|
||||
|
||||
self.addAsyncCleanup(release)
|
||||
return job
|
||||
|
||||
async def test_thumbnail_generate_refused_while_scan_runs(self):
|
||||
make_png(self.output_dir / "a.png")
|
||||
self._run_scan()
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
||||
)
|
||||
self.assertEqual(resp.status, 409)
|
||||
self.assertEqual(
|
||||
await resp.json(),
|
||||
{"success": False, "error": "Scan for metadata already running"},
|
||||
)
|
||||
|
||||
resp = await self.client.request("POST", self.PROGRESS_URL)
|
||||
text = await resp.text()
|
||||
self.assertEqual(resp.status, 200) # the gallery reads an SSE stream
|
||||
self.assertEqual(
|
||||
text,
|
||||
"event: error\ndata: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"error": "Scan for metadata already running",
|
||||
"message": "Scan for metadata already running",
|
||||
}
|
||||
)
|
||||
+ "\n\n",
|
||||
)
|
||||
|
||||
thumbs = self.output_dir / "thumbnails"
|
||||
make_png(thumbs / "b_thumb.png")
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/images/clear-thumbnails"
|
||||
)
|
||||
self.assertEqual(resp.status, 409)
|
||||
self.assertFalse(self._thumb_path("a").exists())
|
||||
self.assertTrue((thumbs / "b_thumb.png").exists())
|
||||
|
||||
async def test_scan_refused_while_thumbnails_generate(self):
|
||||
make_png(self.output_dir / "a.png")
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
real = self.api._generate_gallery_thumbnails_sync
|
||||
|
||||
def gated(*args):
|
||||
entered.set()
|
||||
release.wait(10)
|
||||
return real(*args)
|
||||
|
||||
self.api._generate_gallery_thumbnails_sync = gated
|
||||
loop = asyncio.get_running_loop()
|
||||
generate = asyncio.ensure_future(
|
||||
self.client.request(
|
||||
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
||||
)
|
||||
)
|
||||
self.assertTrue(await loop.run_in_executor(None, entered.wait, 10))
|
||||
resp = await self.client.request("POST", "/prompt_manager/jobs/scan")
|
||||
body = await resp.json()
|
||||
self.assertEqual((resp.status, body["code"]), (409, "busy"))
|
||||
self.assertEqual(
|
||||
body["running"]["title"], "Thumbnail generation (gallery page)"
|
||||
)
|
||||
release.set()
|
||||
self.assertEqual((await generate).status, 200)
|
||||
self.assertTrue(self._thumb_path("a").exists())
|
||||
|
||||
|
||||
class TestRecentImagesCaps(ImageAPITestCase):
|
||||
"""Every path through /images/recent hands the DB a bounded window."""
|
||||
|
||||
|
||||
@@ -420,6 +420,23 @@ class TestStartAndExclusion(ImageTagJobTestCase):
|
||||
# A refused start never leaves the gate taken.
|
||||
self.assertIsNone(model_gate.holder())
|
||||
|
||||
async def test_image_tag_job_refused_while_legacy_restore_holds(self):
|
||||
self.png("a.png")
|
||||
async with self.api._jobs.hold(
|
||||
"restore", "Database restore (admin page)", admission="records"
|
||||
):
|
||||
status, body = await self.start(self.add_body(["a.png"], ["x"]))
|
||||
self.assertEqual((status, body["code"]), (409, "busy"))
|
||||
self.assertEqual(body["running"]["title"], "Database restore (admin page)")
|
||||
status, body = await self.start(self.auto_body(["a.png"]))
|
||||
self.assertEqual((status, body["code"]), (409, "busy"))
|
||||
self.assertIsNone(model_gate.holder())
|
||||
async with self.api._jobs.hold(
|
||||
"thumbnails", "Thumbnail generation (gallery page)", admission="guarded"
|
||||
):
|
||||
status, body = await self.start(self.add_body(["a.png"], ["x"]))
|
||||
self.assertEqual((status, body["code"]), (409, "busy"))
|
||||
|
||||
async def test_record_changing_excludes_running_scan(self):
|
||||
self.png("a.png")
|
||||
gate = asyncio.Event()
|
||||
|
||||
@@ -0,0 +1,395 @@
|
||||
"""Registry admission classes and the legacy routes that hold the registry.
|
||||
|
||||
Record-changing work never runs together (Activity design): a ``records``
|
||||
job is refused while any ``records`` or ``guarded`` job runs, a ``guarded``
|
||||
job (thumbnails) while a ``records`` job runs, and ``read`` jobs only share
|
||||
the one-per-kind rule. The legacy admin and gallery routes that change
|
||||
records or files take a registry hold for their duration, so a workspace job
|
||||
and a legacy-page operation refuse each other, naming the one that runs.
|
||||
|
||||
Fake runners and gated executor calls only; ordering is driven by events.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
_mock_server = MagicMock()
|
||||
_mock_server.PromptServer.instance.routes = MagicMock()
|
||||
sys.modules.setdefault("server", _mock_server)
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from py.api.jobs import ( # noqa: E402
|
||||
ADMISSIONS,
|
||||
Job,
|
||||
JobBusy,
|
||||
JobRegistry,
|
||||
)
|
||||
|
||||
try:
|
||||
from tests.test_api_admin import AdminAPITestCase # noqa: E402
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from test_api_admin import AdminAPITestCase # noqa: E402
|
||||
|
||||
|
||||
def _gated(gate):
|
||||
async def runner(job):
|
||||
await gate.wait()
|
||||
|
||||
return runner
|
||||
|
||||
|
||||
async def _noop(job):
|
||||
return None
|
||||
|
||||
|
||||
class TestAdmissionClasses(unittest.IsolatedAsyncioTestCase):
|
||||
|
||||
def _start(self, registry, kind, admission, gate=None, **kwargs):
|
||||
return registry.start(
|
||||
kind,
|
||||
kind.title(),
|
||||
admission=admission,
|
||||
scope={},
|
||||
runner=_gated(gate) if gate is not None else _noop,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async def test_admission_names_and_default(self):
|
||||
self.assertEqual(ADMISSIONS, ("read", "guarded", "records"))
|
||||
registry = JobRegistry()
|
||||
job = registry.start("report", "Report", scope={}, runner=_noop)
|
||||
await asyncio.wait_for(job.task, 2)
|
||||
self.assertEqual(job.admission, "read")
|
||||
self.assertFalse(job.record_changing)
|
||||
with self.assertRaises(ValueError):
|
||||
registry.start("x", "X", admission="exclusive", scope={}, runner=_noop)
|
||||
with self.assertRaises(ValueError):
|
||||
registry.start(
|
||||
"x",
|
||||
"X",
|
||||
admission="read",
|
||||
record_changing=True,
|
||||
scope={},
|
||||
runner=_noop,
|
||||
)
|
||||
|
||||
async def test_record_changing_is_an_alias_for_records(self):
|
||||
registry = JobRegistry()
|
||||
gate = asyncio.Event()
|
||||
job = registry.start(
|
||||
"scan", "Scan", record_changing=True, scope={}, runner=_gated(gate)
|
||||
)
|
||||
self.assertEqual(job.admission, "records")
|
||||
self.assertTrue(job.record_changing)
|
||||
self.assertEqual(job.summary()["admission"], "records")
|
||||
self.assertIs(
|
||||
registry.blocker("image-tags", record_changing=True), job
|
||||
) # the existing call shape keeps working
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 2)
|
||||
|
||||
async def test_records_job_refused_while_guarded_runs(self):
|
||||
registry = JobRegistry()
|
||||
gate = asyncio.Event()
|
||||
thumbs = self._start(registry, "thumbnails", "guarded", gate)
|
||||
with self.assertRaises(JobBusy) as raised:
|
||||
self._start(registry, "scan", "records")
|
||||
self.assertIs(raised.exception.running, thumbs)
|
||||
self.assertIs(registry.blocker("restore", admission="records"), thumbs)
|
||||
gate.set()
|
||||
await asyncio.wait_for(thumbs.task, 2)
|
||||
scan = self._start(registry, "scan", "records")
|
||||
await asyncio.wait_for(scan.task, 2)
|
||||
self.assertEqual(scan.state, "finished")
|
||||
|
||||
async def test_guarded_refused_while_records_runs(self):
|
||||
registry = JobRegistry()
|
||||
gate = asyncio.Event()
|
||||
scan = self._start(registry, "scan", "records", gate)
|
||||
with self.assertRaises(JobBusy) as raised:
|
||||
self._start(registry, "thumbnails", "guarded")
|
||||
self.assertIs(raised.exception.running, scan)
|
||||
# A second guarded kind may run beside a guarded one, not beside records.
|
||||
gate.set()
|
||||
await asyncio.wait_for(scan.task, 2)
|
||||
hold = asyncio.Event()
|
||||
thumbs = self._start(registry, "thumbnails", "guarded", hold)
|
||||
backup = self._start(registry, "backup", "guarded")
|
||||
await asyncio.wait_for(backup.task, 2)
|
||||
self.assertEqual(backup.state, "finished")
|
||||
hold.set()
|
||||
await asyncio.wait_for(thumbs.task, 2)
|
||||
|
||||
async def test_read_jobs_run_in_parallel(self):
|
||||
registry = JobRegistry()
|
||||
gate = asyncio.Event()
|
||||
scan = self._start(registry, "scan", "records", gate)
|
||||
health = self._start(registry, "health", "read", gate)
|
||||
duplicates = self._start(registry, "duplicates", "read", gate)
|
||||
self.assertTrue(all(j.running for j in (scan, health, duplicates)))
|
||||
with self.assertRaises(JobBusy) as raised: # one job of a kind
|
||||
self._start(registry, "health", "read")
|
||||
self.assertIs(raised.exception.running, health)
|
||||
gate.set()
|
||||
await asyncio.wait_for(
|
||||
asyncio.gather(scan.task, health.task, duplicates.task), 2
|
||||
)
|
||||
|
||||
async def test_non_cancellable_job_ignores_cancel(self):
|
||||
registry = JobRegistry()
|
||||
gate = asyncio.Event()
|
||||
job = self._start(registry, "restore", "records", gate, cancellable=False)
|
||||
self.assertIs(job.summary()["cancellable"], False)
|
||||
self.assertFalse(registry.cancel(job.id))
|
||||
self.assertFalse(job.cancelled)
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 2)
|
||||
self.assertEqual(job.state, "finished")
|
||||
|
||||
async def test_hold_registers_a_running_non_cancellable_job(self):
|
||||
registry = JobRegistry()
|
||||
held = registry.hold(
|
||||
"restore", "Database restore (admin page)", admission="records"
|
||||
)
|
||||
async with held as job:
|
||||
self.assertIs(registry.active("restore"), job)
|
||||
self.assertEqual(job.state, "running")
|
||||
self.assertEqual(
|
||||
(job.summary()["cancellable"], job.summary()["admission"]),
|
||||
(False, "records"),
|
||||
)
|
||||
self.assertFalse(registry.cancel(job.id))
|
||||
with self.assertRaises(JobBusy) as raised:
|
||||
self._start(registry, "scan", "records")
|
||||
self.assertIn("Database restore (admin page)", str(raised.exception))
|
||||
self.assertEqual(job.state, "finished")
|
||||
self.assertIsNone(registry.active("restore"))
|
||||
self.assertIs(registry.list()[0], job)
|
||||
|
||||
async def test_hold_refused_without_waiting(self):
|
||||
registry = JobRegistry()
|
||||
gate = asyncio.Event()
|
||||
scan = self._start(registry, "scan", "records", gate)
|
||||
with self.assertRaises(JobBusy) as raised:
|
||||
registry.hold("maintenance", "Maintenance", admission="records")
|
||||
self.assertIs(raised.exception.running, scan)
|
||||
self.assertEqual([j.kind for j in registry.list()], ["scan"])
|
||||
gate.set()
|
||||
await asyncio.wait_for(scan.task, 2)
|
||||
|
||||
async def test_hold_released_on_exception(self):
|
||||
registry = JobRegistry()
|
||||
with self.assertRaises(RuntimeError):
|
||||
async with registry.hold(
|
||||
"maintenance", "Maintenance", admission="records"
|
||||
) as job:
|
||||
raise RuntimeError("boom")
|
||||
self.assertEqual(job.state, "failed")
|
||||
again = registry.hold("maintenance", "Maintenance", admission="records")
|
||||
async with again as second:
|
||||
second.finish("failed") # the body may end it itself
|
||||
self.assertEqual(second.state, "failed")
|
||||
|
||||
def test_job_defaults_stay_read_and_cancellable(self):
|
||||
job = Job("duplicates", "Duplicate scan")
|
||||
self.assertEqual((job.admission, job.cancellable), ("read", True))
|
||||
|
||||
|
||||
class _LegacyHoldCase(AdminAPITestCase):
|
||||
"""The admin test app plus a running workspace job on demand."""
|
||||
|
||||
def _run_job(self, kind="scan", title="Scan for metadata", admission="records"):
|
||||
gate = asyncio.Event()
|
||||
job = self.api._jobs.start(
|
||||
kind, title, admission=admission, scope={}, runner=_gated(gate)
|
||||
)
|
||||
self.addAsyncCleanup(self._release, gate, job)
|
||||
return gate, job
|
||||
|
||||
@staticmethod
|
||||
async def _release(gate, job):
|
||||
gate.set()
|
||||
await asyncio.wait_for(job.task, 5)
|
||||
|
||||
def _db_digest(self):
|
||||
path = Path(self.api.db.model.db_path)
|
||||
return {
|
||||
suffix: hashlib.sha256(Path(f"{path}{suffix}").read_bytes()).hexdigest()
|
||||
for suffix in ("", "-wal")
|
||||
if Path(f"{path}{suffix}").exists()
|
||||
}
|
||||
|
||||
async def _busy(self, resp, title="Scan for metadata"):
|
||||
body = await resp.json()
|
||||
self.assertEqual(resp.status, 409)
|
||||
self.assertEqual(body, {"success": False, "error": f"{title} already running"})
|
||||
|
||||
|
||||
class TestLegacyRoutesHoldTheRegistry(_LegacyHoldCase):
|
||||
|
||||
async def test_hold_is_listed_as_running_job(self):
|
||||
held = self.api._jobs.hold(
|
||||
"restore", "Database restore (admin page)", admission="records"
|
||||
)
|
||||
async with held:
|
||||
resp = await self.client.request("GET", "/prompt_manager/jobs")
|
||||
jobs = (await resp.json())["jobs"]
|
||||
self.assertEqual(len(jobs), 1)
|
||||
job = jobs[0]
|
||||
self.assertEqual(
|
||||
(job["kind"], job["title"], job["state"]),
|
||||
("restore", "Database restore (admin page)", "running"),
|
||||
)
|
||||
self.assertIs(job["cancellable"], False)
|
||||
self.assertEqual(job["admission"], "records")
|
||||
|
||||
async def test_legacy_maintenance_and_cleanup_refused_while_scan_runs(self):
|
||||
self.api.db.save_prompt(text="keep me", prompt_hash="h1")
|
||||
before = self._db_digest()
|
||||
self._run_job()
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/maintenance",
|
||||
json={"operations": ["cleanup_duplicates", "vacuum"]},
|
||||
)
|
||||
await self._busy(resp)
|
||||
resp = await self.client.request("POST", "/prompt_manager/cleanup")
|
||||
await self._busy(resp)
|
||||
self.assertEqual(self._db_digest(), before)
|
||||
|
||||
async def test_legacy_maintenance_refused_while_image_tags_runs(self):
|
||||
self._run_job("image-tags", "Add tags to 3 images")
|
||||
resp = await self.client.request("POST", "/prompt_manager/maintenance")
|
||||
await self._busy(resp, "Add tags to 3 images")
|
||||
|
||||
async def test_scan_refused_while_legacy_maintenance_holds(self):
|
||||
entered = threading.Event()
|
||||
release = threading.Event()
|
||||
real = self.api._run_maintenance_operations
|
||||
|
||||
def gated(operations):
|
||||
entered.set()
|
||||
release.wait(10)
|
||||
return real(operations)
|
||||
|
||||
self.api._run_maintenance_operations = gated
|
||||
loop = asyncio.get_running_loop()
|
||||
maintenance = asyncio.ensure_future(
|
||||
self.client.request(
|
||||
"POST", "/prompt_manager/maintenance", json={"operations": ["vacuum"]}
|
||||
)
|
||||
)
|
||||
self.assertTrue(await loop.run_in_executor(None, entered.wait, 10))
|
||||
resp = await self.client.request("POST", "/prompt_manager/jobs/scan")
|
||||
body = await resp.json()
|
||||
self.assertEqual((resp.status, body["code"]), (409, "busy"))
|
||||
self.assertEqual(body["running"]["title"], "Maintenance (admin page)")
|
||||
release.set()
|
||||
self.assertEqual((await maintenance).status, 200)
|
||||
listed = (
|
||||
await (await self.client.request("GET", "/prompt_manager/jobs")).json()
|
||||
)["jobs"]
|
||||
self.assertEqual(
|
||||
[(j["kind"], j["state"]) for j in listed], [("maintenance", "finished")]
|
||||
)
|
||||
|
||||
async def test_legacy_duplicate_delete_refused_while_scan_runs(self):
|
||||
victim = self.output_dir / "dup.png"
|
||||
victim.write_bytes(b"x")
|
||||
self._run_job()
|
||||
resp = await self.client.request(
|
||||
"POST",
|
||||
"/prompt_manager/delete_duplicate_images",
|
||||
json={"image_paths": [str(victim)]},
|
||||
)
|
||||
await self._busy(resp)
|
||||
self.assertTrue(victim.exists())
|
||||
# The duplicate family lock was not taken by the refused request.
|
||||
self.assertIsNotNone(self.api._acquire_long_job(self.api.LONG_JOB_DUPLICATES))
|
||||
|
||||
async def test_lora_scan_refused_while_scan_runs(self):
|
||||
from py.config import IntegrationConfig
|
||||
|
||||
lm = Path(self.tmpdir) / "custom_nodes" / "ComfyUI-Lora-Manager"
|
||||
(lm / "py").mkdir(parents=True)
|
||||
self._run_job()
|
||||
imported = []
|
||||
with (
|
||||
patch.multiple(
|
||||
IntegrationConfig, LORA_MANAGER_ENABLED=True, LORA_MANAGER_PATH=str(lm)
|
||||
),
|
||||
patch("py.lora_utils.resolve_lora_manager_path", return_value=str(lm)),
|
||||
patch.object(
|
||||
self.api, "_lora_scan_stream", side_effect=lambda *a: imported.append(a)
|
||||
),
|
||||
):
|
||||
resp = await self.client.request("POST", "/prompt_manager/lora/scan")
|
||||
text = await resp.text()
|
||||
self.assertEqual(resp.status, 200) # the legacy page reads a stream
|
||||
events = [json.loads(line[6:]) for line in text.split("\n\n") if line]
|
||||
self.assertEqual(
|
||||
events,
|
||||
[{"type": "error", "message": "Scan for metadata already running"}],
|
||||
)
|
||||
self.assertEqual(imported, [])
|
||||
|
||||
async def test_legacy_autotag_refused_while_scan_runs(self):
|
||||
self._run_job()
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/autotag/start", json={}
|
||||
)
|
||||
await self._busy(resp)
|
||||
|
||||
async def test_legacy_payloads_unchanged_when_idle(self):
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/maintenance", json={"operations": ["vacuum"]}
|
||||
)
|
||||
body = await resp.json()
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(
|
||||
set(body),
|
||||
{"success", "operations_completed", "all_successful", "results", "message"},
|
||||
)
|
||||
resp = await self.client.request("POST", "/prompt_manager/cleanup")
|
||||
body = await resp.json()
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(set(body), {"success", "message", "duplicates_removed"})
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/delete_duplicate_images", json={"image_paths": []}
|
||||
)
|
||||
self.assertEqual(
|
||||
(resp.status, await resp.json()),
|
||||
(400, {"success": False, "error": "No image paths provided"}),
|
||||
)
|
||||
# Each hold ended with the request; a refused input shows as failed.
|
||||
listed = (
|
||||
await (await self.client.request("GET", "/prompt_manager/jobs")).json()
|
||||
)["jobs"]
|
||||
self.assertEqual(
|
||||
[(j["kind"], j["state"]) for j in listed],
|
||||
[
|
||||
("duplicate_cleanup", "failed"),
|
||||
("maintenance", "finished"),
|
||||
("maintenance", "finished"),
|
||||
],
|
||||
)
|
||||
self.assertEqual(
|
||||
[j["title"] for j in listed],
|
||||
[
|
||||
"Duplicate image delete (gallery page)",
|
||||
"Prompt cleanup (admin page)",
|
||||
"Maintenance (admin page)",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -182,8 +182,13 @@ class TestJobRegistry(unittest.IsolatedAsyncioTestCase):
|
||||
"scope",
|
||||
"listed",
|
||||
"log_length",
|
||||
"cancellable",
|
||||
"admission",
|
||||
},
|
||||
)
|
||||
self.assertEqual(
|
||||
(summary["cancellable"], summary["admission"]), (True, "records")
|
||||
)
|
||||
self.assertEqual(summary["scope"], {"files": "new", "output_folders": 2})
|
||||
# What the detail would list, so a view reads it only when that changed.
|
||||
self.assertEqual((summary["listed"], summary["log_length"]), ({}, 0))
|
||||
|
||||
Reference in New Issue
Block a user