refactored load ml service

app/main.py

@@ -24,7 +24,7 @@ from app.deps import get_current_user, rate_limit_org, verify_api_key, check_all
 from app.tasks.analytics_worker import redis_listener, trigger_kpi_computation
 from app.service.vector_service import cleanup_expired_vectors
 from app.routers import health, datasources, reports, flags, scheduler, run, socket, analytics_stream,ai_query,schema
-
+from app.service.llm_service import load_llm_service
 # ─── Logger Configuration ───────────────────────────────────────────────────────
 logging.basicConfig(
     level=logging.INFO,
@@ -107,9 +107,12 @@ async def lifespan(app: FastAPI):
     logger.info("⏰ Starting KPI refresh scheduler...")
     asyncio.create_task(continuous_kpi_refresh(), name="kpi_scheduler")
     # Now load LLM service - it will use persistent cache
-    from app.service.llm_service import LocalLLMService
-    logger.info("🤖 LLM service initialized (will use persistent cache)")
-    
+    try:
+        load_llm_service()  # Starts background loading
+        logger.info("🤖 LLM service loading in background...")
+    except Exception as e:
+        logger.error(f"❌ LLM load failed: {e}")
+        # Continue anyway - LLM is optional for some features
     yield
     
     # ─── Shutdown ──────────────────────────────────────────────────────────────

app/service/llm_service.py

@@ -6,6 +6,7 @@ import logging
 from threading import Thread, Lock
 import time
 import json
+import os 
 
 logger = logging.getLogger(__name__)
 
@@ -21,33 +22,36 @@ class LocalLLMService:
         self._lock = Lock()
         
         # ✅ Use persistent cache
-        cache_dir = "/data/hf_cache"
-        os.makedirs(cache_dir, exist_ok=True)
+        self.cache_dir = "/data/hf_cache"
+        os.makedirs(self.cache_dir, exist_ok=True)
         
-        logger.info("🚀 Starting background LLM load...")
-        Thread(target=self._load_model_background, daemon=True).start()
+        # ❌ DON'T start loading here - truly lazy
+        self._load_thread = None
     
-    def _load_model_background(self):
-        """Load model in background thread with persistent cache"""
+    def load(self):
+        """Explicitly start loading the model - call this ONLY after build is verified"""
         with self._lock:
             if self._is_loading or self._is_loaded:
+                logger.info("Model already loading or loaded")
                 return
+            
             self._is_loading = True
-        
+            logger.info("🚀 Starting LLM load...")
+            self._load_thread = Thread(target=self._load_model_background, daemon=True)
+            self._load_thread.start()
+    
+    def _load_model_background(self):
+        """Load model in background thread with persistent cache"""
         try:
             logger.info(f"🤖 [BACKGROUND] Loading LLM: {self.model_id}...")
             
-            # ✅ Use persistent cache directory
-            cache_dir = "/data/hf_cache"
-            
             # Phi-3 tokenizer
             self._tokenizer = AutoTokenizer.from_pretrained(
                 self.model_id,
                 token=HF_API_TOKEN,
                 trust_remote_code=True,
-                cache_dir=cache_dir  # ✅ Persistent cache
+                cache_dir=self.cache_dir
             )
-            # .
             self._tokenizer.pad_token = self._tokenizer.eos_token
             
             # Phi-3 model - OPTIMIZED for speed
@@ -58,7 +62,8 @@ class LocalLLMService:
                 device_map="auto",
                 low_cpu_mem_usage=True,
                 trust_remote_code=True,
-                attn_implementation="eager"  # ✅ No flash-attn warnings
+                attn_implementation="eager",
+                cache_dir=self.cache_dir  # ✅ Persistent cache
             )
             
             # ✅ FASTER pipeline settings
@@ -69,8 +74,8 @@ class LocalLLMService:
                 device_map="auto",
                 torch_dtype=torch.float16,
                 trust_remote_code=True,
-                # ✅ SPEED UP: Use pad_token_id
-                pad_token_id=self._tokenizer.eos_token_id
+                pad_token_id=self._tokenizer.eos_token_id,
+                cache_dir=self.cache_dir
             )
             
             with self._lock:
@@ -105,9 +110,7 @@ class LocalLLMService:
             raise TimeoutError("LLM loading in progress")
         
         # ✅ Phi-3 prompt format (TESTED to work)
-        messages = [
-            {"role": "user", "content": prompt}
-        ]
+        messages = [{"role": "user", "content": prompt}]
         
         formatted_prompt = self._tokenizer.apply_chat_template(
             messages,
@@ -122,7 +125,7 @@ class LocalLLMService:
             temperature=temperature,
             do_sample=False,
             pad_token_id=self._tokenizer.eos_token_id,
-            return_full_text=False  # ✅ Only return generated text
+            return_full_text=False
         )
         
         # ✅ SAFE extraction
@@ -136,12 +139,35 @@ class LocalLLMService:
         
         # ✅ VALIDATE JSON before returning
         try:
-            json.loads(response_text)  # Test parse
+            json.loads(response_text)
             logger.info(f"[llm] Valid JSON generated: {response_text[:50]}...")
             return response_text
         except json.JSONDecodeError:
             logger.error(f"[llm] Invalid JSON from LLM: {response_text}")
             raise ValueError(f"LLM returned invalid JSON: {response_text}")
 
-# Singleton
-llm_service = LocalLLMService()
+
+# ✅ LAZY singleton creation - instance created ONLY when first requested
+_llm_service_instance = None
+
+def get_llm_service():
+    """Get or create the singleton LLM service (lazy initialization)"""
+    global _llm_service_instance
+    
+    if _llm_service_instance is None:
+        logger.info("🆕 Creating LLM service instance (lazy)")
+        _llm_service_instance = LocalLLMService()
+    
+    return _llm_service_instance
+
+
+def load_llm_service():
+    """
+    Explicitly load the LLM service.
+    Call this AFTER startup sequence to ensure build is successful.
+    """
+    service = get_llm_service()
+    if not service.is_loaded and not service.is_loading:
+        service.load()
+        logger.info("🤖 LLM service loading triggered")
+    return service