Files
ComfyAssets-ComfyUI_PromptM…/py/api/thumbnail_move.py
T
Vito Sansevero a33b0f4648 feat(thumbnails): move thumbnails between the sources and one folder (#320)
POST /images/thumbnails/plan reviews a move (what is copied, removed and
left, free space, too-long Windows paths, sources no longer listed) and
hands out a token; POST /jobs/thumbnails {mode: move} runs it as a
thumbnails job: hold automatic previews, copy and check every preview on
the job pool, switch with one config commit that records the old copies,
then remove them. Per-file errors still switch; a folder-wide error, a
cancel before the switch or a crash leaves nothing changed (a journal is
undone at the next start); old copies that could not go are removed by a
finishing job at the next start. Leaving a folder removes the previews of
this install's sources that are no longer listed (Q320-8).

Recreate folder and Use this folder answer an unavailable folder; the
Thumbnail maintenance scope prune removes the thumbnails of removed
sources.
2026-10-31 20:29:06 -07:00

524 lines
21 KiB
Python

"""Move thumbnails (#320 Part 3): the reviewed job that changes where previews
live, the removal of old copies after it, recovery at startup and the
"Remove thumbnails of removed sources" scope (Part 5).
Kind ``thumbnails`` with ``mode: "move"`` (D1A): admission ``guarded``, so
scans, deletes, tagging, source changes and every other thumbnail job wait
for it, while the monitor's Needs review recording runs beside it.
Automatic previews are held for the whole job (D2A, ``_auto_thumbnail_blocker``).
Phases: review (``POST /images/thumbnails/plan``, no writes but the
validation probe), start (a plan token), hold (generation bump, journal,
target folder), copy (cancellable at every file boundary, on the job pool
in small chunks, D3A), switch (one config commit sets the location and
records the old stores under ``cleanup``; not cancellable from here),
remove old copies, end. Per-file errors switch anyway (D10A); a folder-wide
error or a cancel before the switch removes what the job wrote and changes
nothing. The move is the only writer of the setting (D4B).
Moving away from a thumbnail folder also removes the previews of this
install's sources that are no longer listed (Q320-8: remove).
"""
import functools
import os
import time
from pathlib import Path
from aiohttp import web
from . import preview_location, preview_store
from . import thumbnail_move_files as files
from .jobs import THUMBNAILS_JOB_KIND
from .thumbnail_move_plan import (
MOVE_GROUPS,
MOVE_TITLE,
PLAN_TTL,
_number,
_plural,
_remove,
_roots_revision,
_volume,
_is_network,
_write_json,
)
from . import thumbnail_move_plan as move_plan
try:
from ...utils import exit_signal
from ...utils.parallel import map_parallel
except ImportError: # tests import the package top-level
from utils import exit_signal
from utils.parallel import map_parallel
class ThumbnailMoveMixin:
"""Start and run the move job; the removal of old copies."""
# ── start ─────────────────────────────────────────────────────────
def _take_move_plan(self, body):
"""The plan a start names, or ``(status, code)``."""
token = body.get("token")
plan = (self._move_plans or {}).get(token) if isinstance(token, str) else None
if plan is None or time.monotonic() - plan.created > PLAN_TTL:
return 409, "plan-stale"
roots = [str(root) for root in self._get_all_output_dirs()]
if _roots_revision(roots) != plan.revision:
return 409, "plan-stale"
if preview_store.setting() != plan.old:
return 409, "plan-stale"
if plan.confirm_required and body.get("confirmed") is not True:
return 400, "confirmation-required"
return plan
def start_thumbnail_move(self, body):
"""Start the move job (sync, on the loop); raises JobBusy. Returns
the job or a refusal response."""
if not self._move_recovered:
blocker = self._jobs.active(THUMBNAILS_JOB_KIND)
return web.json_response(
{
"success": False,
"code": "busy",
"error": "Prompt Manager is still checking an earlier move. "
"Try again in a moment.",
"running": None if blocker is None else {"id": blocker.id},
},
status=409,
)
taken = self._take_move_plan(body)
if isinstance(taken, tuple):
status, code = taken
text = {
"plan-stale": "The review is out of date. Review the move again.",
"confirmation-required": "Confirm that Prompt Manager may "
"create its folders there.",
}[code]
return web.json_response(
{"success": False, "code": code, "error": text}, status=status
)
plan = taken
self._move_plans.pop(plan.token, None)
return self._jobs.start(
THUMBNAILS_JOB_KIND,
MOVE_TITLE,
admission="guarded",
cancellable=True,
scope={
"mode": "move",
"from": plan.old.mode,
"to": plan.target["mode"],
"to_path": plan.target["path"],
},
groups=MOVE_GROUPS,
runner=functools.partial(self._run_thumbnail_move, plan=plan),
)
# ── the job ───────────────────────────────────────────────────────
async def _run_thumbnail_move(self, job, *, plan):
state = {"written": [], "created": [], "created_folder": False}
try:
await self._run_in_executor(self._move_hold, job, plan, state)
except (OSError, ValueError) as error:
await self._run_in_executor(self._move_rollback, plan, state)
self._end_move(
job, "failed", "The new folder could not be made; nothing changed."
)
self.logger.warning(f"Move thumbnails could not start: {error}")
return
items = [
(source, item)
for source in plan.sources
if source.scan is not None
for item in source.scan.copy
]
total = len(items)
job.update(total=total, processed=0, progress=0)
chunk = self._move_chunk(plan)
workers = self._move_workers()
done = 0
for start in range(0, total, chunk):
if job.cancelled:
return await self._move_cancelled(job, plan, state)
batch = items[start : start + chunk]
try:
results = await self._run_in_executor(
map_parallel,
functools.partial(self._move_copy_one, state=state),
batch,
workers,
)
except exit_signal.ExitRequested:
return # exiting: the next start undoes the copies (journal)
except files.FolderGone as error:
self.logger.warning(f"Move thumbnails stopped: {error}")
await self._run_in_executor(self._move_rollback, plan, state)
self._end_move(
job,
"failed",
"The new folder could not be reached or is full; nothing changed.",
)
return
for (source, item), (outcome, detail) in zip(batch, results):
row = {
"path": "/".join(item[0]),
"root": source.label,
"detail": detail,
}
job.record(
"moved" if outcome == "moved" else outcome,
row,
f"{outcome}: {source.label} · {'/'.join(item[0])}",
)
done += len(batch)
label = batch[-1][0].label
job.update(
processed=done,
progress=int(done * 90 / total) if total else 90,
status_line=(
f"Copying {_number(done)} of {_number(total)} previews · {label}"
),
)
job.publish({"type": "progress", "summary": job.summary()})
if job.cancelled:
return await self._move_cancelled(job, plan, state)
# The switch: not cancellable from here (an exit still stops between
# files and leaves the cleanup to the next start).
job.cancellable = False
job.update(status_line="Switching to the new location", progress=92)
job.publish({"type": "progress", "summary": job.summary()})
try:
await self._run_in_executor(self._move_switch, plan)
except Exception:
self.logger.exception("Move thumbnails could not switch")
await self._run_in_executor(self._move_rollback, plan, state)
self._end_move(
job, "failed", "The setting could not be saved; nothing changed."
)
return
await self._remove_old_copies(job, plan)
moved = job.counters.get("moved", 0) + job.counters.get("skipped", 0)
errors = job.counters.get("errors", 0)
where = plan.target["path"] or "inside each source"
parts = [f"{_number(moved)} moved"]
if errors:
parts.append(_plural(errors, "error"))
parts.append(
"old copies removed"
if not job.scope.get("cleanup_left")
else "some old copies are removed at the next start"
)
self._end_move(
job, "finished", f"Thumbnails moved to {where}: " + " · ".join(parts) + "."
)
def _end_move(self, job, state, line):
job.update(status_line=line, progress=100)
job.log_line(line)
job.publish({"type": "complete", "summary": job.summary()})
job.finish(state)
def _move_workers(self):
from .images import _worker_threads
return max(1, int(_worker_threads() or 1))
def _move_chunk(self, plan):
"""2 files at a time across volumes or to/from a share, else 4, never
more than Worker threads (D3A: the pool keeps its size)."""
sides = [plan.target["path"] or ""] + [s.old_store or "" for s in plan.sources]
sides += [s.new_store for s in plan.sources]
network = any(_is_network(side) for side in sides if side)
volumes = {_volume(side) for side in sides if side}
size = 2 if network or len(volumes) > 1 else 4
return max(1, min(size, self._move_workers()))
def _move_hold(self, job, plan, state):
"""Phase 3: bump the generation, write the journal, make the target."""
preview_store.bump_generation()
if plan.to_folder:
folder = plan.target["path"]
if not os.path.isdir(folder):
if plan.target_default:
os.makedirs(folder, exist_ok=True)
else:
os.mkdir(folder)
state["created_folder"] = True
plan.target_id = preview_location.init_folder(
folder, created_by_pm=state["created_folder"]
)
install = plan.install
def add_keys(marker):
for source in plan.sources:
if source.key and source.key not in marker["keys"]:
marker["keys"][source.key] = {
"path": source.root,
"identity": os.path.normcase(os.path.realpath(source.root)),
"dev_ino": preview_location.file_identity(source.root),
"install": install,
}
with preview_store.write_lock:
preview_location.write_marker(folder, add_keys)
_write_json(
move_plan.journal_path(),
{
"state": "copying",
"from": plan.old.to_config(),
"to": plan.target,
"stores": [s.new_store for s in plan.sources],
# Stores that do not exist yet are made by this job only:
# after a crash everything in them goes (B4).
"created_dirs": [
s.new_store for s in plan.sources if not os.path.isdir(s.new_store)
],
"created_folder": state["created_folder"],
"install": plan.install,
},
)
job.log_line(
f"Move thumbnails started: {plan.old.path or 'inside each source'} → "
f"{plan.target['path'] or 'inside each source'}"
)
def _move_copy_one(self, pair, *, state):
"""Copy one preview (a pool thread); ``(group, detail)``."""
source, (relative, size, mtime_ns) = pair
destination = Path(source.new_store).joinpath(*relative)
try:
if not os.path.isdir(destination.parent):
base = Path(source.new_store)
if not base.is_dir():
if plan_parent_missing(base):
raise files.FolderGone("the new folder is gone")
created = []
files._make_folders(base.parent, (base.name,), created)
state["created"].extend(created)
created = []
files._make_folders(base, relative[:-1], created)
state["created"].extend(created)
outcome = files.copy_one(
Path(source.old_store).joinpath(*relative),
destination,
size,
mtime_ns,
state["created"],
)
except files.CopyFailed as error:
return "errors", str(error)
if outcome == "moved":
state["written"].append(str(destination))
return "moved", "Copied and checked"
return "skipped", "Already there"
async def _move_cancelled(self, job, plan, state):
if exit_signal.requested():
return # the next start removes the partial copies (journal)
await self._run_in_executor(self._move_rollback, plan, state)
self._end_move(job, "cancelled", "Move cancelled: nothing changed.")
def _move_rollback(self, plan, state):
"""Remove what this job wrote (files, folders it made, a folder and
marker it created) and the journal: nothing changed."""
for path in state["written"]:
files.remove_file(path)
for folder in reversed(state["created"]):
files.remove_empty_folders(folder)
if plan.to_folder and state.get("created_folder"):
folder = plan.target["path"]
try:
names = os.listdir(folder)
except OSError:
names = None
if names is not None and set(names) <= {preview_location.MARKER_NAME}:
_remove(preview_location.marker_path(folder))
files.remove_empty_folders(folder)
_remove(move_plan.journal_path())
def _move_switch(self, plan):
"""Phase 5: one config commit sets the location and records the old
stores under ``cleanup``; the runtime follows under the write lock."""
from ..config import PromptManagerConfig
from ..config_store import get_config_store
entry = {
"from": plan.old.mode,
"folder": plan.old.path if plan.from_folder else None,
"stores": [s.old_store for s in plan.sources if s.old_store],
"drop_keys": self._old_keys_to_drop(plan),
"install": plan.install,
}
# An older entry never removes a store the move just filled (a move
# back while old copies still waited).
filled = {os.path.normcase(s.new_store) for s in plan.sources}
cleanup = []
for old_entry in plan.old.cleanup:
stores = [
store
for store in old_entry.get("stores") or []
if isinstance(store, str) and os.path.normcase(store) not in filled
]
cleanup.append({**old_entry, "stores": stores})
cleanup.append(entry)
if plan.to_folder:
value = {
"mode": preview_store.FOLDER_MODE,
"path": plan.target["path"],
"id": plan.target_id,
"install": plan.install,
"cleanup": cleanup,
}
else:
value = {
"mode": preview_store.SOURCE_MODE,
"install": plan.install,
"cleanup": cleanup,
}
def mutate(config):
section = config.get("thumbnails")
if not isinstance(section, dict):
section = {}
section["location"] = value
config["thumbnails"] = section
return config
def publish(_config):
PromptManagerConfig.THUMBNAIL_LOCATION = value
preview_store.configure(value, check=False)
get_config_store().commit(mutate, expected_revision=None, publish=publish)
preview_store.refresh()
_remove(move_plan.journal_path())
def _old_keys_to_drop(self, plan):
if not plan.from_folder:
return []
keys = [s.old_store and os.path.basename(s.old_store) for s in plan.sources]
keys = [key for key in keys if key]
if plan.to_folder and os.path.normcase(plan.old.path or "") == os.path.normcase(
plan.target["path"] or ""
):
return []
return keys + [key for key, _store, _scan in plan.unlisted]
async def _remove_old_copies(self, job, plan):
"""Phase 6: run the cleanup the switch recorded (and earlier ones)."""
job.update(status_line="Removing old copies", progress=95)
job.publish({"type": "progress", "summary": job.summary()})
left = await self._run_in_executor(self._run_cleanup, job)
if left:
job.scope["cleanup_left"] = True
def _run_cleanup(self, job=None):
"""Remove every recorded old store's previews, temps and empty
folders; keep in ``cleanup`` what could not go. Returns how many
entries are left (blocking)."""
setting = preview_store.setting()
remaining = []
for entry in setting.cleanup:
if not self._cleanup_entry(entry, job):
remaining.append(entry)
if list(setting.cleanup) != remaining:
self._commit_cleanup(remaining)
return len(remaining)
def _cleanup_entry(self, entry, job):
"""One recorded removal; True when it is complete."""
folder = entry.get("folder")
stores = [
store for store in entry.get("stores") or [] if isinstance(store, str)
]
if folder:
if not os.path.isdir(folder):
return False # the old folder is not reachable: kept for later
folder_marker = preview_location.read_marker(folder) or {}
keys = folder_marker.get("keys") or {}
stores += [
os.path.join(folder, key)
for key in entry.get("drop_keys") or []
if preview_location.valid_key(key) and key in keys
]
complete = True
for store in dict.fromkeys(stores):
if not self._remove_store(store, folder, job):
complete = False
if folder and complete:
drop = set(entry.get("drop_keys") or [])
def forget(marker):
for key in drop:
marker["keys"].pop(key, None)
try:
with preview_store.write_lock:
marker = preview_location.write_marker(folder, forget)
except (OSError, ValueError):
return False
if not marker["keys"] and marker.get("created_by_pm"):
names = set(os.listdir(folder))
rest = {n for n in names if not files.is_os_metadata(n)}
if rest <= {preview_location.MARKER_NAME}:
_remove(preview_location.marker_path(folder))
files.remove_empty_folders(folder)
return complete
def _remove_store(self, store, folder, job):
"""Remove the previews and temps of one old store; True when done."""
store = Path(store)
if folder and not preview_location.valid_key(store.name):
return True
if not store.exists():
return True
if preview_location.is_link_like(store):
return True # never through a link; nothing of ours there
complete = True
removed = 0
for current, dirs, names in os.walk(store):
dirs[:] = [
d for d in dirs if not preview_location.is_link_like(Path(current) / d)
]
for name in names:
path = Path(current) / name
if files.is_temp_name(name) or files._preview_parts(name):
if preview_location.is_link_like(path):
continue
if files.remove_file(path):
removed += 1
else:
complete = False
if job is not None:
job.log_line(
"In use by another program; removed at the next "
f"start: {name}"
)
files.remove_empty_folders(store)
if job is not None and removed:
job.log_line(f"Removed {_number(removed)} old files from {store.name}")
return complete
def _commit_cleanup(self, remaining):
from ..config import PromptManagerConfig
from ..config_store import get_config_store
def mutate(config):
section = config.get("thumbnails")
if isinstance(section, dict) and isinstance(section.get("location"), dict):
section["location"]["cleanup"] = remaining
return config
def publish(config):
value = (config.get("thumbnails") or {}).get("location")
PromptManagerConfig.THUMBNAIL_LOCATION = value
preview_store.set_cleanup(remaining)
get_config_store().commit(mutate, expected_revision=None, publish=publish)
def plan_parent_missing(store):
"""A key folder's parent (the thumbnail folder) is gone mid-copy."""
return not Path(store).parent.is_dir()