From 20767c2300b5731420e8deae0dffcb5eb98111fc Mon Sep 17 00:00:00 2001
From: Folzo <folzo@hotmail.com>
Date: Mon, 20 Jul 2026 10:46:58 +0200
Subject: [PATCH] feat: keep Ollama models in memory for 5 minutes

- Add OLLAMA_KEEP_ALIVE_MINUTES config (default: 5)
- schedule_ollama_unload(): delays Ollama unload when switching to Jean
- cancel_ollama_unload(): cancels timer if returning to Ollama
- Ollama models stay in VRAM for X minutes after last use
- Timer resets on each Ollama request
---
 jean-router/app.py | 42 ++++++++++++++++++++++++++++++++++++++----
 1 file changed, 38 insertions(+), 4 deletions(-)

diff --git a/jean-router/app.py b/jean-router/app.py
index 09eb430..4e522c8 100644
--- a/jean-router/app.py
+++ b/jean-router/app.py
@@ -24,6 +24,7 @@ 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"
+OLLAMA_KEEP_ALIVE_MINUTES = int(os.getenv("OLLAMA_KEEP_ALIVE_MINUTES", "5"))
 LISTEN_PORT = int(os.getenv("LISTEN_PORT", "3014"))
 SWITCH_TIMEOUT = int(os.getenv("SWITCH_TIMEOUT", "90"))
 CHECK_INTERVAL = int(os.getenv("CHECK_INTERVAL", "2"))
@@ -41,6 +42,7 @@ CACHE_TTL = 30
 switch_lock = threading.Lock()
 vram_lock = threading.Lock()
 vram_state = {"active_backend": "jean"}
+ollama_unload_timer = None
 
 
 def get_presets(force_refresh=False):
@@ -246,11 +248,41 @@ def unload_ollama() -> bool:
     return True
 
 
+def schedule_ollama_unload():
+    """Programme le déchargement d'Ollama après OLLAMA_KEEP_ALIVE_MINUTES"""
+    global ollama_unload_timer
+
+    def delayed_unload():
+        log.info(f"[VRAM] Timer écoulé ({OLLAMA_KEEP_ALIVE_MINUTES}min), déchargement Ollama...")
+        unload_ollama()
+
+    if ollama_unload_timer:
+        ollama_unload_timer.cancel()
+        log.info("[VRAM] Timer précédent annulé")
+
+    ollama_unload_timer = threading.Timer(
+        OLLAMA_KEEP_ALIVE_MINUTES * 60,
+        delayed_unload
+    )
+    ollama_unload_timer.daemon = True
+    ollama_unload_timer.start()
+    log.info(f"[VRAM] Déchargement Ollama programmé dans {OLLAMA_KEEP_ALIVE_MINUTES} minutes")
+
+
+def cancel_ollama_unload():
+    """Annule le déchargement programmé d'Ollama"""
+    global ollama_unload_timer
+    if ollama_unload_timer:
+        ollama_unload_timer.cancel()
+        ollama_unload_timer = None
+        log.info("[VRAM] Timer de déchargement Ollama annulé")
+
+
 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
+    - Si target=ollama et jean actif → arrête Jean, annule timer unload Ollama
+    - Si target=jean et ollama actif → programme déchargement Ollama dans X min, démarre Jean
     """
     if not ENABLE_VRAM_MANAGEMENT:
         return True
@@ -263,6 +295,7 @@ def ensure_vram_for_backend(target_backend: str) -> bool:
 
         if target_backend == "ollama" and current == "jean":
             log.info("[VRAM] Switch Jean → Ollama")
+            cancel_ollama_unload()
             if not stop_jean():
                 return False
             vram_state["active_backend"] = "ollama"
@@ -270,8 +303,7 @@ def ensure_vram_for_backend(target_backend: str) -> bool:
 
         elif target_backend == "jean" and current == "ollama":
             log.info("[VRAM] Switch Ollama → Jean")
-            if not unload_ollama():
-                return False
+            schedule_ollama_unload()
             if not start_jean():
                 return False
             vram_state["active_backend"] = "jean"
@@ -384,6 +416,8 @@ def chat_completions():
             if not ensure_vram_for_backend("ollama"):
                 return jsonify({"error": "Failed to prepare VRAM for Ollama"}), 503
 
+            cancel_ollama_unload()
+
             upstream_url = f"http://{OLLAMA_PROXY_HOST}:{OLLAMA_PROXY_PORT}/v1/chat/completions"
             log.info(f"[Router] → Ollama ({model_id})")
         else: