alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-08-14 09:43:22 -03:00
Code Issues Packages Projects Releases Wiki Activity

fix(delete): track staged batches in-process; reconcile at startup

Browse Source
This commit is contained in:
Will Miao
2026-08-11 18:36:30 +08:00
parent f1d3ac0cdc
commit 5fd4946b1f
3 changed files with 542 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/services/pending_delete_service.py
+199 -38
View File
@@ -97,6 +97,11 @@ class PendingDeleteService:
# ServiceRegistry roots during sweeps so undo/purge work even before
# every scanner is registered.
self._known_roots: List[str] = []
# MODEL batches only; recipe batches live in the fixed settings-dir
# parent. Registry access is short critical sections; the lock-free
# reconciliation scan registers concurrently.
self._known_batch_dirs: Dict[str, str] = {}
self._registry_lock = asyncio.Lock()
# ------------------------------------------------------------------
# Public API
@@ -176,6 +181,7 @@ class PendingDeleteService:
)
self._write_manifest_atomic(batch_dir, manifest)
self._remember_root(root)
await self._remember_batch(batch_id, batch_dir)
# Arm the per-batch purge timer. Safe inside the lock: task
# creation does not await, and purge_batch re-reads the
# manifest's expires_at at fire time, so stale timers no-op.
@@ -296,7 +302,7 @@ class PendingDeleteService:
# Track (entry, original_staged_path, loser_dir) for rollback.
moved: List[Tuple[Dict[str, Any], str, str]] = []
processed_losers: List[str] = []
processed_losers: List[Tuple[str, str]] = [] # (loser_id, loser_dir)
try:
for loser_id in batch_ids[1:]:
@@ -332,7 +338,7 @@ class PendingDeleteService:
entry["staged"] = os.path.abspath(new_staged)
winner_manifest["entries"].append(entry)
moved.append((entry, original_staged, loser_dir))
processed_losers.append(loser_dir)
processed_losers.append((loser_id, loser_dir))
except OSError as exc:
logger.warning(
"Merge of %s failed after moving files: %s; rolling back",
@@ -357,10 +363,15 @@ class PendingDeleteService:
self._rollback_merge_moves(moved)
return None
# All moves committed: remove loser dirs (must be empty by now).
for loser_dir in processed_losers:
# All moves committed: remove loser dirs (must be empty by now)
# and drop them from the registry. Skipped losers (missing /
# corrupted / same-dir) stay registered so the sweep still
# quarantines them, exactly as before the registry existed.
for loser_id, loser_dir in processed_losers:
self._remove_manifest(loser_dir)
self._remove_empty_dir(loser_dir)
await self._forget_batch(loser_id)
await self._remember_batch(winner_id, winner_dir)
# Arm a fresh purge timer for the winner with the re-anchored
# expiry (the winner's original timer fires at the OLD expiry and
@@ -438,33 +449,101 @@ class PendingDeleteService:
# Remove the manifest + batch dir only after all entries restored.
self._remove_manifest(batch_dir)
self._remove_empty_dir(batch_dir)
await self._forget_batch(batch_id)
logger.info("Restored pending-delete batch %s", batch_id)
return self._undo_result(manifest)
async def purge_expired(self) -> int:
"""Purge every expired batch across ALL model roots and the recipe dir.
async def purge_expired(self, scan_roots: bool = False) -> int:
"""Purge every expired batch.
Lock-free by design: enumerates staging parents (all scanner types via
the ServiceRegistry plus the global recipe staging dir) and delegates
each batch to :meth:`purge_batch`, which acquires the ops lock. Never
call this while holding the ops lock.
Default (registry-only): iterates a SNAPSHOT of the in-process MODEL
batch registry plus a shallow check of the fixed recipe staging
parent - cheap, no tree walk per delete. With ``scan_roots=True``
(startup sweep only) a reconciliation pass re-discovers every batch
on disk under the model roots and registers it FIRST, so crash
leftovers and externally created batches are covered too.
Lock-free by design: delegates each batch to :meth:`purge_batch`,
which acquires the ops lock. Never call this while holding the ops
lock.
"""
purged = 0
for parent in await self._get_all_staging_parents():
if not os.path.isdir(parent):
if scan_roots:
await self._reconcile_scan_roots()
# MODEL batches: snapshot so purge_batch can remove entries
# mid-iteration without a dict-changed-size error.
batch_ids: List[str] = [
batch_id for batch_id, _dir in await self._registered_batch_dirs()
]
# RECIPE batches: fixed settings-dir parent, shallow check as before.
recipe_parent = self._recipe_staging_parent()
for name in self._list_dir_names(recipe_parent):
if name.endswith(ORPHANED_SUFFIX):
# Quarantine is terminal - never re-rename or delete.
continue
for name in self._list_dir_names(parent):
if name.endswith(ORPHANED_SUFFIX):
# Quarantine is terminal - never re-rename or delete.
continue
try:
await self.purge_batch(name)
purged += 1
except Exception as exc: # defensive - sweep must not crash
logger.warning("Failed to purge batch %s: %s", name, exc)
if name not in batch_ids:
batch_ids.append(name)
for batch_id in batch_ids:
try:
await self.purge_batch(batch_id)
purged += 1
except Exception as exc: # defensive - sweep must not crash
logger.warning("Failed to purge batch %s: %s", batch_id, exc)
return purged
async def _reconcile_scan_roots(self) -> None:
"""Register every pending-delete batch found under the model roots.
Runs at startup (``purge_expired(scan_roots=True)``) to re-discover
batches left over from a previous process or created externally.
Registers ALL non-orphaned batch dirs regardless of manifest validity:
malformed/manifest-less dirs must reach ``_purge_batch_dir`` so it can
QUARANTINE them (preserving the pre-registry sweep semantics). The
walk only descends into dirs literally named ``.lm-pending-delete``,
so false positives are structurally limited.
"""
from .model_scanner import _is_excluded_dir
for root in await self._get_all_model_roots():
if not os.path.isdir(root):
continue
visited: Set[str] = set()
for dirpath, dirnames, _files in os.walk(
root, followlinks=True, topdown=True
):
real_dir = os.path.realpath(dirpath)
if real_dir in visited:
# Symlink cycle: prune descent and move on.
dirnames[:] = []
continue
visited.add(real_dir)
if os.path.basename(dirpath) == PENDING_DELETE_DIR_NAME:
# The current dir IS a staging parent (reachable only when
# a model root itself is one): register its batches.
await self._register_batch_candidates(dirpath)
dirnames[:] = []
continue
next_dirs: List[str] = []
for name in dirnames:
if name == PENDING_DELETE_DIR_NAME:
await self._register_batch_candidates(
os.path.join(dirpath, name)
)
elif _is_excluded_dir(name):
continue
else:
next_dirs.append(name)
dirnames[:] = next_dirs
async def _register_batch_candidates(self, staging_parent: str) -> None:
"""Register every non-orphaned batch subdir of a staging parent."""
for name in self._list_dir_names(staging_parent):
if name.endswith(ORPHANED_SUFFIX):
# Quarantine is terminal - never re-register.
continue
await self._remember_batch(name, os.path.join(staging_parent, name))
async def purge_batch(self, batch_id: str) -> None:
"""Purge one batch. Silent no-op for missing/undone/not-yet-expired.
@@ -476,7 +555,10 @@ class PendingDeleteService:
batch_dir = await self._find_batch_dir(batch_id)
if not batch_dir:
return
self._purge_batch_dir(batch_dir)
if self._purge_batch_dir(batch_dir):
# Both purge and quarantine remove the batch dir (quarantine
# renames it to *.orphaned), so the registry entry is stale.
await self._forget_batch(batch_id)
# ------------------------------------------------------------------
# Internals
@@ -500,6 +582,29 @@ class PendingDeleteService:
if root and root not in self._known_roots:
self._known_roots.append(root)
async def _remember_batch(self, batch_id: str, batch_dir: str) -> None:
"""Register a MODEL batch in the in-process registry (idempotent).
Short critical section (dict mutation only, no I/O while holding the
lock) so the lock-free reconciliation scan can register concurrently.
"""
async with self._registry_lock:
self._known_batch_dirs[batch_id] = batch_dir
async def _forget_batch(self, batch_id: str) -> None:
"""Remove a MODEL batch from the in-process registry (idempotent)."""
async with self._registry_lock:
self._known_batch_dirs.pop(batch_id, None)
async def _registered_batch_dirs(self) -> List[Tuple[str, str]]:
"""Return a SNAPSHOT of (batch_id, batch_dir) registry pairs.
The snapshot lets purge iterate safely while purge_batch removes
entries mid-loop (no dict-changed-size error).
"""
async with self._registry_lock:
return list(self._known_batch_dirs.items())
def _find_model_root(self, scanner: Any, original_file_path: Optional[str]) -> Optional[str]:
"""Return the configured root containing ``original_file_path``."""
finder = getattr(scanner, "_find_root_for_file", None)
@@ -818,18 +923,6 @@ class PendingDeleteService:
settings_paths.get_settings_dir(create=True), PENDING_DELETE_DIR_NAME
)
async def _get_all_staging_parents(self) -> List[str]:
"""Model staging parents for every scanner type + the recipe parent."""
parents: List[str] = []
for root in await self._get_all_model_roots():
parent = os.path.join(root, PENDING_DELETE_DIR_NAME)
if parent not in parents:
parents.append(parent)
recipe_parent = self._recipe_staging_parent()
if recipe_parent not in parents:
parents.append(recipe_parent)
return parents
async def _get_all_model_roots(self) -> List[str]:
"""Collect every configured model root across all scanner types.
@@ -878,13 +971,81 @@ class PendingDeleteService:
return roots
async def _find_batch_dir(self, batch_id: str) -> Optional[str]:
"""Locate a batch directory across every staging parent."""
"""Locate a batch directory.
Registry lookup first (fast path; stale entries are forgotten when
their dir vanished); then a targeted scan of the model roots for a
batch dir named exactly ``batch_id`` under a ``.lm-pending-delete``
parent (restart / externally created batches; manifest verification
applies so random uuid-named user dirs are never registered); finally
the fixed recipe staging parent. Returns ``None`` (404 semantics)
when not found.
"""
if not batch_id:
return None
for parent in await self._get_all_staging_parents():
candidate = os.path.join(parent, batch_id)
if os.path.isdir(candidate):
# 1) Registry fast path.
async with self._registry_lock:
known = self._known_batch_dirs.get(batch_id)
if known is not None:
if os.path.isdir(known):
return known
await self._forget_batch(batch_id) # stale entry - dir is gone
# 2) Targeted scan fallback across the model roots.
for root in await self._get_all_model_roots():
if not os.path.isdir(root):
continue
candidate = await self._scan_root_for_batch(root, batch_id)
if candidate is not None:
await self._remember_batch(batch_id, candidate)
return candidate
# 3) Recipe batches: fixed settings-dir parent, shallow check.
candidate = os.path.join(self._recipe_staging_parent(), batch_id)
if os.path.isdir(candidate):
return candidate
return None
async def _scan_root_for_batch(self, root: str, batch_id: str) -> Optional[str]:
"""Search one model root for a batch dir named exactly ``batch_id``.
Walks the root (``followlinks=True``) with a realpath cycle guard,
looking for ``.lm-pending-delete`` parents whose subdir matches
``batch_id`` AND has a parseable manifest. The manifest check prevents
random uuid-named user dirs from being treated as batches (a batch
with no parseable manifest cannot be undone anyway).
"""
from .model_scanner import _is_excluded_dir
visited: Set[str] = set()
for dirpath, dirnames, _files in os.walk(root, followlinks=True, topdown=True):
real_dir = os.path.realpath(dirpath)
if real_dir in visited:
# Symlink cycle: prune descent and move on.
dirnames[:] = []
continue
visited.add(real_dir)
if os.path.basename(dirpath) == PENDING_DELETE_DIR_NAME:
candidate = os.path.join(dirpath, batch_id)
if (
os.path.isdir(candidate)
and self._read_manifest(candidate) is not None
):
return candidate
dirnames[:] = []
continue
next_dirs: List[str] = []
for name in dirnames:
if name == PENDING_DELETE_DIR_NAME:
candidate = os.path.join(dirpath, name, batch_id)
if (
os.path.isdir(candidate)
and self._read_manifest(candidate) is not None
):
return candidate
continue
if _is_excluded_dir(name):
continue
next_dirs.append(name)
dirnames[:] = next_dirs
return None
def _list_dir_names(self, parent: str) -> List[str]: