src/modeling/mlflow_utils.py

"""Module de gestion des expérimentations avec MLflow.
Ce module contient des fonctions pour configurer MLflow, lancer des expérimentations complètes avec logging des paramètres, métriques et artefacts, et pour logguer le modèle final dans le registry.
Il gère les spécificités des modèles (notamment CatBoost) et organise les artefacts de manière structurée pour faciliter l'analyse et la comparaison des modèles. Les fonctions sont conçues pour être flexibles et réutilisables dans différents contextes d'expérimentation.
"""
import mlflow
import mlflow.sklearn
from pathlib import Path
from catboost import CatBoostRegressor

import time
from src.features.preprocessing import get_categorical_features
from src.modeling.cv import cross_validate_catboost,cross_validate_model

def log_model_to_mlflow(model, model_name, model_artifact_name="model"):
    """
    Log le modèle dans MLflow selon sa librairie.
    - XGBoost : mlflow.xgboost
    - LightGBM : mlflow.lightgbm
    - CatBoost : mlflow.catboost
    - Autres : mlflow.sklearn
    Args:
        - model : modèle entraîné
        - model_name : str, nom du modèle (pour les logs)
        - model_artifact_name : str, nom de l'artifact dans MLflow (par défaut "model")

    """
    model_class = model.__class__.__name__

    try:
        if "XGB" in model_class:
            import mlflow.xgboost
            mlflow.xgboost.log_model(model, name=model_artifact_name)

        elif "LGBM" in model_class:
            import mlflow.lightgbm
            mlflow.lightgbm.log_model(model, name=model_artifact_name)

        elif "CatBoost" in model_class:
            import mlflow.catboost
            mlflow.catboost.log_model(model, name=model_artifact_name)

        else:
            import mlflow.sklearn
            mlflow.sklearn.log_model(model, name=model_artifact_name)

    except Exception as e:
        print(f"Erreur lors du log du modèle {model_name} : {e}")

def run_experiment(
    model_name,
    model,
    X_train,
    X_test,
    y_train,
    y_test,
    build_model_pipeline,
    cross_validate_model,
    train_model,
    evaluate_regression_model,
    cv=5
):
    """
    Lance une expérimentation complète avec MLflow :
    - pipeline
    - cross-validation
    - entraînement
    - évaluation test
    - logging MLflow
    Le logging MLflow inclut :
    - paramètres de l'expérience (model_name, cv, train_size, test_size, training_time)
    - métriques de la CV (MAE, RMSE, R²)
    - métriques du test (MAE, RMSE, R² en log et en réel)
    - modèle entraîné dans les artifacts
    Le pipeline est construit avec build_model_pipeline, qui intègre le préprocessing adapté au type de modèle.
    La validation croisée est réalisée avec cross_validate_model, qui gère les spécificités des modèles (notamment CatBoost).
    L'entraînement final est réalisé avec train_model, qui prend en compte les particularités de chaque modèle.
    L'évaluation est réalisée avec evaluate_regression_model, qui calcule les métriques à la fois en log et en réel.
    Args:
        - model_name : str, nom du modèle pour les logs
        - model : instance du modèle à entraîner
        - X_train, X_test, y_train, y_test : données d'entraînement et  de test
        - build_model_pipeline : fonction pour construire le pipeline avec préprocessing
        - cross_validate_model : fonction pour réaliser la validation croisée
        - train_model : fonction pour entraîner le modèle final
        - evaluate_regression_model : fonction pour évaluer les performances du modèle
        - cv : int, nombre de folds pour la validation croisée (par défaut 5)
    Returns:
    - dict avec les résultats de l'expérimentation (métriques, modèle entraîné, temps d'entraînement)
    """

    try:
        with mlflow.start_run(run_name=model_name):
            start_time = time.time()
            if isinstance(model, CatBoostRegressor):

                # 1. colonnes catégorielles
                categorical_features = get_categorical_features(X_train)

                # 2. CV spécifique
                cv_results = cross_validate_catboost(
                    model=model,
                    X_train=X_train,
                    y_train=y_train,
                    cv=cv
                )

                # 3. train final
                trained_pipeline = model.fit(
                    X_train,
                    y_train,
                    cat_features=categorical_features,
                    verbose=0
                )

                # 4. prédiction
                y_pred = trained_pipeline.predict(X_test)

                # 5. log spécifique
                mlflow.log_param("native_catboost", True)
                mlflow.log_param("cat_features", categorical_features)


            else:
                pipeline = build_model_pipeline(model, X_train)

                cv_results = cross_validate_model(
                    pipeline,
                    X_train,
                    y_train,
                    cv=cv
                )

                trained_pipeline = train_model(
                    pipeline,
                    X_train,
                    y_train
                )

                y_pred = trained_pipeline.predict(X_test)

                mlflow.log_param("native_catboost", False)

            # Metrics test
            test_metrics = evaluate_regression_model(
                y_test,
                y_pred
            )

            #Params
            mlflow.log_param("model_name", model_name)
            mlflow.log_param("cv", cv)
            mlflow.log_param("train_size", X_train.shape[0])
            mlflow.log_param("test_size", X_test.shape[0])
            training_time = time.time() - start_time
            mlflow.log_param("training_time", training_time)

            #  Metrics CV
            for key, value in cv_results.items():
                mlflow.log_metric(key, value)

            #  Metrics test
            for key, value in test_metrics.items():
                mlflow.log_metric(key, value)

            #  Model
            log_model_to_mlflow(trained_pipeline, model_name, model_artifact_name="model")

            return {
                "model": model_name,
                "trained_pipeline": trained_pipeline,
                "training_time": training_time,
                **cv_results,
                **test_metrics
            }

    except Exception as e:
        print(f"Erreur MLflow run ({model_name}) : {e}")
        return None
    



def setup_mlflow(
    tracking_uri,
    experiment_name,
    artifact_root
):
    """
    Configure MLflow :
    - base sqlite
    - dossier artifacts
    - création de l'expérience si besoin

    Args:
        tracking_uri (str): ex "sqlite:///mlflow.db"
        experiment_name (str): nom de l'expérience
        artifact_root (Path): dossier artifacts
    """

    try:
        # Création du dossier artifacts
        artifact_root.mkdir(exist_ok=True, parents=True)

        # Connexion MLflow
        mlflow.set_tracking_uri(tracking_uri)

        # Vérifie si l'expérience existe
        experiment = mlflow.get_experiment_by_name(experiment_name)

        if experiment is None:
            mlflow.create_experiment(
                name=experiment_name,
                artifact_location=artifact_root.as_uri()
            )

        # Active l'expérience
        mlflow.set_experiment(experiment_name)

        print(f"MLflow configuré sur : {experiment_name}")

    except Exception as e:
        print(f"Erreur configuration MLflow : {e}")

import mlflow
import mlflow.sklearn
from pathlib import Path


def log_final_model_to_mlflow(
    model,
    model_name,
    registered_model_name,
    metrics,
    params=None,
    cv_metrics=None,
    artifacts_dir=None,
    artifact_path="model"
):
    """
    Log le modèle final dans MLflow :
    - paramètres
    - métriques test
    - métriques CV
    - artefacts CSV / PNG
    - modèle enregistré dans le registry
    Args:
        - model : modèle entraîné
        - model_name : str, nom du modèle pour les logs
        - registered_model_name : str, nom pour l'enregistrement dans le registry
        - metrics : dict, métriques de test à logguer
        - params : dict, paramètres à logguer (optionnel)
        - cv_metrics : dict, métriques de validation croisée à logguer (optionnel)
        - artifacts_dir : str ou Path, dossier contenant les artefacts à logguer (optionnel)
        - artifact_path : str, chemin dans MLflow pour stocker le modèle (par défaut "model")
    Le logging inclut :
        - paramètres de l'expérience (model_name, cv, train_size, test_size, training_time)
        - métriques de la CV (MAE, RMSE, R²)
        - métriques du test (MAE, RMSE, R² en log et en réel)
        - modèle entraîné dans les artifacts
    Le modèle est enregistré dans MLflow Registry avec le nom spécifié.

    """

    with mlflow.start_run(run_name=model_name):

        mlflow.log_param("model_name", model_name)

        if params is not None:
            mlflow.log_params(params)

        # Metrics test
        for key, value in metrics.items():
            mlflow.log_metric(key, float(value))

        # Metrics CV
        if cv_metrics is not None:
            for key, value in cv_metrics.items():
                mlflow.log_metric(f"cv_{key}", float(value))

        # Artifacts
        if artifacts_dir is not None:
            artifacts_dir = Path(artifacts_dir)

            if artifacts_dir.exists():
                mlflow.log_artifacts(
                    local_dir=str(artifacts_dir),
                    artifact_path="artifacts"
                )

        # Model + registry
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path=artifact_path,
            registered_model_name=registered_model_name
        )

        print(f"Modèle enregistré dans MLflow Registry : {registered_model_name}")