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

feat(delete): stage model and recipe deletes for 30s undo

Browse Source
This commit is contained in:
Will Miao
2026-08-11 14:08:15 +08:00
parent 6a259a14fa
commit 2d6cf545b9
12 changed files with 3444 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/services/checkpoint_scanner.py
+3 -2
View File
@@ -13,7 +13,7 @@ from ..utils.models import CheckpointMetadata
from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3
from ..utils.metadata_manager import MetadataManager
from ..config import config
from .model_scanner import ModelScanner
from .model_scanner import ModelScanner, _is_excluded_dir
from .model_hash_index import ModelHashIndex
logger = logging.getLogger(__name__)
@@ -328,7 +328,8 @@ class CheckpointScanner(ModelScanner):
if not os.path.exists(root_path):
continue
for dirpath, _dirnames, filenames in os.walk(root_path):
for dirpath, dirnames, filenames in os.walk(root_path):
dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)]
for filename in filenames:
if not filename.endswith(".metadata.json"):
continue
py/services/model_lifecycle_service.py
+23 -3
View File
@@ -7,6 +7,7 @@ import os
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, cast
from ..services.service_registry import ServiceRegistry
from ..services.pending_delete_service import get_pending_delete_service
from ..utils.constants import PREVIEW_EXTENSIONS
from ..utils.metadata_manager import MetadataManager
@@ -129,9 +130,24 @@ class ModelLifecycleService:
target_dir = os.path.dirname(file_path)
base_name = os.path.basename(file_path)
file_name, main_extension = os.path.splitext(base_name)
deleted_files = await delete_model_artifacts(
target_dir, file_name, main_extension=main_extension
# Stage the delete into the pending-delete service when undo is
# enabled; a successful stage renames the artifacts away, otherwise
# fall back to the direct hard delete.
pending_delete_service = await get_pending_delete_service()
batch_id = await pending_delete_service.stage_model_delete(
scanner=self._scanner,
target_dir=target_dir,
file_name=file_name,
main_extension=main_extension,
original_file_path=file_path,
cached_entry=cached_entry,
)
deleted_files: List[str] = []
if batch_id is None:
deleted_files = await delete_model_artifacts(
target_dir, file_name, main_extension=main_extension
)
if cache:
cache.raw_data = [
@@ -151,7 +167,11 @@ class ModelLifecycleService:
if callable(persist_current_cache):
await cast(Awaitable[Any], persist_current_cache())
return {"success": True, "deleted_files": deleted_files}
return {
"success": True,
"deleted_files": deleted_files,
"batch_id": batch_id,
}
@staticmethod
def _extract_model_id_from_payload(payload: Any) -> Optional[int]:
py/services/model_scanner.py
+73 -5
View File
@@ -19,12 +19,28 @@ from .service_registry import ServiceRegistry
from .websocket_manager import ws_manager
from .persistent_model_cache import get_persistent_cache
from .settings_manager import get_settings_manager
from .pending_delete_service import PENDING_DELETE_DIR_NAME, get_pending_delete_service
from .cache_entry_validator import CacheEntryValidator
from .cache_health_monitor import CacheHealthMonitor, CacheHealthStatus
logger = logging.getLogger(__name__)
def _is_excluded_dir(name: str) -> bool:
"""Return True when a directory entry must be skipped during model walks.
The pending-delete staging directory is excluded so staged files never
appear in the library as ghost model entries.
"""
return name == PENDING_DELETE_DIR_NAME
def _is_pending_delete_path(path: str) -> bool:
"""Return True when any path component is the pending-delete staging dir."""
normalized = str(path).replace(os.sep, "/")
return any(part == PENDING_DELETE_DIR_NAME for part in normalized.split("/"))
@dataclass
class CacheBuildResult:
"""Represents the outcome of scanning model files for cache building."""
@@ -711,6 +727,8 @@ class ModelScanner:
if ext in self.file_extensions:
total_files += 1
elif entry.is_dir(follow_symlinks=True):
if _is_excluded_dir(entry.name):
continue
count_recursive(entry.path)
except Exception as e:
logger.error(f"Error counting files in entry {entry.path}: {e}")
@@ -864,7 +882,8 @@ class ModelScanner:
continue
# Recursively scan directory
for root, _, files in os.walk(root_path, followlinks=True):
for root, dirnames, files in os.walk(root_path, followlinks=True):
dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)]
real_root = os.path.realpath(root)
if real_root in visited_real_paths:
continue
@@ -1137,6 +1156,11 @@ class ModelScanner:
hash_index = hash_index or self._hash_index
excluded_models = excluded_models if excluded_models is not None else self._excluded_models
# Belt-and-braces: staged files must never become library entries even
# if a caller invokes this method directly with a staging path.
if _is_pending_delete_path(file_path):
return None
metadata, should_skip = await MetadataManager.load_metadata(file_path, self.model_class)
if should_skip:
@@ -1456,6 +1480,8 @@ class ModelScanner:
if self.is_cancelled():
return
elif entry.is_dir(follow_symlinks=True):
if _is_excluded_dir(entry.name):
continue
await scan_recursive(entry.path, root_path, visited_paths)
except Exception as entry_error:
logger.error(f"Error processing entry {entry.path}: {entry_error}")
@@ -2206,6 +2232,11 @@ class ModelScanner:
# Track deleted models to update cache once
deleted_models = []
# Stage each file into the pending-delete staging area and merge
# all per-file batches into ONE batch for the whole bulk action.
pending_delete_service = await get_pending_delete_service()
batch_ids: List[str] = []
for file_path in file_paths:
if self.is_cancelled():
logger.info(f"{self.model_type.capitalize()} Scanner: Bulk delete cancelled by user")
@@ -2218,11 +2249,35 @@ class ModelScanner:
base_name = os.path.basename(file_path)
file_name, main_extension = os.path.splitext(base_name)
deleted_files = await delete_model_artifacts(
target_dir,
file_name,
# Snapshot the cache entry BEFORE the cache mutation that
# runs after the loop - the manifest needs it for undo.
cached_entry = None
if cache is not None:
cached_entry = next(
(item for item in cache.raw_data if item.get('file_path') == file_path),
None,
)
batch_id = await pending_delete_service.stage_model_delete(
scanner=self,
target_dir=target_dir,
file_name=file_name,
main_extension=main_extension,
original_file_path=file_path,
cached_entry=cached_entry,
)
if batch_id is not None:
# Artifacts were renamed into staging: the main file is
# gone from its original location.
batch_ids.append(batch_id)
deleted_files = [file_path]
else:
deleted_files = await delete_model_artifacts(
target_dir,
file_name,
main_extension=main_extension,
)
if deleted_files:
deleted_models.append(file_path)
@@ -2246,6 +2301,18 @@ class ModelScanner:
'error': str(e)
})
# Merge every staged per-file batch into ONE undoable batch. On a
# merge failure (cross-volume EXDEV etc.) the response falls back
# to the constituent batch_ids array so the frontend can undo them
# sequentially.
batch_field: Dict[str, Any] = {}
if batch_ids:
merged_id = await pending_delete_service.merge_batches(batch_ids)
if merged_id is not None:
batch_field['batch_id'] = merged_id
else:
batch_field['batch_ids'] = list(batch_ids)
# Batch update cache if any models were deleted
if deleted_models:
# Update the cache in a batch operation
@@ -2257,7 +2324,8 @@ class ModelScanner:
'total_deleted': total_deleted,
'total_attempted': len(file_paths),
'cache_updated': cache_updated,
'results': results
'results': results,
**batch_field
}
except Exception as e:
py/services/pending_delete_service.py
+974
View File
@@ -0,0 +1,974 @@
"""Pending-delete staging service.
Stages model/recipe deletes into hidden per-root staging directories so a
30-second undo window can restore them before the physical purge runs. The
service is the foundation for the delete-undo feature: every staged batch is
described by a ``manifest.json`` which is the ONLY source of truth.
LOCK HIERARCHY (critical - asyncio.Lock is NOT re-entrant):
``_ops_lock`` is acquired ONLY by stage_model_delete, stage_recipe_delete,
merge_batches, undo and purge_batch. ``purge_expired()`` NEVER acquires it -
it enumerates staging dirs and delegates each batch to ``purge_batch`` (which
locks). The opportunistic ``await self.purge_expired()`` at the start of
stage_*/undo MUST therefore run BEFORE those methods acquire the lock.
"""
from __future__ import annotations
import asyncio
import errno
import json
import logging
import os
import shutil
import tempfile
import time
import uuid
from typing import (
Any,
Awaitable,
Callable,
Dict,
List,
Optional,
Sequence,
Set,
Tuple,
cast,
)
from ..utils.constants import PREVIEW_EXTENSIONS
from ..utils import settings_paths
from .settings_manager import get_settings_manager
logger = logging.getLogger(__name__)
# Undo window in seconds before a staged batch becomes purge-eligible.
PENDING_DELETE_TTL_SECONDS = 30
# Hidden staging directory name placed under each model root (and the settings
# dir for recipes).
PENDING_DELETE_DIR_NAME = ".lm-pending-delete"
# Manifest file name inside every batch directory.
MANIFEST_FILE_NAME = "manifest.json"
# Suffix appended when quarantining malformed/manifest-less batch dirs. The
# quarantine is terminal: never re-renamed, never re-quarantined, never
# deleted by the sweep.
ORPHANED_SUFFIX = ".orphaned"
# Map scanner.model_type (singular) to the manifest page type values.
_MODEL_TYPE_PAGE_MAP = {
"lora": "loras",
"checkpoint": "checkpoints",
"embedding": "embeddings",
}
# Module-level alias so tests can spy on timer task creation without patching
# the global asyncio module.
_create_task = asyncio.create_task
class PendingDeleteService:
"""Stage, undo and purge pending model/recipe deletions.
Singleton + asyncio.Lock pattern (mirrors py/services/model_scanner.py).
"""
_instance: Optional["PendingDeleteService"] = None
_lock: asyncio.Lock = asyncio.Lock()
@classmethod
async def get_instance(cls) -> "PendingDeleteService":
"""Return the lazily initialised singleton instance."""
async with cls._lock:
if cls._instance is None:
cls._instance = cls()
return cls._instance
def __init__(self) -> None:
if hasattr(self, "_initialized"):
return
self._initialized = True
# Serialises stage/merge/undo/purge_batch. purge_expired never locks.
self._ops_lock = asyncio.Lock()
# Track fire-and-forget purge timer tasks to keep them alive and to
# cancel them on shutdown / singleton reset.
self._purge_tasks: Set[Any] = set()
# Roots the service has staged into (in-process). Combined with the
# ServiceRegistry roots during sweeps so undo/purge work even before
# every scanner is registered.
self._known_roots: List[str] = []
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
async def stage_model_delete(
self,
*,
scanner: Any,
target_dir: str,
file_name: str,
main_extension: Optional[str],
original_file_path: str,
cached_entry: Optional[Dict[str, Any]],
) -> Optional[str]:
"""Rename a model's artifacts into a per-root staging batch.
Returns the batch id, or ``None`` when undo is disabled, the staging
root cannot be resolved, or staging failed (caller falls back to a
hard delete).
"""
# LOCK-FREE section: opportunistic purge must never run while holding
# the ops lock (the lock is not re-entrant).
await self._opportunistic_purge()
if not self._undo_enabled():
return None
async with self._ops_lock:
batch_dir: Optional[str] = None
staged_pairs: List[Dict[str, Any]] = []
try:
root = self._find_model_root(scanner, original_file_path)
if not root:
logger.warning(
"No model root contains %s; skipping staging",
original_file_path,
)
return None
artifacts = self._enumerate_model_artifacts(
target_dir, file_name, main_extension
)
if not artifacts:
logger.warning(
"No existing artifacts for %s; skipping staging",
original_file_path,
)
return None
batch_id = self._new_batch_id()
batch_dir = os.path.join(
os.path.join(root, PENDING_DELETE_DIR_NAME), batch_id
)
os.makedirs(batch_dir, exist_ok=True)
staged_pairs = self._rename_artifacts_into_batch(
batch_dir, artifacts, staged_pairs
)
# Attach the model snapshot to the MAIN-file entry (the one
# whose original path is the model file itself, not the
# metadata/preview sidecars). Merged bulk manifests therefore
# carry EVERY deleted model's snapshot on its entry; the
# top-level model_snapshot is kept for backward compatibility
# and the single-delete path.
main_abs = os.path.abspath(original_file_path)
for entry in staged_pairs:
if entry.get("original") == main_abs:
entry["snapshot"] = cached_entry
break
manifest = self._build_manifest(
batch_id=batch_id,
kind="model",
model_type=self._resolve_model_type(scanner),
expires_at=int(time.time()) + PENDING_DELETE_TTL_SECONDS,
entries=staged_pairs,
model_snapshot=cached_entry,
)
self._write_manifest_atomic(batch_dir, manifest)
self._remember_root(root)
# 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.
self._arm_purge_timer(batch_id)
logger.info(
"Staged model delete batch %s with %d file(s)",
batch_id,
len(staged_pairs),
)
return batch_id
except OSError as exc:
logger.warning(
"Staging model %s failed: %s; rolling back", original_file_path, exc
)
if batch_dir:
self._rollback_model_staging(batch_dir, staged_pairs)
self._remove_empty_dir(batch_dir)
return None
except Exception as exc: # defensive - never block the delete flow
logger.warning(
"Unexpected error staging model %s: %s", original_file_path, exc
)
return None
async def stage_recipe_delete(
self,
*,
recipe_json_path: str,
image_path: Optional[str],
recipe_data: Optional[Dict[str, Any]],
) -> Optional[str]:
"""Copy a recipe JSON (and, when it exists, its image) into staging.
Returns the batch id, or ``None`` when undo is disabled / staging
failed. Missing or shared preview images are skipped.
"""
await self._opportunistic_purge()
if not self._undo_enabled():
return None
async with self._ops_lock:
batch_dir: Optional[str] = None
staged_pairs: List[Dict[str, Any]] = []
try:
json_path = os.path.abspath(os.path.normpath(recipe_json_path))
if not os.path.exists(json_path):
logger.warning(
"Recipe JSON %s does not exist; skipping staging", json_path
)
return None
batch_id = self._new_batch_id()
batch_dir = os.path.join(self._recipe_staging_parent(), batch_id)
os.makedirs(batch_dir, exist_ok=True)
staged_pairs = self._copy_recipe_artifacts(
batch_dir, json_path, image_path, staged_pairs
)
manifest = self._build_manifest(
batch_id=batch_id,
kind="recipe",
model_type=None,
expires_at=int(time.time()) + PENDING_DELETE_TTL_SECONDS,
entries=staged_pairs,
recipe_snapshot=recipe_data,
)
self._write_manifest_atomic(batch_dir, manifest)
self._arm_purge_timer(batch_id)
logger.info(
"Staged recipe delete batch %s with %d file(s)",
batch_id,
len(staged_pairs),
)
return batch_id
except OSError as exc:
logger.warning(
"Staging recipe %s failed: %s; rolling back",
recipe_json_path,
exc,
)
if batch_dir:
self._rollback_recipe_staging(batch_dir, staged_pairs)
self._remove_empty_dir(batch_dir)
return None
except Exception as exc: # defensive - never block the delete flow
logger.warning(
"Unexpected error staging recipe %s: %s", recipe_json_path, exc
)
return None
async def merge_batches(self, batch_ids: Sequence[str]) -> Optional[str]:
"""Merge several batches into the first batch's manifest.
Winner is ``batch_ids[0]``. The staged files of losing batches are
MOVED (os.rename) into the winner's batch dir and their ``staged``
paths rewritten in the merged manifest BEFORE any loser dir is
removed. ``expires_at`` is re-anchored to ``now + TTL`` at merge time
and a FRESH purge timer is armed for the winner.
On any move failure every already-moved file is moved BACK and the
original batch dirs/manifests are left intact; ``None`` is returned so
callers fall back to the ``batch_ids`` array contract. Cross-volume
merges hit EXDEV here - expected and fine (the fallback is the normal
path for those bulks).
"""
if not batch_ids:
return None
async with self._ops_lock:
winner_id = batch_ids[0]
winner_dir = await self._find_batch_dir(winner_id)
if not winner_dir:
return None
winner_manifest = self._read_manifest(winner_dir)
if winner_manifest is None:
return None
# Track (entry, original_staged_path, loser_dir) for rollback.
moved: List[Tuple[Dict[str, Any], str, str]] = []
processed_losers: List[str] = []
try:
for loser_id in batch_ids[1:]:
loser_dir = await self._find_batch_dir(loser_id)
if not loser_dir or os.path.normpath(loser_dir) == os.path.normpath(
winner_dir
):
continue
loser_manifest = self._read_manifest(loser_dir)
if loser_manifest is None:
# Corrupted loser: leave it for the sweep to quarantine.
continue
for entry in loser_manifest.get("entries") or []:
if entry.get("restored"):
continue
staged_path = entry.get("staged")
if not staged_path or not os.path.exists(staged_path):
continue
new_staged = os.path.join(
winner_dir, os.path.basename(staged_path)
)
if os.path.exists(new_staged):
# os.rename would silently overwrite the existing
# staged file on POSIX - never drop a staged file.
# Abort the merge so callers fall back to the
# batch_ids array contract.
raise OSError(
f"Merge collision: {os.path.basename(staged_path)} "
"already staged in winner batch"
)
os.rename(staged_path, new_staged)
original_staged = entry["staged"]
entry["staged"] = os.path.abspath(new_staged)
winner_manifest["entries"].append(entry)
moved.append((entry, original_staged, loser_dir))
processed_losers.append(loser_dir)
except OSError as exc:
logger.warning(
"Merge of %s failed after moving files: %s; rolling back",
list(batch_ids),
exc,
)
self._rollback_merge_moves(moved)
return None
# Re-anchor expiry and persist the merged manifest atomically.
winner_manifest["expires_at"] = (
int(time.time()) + PENDING_DELETE_TTL_SECONDS
)
try:
self._write_manifest_atomic(winner_dir, winner_manifest)
except OSError as exc:
logger.warning(
"Failed to write merged manifest for %s: %s; rolling back",
winner_id,
exc,
)
self._rollback_merge_moves(moved)
return None
# All moves committed: remove loser dirs (must be empty by now).
for loser_dir in processed_losers:
self._remove_manifest(loser_dir)
self._remove_empty_dir(loser_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
# no-ops after re-reading the manifest - without this fresh timer
# an idle server would never purge the merged batch).
self._arm_purge_timer(winner_id)
logger.info("Merged batches %s into %s", list(batch_ids), winner_id)
return winner_id
async def undo(self, batch_id: str) -> Dict[str, Any]:
"""Restore every staged file of a batch to its original path.
Raises ``ValueError`` for unknown batches, expired batches ("Undo
window expired") and occupied target paths ("Target path occupied").
Restores entries one at a time, persisting the manifest after each, so
a mid-undo failure leaves a retry-able state.
"""
await self._opportunistic_purge()
async with self._ops_lock:
batch_dir = await self._find_batch_dir(batch_id)
if not batch_dir:
raise ValueError(f"Unknown batch id: {batch_id}")
manifest = self._read_manifest(batch_dir)
if manifest is None:
raise ValueError(f"Manifest missing for batch {batch_id}")
if manifest.get("state") == "restored":
return self._undo_result(manifest)
now = time.time()
expires_at = manifest.get("expires_at")
if isinstance(expires_at, (int, float)) and expires_at < now:
raise ValueError("Undo window expired")
entries = manifest.get("entries") or []
# Pre-check ALL target paths (except already-restored entries) so
# an occupied original path protects the new file and leaves the
# whole batch intact.
for entry in entries:
if entry.get("restored"):
continue
original_path = entry.get("original")
if original_path and os.path.exists(original_path):
raise ValueError("Target path occupied")
for entry in entries:
if entry.get("restored"):
continue
staged_path = entry.get("staged")
original_path = entry.get("original")
if not staged_path or not original_path:
entry["restored"] = True
continue
if not os.path.exists(staged_path):
# Staged file already gone (purged or manually removed):
# treat as restored and finish the rest of the batch.
entry["restored"] = True
self._write_manifest_atomic(batch_dir, manifest)
continue
self._restore_file(staged_path, original_path)
entry["restored"] = True
# Persist after each entry so a mid-undo failure is retry-able.
self._write_manifest_atomic(batch_dir, manifest)
manifest["state"] = "restored"
try:
self._write_manifest_atomic(batch_dir, manifest)
except OSError as exc:
logger.warning(
"Failed to mark manifest restored for %s: %s", batch_id, exc
)
# Remove the manifest + batch dir only after all entries restored.
self._remove_manifest(batch_dir)
self._remove_empty_dir(batch_dir)
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.
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.
"""
purged = 0
for parent in await self._get_all_staging_parents():
if not os.path.isdir(parent):
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)
return purged
async def purge_batch(self, batch_id: str) -> None:
"""Purge one batch. Silent no-op for missing/undone/not-yet-expired.
Missing staged files (already-restored / partially-restored batches)
are treated as already-purged. A per-file purge failure (locked file)
skips only that file and keeps the batch dir for the next round.
"""
async with self._ops_lock:
batch_dir = await self._find_batch_dir(batch_id)
if not batch_dir:
return
self._purge_batch_dir(batch_dir)
# ------------------------------------------------------------------
# Internals
# ------------------------------------------------------------------
async def _opportunistic_purge(self) -> None:
"""Fire the opportunistic sweep. Cheap when empty; never locked."""
try:
await self.purge_expired()
except Exception as exc: # defensive - staging/undo must still proceed
logger.warning("Opportunistic pending-delete purge failed: %s", exc)
def _undo_enabled(self) -> bool:
try:
return bool(get_settings_manager().get("delete_undo_enabled", True))
except Exception as exc: # defensive - default to enabled
logger.warning("Failed to read delete_undo_enabled setting: %s", exc)
return True
def _remember_root(self, root: str) -> None:
"""Record a root the service has staged into (in-process registry)."""
if root and root not in self._known_roots:
self._known_roots.append(root)
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)
if callable(finder):
try:
root = cast(Optional[str], finder(original_file_path))
if root:
return os.path.abspath(root)
except Exception as exc: # defensive - fall back to roots scan
logger.debug("_find_root_for_file failed: %s", exc)
if not original_file_path:
return None
roots_getter = getattr(scanner, "get_model_roots", None)
if not callable(roots_getter):
return None
try:
normalized = os.path.abspath(os.path.normpath(original_file_path))
for root in cast(Sequence[str], roots_getter()) or []:
root_abs = os.path.abspath(os.path.normpath(root))
if normalized == root_abs or normalized.startswith(root_abs + os.sep):
return root_abs
except Exception as exc: # defensive - never block the delete flow
logger.debug("get_model_roots fallback failed: %s", exc)
return None
def _resolve_model_type(self, scanner: Any) -> Optional[str]:
raw = getattr(scanner, "model_type", None)
if not raw:
return None
return _MODEL_TYPE_PAGE_MAP.get(raw, raw)
def _enumerate_model_artifacts(
self, target_dir: str, file_name: str, main_extension: Optional[str]
) -> List[str]:
"""Enumerate existing artifacts exactly like delete_model_artifacts."""
main_extension = ".safetensors" if main_extension is None else main_extension
main_file = f"{file_name}{main_extension}" if main_extension else file_name
patterns = [main_file, f"{file_name}.metadata.json"]
for ext in PREVIEW_EXTENSIONS:
patterns.append(f"{file_name}{ext}")
artifacts: List[str] = []
for pattern in patterns:
path = os.path.abspath(os.path.join(target_dir, pattern))
if os.path.exists(path):
artifacts.append(path)
return artifacts
def _rename_artifacts_into_batch(
self,
batch_dir: str,
artifacts: Sequence[str],
staged_pairs: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""Rename artifacts into the batch dir, recording progress per file.
Progress is appended to ``staged_pairs`` before the next move so a
mid-way OSError leaves the caller with the already-moved files for
rollback.
"""
for original_path in artifacts:
staged_path = os.path.join(batch_dir, os.path.basename(original_path))
os.rename(original_path, staged_path)
staged_pairs.append(
{
"staged": os.path.abspath(staged_path),
"original": os.path.abspath(original_path),
"restored": False,
}
)
return staged_pairs
def _copy_recipe_artifacts(
self,
batch_dir: str,
json_path: str,
image_path: Optional[str],
staged_pairs: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""Copy the recipe JSON and, when it exists, the image into the batch."""
staged_json = os.path.join(batch_dir, os.path.basename(json_path))
shutil.copy2(json_path, staged_json)
staged_pairs.append(
{
"staged": os.path.abspath(staged_json),
"original": json_path,
"restored": False,
}
)
if image_path:
image_abs = os.path.abspath(os.path.normpath(image_path))
if os.path.exists(image_abs):
staged_image = os.path.join(batch_dir, os.path.basename(image_abs))
shutil.copy2(image_abs, staged_image)
staged_pairs.append(
{
"staged": os.path.abspath(staged_image),
"original": image_abs,
"restored": False,
}
)
return staged_pairs
def _restore_file(self, staged_path: str, original_path: str) -> None:
"""Restore a staged file to its original path, tolerating EXDEV.
``os.rename`` is atomic and preferred (model staging and most recipe
restores are same-volume). Recipe staging copies into the settings-dir
staging parent, which may live on a DIFFERENT filesystem than the
recipes dir; rename then raises EXDEV. Fall back to ``shutil.copy2`` +
``os.remove`` so the bytes are restored and the staged copy removed.
"""
try:
os.rename(staged_path, original_path)
except OSError as exc:
if exc.errno != errno.EXDEV:
raise
shutil.copy2(staged_path, original_path)
os.remove(staged_path)
def _rollback_model_staging(
self, batch_dir: str, staged_pairs: Sequence[Dict[str, Any]]
) -> None:
"""Rename already-staged files back to their originals."""
for pair in reversed(list(staged_pairs)):
staged_path = pair.get("staged")
original_path = pair.get("original")
if not staged_path or not original_path:
continue
if not os.path.exists(staged_path):
continue
try:
os.rename(staged_path, original_path)
except OSError as exc: # pragma: no cover - best-effort rollback
logger.warning(
"Failed to roll back staged file %s -> %s: %s",
staged_path,
original_path,
exc,
)
def _rollback_recipe_staging(
self, batch_dir: str, staged_pairs: Sequence[Dict[str, Any]]
) -> None:
"""Remove staged copies (recipe originals were never moved)."""
for pair in staged_pairs:
staged_path = pair.get("staged")
if not staged_path:
continue
try:
if os.path.exists(staged_path):
os.remove(staged_path)
except OSError as exc: # pragma: no cover - best-effort rollback
logger.warning(
"Failed to remove staged copy %s: %s", staged_path, exc
)
def _rollback_merge_moves(
self, moved: Sequence[Tuple[Dict[str, Any], str, str]]
) -> None:
"""Move already-merged files back to their original loser batch dirs."""
for _entry, original_staged, _loser_dir in reversed(list(moved)):
current = _entry.get("staged")
if not current or not original_staged:
continue
if not os.path.exists(current):
continue
try:
os.rename(current, original_staged)
except OSError as exc: # pragma: no cover - best-effort rollback
logger.warning(
"Failed to roll back merge move %s -> %s: %s",
current,
original_staged,
exc,
)
def _purge_batch_dir(self, batch_dir: str) -> bool:
"""Purge one batch dir. Returns True when the batch was purged/removed."""
if not os.path.isdir(batch_dir):
return False
manifest = self._read_manifest(batch_dir)
if manifest is None:
# Corrupted or manifest-less batch: quarantine, NEVER delete the
# staged files (they may be the only copy of the user's data).
self._quarantine_batch_dir(batch_dir)
return True
if manifest.get("state") == "restored":
return False
expires_at = manifest.get("expires_at")
if not isinstance(expires_at, (int, float)) or expires_at >= time.time():
# Not yet expired - stale timers from merged-away/undone batches
# are harmless.
return False
entries = manifest.get("entries") or []
remaining: List[Dict[str, Any]] = []
for entry in entries:
staged_path = entry.get("staged")
if not staged_path or not os.path.exists(staged_path):
# Missing staged file = already restored / already purged.
continue
try:
os.remove(staged_path)
except OSError as exc:
logger.warning(
"Skipping locked staged file %s: %s", staged_path, exc
)
remaining.append(entry)
if remaining:
# Never remove the batch dir past per-file errors; the batch is
# retried by the next opportunistic purge.
return False
self._remove_manifest(batch_dir)
self._remove_empty_dir(batch_dir)
return True
def _quarantine_batch_dir(self, batch_dir: str) -> str:
"""Rename a malformed batch dir to ``<batch_id>.orphaned`` (terminal)."""
orphaned_dir = f"{batch_dir}{ORPHANED_SUFFIX}"
if os.path.exists(orphaned_dir):
orphaned_dir = f"{batch_dir}-{int(time.time())}{ORPHANED_SUFFIX}"
try:
os.rename(batch_dir, orphaned_dir)
except OSError as exc: # pragma: no cover - defensive
logger.warning("Failed to quarantine %s: %s", batch_dir, exc)
return batch_dir
logger.warning(
"Quarantined malformed/manifest-less pending-delete batch %s",
os.path.basename(batch_dir),
)
return orphaned_dir
def _build_manifest(
self,
*,
batch_id: str,
kind: str,
model_type: Optional[str],
expires_at: int,
entries: Sequence[Dict[str, Any]],
model_snapshot: Any = None,
recipe_snapshot: Any = None,
) -> Dict[str, Any]:
is_model = kind == "model"
return {
"batch_id": batch_id,
"kind": kind,
"model_type": model_type if is_model else None,
"state": "staged",
"expires_at": int(expires_at),
"entries": list(entries),
"model_snapshot": model_snapshot if is_model else None,
"recipe_snapshot": recipe_snapshot if not is_model else None,
}
def _undo_result(self, manifest: Dict[str, Any]) -> Dict[str, Any]:
restored_paths = [
entry["original"]
for entry in manifest.get("entries") or []
if entry.get("restored") and entry.get("original")
]
return {
"batch_id": manifest.get("batch_id"),
"kind": manifest.get("kind"),
"model_type": manifest.get("model_type"),
"restored": restored_paths,
}
def _write_manifest_atomic(
self, batch_dir: str, manifest: Dict[str, Any]
) -> None:
"""Write manifest.json atomically (temp file + os.replace)."""
manifest_path = os.path.join(batch_dir, MANIFEST_FILE_NAME)
fd, temp_path = tempfile.mkstemp(
dir=batch_dir, prefix=".manifest-", suffix=".tmp"
)
try:
with os.fdopen(fd, "w", encoding="utf-8") as handle:
json.dump(manifest, handle, indent=2, ensure_ascii=False)
os.replace(temp_path, manifest_path)
except BaseException:
try:
os.remove(temp_path)
except OSError:
pass
raise
def _read_manifest(self, batch_dir: str) -> Optional[Dict[str, Any]]:
manifest_path = os.path.join(batch_dir, MANIFEST_FILE_NAME)
try:
with open(manifest_path, "r", encoding="utf-8") as handle:
payload = json.load(handle)
except FileNotFoundError:
return None
except (json.JSONDecodeError, OSError) as exc:
logger.warning(
"Corrupted pending-delete manifest at %s: %s", manifest_path, exc
)
return None
if not isinstance(payload, dict):
logger.warning("Invalid pending-delete manifest at %s", manifest_path)
return None
return payload
def _recipe_staging_parent(self) -> str:
# Resolve through the module namespace so the conftest settings-dir
# isolation patch takes effect at call time.
return os.path.join(
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.
Combines the in-process staging roots with the ServiceRegistry's
per-type scanners so sweeps cover every scanner type while undo/purge
still resolve batches staged before the registry was populated.
"""
from .service_registry import ServiceRegistry
roots: List[str] = []
for root in self._known_roots:
if root and root not in roots:
roots.append(root)
for getter_name in (
"get_lora_scanner",
"get_checkpoint_scanner",
"get_embedding_scanner",
):
getter = getattr(ServiceRegistry, getter_name, None)
if not callable(getter):
continue
try:
scanner = await cast(Callable[[], Awaitable[Any]], getter)()
except Exception as exc: # defensive - keep sweeping other types
logger.debug(
"Failed to resolve %s for purge enumeration: %s",
getter_name,
exc,
)
continue
if scanner is None:
continue
get_roots = getattr(scanner, "get_model_roots", None)
if not callable(get_roots):
continue
try:
scanner_roots = cast(Sequence[Any], get_roots())
except Exception as exc: # defensive
logger.debug(
"get_model_roots failed for %s: %s", getter_name, exc
)
continue
for root in scanner_roots or []:
if root and root not in roots:
roots.append(root)
return roots
async def _find_batch_dir(self, batch_id: str) -> Optional[str]:
"""Locate a batch directory across every staging parent."""
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):
return candidate
return None
def _list_dir_names(self, parent: str) -> List[str]:
try:
return [
name
for name in os.listdir(parent)
if os.path.isdir(os.path.join(parent, name))
]
except OSError as exc: # pragma: no cover - defensive
logger.debug("Failed to list staging parent %s: %s", parent, exc)
return []
def _remove_manifest(self, batch_dir: str) -> None:
try:
os.remove(os.path.join(batch_dir, MANIFEST_FILE_NAME))
except OSError as exc: # pragma: no cover - best-effort
logger.debug("Failed to remove manifest in %s: %s", batch_dir, exc)
def _remove_empty_dir(self, directory: str) -> None:
try:
os.rmdir(directory)
except OSError as exc:
logger.debug("Directory %s not empty or missing: %s", directory, exc)
def _new_batch_id(self) -> str:
return uuid.uuid4().hex
def _arm_purge_timer(self, batch_id: str) -> None:
"""Spawn a fire-and-forget purge timer for a batch.
The timer sleeps until the batch's current expiry and then calls
purge_batch, which re-reads the manifest's ``expires_at`` at fire time
so merged-away/undone/not-yet-expired batches are silent no-ops.
"""
try:
asyncio.get_running_loop()
except RuntimeError:
return
task = _create_task(
self._purge_batch_after_ttl(batch_id),
name=f"pending_delete_purge_{batch_id}",
)
self._purge_tasks.add(task)
task.add_done_callback(self._purge_tasks.discard)
async def _purge_batch_after_ttl(self, batch_id: str) -> None:
try:
delay = await self._seconds_until_expiry(batch_id)
if delay is None:
return
await asyncio.sleep(max(0.0, delay))
await self.purge_batch(batch_id)
except asyncio.CancelledError:
raise
except Exception as exc: # defensive - a timer must never crash the loop
logger.warning("Pending-delete purge timer for %s failed: %s", batch_id, exc)
async def _seconds_until_expiry(self, batch_id: str) -> Optional[float]:
batch_dir = await self._find_batch_dir(batch_id)
if not batch_dir:
return None
manifest = self._read_manifest(batch_dir)
if manifest is None:
return None
expires_at = manifest.get("expires_at")
if not isinstance(expires_at, (int, float)):
return None
return float(expires_at) - time.time()
def _cancel_purge_tasks(self) -> None:
for task in list(self._purge_tasks):
task.cancel()
self._purge_tasks.clear()
def _reset_pending_delete_service() -> None:
"""Reset the singleton and cancel in-flight purge timers (tests/shutdown)."""
instance = PendingDeleteService._instance
if instance is not None:
instance._cancel_purge_tasks()
PendingDeleteService._instance = None
async def get_pending_delete_service() -> PendingDeleteService:
"""Return the lazily initialised global :class:`PendingDeleteService`."""
return await PendingDeleteService.get_instance()
py/services/recipes/persistence_service.py
+56 -10
View File
@@ -14,6 +14,7 @@ from typing import Any, Awaitable, Dict, Iterable, Optional, cast
from ...config import config
from ...recipes.constants import GEN_PARAM_KEYS
from ...utils.utils import calculate_recipe_fingerprint
from ..pending_delete_service import get_pending_delete_service
from .errors import RecipeNotFoundError, RecipeValidationError
@@ -201,12 +202,31 @@ class RecipePersistenceService:
recipe_data = json.load(file_obj)
image_path = recipe_data.get("file_path")
# Stage the delete so the recipe can be undone within the undo window.
# The staging service COPIES the JSON (and existing image) into the
# global staging dir and stores recipe_data as the manifest snapshot;
# the originals are removed below as before. When staging is skipped
# (undo disabled / staging failure) the existing hard delete runs.
pending_delete_service = await get_pending_delete_service()
batch_id = await pending_delete_service.stage_recipe_delete(
recipe_json_path=recipe_json_path,
image_path=image_path,
recipe_data=recipe_data,
)
os.remove(recipe_json_path)
if image_path and os.path.exists(image_path):
os.remove(image_path)
await recipe_scanner.remove_recipe(recipe_id)
return PersistenceResult({"success": True, "message": "Recipe deleted successfully"})
return PersistenceResult(
{
"success": True,
"message": "Recipe deleted successfully",
"batch_id": batch_id,
}
)
async def update_recipe(self, *, recipe_scanner, recipe_id: str, updates: dict[str, Any]) -> PersistenceResult:
"""Update persisted metadata for a recipe."""
@@ -450,6 +470,9 @@ class RecipePersistenceService:
deleted_recipes: list[str] = []
failed_recipes: list[dict[str, Any]] = []
batch_ids: list[str] = []
pending_delete_service = await get_pending_delete_service()
for recipe_id in recipe_ids:
recipe_json_path = await recipe_scanner.get_recipe_json_path(recipe_id)
@@ -461,6 +484,17 @@ class RecipePersistenceService:
with open(recipe_json_path, "r", encoding="utf-8") as file_obj:
recipe_data = json.load(file_obj)
image_path = recipe_data.get("file_path")
# Stage each recipe into its own batch; collect the ids so the
# whole bulk action can be merged into ONE undoable batch.
batch_id = await pending_delete_service.stage_recipe_delete(
recipe_json_path=recipe_json_path,
image_path=image_path,
recipe_data=recipe_data,
)
if batch_id:
batch_ids.append(batch_id)
os.remove(recipe_json_path)
if image_path and os.path.exists(image_path):
os.remove(image_path)
@@ -471,15 +505,27 @@ class RecipePersistenceService:
if deleted_recipes:
await recipe_scanner.bulk_remove(deleted_recipes)
return PersistenceResult(
{
"success": True,
"deleted": deleted_recipes,
"failed": failed_recipes,
"total_deleted": len(deleted_recipes),
"total_failed": len(failed_recipes),
}
)
payload: dict[str, Any] = {
"success": True,
"deleted": deleted_recipes,
"failed": failed_recipes,
"total_deleted": len(deleted_recipes),
"total_failed": len(failed_recipes),
}
if batch_ids:
merged_batch_id = await pending_delete_service.merge_batches(batch_ids)
if merged_batch_id:
# Merge succeeded: one undo action covers the whole bulk.
payload["batch_id"] = merged_batch_id
else:
# Merge failure (e.g. cross-volume move): expose the constituent
# batches so the caller can undo them one at a time.
payload["batch_ids"] = batch_ids
else:
payload["batch_id"] = None
return PersistenceResult(payload)
async def save_recipe_from_widget(
self,
py/services/settings_manager.py
+1
View File
@@ -111,6 +111,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
"backup_retention_count": 5,
"use_new_license_icons": True,
"group_by_model": False,
"delete_undo_enabled": True,
# AI / LLM provider configuration (BYOK)
"llm_provider": "openai", # "openai" | "ollama" | "custom"
"llm_api_key": "",
py/utils/usage_stats.py
+3 -1
View File
@@ -10,6 +10,7 @@ from typing import Any, Awaitable, Dict, Set, cast
from ..config import config
from ..services.service_registry import ServiceRegistry
from ..services.model_scanner import _is_excluded_dir
from ..utils.settings_paths import get_settings_dir
# Check if running in standalone mode
@@ -421,7 +422,8 @@ class UsageStats:
if not os.path.exists(root_path):
continue
for dirpath, _dirnames, filenames in os.walk(root_path):
for dirpath, dirnames, filenames in os.walk(root_path):
dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)]
for filename in filenames:
extension = os.path.splitext(filename)[1].lower()
if extension not in supported_extensions: