lazy model loading

app.py

@@ -13,7 +13,6 @@ import pandas as pd
 from src.utils.ml_utils.model.estimator import NetworkSecurityModel
 from contextlib import asynccontextmanager
 import mlflow
-import dagshub
 
 ca = certifi.where()
 load_dotenv()
@@ -30,17 +29,9 @@ from src.utils.main_utils.utils import load_object, save_object
 from fastapi.templating import Jinja2Templates
 templates = Jinja2Templates(directory="./templates")
 
-# Initialize DagHub for MLflow tracking
-try:
-    dagshub.init(repo_owner='kshitijk146', repo_name='MLOPS_project_network_Security_system', mlflow=True)
-    MLFLOW_AVAILABLE = True
-    logging.info("✅ MLflow tracking initialized")
-except Exception as e:
-    logging.warning(f"⚠️ MLflow initialization failed: {e}")
-    MLFLOW_AVAILABLE = False
-
 # Cache for loaded models
 MODEL_CACHE = {"model": None, "preprocessor": None}
+MLFLOW_AVAILABLE = True  # Assume available, model_trainer.py handles initialization
 
 def load_models_from_mlflow():
     """Load latest models from MLflow"""
@@ -152,11 +143,11 @@ async def training_route():
         training_pipeline = Trainingpipeline()
         training_pipeline.run_pipeline()
         
-        # Reload models from MLflow after training
-        if MLFLOW_AVAILABLE:
-            load_models_from_mlflow()
+        # Clear model cache so next prediction loads fresh models
+        MODEL_CACHE["model"] = None
+        MODEL_CACHE["preprocessor"] = None
         
-        return Response("✅ Training completed and models loaded from MLflow!")
+        return Response("✅ Training completed! Models logged to MLflow. Call /predict to use them.")
     except Exception as e:
         raise NetworkSecurityException(e, sys)