mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-14 09:43:22 -03:00
fix(delete): track staged batches in-process; reconcile at startup
This commit is contained in:
+6
-3
@@ -251,11 +251,14 @@ class LoraManager:
|
|||||||
# Startup sweep: purge pending-delete batches that expired during a
|
# Startup sweep: purge pending-delete batches that expired during a
|
||||||
# previous run. Non-blocking (fire-and-forget); purge_expired only
|
# previous run. Non-blocking (fire-and-forget); purge_expired only
|
||||||
# removes already-expired batches, so a staged undo that survived a
|
# removes already-expired batches, so a staged undo that survived a
|
||||||
# restart stays restorable. Covers both plugin and standalone modes
|
# restart stays restorable. scan_roots=True runs the reconciliation
|
||||||
# (StandaloneLoraManager reuses this classmethod).
|
# pass first so leftover batches (the in-process registry is empty
|
||||||
|
# after a restart) are re-discovered on disk. Covers both plugin
|
||||||
|
# and standalone modes (StandaloneLoraManager reuses this
|
||||||
|
# classmethod).
|
||||||
pending_delete_service = await get_pending_delete_service()
|
pending_delete_service = await get_pending_delete_service()
|
||||||
asyncio.create_task(
|
asyncio.create_task(
|
||||||
pending_delete_service.purge_expired(),
|
pending_delete_service.purge_expired(scan_roots=True),
|
||||||
name="pending_delete_startup_sweep",
|
name="pending_delete_startup_sweep",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -97,6 +97,11 @@ class PendingDeleteService:
|
|||||||
# ServiceRegistry roots during sweeps so undo/purge work even before
|
# ServiceRegistry roots during sweeps so undo/purge work even before
|
||||||
# every scanner is registered.
|
# every scanner is registered.
|
||||||
self._known_roots: List[str] = []
|
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
|
# Public API
|
||||||
@@ -176,6 +181,7 @@ class PendingDeleteService:
|
|||||||
)
|
)
|
||||||
self._write_manifest_atomic(batch_dir, manifest)
|
self._write_manifest_atomic(batch_dir, manifest)
|
||||||
self._remember_root(root)
|
self._remember_root(root)
|
||||||
|
await self._remember_batch(batch_id, batch_dir)
|
||||||
# Arm the per-batch purge timer. Safe inside the lock: task
|
# Arm the per-batch purge timer. Safe inside the lock: task
|
||||||
# creation does not await, and purge_batch re-reads the
|
# creation does not await, and purge_batch re-reads the
|
||||||
# manifest's expires_at at fire time, so stale timers no-op.
|
# 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.
|
# Track (entry, original_staged_path, loser_dir) for rollback.
|
||||||
moved: List[Tuple[Dict[str, Any], str, str]] = []
|
moved: List[Tuple[Dict[str, Any], str, str]] = []
|
||||||
processed_losers: List[str] = []
|
processed_losers: List[Tuple[str, str]] = [] # (loser_id, loser_dir)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
for loser_id in batch_ids[1:]:
|
for loser_id in batch_ids[1:]:
|
||||||
@@ -332,7 +338,7 @@ class PendingDeleteService:
|
|||||||
entry["staged"] = os.path.abspath(new_staged)
|
entry["staged"] = os.path.abspath(new_staged)
|
||||||
winner_manifest["entries"].append(entry)
|
winner_manifest["entries"].append(entry)
|
||||||
moved.append((entry, original_staged, loser_dir))
|
moved.append((entry, original_staged, loser_dir))
|
||||||
processed_losers.append(loser_dir)
|
processed_losers.append((loser_id, loser_dir))
|
||||||
except OSError as exc:
|
except OSError as exc:
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"Merge of %s failed after moving files: %s; rolling back",
|
"Merge of %s failed after moving files: %s; rolling back",
|
||||||
@@ -357,10 +363,15 @@ class PendingDeleteService:
|
|||||||
self._rollback_merge_moves(moved)
|
self._rollback_merge_moves(moved)
|
||||||
return None
|
return None
|
||||||
|
|
||||||
# All moves committed: remove loser dirs (must be empty by now).
|
# All moves committed: remove loser dirs (must be empty by now)
|
||||||
for loser_dir in processed_losers:
|
# 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_manifest(loser_dir)
|
||||||
self._remove_empty_dir(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
|
# Arm a fresh purge timer for the winner with the re-anchored
|
||||||
# expiry (the winner's original timer fires at the OLD expiry and
|
# 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.
|
# Remove the manifest + batch dir only after all entries restored.
|
||||||
self._remove_manifest(batch_dir)
|
self._remove_manifest(batch_dir)
|
||||||
self._remove_empty_dir(batch_dir)
|
self._remove_empty_dir(batch_dir)
|
||||||
|
await self._forget_batch(batch_id)
|
||||||
|
|
||||||
logger.info("Restored pending-delete batch %s", batch_id)
|
logger.info("Restored pending-delete batch %s", batch_id)
|
||||||
return self._undo_result(manifest)
|
return self._undo_result(manifest)
|
||||||
|
|
||||||
async def purge_expired(self) -> int:
|
async def purge_expired(self, scan_roots: bool = False) -> int:
|
||||||
"""Purge every expired batch across ALL model roots and the recipe dir.
|
"""Purge every expired batch.
|
||||||
|
|
||||||
Lock-free by design: enumerates staging parents (all scanner types via
|
Default (registry-only): iterates a SNAPSHOT of the in-process MODEL
|
||||||
the ServiceRegistry plus the global recipe staging dir) and delegates
|
batch registry plus a shallow check of the fixed recipe staging
|
||||||
each batch to :meth:`purge_batch`, which acquires the ops lock. Never
|
parent - cheap, no tree walk per delete. With ``scan_roots=True``
|
||||||
call this while holding the ops lock.
|
(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
|
purged = 0
|
||||||
for parent in await self._get_all_staging_parents():
|
if scan_roots:
|
||||||
if not os.path.isdir(parent):
|
await self._reconcile_scan_roots()
|
||||||
continue
|
# MODEL batches: snapshot so purge_batch can remove entries
|
||||||
for name in self._list_dir_names(parent):
|
# 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):
|
if name.endswith(ORPHANED_SUFFIX):
|
||||||
# Quarantine is terminal - never re-rename or delete.
|
# Quarantine is terminal - never re-rename or delete.
|
||||||
continue
|
continue
|
||||||
|
if name not in batch_ids:
|
||||||
|
batch_ids.append(name)
|
||||||
|
for batch_id in batch_ids:
|
||||||
try:
|
try:
|
||||||
await self.purge_batch(name)
|
await self.purge_batch(batch_id)
|
||||||
purged += 1
|
purged += 1
|
||||||
except Exception as exc: # defensive - sweep must not crash
|
except Exception as exc: # defensive - sweep must not crash
|
||||||
logger.warning("Failed to purge batch %s: %s", name, exc)
|
logger.warning("Failed to purge batch %s: %s", batch_id, exc)
|
||||||
return purged
|
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:
|
async def purge_batch(self, batch_id: str) -> None:
|
||||||
"""Purge one batch. Silent no-op for missing/undone/not-yet-expired.
|
"""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)
|
batch_dir = await self._find_batch_dir(batch_id)
|
||||||
if not batch_dir:
|
if not batch_dir:
|
||||||
return
|
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
|
# Internals
|
||||||
@@ -500,6 +582,29 @@ class PendingDeleteService:
|
|||||||
if root and root not in self._known_roots:
|
if root and root not in self._known_roots:
|
||||||
self._known_roots.append(root)
|
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]:
|
def _find_model_root(self, scanner: Any, original_file_path: Optional[str]) -> Optional[str]:
|
||||||
"""Return the configured root containing ``original_file_path``."""
|
"""Return the configured root containing ``original_file_path``."""
|
||||||
finder = getattr(scanner, "_find_root_for_file", None)
|
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
|
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]:
|
async def _get_all_model_roots(self) -> List[str]:
|
||||||
"""Collect every configured model root across all scanner types.
|
"""Collect every configured model root across all scanner types.
|
||||||
|
|
||||||
@@ -878,15 +971,83 @@ class PendingDeleteService:
|
|||||||
return roots
|
return roots
|
||||||
|
|
||||||
async def _find_batch_dir(self, batch_id: str) -> Optional[str]:
|
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:
|
if not batch_id:
|
||||||
return None
|
return None
|
||||||
for parent in await self._get_all_staging_parents():
|
# 1) Registry fast path.
|
||||||
candidate = os.path.join(parent, batch_id)
|
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):
|
if os.path.isdir(candidate):
|
||||||
return candidate
|
return candidate
|
||||||
return None
|
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]:
|
def _list_dir_names(self, parent: str) -> List[str]:
|
||||||
try:
|
try:
|
||||||
return [
|
return [
|
||||||
|
|||||||
@@ -451,7 +451,9 @@ async def test_g_manifestless_dir_quarantined(tmp_path: Path, monkeypatch) -> No
|
|||||||
await _register_model_root(monkeypatch, lora_roots=[root])
|
await _register_model_root(monkeypatch, lora_roots=[root])
|
||||||
|
|
||||||
service = await PendingDeleteService.get_instance()
|
service = await PendingDeleteService.get_instance()
|
||||||
await service.purge_expired()
|
# The batch is hand-created (unregistered): the reconciliation pass is
|
||||||
|
# required for the default registry-only purge to discover it.
|
||||||
|
await service.purge_expired(scan_roots=True)
|
||||||
|
|
||||||
orphaned = staging / "batch1.orphaned"
|
orphaned = staging / "batch1.orphaned"
|
||||||
assert orphaned.is_dir()
|
assert orphaned.is_dir()
|
||||||
@@ -474,7 +476,8 @@ async def test_h_corrupted_manifest_quarantined(tmp_path: Path, monkeypatch) ->
|
|||||||
await _register_model_root(monkeypatch, lora_roots=[root])
|
await _register_model_root(monkeypatch, lora_roots=[root])
|
||||||
service = await PendingDeleteService.get_instance()
|
service = await PendingDeleteService.get_instance()
|
||||||
|
|
||||||
await service.purge_expired() # must not crash
|
# Hand-created (unregistered) batch: reconciliation discovers it.
|
||||||
|
await service.purge_expired(scan_roots=True) # must not crash
|
||||||
|
|
||||||
orphaned = staging / "batch2.orphaned"
|
orphaned = staging / "batch2.orphaned"
|
||||||
assert orphaned.is_dir()
|
assert orphaned.is_dir()
|
||||||
@@ -1055,7 +1058,8 @@ async def test_r_purge_expired_enumerates_all_scanner_types_and_recipe_dir(
|
|||||||
)
|
)
|
||||||
|
|
||||||
service = await PendingDeleteService.get_instance()
|
service = await PendingDeleteService.get_instance()
|
||||||
purged = await service.purge_expired()
|
# Hand-created (unregistered) batches: reconciliation pass discovers them.
|
||||||
|
purged = await service.purge_expired(scan_roots=True)
|
||||||
|
|
||||||
assert purged >= 4
|
assert purged >= 4
|
||||||
for root in (lora_root, ckpt_root, emb_root):
|
for root in (lora_root, ckpt_root, emb_root):
|
||||||
@@ -1117,12 +1121,13 @@ async def test_t_quarantine_is_terminal(tmp_path: Path, monkeypatch) -> None:
|
|||||||
await _register_model_root(monkeypatch, lora_roots=[root])
|
await _register_model_root(monkeypatch, lora_roots=[root])
|
||||||
service = await PendingDeleteService.get_instance()
|
service = await PendingDeleteService.get_instance()
|
||||||
|
|
||||||
await service.purge_expired()
|
# Hand-created (unregistered) batch: reconciliation discovers it.
|
||||||
|
await service.purge_expired(scan_roots=True)
|
||||||
orphaned = staging / "qbatch.orphaned"
|
orphaned = staging / "qbatch.orphaned"
|
||||||
assert orphaned.is_dir()
|
assert orphaned.is_dir()
|
||||||
|
|
||||||
# Second sweep must NOT re-rename or delete the quarantined dir.
|
# Second sweep must NOT re-rename or delete the quarantined dir.
|
||||||
await service.purge_expired()
|
await service.purge_expired(scan_roots=True)
|
||||||
assert orphaned.is_dir()
|
assert orphaned.is_dir()
|
||||||
assert (orphaned / "model.safetensors").read_bytes() == b"data"
|
assert (orphaned / "model.safetensors").read_bytes() == b"data"
|
||||||
assert not batch_dir.exists()
|
assert not batch_dir.exists()
|
||||||
@@ -1169,7 +1174,9 @@ async def test_u_lock_no_deadlock_with_concurrent_purge(tmp_path: Path, monkeypa
|
|||||||
cached_entry=None,
|
cached_entry=None,
|
||||||
)
|
)
|
||||||
|
|
||||||
purge_task = asyncio.create_task(service.purge_expired())
|
# Hand-created (unregistered) "expired" batch: the purge task must run the
|
||||||
|
# reconciliation pass to discover it alongside the staged "new" batch.
|
||||||
|
purge_task = asyncio.create_task(service.purge_expired(scan_roots=True))
|
||||||
stage_task = asyncio.create_task(do_stage())
|
stage_task = asyncio.create_task(do_stage())
|
||||||
results = await asyncio.gather(purge_task, stage_task, return_exceptions=True)
|
results = await asyncio.gather(purge_task, stage_task, return_exceptions=True)
|
||||||
|
|
||||||
@@ -1533,3 +1540,327 @@ async def test_snap2_merge_keeps_both_snapshots(
|
|||||||
snap_entries = [e for e in manifest["entries"] if e.get("snapshot")]
|
snap_entries = [e for e in manifest["entries"] if e.get("snapshot")]
|
||||||
assert len(snap_entries) == 2
|
assert len(snap_entries) == 2
|
||||||
assert {e["snapshot"]["file_path"] for e in snap_entries} == {str(a1), str(b1)}
|
assert {e["snapshot"]["file_path"] for e in snap_entries} == {str(a1), str(b1)}
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# Batch-registry lifecycle (todo 1: in-process _known_batch_dirs)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
# (a) stage_model_delete registers in _known_batch_dirs
|
||||||
|
async def test_reg_a_stage_model_registers_batch(tmp_path: Path) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
batch_id = await _stage_simple(service, root, "model")
|
||||||
|
|
||||||
|
assert service._known_batch_dirs.get(batch_id) == str(
|
||||||
|
root / PENDING_DELETE_DIR_NAME / batch_id
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# (b) undo success removes the entry
|
||||||
|
async def test_reg_b_undo_success_removes_registry_entry(tmp_path: Path) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
batch_id = await _stage_simple(service, root, "model")
|
||||||
|
assert batch_id in service._known_batch_dirs
|
||||||
|
|
||||||
|
await service.undo(batch_id)
|
||||||
|
|
||||||
|
assert batch_id not in service._known_batch_dirs
|
||||||
|
|
||||||
|
|
||||||
|
# (c) purge_batch removes the entry after a real purge
|
||||||
|
async def test_reg_c_purge_batch_removes_registry_entry(tmp_path: Path) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
batch_id = await _stage_simple(service, root, "model")
|
||||||
|
batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id
|
||||||
|
manifest_path = batch_dir / "manifest.json"
|
||||||
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||||||
|
manifest["expires_at"] = int(time.time()) - 10
|
||||||
|
manifest_path.write_text(json.dumps(manifest))
|
||||||
|
|
||||||
|
await service.purge_batch(batch_id)
|
||||||
|
|
||||||
|
assert batch_id not in service._known_batch_dirs
|
||||||
|
assert not batch_dir.exists()
|
||||||
|
|
||||||
|
|
||||||
|
# (d) quarantine (corrupted manifest) removes the entry
|
||||||
|
async def test_reg_d_quarantine_removes_registry_entry(tmp_path: Path) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
batch_id = await _stage_simple(service, root, "model")
|
||||||
|
batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id
|
||||||
|
assert batch_id in service._known_batch_dirs
|
||||||
|
|
||||||
|
# Corrupt the manifest: purge_batch quarantines the dir (returns True).
|
||||||
|
(batch_dir / "manifest.json").write_text("{ not valid json !!!")
|
||||||
|
|
||||||
|
await service.purge_batch(batch_id)
|
||||||
|
|
||||||
|
assert batch_id not in service._known_batch_dirs
|
||||||
|
assert not batch_dir.exists()
|
||||||
|
assert (batch_dir.with_name(f"{batch_id}.orphaned")).is_dir()
|
||||||
|
|
||||||
|
|
||||||
|
# (e) merge success: winner present + losers removed; EXDEV-abort: unchanged
|
||||||
|
async def test_reg_e_merge_registry_lifecycle(
|
||||||
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||||
|
) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
_spy_purge_timers(monkeypatch)
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
bid_a = await _stage_simple(service, root, "alpha")
|
||||||
|
bid_b = await _stage_simple(service, root, "beta")
|
||||||
|
bid_c = await _stage_simple(service, root, "gamma")
|
||||||
|
assert set(service._known_batch_dirs) == {bid_a, bid_b, bid_c}
|
||||||
|
|
||||||
|
# Merge success: winner stays, processed loser forgotten, untouched batch stays.
|
||||||
|
assert await service.merge_batches([bid_a, bid_b]) == bid_a
|
||||||
|
assert bid_a in service._known_batch_dirs
|
||||||
|
assert bid_b not in service._known_batch_dirs
|
||||||
|
assert bid_c in service._known_batch_dirs
|
||||||
|
|
||||||
|
# EXDEV-abort: registry untouched.
|
||||||
|
def exdev_rename(src: str, dst: str) -> None:
|
||||||
|
raise OSError(errno.EXDEV, "Invalid cross-device link", src, dst)
|
||||||
|
|
||||||
|
monkeypatch.setattr("py.services.pending_delete_service.os.rename", exdev_rename)
|
||||||
|
before = dict(service._known_batch_dirs)
|
||||||
|
assert await service.merge_batches([bid_a, bid_c]) is None
|
||||||
|
assert dict(service._known_batch_dirs) == before
|
||||||
|
|
||||||
|
|
||||||
|
# (f) _reset_pending_delete_service clears the registry
|
||||||
|
async def test_reg_f_reset_clears_registry(tmp_path: Path) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
batch_id = await _stage_simple(service, root, "model")
|
||||||
|
assert service._known_batch_dirs
|
||||||
|
|
||||||
|
_reset_pending_delete_service()
|
||||||
|
|
||||||
|
fresh = await PendingDeleteService.get_instance()
|
||||||
|
assert fresh is not service
|
||||||
|
assert fresh._known_batch_dirs == {}
|
||||||
|
|
||||||
|
|
||||||
|
# (g) scan_roots=True reconciles externally created batches (expired purged,
|
||||||
|
# non-expired registered); the registry-only default does NOT find them
|
||||||
|
async def test_reg_g_reconciliation_finds_external_batches(
|
||||||
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||||
|
) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
await _register_model_root(monkeypatch, lora_roots=[root])
|
||||||
|
|
||||||
|
expired_dir = root / PENDING_DELETE_DIR_NAME / "ext-expired"
|
||||||
|
expired_dir.mkdir(parents=True)
|
||||||
|
(expired_dir / "old.safetensors").write_bytes(b"old")
|
||||||
|
_write_batch_manifest(
|
||||||
|
expired_dir,
|
||||||
|
batch_id="ext-expired",
|
||||||
|
kind="model",
|
||||||
|
model_type="loras",
|
||||||
|
expires_at=int(time.time()) - 10,
|
||||||
|
entries=[
|
||||||
|
{
|
||||||
|
"staged": str(expired_dir / "old.safetensors"),
|
||||||
|
"original": str(root / "old.safetensors"),
|
||||||
|
"restored": False,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
fresh_dir = root / PENDING_DELETE_DIR_NAME / "ext-fresh"
|
||||||
|
fresh_dir.mkdir(parents=True)
|
||||||
|
(fresh_dir / "new.safetensors").write_bytes(b"new")
|
||||||
|
_write_batch_manifest(
|
||||||
|
fresh_dir,
|
||||||
|
batch_id="ext-fresh",
|
||||||
|
kind="model",
|
||||||
|
model_type="loras",
|
||||||
|
expires_at=int(time.time()) + 100,
|
||||||
|
entries=[
|
||||||
|
{
|
||||||
|
"staged": str(fresh_dir / "new.safetensors"),
|
||||||
|
"original": str(root / "new.safetensors"),
|
||||||
|
"restored": False,
|
||||||
|
}
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
|
||||||
|
# Registry-only default: the externally created batches are invisible.
|
||||||
|
await service.purge_expired()
|
||||||
|
assert expired_dir.is_dir()
|
||||||
|
assert fresh_dir.is_dir()
|
||||||
|
assert "ext-expired" not in service._known_batch_dirs
|
||||||
|
assert "ext-fresh" not in service._known_batch_dirs
|
||||||
|
|
||||||
|
# Reconciliation pass: expired one purged, non-expired one registered.
|
||||||
|
await service.purge_expired(scan_roots=True)
|
||||||
|
|
||||||
|
assert not expired_dir.exists()
|
||||||
|
assert not (root / "old.safetensors").exists()
|
||||||
|
assert fresh_dir.is_dir()
|
||||||
|
assert (fresh_dir / "new.safetensors").exists()
|
||||||
|
assert "ext-expired" not in service._known_batch_dirs
|
||||||
|
assert service._known_batch_dirs.get("ext-fresh") == str(fresh_dir)
|
||||||
|
|
||||||
|
|
||||||
|
# (h) _find_batch_dir with cleared registry locates + registers (restart sim)
|
||||||
|
async def test_reg_h_find_batch_dir_restart_simulation(
|
||||||
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||||
|
) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
await _register_model_root(monkeypatch, lora_roots=[root])
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
batch_id = await _stage_simple(service, root, "model")
|
||||||
|
assert batch_id in service._known_batch_dirs
|
||||||
|
|
||||||
|
# Simulate a restart: the in-process registry is empty but the batch dir
|
||||||
|
# is still on disk.
|
||||||
|
service._known_batch_dirs.clear()
|
||||||
|
|
||||||
|
found = await service._find_batch_dir(batch_id)
|
||||||
|
|
||||||
|
assert found == str(root / PENDING_DELETE_DIR_NAME / batch_id)
|
||||||
|
assert service._known_batch_dirs.get(batch_id) == found
|
||||||
|
|
||||||
|
# Undo works after the restart simulation.
|
||||||
|
await service.undo(batch_id)
|
||||||
|
assert (root / "model.safetensors").read_bytes() == b"model-data"
|
||||||
|
assert batch_id not in service._known_batch_dirs
|
||||||
|
|
||||||
|
|
||||||
|
# (i) purge iteration uses a snapshot: no dict-changed-size when entries are
|
||||||
|
# removed mid-iteration
|
||||||
|
async def test_reg_i_purge_iteration_uses_snapshot(
|
||||||
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||||
|
) -> None:
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
await _register_model_root(monkeypatch, lora_roots=[root])
|
||||||
|
service = await PendingDeleteService.get_instance()
|
||||||
|
ids = [await _stage_simple(service, root, f"m{i}") for i in range(5)]
|
||||||
|
|
||||||
|
for batch_id in ids:
|
||||||
|
manifest_path = root / PENDING_DELETE_DIR_NAME / batch_id / "manifest.json"
|
||||||
|
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||||||
|
manifest["expires_at"] = int(time.time()) - 10
|
||||||
|
manifest_path.write_text(json.dumps(manifest))
|
||||||
|
|
||||||
|
# Every purge removes its registry entry mid-loop; the snapshot makes this
|
||||||
|
# safe (iterating the dict directly would raise RuntimeError).
|
||||||
|
await service.purge_expired()
|
||||||
|
|
||||||
|
assert service._known_batch_dirs == {}
|
||||||
|
for batch_id in ids:
|
||||||
|
assert not (root / PENDING_DELETE_DIR_NAME / batch_id).exists()
|
||||||
|
|
||||||
|
|
||||||
|
# (j) STARTUP SWEEP PIN: the startup sweep task passes scan_roots=True
|
||||||
|
async def test_reg_j_startup_sweep_passes_scan_roots_true(
|
||||||
|
monkeypatch: pytest.MonkeyPatch,
|
||||||
|
) -> None:
|
||||||
|
from py import lora_manager
|
||||||
|
|
||||||
|
sweep_calls: List[Dict[str, Any]] = []
|
||||||
|
|
||||||
|
class _SpySweepService:
|
||||||
|
async def purge_expired(self, scan_roots: bool = False) -> int:
|
||||||
|
sweep_calls.append({"scan_roots": scan_roots})
|
||||||
|
return 0
|
||||||
|
|
||||||
|
async def _fake_get_service() -> _SpySweepService:
|
||||||
|
return _SpySweepService()
|
||||||
|
|
||||||
|
monkeypatch.setattr(lora_manager, "get_pending_delete_service", _fake_get_service)
|
||||||
|
|
||||||
|
async def _stub(*args: Any, **_kwargs: Any) -> Any:
|
||||||
|
return args[0] if args else None
|
||||||
|
|
||||||
|
class _DummyScanner:
|
||||||
|
async def initialize_in_background(self) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
dummy = _DummyScanner()
|
||||||
|
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_civitai_client", lambda: _stub())
|
||||||
|
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_download_manager", lambda: _stub())
|
||||||
|
monkeypatch.setattr(
|
||||||
|
lora_manager.ServiceRegistry, "get_download_queue_service", lambda: _stub()
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_backup_service", lambda: _stub())
|
||||||
|
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_websocket_manager", lambda: _stub())
|
||||||
|
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_lora_scanner", lambda: _stub(dummy))
|
||||||
|
monkeypatch.setattr(
|
||||||
|
lora_manager.ServiceRegistry, "get_checkpoint_scanner", lambda: _stub(dummy)
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
lora_manager.ServiceRegistry, "get_embedding_scanner", lambda: _stub(dummy)
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(lora_manager.ServiceRegistry, "get_recipe_scanner", lambda: _stub(dummy))
|
||||||
|
|
||||||
|
from py.services import metadata_service as metadata_service_module
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
metadata_service_module,
|
||||||
|
"initialize_metadata_providers",
|
||||||
|
_stub,
|
||||||
|
)
|
||||||
|
|
||||||
|
from py.services.llm_service import LLMService
|
||||||
|
|
||||||
|
monkeypatch.setattr(LLMService, "get_instance", _stub)
|
||||||
|
|
||||||
|
async def _fake_migration() -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
lora_manager.ExampleImagesMigration,
|
||||||
|
"check_and_run_migrations",
|
||||||
|
staticmethod(_fake_migration),
|
||||||
|
)
|
||||||
|
|
||||||
|
captured: List[Any] = []
|
||||||
|
|
||||||
|
class _DummyTask:
|
||||||
|
def add_done_callback(self, _cb: Any) -> None: # pragma: no cover - stub
|
||||||
|
pass
|
||||||
|
|
||||||
|
def done(self) -> bool: # pragma: no cover - stub
|
||||||
|
return False
|
||||||
|
|
||||||
|
def _capture_task(coro: Any, *args: Any, **kwargs: Any) -> _DummyTask:
|
||||||
|
captured.append(coro)
|
||||||
|
return _DummyTask()
|
||||||
|
|
||||||
|
monkeypatch.setattr(asyncio, "create_task", _capture_task)
|
||||||
|
|
||||||
|
try:
|
||||||
|
await lora_manager.LoraManager._initialize_services()
|
||||||
|
finally:
|
||||||
|
sweep_coro: Any = None
|
||||||
|
for coro in captured:
|
||||||
|
qualname = getattr(coro.cr_code, "co_qualname", "")
|
||||||
|
if "_SpySweepService.purge_expired" in qualname:
|
||||||
|
sweep_coro = coro
|
||||||
|
else:
|
||||||
|
coro.close()
|
||||||
|
if sweep_coro is not None:
|
||||||
|
# The sweep task body only runs when awaited; execute just the
|
||||||
|
# spy's purge_expired so it records its invocation arguments.
|
||||||
|
await sweep_coro
|
||||||
|
|
||||||
|
# The startup sweep must invoke purge_expired with scan_roots=True (the
|
||||||
|
# reconciliation flag) - forgetting it would break restart cleanup.
|
||||||
|
assert sweep_calls == [{"scan_roots": True}]
|
||||||
|
|||||||
Reference in New Issue
Block a user