From a8b926649a53724679fea5cfeccaf3ecc6796faa Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 6 Apr 2026 10:38:24 -0300 Subject: [PATCH] feat: enhance training reporting and directory management - Added functionality to store run name and ID in the training results. - Implemented report generation in the `DataManagerRepository`, including methods to create a run directory and generate comprehensive reports. - Updated `TrainModelResult` to include `run_id` and `run_dir` attributes for better tracking of training sessions. --- model_manager/activities/training.py | 8 + .../utils/models/train_model_result.py | 3 + .../repository/data_manager_repository.py | 115 ++++- .../utils/repository/model_repository.py | 471 ++++++++++++++++++ 4 files changed, 596 insertions(+), 1 deletion(-) create mode 100644 model_manager/utils/repository/model_repository.py diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index 0d6f31c..b56feb6 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -248,6 +248,14 @@ class Training(SientiaMonitoring): ) as run_info: wrapper.store_model(name=train_params.model_name) + train_result.run_name = run_info.run_name + train_result.run_id = run_info.run_id + + train_result = self.data_manager_repository.generate_report( + train_result, + metadata=metadata, + ) + model_saved = True return { diff --git a/model_manager/utils/models/train_model_result.py b/model_manager/utils/models/train_model_result.py index e4031d3..98edc56 100644 --- a/model_manager/utils/models/train_model_result.py +++ b/model_manager/utils/models/train_model_result.py @@ -44,6 +44,9 @@ class TrainModelResult: equation: dict | None = None equation_path: str | None = None run_name: str | None = None + run_id: str | None = None report_path: str | None = None train_data_path: str | None = None test_data_path: str | None = None + + run_dir: str | None = None diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py index 807361e..90061d2 100644 --- a/model_manager/utils/repository/data_manager_repository.py +++ b/model_manager/utils/repository/data_manager_repository.py @@ -13,7 +13,9 @@ integration. Models are trained elsewhere (e.g., via SientiaModel wrappers), and this repository focuses solely on preparing data structures for them. """ +from datetime import datetime from io import BytesIO +from os import makedirs, path from typing import Any import numpy as np @@ -25,7 +27,7 @@ from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ from model_manager.sientia.metrics import mae, mse, r2 from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.models.train_model_result import TrainModelResult - +from model_manager.sientia.reports import Reports # type: ignore[import-untyped] def train_test_split(data: pd.DataFrame | pd.Series, train_size: float, random_state: int | None = None, shuffle: bool = True) -> tuple[pd.DataFrame, pd.DataFrame]: # 1. Definir a semente (seed) para reprodutibilidade @@ -320,3 +322,114 @@ class DataManagerRepository(SientiaMonitoring): ) return data + def _get_reports_directory(self) -> str: + """ + Get the absolute path to the reports directory. + + Returns: + str: Absolute path to model_manager/reports directory. + """ + # Get the directory where this file is located (model_manager/utils/repository/) + current_file_dir = path.dirname(path.abspath(__file__)) + # Navigate up to model_manager/ and then to reports/ + model_manager_dir = path.dirname(path.dirname(current_file_dir)) + reports_dir = path.join(model_manager_dir, 'reports') + return reports_dir + + def _create_run_directory(self, base_path: str, run_name: str, metadata: dict[str, Any] | None = None) -> str: + """ + Creates a directory inside the 'reports' folder with the run name and a timestamp. + + Uses microsecond precision in timestamp to minimize collision probability + in high-concurrency scenarios. + + Args: + base_path (str): The path to the 'reports' folder. + run_name (str): The name of the run. + + Returns: + str: The path to the created directory. + + Raises: + PermissionError: If there are insufficient permissions to create the directory. + OSError: If directory creation fails for any other reason. + """ + # Use microsecond precision to reduce collision probability + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f') + run_dir = path.join(base_path, 'temp', f'{run_name}_{timestamp}') + + try: + makedirs(run_dir, exist_ok=True) + return run_dir + except PermissionError as e: + error_msg = f'Permission denied when creating directory: {run_dir}' + self.error(error_msg, metadata) + raise PermissionError(error_msg) from e + except OSError as e: + error_msg = f'Failed to create directory {run_dir}: {str(e)}' + self.error(error_msg, metadata) + raise OSError(error_msg) from e + + def generate_report( + self, data: TrainModelResult, metadata: dict[str, Any] | None = None + ) -> TrainModelResult: + """ + Generates a comprehensive report summarizing data quality, data drift, and regression analysis. + + Args: + reference_data (pd.DataFrame): The training dataset with predictions added. + current_data (pd.DataFrame): The testing dataset with predictions added. + data: The training model result containing the datasets, model, and parameters. + + Returns: + The updated result object with paths to the generated report and data files. + + Raises: + ValueError: If data conversion to float64 fails or DataFrames are invalid. + PermissionError: If there are insufficient permissions to write files. + OSError: If file writing fails for any other reason. + """ + + if data.run_name is None: + raise ValueError('run_name is not set, cannot generate report') + + # Convert data to float64 for report generation + # This may raise ValueError if data contains non-numeric values + reference_data = data.train_data + current_data = data.val_data + reference_data_float = reference_data.astype(np.float64) + current_data_float = current_data.astype(np.float64) + + + # Initialize report generator + data.run_dir = self._create_run_directory(self._get_reports_directory(), data.run_name) + report = Reports( + reference_data=reference_data_float, + current_data=current_data_float, + base_path=data.run_dir, + ) + + # Generate report sections + report.add_data_quality_section(columns=data.params.variable_columns + ['target']) + report.add_data_drift_section(columns=data.params.variable_columns + ['target']) + report.add_regression_section() + + # Save HTML report + data.report_path = path.join(data.run_dir, 'report.html') + report.save_all_sections_html(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) + + # 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) + + # Save equation as JSON + if data.equation is not None and data.params.model_type == 'linear_regression': + data.equation_path = path.join(data.run_dir, 'model_equation.json') + with open(data.equation_path, 'w', encoding='utf-8') as f: + json.dump(data.equation, f, indent=2, ensure_ascii=False) + + return data diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py new file mode 100644 index 0000000..dcd1ba8 --- /dev/null +++ b/model_manager/utils/repository/model_repository.py @@ -0,0 +1,471 @@ +""" +MLFlow Repository + +This module contains the MLFlowRepository class, which is responsible for +handling model training artifacts and MLFlow operations for the Model Manager system. + +It includes methods for generating training reports, managing artifacts, +and logging model runs to MLFlow. + +""" + +import json +import os +import shutil +import warnings +from datetime import datetime +from os import makedirs, path + +import numpy as np +import pandas as pd +from sientia_do.observability.logger import Logger + +from model_manager.sientia.model_serving import ModelServing # type: ignore[import-untyped] +from model_manager.sientia.reports import Reports # type: ignore[import-untyped] +from model_manager.utils.models.train_model_result import TrainModelResult + +# Suppress sklearn FutureWarning about 'squared' deprecation without changing business logic +warnings.filterwarnings('ignore', category=FutureWarning, message=".*'squared' is deprecated.*") + + +class ModelRepository: + def __init__(self, url, username, password, logger: Logger): + self.model_serving = ModelServing(tracking_uri=url, username=username, password=password) + self.logger = logger + self.logger.info(f'MLFlow client initialized at {url}') + + 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 + train_result.run_name = self._get_next_run_name(experiment_name) + train_result = self._generate_artifacts(train_result) + self._save_run(train_result) + + self.logger.info( + f'Model saved successfully - experiment run id: {train_result.params.experiment_run_id}, ' + f'experiment name: {experiment_name}, ' + f'run name: {train_result.run_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 + + 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. + + Args: + experiment_name (str): The name of the experiment for which the next run name is being generated. + + Returns: + str: A unique run name in the format "-". + """ + runs = self.model_serving.search_runs_by_name( + experiment_names=[experiment_name], order_by=['start_time desc'] + ) + + next_run_number = len(runs) + 1 + return f'{experiment_name}-{next_run_number}' + + def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult: + """ + Generates and organizes artifacts related to the training process, such as reports and data files. + + Args: + data: The training model result containing the datasets, model, and parameters. + + Returns: + The updated result object with paths to the generated artifacts. + + Raises: + FileNotFoundError: If the reports directory or header.html file does not exist. + ValueError: If run_name is not set. + """ + # Validate that run_name is set + if not data.run_name: + error_msg = 'run_name must be set before generating artifacts' + self.logger.error(error_msg) + raise ValueError(error_msg) + + reference_data, current_data = self._init_artifacts_data(data) + base_path = self._get_reports_directory() + + # Validate that reports directory exists + if not path.exists(base_path): + error_msg = f'Reports directory does not exist: {base_path}' + self.logger.error(error_msg) + raise FileNotFoundError(error_msg) + + data.run_dir = self._create_run_directory(base_path, data.run_name) + header_file_path = path.join(base_path, 'header.html') + + # Validate that header.html exists + if not path.exists(header_file_path): + error_msg = f'Header file does not exist: {header_file_path}' + self.logger.error(error_msg) + raise FileNotFoundError(error_msg) + + 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): + """ + Logs the details of a machine learning run, including parameters, metrics, models, and artifacts, + to the Sientia tracking system. + + Args: + data: The training model result containing the datasets, model, parameters, + and evaluation metrics. + + Raises: + ValueError: If required metrics or artifacts are missing. + Exception: If MLflow logging fails for any reason. + """ + # Validate that required artifacts exist before attempting to log + if not data.report_path or not path.exists(data.report_path): + error_msg = f'Report file does not exist: {data.report_path}' + self.logger.error(error_msg) + raise ValueError(error_msg) + + if not data.train_data_path or not path.exists(data.train_data_path): + error_msg = f'Training data file does not exist: {data.train_data_path}' + self.logger.error(error_msg) + raise ValueError(error_msg) + + if not data.test_data_path or not path.exists(data.test_data_path): + error_msg = f'Test data file does not exist: {data.test_data_path}' + self.logger.error(error_msg) + raise ValueError(error_msg) + + # Validate that metrics are present + if data.mse_val is None or data.r2_val is None or data.mae_val is None: + error_msg = 'One or more metrics (MSE, R2, MAE) are None' + self.logger.error(error_msg) + raise ValueError(error_msg) + + # Prepare parameters + interval_strs = [ + (str(interval[0]), str(interval[1])) + for interval in (data.params.removed_intervals or []) + ] + + # Set experiment and create run + self.model_serving.set_experiment(data.params.experiment_name) + + 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_name', data.params.model_name) + self.model_serving.log_param( + 'models_params', + {'degree': data.params.degree, 'interaction_only': data.params.interaction_only}, + ) + 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('nan_treatment', data.params.nan_treatment) + self.model_serving.log_param('lag_train', data.params.lag_train) + self.model_serving.log_param('lag_transform', data.params.lag_val) + static_threshold_value = None + if data.params.rem_static_win: + static_threshold_value = ( + data.params.static_threshold if data.params.static_threshold is not None else 1 + ) + self.model_serving.log_param('static_threshold', static_threshold_value) + self.model_serving.log_param('lower_limits', data.params.low_lim) + self.model_serving.log_param('upper_limits', data.params.upp_lim) + self.model_serving.log_param('scaler_name', data.params.scaler_name) + self.model_serving.log_param('scaler_params', data.scaler_dict) + self.model_serving.log_param('include_ar', data.params.include_ar) + self.model_serving.log_param('train_size', round(data.params.train_size / 100, 2)) + self.model_serving.log_param('test_size', round(1 - (data.params.train_size / 100), 2)) + self.model_serving.log_param('start_date', data.params.start_date) + self.model_serving.log_param('end_date', data.params.end_date) + self.model_serving.log_param('removed_intervals', interval_strs) + self.model_serving.log_param('retrain', False) + self.model_serving.log_param('support_filters', data.params.support_filters) + + # 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 models + self.model_serving.log_model(data.process_data, 'data_model') + self.model_serving.log_model(data.regr, 'prediction_model') + + # 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 equation artifact if available + if data.equation_path and path.exists(data.equation_path): + self.model_serving.log_artifact(data.equation_path) + + def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]: + """ + Prepares the reference and current datasets for artifact generation. + + Args: + data: The training model result containing the datasets and model. + + Returns: + tuple: A tuple containing: + - reference_data: The training dataset with predictions added. + - current_data: The testing dataset with predictions added. + + Raises: + ValueError: If training or test datasets are empty or invalid. + AttributeError: If required attributes are missing from the data object. + """ + # Validate that required DataFrames are not empty + # Note: x_train, y_train, x_test, y_test, and regr are required fields in TrainModelResult + # so we only check if they are empty, not None + if data.x_train.empty: + error_msg = 'Training features (x_train) are empty' + self.logger.error(error_msg) + raise ValueError(error_msg) + + if data.y_train.empty: + error_msg = 'Training target (y_train) is empty' + self.logger.error(error_msg) + raise ValueError(error_msg) + + if data.x_test.empty: + error_msg = 'Test features (x_test) are empty' + self.logger.error(error_msg) + raise ValueError(error_msg) + + if data.y_test.empty: + error_msg = 'Test target (y_test) is empty' + self.logger.error(error_msg) + raise ValueError(error_msg) + + # Validate that predictions exist (y_pred is optional, so check for None) + if data.y_pred is None: + error_msg = 'Test predictions (y_pred) are None' + self.logger.error(error_msg) + raise ValueError(error_msg) + + if data.y_train_pred is None: + error_msg = 'Training predictions (y_train_pred) are None' + self.logger.error(error_msg) + raise ValueError(error_msg) + + # Prepare reference data (training set) + reference_data = pd.concat([data.x_train, data.y_train], axis=1) + reference_data = reference_data.rename(columns={data.params.target_variable: 'target'}) + # Use pre-calculated predictions (calculated before denormalization to avoid overflow) + reference_data['prediction'] = data.y_train_pred + + # Prepare current data (test set) + current_data = pd.concat([data.x_test, data.y_test], axis=1) + current_data = current_data.rename(columns={data.params.target_variable: 'target'}) + current_data['prediction'] = data.y_pred + + return reference_data, current_data + + def _create_run_directory(self, base_path: str, run_name: str) -> str: + """ + Creates a directory inside the 'reports' folder with the run name and a timestamp. + + Uses microsecond precision in timestamp to minimize collision probability + in high-concurrency scenarios. + + Args: + base_path (str): The path to the 'reports' folder. + run_name (str): The name of the run. + + Returns: + str: The path to the created directory. + + Raises: + PermissionError: If there are insufficient permissions to create the directory. + OSError: If directory creation fails for any other reason. + """ + # Use microsecond precision to reduce collision probability + timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f') + run_dir = path.join(base_path, 'temp', f'{run_name}_{timestamp}') + + try: + makedirs(run_dir, exist_ok=True) + return run_dir + except PermissionError as e: + error_msg = f'Permission denied when creating directory: {run_dir}' + self.logger.error(error_msg) + raise PermissionError(error_msg) from e + except OSError as e: + error_msg = f'Failed to create directory {run_dir}: {str(e)}' + self.logger.error(error_msg) + raise OSError(error_msg) from e + + def _setup_run_directory(self, run_dir: str, header_file_path: str): + """ + Creates empty files and copies a header file into the specified run directory. + + Note: Lock removed as each run has its own unique directory, so no synchronization + is needed between different runs. File operations within the same directory are + atomic at the OS level. + + Args: + run_dir (str): The path to the run directory where the files will be created. + header_file_path (str): The path to the header.html file to be copied. + + Raises: + FileNotFoundError: If the header file does not exist. + PermissionError: If there are insufficient permissions to create files. + OSError: If file creation or copying fails for any other reason. + """ + empty_files = ['data_drift.html', 'data_quality.html', 'regression.html'] + + try: + # Create empty placeholder files + for file_name in empty_files: + file_path = path.join(run_dir, file_name) + with open(file_path, 'w'): + pass # Create empty file + + # Copy header file to run directory + header_dest = path.join(run_dir, 'header.html') + shutil.copy(header_file_path, header_dest) + except FileNotFoundError as e: + error_msg = f'Header file not found: {header_file_path}' + self.logger.error(error_msg) + raise FileNotFoundError(error_msg) from e + except PermissionError as e: + error_msg = f'Permission denied when setting up directory: {run_dir}' + self.logger.error(error_msg) + raise PermissionError(error_msg) from e + except OSError as e: + error_msg = f'Failed to setup run directory {run_dir}: {str(e)}' + self.logger.error(error_msg) + raise OSError(error_msg) from e + + def _generate_report( + self, reference_data: pd.DataFrame, current_data: pd.DataFrame, data: TrainModelResult + ) -> TrainModelResult: + """ + Generates a comprehensive report summarizing data quality, data drift, and regression analysis. + + Args: + reference_data (pd.DataFrame): The training dataset with predictions added. + current_data (pd.DataFrame): The testing dataset with predictions added. + data: The training model result containing the datasets, model, and parameters. + + Returns: + The updated result object with paths to the generated report and data files. + + Raises: + ValueError: If data conversion to float64 fails or DataFrames are invalid. + PermissionError: If there are insufficient permissions to write files. + OSError: If file writing fails for any other reason. + """ + try: + # Convert data to float64 for report generation + # This may raise ValueError if data contains non-numeric values + reference_data_float = reference_data.astype(np.float64) + current_data_float = current_data.astype(np.float64) + + # Initialize report generator + report = Reports( + reference_data=reference_data_float, + current_data=current_data_float, + base_path=data.run_dir, + ) + + # Generate report sections + report.add_data_quality_section(columns=data.params.variable_columns + ['target']) + report.add_data_drift_section(columns=data.params.variable_columns + ['target']) + report.add_regression_section() + + # Validate that run_dir is set (should be set by _create_run_directory) + if not data.run_dir: + error_msg = 'run_dir is not set after directory creation' + self.logger.error(error_msg) + raise ValueError(error_msg) + + # Save HTML report + data.report_path = path.join(data.run_dir, 'report.html') + report.save_all_sections_html(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) + + # 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) + + # Save equation as JSON + if data.equation is not None: + data.equation_path = path.join(data.run_dir, 'model_equation.json') + with open(data.equation_path, 'w', encoding='utf-8') as f: + json.dump(data.equation, f, indent=2, ensure_ascii=False) + + 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) + raise ValueError(error_msg) from e + except PermissionError as e: + error_msg = f'Permission denied when writing report files to: {data.run_dir}' + self.logger.error(error_msg) + raise PermissionError(error_msg) from e + except OSError as e: + error_msg = f'Failed to generate report in {data.run_dir}: {str(e)}' + self.logger.error(error_msg) + raise OSError(error_msg) from e + + def _get_reports_directory(self) -> str: + """ + Get the absolute path to the reports directory. + + Returns: + str: Absolute path to model_manager/reports directory. + """ + # Get the directory where this file is located (model_manager/utils/repository/) + current_file_dir = path.dirname(path.abspath(__file__)) + # Navigate up to model_manager/ and then to reports/ + model_manager_dir = path.dirname(path.dirname(current_file_dir)) + reports_dir = path.join(model_manager_dir, 'reports') + return reports_dir