Vito
88fa13689a
* feat(autotag): model status without paths, and a busy answer that names the operation * feat(autotag): model downloads run as Activity jobs * fix(autotag): keep each download route's error shape and harden the fake download
553 lines
22 KiB
Python
553 lines
22 KiB
Python
"""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_explicit_read(self):
|
|
self.assertEqual(ADMISSIONS, ("read", "guarded", "records", "exclusive"))
|
|
registry = JobRegistry()
|
|
job = registry.start(
|
|
"report", "Report", admission="read", 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="bogus", scope={}, runner=_noop)
|
|
with self.assertRaises(ValueError):
|
|
registry.start(
|
|
"x",
|
|
"X",
|
|
admission="read",
|
|
record_changing=True,
|
|
scope={},
|
|
runner=_noop,
|
|
)
|
|
|
|
async def test_start_requires_an_admission_class(self):
|
|
# A job that forgets its class must not silently run as "read" beside
|
|
# a scan or restore: every start call decides.
|
|
registry = JobRegistry()
|
|
with self.assertRaises(TypeError) as raised:
|
|
registry.start("review", "Review", scope={}, runner=_noop)
|
|
self.assertIn("admission", str(raised.exception))
|
|
self.assertEqual(registry.list(), [])
|
|
job = registry.start(
|
|
"review", "Review", record_changing=False, scope={}, runner=_noop
|
|
)
|
|
await asyncio.wait_for(job.task, 2)
|
|
self.assertEqual(job.admission, "read")
|
|
|
|
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:
|
|
async with registry.hold("maintenance", "Maintenance", admission="records"):
|
|
self.fail("a refused hold must not run its body")
|
|
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")
|
|
|
|
async def test_hold_admits_on_entry_not_at_the_call(self):
|
|
# A hold created but never entered (an await or an exception between
|
|
# the call and ``async with``) must not leave a running records job
|
|
# behind that blocks all record work until a restart.
|
|
registry = JobRegistry()
|
|
held = registry.hold("restore", "Restore", admission="records")
|
|
self.assertEqual(registry.list(), [])
|
|
self.assertIsNone(registry.blocker("scan", admission="records"))
|
|
del held
|
|
async with registry.hold("restore", "Restore", admission="records") as job:
|
|
self.assertIs(registry.active("restore"), job)
|
|
self.assertEqual(job.state, "finished")
|
|
|
|
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_read_only_maintenance_runs_during_scan(self):
|
|
# The legacy statistics and consistency checks only read: they hold
|
|
# the registry as "read", so a long scan does not block them.
|
|
self.api.db.save_prompt(text="keep me", prompt_hash="h1")
|
|
_gate, scan = self._run_job()
|
|
listed = []
|
|
real = self.api._run_maintenance_operations
|
|
|
|
def observed(operations):
|
|
listed.extend(
|
|
(job.kind, job.admission, job.state) for job in self.api._jobs.list()
|
|
)
|
|
return real(operations)
|
|
|
|
self.api._run_maintenance_operations = observed
|
|
resp = await self.client.request(
|
|
"POST",
|
|
"/prompt_manager/maintenance",
|
|
json={
|
|
"operations": [
|
|
"statistics",
|
|
"check_consistency",
|
|
"check_hash_duplicates",
|
|
]
|
|
},
|
|
)
|
|
body = await resp.json()
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertIs(body["success"], True)
|
|
self.assertEqual(
|
|
set(body["results"]),
|
|
{"statistics", "check_consistency", "check_hash_duplicates"},
|
|
)
|
|
self.assertIn(("maintenance", "read", "running"), listed)
|
|
self.assertEqual(scan.state, "running")
|
|
maintenance = self.api._jobs.list()[-1]
|
|
self.assertEqual(
|
|
(maintenance.kind, maintenance.admission, maintenance.state),
|
|
("maintenance", "read", "finished"),
|
|
)
|
|
|
|
async def test_mixed_maintenance_selection_is_refused_during_scan(self):
|
|
# One write operation among read-only ones makes the run "records".
|
|
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": ["statistics", "vacuum"]},
|
|
)
|
|
await self._busy(resp)
|
|
self.assertEqual(self._db_digest(), before)
|
|
|
|
async def test_default_maintenance_selection_is_refused_during_scan(self):
|
|
# No "operations" key means the default list, which vacuums.
|
|
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={})
|
|
await self._busy(resp)
|
|
self.assertEqual(self._db_digest(), before)
|
|
self.assertEqual([j.kind for j in self.api._jobs.list()], ["scan"])
|
|
|
|
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]
|
|
# The admin page's modal shows ``status`` and closes on ``type: error``.
|
|
message = "Scan for metadata already running"
|
|
self.assertEqual(
|
|
events,
|
|
[{"type": "error", "status": message, "progress": 0, "message": message}],
|
|
)
|
|
self.assertEqual(imported, [])
|
|
|
|
async def test_legacy_autotag_refused_while_scan_runs(self):
|
|
# Both pages read this route as a stream (a 409 reads as a lost
|
|
# connection there), so a refusal by another job is one error frame:
|
|
# the admin page keys on ``type``, the gallery on ``status``.
|
|
self._run_job()
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/autotag/start", json={}
|
|
)
|
|
text = await resp.text()
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertEqual(resp.headers["Content-Type"], "text/event-stream")
|
|
events = [json.loads(line[6:]) for line in text.split("\n\n") if line]
|
|
message = "Scan for metadata already running"
|
|
self.assertEqual(
|
|
events, [{"type": "error", "status": "error", "message": message}]
|
|
)
|
|
self.assertEqual(
|
|
[j.kind for j in self.api._jobs.list() if j.kind == "autotag"], []
|
|
)
|
|
|
|
async def test_second_legacy_autotag_batch_keeps_its_json_409(self):
|
|
from py.api.autotag_routes import _BUSY_ERROR
|
|
|
|
self._run_job("autotag", "Tag output images (legacy page)")
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/autotag/start", json={}
|
|
)
|
|
self.assertEqual(resp.status, 409)
|
|
self.assertEqual(await resp.json(), {"success": False, "error": _BUSY_ERROR})
|
|
|
|
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; refused input never took one.
|
|
listed = await self._listed()
|
|
self.assertEqual(
|
|
[(j["kind"], j["state"]) for j in listed],
|
|
[("maintenance", "finished"), ("maintenance", "finished")],
|
|
)
|
|
self.assertEqual(
|
|
[j["title"] for j in listed],
|
|
["Prompt cleanup (admin page)", "Maintenance (admin page)"],
|
|
)
|
|
|
|
async def _listed(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/jobs")
|
|
return (await resp.json())["jobs"]
|
|
|
|
|
|
class TestHeldJobStates(_LegacyHoldCase):
|
|
"""Activity lists what a held legacy operation really did (Minor 7)."""
|
|
|
|
async def _listed(self):
|
|
resp = await self.client.request("GET", "/prompt_manager/jobs")
|
|
return [(j["kind"], j["state"]) for j in (await resp.json())["jobs"]]
|
|
|
|
async def test_refused_input_takes_no_hold(self):
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/delete_duplicate_images", json={}
|
|
)
|
|
self.assertEqual(resp.status, 400)
|
|
resp = await self.client.request(
|
|
"POST",
|
|
"/prompt_manager/images/generate-thumbnails",
|
|
data=b"{not json",
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
self.assertEqual(resp.status, 400)
|
|
self.assertEqual(await self._listed(), [])
|
|
|
|
async def test_thumbnail_stream_that_ends_in_error_is_failed(self):
|
|
with patch.object(self.api, "_get_all_output_dirs", return_value=[]):
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/images/generate-thumbnails/progress"
|
|
)
|
|
text = await resp.text()
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertIn("event: error", text)
|
|
self.assertEqual(await self._listed(), [("thumbnails", "failed")])
|
|
|
|
async def test_legacy_autotag_that_ends_in_error_is_failed(self):
|
|
service = MagicMock()
|
|
service.get_models_status.return_value = {} # model not downloaded
|
|
with patch("py.autotag.get_autotag_service", return_value=service):
|
|
resp = await self.client.request(
|
|
"POST", "/prompt_manager/autotag/start", json={}
|
|
)
|
|
text = await resp.text()
|
|
self.assertEqual(resp.status, 200)
|
|
self.assertIn('"type": "error"', text)
|
|
self.assertEqual(await self._listed(), [("autotag", "failed")])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|