301 lines
13 KiB
Python
301 lines
13 KiB
Python
"""
|
|
Model Management Extensions for MultiGPU
|
|
Extends ComfyUI's model management with multi-device capabilities and lifecycle tracking.
|
|
"""
|
|
|
|
import torch
|
|
import logging
|
|
import hashlib
|
|
import psutil
|
|
import comfy.model_management as mm
|
|
import gc
|
|
from datetime import datetime, timezone
|
|
import server
|
|
import weakref
|
|
import platform
|
|
import ctypes
|
|
import comfy.model_patcher
|
|
from collections import defaultdict
|
|
|
|
|
|
|
|
logger = logging.getLogger("MultiGPU")
|
|
|
|
# ==========================================================================================
|
|
# Model Analysis and Store Management (DisTorch V1 & V2)
|
|
# ==========================================================================================
|
|
|
|
# DisTorch V2 SafeTensor stores
|
|
safetensor_allocation_store = {}
|
|
safetensor_settings_store = {}
|
|
|
|
# DisTorch V1 GGUF stores (backwards compatibility)
|
|
model_allocation_store = {}
|
|
|
|
def create_safetensor_model_hash(model, caller):
|
|
"""Create a unique hash for a safetensor model to track allocations"""
|
|
if hasattr(model, 'model'):
|
|
actual_model = model.model
|
|
model_type = type(actual_model).__name__
|
|
model_size = model.model_size() if hasattr(model, 'model_size') else sum(p.numel() * p.element_size() for p in actual_model.parameters())
|
|
first_layers = str(list(model.model_state_dict().keys() if hasattr(model, 'model_state_dict') else actual_model.state_dict().keys())[:3])
|
|
else:
|
|
model_type = type(model).__name__
|
|
model_size = sum(p.numel() * p.element_size() for p in model.parameters())
|
|
first_layers = str(list(model.state_dict().keys())[:3])
|
|
|
|
identifier = f"{model_type}_{model_size}_{first_layers}"
|
|
final_hash = hashlib.sha256(identifier.encode()).hexdigest()
|
|
logger.debug(f"[MultiGPU DisTorch V2] Created hash for {caller}: {final_hash[:8]}...")
|
|
return final_hash
|
|
|
|
def create_model_hash(model, caller):
|
|
"""Create a unique hash for a GGUF model to track allocations (DisTorch V1)"""
|
|
model_type = type(model.model).__name__
|
|
model_size = model.model_size()
|
|
first_layers = str(list(model.model_state_dict().keys())[:3])
|
|
identifier = f"{model_type}_{model_size}_{first_layers}"
|
|
final_hash = hashlib.sha256(identifier.encode()).hexdigest()
|
|
logger.debug(f"[MultiGPU_DisTorch_HASH] Created hash for {caller}: {final_hash[:8]}...")
|
|
return final_hash
|
|
|
|
# ==========================================================================================
|
|
# Memory Logging Infrastructure
|
|
# ==========================================================================================
|
|
|
|
_MEM_SNAPSHOT_LAST = {}
|
|
_MEM_SNAPSHOT_SERIES = {}
|
|
|
|
def _capture_memory_snapshot():
|
|
"""Capture memory snapshot for CPU and all devices"""
|
|
# Import here to avoid circular dependency
|
|
from .device_utils import get_device_list
|
|
|
|
snapshot = {}
|
|
|
|
# CPU
|
|
vm = psutil.virtual_memory()
|
|
snapshot["cpu"] = (vm.used, vm.total)
|
|
|
|
# GPU devices
|
|
devices = [d for d in get_device_list() if d != "cpu"]
|
|
for dev_str in devices:
|
|
device = torch.device(dev_str)
|
|
total = mm.get_total_memory(device)
|
|
free_info = mm.get_free_memory(device, torch_free_too=True)
|
|
system_free = free_info[0] if isinstance(free_info, tuple) else free_info
|
|
used = max(0, total - system_free)
|
|
snapshot[dev_str] = (used, total)
|
|
|
|
return snapshot
|
|
|
|
def multigpu_memory_log(identifier, tag):
|
|
"""Record timestamped memory snapshot with clean aligned logging"""
|
|
if identifier == "print_summary":
|
|
for id_key in sorted(_MEM_SNAPSHOT_SERIES.keys()):
|
|
series = _MEM_SNAPSHOT_SERIES[id_key]
|
|
logger.mgpu_mm_log(f"=== memory summary: {id_key} ===")
|
|
for ts, tag_name, snap in series:
|
|
parts = []
|
|
cpu_used, cpu_total = snap.get("cpu", (0, 0))
|
|
parts.append(f"cpu|{cpu_used/(1024**3):.2f}")
|
|
for dev in sorted([k for k in snap.keys() if k != "cpu"]):
|
|
used, total = snap[dev]
|
|
parts.append(f"{dev}|{used/(1024**3):.2f}")
|
|
ts_str = ts.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
|
|
tag_padded = f"{id_key}_{tag_name}".ljust(35)
|
|
logger.mgpu_mm_log(f"{ts_str} {tag_padded} {' '.join(parts)}")
|
|
return
|
|
|
|
ts = datetime.now(timezone.utc)
|
|
curr = _capture_memory_snapshot()
|
|
|
|
# Store in series
|
|
if identifier not in _MEM_SNAPSHOT_SERIES:
|
|
_MEM_SNAPSHOT_SERIES[identifier] = []
|
|
_MEM_SNAPSHOT_SERIES[identifier].append((ts, tag, curr))
|
|
|
|
# Clean aligned format: timestamp + padded tag + memory values
|
|
ts_str = ts.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "Z"
|
|
tag_padded = f"{identifier}_{tag}".ljust(35)
|
|
|
|
parts = []
|
|
cpu_used, _ = curr.get("cpu", (0, 0))
|
|
parts.append(f"cpu|{cpu_used/(1024**3):.2f}")
|
|
|
|
for dev in sorted([k for k in curr.keys() if k != "cpu"]):
|
|
used, _ = curr[dev]
|
|
parts.append(f"{dev}|{used/(1024**3):.2f}")
|
|
|
|
logger.mgpu_mm_log(f"{ts_str} {tag_padded} {' '.join(parts)}")
|
|
|
|
_MEM_SNAPSHOT_LAST[identifier] = (tag, curr)
|
|
|
|
|
|
# ==========================================================================================
|
|
# Memory Management and Cleanup
|
|
# ==========================================================================================
|
|
|
|
CPU_MEMORY_THRESHOLD_PERCENT = 85.0
|
|
CPU_RESET_HYSTERESIS_PERCENT = 5.0
|
|
_last_cpu_usage_at_reset = 0.0
|
|
|
|
|
|
def trigger_executor_cache_reset(reason="policy", force=False):
|
|
"""Trigger PromptExecutor.reset() by setting 'free_memory' flag"""
|
|
global _last_cpu_usage_at_reset
|
|
|
|
prompt_server = server.PromptServer.instance
|
|
if prompt_server is None:
|
|
logger.debug("[MultiGPU_Memory_Management] PromptServer not initialized")
|
|
return
|
|
|
|
if prompt_server.prompt_queue.currently_running and not force:
|
|
logger.debug(f"[MultiGPU_Memory_Management] Skipping reset during execution (reason: {reason})")
|
|
return
|
|
|
|
multigpu_memory_log("executor_reset", f"pre-trigger ({reason})")
|
|
logger.info(f"[MultiGPU_Memory_Management] Triggering PromptExecutor cache reset. Reason: {reason}")
|
|
|
|
prompt_server.prompt_queue.set_flag("free_memory", True)
|
|
logger.debug("[MultiGPU_Memory_Management] 'free_memory' flag set")
|
|
|
|
vm = psutil.virtual_memory()
|
|
_last_cpu_usage_at_reset = vm.percent
|
|
|
|
multigpu_memory_log("executor_reset", f"post-trigger ({reason})")
|
|
|
|
def check_cpu_memory_threshold(threshold_percent=CPU_MEMORY_THRESHOLD_PERCENT):
|
|
"""Check CPU memory and trigger reset if threshold exceeded"""
|
|
if server.PromptServer.instance is None:
|
|
return
|
|
|
|
if server.PromptServer.instance.prompt_queue.currently_running:
|
|
return
|
|
|
|
vm = psutil.virtual_memory()
|
|
current_usage = vm.percent
|
|
|
|
if current_usage > threshold_percent:
|
|
if current_usage > (_last_cpu_usage_at_reset + CPU_RESET_HYSTERESIS_PERCENT):
|
|
logger.warning(f"[MultiGPU_Memory_Monitor] CPU usage ({current_usage:.1f}%) exceeds threshold ({threshold_percent:.1f}%)")
|
|
multigpu_memory_log("cpu_monitor", f"trigger:{current_usage:.1f}pct")
|
|
trigger_executor_cache_reset(reason="cpu_threshold_exceeded", force=False)
|
|
else:
|
|
logger.debug(f"[MultiGPU_Memory_Monitor] CPU usage high ({current_usage:.1f}%) but within hysteresis")
|
|
multigpu_memory_log("cpu_monitor", f"skip_hysteresis:{current_usage:.1f}pct")
|
|
|
|
def force_full_system_cleanup(reason="manual", force=True):
|
|
"""
|
|
Mirror ComfyUI-Manager 'Free model and node cache' by setting both flags:
|
|
unload_models=True and free_memory=True
|
|
"""
|
|
vm = psutil.virtual_memory()
|
|
pre_cpu = vm.used
|
|
pre_models = len(mm.current_loaded_models)
|
|
|
|
multigpu_memory_log("full_cleanup", f"start:{reason}")
|
|
logger.mgpu_mm_log(f"[ManagerMatch] Requesting cleanup (reason={reason}) | pre_models={pre_models}, cpu_used_gib={pre_cpu/(1024**3):.2f}")
|
|
|
|
if server.PromptServer.instance is not None:
|
|
pq = server.PromptServer.instance.prompt_queue
|
|
if (not pq.currently_running) or force:
|
|
pq.set_flag("unload_models", True)
|
|
pq.set_flag("free_memory", True)
|
|
logger.mgpu_mm_log("[ManagerMatch] Flags set: unload_models=True, free_memory=True")
|
|
else:
|
|
logger.mgpu_mm_log("[ManagerMatch] Skipped - execution active and force=False")
|
|
|
|
vm = psutil.virtual_memory()
|
|
post_cpu = vm.used
|
|
post_models = len(mm.current_loaded_models)
|
|
delta_cpu_mb = (post_cpu - pre_cpu) / (1024**2)
|
|
|
|
multigpu_memory_log("full_cleanup", f"requested:{reason}")
|
|
summary = f"[ManagerMatch] Cleanup requested (reason={reason}) | models {pre_models}->{post_models}, cpu_delta_mb={delta_cpu_mb:.2f}"
|
|
logger.mgpu_mm_log(summary)
|
|
return summary
|
|
|
|
# ==========================================================================================
|
|
# Core Patching: unload_all_models
|
|
# ==========================================================================================
|
|
|
|
if not hasattr(mm.unload_all_models, '_mgpu_eject_distorch_patched'):
|
|
logger.info("[MultiGPU Core Patching] Patching mm.unload_all_models for DisTorch2 ejection support")
|
|
|
|
_mgpu_original_unload_all_models = mm.unload_all_models
|
|
|
|
def _mgpu_patched_unload_all_models():
|
|
"""
|
|
Patched mm.unload_all_models that checks to see if the .
|
|
All other models (including DisTorch models without the flag) unload normally.
|
|
"""
|
|
from . import DISTORCH2_UNLOAD_MODEL
|
|
|
|
logger.mgpu_mm_log(f"[Phase 2 Debug] Patched unload_all_models called - initial model count: {len(mm.current_loaded_models)}")
|
|
logger.mgpu_mm_log(f"[Phase 2 Debug] DISTORCH2_UNLOAD_MODEL={DISTORCH2_UNLOAD_MODEL}")
|
|
|
|
if DISTORCH2_UNLOAD_MODEL == False:
|
|
logger.mgpu_mm_log("[Phase 2 Debug] Standard unload_all_models() called from Comfy Core")
|
|
_mgpu_original_unload_all_models()
|
|
return
|
|
|
|
# Direct approach: iterate through loaded models and selectively unload
|
|
models_to_unload = []
|
|
kept_models = []
|
|
|
|
for i, lm in enumerate(mm.current_loaded_models):
|
|
mp = lm.model # weakref call to ModelPatcher
|
|
if mp is not None and hasattr(mp, 'model'):
|
|
# Check if this is a DisTorch model with keep_loaded flag
|
|
should_retain = getattr(mp.model, '_mgpu_keep_loaded', True)
|
|
model_name = type(getattr(mp, 'model', mp)).__name__
|
|
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={should_retain}")
|
|
|
|
# Retain models that either:
|
|
# 1. Are non-DisTorch models (missing _mgpu_keep_loaded attribute)
|
|
# 2. Are DisTorch models with keep_loaded=True
|
|
|
|
if should_retain:
|
|
kept_models.append(lm)
|
|
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Adding to kept_models: {model_name}")
|
|
# GC ANCHOR TEST: Prevent premature GC of clone patchers
|
|
add_retention_anchor(mp, "keep_loaded_test")
|
|
else:
|
|
models_to_unload.append(lm)
|
|
else:
|
|
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: ModelPatcher is None or missing model attribute")
|
|
models_to_unload.append(lm)
|
|
|
|
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Final counts - kept_models: {len(kept_models)}, models_to_unload: {len(models_to_unload)}")
|
|
|
|
if kept_models:
|
|
logger.mgpu_mm_log(f"Found {len(kept_models)} model(s) to retain, unloading {len(models_to_unload)} model(s)")
|
|
|
|
# Unload models that don't have keep_loaded flag
|
|
for lm in models_to_unload:
|
|
try:
|
|
lm.model_unload(unpatch_weights=True)
|
|
logger.debug(f"Unloaded model: {type(lm.model.model).__name__ if lm.model else 'Unknown'}")
|
|
except Exception as e:
|
|
logger.warning(f"Error unloading model: {e}")
|
|
|
|
# Remove unloaded models from current_loaded_models
|
|
mm.current_loaded_models = kept_models
|
|
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Updated mm.current_loaded_models, new count: {len(mm.current_loaded_models)}")
|
|
logger.mgpu_mm_log(f"Successfully retained {len(kept_models)} model(s) during unload")
|
|
else:
|
|
logger.mgpu_mm_log("No models with keep_loaded=True found - delegating to original unload_all_models")
|
|
_mgpu_original_unload_all_models()
|
|
|
|
# Phase 1: Reset DISTORCH2_UNLOAD_MODEL flag at end of unload (REGARDLESS)
|
|
logger.mgpu_mm_log("[PHASE1_DEBUG] Setting DISTORCH2_UNLOAD_MODEL=False at end of unload")
|
|
multigpu_memory_log("distorch_flag", "reset_false")
|
|
DISTORCH2_UNLOAD_MODEL = False
|
|
|
|
mm.unload_all_models = _mgpu_patched_unload_all_models
|
|
mm.unload_all_models._mgpu_eject_distorch_patched = True
|
|
logger.info("[MultiGPU Core Patching] mm.unload_all_models patched successfully")
|
|
else:
|
|
logger.debug("[MultiGPU Core Patching] mm.unload_all_models already patched - skipping")
|