propre et fonctionnel

src/main.py

@@ -6,7 +6,10 @@
 import os
 import logging
 import hydra
+import mlflow
+
 from omegaconf import DictConfig, OmegaConf
+from config import Config
 from trainers.cuml.svm_trainer import SvmTrainer
 from trainers.cuml.random_forest_trainer import RandomForestTrainer
 from trainers.cuml.logistic_regression_trainer import LogisticRegressionTrainer
@@ -15,19 +18,18 @@ from trainers.huggingface.huggingface_transformer_trainer import HuggingFaceTran
 from optimizers.optuna_optimizer import OptunaOptimizer
 from optimizers.ray_tune_optimizer import RayTuneOptimizer
 from mlflow_integration.mlflow_decorator import MLflowDecorator
-from config import Config
+
+import tempfile
+import pickle
+import pandas as pd
+from utilities.cuml_pyfunc_wrapper import CuMLPyFuncWrapper
+
 logger = logging.getLogger(__name__)
 
 
 def get_trainer(config: Config):
     """
     Crée et retourne l'instance du trainer approprié en fonction de la configuration.
-    
-    Args:
-        config: Objet de configuration
-        
-    Returns:
-        Une instance concrète de BaseTrainer
     """
     model_type = config.model.type.lower()
 
@@ -43,7 +45,6 @@ def get_trainer(config: Config):
     if model_type not in trainer_map:
         raise ValueError(f"Type de modèle non supporté: {model_type}")
 
-    # Création de l'instance du trainer avec la configuration
     trainer_class = trainer_map[model_type]
     return trainer_class(
         config=config,
@@ -55,16 +56,8 @@ def get_trainer(config: Config):
 def get_optimizer(config: Config):
     """
     Crée et retourne l'instance d'optimizer appropriée en fonction de la configuration.
-    
-    Args:
-        config: Objet de configuration
-        
-    Returns:
-        Une instance concrète de HyperparameterOptimizer
     """
     optimizer_type = config.hyperparameters.optimizer.lower()
-
-    # Mapping des types d'optimizers
     optimizer_map = {
         "optuna": OptunaOptimizer,
         "raytune": RayTuneOptimizer,
@@ -73,23 +66,56 @@ def get_optimizer(config: Config):
     if optimizer_type not in optimizer_map:
         raise ValueError(f"Type d'optimizer non supporté: {optimizer_type}")
 
-    # Création de l'instance de l'optimizer
-    optimizer_class = optimizer_map[optimizer_type]
-    return optimizer_class()
+    return optimizer_map[optimizer_type]()
+
+
+def log_cuml_model_to_mlflow(trainer_instance, run_id=None):
+    """
+    Sérialise le vectorizer et le classifier dans un répertoire temporaire
+    puis log le tout dans MLflow en tant que modèle PyFunc.
+
+    Les artifacts sont ainsi stockés dans mlruns, liés au run en cours.
+    """
+    logger.info("Logging du modèle CuML via mlflow.pyfunc.log_model...")
+    input_example = pd.DataFrame({"example_text": ["exemple"]})
+
+    # On va utiliser mlflow.pyfunc.log_model pour stocker le wrapper PyFunc + nos artifacts
+    with tempfile.TemporaryDirectory() as tmpdir:
+        vectorizer_path = os.path.join(tmpdir, "vectorizer.pkl")
+        classifier_path = os.path.join(tmpdir, "classifier.pkl")
+
+        # Sauvegarde sur disque
+        with open(vectorizer_path, "wb") as vf:
+            pickle.dump(trainer_instance.vectorizer, vf)
+        with open(classifier_path, "wb") as cf:
+            pickle.dump(trainer_instance.classifier, cf)
+
+        # PyFunc wrapper (placeholder, héberge la logique de load_model)
+        pyfunc_wrapper = CuMLPyFuncWrapper(
+            vectorizer=None,
+            classifier=None
+        )
+
+        # Log en PyFunc; "cuml_model" est le chemin (artifact_path) où sera stocké le modèle dans MLflow
+        mlflow.pyfunc.log_model(
+            artifact_path="cuml_model",
+            python_model=pyfunc_wrapper,
+            artifacts={
+                "vectorizer": vectorizer_path,
+                "classifier": classifier_path
+            },
+            input_example=input_example
+        )
+
+    logger.info("Le modèle et ses artifacts ont été enregistrés dans MLflow.")
 
 
 @hydra.main(config_path="conf", config_name="config", version_base=None)
 def main(cfg: DictConfig) -> None:
     """
     Point d'entrée principal de l'application.
-    
-    Args:
-        cfg: Configuration Hydra sous forme de DictConfig
     """
-    # La configuration Hydra 'cfg' est maintenant correctement structurée grâce aux defaults dans config.yaml.
-    # Conversion directe en configuration Pydantic.
     try:
-        # Convertit directement cfg en dict et l'utilise pour instancier Config
         config = Config(**OmegaConf.to_container(cfg, resolve=True))
     except Exception as e:
         logger.error(f"Erreur lors de la validation Pydantic de la configuration: {e}")
@@ -98,14 +124,20 @@ def main(cfg: DictConfig) -> None:
 
     logger.info(f"Configuration Pydantic finale chargée: {config}")
 
-    # Création du trainer approprié
-    trainer = get_trainer(config)
+    # Sélection du tracker MLflow
+    mlflow.set_tracking_uri(config.mlflow.tracking_uri)
 
-    # Construction des composants (vectorizer, classifier, etc.)
+    trainer = get_trainer(config)
     trainer.build_components()
 
-    # Définir la séquence d'opérations à exécuter dans un seul run MLflow
     def run_pipeline(trainer_instance):
+        """
+        Exécute la séquence complète :
+          - Optimisation hyperparamètres (si besoin)
+          - Entraînement
+          - Évaluation
+          - Logging MLflow (paramètres, métriques, et modèles)
+        """
         logger.info("Vérification et lancement de l'optimisation des hyperparamètres si nécessaire...")
         trainer_instance.optimize_if_needed()
 
@@ -113,26 +145,29 @@ def main(cfg: DictConfig) -> None:
         trainer_instance.train()
 
         logger.info("Lancement de l'évaluation...")
-        trainer_instance.evaluate()
+        metrics = trainer_instance.evaluate()
+        logger.info(f"Metrics calculés: {metrics}")
+
+        # Log des métriques
+        mlflow.log_metrics(metrics)
 
         logger.info("Logging des paramètres...")
         trainer_instance.log_parameters_to_mlflow()
 
-        logger.info("Pipeline complet (optimisation, entraînement, évaluation, logging) terminé.")
+        # Log du modèle final (vectorizer+classifier) sous forme PyFunc
+        log_cuml_model_to_mlflow(trainer_instance)
+
+        logger.info("Pipeline MLflow complet terminé.")
 
-    # Initialiser le décorateur MLflow
+    # On utilise un décorateur défini pour centraliser le démarrage/arrêt du run
     mlflow_decorator = MLflowDecorator(
         experiment_name=config.mlflow.experiment_name,
         tracking_uri=config.mlflow.tracking_uri
     )
 
-    # Appliquer le décorateur à la fonction pipeline complète
     run_pipeline_with_mlflow = mlflow_decorator(run_pipeline)
-
-    # Exécuter le pipeline complet sous un seul run MLflow
     logger.info("Lancement du pipeline complet avec MLflow...")
     run_pipeline_with_mlflow(trainer)
-
     logger.info("Pipeline MLflow terminé avec succès.")