SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.

model_manager/utils/repository/model_repository.py

@@ -9,6 +9,7 @@ and logging model runs to MLFlow.
 
 """
 
+import os
 import shutil
 from datetime import datetime
 from os import makedirs, path
@@ -22,14 +23,81 @@ from model_manager.sientia.reports import Reports  # type: ignore[import-untyped
 from model_manager.utils.models.train_model_result import TrainModelResult
 
 
-class MLFlowRepository:
-    def __init__(self, host, username, password, logger: Logger):
-        self.model_serving = ModelServing(
-            tracking_uri=host, username=username, password=password, logger=logger
-        )
+class ModelRepository:
+    def __init__(self, url, username, password, logger: Logger):
+        self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
+
         self.logger = logger
 
-    def get_next_run_name(self, experiment_name: str) -> str:
+    def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
+        """
+        Save a trained ML model and its artifacts to MLflow.
+
+        This activity orchestrates the complete model saving pipeline:
+        1. Generates the next run name for the experiment
+        2. Creates and organizes artifacts (reports, data files)
+        3. Logs model, parameters, metrics, and artifacts to MLflow
+
+        Args:
+            input_data: Configuration for model saving operation
+                Required keys:
+                    - metadata (dict): Workflow execution metadata
+                    - train_result (TrainModelResult): Training result with model and metrics
+
+        Returns:
+            TrainModelResult: Updated training result with run_name and artifacts
+
+        Raises:
+            Exception: If model saving fails (after sending notification)
+        """
+        experiment_name = train_result.params.experiment_name
+        self.logger.info(f'Starting model save for experiment: {experiment_name}')
+
+        # Step 1: Generate next run name
+        self.logger.info('Generating run name')
+        train_result.run_name = self._get_next_run_name(experiment_name)
+        self.logger.info(f'Generated run name: {train_result.run_name}')
+
+        # Step 2: Generate artifacts (reports, CSV files)
+        self.logger.info('Generating artifacts')
+        train_result = self._generate_artifacts(train_result)
+        self.logger.info('Artifacts generated successfully')
+
+        # Step 3: Save run to MLflow
+        self.logger.info('Saving run to MLflow')
+        self._save_run(train_result)
+
+        self.logger.info(
+            f'Model saved successfully - Run: {train_result.run_name}, '
+            f'Experiment: {experiment_name}'
+        )
+
+        return train_result
+
+    def cleanup_run_directory(self, run_dir: str) -> None:
+        """
+        Clean up temporary run directory after model training.
+
+        This activity deletes the temporary directory created during model training
+        and artifact generation. It implements idempotent cleanup to handle cases
+        where the directory may have already been deleted.
+
+        Args:
+            run_dir (str): Path to the run directory to delete
+        """
+        if not run_dir:
+            self.logger.info('No run directory specified, skipping cleanup')
+            return
+
+        self.logger.info(f'Cleaning up run directory: {run_dir}')
+
+        if os.path.exists(run_dir):
+            shutil.rmtree(run_dir)
+            self.logger.info(f'Run directory deleted successfully: {run_dir}')
+        else:
+            self.logger.info(f'Run directory already deleted: {run_dir}')
+
+    def _get_next_run_name(self, experiment_name: str) -> str:
         """
         Generates the next run name for a given experiment.
 
@@ -46,7 +114,7 @@ class MLFlowRepository:
         next_run_number = len(runs) + 1
         return f'{experiment_name}-{next_run_number}'
 
-    def generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
+    def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
         """
         Generates and organizes artifacts related to the training process, such as reports and data files.
 
@@ -87,7 +155,7 @@ class MLFlowRepository:
         self._setup_run_directory(data.run_dir, header_file_path)
         return self._generate_report(reference_data, current_data, data)
 
-    def save_run(self, data: TrainModelResult):
+    def _save_run(self, data: TrainModelResult):
         """
         Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
         to the Sientia tracking system.
@@ -122,60 +190,48 @@ class MLFlowRepository:
             self.logger.error(error_msg)
             raise ValueError(error_msg)
 
-        try:
-            # Prepare parameters
-            train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
-            interval_strs = [
-                (str(interval[0]), str(interval[1]))
-                for interval in (data.params.removed_intervals or [])
-            ]
+        # Prepare parameters
+        train_test_split = f'{data.params.train_size}-{100 - data.params.train_size}'
 
-            # Set experiment and create run
-            self.model_serving.set_experiment(data.params.experiment_name)
-            self.logger.info(
-                f"Logging run '{data.run_name}' to experiment '{data.params.experiment_name}'"
-            )
+        interval_strs = [
+            (str(interval[0]), str(interval[1]))
+            for interval in (data.params.removed_intervals or [])
+        ]
 
-            with self.model_serving.save_experiment(
-                run_name=data.run_name, description=data.params.experiment_name
-            ):
-                # Log model parameters
-                self.model_serving.log_param('model_type', 'Linear Regression')
-                self.model_serving.log_param('target_variable', data.params.target_variable)
-                self.model_serving.log_param('input_variables', data.params.variable_columns)
-                self.model_serving.log_param('lag_train', data.params.lag_train)
-                self.model_serving.log_param('lag_val', data.params.lag_val)
-                self.model_serving.log_param('ma', data.params.window)
-                self.model_serving.log_param('low_lim', data.params.low_lim)
-                self.model_serving.log_param('upp_lim', data.params.upp_lim)
-                self.model_serving.log_param('normalized', data.scaler_dict)
-                self.model_serving.log_param('ar', data.params.include_ar)
-                self.model_serving.log_param('Train_test_split', train_test_split)
-                self.model_serving.log_param('Removed_intervals', interval_strs)
-                self.model_serving.log_param('Retrain', False)
+        # Set experiment and create run
+        self.model_serving.set_experiment(data.params.experiment_name)
 
-                # Log evaluation metrics
-                self.model_serving.log_metric('MSE', data.mse_val)
-                self.model_serving.log_metric('R2', data.r2_val)
-                self.model_serving.log_metric('MAE', data.mae_val)
+        with self.model_serving.save_experiment(
+            run_name=data.run_name, description=data.params.experiment_name
+        ):
+            # Log model parameters
+            self.model_serving.log_param('model_type', 'Linear Regression')
+            self.model_serving.log_param('target_variable', data.params.target_variable)
+            self.model_serving.log_param('input_variables', data.params.variable_columns)
+            self.model_serving.log_param('lag_train', data.params.lag_train)
+            self.model_serving.log_param('lag_val', data.params.lag_val)
+            self.model_serving.log_param('ma', data.params.window)
+            self.model_serving.log_param('low_lim', data.params.low_lim)
+            self.model_serving.log_param('upp_lim', data.params.upp_lim)
+            self.model_serving.log_param('normalized', data.scaler_dict)
+            self.model_serving.log_param('ar', data.params.include_ar)
+            self.model_serving.log_param('Train_test_split', train_test_split)
+            self.model_serving.log_param('Removed_intervals', interval_strs)
+            self.model_serving.log_param('Retrain', False)
 
-                # Log models
-                self.model_serving.log_model(data.process_data, 'data_model')
-                self.model_serving.log_model(data.regr, 'prediction_model')
+            # Log evaluation metrics
+            self.model_serving.log_metric('MSE', data.mse_val)
+            self.model_serving.log_metric('R2', data.r2_val)
+            self.model_serving.log_metric('MAE', data.mae_val)
 
-                # Log artifacts
-                self.model_serving.log_artifact(data.report_path)
-                self.model_serving.log_artifact(data.train_data_path)
-                self.model_serving.log_artifact(data.test_data_path)
+            # Log models
+            self.model_serving.log_model(data.process_data, 'data_model')
+            self.model_serving.log_model(data.regr, 'prediction_model')
 
-                self.logger.info(
-                    f"Successfully logged run '{data.run_name}' with metrics: MSE={data.mse_val:.4f}, R2={data.r2_val:.4f}, MAE={data.mae_val:.4f}"
-                )
-
-        except Exception as e:
-            error_msg = f"Failed to save run '{data.run_name}' to MLflow: {str(e)}"
-            self.logger.error(error_msg)
-            raise RuntimeError(error_msg) from e
+            # Log artifacts
+            self.model_serving.log_artifact(data.report_path)
+            self.model_serving.log_artifact(data.train_data_path)
+            self.model_serving.log_artifact(data.test_data_path)
 
     def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
         """
@@ -258,7 +314,6 @@ class MLFlowRepository:
 
         try:
             makedirs(run_dir, exist_ok=True)
-            self.logger.info(f'Created run directory: {run_dir}')
             return run_dir
         except PermissionError as e:
             error_msg = f'Permission denied when creating directory: {run_dir}'
@@ -298,9 +353,6 @@ class MLFlowRepository:
             # Copy header file to run directory
             header_dest = path.join(run_dir, 'header.html')
             shutil.copy(header_file_path, header_dest)
-
-            self.logger.info(f'Run directory setup completed successfully in: {run_dir}')
-
         except FileNotFoundError as e:
             error_msg = f'Header file not found: {header_file_path}'
             self.logger.error(error_msg)
@@ -360,20 +412,16 @@ class MLFlowRepository:
             # Save HTML report
             data.report_path = path.join(data.run_dir, 'report.html')
             report.save_all_sections_html(data.report_path)
-            self.logger.info(f'Generated HTML report: {data.report_path}')
 
             # Save training data CSV
             data.train_data_path = path.join(data.run_dir, 'train_data.csv')
             reference_data.to_csv(data.train_data_path, index=False)
-            self.logger.info(f'Saved training data: {data.train_data_path}')
 
             # Save test data CSV
             data.test_data_path = path.join(data.run_dir, 'test_data.csv')
             current_data.to_csv(data.test_data_path, index=False)
-            self.logger.info(f'Saved test data: {data.test_data_path}')
 
             return data
-
         except ValueError as e:
             error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
             self.logger.error(error_msg)