Vito Sansevero
96cd973ebd
Final review fixes M1, M2, M3 and M5 for the workflow storage dedupe: - M1: a batch also ends after 0.5 s, and the loop leaves the write lock free for 0.15 s after each committed batch (longer than SQLite's longest busy poll), so the image monitor and run counting on another connection get in between. RED: a threaded writer with a 2 s busy timeout hit "database is locked" (4 of 5 links landed). - M2: each row runs in its own savepoint. A row that cannot be rewritten (undecodable text, a text that would not come back) is rolled back alone, stays inline and is listed in errors; the rest of the batch goes on. - M3: a reused document is decompressed and checked before its id is handed out, and repaired from the text (whose hash is its key) when damaged; verify_database_file reads every document and refuses a damaged one. - M5: the two legacy generated_images rebuilds refuse a file that already has the document id columns instead of dropping them. - The preview names the pre-migration backup and says it is always taken.
403 lines
15 KiB
Python
403 lines
15 KiB
Python
"""The workflow storage migration and Reclaim space (#249 S6 carry-over).
|
|
|
|
``migrate_workflow_storage`` rewrites old-form image rows (workflow and
|
|
prompt graph text inline) to the new form (a ``workflow_documents`` id,
|
|
``database/workflow_documents.py``) in batches:
|
|
|
|
- each batch is one ``BEGIN IMMEDIATE`` transaction over at most
|
|
``BATCH_ROWS`` rows or about ``BATCH_BYTES`` of text, found by a rowid
|
|
range search (``NEXT_IDS_SQL``) and read one row at a time, so memory is
|
|
bounded by one row plus the batch's distinct documents;
|
|
- each text is stored (or its document reused) and read back through the
|
|
one accessor before its row changes; a text that would not come back
|
|
identical fails the batch (``DocumentMismatch``) and rolls it back whole;
|
|
- a value that is not text (or is over ``MAX_DOCUMENT_BYTES``) stays inline
|
|
and is counted as protected;
|
|
- each row runs in its own savepoint: one that cannot be rewritten (text
|
|
SQLite cannot decode, a text that would not come back) stays inline and
|
|
is listed, and the rest of the batch goes on;
|
|
- a batch ends after ``BATCH_SECONDS``, and the write lock is then left
|
|
free for ``BATCH_PAUSE``, so other writers are never starved.
|
|
|
|
A crash or a stop between batches leaves committed batches in the new form
|
|
and the rest in the old form; both read the same, and a rerun continues
|
|
from the first old-form row. Only database rows change: no image or video
|
|
file is opened.
|
|
|
|
``reclaim_space`` is the VACUUM afterwards: it sweeps unreferenced
|
|
documents, rebuilds the file and checkpoints the WAL with TRUNCATE, so the
|
|
freed pages go back to the disk.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import sqlite3
|
|
import time
|
|
from contextlib import closing
|
|
|
|
from .workflow_documents import (
|
|
TEXT_COLUMNS,
|
|
DocumentMismatch,
|
|
read_stored_text,
|
|
release_documents,
|
|
storable,
|
|
store_text,
|
|
sweep_documents,
|
|
)
|
|
|
|
OPERATION = "workflow_storage"
|
|
BATCH_ROWS = 200
|
|
BATCH_BYTES = 16 * 2**20 # characters of text rewritten per transaction
|
|
# A batch holds the write lock: it ends after this long, and the loop then
|
|
# leaves the lock free for BATCH_PAUSE, longer than SQLite's longest busy
|
|
# poll (100 ms), so the image monitor and the run counting get in between.
|
|
BATCH_SECONDS = 0.5
|
|
BATCH_PAUSE = 0.15
|
|
# What one row's rewrite may raise without failing its batch: the row stays
|
|
# inline and is reported (text SQLite cannot decode, a text that would not
|
|
# come back identical).
|
|
ROW_ERRORS = (sqlite3.Error, DocumentMismatch, UnicodeError, ValueError, TypeError)
|
|
|
|
logger = logging.getLogger("prompt_manager.database.workflow_storage")
|
|
_clock = time.monotonic
|
|
|
|
# typeof() reads a column's type from the record header: no overflow page.
|
|
NEXT_IDS_SQL = (
|
|
"SELECT id FROM generated_images WHERE id>? AND "
|
|
"(typeof(workflow_data)!='null' OR typeof(prompt_metadata)!='null') "
|
|
"ORDER BY id LIMIT ?"
|
|
)
|
|
ROW_SQL = (
|
|
"SELECT workflow_data, prompt_metadata, workflow_document_id, "
|
|
"prompt_document_id FROM generated_images WHERE id=?"
|
|
)
|
|
UPDATE_SQL = (
|
|
"UPDATE generated_images SET {assignments} WHERE id=?" # columns are constants
|
|
)
|
|
# One pass over the record headers: rows with inline text, and inline
|
|
# values that are not text (they stay inline).
|
|
INLINE_COUNTS_SQL = (
|
|
"SELECT "
|
|
"SUM(typeof(workflow_data)='text' OR typeof(prompt_metadata)='text'), "
|
|
"SUM(typeof(workflow_data) NOT IN ('null','text')) "
|
|
"+ SUM(typeof(prompt_metadata) NOT IN ('null','text')) "
|
|
"FROM generated_images"
|
|
)
|
|
# octet_length() (SQLite 3.43+) reads no overflow page either.
|
|
INLINE_BYTES_SQL = (
|
|
"SELECT SUM(octet_length(workflow_data)), SUM(octet_length(prompt_metadata)) "
|
|
"FROM generated_images"
|
|
)
|
|
|
|
QUERY_PLAN_STATEMENTS = {
|
|
"next old-form rows": (NEXT_IDS_SQL, (0, BATCH_ROWS)),
|
|
"one row": (ROW_SQL, (1,)),
|
|
}
|
|
|
|
|
|
def _rewrite_row(conn, image_id, row, verified, checked):
|
|
"""Store the row's inline texts as documents; returns (changed, protected, size).
|
|
|
|
``verified`` maps document ids already read back in this batch to their
|
|
text, and ``checked`` holds the reused ids ``store_text`` already
|
|
checked, so a document shared by many rows is decompressed once.
|
|
"""
|
|
changes, replaced = {}, []
|
|
protected = size = 0
|
|
for index, (column, reference) in enumerate(TEXT_COLUMNS):
|
|
value, current = row[index], row[index + len(TEXT_COLUMNS)]
|
|
if value is None:
|
|
continue
|
|
if not storable(value):
|
|
protected += 1
|
|
continue
|
|
document_id = store_text(conn, value, checked)
|
|
if document_id not in verified:
|
|
verified[document_id] = read_stored_text(conn, None, document_id)
|
|
if verified[document_id] != value:
|
|
raise DocumentMismatch(f"Image {image_id}: {column} would change")
|
|
changes[column], changes[reference] = None, document_id
|
|
if current is not None and current != document_id:
|
|
replaced.append(current)
|
|
size += len(value)
|
|
if changes:
|
|
assignments = ",".join(f"{key}=?" for key in changes)
|
|
conn.execute(
|
|
UPDATE_SQL.format(assignments=assignments),
|
|
[*changes.values(), image_id],
|
|
)
|
|
release_documents(conn, replaced)
|
|
return bool(changes), protected, size
|
|
|
|
|
|
def migrate_batch(
|
|
conn,
|
|
after_id,
|
|
max_rows=BATCH_ROWS,
|
|
max_bytes=BATCH_BYTES,
|
|
max_seconds=BATCH_SECONDS,
|
|
):
|
|
"""Rewrite the next batch after ``after_id`` in one transaction.
|
|
|
|
The batch ends after ``max_rows`` rows, about ``max_bytes`` of text or
|
|
``max_seconds``, whichever comes first. Each row runs in its own
|
|
savepoint: a row that cannot be rewritten (``ROW_ERRORS``) is rolled
|
|
back alone, stays inline and is listed in ``errors``; anything else rolls
|
|
the whole batch back and propagates. Returns ``{last_id, rows,
|
|
protected, errors, done}``: the last row looked at, rows rewritten,
|
|
inline values kept (not text), the ids kept by an error and whether
|
|
there was nothing left to look at.
|
|
"""
|
|
if conn.in_transaction:
|
|
raise RuntimeError("The migration needs its own transaction")
|
|
started = _clock()
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
try:
|
|
ids = [row[0] for row in conn.execute(NEXT_IDS_SQL, (after_id, max_rows))]
|
|
result = {
|
|
"last_id": after_id,
|
|
"rows": 0,
|
|
"protected": 0,
|
|
"errors": [],
|
|
"done": not ids,
|
|
}
|
|
verified, checked, size = {}, set(), 0
|
|
for image_id in ids:
|
|
conn.execute("SAVEPOINT workflow_row")
|
|
try:
|
|
row = conn.execute(ROW_SQL, (image_id,)).fetchone()
|
|
changed, protected, row_size = _rewrite_row(
|
|
conn, image_id, row, verified, checked
|
|
)
|
|
except ROW_ERRORS as error:
|
|
conn.execute("ROLLBACK TO workflow_row")
|
|
conn.execute("RELEASE workflow_row")
|
|
verified.clear() # the row's documents were rolled back too
|
|
checked.clear()
|
|
logger.warning(f"Image record {image_id} kept inline: {error}")
|
|
result["errors"].append(image_id)
|
|
changed, protected, row_size = False, 0, 0
|
|
else:
|
|
conn.execute("RELEASE workflow_row")
|
|
result["rows"] += int(changed)
|
|
result["protected"] += protected
|
|
result["last_id"] = image_id
|
|
size += row_size
|
|
if size >= max_bytes or _clock() - started >= max_seconds:
|
|
break
|
|
conn.commit()
|
|
return result
|
|
except BaseException:
|
|
conn.rollback()
|
|
raise
|
|
|
|
|
|
def _file_bytes(db_path):
|
|
return sum(
|
|
os.path.getsize(path)
|
|
for path in (db_path, db_path + "-wal")
|
|
if os.path.exists(path)
|
|
)
|
|
|
|
|
|
class WorkflowStorageOperations:
|
|
"""PromptDatabase mixin: the workflow storage migration and Reclaim space."""
|
|
|
|
def migrate_workflow_storage(
|
|
self,
|
|
*,
|
|
on_batch=None,
|
|
should_stop=None,
|
|
max_rows=BATCH_ROWS,
|
|
max_bytes=BATCH_BYTES,
|
|
) -> dict:
|
|
"""Rewrite every old-form row, batch by batch (blocking).
|
|
|
|
``on_batch(progress)`` runs after each committed batch with the
|
|
running totals ``{rows, protected, batches, last_id}``;
|
|
``should_stop()`` is asked before each batch. Each committed batch is
|
|
followed by ``BATCH_PAUSE`` with the write lock free. Returns
|
|
``{rows, protected, batches, stopped, errors, documents_removed}``
|
|
(``errors``: the ids of rows kept inline by an error). Unreferenced
|
|
documents are swept once every row was looked at.
|
|
"""
|
|
totals = {"rows": 0, "protected": 0, "batches": 0, "last_id": 0}
|
|
errors = []
|
|
stopped = False
|
|
while True:
|
|
if should_stop is not None and should_stop():
|
|
stopped = True
|
|
break
|
|
batch = self.migrate_workflow_batch(totals["last_id"], max_rows, max_bytes)
|
|
if batch["done"]:
|
|
break
|
|
time.sleep(BATCH_PAUSE) # leave the write lock to other writers
|
|
errors += batch["errors"]
|
|
totals["rows"] += batch["rows"]
|
|
totals["protected"] += batch["protected"]
|
|
totals["batches"] += 1
|
|
totals["last_id"] = batch["last_id"]
|
|
if on_batch is not None:
|
|
on_batch(dict(totals))
|
|
removed = 0 if stopped else self.sweep_workflow_documents()
|
|
return {
|
|
"rows": totals["rows"],
|
|
"protected": totals["protected"],
|
|
"batches": totals["batches"],
|
|
"stopped": stopped,
|
|
"errors": errors,
|
|
"documents_removed": removed,
|
|
}
|
|
|
|
def migrate_workflow_batch(
|
|
self, after_id, max_rows=BATCH_ROWS, max_bytes=BATCH_BYTES
|
|
) -> dict:
|
|
"""One batch on the calling thread's connection (``migrate_batch``)."""
|
|
return migrate_batch(self.model.get_connection(), after_id, max_rows, max_bytes)
|
|
|
|
def sweep_workflow_documents(self) -> int:
|
|
"""Remove every document no image row references (one transaction)."""
|
|
conn = self.model.get_connection()
|
|
with conn:
|
|
return sweep_documents(conn)
|
|
|
|
def workflow_storage_status(self) -> dict:
|
|
"""``workflow_storage_counts`` on one read-only snapshot (blocking)."""
|
|
from .maintenance import _read_snapshot
|
|
|
|
return _read_snapshot(self.model.db_path, self.workflow_storage_counts)
|
|
|
|
def workflow_storage_counts(self, conn) -> dict:
|
|
"""``{rows, protected, documents, inline_bytes}`` on ``conn`` (reads only).
|
|
|
|
``inline_bytes`` is None where SQLite has no ``octet_length``.
|
|
"""
|
|
rows, protected = conn.execute(INLINE_COUNTS_SQL).fetchone()
|
|
documents = conn.execute("SELECT COUNT(*) FROM workflow_documents").fetchone()
|
|
inline_bytes = None
|
|
if sqlite3.sqlite_version_info >= (3, 43, 0):
|
|
row = conn.execute(INLINE_BYTES_SQL).fetchone()
|
|
inline_bytes = sum(value or 0 for value in row)
|
|
return {
|
|
"rows": rows or 0,
|
|
"protected": protected or 0,
|
|
"documents": documents[0],
|
|
"inline_bytes": inline_bytes,
|
|
}
|
|
|
|
def reclaim_space(self) -> dict:
|
|
"""Sweep, VACUUM and truncate the WAL (blocking); file bytes before/after.
|
|
|
|
The caller holds the registry's ``exclusive`` slot. Raises whatever
|
|
stopped VACUUM; the database is then unchanged.
|
|
"""
|
|
db_path = self.model.db_path
|
|
removed = self.sweep_workflow_documents()
|
|
before = _file_bytes(db_path)
|
|
self.model.vacuum_database()
|
|
with closing(sqlite3.connect(db_path)) as checkpoint:
|
|
busy = checkpoint.execute("PRAGMA wal_checkpoint(TRUNCATE)").fetchone()[0]
|
|
return {
|
|
"before_bytes": before,
|
|
"after_bytes": _file_bytes(db_path),
|
|
"documents_removed": removed,
|
|
"checkpoint_busy": bool(busy),
|
|
}
|
|
|
|
# ── preview ───────────────────────────────────────────────────────
|
|
|
|
def _plan_workflow_storage(self):
|
|
"""The reviewed plan of the migration (one read-only snapshot)."""
|
|
from .maintenance import (
|
|
MaintenanceDraft,
|
|
_Fingerprint,
|
|
_megabytes,
|
|
_plural,
|
|
_read_snapshot,
|
|
)
|
|
|
|
def read(conn):
|
|
sizes = [
|
|
conn.execute(f"PRAGMA {name}").fetchone()[0]
|
|
for name in ("page_size", "page_count", "freelist_count")
|
|
]
|
|
return sizes, self.workflow_storage_counts(conn)
|
|
|
|
(page_size, page_count, free_pages), counts = _read_snapshot(
|
|
self.model.db_path, read
|
|
)
|
|
size = page_size * page_count
|
|
details = [
|
|
{
|
|
"label": "Database file",
|
|
"detail": f"{_megabytes(size)}, of which "
|
|
f"{_megabytes(free_pages * page_size)} is free pages",
|
|
},
|
|
{
|
|
"label": "Image records",
|
|
"detail": f"{_plural(counts['rows'], 'record')} keep their workflow "
|
|
"or prompt graph inline; each distinct text is stored once, "
|
|
"compressed",
|
|
},
|
|
]
|
|
if counts["inline_bytes"] is not None:
|
|
details.append(
|
|
{
|
|
"label": "Inline text",
|
|
"detail": f"{_megabytes(counts['inline_bytes'])} now",
|
|
}
|
|
)
|
|
if counts["documents"]:
|
|
details.append(
|
|
{
|
|
"label": "Already stored once",
|
|
"detail": _plural(counts["documents"], "workflow document"),
|
|
}
|
|
)
|
|
name = os.path.basename(self.model.db_path)
|
|
details += [
|
|
{
|
|
"label": "Backup",
|
|
"detail": BACKUP_NOTE.format(
|
|
name=f"{name}.pre-workflow-storage_<date and time>"
|
|
),
|
|
},
|
|
{"label": "Reclaim space", "detail": RECLAIM_NOTE},
|
|
{"label": "Older versions", "detail": OLDER_VERSIONS_NOTE},
|
|
]
|
|
protected = counts["protected"]
|
|
return MaintenanceDraft(
|
|
operation=OPERATION,
|
|
fingerprint=_Fingerprint(OPERATION).hexdigest(),
|
|
targets=[],
|
|
scope={
|
|
"count": counts["rows"],
|
|
"label": _plural(counts["rows"], "image record") + " to rewrite",
|
|
},
|
|
protected={
|
|
"count": protected,
|
|
"reasons": (
|
|
[f"{_plural(protected, 'stored value')} not text; kept as is"]
|
|
if protected
|
|
else []
|
|
),
|
|
},
|
|
media={"deletes_files": False, "label": MEDIA_NOTE},
|
|
details=details,
|
|
)
|
|
|
|
|
|
MEDIA_NOTE = "No image or video file is opened or changed; only the database."
|
|
RECLAIM_NOTE = (
|
|
"When every record is rewritten, Reclaim space runs alone to give the freed "
|
|
"space back to the disk (VACUUM). The database is locked while it runs."
|
|
)
|
|
BACKUP_NOTE = (
|
|
"Always taken first, whatever the safety backup setting: {name} beside "
|
|
"the database, verified, and kept until you remove it (it is never pruned)."
|
|
)
|
|
OLDER_VERSIONS_NOTE = (
|
|
"An older Prompt Manager version shows rewritten images without their "
|
|
"workflow. Restoring the pre-migration backup brings back the old form."
|
|
)
|