From 2ca4235b8936da1eadd8c08e4e3809a9988447b2 Mon Sep 17 00:00:00 2001 From: Folzo Date: Sun, 12 Apr 2026 12:14:59 +0200 Subject: [PATCH] feat: Add ComfyUI memory manager that restarts container after inactivity - Monitors VRAM, queue status, and history - Restarts ComfyUI container after timeout if RAM > 5GB and VRAM < 1% - More effective at freeing RAM than model unloading alone --- scripts/comfyui-memory-manager.py | 157 +++++++++++++++++++++++++ systemd/comfyui-memory-manager.service | 16 +++ 2 files changed, 173 insertions(+) create mode 100644 scripts/comfyui-memory-manager.py create mode 100644 systemd/comfyui-memory-manager.service diff --git a/scripts/comfyui-memory-manager.py b/scripts/comfyui-memory-manager.py new file mode 100644 index 0000000..a33df45 --- /dev/null +++ b/scripts/comfyui-memory-manager.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +# ComfyUI Memory Manager - Redémarre ComfyUI après inactivité +# Usage: python3 comfyui-memory-manager.py [--interval 60] [--timeout 600] + +import requests +import time +import argparse +import logging +import subprocess +import sys +import signal +from datetime import datetime + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) + +DEFAULT_INTERVAL = 60 +DEFAULT_TIMEOUT = 600 +COMFYUI_URL = "http://localhost:8188" +CONTAINER_NAME = "comfyui" + +class ComfyUIMemoryManager: + def __init__(self, interval=DEFAULT_INTERVAL, timeout=DEFAULT_TIMEOUT): + self.interval = interval + self.timeout = timeout + self.last_activity = datetime.now() + self.running = True + + signal.signal(signal.SIGINT, self.signal_handler) + signal.signal(signal.SIGTERM, self.signal_handler) + + def signal_handler(self, signum, frame): + logger.info("Signal reçu, arrêt...") + self.running = False + sys.exit(0) + + def get_stats(self): + try: + response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=10) + if response.status_code == 200: + return response.json() + except: + return None + + def get_vram_percent(self): + stats = self.get_stats() + if stats and "devices" in stats: + device = stats["devices"][0] + vram_total = device.get("vram_total", 0) + vram_free = device.get("vram_free", 0) + if vram_total > 0: + vram_used = vram_total - vram_free + return (vram_used / vram_total) * 100 + return 0 + + def get_queue_status(self): + try: + response = requests.get(f"{COMFYUI_URL}/queue", timeout=10) + if response.status_code == 200: + data = response.json() + running = len(data.get("queue_running", [])) + pending = len(data.get("queue_pending", [])) + return running, pending + except: + pass + return 0, 0 + + def get_history_count(self): + try: + response = requests.get(f"{COMFYUI_URL}/history?max_items=1", timeout=10) + if response.status_code == 200: + history = response.json() + return len(history) + except: + pass + return 0 + + def restart_container(self): + logger.info("🔄 Redémarrage du container ComfyUI...") + try: + subprocess.run(["docker", "restart", CONTAINER_NAME], check=True) + logger.info("✅ Container redémarré") + time.sleep(10) + return True + except Exception as e: + logger.error(f"❌ Erreur redémarrage: {e}") + return False + + def run(self): + logger.info(f"🚀 ComfyUI Memory Manager démarré") + logger.info(f" Vérification: {self.interval}s | Timeout inactivité: {self.timeout}s") + + while self.running: + try: + vram_percent = self.get_vram_percent() + running, pending = self.get_queue_status() + history_count = self.get_history_count() + + idle_seconds = (datetime.now() - self.last_activity).total_seconds() + + if running > 0 or pending > 0: + self.last_activity = datetime.now() + status = "active" + elif history_count > 0: + self.last_activity = datetime.now() + status = "recent" + else: + status = "idle" + + logger.info(f"VRAM: {vram_percent:.1f}% | Queue: {running}r/{pending}p | Idle: {idle_seconds:.0f}s | Status: {status}") + + if status == "idle" and vram_percent < 1 and idle_seconds >= self.timeout: + logger.info(f"⏰ Inactivité longue ({idle_seconds:.0f}s >= {self.timeout}s), VRAM libre") + + stats = self.get_stats() + if stats: + ram_total = stats["system"]["ram_total"] / (1024**3) + ram_free = stats["system"]["ram_free"] / (1024**3) + ram_used = ram_total - ram_free + logger.info(f"📊 RAM: {ram_used:.1f}GB/{ram_total:.1f}GB") + + if ram_used > 5: + logger.info("📊 RAM utilisée élevée, redémarrage pour libérer...") + self.restart_container() + else: + logger.info("📊 RAM déjà basse, pas de redémarrage nécessaire") + + self.last_activity = datetime.now() + + time.sleep(self.interval) + + except Exception as e: + logger.error(f"Erreur: {e}") + time.sleep(self.interval) + + logger.info("👋 Manager arrêté") + sys.exit(0) + +def main(): + parser = argparse.ArgumentParser(description="ComfyUI Memory Manager") + parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL) + parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT) + parser.add_argument("--debug", action="store_true") + + args = parser.parse_args() + + if args.debug: + logging.getLogger().setLevel(logging.DEBUG) + + manager = ComfyUIMemoryManager(interval=args.interval, timeout=args.timeout) + manager.run() + +if __name__ == "__main__": + main() diff --git a/systemd/comfyui-memory-manager.service b/systemd/comfyui-memory-manager.service new file mode 100644 index 0000000..730ca28 --- /dev/null +++ b/systemd/comfyui-memory-manager.service @@ -0,0 +1,16 @@ +[Unit] +Description=ComfyUI Memory Manager +After=network.target docker.service +Before=comfyui.service + +[Service] +Type=simple +User=folzo +WorkingDirectory=/home/folzo/Documents/OpenCode/Projets/Effi +ExecStart=/usr/bin/python3 /home/folzo/Documents/OpenCode/Projets/Effi/scripts/comfyui-memory-manager.py +Restart=always +RestartSec=10 +TimeoutStopSec=30 + +[Install] +WantedBy=multi-user.target