adding mlfow registered model loading

app.py

@@ -49,9 +49,17 @@ def load_models_from_mlflow():
             logging.error("MLflow not available")
             return False
         
+        logging.info("Searching for latest MLflow run...")
+        
         # Get the latest run from the experiment
         client = mlflow.tracking.MlflowClient()
-        experiment = client.get_experiment_by_name("Default")
+        
+        # Try to get experiment, if it doesn't exist, no models are trained yet
+        try:
+            experiment = client.get_experiment_by_name("Default")
+        except Exception as e:
+            logging.warning(f"Could not get experiment: {e}")
+            return False
         
         if experiment is None:
             logging.warning("No MLflow experiment found. Train model first.")
@@ -89,20 +97,30 @@ def load_models_from_mlflow():
         
     except Exception as e:
         logging.error(f"Error loading models from MLflow: {e}")
+        import traceback
+        logging.error(traceback.format_exc())
         return False
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Load models on startup"""
-    logging.info("===== Application Startup - Loading models from MLflow =====")
+    logging.info("===== Application Startup - Checking for models =====")
     
     if MLFLOW_AVAILABLE:
-        success = load_models_from_mlflow()
-        if not success:
-            logging.warning("⚠️ Could not load models from MLflow. Please train first via /train endpoint.")
+        try:
+            # Try to load models but don't block startup if it fails
+            logging.info("Attempting to load models from MLflow...")
+            success = load_models_from_mlflow()
+            if success:
+                logging.info("✅ Models loaded successfully from MLflow")
+            else:
+                logging.warning("⚠️ No models found in MLflow. Train via /train endpoint.")
+        except Exception as e:
+            logging.warning(f"⚠️ Could not load from MLflow: {e}. Train via /train endpoint.")
     else:
         logging.warning("⚠️ MLflow not available. Please train via /train endpoint.")
     
+    logging.info("✅ Application ready to serve requests")
     yield
     logging.info("===== Application Shutdown =====")
 

src/components/model_trainer.py

@@ -21,7 +21,16 @@ import dagshub
 import os 
 from dotenv import load_dotenv
 load_dotenv()
-dagshub.init(repo_owner='kshitijk146', repo_name='MLOPS_project_network_Security_system', mlflow=True)
+
+# Only initialize DagHub once, and make it optional
+try:
+    if not os.getenv("DAGSHUB_INITIALIZED"):
+        dagshub.init(repo_owner='kshitijk146', repo_name='MLOPS_project_network_Security_system', mlflow=True)
+        os.environ["DAGSHUB_INITIALIZED"] = "1"
+        logging.info("✅ DagHub/MLflow initialized in model_trainer")
+except Exception as e:
+    logging.warning(f"⚠️ DagHub initialization failed: {e}. Continuing without MLflow tracking.")
+
 class ModelTrainer:
     def __init__(self, model_trainer_config: Model_trainer_config, data_transformation_artifact: DataTransformationArtifact):
         try: