feat(memory): Ajout déchargement automatique mémoire sur FolzoCorp
- effi-stt: Déchargement Whisper après 300s inactivité (INACTIVITY_TIMEOUT) * Nouveaux endpoints: /unload, /reload, /stats * Modelo rechargé à la première requête - ComfyUI: Script monitor qui décharge VRAM après inactivité * Vérifie usage VRAM toutes les 60s * Exécute workflow UnloadAllModelsNode après 300s inactivité - Kreuzberg: Manager qui redémarre le container après seuil mémoire + inactivité * Restart si > 2GB pendant > 10min inactif - Documentation: docs/MEMOIRE-GESTION.md
This commit is contained in:
@@ -0,0 +1,202 @@
|
||||
# Gestion Mémoire - FolzoCorp
|
||||
|
||||
## Vue d'ensemble
|
||||
|
||||
Ce document décrit les mécanismes de gestion mémoire implementés sur **FolzoCorp** (192.168.1.11) pour optimiser l'utilisation RAM.
|
||||
|
||||
## Problème initial
|
||||
|
||||
- **RAM Totale**: 32 GB
|
||||
- **RAM Utilisée**: ~29 GB (93%)
|
||||
- **Swap utilisé**: 8 GB (100%)
|
||||
- **Cause principale**: ComfyUI (15.2 GB) + services permanentents
|
||||
|
||||
## Solutions implémentées
|
||||
|
||||
### 1. effi-stt (Whisper STT)
|
||||
|
||||
**Fichier**: `effi-stt/main.py`
|
||||
|
||||
**Mécanisme**: Déchargement automatique après inactivité
|
||||
|
||||
**Fonctionnement**:
|
||||
- Le modèle Whisper (`large-v3`) est chargé à la première utilisation
|
||||
- Après chaque requête, un timer de 300s (5 minutes) est démarré
|
||||
- Si aucune requête n'est reçue pendant ce délai, le modèle est déchargé
|
||||
- La mémoire est libérée via `gc.collect()`
|
||||
|
||||
**Nouveaux endpoints**:
|
||||
| Endpoint | Méthode | Description |
|
||||
|----------|---------|-------------|
|
||||
| `/unload` | POST | Déchargement manuel du modèle |
|
||||
| `/reload` | POST | Rechargement manuel du modèle |
|
||||
| `/stats` | GET | Statistiques (model_loaded, last_request_age) |
|
||||
|
||||
**Variables d'environnement**:
|
||||
| Variable | Défaut | Description |
|
||||
|----------|--------|-------------|
|
||||
| `INACTIVITY_TIMEOUT` | 300 | Timeout en secondes avant déchargement |
|
||||
|
||||
**Commandes**:
|
||||
```bash
|
||||
# Redémarrer avec les nouvelles variables
|
||||
cd effi-stt
|
||||
docker compose down
|
||||
docker compose up -d --build
|
||||
|
||||
# Vérifier le statut
|
||||
curl http://localhost:8200/stats
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 2. ComfyUI (Image + Audio Generation)
|
||||
|
||||
**Fichier**: `scripts/comfyui-unload-monitor.py`
|
||||
|
||||
**Mécanisme**: Monitoring externe + déchargement via API
|
||||
|
||||
**Fonctionnement**:
|
||||
- Script Python qui tourne en arrière-plan
|
||||
- Vérifie l'utilisation VRAM toutes les 60s
|
||||
- Après 300s (5 minutes) d'inactivité + VRAM utilisée:
|
||||
- Exécute un workflow pour décharger tous les modèles
|
||||
- Appelle `gc.collect()` et `torch.cuda.empty_cache()`
|
||||
|
||||
**Utilisation**:
|
||||
```bash
|
||||
# Lancer le monitor
|
||||
python3 scripts/comfyui-unload-monitor.py
|
||||
|
||||
# Avec paramètres personnalisés
|
||||
python3 scripts/comfyui-unload-monitor.py --timeout 600 --interval 30
|
||||
|
||||
# Mode debug
|
||||
python3 scripts/comfyui-unload-monitor.py --debug
|
||||
|
||||
# En arrière-plan (systemd)
|
||||
nohup python3 scripts/comfyui-unload-monitor.py > /var/log/comfyui-unload.log 2>&1 &
|
||||
```
|
||||
|
||||
**Démarrage automatique (systemd)**:
|
||||
```bash
|
||||
# Créer le service
|
||||
sudo tee /etc/systemd/system/comfyui-unload.service << EOF
|
||||
[Unit]
|
||||
Description=ComfyUI Model Unload Monitor
|
||||
After=network.target docker.service
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
User=folzo
|
||||
WorkingDirectory=/home/folzo/Documents/OpenCode/Projets/Effi
|
||||
ExecStart=/usr/bin/python3 scripts/comfyui-unload-monitor.py
|
||||
Restart=always
|
||||
RestartSec=10
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
EOF
|
||||
|
||||
# Activer le service
|
||||
sudo systemctl daemon-reload
|
||||
sudo systemctl enable comfyui-unload
|
||||
sudo systemctl start comfyui-unload
|
||||
|
||||
# Vérifier le statut
|
||||
sudo systemctl status comfyui-unload
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### 3. Kreuzberg (Document Extraction)
|
||||
|
||||
**Fichier**: `scripts/kreuzberg-memory-manager.py`
|
||||
|
||||
**Mécanisme**: Restart automatique après inactivité + seuil mémoire
|
||||
|
||||
**Fonctionnement**:
|
||||
- Vérifie la mémoire du container toutes les 60s
|
||||
- Si mémoire > 2GB pendant > 10 minutes d'inactivité:
|
||||
- Redémarre le container kreuzberg
|
||||
- Libère toute la mémoire accumulée
|
||||
|
||||
**Note**: Kreuzberg est un service pre-built (image Rust), donc le déchargement interne n'est pas possible. Le restart container est la seule option pour libérer la mémoire.
|
||||
|
||||
**Utilisation**:
|
||||
```bash
|
||||
# Lancer le manager
|
||||
python3 scripts/kreuzberg-memory-manager.py
|
||||
|
||||
# Avec paramètres personnalisés
|
||||
python3 scripts/kreuzberg-memory-manager.py --timeout 600 --max-memory 1.5
|
||||
|
||||
# En arrière-plan (systemd)
|
||||
nohup python3 scripts/kreuzberg-memory-manager.py > /var/log/kreuzberg-mem.log 2>&1 &
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Impact mémoire attendu
|
||||
|
||||
| Service | Avant | Après (inactif) |
|
||||
|---------|-------|-----------------|
|
||||
| ComfyUI | 15.2 GB | ~500 MB (VRAM libérée) |
|
||||
| effi-stt | ~3 GB | ~100 MB (modèle déchargé) |
|
||||
| kreuzberg | 1 GB | ~100 MB (container redémarré) |
|
||||
| **Total** | **~20 GB** | **< 1 GB** |
|
||||
|
||||
---
|
||||
|
||||
## Surveillance
|
||||
|
||||
### Commandes utiles
|
||||
|
||||
```bash
|
||||
# Vérifier mémoire containers
|
||||
docker stats --no-stream
|
||||
|
||||
# Vérifier VRAM GPU
|
||||
rocm-smi --showmeminfo vram
|
||||
|
||||
# Statuts effi-stt
|
||||
curl http://localhost:8200/stats
|
||||
|
||||
# Logs du monitor ComfyUI
|
||||
tail -f /var/log/comfyui-unload.log
|
||||
|
||||
# Logs du manager Kreuzberg
|
||||
tail -f /var/log/kreuzberg-mem.log
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Fichiers modifiés/créés
|
||||
|
||||
### Modifiés
|
||||
- `effi-stt/main.py` - Ajout timeout + déchargement
|
||||
- `effi-stt/Dockerfile` - Variable INACTIVITY_TIMEOUT
|
||||
- `effi-stt/docker-compose.yml` - Variable INACTIVITY_TIMEOUT
|
||||
|
||||
### Créés
|
||||
- `scripts/comfyui-unload-monitor.py` - Monitor déchargement ComfyUI
|
||||
- `scripts/kreuzberg-memory-manager.py` - Manager mémoire Kreuzberg
|
||||
- `docs/MEMOIRE-GESTION.md` - Cette documentation
|
||||
|
||||
---
|
||||
|
||||
## Limitations connues
|
||||
|
||||
1. **ComfyUI**: Le monitor utilise l'API ComfyUI. Si ComfyUI est complètement bloqué, le monitor pourrait ne pas fonctionner.
|
||||
|
||||
2. **Kreuzberg**: Le restart du container interrompt brièvement le service (~5-10s).
|
||||
|
||||
3. **effi-stt**: Le premier appel après déchargement sera plus lent (rechargement du modèle).
|
||||
|
||||
---
|
||||
|
||||
## Améliorations futures possibles
|
||||
|
||||
1. Intégrer les monitors dans Docker Compose avec des sidecars
|
||||
2. Ajouter des métriques Prometheus pour le monitoring
|
||||
3. Implémenter un agente Effi qui surveille et actionne selon les besoins
|
||||
@@ -23,6 +23,7 @@ ENV DEVICE=cpu
|
||||
ENV COMPUTE_TYPE=int8
|
||||
ENV LANGUAGE=fr
|
||||
ENV CPU_THREADS=4
|
||||
ENV INACTIVITY_TIMEOUT=300
|
||||
ENV PORT=8200
|
||||
ENV HOST=0.0.0.0
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ services:
|
||||
- COMPUTE_TYPE=int8
|
||||
- LANGUAGE=fr
|
||||
- CPU_THREADS=4
|
||||
- INACTIVITY_TIMEOUT=300
|
||||
volumes:
|
||||
# Cache pour les modèles (évite de re-télécharger)
|
||||
- effi-stt-cache:/root/.cache/huggingface
|
||||
|
||||
+119
-63
@@ -1,5 +1,6 @@
|
||||
# 🎙️ Effi-STT - Speech-to-Text pour Effi
|
||||
# Effi-STT - Speech-to-Text pour Effi
|
||||
# API OpenAI-compatible avec faster-whisper
|
||||
# Inclut déchargement automatique après inactivité
|
||||
|
||||
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
@@ -9,39 +10,100 @@ import tempfile
|
||||
import os
|
||||
from typing import Optional
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
|
||||
# Configuration du logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
app = FastAPI(
|
||||
title="Effi-STT",
|
||||
description="Speech-to-Text API compatible OpenAI pour Effi",
|
||||
version="1.0.0"
|
||||
version="1.1.0"
|
||||
)
|
||||
|
||||
# Configuration
|
||||
MODEL_SIZE = os.getenv("MODEL_SIZE", "large-v3")
|
||||
DEVICE = os.getenv("DEVICE", "cpu")
|
||||
COMPUTE_TYPE = os.getenv("COMPUTE_TYPE", "int8")
|
||||
LANGUAGE = os.getenv("LANGUAGE", "fr")
|
||||
INACTIVITY_TIMEOUT = int(os.getenv("INACTIVITY_TIMEOUT", "300"))
|
||||
PORT = int(os.getenv("PORT", "8200"))
|
||||
HOST = os.getenv("HOST", "0.0.0.0")
|
||||
|
||||
logger.info(f"🎙️ Chargement du modèle Whisper: {MODEL_SIZE}")
|
||||
logger.info(f" Device: {DEVICE}, Compute type: {COMPUTE_TYPE}")
|
||||
model = None
|
||||
model_loaded = False
|
||||
model_lock = threading.Lock()
|
||||
last_request_time = time.time()
|
||||
unload_timer = None
|
||||
|
||||
# Chargement du modèle au démarrage
|
||||
model = WhisperModel(
|
||||
MODEL_SIZE,
|
||||
device=DEVICE,
|
||||
compute_type=COMPUTE_TYPE,
|
||||
cpu_threads=int(os.getenv("CPU_THREADS", "4"))
|
||||
)
|
||||
logger.info("✅ Modèle chargé avec succès!")
|
||||
def load_model():
|
||||
global model, model_loaded
|
||||
with model_lock:
|
||||
if not model_loaded:
|
||||
logger.info(f"Chargement du modèle Whisper: {MODEL_SIZE}")
|
||||
logger.info(f" Device: {DEVICE}, Compute type: {COMPUTE_TYPE}")
|
||||
model = WhisperModel(
|
||||
MODEL_SIZE,
|
||||
device=DEVICE,
|
||||
compute_type=COMPUTE_TYPE,
|
||||
cpu_threads=int(os.getenv("CPU_THREADS", "4"))
|
||||
)
|
||||
model_loaded = True
|
||||
logger.info("✅ Modèle chargé avec succès!")
|
||||
|
||||
def unload_model():
|
||||
global model, model_loaded
|
||||
with model_lock:
|
||||
if model_loaded:
|
||||
logger.info("🕐 Inactivité détectée - Déchargement du modèle Whisper...")
|
||||
del model
|
||||
model = None
|
||||
model_loaded = False
|
||||
import gc
|
||||
gc.collect()
|
||||
logger.info("✅ Modèle déchargé, mémoire libérée")
|
||||
|
||||
def reset_unload_timer():
|
||||
global last_request_time, unload_timer
|
||||
last_request_time = time.time()
|
||||
if unload_timer is not None:
|
||||
unload_timer.cancel()
|
||||
if model_loaded:
|
||||
timer = threading.Timer(INACTIVITY_TIMEOUT, unload_model)
|
||||
timer.daemon = True
|
||||
timer.start()
|
||||
unload_timer = timer
|
||||
logger.debug(f"⏱️ Timer déchargement programmé dans {INACTIVITY_TIMEOUT}s")
|
||||
|
||||
def check_and_load_model():
|
||||
global model_loaded
|
||||
if not model_loaded:
|
||||
logger.info("📥 Chargement du modèle Whisper à la demande...")
|
||||
load_model()
|
||||
reset_unload_timer()
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup():
|
||||
logger.info(f"🚀 Démarrage effi-stt v1.1.0 sur http://{HOST}:{PORT}")
|
||||
logger.info(f" Timeout d'inactivité: {INACTIVITY_TIMEOUT}s")
|
||||
load_model()
|
||||
reset_unload_timer()
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown():
|
||||
if unload_timer:
|
||||
unload_timer.cancel()
|
||||
unload_model()
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
"""Endpoint de vérification de santé"""
|
||||
return {"status": "healthy", "model": MODEL_SIZE}
|
||||
check_and_load_model()
|
||||
return {
|
||||
"status": "healthy",
|
||||
"model": MODEL_SIZE,
|
||||
"model_loaded": model_loaded,
|
||||
"inactivity_timeout": INACTIVITY_TIMEOUT
|
||||
}
|
||||
|
||||
@app.post("/v1/audio/transcriptions")
|
||||
async def transcribe(
|
||||
@@ -51,88 +113,82 @@ async def transcribe(
|
||||
prompt: Optional[str] = Form(None, description="Prompt optionnel"),
|
||||
response_format: Optional[str] = Form("json", description="Format de réponse")
|
||||
):
|
||||
"""
|
||||
Endpoint de transcription compatible API OpenAI
|
||||
"""
|
||||
try:
|
||||
logger.info(f"📝 Transcription demandée: lang={language}, file={file.filename}")
|
||||
check_and_load_model()
|
||||
reset_unload_timer()
|
||||
logger.info(f"📝 Transcription: lang={language}, file={file.filename}")
|
||||
|
||||
# Sauvegarder le fichier temporairement
|
||||
suffix = os.path.splitext(file.filename)[1] or ".mp3"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
content = await file.read()
|
||||
tmp.write(content)
|
||||
tmp_path = tmp.name
|
||||
logger.info(f" Fichier sauvegardé: {tmp_path} ({len(content)} bytes)")
|
||||
logger.info(f" Fichier: {tmp_path} ({len(content)} bytes)")
|
||||
|
||||
try:
|
||||
# Transcription avec faster-whisper
|
||||
segments, info = model.transcribe(
|
||||
tmp_path,
|
||||
language=language if language != "auto" else None,
|
||||
task="transcribe",
|
||||
beam_size=5,
|
||||
best_of=5,
|
||||
condition_on_previous_text=True
|
||||
)
|
||||
with model_lock:
|
||||
segments, info = model.transcribe(
|
||||
tmp_path,
|
||||
language=language if language != "auto" else None,
|
||||
task="transcribe",
|
||||
beam_size=5,
|
||||
best_of=5,
|
||||
condition_on_previous_text=True
|
||||
)
|
||||
|
||||
# Concaténer tous les segments
|
||||
text_parts = []
|
||||
for segment in segments:
|
||||
text_parts.append(segment.text)
|
||||
|
||||
full_text = " ".join(text_parts).strip()
|
||||
logger.info(f"✅ Transcription: '{full_text[:50]}...'")
|
||||
logger.info(f" Langue: {info.language} ({info.language_probability:.2f})")
|
||||
|
||||
logger.info(f"✅ Transcription réussie: '{full_text[:50]}...'")
|
||||
logger.info(f" Langue détectée: {info.language}, Probabilité: {info.language_probability:.2f}")
|
||||
|
||||
# Format OpenAI standard
|
||||
if response_format == "text":
|
||||
return full_text
|
||||
elif response_format == "srt":
|
||||
# TODO: Implémenter format SRT
|
||||
return full_text
|
||||
elif response_format == "vtt":
|
||||
# TODO: Implémenter format VTT
|
||||
return full_text
|
||||
else:
|
||||
# Format JSON (défaut)
|
||||
return JSONResponse(content={"text": full_text})
|
||||
|
||||
finally:
|
||||
# Nettoyage
|
||||
os.unlink(tmp_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Erreur transcription: {str(e)}")
|
||||
logger.error(f"❌ Erreur: {str(e)}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post("/unload")
|
||||
async def unload():
|
||||
unload_model()
|
||||
return {"status": "unloaded"}
|
||||
|
||||
@app.post("/reload")
|
||||
async def reload():
|
||||
load_model()
|
||||
reset_unload_timer()
|
||||
return {"status": "loaded"}
|
||||
|
||||
@app.get("/v1/models")
|
||||
async def list_models():
|
||||
"""Liste les modèles disponibles (compatibilité OpenAI)"""
|
||||
return {
|
||||
"data": [
|
||||
{
|
||||
"id": "whisper-1",
|
||||
"object": "model",
|
||||
"created": 1677610602,
|
||||
"owned_by": "openai",
|
||||
"permission": [],
|
||||
"root": "whisper-1",
|
||||
"parent": None
|
||||
}
|
||||
],
|
||||
"data": [{"id": "whisper-1", "object": "model", "created": 1677610602, "owned_by": "openai", "permission": [], "root": "whisper-1", "parent": None}],
|
||||
"object": "list"
|
||||
}
|
||||
|
||||
@app.get("/stats")
|
||||
async def stats():
|
||||
return {
|
||||
"model_loaded": model_loaded,
|
||||
"last_request_age_seconds": int(time.time() - last_request_time),
|
||||
"inactivity_timeout": INACTIVITY_TIMEOUT
|
||||
}
|
||||
|
||||
if __name__ == "__main__":
|
||||
port = int(os.getenv("PORT", "8200"))
|
||||
host = os.getenv("HOST", "0.0.0.0")
|
||||
|
||||
logger.info(f"🚀 Démarrage du serveur sur http://{host}:{port}")
|
||||
logger.info(f"🚀 Serveur sur http://{HOST}:{PORT}")
|
||||
logger.info(f" Endpoints:")
|
||||
logger.info(f" - GET /health")
|
||||
logger.info(f" - POST /v1/audio/transcriptions")
|
||||
logger.info(f" - GET /v1/models")
|
||||
|
||||
uvicorn.run(app, host=host, port=port, log_level="info")
|
||||
logger.info(f" - POST /unload (déchargement manuel)")
|
||||
logger.info(f" - POST /reload (rechargement manuel)")
|
||||
logger.info(f" - GET /stats")
|
||||
uvicorn.run(app, host=HOST, port=PORT, log_level="info")
|
||||
|
||||
Executable
+182
@@ -0,0 +1,182 @@
|
||||
#!/usr/bin/env python3
|
||||
# ComfyUI Model Unloader - Décharge les modèles après inactivité
|
||||
# Usage: python3 comfyui-unload-monitor.py [--interval 60] [--timeout 300]
|
||||
|
||||
import requests
|
||||
import time
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import signal
|
||||
import subprocess
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_INTERVAL = 60
|
||||
DEFAULT_TIMEOUT = 300
|
||||
COMFYUI_URL = "http://localhost:8188"
|
||||
|
||||
class ComfyUIUnloadMonitor:
|
||||
def __init__(self, interval=DEFAULT_INTERVAL, timeout=DEFAULT_TIMEOUT):
|
||||
self.interval = interval
|
||||
self.timeout = timeout
|
||||
self.last_activity = datetime.now()
|
||||
self.running = True
|
||||
self.models_loaded = False
|
||||
|
||||
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...")
|
||||
self.running = False
|
||||
|
||||
def check_comfyui_activity(self):
|
||||
try:
|
||||
response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
|
||||
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
|
||||
except Exception as e:
|
||||
logger.debug(f"Erreur check_queue: {e}")
|
||||
return False
|
||||
|
||||
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 unload_all_models(self):
|
||||
logger.info("🔄 Déchargement de tous les modèles ComfyUI...")
|
||||
|
||||
workflow = {
|
||||
"3": {
|
||||
"inputs": {"value": "unload_trigger"},
|
||||
"class_type": "UnloadAllModelsNode"
|
||||
}
|
||||
}
|
||||
|
||||
try:
|
||||
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
|
||||
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}%)")
|
||||
|
||||
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}")
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
vram, ram, vram_percent = self.get_memory_usage()
|
||||
|
||||
if vram is not None and vram_percent > 1:
|
||||
self.models_loaded = True
|
||||
idle_time = (datetime.now() - self.last_activity).total_seconds()
|
||||
|
||||
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")
|
||||
|
||||
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()
|
||||
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}")
|
||||
time.sleep(self.interval)
|
||||
|
||||
logger.info("👋 Monitor arrêté")
|
||||
|
||||
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")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.debug:
|
||||
logging.getLogger().setLevel(logging.DEBUG)
|
||||
|
||||
monitor = ComfyUIUnloadMonitor(interval=args.interval, timeout=args.timeout)
|
||||
monitor.run()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+174
@@ -0,0 +1,174 @@
|
||||
#!/usr/bin/env python3
|
||||
# Kreuzberg Memory Manager - Redémarre le container après inactivité
|
||||
# Usage: python3 kreuzberg-memory-manager.py [--interval 60] [--timeout 600] [--max-memory 2]
|
||||
|
||||
import requests
|
||||
import time
|
||||
import argparse
|
||||
import logging
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_INTERVAL = 60
|
||||
DEFAULT_TIMEOUT = 600
|
||||
DEFAULT_MAX_MEMORY = 2.0
|
||||
KREUZBERG_API = "http://localhost:8300"
|
||||
CONTAINER_NAME = "kreuzberg"
|
||||
|
||||
class KreuzbergMemoryManager:
|
||||
def __init__(self, interval=DEFAULT_INTERVAL, timeout=DEFAULT_TIMEOUT, max_memory_gb=DEFAULT_MAX_MEMORY):
|
||||
self.interval = interval
|
||||
self.timeout = timeout
|
||||
self.max_memory_gb = max_memory_gb
|
||||
self.last_activity = time.time()
|
||||
self.running = True
|
||||
self.container_id = None
|
||||
|
||||
import signal
|
||||
signal.signal(signal.SIGINT, self.signal_handler)
|
||||
signal.signal(signal.SIGTERM, self.signal_handler)
|
||||
|
||||
self.find_container()
|
||||
|
||||
def signal_handler(self, signum, frame):
|
||||
logger.info("Signal reçu, arrêt du manager...")
|
||||
self.running = False
|
||||
|
||||
def find_container(self):
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "ps", "-q", "--filter", f"name={CONTAINER_NAME}"],
|
||||
capture_output=True, text=True
|
||||
)
|
||||
if result.stdout.strip():
|
||||
self.container_id = result.stdout.strip()
|
||||
logger.info(f"Container trouvé: {self.container_id[:12]}")
|
||||
else:
|
||||
logger.warning(f"Container {CONTAINER_NAME} non trouvé")
|
||||
except Exception as e:
|
||||
logger.error(f"Erreur recherche container: {e}")
|
||||
|
||||
def get_container_memory(self):
|
||||
if not self.container_id:
|
||||
return None
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["docker", "stats", "--no-stream", "--format", "{{.MemUsage}}", self.container_id],
|
||||
capture_output=True, text=True
|
||||
)
|
||||
if result.stdout.strip():
|
||||
mem_str = result.stdout.strip()
|
||||
return self.parse_memory(mem_str)
|
||||
except Exception as e:
|
||||
logger.debug(f"Erreur stats: {e}")
|
||||
return None
|
||||
|
||||
def parse_memory(self, mem_str):
|
||||
try:
|
||||
parts = mem_str.split("/")
|
||||
if len(parts) >= 1:
|
||||
used = parts[0].strip()
|
||||
if "GiB" in used:
|
||||
return float(used.replace("GiB", "")) * 1024
|
||||
elif "MiB" in used:
|
||||
return float(used.replace("MiB", ""))
|
||||
elif "GB" in used:
|
||||
return float(used.replace("GB", "")) * 1024
|
||||
elif "MB" in used:
|
||||
return float(used.replace("MB", ""))
|
||||
except:
|
||||
pass
|
||||
return None
|
||||
|
||||
def check_api_activity(self):
|
||||
try:
|
||||
response = requests.get(f"{KREUZBERG_API}/health", timeout=5)
|
||||
if response.status_code in [200, 404]:
|
||||
self.last_activity = time.time()
|
||||
return True
|
||||
except:
|
||||
pass
|
||||
return False
|
||||
|
||||
def restart_container(self):
|
||||
logger.info(f"🔄 Redémarrage du container {CONTAINER_NAME}...")
|
||||
try:
|
||||
subprocess.run(["docker", "restart", CONTAINER_NAME], check=True)
|
||||
logger.info("✅ Container redémarré")
|
||||
time.sleep(5)
|
||||
self.find_container()
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Erreur redémarrage: {e}")
|
||||
return False
|
||||
|
||||
def run(self):
|
||||
logger.info(f"🚀 Démarrage Kreuzberg Memory Manager")
|
||||
logger.info(f" Intervalle: {self.interval}s")
|
||||
logger.info(f" Timeout inactivité: {self.timeout}s")
|
||||
logger.info(f" Seuil mémoire: {self.max_memory_gb}GB")
|
||||
|
||||
idle_time = 0
|
||||
|
||||
while self.running:
|
||||
try:
|
||||
memory_mb = self.get_container_memory()
|
||||
|
||||
if memory_mb is not None:
|
||||
memory_gb = memory_mb / 1024
|
||||
|
||||
is_active = self.check_api_activity()
|
||||
|
||||
if is_active:
|
||||
idle_time = 0
|
||||
self.last_activity = time.time()
|
||||
logger.debug(f"📊 Mémoire: {memory_gb:.2f}GB, API active")
|
||||
else:
|
||||
idle_time += self.interval
|
||||
logger.debug(f"📊 Mémoire: {memory_gb:.2f}GB, Inactif: {idle_time}s")
|
||||
|
||||
if memory_gb > self.max_memory_gb:
|
||||
logger.warning(f"⚠️ Mémoire ({memory_gb:.2f}GB) > seuil ({self.max_memory_gb}GB)")
|
||||
if idle_time >= self.timeout:
|
||||
logger.info(f"⏰ Seuil timeout atteint ({idle_time}s >= {self.timeout}s)")
|
||||
self.restart_container()
|
||||
idle_time = 0
|
||||
|
||||
time.sleep(self.interval)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Erreur: {e}")
|
||||
time.sleep(self.interval)
|
||||
|
||||
logger.info("👋 Manager arrêté")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Kreuzberg Memory Manager")
|
||||
parser.add_argument("--interval", type=int, default=DEFAULT_INTERVAL,
|
||||
help=f"Intervalle (défaut: {DEFAULT_INTERVAL}s)")
|
||||
parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT,
|
||||
help=f"Timeout inactivité (défaut: {DEFAULT_TIMEOUT}s)")
|
||||
parser.add_argument("--max-memory", type=float, default=DEFAULT_MAX_MEMORY,
|
||||
help=f"Seuil mémoire GB (défaut: {DEFAULT_MAX_MEMORY}GB)")
|
||||
parser.add_argument("--debug", action="store_true", help="Mode debug")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.debug:
|
||||
logging.getLogger().setLevel(logging.DEBUG)
|
||||
|
||||
manager = KreuzbergMemoryManager(
|
||||
interval=args.interval,
|
||||
timeout=args.timeout,
|
||||
max_memory_gb=args.max_memory
|
||||
)
|
||||
manager.run()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user