SIENTIAPDE-1249: Implement data models for Model Manager and add unit tests. This commit introduces data transfer objects (DTOs) and model classes for experiment status, training parameters, and training results, along with corresponding unit tests to ensure their correct behavior.

model_manager/utils/models/__init__.py

@@ -0,0 +1,16 @@
+"""
+Models and DTOs for the Model Manager system.
+
+This module contains data transfer objects (DTOs) and model classes used
+throughout the Model Manager workflows and activities.
+"""
+
+from model_manager.utils.models.experiment_status import ExperimentStatus
+from model_manager.utils.models.train_model_params import TrainModelParams
+from model_manager.utils.models.train_model_result import TrainModelResult
+
+__all__ = [
+    'ExperimentStatus',
+    'TrainModelParams',
+    'TrainModelResult',
+]

model_manager/utils/models/experiment_status.py

@@ -0,0 +1,32 @@
+from enum import Enum
+
+
+class ExperimentStatus(str, Enum):
+    """
+    Status values for experiment run lifecycle.
+
+    This enum defines all possible status values that an experiment run can have
+    throughout its lifecycle, from initialization through training, model saving,
+    and cleanup. These statuses are used to track progress and identify failures
+    in the training pipeline.
+
+    The status values follow the naming convention from the original Mage pipeline
+    to maintain compatibility with existing database records and monitoring systems.
+
+    Attributes:
+        MAGE_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing.
+        TRAINING_SUCCESS: Training completed successfully with model and metrics calculated.
+        TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions.
+        MLFLOW_SENT: Model successfully saved to MLFlow.
+        MLFLOW_SEND_ERROR: Model saving to MLFlow failed due to connection or serialization errors.
+        FILE_DELETED: Cleanup completed successfully with all artifacts removed.
+        FILE_DELETE_ERROR: Cleanup failed due to file system or MinIO errors.
+    """
+
+    MAGE_WAITING_PROC = 'MAGE_WAITING_PROC'
+    TRAINING_SUCCESS = 'TRAINING_SUCCESS'
+    TRAINING_ERROR = 'TRAINING_ERROR'
+    MLFLOW_SENT = 'MLFLOW_SENT'
+    MLFLOW_SEND_ERROR = 'MLFLOW_SEND_ERROR'
+    FILE_DELETED = 'FILE_DELETED'
+    FILE_DELETE_ERROR = 'FILE_DELETE_ERROR'

model_manager/utils/models/train_model_params.py

@@ -0,0 +1,166 @@
+from typing import Any
+
+
+class TrainModelParams:
+    """
+    Parameters for machine learning model training.
+
+    This class encapsulates all configuration parameters required for the training
+    pipeline, including data processing settings, model configuration, and experiment
+    tracking information. All parameters are validated upon initialization to ensure
+    data integrity and prevent runtime errors.
+
+    Attributes:
+        variable_columns (List[str]): List of variable column names to use as features.
+        lag_train (int): Number of lags to apply during training phase.
+        lag_val (int): Number of lags to apply during validation phase.
+        target_variable (str): Name of the target variable to predict.
+        rem_static_win (bool): Whether to remove static windows from data.
+        low_lim (Dict[str, float]): Dictionary of lower limits for each variable.
+        upp_lim (Dict[str, float]): Dictionary of upper limits for each variable.
+        window (int): Window size for rolling operations.
+        use_scaler (bool): Whether to use a scaler for data normalization.
+        include_ar (bool): Whether to include autoregressive variables.
+        bucket_name (str): Name of the MinIO bucket containing training data.
+        file_name (str): Name of the file in the MinIO bucket.
+        line_separator (str): Line separator used in the CSV file.
+        decimal_separator (str): Decimal separator used in the CSV file.
+        train_size (int): Percentage of data to use for training (0-100).
+        shuffle (bool): Whether to shuffle the data during train/test split.
+        experiment_run_id (int): Unique identifier for the experiment run.
+        experiment_name (str): Name of the experiment for tracking.
+        experiment_description (str): Description of the experiment.
+        removed_intervals (list): List of time intervals to remove from the data.
+    """
+
+    def __init__(
+        self,
+        variable_columns: Any | None,
+        lag_train: Any | None,
+        lag_val: Any | None,
+        target_variable: Any | None,
+        rem_static_win: Any | None,
+        low_lim: Any | None,
+        upp_lim: Any | None,
+        window: Any | None,
+        use_scaler: Any | None,
+        include_ar: Any | None,
+        bucket_name: Any | None,
+        file_name: Any | None,
+        line_separator: Any | None,
+        decimal_separator: Any | None,
+        train_size: Any | None,
+        shuffle: Any | None,
+        experiment_run_id: Any | None,
+        experiment_name: Any | None,
+        experiment_description: Any | None,
+        removed_intervals: Any | None,
+    ):
+        """
+        Initialize the TrainModelParams object and validate all input parameters.
+
+        This constructor validates that all required parameters are provided and
+        have the correct types. It raises descriptive errors if validation fails,
+        helping to catch configuration issues early in the pipeline.
+
+        Args:
+            variable_columns (list | None): List of variable column names.
+            lag_train (int | None): Number of lags to apply during training.
+            lag_val (int | None): Number of lags to apply during validation.
+            target_variable (str | None): Name of the target variable.
+            rem_static_win (bool | None): Whether to remove static windows.
+            low_lim (dict | None): Dictionary of lower limits for variables.
+            upp_lim (dict | None): Dictionary of upper limits for variables.
+            window (int | None): Window size for rolling operations.
+            use_scaler (bool | None): Whether to use a scaler for normalization.
+            include_ar (bool | None): Whether to include autoregressive variables.
+            bucket_name (str | None): Name of the MinIO bucket.
+            file_name (str | None): Name of the file in the MinIO bucket.
+            line_separator (str | None): Line separator used in the file.
+            decimal_separator (str | None): Decimal separator used in the file.
+            train_size (int | None): Percentage of data to use for training.
+            shuffle (bool | None): Whether to shuffle the data during splitting.
+            experiment_run_id (int | None): ID of the experiment run.
+            experiment_name (str | None): Name of the experiment.
+            experiment_description (str | None): Description of the experiment.
+            removed_intervals (list | None): List of intervals to remove from the data.
+
+        Raises:
+            ValueError: If any required parameter is None.
+            TypeError: If any parameter is not of the expected type.
+        """
+        self.variable_columns: list[str] = self._check_none(
+            variable_columns, list, 'variable_columns'
+        )
+        self.lag_train: int = self._check_none(lag_train, int, 'lag_train')
+        self.lag_val: int = self._check_none(lag_val, int, 'lag_val')
+        self.target_variable: str = self._check_none(target_variable, str, 'target_variable')
+        self.rem_static_win: bool = self._check_none(rem_static_win, bool, 'rem_static_win')
+        self.low_lim: dict[str, float] = self._check_none(low_lim, dict, 'low_lim')
+        self.upp_lim: dict[str, float] = self._check_none(upp_lim, dict, 'upp_lim')
+        self.window: int = self._check_none(window, int, 'window')
+        self.use_scaler: bool = self._check_none(use_scaler, bool, 'use_scaler')
+        self.include_ar: bool = self._check_none(include_ar, bool, 'include_ar')
+        self.bucket_name: str = self._check_none(bucket_name, str, 'bucket_name')
+        self.file_name: str = self._check_none(file_name, str, 'file_name')
+        self.line_separator: str = self._check_none(line_separator, str, 'line_separator')
+        self.decimal_separator: str = self._check_none(decimal_separator, str, 'decimal_separator')
+        self.train_size: int = self._check_none(train_size, int, 'train_size')
+        self.shuffle: bool = self._check_none(shuffle, bool, 'shuffle')
+        self.experiment_run_id: int = self._check_none(experiment_run_id, int, 'experiment_run_id')
+        self.experiment_name: str = self._check_none(experiment_name, str, 'experiment_name')
+        self.experiment_description: str = self._check_none(
+            experiment_description, str, 'experiment_description'
+        )
+        self.removed_intervals: list = self._check_type(
+            removed_intervals, list, 'removed_intervals'
+        )
+
+    def _check_none(self, value: Any | None, expected_type: type, field_name: str):
+        """
+        Validate that a value is not None and check its type.
+
+        This method ensures that required parameters are provided and have the
+        correct type, raising descriptive errors if validation fails.
+
+        Args:
+            value (Any | None): The value to validate.
+            expected_type (type): The expected type of the value.
+            field_name (str): The name of the field being validated (for error messages).
+
+        Returns:
+            Any: The validated value if it is not None and matches the expected type.
+
+        Raises:
+            ValueError: If the value is None.
+            TypeError: If the value is not of the expected type.
+        """
+        if value is None:
+            error = f'{field_name} is required and cannot be None.'
+            raise ValueError(error)
+
+        return self._check_type(value, expected_type, field_name)
+
+    def _check_type(self, value: Any | None, expected_type: type, field_name: str):
+        """
+        Validate that a value matches the expected type.
+
+        This method checks type compatibility and raises a descriptive error
+        if the value does not match the expected type.
+
+        Args:
+            value (Any | None): The value to validate.
+            expected_type (type): The expected type of the value.
+            field_name (str): The name of the field being validated (for error messages).
+
+        Returns:
+            Any: The validated value if it matches the expected type.
+
+        Raises:
+            TypeError: If the value is not of the expected type.
+        """
+        if value is not None and not isinstance(value, expected_type):
+            error = f'{field_name} must be of type {expected_type.__name__}, but got {type(value).__name__}.'
+            raise TypeError(error)
+
+        return value

model_manager/utils/models/train_model_result.py

@@ -0,0 +1,55 @@
+from dataclasses import dataclass
+
+import pandas as pd
+from sientia.linear_models import LinearRegressionModel
+from sientia.preprocessing import DataPreprocessor
+
+from model_manager.utils.models.train_model_params import TrainModelParams
+
+
+@dataclass
+class TrainModelResult:
+    """
+    A data container for storing the results of a machine learning training process.
+
+    This dataclass encapsulates all outputs from the training pipeline, including
+    the trained model, datasets, evaluation metrics, and paths to generated artifacts.
+    It is used to pass results between activities in the training workflow.
+
+    Attributes:
+        params (TrainModelParams): The parameters used to train the model.
+        process_data (DataPreprocessor): The data preprocessor object used to process the input data.
+        X_train (pd.DataFrame): The training dataset features.
+        X_test (pd.DataFrame): The testing dataset features.
+        y_train (pd.DataFrame): The training dataset target values.
+        y_test (pd.DataFrame): The testing dataset target values.
+        regr (LinearRegressionModel): The trained linear regression model.
+        scaler_dict (dict): A dictionary containing the scalers used to scale the features and target values.
+        y_pred (Optional[pd.Series]): The predicted target values for the testing dataset. Default is None.
+        mse_val (Optional[float]): The Mean Squared Error (MSE) of the predictions. Default is None.
+        mae_val (Optional[float]): The Mean Absolute Error (MAE) of the predictions. Default is None.
+        r2_val (Optional[float]): The R-squared (R²) value of the predictions. Default is None.
+        run_name (Optional[str]): The name of the MLFlow run. Default is None.
+        report_path (Optional[str]): The path to the generated HTML report file. Default is None.
+        train_data_path (Optional[str]): The path to the training dataset CSV file. Default is None.
+        test_data_path (Optional[str]): The path to the testing dataset CSV file. Default is None.
+        run_dir (Optional[str]): The path to the run directory containing all artifacts. Default is None.
+    """
+
+    params: TrainModelParams
+    process_data: DataPreprocessor
+    X_train: pd.DataFrame
+    X_test: pd.DataFrame
+    y_train: pd.DataFrame
+    y_test: pd.DataFrame
+    regr: LinearRegressionModel
+    scaler_dict: dict
+    y_pred: pd.Series | None = None
+    mse_val: float | None = None
+    mae_val: float | None = None
+    r2_val: float | None = None
+    run_name: 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