432 lines
18 KiB
Python
432 lines
18 KiB
Python
"""Store each workflow once, as a maintenance job (py/api/workflow_storage.py).
|
|
|
|
``POST /prompt_manager/maintenance/preview {operation:"workflow_storage"}``
|
|
reviews the rewrite; applying it starts a ``records`` job that takes the
|
|
verified safety backup first, rewrites old-form rows batch by batch (Cancel
|
|
stops after the current batch) and then starts "Reclaim space", an
|
|
``exclusive`` job with its own disk check. No image file is ever opened.
|
|
"""
|
|
|
|
import asyncio
|
|
import builtins
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
from unittest.mock import patch
|
|
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
from database import workflow_storage as ws
|
|
from database.workflow_documents import DocumentReader, resolve_image_row, store_text
|
|
from py.api import maintenance
|
|
from py.api.jobs import MAINTENANCE_JOB_KIND
|
|
|
|
try:
|
|
from tests.test_api_maintenance import ApplyTestCase
|
|
from tests.test_workflow_documents import sample_prompt_graph, sample_workflow
|
|
except ImportError: # discovered with tests/ as the top-level directory
|
|
from test_api_maintenance import ApplyTestCase
|
|
from test_workflow_documents import sample_prompt_graph, sample_workflow
|
|
|
|
WAIT = 20
|
|
OPERATION = "workflow_storage"
|
|
|
|
|
|
def _digest(path):
|
|
with open(path, "rb") as handle:
|
|
return hashlib.sha256(handle.read()).hexdigest()
|
|
|
|
|
|
class WorkflowStorageTestCase(ApplyTestCase):
|
|
def old_library(self, count=12, media=False):
|
|
"""Old-form rows (text inline), four images per run."""
|
|
prompt_id = self.save("a cat")
|
|
conn = self.api.db.model.get_connection()
|
|
for index in range(count):
|
|
run = index // 4
|
|
path = self.output_dir / f"{index}.png"
|
|
workflow = json.dumps(sample_workflow(run, nodes=120))
|
|
if media:
|
|
info = PngImagePlugin.PngInfo()
|
|
info.add_text("workflow", workflow)
|
|
with Image.new("RGB", (4, 4)) as image:
|
|
image.save(path, pnginfo=info)
|
|
with conn:
|
|
conn.execute(
|
|
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
|
"file_path, workflow_data, prompt_metadata, parameters) "
|
|
"VALUES (?, ?, ?, ?, ?, ?, '{}')",
|
|
(
|
|
prompt_id,
|
|
str(path),
|
|
path.name,
|
|
str(path),
|
|
workflow,
|
|
json.dumps(sample_prompt_graph(run)),
|
|
),
|
|
)
|
|
return prompt_id
|
|
|
|
def texts(self):
|
|
conn = self.api.db.model.get_connection()
|
|
reader = DocumentReader(conn)
|
|
return {
|
|
row["id"]: resolve_image_row(reader, row)
|
|
for row in conn.execute(
|
|
"SELECT id, workflow_data, prompt_metadata, workflow_document_id, "
|
|
"prompt_document_id FROM generated_images ORDER BY id"
|
|
)
|
|
}
|
|
|
|
def inline_count(self):
|
|
conn = self.api.db.model.get_connection()
|
|
return conn.execute(
|
|
"SELECT COUNT(*) FROM generated_images WHERE workflow_data IS NOT NULL"
|
|
).fetchone()[0]
|
|
|
|
def reclaim_job(self):
|
|
jobs = [job for job in self.api._jobs.list() if job.title == "Reclaim space"]
|
|
return jobs[0] if jobs else None
|
|
|
|
async def finish(self, job):
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
reclaim = self.reclaim_job()
|
|
if reclaim is not None and reclaim.task is not None:
|
|
await asyncio.wait_for(asyncio.shield(reclaim.task), WAIT)
|
|
return reclaim
|
|
|
|
async def start(self):
|
|
plan = await self.preview(OPERATION)
|
|
body = await self.apply(plan["token"])
|
|
return plan, body["job"], self.api._jobs.get(body["job"]["id"])
|
|
|
|
|
|
class TestWorkflowStorageJob(WorkflowStorageTestCase):
|
|
async def test_preview_reviews_the_rewrite_and_writes_nothing(self):
|
|
self.old_library()
|
|
db_path = self.api.db.model.db_path
|
|
before = _digest(db_path), _digest(db_path + "-wal")
|
|
plan = await self.preview(OPERATION)
|
|
self.assertEqual((_digest(db_path), _digest(db_path + "-wal")), before)
|
|
self.assertEqual(plan["operation"], OPERATION)
|
|
self.assertEqual(plan["scope"]["count"], 12)
|
|
self.assertEqual(plan["scope"]["label"], "12 image records to rewrite")
|
|
self.assertFalse(plan["media"]["deletes_files"])
|
|
self.assertEqual(
|
|
plan["media"]["label"],
|
|
"No image or video file is opened or changed; only the database.",
|
|
)
|
|
self.assertEqual(plan["backup"]["policy"], "required")
|
|
self.assertEqual(
|
|
plan["backup"]["destination_label"],
|
|
"prompts.db.pre-workflow-storage_<date and time>",
|
|
)
|
|
self.assertGreater(plan["backup"]["expected_bytes"], 0)
|
|
self.assertTrue(plan["disk"]["sufficient"])
|
|
|
|
async def test_apply_backs_up_rewrites_then_reclaims_space(self):
|
|
prompt_id = self.old_library(count=40)
|
|
before = self.texts()
|
|
listed = await self.client.request(
|
|
"GET", f"/prompt_manager/prompts/{prompt_id}/images"
|
|
)
|
|
listed_before = await listed.json()
|
|
size_before = self.database_bytes()
|
|
plan, summary, job = await self.start()
|
|
self.assertEqual(summary["kind"], MAINTENANCE_JOB_KIND)
|
|
self.assertEqual(summary["title"], "Store each workflow once")
|
|
self.assertEqual(summary["admission"], "records")
|
|
self.assertTrue(summary["cancellable"])
|
|
self.assertEqual(summary["scope"], {"operation": OPERATION})
|
|
reclaim = await self.finish(job)
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertEqual(job.counters, {"changed": 40, "protected": 0, "errors": 0})
|
|
self.assertEqual(job.total, 40)
|
|
self.assertEqual(job.processed, 40)
|
|
self.assertEqual(self.inline_count(), 0)
|
|
self.assertEqual(self.texts(), before)
|
|
self.assertEqual(self.backups(), []) # never in the pruned pool
|
|
[backup] = [
|
|
name
|
|
for name in os.listdir(self.tmpdir)
|
|
if name.startswith("prompts.db.pre-workflow-storage_")
|
|
]
|
|
self.assertIn(f"Pre-migration backup {backup} verified", "\n".join(job.log))
|
|
self.assertEqual(
|
|
job.status_line,
|
|
f"Rewrote 40 image records · Pre-migration backup {backup}",
|
|
)
|
|
# The safety backup holds the old form, verified.
|
|
conn = sqlite3.connect(os.path.join(self.tmpdir, backup))
|
|
try:
|
|
inline = conn.execute(
|
|
"SELECT COUNT(*) FROM generated_images WHERE workflow_data IS NOT NULL"
|
|
).fetchone()[0]
|
|
finally:
|
|
conn.close()
|
|
self.assertEqual(inline, 40)
|
|
# Reclaim space ran alone afterwards and gave the space back.
|
|
self.assertIsNotNone(reclaim)
|
|
self.assertEqual(reclaim.admission, "exclusive")
|
|
self.assertFalse(reclaim.cancellable)
|
|
self.assertEqual(reclaim.state, "finished", reclaim.log)
|
|
self.assertIn("Reclaim space started", "\n".join(job.log))
|
|
self.assertLess(self.database_bytes(), size_before / 3)
|
|
wal = self.api.db.model.db_path + "-wal"
|
|
self.assertEqual(os.path.getsize(wal) if os.path.exists(wal) else 0, 0)
|
|
self.assertRegex(reclaim.status_line, r"^Reclaimed [\d.,]+ MB · ")
|
|
listed = await self.client.request(
|
|
"GET", f"/prompt_manager/prompts/{prompt_id}/images"
|
|
)
|
|
self.assertEqual(await listed.json(), listed_before)
|
|
self.assertEqual(self.strays(), [backup]) # only the kept backup
|
|
self.assertNotIn(self.tmpdir, json.dumps(job.report()))
|
|
|
|
async def test_a_backup_failure_rewrites_nothing(self):
|
|
self.old_library()
|
|
from py.api import workflow_storage as api_storage
|
|
|
|
with patch.object(
|
|
api_storage,
|
|
"take_safety_backup",
|
|
side_effect=maintenance.SafetyBackupFailed("no"),
|
|
):
|
|
_plan, _summary, job = await self.start()
|
|
await self.finish(job)
|
|
self.assertEqual(job.state, "failed")
|
|
self.assertEqual(job.last_event["code"], "backup-failed")
|
|
self.assertEqual(self.inline_count(), 12)
|
|
self.assertIsNone(self.reclaim_job())
|
|
|
|
async def test_cancel_stops_after_the_current_batch_and_a_rerun_continues(self):
|
|
self.old_library()
|
|
before = self.texts()
|
|
real = self.api.db.migrate_workflow_batch
|
|
holder = {}
|
|
|
|
def first_then_cancel(*args, **kwargs):
|
|
result = real(*args, **kwargs)
|
|
holder["job"].cancelled = True # Cancel pressed during batch one
|
|
return result
|
|
|
|
with (
|
|
patch.object(ws, "BATCH_ROWS", 5),
|
|
patch.object(
|
|
self.api.db, "migrate_workflow_batch", side_effect=first_then_cancel
|
|
),
|
|
):
|
|
plan = await self.preview(OPERATION)
|
|
body = await self.apply(plan["token"])
|
|
holder["job"] = job = self.api._jobs.get(body["job"]["id"])
|
|
await self.finish(job)
|
|
self.assertEqual(job.state, "cancelled")
|
|
self.assertEqual(job.counters["changed"], 5)
|
|
self.assertEqual(self.inline_count(), 7)
|
|
self.assertIn("run it again to continue", job.status_line)
|
|
self.assertIsNone(self.reclaim_job())
|
|
self.assertEqual(self.texts(), before)
|
|
_plan, _summary, again = await self.start()
|
|
await self.finish(again)
|
|
self.assertEqual(again.state, "finished", again.log)
|
|
self.assertEqual(again.counters["changed"], 7)
|
|
self.assertEqual(self.inline_count(), 0)
|
|
self.assertEqual(self.texts(), before)
|
|
|
|
async def test_a_failed_batch_fails_the_job_and_every_record_still_reads(self):
|
|
self.old_library()
|
|
before = self.texts()
|
|
real = self.api.db.migrate_workflow_batch
|
|
calls = []
|
|
|
|
def failing(*args, **kwargs):
|
|
calls.append(1)
|
|
if len(calls) == 2:
|
|
raise sqlite3.OperationalError("disk I/O error")
|
|
return real(*args, **kwargs)
|
|
|
|
with (
|
|
patch.object(ws, "BATCH_ROWS", 5),
|
|
patch.object(self.api.db, "migrate_workflow_batch", side_effect=failing),
|
|
):
|
|
_plan, _summary, job = await self.start()
|
|
await self.finish(job)
|
|
self.assertEqual(job.state, "failed")
|
|
self.assertIn(
|
|
"The batch in progress was rolled back; every record reads as before.",
|
|
job.status_line,
|
|
)
|
|
self.assertNotIn("disk I/O", job.status_line)
|
|
self.assertEqual(job.counters["changed"], 5)
|
|
self.assertEqual(job.counters["errors"], 1)
|
|
self.assertEqual(self.texts(), before)
|
|
self.assertIsNone(self.reclaim_job())
|
|
|
|
async def test_reclaim_space_that_cannot_start_says_why(self):
|
|
self.old_library()
|
|
release = asyncio.Event()
|
|
|
|
async def held(job):
|
|
await release.wait()
|
|
|
|
reader = self.api._jobs.start(
|
|
"health", "Health check", admission="read", runner=held
|
|
)
|
|
try:
|
|
_plan, _summary, job = await self.start()
|
|
await self.finish(job)
|
|
finally:
|
|
release.set()
|
|
await asyncio.wait_for(reader.task, WAIT)
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertIsNone(self.reclaim_job())
|
|
self.assertIn(
|
|
"Reclaim space did not start: Health check is running. "
|
|
"Run Optimize database to give the space back.",
|
|
"\n".join(job.log),
|
|
)
|
|
self.assertIn("Reclaim space did not start", job.status_line)
|
|
|
|
async def test_reclaim_space_checks_the_disk_first(self):
|
|
self.old_library()
|
|
before = self.texts()
|
|
_plan, _summary, job = await self.start()
|
|
free = {"bytes": 10**12}
|
|
|
|
def free_bytes(_path):
|
|
return free["bytes"]
|
|
|
|
with patch.object(maintenance, "_free_bytes", side_effect=free_bytes):
|
|
await asyncio.wait_for(asyncio.shield(job.task), WAIT)
|
|
reclaim = self.reclaim_job()
|
|
free["bytes"] = 1024
|
|
# The reclaim job checks when it runs, after the rewrite.
|
|
await asyncio.wait_for(asyncio.shield(reclaim.task), WAIT)
|
|
self.assertEqual(reclaim.state, "failed")
|
|
self.assertEqual(reclaim.last_event["code"], "insufficient-space")
|
|
self.assertIn("Not enough free disk space", reclaim.status_line)
|
|
self.assertEqual(self.texts(), before)
|
|
|
|
async def test_apply_refuses_up_front_when_the_backup_does_not_fit(self):
|
|
self.old_library()
|
|
plan = await self.preview(OPERATION)
|
|
with patch.object(maintenance, "_free_bytes", return_value=1024):
|
|
body = await self.apply(plan["token"], status=507)
|
|
self.assertEqual(body["code"], "insufficient-space")
|
|
self.assertEqual(self.inline_count(), 12)
|
|
|
|
async def test_the_jobs_never_open_a_media_file(self):
|
|
self.old_library(media=True)
|
|
media = str(self.output_dir)
|
|
before = {
|
|
name: (
|
|
_digest(os.path.join(media, name)),
|
|
os.stat(os.path.join(media, name)),
|
|
)
|
|
for name in os.listdir(media)
|
|
}
|
|
opened = []
|
|
real_open, real_os_open = builtins.open, os.open
|
|
|
|
def tracking_open(file, *args, **kwargs):
|
|
opened.append(str(file))
|
|
return real_open(file, *args, **kwargs)
|
|
|
|
def tracking_os_open(path, *args, **kwargs):
|
|
opened.append(str(path))
|
|
return real_os_open(path, *args, **kwargs)
|
|
|
|
with patch("builtins.open", tracking_open), patch("os.open", tracking_os_open):
|
|
with patch.object(Image, "open", side_effect=AssertionError("opened")):
|
|
_plan, _summary, job = await self.start()
|
|
reclaim = await self.finish(job)
|
|
self.assertEqual(job.state, "finished", job.log)
|
|
self.assertEqual(reclaim.state, "finished", reclaim.log)
|
|
self.assertEqual([path for path in opened if path.startswith(media)], [])
|
|
for name, (digest, stat) in before.items():
|
|
path = os.path.join(media, name)
|
|
self.assertEqual(_digest(path), digest)
|
|
self.assertEqual(os.stat(path).st_mtime_ns, stat.st_mtime_ns)
|
|
|
|
async def test_a_second_apply_of_the_token_is_plan_used(self):
|
|
self.old_library(count=4)
|
|
plan = await self.preview(OPERATION)
|
|
body = await self.apply(plan["token"])
|
|
await self.finish(self.api._jobs.get(body["job"]["id"]))
|
|
again = await self.apply(plan["token"], status=409)
|
|
self.assertEqual(again["code"], "plan-used")
|
|
|
|
|
|
class TestWorkflowStorageOutcomes(WorkflowStorageTestCase):
|
|
async def test_kept_values_and_swept_documents_are_logged(self):
|
|
prompt_id = self.old_library(count=4)
|
|
conn = self.api.db.model.get_connection()
|
|
with conn:
|
|
image_id = conn.execute(
|
|
"INSERT INTO generated_images (prompt_id, image_path, filename, "
|
|
"file_path, workflow_data) VALUES (?, '/x/b', 'b', '/x/b', x'00')",
|
|
(prompt_id,),
|
|
).lastrowid
|
|
store_text(conn, "a document no record references")
|
|
_plan, _summary, job = await self.start()
|
|
await self.finish(job)
|
|
log = "\n".join(job.log)
|
|
self.assertEqual(job.counters["protected"], 1)
|
|
self.assertIn(
|
|
"Rewrote 4 image records; 1 stored value not text, kept as is", log
|
|
)
|
|
self.assertIn("Removed 1 unused workflow document", log)
|
|
# #586: the kept value has its row, in the review's wording.
|
|
self.assertEqual(
|
|
[row["detail"] for row in job.outcomes["protected"]],
|
|
[f"Image record {image_id}: workflow, stored value not text; kept as is"],
|
|
)
|
|
report = job.report()["outcomes"]["protected"]
|
|
self.assertEqual((report["count"], report["truncated"]), (1, False))
|
|
|
|
async def reclaim_with(self, **patches):
|
|
self.old_library(count=4)
|
|
before = self.texts()
|
|
with contextlib.ExitStack() as stack:
|
|
for name, value in patches.items():
|
|
stack.enter_context(patch.object(self.api.db, name, value))
|
|
_plan, _summary, job = await self.start()
|
|
reclaim = await self.finish(job)
|
|
self.assertEqual(self.texts(), before)
|
|
return reclaim
|
|
|
|
async def test_a_failed_reclaim_says_the_database_is_unchanged(self):
|
|
reclaim = await self.reclaim_with(
|
|
reclaim_space=lambda: (_ for _ in ()).throw(
|
|
sqlite3.OperationalError("database is locked")
|
|
)
|
|
)
|
|
self.assertEqual(reclaim.state, "failed")
|
|
self.assertEqual(
|
|
reclaim.status_line, "Reclaim space failed; the database is unchanged."
|
|
)
|
|
|
|
async def test_a_reclaim_whose_disk_cannot_be_measured_fails_cleanly(self):
|
|
from py.api import workflow_storage as api_storage
|
|
|
|
with patch.object(api_storage, "live_bytes", side_effect=OSError("gone")):
|
|
reclaim = await self.reclaim_with()
|
|
self.assertEqual(reclaim.state, "failed")
|
|
self.assertNotIn("gone", reclaim.status_line)
|
|
|
|
async def test_a_busy_checkpoint_says_the_rest_follows_a_restart(self):
|
|
reclaim = await self.reclaim_with(
|
|
reclaim_space=lambda: {
|
|
"before_bytes": 3 * 2**20,
|
|
"after_bytes": 2**20,
|
|
"documents_removed": 0,
|
|
"checkpoint_busy": True,
|
|
}
|
|
)
|
|
self.assertEqual(
|
|
reclaim.status_line,
|
|
"Reclaimed 2.0 MB · the database is now 1.0 MB "
|
|
"(the rest after the next restart)",
|
|
)
|