diff --git a/docker-compose.litellm.yml b/docker-compose.litellm.yml index 869c313..e8f54be 100644 --- a/docker-compose.litellm.yml +++ b/docker-compose.litellm.yml @@ -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 diff --git a/jean-router/app.py b/jean-router/app.py index 00acc99..09eb430 100644 --- a/jean-router/app.py +++ b/jean-router/app.py @@ -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: