adding mlfow registered model loading

Dockerfile

@@ -26,8 +26,6 @@ RUN mkdir -p /app/data /app/final_model /app/templates
 
 # run the load_data_to_sqlite.py script to initialize the database
 RUN python load_data_to_sqlite.py
-# Train the model during build (this persists across container restarts)
-RUN python -c "from src.pipeline.training_pipeline import Trainingpipeline; Trainingpipeline().run_pipeline()"
 
 # Expose port 7860 (HF Space requirement)
 EXPOSE 7860

app.py

@@ -11,6 +11,9 @@ from fastapi.responses import Response
 from starlette.responses import RedirectResponse
 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()
@@ -18,15 +21,92 @@ mongo_db_uri = os.getenv("MONGO_DB_URI")
 
 from src.constant.training_pipeline import DATA_INGESTION_COLLECTION_NAME
 from src.constant.training_pipeline import DATA_INGESTION_DATBASE_NANE
-from src.utils.main_utils.utils import load_object
+from src.utils.main_utils.utils import load_object, save_object
 # import pymongo
 
 # client = pymongo.MongoClient(mongo_db_uri,tlsCAFile=ca)
 # database = client[DATA_INGESTION_DATBASE_NANE]
 # collection = database[DATA_INGESTION_COLLECTION_NAME]
 from fastapi.templating import Jinja2Templates
-templates = Jinja2Templates(directory="./templates")  
-app = FastAPI()
+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}
+
+def load_models_from_mlflow():
+    """Load latest models from MLflow"""
+    try:
+        if not MLFLOW_AVAILABLE:
+            logging.error("MLflow not available")
+            return False
+        
+        # Get the latest run from the experiment
+        client = mlflow.tracking.MlflowClient()
+        experiment = client.get_experiment_by_name("Default")
+        
+        if experiment is None:
+            logging.warning("No MLflow experiment found. Train model first.")
+            return False
+        
+        runs = client.search_runs(
+            experiment_ids=[experiment.experiment_id],
+            order_by=["start_time DESC"],
+            max_results=1
+        )
+        
+        if not runs:
+            logging.warning("No MLflow runs found. Train model first.")
+            return False
+        
+        latest_run = runs[0]
+        run_id = latest_run.info.run_id
+        
+        logging.info(f"Loading models from MLflow run: {run_id}")
+        
+        # Load model and preprocessor
+        model_uri = f"runs:/{run_id}/model"
+        preprocessor_uri = f"runs:/{run_id}/preprocessor"
+        
+        MODEL_CACHE["model"] = mlflow.sklearn.load_model(model_uri)
+        MODEL_CACHE["preprocessor"] = mlflow.sklearn.load_model(preprocessor_uri)
+        
+        # Save to local directory as backup
+        os.makedirs("final_model", exist_ok=True)
+        save_object("final_model/model.pkl", MODEL_CACHE["model"])
+        save_object("final_model/preprocessor.pkl", MODEL_CACHE["preprocessor"])
+        
+        logging.info("✅ Models loaded from MLflow and cached locally")
+        return True
+        
+    except Exception as e:
+        logging.error(f"Error loading models from MLflow: {e}")
+        return False
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Load models on startup"""
+    logging.info("===== Application Startup - Loading models from MLflow =====")
+    
+    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.")
+    else:
+        logging.warning("⚠️ MLflow not available. Please train via /train endpoint.")
+    
+    yield
+    logging.info("===== Application Shutdown =====")
+
+app = FastAPI(lifespan=lifespan)
 
 orgin = ["*"]
 
@@ -45,13 +125,17 @@ app.add_middleware(
 @app.get("/")
 async def root():
     """Root endpoint with system status"""
+    model_status = "✅ Ready" if MODEL_CACHE["model"] is not None else "⚠️ Not trained - call /train first"
+    
     return {
         "status": "running",
         "service": "Network Security System - Phishing Detection",
+        "model_status": model_status,
+        "mlflow_enabled": MLFLOW_AVAILABLE,
         "endpoints": {
             "docs": "/docs",
-            "train": "/train",
-            "predict": "/predict"
+            "train": "/train (trains and logs to MLflow)",
+            "predict": "/predict (loads from MLflow)"
         }
     }
 
@@ -61,28 +145,45 @@ async def training_route():
         logging.info("Starting training pipeline...")
         training_pipeline = Trainingpipeline()
         training_pipeline.run_pipeline()
-        return Response("Training completed successfully!")
+        
+        # Reload models from MLflow after training
+        if MLFLOW_AVAILABLE:
+            load_models_from_mlflow()
+        
+        return Response("✅ Training completed and models loaded from MLflow!")
     except Exception as e:
         raise NetworkSecurityException(e, sys)
 
 @app.post("/predict") # predict route
 async def predict_route(request: Request, file: UploadFile =File(...)):
     try:
+        # Check if models are loaded
+        if MODEL_CACHE["model"] is None or MODEL_CACHE["preprocessor"] is None:
+            # Try to load from MLflow
+            if not load_models_from_mlflow():
+                return Response(
+                    "❌ No trained model available. Please call /train endpoint first.",
+                    status_code=400
+                )
+        
         df = pd.read_csv(file.file)
         # Remove target column if it exists
         if 'Result' in df.columns:
             df = df.drop(columns=['Result'])
-        preprocessor = load_object(file_path = "final_model/preprocessor.pkl") 
-        model = load_object(file_path= "final_model/model.pkl")
-        NSmodel = NetworkSecurityModel(preprocessing_object= preprocessor, trained_model_object= model)
-        print(df.iloc[0])
+        
+        # Use cached models from MLflow
+        preprocessor = MODEL_CACHE["preprocessor"]
+        model = MODEL_CACHE["model"]
+        
+        NSmodel = NetworkSecurityModel(preprocessing_object=preprocessor, trained_model_object=model)
         y_pred = NSmodel.predict(df)
-        print(y_pred)
         df['predicted_column'] = y_pred
-        print(df['predicted_column'])
+        
+        # Save predictions
+        os.makedirs("final_model", exist_ok=True)
         df.to_csv("final_model/predicted.csv")
         
-        table_html = df.to_html(classes = 'table table-striped')
+        table_html = df.to_html(classes='table table-striped')
         return templates.TemplateResponse("table.html", {"request": request, "table": table_html})
     
     except Exception as e:

src/components/model_trainer.py

@@ -30,8 +30,9 @@ class ModelTrainer:
         except Exception as e:
             raise NetworkSecurityException(e, sys) from e
     
-    def track_mlflow(self,best_model, classificationmetric):
-        with mlflow.start_run():
+    def track_mlflow(self, best_model, preprocessor, classificationmetric):
+        """Log model, preprocessor, and metrics to MLflow"""
+        with mlflow.start_run() as run:
             f1_score = classificationmetric.f1_score
             precision_score = classificationmetric.precision_score
             recall_score = classificationmetric.recall_score
@@ -39,7 +40,14 @@ class ModelTrainer:
             mlflow.log_metric("f1_score", f1_score)
             mlflow.log_metric("precision_score", precision_score)
             mlflow.log_metric("recall_score", recall_score)
+            
+            # Log both model and preprocessor
             mlflow.sklearn.log_model(best_model, "model")
+            mlflow.sklearn.log_model(preprocessor, "preprocessor")
+            
+            # Log run ID for easy retrieval
+            logging.info(f"✅ Models logged to MLflow - Run ID: {run.info.run_id}")
+            return run.info.run_id
         
     def train_model(self, x_train, y_train,x_test, y_test):
         models = {
@@ -104,15 +112,13 @@ class ModelTrainer:
         y_train_pred = best_model.predict(x_train)
         classification_train_metric= classification_score(y_true = y_train, y_pred=y_train_pred)
         
-        # track mlfow 
-        self.track_mlflow(best_model, classification_train_metric)
-        
-        
-        
         y_test_pred = best_model.predict(x_test)
         classification_test_metric = classification_score(y_true = y_test, y_pred=y_test_pred)
         
         preprocessor  = load_object(file_path=self.data_transformation_artifact.transformed_object_file_path)
+        
+        # Track to MLflow (logs model + preprocessor)
+        self.track_mlflow(best_model, preprocessor, classification_train_metric)
         model_dir_path = os.path.dirname(self.model_trainer_config.trained_model_file_path)
         os.makedirs(model_dir_path, exist_ok=True)