diff --git a/core/leak_report.py b/core/leak_report.py index 16110ffc..905b4f7b 100644 --- a/core/leak_report.py +++ b/core/leak_report.py @@ -55,16 +55,59 @@ def report(mm, when) -> int: model = m.real_model() if model is None: continue - log.warning("FunUI Leak", f"{when}: {type(model).__name__} outlived its patcher. What holds it:\n" + "\n".join(chains(model))) + held = [r for r in gc.get_referrers(model) if type(r).__name__ == "dict" and "patches" in r] + owners = [o for h in held for o in gc.get_referrers(h) if type(o).__name__.startswith("ModelPatcher")] + + def depth(p): + n = 0 + while getattr(p, "parent", None) is not None and n < 50: + p, n = p.parent, n + 1 + return n + lineage = ", ".join(f"{type(o).__name__}(parents {depth(o)}, tracked {any(l.model is o for l in mm.current_loaded_models)})" for o in owners) + log.warning("FunUI Leak", f"{when}: {type(model).__name__} outlived its patcher. Patchers holding it: {lineage or 'none'}. What holds it:\n" + "\n".join(chains(model))) return len(dead) +_in_gc = [False] + + +def _gc_flag(phase, info): + _in_gc[0] = phase == "start" + + +def watch_switch(mm) -> bool: + """ComfyUI follows a model through its patcher's parent when the patcher dies. Say when that hand-over finds no parent, and whether + Python's own gc was running at that moment (it clears weak links to every patcher in a cycle before the hand-over can use them).""" + original = getattr(mm.LoadedModel, "_switch_parent", None) + if original is None or getattr(original, TAG, False): + return False + if _gc_flag not in gc.callbacks: + gc.callbacks.append(_gc_flag) + + def _switch_parent(self): + try: + parent = self._parent_model() + real = self.real_model() if callable(self.real_model) else None + if parent is None and real is not None: + live = [type(o).__name__ for o in gc.get_referrers(real) if type(o).__name__.startswith("ModelPatcher")] + log.warning("FunUI Leak", f"a {type(real).__name__} patcher died and its parent was already gone, so ComfyUI loses track of the model. " + f"Python's gc was running: {_in_gc[0]}. Patchers still holding the model: {len(live)}") + except Exception as exc: # noqa: BLE001 -- a diagnostic must never stop a run + log.failed("leak report", exc) + return original(self) + + setattr(_switch_parent, TAG, True) + mm.LoadedModel._switch_parent = _switch_parent + return True + + def install(mm=None) -> bool: if mm is None: import comfy.model_management as mm original = getattr(mm, "cleanup_models_gc", None) if original is None or getattr(original, TAG, False): return False + watch_switch(mm) def cleanup_models_gc(): try: diff --git a/core/tests/test_leak_report.py b/core/tests/test_leak_report.py index e6fa8c43..85828280 100644 --- a/core/tests/test_leak_report.py +++ b/core/tests/test_leak_report.py @@ -37,10 +37,40 @@ def test_a_model_kept_by_something_else_is_named_before_and_after_gc(comfyui, ca def test_install_wraps_once_and_a_clean_state_logs_nothing(comfyui, capsys): import comfy.model_management as mm - original = mm.cleanup_models_gc + original, switch = mm.cleanup_models_gc, mm.LoadedModel._switch_parent try: assert leak_report.install(mm) and not leak_report.install(mm) mm.cleanup_models_gc() assert "FunUI Leak" not in capsys.readouterr().err finally: - mm.cleanup_models_gc = original + mm.cleanup_models_gc, mm.LoadedModel._switch_parent = original, switch + if leak_report._gc_flag in gc.callbacks: + gc.callbacks.remove(leak_report._gc_flag) + + +def test_a_hand_over_that_finds_no_parent_is_said(comfyui, capsys): + import weakref + import comfy.model_management as mm + import comfy.model_patcher as mp + original = mm.LoadedModel._switch_parent + try: + assert leak_report.watch_switch(mm) and not leak_report.watch_switch(mm) + base = torch.nn.Linear(2, 2) + base.model_config = type("C", (), {})() + parent = mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu")) + child = parent.clone() + loaded = mm.LoadedModel(child) + loaded.real_model = weakref.ref(base) + keep = Holder(base) + del parent + gc.collect() # the parent is held by the child: still there + assert loaded._parent_model() is not None + loaded._parent_model = lambda: None # what a gc that cleared the weak link leaves behind + loaded._switch_parent() + shown = capsys.readouterr().err + assert "parent was already gone" in shown and "Python's gc was running: False" in shown + del keep + finally: + mm.LoadedModel._switch_parent = original + if leak_report._gc_flag in gc.callbacks: + gc.callbacks.remove(leak_report._gc_flag)