Add MLflow tracking integration

.env.example

@@ -0,0 +1,6 @@
+# DagsHub MLflow Tracking Configuration
+# Get your token from: https://dagshub.com/user/settings/tokens
+
+DAGSHUB_USERNAME=your_dagshub_username
+DAGSHUB_TOKEN=your_dagshub_token_here
+MLFLOW_TRACKING_URI=https://dagshub.com/your_username/your_repo.mlflow/

config/config.yaml

@@ -47,7 +47,7 @@ model_pusher:
   model_registry_path: models/production
 
 mlflow:
-  tracking_uri: http://localhost:5000
+  tracking_uri: https://dagshub.com/abheshith7/AutoML-MLOps-PipeLine.mlflow/
   experiment_name: automl_experiment
   run_name: null
   registry_uri: null

src/mlpipeline/automl/autogluon_trainer.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Tuple
 import pandas as pd
 from autogluon.tabular import TabularPredictor
 
@@ -34,13 +34,24 @@ class AutoGluonTrainer:
         leaderboard = self.predictor.leaderboard(silent=True)
         best_model = leaderboard.iloc[0]
         
+        # Get feature importance if available
+        try:
+            feature_importance = self.predictor.feature_importance(data=train_data)
+        except:
+            feature_importance = None
+        
         metrics = {
-            'score': float(best_model['score_val']),
+            'validation_accuracy': float(best_model['score_val']),
+            'score': float(best_model['score_val']),  # Keep for backward compatibility
             'score_test': float(best_model.get('score_test', 0.0)),
+            'fit_time': float(best_model.get('fit_time', 0.0)),
+            'pred_time_val': float(best_model.get('pred_time_val', 0.0)),
+            'num_models_trained': len(leaderboard),
+            'best_model_name': str(best_model['model']),
         }
         
         logger.info(f"AutoGluon training completed. Best score: {metrics['score']}")
-        return metrics
+        return metrics, feature_importance
     
     def predict(self, data: pd.DataFrame) -> pd.Series:
         if self.predictor is None:

src/mlpipeline/components/automl_trainer.py

@@ -1,4 +1,6 @@
 import pandas as pd
+import mlflow
+import os
 from mlpipeline.entity import ModelTrainerConfig, ModelTrainerArtifact
 from mlpipeline.automl import AutoMLFactory
 from mlpipeline.logging.logger import get_logger
@@ -7,7 +9,6 @@ from mlpipeline.constants import AUTOML_CONFIG_FILE_PATH
 from mlpipeline.utils.common import read_yaml
 from pathlib import Path
 import sys
-import os
 
 logger = get_logger(__name__)
 
@@ -20,11 +21,38 @@ class AutoMLTrainer:
         try:
             logger.info("Starting model training")
             
+            tracking_uri = os.getenv("MLFLOW_TRACKING_URI", "https://dagshub.com/abheshith7/AutoML-MLOps-PipeLine.mlflow/")
+            dagshub_token = os.getenv("DAGSHUB_TOKEN")
+            
+            mlflow_enabled = False
+            if dagshub_token and "dagshub.com" in tracking_uri:
+                try:
+                    os.environ["MLFLOW_TRACKING_USERNAME"] = os.getenv("DAGSHUB_USERNAME", "abheshith7")
+                    os.environ["MLFLOW_TRACKING_PASSWORD"] = dagshub_token
+                    mlflow.set_tracking_uri(tracking_uri)
+                    mlflow.set_experiment("automl_experiment")
+                    mlflow_enabled = True
+                    logger.info(f"MLflow tracking enabled: {tracking_uri}")
+                except Exception as e:
+                    logger.warning(f"MLflow tracking disabled: {str(e)}")
+                    mlflow_enabled = False
+            else:
+                logger.warning("MLflow tracking disabled: DAGSHUB_TOKEN not set")
+            
+            if mlflow_enabled:
+                mlflow.start_run()
+            
             train_df = pd.read_csv(self.config.train_data_path)
             
             automl_config = read_yaml(Path(AUTOML_CONFIG_FILE_PATH))
             library_config = automl_config[self.config.automl_library]
             
+            if mlflow_enabled:
+                mlflow.log_param("automl_library", self.config.automl_library)
+                mlflow.log_param("target_column", self.config.target_column)
+                mlflow.log_param("train_samples", len(train_df))
+                mlflow.log_params(library_config)
+            
             trainer = AutoMLFactory.create_trainer(
                 self.config.automl_library,
                 library_config
@@ -33,11 +61,64 @@ class AutoMLTrainer:
             os.makedirs(self.config.root_dir, exist_ok=True)
             
             if self.config.automl_library == 'autogluon':
-                metrics = trainer.train(train_df, self.config.target_column, self.config.model_path)
+                result = trainer.train(train_df, self.config.target_column, self.config.model_path)
+                if isinstance(result, tuple):
+                    metrics, feature_importance = result
+                else:
+                    metrics = result
+                    feature_importance = None
             else:
                 X_train = train_df.drop(columns=[self.config.target_column])
                 y_train = train_df[self.config.target_column]
                 metrics = trainer.train(X_train, y_train, self.config.model_path)
+                feature_importance = None
+            
+            if mlflow_enabled:
+                # Separate numeric metrics from string values
+                numeric_metrics = {}
+                string_values = {}
+                
+                for key, value in metrics.items():
+                    if isinstance(value, (int, float)):
+                        numeric_metrics[key] = value
+                    else:
+                        string_values[key] = str(value)
+                
+                # Log numeric metrics only
+                if numeric_metrics:
+                    mlflow.log_metrics(numeric_metrics)
+                
+                # Log string values as tags
+                for key, value in string_values.items():
+                    mlflow.set_tag(key, value)
+                
+                # Log feature importance as artifact
+                if feature_importance is not None:
+                    import json
+                    fi_dict = feature_importance.to_dict() if hasattr(feature_importance, 'to_dict') else {}
+                    fi_path = Path(self.config.root_dir) / "feature_importance.json"
+                    with open(fi_path, 'w') as f:
+                        json.dump(fi_dict, f, indent=2)
+                    mlflow.log_artifact(str(fi_path))
+                
+                # Log model leaderboard
+                try:
+                    from autogluon.tabular import TabularPredictor
+                    predictor = TabularPredictor.load(str(self.config.model_path))
+                    leaderboard = predictor.leaderboard(silent=True)
+                    lb_path = Path(self.config.root_dir) / "leaderboard.csv"
+                    leaderboard.to_csv(lb_path, index=False)
+                    mlflow.log_artifact(str(lb_path))
+                except:
+                    pass
+                
+                # Set additional tags
+                mlflow.set_tag("model_type", "AutoML")
+                mlflow.set_tag("framework", self.config.automl_library)
+                
+                run_id = mlflow.active_run().info.run_id
+                logger.info(f"MLflow run logged: {run_id}")
+                mlflow.end_run()
             
             logger.info(f"Model trained with metrics: {metrics}")
             
@@ -48,4 +129,6 @@ class AutoMLTrainer:
                 message=f"Model trained successfully with score: {metrics.get('score', 0.0):.4f}"
             )
         except Exception as e:
+            if mlflow.active_run():
+                mlflow.end_run()
             raise ModelTrainingException(str(e), sys)

src/mlpipeline/components/model_evaluation.py

@@ -1,12 +1,13 @@
 import pandas as pd
 import json
+import mlflow
+import os
 from sklearn.metrics import accuracy_score, f1_score, classification_report
 from mlpipeline.entity import ModelEvaluationConfig, ModelEvaluationArtifact
 from autogluon.tabular import TabularPredictor
 from mlpipeline.logging.logger import get_logger
 from mlpipeline.exception import ModelEvaluationException
 import sys
-import os
 
 logger = get_logger(__name__)
 
@@ -19,8 +20,32 @@ class ModelEvaluation:
         try:
             logger.info("Starting model evaluation")
             
+            tracking_uri = os.getenv("MLFLOW_TRACKING_URI", "https://dagshub.com/abheshith7/AutoML-MLOps-PipeLine.mlflow/")
+            dagshub_token = os.getenv("DAGSHUB_TOKEN")
+            
+            mlflow_enabled = False
+            if dagshub_token and "dagshub.com" in tracking_uri:
+                try:
+                    os.environ["MLFLOW_TRACKING_USERNAME"] = os.getenv("DAGSHUB_USERNAME", "abheshith7")
+                    os.environ["MLFLOW_TRACKING_PASSWORD"] = dagshub_token
+                    mlflow.set_tracking_uri(tracking_uri)
+                    mlflow.set_experiment("automl_experiment")
+                    mlflow_enabled = True
+                    logger.info(f"MLflow tracking enabled: {tracking_uri}")
+                except Exception as e:
+                    logger.warning(f"MLflow tracking disabled: {str(e)}")
+                    mlflow_enabled = False
+            else:
+                logger.warning("MLflow tracking disabled: DAGSHUB_TOKEN not set")
+            
+            if mlflow_enabled:
+                mlflow.start_run()
+            
             test_df = pd.read_csv(self.config.test_data_path)
             
+            if mlflow_enabled:
+                mlflow.log_param("test_samples", len(test_df))
+            
             predictor = TabularPredictor.load(str(self.config.model_path))
             
             predictions = predictor.predict(test_df)
@@ -29,24 +54,63 @@ class ModelEvaluation:
             predictions_binary = (predictions > 0).astype(int)
             y_test_binary = (y_test > 0).astype(int)
             
+            # Calculate comprehensive metrics
+            from sklearn.metrics import precision_score, recall_score, roc_auc_score, confusion_matrix
+            
             accuracy = float(accuracy_score(y_test_binary, predictions_binary))
             f1 = float(f1_score(y_test_binary, predictions_binary, average='weighted'))
+            precision = float(precision_score(y_test_binary, predictions_binary, average='weighted'))
+            recall = float(recall_score(y_test_binary, predictions_binary, average='weighted'))
+            
+            # Get prediction probabilities for ROC-AUC
+            try:
+                pred_proba = predictor.predict_proba(test_df)
+                if hasattr(pred_proba, 'iloc'):
+                    pred_proba_positive = pred_proba.iloc[:, 1]
+                else:
+                    pred_proba_positive = pred_proba[:, 1]
+                auc = float(roc_auc_score(y_test_binary, pred_proba_positive))
+            except:
+                auc = 0.0
             
             metrics = {
-                "accuracy": accuracy,
-                "f1_score": f1
+                "test_accuracy": accuracy,
+                "test_f1_score": f1,
+                "test_precision": precision,
+                "test_recall": recall,
+                "test_roc_auc": auc
+            }
+            
+            # Save confusion matrix
+            cm = confusion_matrix(y_test_binary, predictions_binary)
+            cm_dict = {
+                "true_negatives": int(cm[0][0]),
+                "false_positives": int(cm[0][1]),
+                "false_negatives": int(cm[1][0]),
+                "true_positives": int(cm[1][1])
             }
             
+            if mlflow_enabled:
+                mlflow.log_metrics(metrics)
+                mlflow.log_metrics(cm_dict)
+                mlflow.set_tag("evaluation_stage", "test")
+            
             os.makedirs(self.config.root_dir, exist_ok=True)
             
             with open(self.config.metrics_file, "w") as f:
                 json.dump(metrics, f, indent=2)
             
+            if mlflow_enabled:
+                mlflow.log_artifact(str(self.config.metrics_file))
+                run_id = mlflow.active_run().info.run_id
+                logger.info(f"MLflow run logged: {run_id}")
+                mlflow.end_run()
+            
             logger.info(f"Evaluation metrics: {metrics}")
             
             return ModelEvaluationArtifact(
                 is_model_accepted=True,
-                evaluation_metrics=metrics,
+                evaluation_metrics={"accuracy": accuracy, "f1_score": f1},
                 message=f"Model evaluation completed with accuracy: {accuracy:.4f}"
             )
         except Exception as e:

src/mlpipeline/constants.py

@@ -11,5 +11,5 @@ MODELS_DIR = Path("artifacts/models")
 DATA_DIR = Path("artifacts/data")
 REPORTS_DIR = Path("artifacts/reports")
 
-MLFLOW_TRACKING_URI = "http://localhost:5000"
+MLFLOW_TRACKING_URI = "https://dagshub.com/abheshith7/AutoML-MLOps-PipeLine.mlflow/"
 MLFLOW_REGISTRY_URI = None