From ef7a8915e4d57034de3940312514bc7ca1efbeaa Mon Sep 17 00:00:00 2001 From: pixelpaws Date: Sun, 25 Oct 2026 20:16:38 +0800 Subject: [PATCH] fix(scanner): re-freeze GC after incremental cache writes (#1163) Downloads, imports, reconciles, metadata syncs and bulk updates all add freshly allocated objects to caches that were frozen at startup, so ComfyUI's post-prompt gc.collect() grows again within a running session (issue #1161, follow-up to #1147). - add_model_to_cache now stores the slim cache representation: license flags are derived from the full source payload first, then the civitai payload is slimmed to the scan-path shape. Fresh downloads previously also ended up with default license flags. - New schedule_refreeze(): debounced re-freeze with an idle guard that reschedules while the ComfyUI prompt queue is busy, so the freeze never lands mid-prompt. Wired into downloads/imports, reconcile, single-model updates, deferred bulk updates and metadata sync. - freeze_long_lived_objects() logs collect duration and the tracked-object count at DEBUG level to tell 'more objects' apart from 'slower objects'. --- py/services/model_scanner.py | 33 ++++++++- py/utils/gc_utils.py | 105 +++++++++++++++++++++++---- tests/services/test_model_scanner.py | 90 +++++++++++++++++++++++ tests/utils/test_gc_utils.py | 45 +++++++++++- 4 files changed, 255 insertions(+), 18 deletions(-) diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 0f3c4dc6..d4a09394 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -13,7 +13,7 @@ from ..utils.models import BaseModelMetadata, autov3_from_civitai_files, quantiz from ..config import config from ..utils.constants import is_empty_placeholder_hash from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256, calculate_autov3 -from ..utils.gc_utils import schedule_refreeze_after_rebuild +from ..utils.gc_utils import schedule_refreeze, schedule_refreeze_after_rebuild from ..utils.gguf_utils import detect_quantization from ..utils.metadata_manager import MetadataManager from ..utils.sidecar_paths import ( @@ -2588,6 +2588,10 @@ class ModelScanner: elapsed_seconds=time.time() - start_time, **summary, ) + if total_added > 0 or total_removed > 0: + # Reconcile-added entries are GC-tracked again; re-freeze + # once the follow-up work settles (issue #1161). + schedule_refreeze(f"{self.model_type} cache reconcile") return summary except Exception as e: logger.error(f"{self.model_type.capitalize()} Scanner: Error reconciling cache: {e}", exc_info=True) @@ -3889,7 +3893,20 @@ class ModelScanner: # Update folder in metadata metadata_dict['folder'] = folder - + + # Insert the slim cache representation, not the full payload: a + # fresh download/import carries the entire source payload, while + # scan-path entries keep only a handful of civitai keys. Keeping + # the fat payload per download re-tracks thousands of objects for + # the GC until the next restart (issue #1161). License flags must + # be derived before slimming — they read model-level fields the + # slim payload drops. + cache_entry = dict(metadata_dict) + if cache_entry.get('civitai'): + self._ensure_license_flags(cache_entry) + cache_entry['civitai'] = self._slim_civitai_payload(cache_entry['civitai']) + metadata_dict = cache_entry + file_path = metadata_dict.get('file_path', '') if file_path: old_entries = [item for item in self._cache.raw_data if item.get('file_path') == file_path] @@ -3917,6 +3934,9 @@ class ModelScanner: ) await self._persist_current_cache() self.bump_cache_version() + # The new entry is GC-tracked again; re-freeze once the write + # burst settles (issue #1161). + schedule_refreeze(f"{self.model_type} model added to cache") return True except Exception as e: logger.error(f"Error adding model to cache: {e}") @@ -4213,6 +4233,9 @@ class ModelScanner: await self._cache.resort() await self._persist_current_cache(force=True) self.bump_cache_version() + # Bulk updates replace entries en masse; re-freeze once settled + # (issue #1161). + schedule_refreeze(f"{self.model_type} bulk cache update") except Exception: logger.error( "%s Scanner: failed to finalize deferred cache persist", @@ -4334,6 +4357,9 @@ class ModelScanner: if cache_modified: await self._persist_current_cache() self.bump_cache_version() + # The replacement entry is GC-tracked again; re-freeze once + # the write burst settles (issue #1161). + schedule_refreeze(f"{self.model_type} cache update") if metadata and cache_entry is not None: return cache_entry @@ -4539,6 +4565,9 @@ class ModelScanner: old_item_for_sql, ) + # In-place updates still allocate fresh values; re-freeze once the + # sync burst settles (issue #1161). + schedule_refreeze(f"{self.model_type} metadata sync") return True async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool: diff --git a/py/utils/gc_utils.py b/py/utils/gc_utils.py index 3a0eb031..11be8e23 100644 --- a/py/utils/gc_utils.py +++ b/py/utils/gc_utils.py @@ -8,7 +8,10 @@ collection walks hundreds of thousands of objects while holding the GIL, freezing the whole server for seconds (issue #1147). Freezing those objects once they are hydrated moves them into the permanent -generation, which the collector no longer scans. +generation, which the collector no longer scans. Incremental cache writes +(downloads, reconciles, metadata syncs) keep adding tracked objects after the +startup freeze, so they schedule a debounced re-freeze once the write burst +settles (issue #1161). """ from __future__ import annotations @@ -16,6 +19,8 @@ from __future__ import annotations import asyncio import gc import logging +import os +import time from typing import Optional logger = logging.getLogger(__name__) @@ -27,6 +32,11 @@ logger = logging.getLogger(__name__) # can never be collected. _REFREEZE_DELAY_SECONDS = 10.0 +# Upper bound on idle-guard reschedules. A server running prompts back-to-back +# for longer than this many settle delays simply skips the freeze; the next +# cache mutation schedules a fresh attempt. +_REFREEZE_MAX_RETRIES = 6 + _pending_refreeze: Optional[asyncio.TimerHandle] = None @@ -37,13 +47,21 @@ def freeze_long_lived_objects(reason: str) -> None: Best-effort: failures are logged, never raised. """ try: + debug_logging = logger.isEnabledFor(logging.DEBUG) + started = time.perf_counter() + tracked_before = len(gc.get_objects()) if debug_logging else 0 collected = gc.collect() gc.freeze() - logger.debug( - "LoRA Manager: froze GC-tracked objects after %s (%d unreachable collected)", - reason, - collected, - ) + elapsed = time.perf_counter() - started + if debug_logging: + logger.debug( + "LoRA Manager: froze GC-tracked objects after %s " + "(%d unreachable collected in %.2fs; %d objects tracked before collect)", + reason, + collected, + elapsed, + tracked_before, + ) except Exception: # pragma: no cover - defensive guard logger.warning( "LoRA Manager: failed to freeze GC-tracked objects after %s", @@ -52,17 +70,49 @@ def freeze_long_lived_objects(reason: str) -> None: ) -def schedule_refreeze_after_rebuild( - reason: str, delay: float = _REFREEZE_DELAY_SECONDS -) -> None: - """Schedule a debounced re-freeze for after a cache rebuild settles. +def _comfyui_queue_is_idle() -> bool: + """Return True when no prompt is queued or running (or cannot be checked). - A full rebuild replaces the frozen cache with newly allocated objects, + Re-freezing while a prompt executes would move the prompt's transient + object graph (model tensors included) into the permanent generation, + where cyclic garbage can never be reclaimed — so scheduled re-freezes + wait for an idle queue. Standalone mode has no prompt queue and is always + idle. When the queue state cannot be determined (imports unavailable, or + the queue does not report an integer count), assume idle: the startup + freeze already runs under the same unguarded conditions. + """ + if os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1": + return True + try: + from server import PromptServer # ComfyUI host module + + instance = getattr(PromptServer, "instance", None) + queue = getattr(instance, "prompt_queue", None) + if queue is None: + return True + remaining = queue.get_tasks_remaining() + if not isinstance(remaining, int): + return True + return remaining == 0 + except Exception: + return True + + +def schedule_refreeze( + reason: str, + delay: float = _REFREEZE_DELAY_SECONDS, + _retries_left: int = _REFREEZE_MAX_RETRIES, +) -> None: + """Schedule a debounced, idle-guarded re-freeze for after a cache change. + + Cache rebuilds and incremental writes (downloads, reconciles, metadata + syncs) replace or extend the frozen cache with freshly allocated objects, which are GC-tracked again, so ComfyUI's post-prompt gc.collect() gets expensive once more. Re-freezing after a short settle delay moves the new - cache into the permanent generation. Debounced: a library switch re-init - fires this once per scanner, and only the last schedule wins. No-op when - no event loop is running. + objects into the permanent generation. Debounced: only the last schedule + wins. When the timer fires while ComfyUI still has prompts queued or + running, it reschedules instead of freezing mid-prompt (bounded by + ``_REFREEZE_MAX_RETRIES``). No-op when no event loop is running. """ global _pending_refreeze try: @@ -71,4 +121,29 @@ def schedule_refreeze_after_rebuild( return if _pending_refreeze is not None: _pending_refreeze.cancel() - _pending_refreeze = loop.call_later(delay, freeze_long_lived_objects, reason) + + def _fire() -> None: + global _pending_refreeze + _pending_refreeze = None + if _comfyui_queue_is_idle(): + freeze_long_lived_objects(reason) + elif _retries_left > 0: + schedule_refreeze(reason, delay, _retries_left - 1) + else: + logger.debug( + "LoRA Manager: skipped GC re-freeze after %s; prompt queue stayed busy", + reason, + ) + + _pending_refreeze = loop.call_later(delay, _fire) + + +def schedule_refreeze_after_rebuild( + reason: str, delay: float = _REFREEZE_DELAY_SECONDS +) -> None: + """Schedule a debounced re-freeze for after a cache rebuild settles. + + Thin wrapper kept for the original rebuild call sites; equivalent to + :func:`schedule_refreeze`. + """ + schedule_refreeze(reason, delay) diff --git a/tests/services/test_model_scanner.py b/tests/services/test_model_scanner.py index a7443dfd..2048b230 100644 --- a/tests/services/test_model_scanner.py +++ b/tests/services/test_model_scanner.py @@ -3721,3 +3721,93 @@ async def test_process_model_file_keeps_a_zero_byte_placeholder_on_empty_files( assert await scanner._process_model_file(str(model_file), str(tmp_path)) is None assert json.loads(sidecar.read_text(encoding="utf-8"))["sha256"] == EMPTY_HASH_SHA256 + + +async def test_add_model_to_cache_slims_civitai_payload(tmp_path: Path, monkeypatch): + """Downloaded/imported entries must match the slim scan-path shape. + + The full source payload (images, files, description, ...) stays with the + caller; the cache keeps only the frequently used keys, and license flags + are derived before slimming drops the model-level license fields + (issue #1161). + """ + scanner = DummyScanner(tmp_path) + scanner._cache = ModelCache(raw_data=[], folders=[]) + scheduled: List[str] = [] + monkeypatch.setattr( + model_scanner, "schedule_refreeze", lambda reason: scheduled.append(reason) + ) + + full_payload = { + "id": 123, + "modelId": 456, + "name": "v1", + "creator": {"username": "alice", "image": "avatar.png"}, + "trainedWords": ["trigger"], + "model": { + "type": "LORA", + "allowNoCredit": True, + "allowCommercialUse": ["Image"], + "description": "long text", + }, + "images": [{"url": "u", "meta": {"prompt": "p"}}], + "files": [{"name": "f"}], + "stats": {"downloadCount": 1}, + } + metadata = { + "file_path": "x.txt", + "folder": "", + "sha256": "abc", + "tags": [], + "civitai": full_payload, + } + + ok = await scanner.add_model_to_cache(metadata, "") + + assert ok is True + entry = scanner._cache.raw_data[0] + assert entry["civitai"] == { + "id": 123, + "modelId": 456, + "name": "v1", + "creator": {"username": "alice"}, + "trainedWords": ["trigger"], + "model": {"type": "LORA"}, + } + assert entry["license_flags"] == build_license_flags( + {"allowNoCredit": True, "allowCommercialUse": ["Image"]} + ) + # The caller's dict keeps the full payload for downstream use + # (e.g. the filename-template flow). + assert "images" in metadata["civitai"] + assert scheduled == ["dummy model added to cache"] + + +async def test_update_single_model_cache_schedules_refreeze(tmp_path: Path, monkeypatch): + scanner = DummyScanner(tmp_path) + scanner._cache = ModelCache( + raw_data=[ + { + "file_path": "old.txt", + "folder": "", + "sha256": "abc", + "tags": [], + "model_name": "m", + "file_name": "old", + } + ], + folders=[], + ) + scheduled: List[str] = [] + monkeypatch.setattr( + model_scanner, "schedule_refreeze", lambda reason: scheduled.append(reason) + ) + + result = await scanner.update_single_model_cache( + "old.txt", + "old.txt", + {"file_path": "old.txt", "sha256": "abc", "tags": [], "model_name": "m"}, + ) + + assert result + assert "dummy cache update" in scheduled diff --git a/tests/utils/test_gc_utils.py b/tests/utils/test_gc_utils.py index 7e216f9b..3b6345f9 100644 --- a/tests/utils/test_gc_utils.py +++ b/tests/utils/test_gc_utils.py @@ -6,7 +6,11 @@ import logging import pytest from py.utils import gc_utils -from py.utils.gc_utils import freeze_long_lived_objects, schedule_refreeze_after_rebuild +from py.utils.gc_utils import ( + freeze_long_lived_objects, + schedule_refreeze, + schedule_refreeze_after_rebuild, +) @pytest.fixture(autouse=True) @@ -86,3 +90,42 @@ async def test_schedule_refreeze_debounces( await asyncio.sleep(0.2) assert fired == ["second"] + + +async def test_schedule_refreeze_waits_for_idle_queue( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fired: list[str] = [] + monkeypatch.setattr( + gc_utils, "freeze_long_lived_objects", lambda reason: fired.append(reason) + ) + idle = {"value": False} + monkeypatch.setattr( + gc_utils, "_comfyui_queue_is_idle", lambda: idle["value"] + ) + + schedule_refreeze("download", delay=0.02) + await asyncio.sleep(0.07) + assert fired == [] + + idle["value"] = True + await asyncio.sleep(0.1) + + assert fired == ["download"] + + +async def test_schedule_refreeze_gives_up_after_max_retries( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + fired: list[str] = [] + monkeypatch.setattr( + gc_utils, "freeze_long_lived_objects", lambda reason: fired.append(reason) + ) + monkeypatch.setattr(gc_utils, "_comfyui_queue_is_idle", lambda: False) + + with caplog.at_level(logging.DEBUG, logger="py.utils.gc_utils"): + schedule_refreeze("download", delay=0.01, _retries_left=2) + await asyncio.sleep(0.2) + + assert fired == [] + assert "skipped GC re-freeze" in caplog.text