feat: add VRAM management for Jean/Ollama switching

- stop_jean(): stops Jean service to free VRAM before Ollama use - start_jean(): starts Jean service when switching back from Ollama - unload_ollama(): forces Ollama to unload models before Jean use - ensure_vram_for_backend(): orchestrates VRAM switching - Tracks active backend state to avoid unnecessary switches - Configurable timeouts via environment variables - Prevents VRAM contention between Jean and Ollama
2026-07-20 10:33:35 +02:00
parent 856ba4ae36
commit 8b7c4fbd9c
2 changed files with 157 additions and 0 deletions
@@ -63,7 +63,12 @@ services:
      - JEAN_PROXY_PORT=3012
      - OLLAMA_PROXY_HOST=localhost
      - OLLAMA_PROXY_PORT=3011
+      - OLLAMA_DIRECT_HOST=${OLLAMA_HOST}
+      - OLLAMA_DIRECT_PORT=11434
      - ENABLE_OLLAMA=true
+      - ENABLE_VRAM_MANAGEMENT=true
      - LISTEN_PORT=3014
      - SWITCH_TIMEOUT=90
      - CHECK_INTERVAL=2
+      - JEAN_STOP_TIMEOUT=30
+      - OLLAMA_UNLOAD_TIMEOUT=15
@@ -20,10 +20,15 @@ JEAN_PROXY_HOST = os.getenv("JEAN_PROXY_HOST", "localhost")
 JEAN_PROXY_PORT = int(os.getenv("JEAN_PROXY_PORT", "3012"))
 OLLAMA_PROXY_HOST = os.getenv("OLLAMA_PROXY_HOST", "localhost")
 OLLAMA_PROXY_PORT = int(os.getenv("OLLAMA_PROXY_PORT", "3011"))
+OLLAMA_DIRECT_HOST = os.getenv("OLLAMA_DIRECT_HOST", JEAN_WEB_HOST)
+OLLAMA_DIRECT_PORT = int(os.getenv("OLLAMA_DIRECT_PORT", "11434"))
 ENABLE_OLLAMA = os.getenv("ENABLE_OLLAMA", "true").lower() == "true"
+ENABLE_VRAM_MANAGEMENT = os.getenv("ENABLE_VRAM_MANAGEMENT", "true").lower() == "true"
 LISTEN_PORT = int(os.getenv("LISTEN_PORT", "3014"))
 SWITCH_TIMEOUT = int(os.getenv("SWITCH_TIMEOUT", "90"))
 CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "2"))
+JEAN_STOP_TIMEOUT = int(os.getenv("JEAN_STOP_TIMEOUT", "30"))
+OLLAMA_UNLOAD_TIMEOUT = int(os.getenv("OLLAMA_UNLOAD_TIMEOUT", "15"))

 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 log = logging.getLogger(__name__)
@@ -34,6 +39,8 @@ presets_cache = {"data": [], "timestamp": 0}
 ollama_cache = {"data": [], "timestamp": 0}
 CACHE_TTL = 30
 switch_lock = threading.Lock()
+vram_lock = threading.Lock()
+vram_state = {"active_backend": "jean"}


 def get_presets(force_refresh=False):
@@ -134,6 +141,145 @@ def switch_preset(preset_id: str) -> bool:
    return False


+def stop_jean() -> bool:
+    """Arrête le service Jean pour libérer la VRAM"""
+    log.info("[VRAM] Arrêt de Jean pour libérer la VRAM...")
+
+    try:
+        url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/stop"
+        r = requests.post(url, timeout=5)
+        if r.ok:
+            log.info("[VRAM] Jean arrêté via API web")
+        else:
+            log.warning(f"[VRAM] API stop a retourné {r.status_code}, tentative via health check")
+    except Exception as e:
+        log.warning(f"[VRAM] Erreur API stop: {e}")
+
+    deadline = time.time() + JEAN_STOP_TIMEOUT
+    while time.time() < deadline:
+        try:
+            health_url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/health"
+            r = requests.get(health_url, timeout=2)
+            if not r.ok:
+                log.info("[VRAM] Jean confirmé arrêté")
+                return True
+        except Exception:
+            log.info("[VRAM] Jean confirmé arrêté (connexion refusée)")
+            return True
+        time.sleep(CHECK_INTERVAL)
+
+    log.error(f"[VRAM] Timeout: Jean ne s'est pas arrêté après {JEAN_STOP_TIMEOUT}s")
+    return False
+
+
+def start_jean() -> bool:
+    """Démarre le service Jean"""
+    log.info("[VRAM] Démarrage de Jean...")
+
+    try:
+        url = f"http://{JEAN_WEB_HOST}:{JEAN_WEB_PORT}/api/start"
+        r = requests.post(url, timeout=5)
+        if r.ok:
+            log.info("[VRAM] Jean démarré via API web")
+        else:
+            log.warning(f"[VRAM] API start a retourné {r.status_code}")
+    except Exception as e:
+        log.warning(f"[VRAM] Erreur API start: {e}")
+
+    log.info("[VRAM] Attente du chargement du modèle Jean...")
+    deadline = time.time() + SWITCH_TIMEOUT
+    while time.time() < deadline:
+        try:
+            health_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/health"
+            hr = requests.get(health_url, timeout=2)
+            if hr.ok:
+                models_url = f"http://{JEAN_PROXY_HOST}:{JEAN_PROXY_PORT}/v1/models"
+                mr = requests.get(models_url, timeout=2)
+                if mr.ok:
+                    models = mr.json().get("data", [])
+                    if models:
+                        log.info(f"[VRAM] Jean prêt après {int(time.time() - deadline + SWITCH_TIMEOUT)}s")
+                        return True
+        except Exception:
+            pass
+        time.sleep(CHECK_INTERVAL)
+
+    log.error(f"[VRAM] Timeout: Jean n'a pas démarré après {SWITCH_TIMEOUT}s")
+    return False
+
+
+def unload_ollama() -> bool:
+    """Force le déchargement des modèles Ollama pour libérer la VRAM"""
+    log.info("[VRAM] Déchargement des modèles Ollama...")
+
+    try:
+        url = f"http://{OLLAMA_DIRECT_HOST}:{OLLAMA_DIRECT_PORT}/api/generate"
+        payload = {
+            "model": "mistral",
+            "keep_alive": 0,
+            "prompt": ""
+        }
+        r = requests.post(url, json=payload, timeout=5)
+        if r.ok:
+            log.info("[VRAM] Ollama déchargé via API")
+            return True
+        else:
+            log.warning(f"[VRAM] API unload a retourné {r.status_code}")
+    except Exception as e:
+        log.warning(f"[VRAM] Erreur API unload: {e}")
+
+    try:
+        url = f"http://{OLLAMA_DIRECT_HOST}:{OLLAMA_DIRECT_PORT}/api/ps"
+        r = requests.get(url, timeout=3)
+        if r.ok:
+            models = r.json().get("models", [])
+            if not models:
+                log.info("[VRAM] Ollama confirmé déchargé (aucun modèle actif)")
+                return True
+            else:
+                log.warning(f"[VRAM] Ollama a encore {len(models)} modèle(s) chargé(s)")
+    except Exception as e:
+        log.warning(f"[VRAM] Impossible de vérifier /api/ps: {e}")
+
+    time.sleep(2)
+    log.info("[VRAM] Ollama supposé déchargé (timeout)")
+    return True
+
+
+def ensure_vram_for_backend(target_backend: str) -> bool:
+    """
+    S'assure que la VRAM est disponible pour le backend cible.
+    - Si target=ollama et jean actif → arrête Jean
+    - Si target=jean et ollama actif → décharge Ollama
+    """
+    if not ENABLE_VRAM_MANAGEMENT:
+        return True
+
+    with vram_lock:
+        current = vram_state["active_backend"]
+
+        if current == target_backend:
+            return True
+
+        if target_backend == "ollama" and current == "jean":
+            log.info("[VRAM] Switch Jean → Ollama")
+            if not stop_jean():
+                return False
+            vram_state["active_backend"] = "ollama"
+            return True
+
+        elif target_backend == "jean" and current == "ollama":
+            log.info("[VRAM] Switch Ollama → Jean")
+            if not unload_ollama():
+                return False
+            if not start_jean():
+                return False
+            vram_state["active_backend"] = "jean"
+            return True
+
+        return True
+
+
@app.route("/v1/models", methods=["GET"])
 def list_models():
    """Expose tous les presets Jean + modèles Ollama comme modèles OpenAI"""
@@ -218,6 +364,9 @@ def chat_completions():
    jean_preset_ids = [p["id"] for p in presets]

    if model_id in jean_preset_ids:
+        if not ensure_vram_for_backend("jean"):
+            return jsonify({"error": "Failed to prepare VRAM for Jean"}), 503
+
        with switch_lock:
            active = get_active_preset()
            if not active or active["id"] != model_id:
@@ -232,6 +381,9 @@ def chat_completions():
        ollama_ids = [m.get("name", "").replace(":latest", "") for m in ollama_models]

        if model_id in ollama_ids:
+            if not ensure_vram_for_backend("ollama"):
+                return jsonify({"error": "Failed to prepare VRAM for Ollama"}), 503
+
            upstream_url = f"http://{OLLAMA_PROXY_HOST}:{OLLAMA_PROXY_PORT}/v1/chat/completions"
            log.info(f"[Router] → Ollama ({model_id})")
        else: