Update brain_lazy.py

brain_lazy.py

@@ -1,57 +1,208 @@
 """
-brain_lazy.py
-Lazy loader for multimodular_modul_v7
-- Avoids Hugging Face runtime startup timeout
-- Loads brain in background while API/CLI starts instantly
+brain_lazy.py  (rewritten)
+Purpose:
+- Start your app FAST (no startup timeout on Hugging Face).
+- Immediately after startup, auto-load the full multimodular brain in the background
+  (models, weights, heavy imports) so users don’t hit first-use lag.
+- Keep your original multimodular_modul_v7.py completely untouched.
+
+Notes:
+- This does NOT change pip install time for big wheels (e.g., torch/timm). That happens
+  during Space build from requirements.txt. This file prevents slow *runtime* model
+  initialization by deferring it to a background preload right after boot.
 """
 
 import os
-import importlib
-import threading
 import time
+import threading
+import importlib
+from typing import Optional, Any
+
+# -----------------------------------------------------------------------------
+# Persistent cache (Hugging Face → enable Persistent storage in Space settings)
+# -----------------------------------------------------------------------------
+CACHE_DIR = "/home/user/app/cache"
+os.environ.setdefault("TRANSFORMERS_CACHE", CACHE_DIR)
+os.environ.setdefault("HF_HOME", CACHE_DIR)
+os.makedirs(CACHE_DIR, exist_ok=True)
 
-# ------------------------------
-# Persistent storage cache setup
-# ------------------------------
-os.environ["TRANSFORMERS_CACHE"] = "/home/user/app/cache"
-os.environ["HF_HOME"] = "/home/user/app/cache"
-os.makedirs("/home/user/app/cache", exist_ok=True)
-
-# ------------------------------
-# Brain loader
-# ------------------------------
-_brain = None
+# -----------------------------------------------------------------------------
+# Loader flags
+# -----------------------------------------------------------------------------
+_brain = None                 # the real multimodular module (multimodular_modul_v7)
+_lock = threading.Lock()
 _is_loading = False
 _is_ready = False
-_lock = threading.Lock()
+_last_error: Optional[str] = None
 
-def _load_brain():
-    global _brain, _is_ready, _is_loading
+# How long proxies will wait (max) for background preload before returning a
+# graceful “warming up” message. Tweak for your UX.
+PROXY_WAIT_SECONDS = 25
+
+# -----------------------------------------------------------------------------
+# Background preload
+# -----------------------------------------------------------------------------
+def _load_brain_blocking() -> Any:
+    """
+    Imports the heavy brain module and performs a light warm-up.
+    This runs either in the background (at startup) or on-demand if a call arrives early.
+    """
+    global _brain, _is_ready, _is_loading, _last_error
     with _lock:
-        if _brain is None:
-            _is_loading = True
-            print("⏳ Loading multimodular brain (lazy mode)...")
-            start_time = time.time()
-            _brain = importlib.import_module("multimodular_modul_v7")
+        if _brain is not None:
+            return _brain
+        if _is_loading:
+            # Another thread is loading; just return and let caller wait/poll.
+            return None
+
+        _is_loading = True
+        _last_error = None
+        start = time.time()
+        try:
+            print("⏳ [brain_lazy] Importing multimodular_modul_v7 ...")
+            brain = importlib.import_module("multimodular_modul_v7")
+
+            # Optional: if your module exposes an init() or warm_up(), call it.
+            # Otherwise, do a tiny no-op inference to trigger weights load.
+            warm_started = False
+            if hasattr(brain, "init"):
+                try:
+                    brain.init()
+                    warm_started = True
+                    print("✅ [brain_lazy] brain.init() finished.")
+                except Exception as e:
+                    print(f"⚠️  [brain_lazy] brain.init() failed: {e}")
+
+            if hasattr(brain, "warm_up"):
+                try:
+                    brain.warm_up()
+                    warm_started = True
+                    print("✅ [brain_lazy] brain.warm_up() finished.")
+                except Exception as e:
+                    print(f"⚠️  [brain_lazy] brain.warm_up() failed: {e}")
+
+            # Minimal warm-up if none provided
+            if not warm_started and hasattr(brain, "process_input"):
+                try:
+                    _ = brain.process_input("ping")
+                    print("✅ [brain_lazy] minimal warm-up via process_input('ping') done.")
+                except Exception as e:
+                    print(f"⚠️  [brain_lazy] minimal warm-up failed: {e}")
+
+            _brain = brain
             _is_ready = True
+            print(f"✅ [brain_lazy] Brain loaded in {time.time() - start:.2f}s")
+            return _brain
+        except Exception as e:
+            _last_error = str(e)
+            print(f"❌ [brain_lazy] Brain load failed: {e}")
+            return None
+        finally:
             _is_loading = False
-            print(f"✅ Brain loaded in {time.time() - start_time:.2f} seconds.")
-    return _brain
-
-def preload_in_background():
-    threading.Thread(target=_load_brain, daemon=True).start()
-
-# Start background preload at import
-preload_in_background()
-
-# ------------------------------
-# Proxy functions
-# ------------------------------
-def process_input(text): return _load_brain().process_input(text)
-def search_kb(query): return _load_brain().search_kb(query)
-def upload_media(path): return _load_brain().upload_media(path)
-def backup_brain(): return _load_brain().backup_brain()
-def restore_brain(): return _load_brain().restore_brain()
-def show_creative_skills(): return _load_brain().show_creative_skills()
-def sync_status(): return _load_brain().sync_status()
-def is_ready(): return _is_ready
+
+def _preload_thread():
+    """
+    Kicks off immediately at import so HF sees a fast boot,
+    then models load in background right away.
+    """
+    _load_brain_blocking()
+
+# Start background preload now (non-blocking)
+threading.Thread(target=_preload_thread, daemon=True).start()
+
+# -----------------------------------------------------------------------------
+# Helpers
+# -----------------------------------------------------------------------------
+def is_ready() -> bool:
+    return _is_ready
+
+def last_error() -> Optional[str]:
+    return _last_error
+
+def _ensure_loaded_with_wait(timeout_s: float) -> Optional[Any]:
+    """
+    Ensure the brain is loaded. Wait up to `timeout_s` seconds for background preload.
+    If still not ready after timeout, return None (caller can respond with a
+    graceful "warming up" message).
+    """
+    global _brain
+    # Fast path
+    if _brain is not None and _is_ready:
+        return _brain
+
+    # Trigger on-demand load if background hasn’t started (very rare)
+    if not _is_loading and _brain is None:
+        # Start a parallel load but do not block the whole timeout here;
+        # we will poll below.
+        threading.Thread(target=_load_brain_blocking, daemon=True).start()
+
+    waited = 0.0
+    interval = 0.25
+    while waited < timeout_s:
+        if _brain is not None and _is_ready:
+            return _brain
+        time.sleep(interval)
+        waited += interval
+
+    return None
+
+def _warming_up_message(op: str) -> Any:
+    """
+    Graceful response while models finish loading.
+    You can customize this to fit your API schema/UI.
+    """
+    msg = {
+        "status": "warming_up",
+        "operation": op,
+        "detail": "CHB is loading models in the background. Please retry in a few seconds.",
+        "ready": _is_ready,
+        "error": _last_error,
+    }
+    return msg
+
+# -----------------------------------------------------------------------------
+# Public proxy API (mirrors your multimodular_modul_v7 public surface)
+# Each call tries to use the loaded brain; if not ready within PROXY_WAIT_SECONDS,
+# it returns a non-blocking 'warming_up' payload instead of hanging requests.
+# -----------------------------------------------------------------------------
+def process_input(text: str) -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("process_input")
+    return brain.process_input(text)
+
+def search_kb(query: str) -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("search_kb")
+    return brain.search_kb(query)
+
+def upload_media(file_path: str) -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("upload_media")
+    return brain.upload_media(file_path)
+
+def backup_brain() -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("backup_brain")
+    return brain.backup_brain()
+
+def restore_brain() -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("restore_brain")
+    return brain.restore_brain()
+
+def show_creative_skills() -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("show_creative_skills")
+    return brain.show_creative_skills()
+
+def sync_status() -> Any:
+    brain = _ensure_loaded_with_wait(PROXY_WAIT_SECONDS)
+    if brain is None:
+        return _warming_up_message("sync_status")
+    return brain.sync_status()