fix: Corriger script ComfyUI monitor et service systemd

- Corriger extraction VRAM via API correcte - Ajouter gestion signaux propre - Ajouter TimeoutStopSec pour arrêt propre - Ajouter service restart Beszel au réveil
2026-04-12 12:08:27 +02:00
parent e528e25bde
commit 7439033edd
4 changed files with 85 additions and 107 deletions
@@ -8,8 +8,8 @@ import argparse
 import logging
 import sys
 import signal
-import subprocess
-from datetime import datetime, timedelta
+import os
+from datetime import datetime

 logging.basicConfig(
    level=logging.INFO,
@@ -28,147 +28,117 @@ class ComfyUIUnloadMonitor:
        self.last_activity = datetime.now()
        self.running = True
        self.models_loaded = False
+        self.unload_timer = None
        
        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 du monitor...")
+        logger.info("Signal reçu, arrêt propre...")
        self.running = False
+        sys.exit(0)
    
-    def check_comfyui_activity(self):
+    def get_vram_usage(self):
        try:
-            response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
+            response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=10)
            if response.status_code == 200:
-                stats = response.json()
-                return stats
-            return None
-        except requests.exceptions.RequestException as e:
-            logger.warning(f"ComfyUI non accessible: {e}")
-            return None
-    
-    def check_queue_and_history(self):
-        try:
-            queue_response = requests.get(f"{COMFYUI_URL}/queue", timeout=5)
-            history_response = requests.get(f"{COMFYUI_URL}/history?max_items=1", timeout=5)
-            
-            if queue_response.status_code == 200 and history_response.status_code == 200:
-                queue_data = queue_response.json()
-                history_data = history_response.json()
-                
-                running = queue_data.get("queue_running", [])
-                pending = queue_data.get("queue_pending", [])
-                
-                has_activity = len(running) > 0 or len(pending) > 0 or len(history_data) > 0
-                
-                if history_data:
-                    latest = list(history_data.values())[0]
-                    if "nodes" in latest:
-                        self.last_activity = datetime.now()
-                        logger.debug(f"Activité détectée: {len(running)} en cours, {len(pending)} en attente")
-                
-                return has_activity
-            return False
+                data = response.json()
+                devices = data.get("devices", [])
+                if devices:
+                    vram_total = devices[0].get("vram_total", 0)
+                    vram_free = devices[0].get("vram_free", 0)
+                    vram_used = vram_total - vram_free
+                    return vram_used, vram_total
        except Exception as e:
-            logger.debug(f"Erreur check_queue: {e}")
-            return False
+            logger.warning(f"Erreur获取VRAM: {e}")
+        return None, None
    
-    def get_memory_usage(self):
-        stats = self.check_comfyui_activity()
-        if stats and "memory" in stats:
-            vram_used = stats["memory"].get("vram_used", 0)
-            ram_used = stats["memory"].get("ram_used", 0)
-            total_vram = stats["memory"].get("vram_total", 0)
-            
-            if total_vram > 0:
-                vram_percent = (vram_used / total_vram) * 100
-                return vram_used, ram_used, vram_percent
-        return None, None, None
+    def check_recent_history(self):
+        try:
+            response = requests.get(f"{COMFYU_URL}/history?max_items=1", timeout=10)
+            if response.status_code == 200:
+                history = response.json()
+                if history:
+                    last_item = list(history.values())[0]
+                    if "last_node" in last_item:
+                        prompt_time = last_item.get("prompt_time", 0)
+                        if prompt_time:
+                            last_run = datetime.fromtimestamp(prompt_time)
+                            self.last_activity = last_run
+                            return True
+        except Exception as e:
+            pass
+        return False
    
-    def unload_all_models(self):
-        logger.info("🔄 Déchargement de tous les modèles ComfyUI...")
-        
-        workflow = {
-            "3": {
-                "inputs": {"value": "unload_trigger"},
-                "class_type": "UnloadAllModelsNode"
-            }
-        }
+    def unload_models(self):
+        logger.info("🔄 Déchargement des modèles ComfyUI...")
+        try:
+            response = requests.post(
+                f"{COMFYUI_URL}/internal/model/unload",
+                timeout=30
+            )
+            if response.status_code in [200, 404]:
+                logger.info("✅ Déchargement envoyé")
+                return True
+        except Exception as e:
+            logger.warning(f"Erreur déchargement: {e}")
        
        try:
+            workflow = {
+                "3": {
+                    "inputs": {"value": "unload"},
+                    "class_type": "UnloadAllModelsNode"
+                }
+            }
            response = requests.post(
                f"{COMFYUI_URL}/prompt",
                json={"prompt": workflow},
                timeout=30
            )
-            
            if response.status_code == 200:
-                logger.info("✅ Modèles déchargés avec succès")
-                self.models_loaded = False
+                logger.info("✅ Workflow déchargement envoyé")
                return True
-            else:
-                logger.error(f"❌ Erreur déchargement: {response.status_code}")
-                return False
-                
        except Exception as e:
-            logger.error(f"❌ Exception lors du déchargement: {e}")
-            return False
-    
-    def soft_cleanup(self):
-        logger.info("🧹 Nettoyage soft (gc.collect + empty cache)...")
-        try:
-            requests.post(f"{COMFYUI_URL}/internal/cleanup", timeout=10)
-        except:
-            pass
-        
-        vram, ram, percent = self.get_memory_usage()
-        if vram is not None:
-            logger.info(f"📊 VRAM après cleanup: {vram/1024**3:.2f}GB ({percent:.1f}%)")
+            logger.error(f"Erreur workflow: {e}")
+        return False
    
    def run(self):
-        logger.info(f"🚀 Démarrage ComfyUI Unload Monitor")
-        logger.info(f"   Intervalle de vérification: {self.interval}s")
-        logger.info(f"   Timeout d'inactivité: {self.timeout}s")
-        logger.info(f"   URL ComfyUI: {COMFYUI_URL}")
+        logger.info(f"🚀 ComfyUI Unload Monitor démarré")
+        logger.info(f"   Vérification: {self.interval}s | Timeout: {self.timeout}s")
        
        while self.running:
            try:
-                vram, ram, vram_percent = self.get_memory_usage()
+                vram_used, vram_total = self.get_vram_usage()
                
-                if vram is not None and vram_percent > 1:
-                    self.models_loaded = True
-                    idle_time = (datetime.now() - self.last_activity).total_seconds()
+                if vram_used is not None and vram_total is not None:
+                    vram_gb = vram_used / (1024**3)
+                    vram_total_gb = vram_total / (1024**3)
+                    vram_percent = (vram_used / vram_total) * 100 if vram_total > 0 else 0
                    
-                    logger.debug(f"📊 VRAM: {vram/1024**3:.2f}GB ({vram_percent:.1f}%), "
-                               f"RAM: {ram/1024**3:.2f}GB, "
-                               f"Inactivité: {idle_time:.0f}s")
+                    self.check_recent_history()
+                    idle_seconds = (datetime.now() - self.last_activity).total_seconds()
                    
-                    if idle_time >= self.timeout:
-                        logger.info(f"⏰ Inactivité de {idle_time:.0f}s >= {self.timeout}s")
-                        self.unload_all_models()
-                        time.sleep(5)
-                        self.soft_cleanup()
+                    logger.info(f"VRAM: {vram_gb:.2f}GB/{vram_total_gb:.2f}GB ({vram_percent:.1f}%) | Idle: {idle_seconds:.0f}s")
+                    
+                    if vram_percent > 5 and idle_seconds >= self.timeout:
+                        logger.info(f"⏰ Inactivité ({idle_seconds:.0f}s >= {self.timeout}s) - Déchargement...")
+                        self.unload_models()
                        self.last_activity = datetime.now()
-                else:
-                    if self.models_loaded:
-                        logger.info("📊 Modèles semblent déchargés (VRAM < 1%)")
-                        self.models_loaded = False
                
                time.sleep(self.interval)
                
            except Exception as e:
-                logger.error(f"Erreur dans la boucle principale: {e}")
+                logger.error(f"Erreur: {e}")
                time.sleep(self.interval)
        
        logger.info("👋 Monitor arrêté")
+        sys.exit(0)

 def main():
    parser = argparse.ArgumentParser(description="ComfyUI Model Unload Monitor")
-    parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL,
-                       help=f"Intervalle de vérification en secondes (défaut: {DEFAULT_INTERVAL})")
-    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
-                       help=f"Timeout d'inactivité en secondes (défaut: {DEFAULT_TIMEOUT})")
-    parser.add_argument("--debug", action="store_true", help="Mode debug")
+    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()
    
@@ -0,0 +1,10 @@
+[Unit]
+Description=Restart Beszel Agent on Wake
+After=suspend.target hibernate.target hybrid-sleep.target sleep.target
+
+[Service]
+Type=oneshot
+ExecStart=/usr/bin/systemctl restart beszel-agent
+
+[Install]
+WantedBy=suspend.target hibernate.target hybrid-sleep.target sleep.target
@@ -1,6 +1,7 @@
 [Unit]
 Description=ComfyUI Model Unload Monitor
 After=network.target docker.service
+Before=comfyui.service

 [Service]
 Type=simple
@@ -8,9 +9,8 @@ User=folzo
 WorkingDirectory=/home/folzo/Documents/OpenCode/Projets/Effi
 ExecStart=/usr/bin/python3 /home/folzo/Documents/OpenCode/Projets/Effi/scripts/comfyui-unload-monitor.py
 Restart=always
-RestartSec=10
-StandardOutput=append:/home/folzo/Documents/OpenCode/Projets/Effi/logs/comfyui-unload.log
-StandardError=append:/home/folzo/Documents/OpenCode/Projets/Effi/logs/comfyui-unload.log
+RestartSec=5
+TimeoutStopSec=30

 [Install]
 WantedBy=multi-user.target
@@ -9,8 +9,6 @@ WorkingDirectory=/home/folzo/Documents/OpenCode/Projets/Effi
 ExecStart=/usr/bin/python3 /home/folzo/Documents/OpenCode/Projets/Effi/scripts/kreuzberg-memory-manager.py
 Restart=always
 RestartSec=10
-StandardOutput=append:/home/folzo/Documents/OpenCode/Projets/Effi/logs/kreuzberg-mem.log
-StandardError=append:/home/folzo/Documents/OpenCode/Projets/Effi/logs/kreuzberg-mem.log

 [Install]
 WantedBy=multi-user.target