feat: enhance training workflow with model metadata loading and refactor data handling

- Introduced a new activity to load model metadata from the model store.
- Refactored training logic to utilize new model metadata and improved parameter handling.
- Updated the `TrainModelParams` class to include additional fields for model configuration.
- Replaced deprecated utility functions with a custom train-test split implementation.
- Removed unused utility functions and cleaned up the data manager repository.
- Adjusted experiment tracking to include model-specific metadata in notifications.

model_manager/activities/experiment_tracking.py

@@ -246,7 +246,7 @@ class ExperimentTracking(Postgres):
             ValueError: If required parameters are missing for the update type
             RuntimeError: If update operation fails
         """
-        metadata = input_data.get('metadata', {})
+        metadata = input_data.get('metadata')
         experiment_run_id = input_data['experiment_run_id']
         update_type = input_data['update_type']
         status = input_data.get('status')
@@ -272,8 +272,8 @@ class ExperimentTracking(Postgres):
             error_msg = f'Error updating experiment run - ID: {experiment_run_id}, Status: {status}, Error: {str(e)}'
             trace = traceback.format_exc()
 
-            self.send_notification(
-                metadata=metadata,
+            await self.send_notification_async(
+                metadata=metadata or {},
                 notification_id='UPDATE_EXPERIMENT_RUN_ERROR',
                 message=error_msg,
                 block='update_experiment_run',

model_manager/activities/training.py

@@ -21,6 +21,7 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.repository.minio_repository import MinioRepository
     from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
     from sientia_model.model_repository.plugin_store import PluginStore
+    from sientia_model.wrappers.sientia_model import SientiaModel
 
     from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
     from model_manager.utils.exceptions import ModelTrainingError
@@ -60,6 +61,46 @@ class Training(SientiaMonitoring):
         self.plugin_store = plugin_store
         self.minio_repository = minio_repository
 
+    @activity.defn(name='load_model_metadata')
+    async def load_model_metadata(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Load model metadata/schemas from the model store.
+
+        This activity is responsible for fetching model metadata/schemas from the
+        model store index and extracting a serializable `model_metadata` dict that
+        `TrainModelParams.validate_business_rules()` depends on.
+
+        Args:
+            input_data: Workflow input at the same level as `validate_train_params`,
+                including at least `model_name` and the fields required by
+                `TrainModelParams.from_dict` to build wrapper kwargs.
+
+        Return:
+            dict[str, Any]: Updated `input_data` containing `input_data['model_metadata']`.
+        """
+        metadata = input_data.get('metadata', {})
+
+        try:
+            train_params = TrainModelParams.from_dict(input_data)
+            model_metadata = self.plugin_store.get_model_index(
+                model_name=train_params.model_name,
+                metadata=metadata,
+            )
+
+            train_params.model_metadata = model_metadata
+            return train_params.to_dict()
+        except Exception as exc:
+            trace = traceback.format_exc()
+            await self.send_notification_async(
+                metadata=metadata,
+                notification_id='LOAD_MODEL_METADATA_ERROR',
+                message=f'Error loading model metadata: {str(exc)}',
+                block='load_model_metadata',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+            raise
+        
     @activity.defn(name='validate_train_params')
     async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
         """
@@ -81,10 +122,9 @@ class Training(SientiaMonitoring):
             ValueError, TypeError, KeyError: If validation fails (after sending notification)
         """
         metadata = input_data.get('metadata', {})
-        metrics_status = 'success'
-
         try:
             train_params = TrainModelParams.from_dict(input_data)
+            
             train_params.validate_business_rules()
 
             self.info(
@@ -96,11 +136,10 @@ class Training(SientiaMonitoring):
 
             return train_params
         except (ValueError, TypeError, KeyError) as e:
-            metrics_status = 'error'
             error_msg = f'Error validating training parameters: {str(e)}'
             trace = traceback.format_exc()
 
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='VALIDATE_TRAIN_PARAMS_ERROR',
                 message=error_msg,
@@ -109,13 +148,6 @@ class Training(SientiaMonitoring):
                 attachment_content=trace,
             )
             raise
-        finally:
-            await self._emit_metrics(
-                metadata=metadata,
-                metrics_status=metrics_status,
-                activity_name='validate_train_params',
-                emit_workflow_metric=(metrics_status == 'error'),
-            )
 
     @activity.defn(name='train_model')
     async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]:
@@ -146,10 +178,8 @@ class Training(SientiaMonitoring):
         if isinstance(train_params, dict):
             train_params = TrainModelParams.from_dict(train_params)
 
-
         model_trained = False
         model_saved = False
-        metrics_status = 'success'
 
         try:
             # Download training file bytes from MinIO
@@ -161,7 +191,7 @@ class Training(SientiaMonitoring):
 
             # Download optional validation file bytes from the same bucket
             val_bytes: bytes | None = None
-            validation_name = getattr(train_params, 'validation_file_name', None)
+            validation_name = train_params.val_file_name
             if validation_name is not None:
                 val_bytes = await self.minio_repository.download_file(
                     object_name=validation_name,
@@ -179,34 +209,35 @@ class Training(SientiaMonitoring):
             wrapper = await self.plugin_store.get_model(
                 model_name=train_params.model_name,
                 force_download=False,
-                opt_params={},
-                model_kwargs={},
-                data_model_kwargs={},
-                metadata=metadata,
+                opt_params=train_params.opt_params or {},
+                model_kwargs=train_params.model_kwargs or {},
+                data_model_kwargs=train_params.data_model_kwargs or {},
+                metadata=metadata
             )
 
-            train_df = pd.concat([train_result.x_train, train_result.y_train], axis=1)
-            val_df = pd.concat([train_result.x_test, train_result.y_test], axis=1)
+            train_data = train_result.train_data
+            val_data = train_result.val_data
 
             wrapper.train(
-                train_data=train_df,
-                val_data=val_df,
+                train_data=train_data,
+                val_data=val_data,
                 target=train_params.target_variable,
             )
 
             # Generate predictions using the trained wrapper
-            transformed_train, _ = wrapper.transform(train_df)
-            transformed_val, _ = wrapper.transform(val_df)
+            transformed_train, _ = wrapper.transform(train_data)
+            transformed_val, _ = wrapper.transform(val_data)
 
             y_train_pred_df, _ = wrapper.predict({}, transformed_train)
             y_val_pred_df, _ = wrapper.predict({}, transformed_val)
 
-            # Use the first column of the prediction DataFrame as the target prediction
-            train_result.y_train_pred = y_train_pred_df.iloc[:, 0]
-            train_result.y_pred = y_val_pred_df.iloc[:, 0]
+            y_train_pred_df.sort_index(inplace=True, ascending=False)
+            y_val_pred_df.sort_index(inplace=True, ascending=False)
+
+            train_result.y_train_pred = y_train_pred_df
+            train_result.y_pred = y_val_pred_df
 
             train_result = self.data_manager_repository.compute_regression_metrics(
-                train_params,
                 train_result,
             )
 
@@ -215,7 +246,7 @@ class Training(SientiaMonitoring):
             async with self.mlflow_repository.start_run(
                 model_name=train_params.model_name,
                 run_name=None,
-                experiment_name=train_params.experiment_name,
+                experiment_name=f'{train_params.model_name}_experiment',
                 tags=None,
                 metadata=metadata,
             ) as run_info:
@@ -224,7 +255,8 @@ class Training(SientiaMonitoring):
             model_saved = True
 
             return {
-                'run_name': run_info.run_name or run_info.run_id,
+                'run_name': run_info.run_name,
+                'run_id': run_info.run_id,
             }
         except Exception as e:  # noqa: BLE001
             metrics_status = 'error'
@@ -237,8 +269,8 @@ class Training(SientiaMonitoring):
 
             trace = traceback.format_exc()
 
-            self.send_notification(
-                metadata=metadata,
+            await self.send_notification_async(
+                metadata=metadata or {},
                 notification_id='TRAIN_MODEL_ERROR',
                 message=error_msg,
                 block='train_model',
@@ -250,13 +282,6 @@ class Training(SientiaMonitoring):
                 model_trained=model_trained,
                 model_saved=model_saved,
             ) from e
-        finally:
-            await self._emit_metrics(
-                metadata=metadata,
-                metrics_status=metrics_status,
-                activity_name='train_model',
-                emit_workflow_metric=(metrics_status == 'error'),
-            )
 
     @activity.defn(name='cleanup_resources')
     async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
@@ -284,8 +309,6 @@ class Training(SientiaMonitoring):
                 metadata=metadata,
             )
         except Exception as e:  # noqa: BLE001
-            metrics_status = 'error'
-
             error_msg = (
                 'Error cleaning up resources - '
                 f'File: {bucket_name}/{file_name}, Error: {str(e)}'
@@ -303,44 +326,3 @@ class Training(SientiaMonitoring):
             )
 
             raise
-        finally:
-            await self._emit_metrics(
-                metadata=metadata,
-                metrics_status=metrics_status,
-                activity_name='cleanup_resources',
-                emit_workflow_metric=True,
-            )
-
-    async def _emit_metrics(
-        self,
-        metadata: dict[str, Any],
-        metrics_status: str,
-        activity_name: str,
-        emit_workflow_metric: bool,
-    ) -> None:
-        """
-        Emit workflow and activity execution metrics.
-
-        Args:
-            metadata: Activity metadata containing pod_id and workflow_name
-            metrics_status: Execution status ('success' or 'error')
-            activity_name: Name of the activity being executed
-        """
-        if emit_workflow_metric:
-            await self.emit_metric(
-                metric_object=WORKFLOW_EXECUTION_TOTAL,
-                tags={
-                    'pod_id': metadata.get('pod_id'),
-                    'workflow_name': metadata.get('workflow_name'),
-                    'status': metrics_status,
-                },
-            )
-
-        await self.emit_metric(
-            metric_object=ACTIVITY_EXECUTION_TOTAL,
-            tags={
-                'pod_id': metadata.get('pod_id'),
-                'activity_name': activity_name,
-                'status': metrics_status,
-            },
-        )

model_manager/sientia/utils.py

@@ -1,41 +0,0 @@
-from typing import Any
-
-from numpy.typing import ArrayLike
-from sklearn.model_selection import train_test_split
-
-
-def split_train_test(
-    *data: Any,
-    test_size: float | None = None,
-    train_size: float | None = None,
-    random_state: int | None = None,
-    shuffle: bool = True,
-    stratify: ArrayLike | None = None,
-) -> tuple[Any, Any, Any, Any]:
-    """
-    Split arrays or matrices into random train and test subsets.
-
-    Wrapper for sklearn.model_selection.train_test_split.
-
-    Args:
-        *data: data to be split.
-        test_size: size of test subset.
-        train_size: size of train subset.
-        random_state: Seed applied to the data before applying the split.
-        shuffle: Whether or not to shuffle the data before splitting.
-        stratify: If not None, data is split in a stratified fashion, using this as the class labels.
-
-    Returns:
-        X_train, X_test, y_train, y_test
-
-    Thread-safe: This function is stateless and thread-safe.
-    """
-    X_train, X_test, y_train, y_test = train_test_split(
-        *data,
-        test_size=test_size,
-        train_size=train_size,
-        random_state=random_state,
-        shuffle=shuffle,
-        stratify=stratify,
-    )
-    return X_train, X_test, y_train, y_test

model_manager/utils/models/train_model_params.py

@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from typing import Any
 
+from jsonschema import Draft202012Validator, ValidationError  # type: ignore[import-untyped]
+
 from model_manager.sientia.models import validate_frontend_date_format
 
 # Model name constants
@@ -23,18 +25,9 @@ class TrainModelParams:
 
     Attributes:
         variable_columns (list[str]): List of variable column names to use as features.
-        lag_train (dict[str, int]): Dictionary of lags per variable for training phase.
-        lag_val (dict[str, int]): Dictionary of lags per variable for 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 training file in the MinIO bucket.
-        validation_file_name (str | None): Optional name of the validation file in the same MinIO bucket as the training file.
+        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.
         date_column (str | None): Name of the date/time column. If set with date_format, the column is parsed as datetime.
@@ -42,50 +35,41 @@ class TrainModelParams:
         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.
-        removed_intervals (list): List of time intervals to remove from the data.
         model_name (str): Name of the model type ('Linear Regression' or 'Polynomial Regression').
-        degree (int): Degree of polynomial features (1 for linear, >1 for polynomial).
-        interaction_only (bool): If True, only interaction features are produced for polynomial.
-        nan_treatment (str): Treatment for NaN values ('drop' or 'linear interpolation').
-        start_date (str | None): Start date for filtering data.
-        end_date (str | None): End date for filtering data.
-        scaler_name (str): Name of the scaler to use ('Standard Scaler' or 'None').
-        support_filters (dict): Custom support filters per variable.
-        static_threshold (int | None): Threshold for static window removal (1-1000). Only used when rem_static_win is True.
+        val_file_name (str | None): Name of the validation file in the MinIO bucket.
+        data_model_kwargs (dict | None): Keyword arguments for the data model.
+        model_kwargs (dict | None): Keyword arguments for the model.
+        opt_params (dict | None): Keyword optimazation arguments for the wrapper.
+        model_type (str): Type of the model to use (ex.: 'Linear Regression', 'XGBoost').
     """
 
+    # Old Parameters (keep)
+
     variable_columns: list[str]
-    lag_train: dict[str, int]
-    lag_val: dict[str, int]
     target_variable: str
-    rem_static_win: bool
-    low_lim: dict[str, float]
-    upp_lim: dict[str, float]
-    window: int
-    use_scaler: bool
-    include_ar: bool
     bucket_name: str
     file_name: str
-    validation_file_name: str | None
     line_separator: str
     decimal_separator: str
     date_column: str | None
     date_format: str | None
     train_size: int
     shuffle: bool
+    random_state: int
     experiment_run_id: int
-    experiment_name: str
-    removed_intervals: list
     model_name: str
-    degree: int
-    interaction_only: bool
-    nan_treatment: str
-    start_date: str | None
-    end_date: str | None
-    scaler_name: str
-    support_filters: dict
-    static_threshold: int | None
+    experiment_name: str
+
+    # New Parameters
+    val_file_name: str | None
+    data_model_kwargs: dict | None # Removed params used in DataPreprocessor here
+    model_kwargs: dict | None # Removed params used in Linear Regression Model here
+    opt_params: dict | None
+    model_type: str
+    model_id: str | None
+
+    # Context Parameters
+    model_metadata: dict | None
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> 'TrainModelParams':
@@ -98,8 +82,8 @@ class TrainModelParams:
         workflow input data.
 
         Args:
-            data: Dictionary containing training parameters with keys matching
-                  the attribute names (variable_columns, lag_train, etc.)
+            data: Dictionary containing training parameters with keys matching the
+                  attribute names (e.g. variable_columns, data_model_kwargs, model_kwargs, opt_params).
 
         Returns:
             TrainModelParams: Validated instance with all fields populated
@@ -109,53 +93,43 @@ class TrainModelParams:
             TypeError: If any field has an incorrect type
             KeyError: If any required key is missing from the dictionary
         """
+        # `from_dict()` should only build the "raw" object from the input dict.
+        # Semantic validation and defaults must be handled by `validate_business_rules()`
+        # (using `model_metadata` JSON Schemas).
+
+        model_name = cls._check_none(data.get('model_name'), str, 'model_name')
+
         return cls(
-            variable_columns=cls._check_none(
-                data.get('variable_columns'), list, 'variable_columns'
-            ),
-            lag_train=cls._check_none(data.get('lag_train'), dict, 'lag_train'),
-            lag_val=cls._check_none(data.get('lag_val'), dict, 'lag_val'),
+            variable_columns=cls._check_none(data.get('variable_columns'), list, 'variable_columns'),
             target_variable=cls._check_none(data.get('target_variable'), str, 'target_variable'),
-            rem_static_win=cls._check_none(data.get('rem_static_win'), bool, 'rem_static_win'),
-            low_lim=cls._check_none(data.get('low_lim'), dict, 'low_lim'),
-            upp_lim=cls._check_none(data.get('upp_lim'), dict, 'upp_lim'),
-            window=cls._check_none(data.get('window'), int, 'window'),
-            use_scaler=cls._check_none(data.get('use_scaler'), bool, 'use_scaler'),
-            include_ar=cls._check_none(data.get('include_ar'), bool, 'include_ar'),
             bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'),
             file_name=cls._check_none(data.get('file_name'), str, 'file_name'),
-            validation_file_name=cls._check_type(
-                data.get('validation_file_name'), str, 'validation_file_name'
-            ),
             line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'),
-            decimal_separator=cls._check_none(
-                data.get('decimal_separator'), str, 'decimal_separator'
-            ),
+            decimal_separator=cls._check_none(data.get('decimal_separator'), str, 'decimal_separator'),
             date_column=data.get('date_column'),
             date_format=data.get('date_format'),
             train_size=cls._check_none(data.get('train_size'), int, 'train_size'),
             shuffle=cls._check_none(data.get('shuffle'), bool, 'shuffle'),
-            experiment_run_id=cls._check_none(
-                data.get('experiment_run_id'), int, 'experiment_run_id'
-            ),
-            experiment_name=cls._check_none(data.get('experiment_name'), str, 'experiment_name'),
-            removed_intervals=cls._check_type(
-                data.get('removed_intervals'), list, 'removed_intervals'
-            ),
-            model_name=cls._check_none(data.get('model_name'), str, 'model_name'),
-            degree=cls._check_none(data.get('degree'), int, 'degree'),
-            interaction_only=cls._check_none(
-                data.get('interaction_only'), bool, 'interaction_only'
-            ),
-            nan_treatment=cls._check_none(data.get('nan_treatment'), str, 'nan_treatment'),
-            start_date=cls._check_type(data.get('start_date'), str, 'start_date'),
-            end_date=cls._check_type(data.get('end_date'), str, 'end_date'),
-            scaler_name=cls._check_none(data.get('scaler_name'), str, 'scaler_name'),
-            support_filters=cls._check_type(data.get('support_filters'), dict, 'support_filters')
-            or {},
-            static_threshold=cls._check_type(data.get('static_threshold'), int, 'static_threshold'),
+            random_state=cls._check_none(data.get('random_state', 42), int, 'random_state'),
+            experiment_run_id=cls._check_none(data.get('experiment_run_id'), int, 'experiment_run_id'),
+            model_name=model_name,
+            experiment_name=model_name + '_experiment',
+            val_file_name=data.get('val_file_name'),
+            data_model_kwargs=cls._check_none(data.get('data_model_kwargs'), dict, 'data_model_kwargs'),
+            model_kwargs=cls._check_none(data.get('model_kwargs'), dict, 'model_kwargs'),
+            opt_params=cls._check_none(data.get('opt_params'), dict, 'opt_params'),
+            model_type=cls._check_none(data.get('model_type'), str, 'model_type'),
+            model_id=data.get('model_id'),
+            
+            model_metadata=cls._check_none(data.get('model_metadata'), dict, 'model_metadata'),
         )
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert TrainModelParams to a dictionary.
+        """
+        return self.__dict__
+
     @staticmethod
     def _check_none(value: Any | None, expected_type: type, field_name: str) -> Any:
         """
@@ -221,11 +195,8 @@ class TrainModelParams:
         """
         self._validate_numeric_ranges()
         self._validate_model_params()
-        self._validate_intervals_and_dates()
-        self._validate_limits()
         self._validate_required_strings()
         self._validate_date_format()
-        self._validate_validation_file_name()
 
     def _validate_numeric_ranges(self) -> None:
         """Validate numeric parameters are within acceptable ranges."""
@@ -235,93 +206,39 @@ class TrainModelParams:
         if not self.variable_columns:
             raise ValueError('variable_columns cannot be empty')
     
-        for var, lag in self.lag_train.items():
-            if lag < 0:
-                raise ValueError(f'lag_train for {var} must be non-negative, got {lag}')
-
-        for var, lag in self.lag_val.items():
-            if lag < 0:
-                raise ValueError(f'lag_val for {var} must be non-negative, got {lag}')
-
-        if self.window < 0:
-            raise ValueError(f'window must be non-negative, got {self.window}')
-
-        # Validate static_threshold only when rem_static_win is True and value is provided
-        if self.rem_static_win and self.static_threshold is not None:
-            if not 1 <= self.static_threshold <= 1000:
-                raise ValueError(
-                    f'static_threshold must be between 1 and 1000, got {self.static_threshold}'
-                )
 
     def _validate_model_params(self) -> None:
         """Validate model-related parameters."""
-        if self.degree < 1:
-            raise ValueError(f'degree must be at least 1, got {self.degree}')
+        if not self.model_metadata:
+            raise ValueError('model_metadata is required')
         
-        valid_nan_treatments = ['drop', 'linear interpolation', 'fill linear']
-        if self.nan_treatment not in valid_nan_treatments:
-            raise ValueError(
-                f'nan_treatment must be one of {valid_nan_treatments}, got {self.nan_treatment}'
-            )
+        schemas = self.model_metadata.get('schemas', {}).get("components", {}).get("schemas")
 
-        valid_scalers = ['Standard Scaler', 'None']
-        if self.scaler_name not in valid_scalers:
-            raise ValueError(f'scaler_name must be one of {valid_scalers}, got {self.scaler_name}')
+        if not schemas:
+            return
 
-        valid_models = [MODEL_LINEAR_REGRESSION, MODEL_POLYNOMIAL_REGRESSION]
-        if self.model_name not in valid_models:
-            raise ValueError(f'model_name must be one of {valid_models}, got {self.model_name}')
+        data_model_schema = schemas.get("data_model")
+        model_schema = schemas.get("model")
+        opt_params_schema = schemas.get("opt_params")
 
-        if self.model_name == MODEL_POLYNOMIAL_REGRESSION and self.degree < 2:
-            raise ValueError(
-                f'degree must be at least 2 for {MODEL_POLYNOMIAL_REGRESSION}, got {self.degree}'
-            )
+        if data_model_schema:
+            self._validate_model_param(data_model_schema, self.data_model_kwargs)
+        if model_schema:
+            self._validate_model_param(model_schema, self.model_kwargs)
+        if opt_params_schema:
+            self._validate_model_param(opt_params_schema, self.opt_params)
         
-        if self.model_name == MODEL_POLYNOMIAL_REGRESSION and self.scaler_name == 'None':
-            raise ValueError(
-                f'scaler_name must be set (e.g., "Standard Scaler") for {MODEL_POLYNOMIAL_REGRESSION} '
-                'to avoid numerical overflow with large feature values'
-            )
         
-        if self.model_name == MODEL_LINEAR_REGRESSION and self.degree != 1:
-            raise ValueError(f'degree must be 1 for {MODEL_LINEAR_REGRESSION}, got {self.degree}')
 
-    def _validate_intervals_and_dates(self) -> None:
-        """Validate removed_intervals format and date parameters."""
-        if self.removed_intervals:
-            for i, interval in enumerate(self.removed_intervals):
-                if not isinstance(interval, (list, tuple)):
-                    raise ValueError(
-                        f'removed_intervals[{i}] must be a list or tuple, '
-                        f'got {type(interval).__name__}'
-                    )
-                if len(interval) < 2:
-                    raise ValueError(
-                        f'removed_intervals[{i}] must have at least 2 elements (start, end), '
-                        f'got {len(interval)}'
-                    )
-
-        if self.start_date is not None and not isinstance(self.start_date, str):
-            raise TypeError(f'start_date must be a string, got {type(self.start_date).__name__}')
-
-        if self.end_date is not None and not isinstance(self.end_date, str):
-            raise TypeError(f'end_date must be a string, got {type(self.end_date).__name__}')
-
-    def _validate_limits(self) -> None:
-        """Validate low_lim and upp_lim consistency."""
-        if set(self.low_lim.keys()) != set(self.upp_lim.keys()):
-            raise ValueError(
-                f'low_lim and upp_lim must have the same keys. '
-                f'low_lim keys: {set(self.low_lim.keys())}, '
-                f'upp_lim keys: {set(self.upp_lim.keys())}'
-            )
-
-        for var in self.low_lim:
-            if self.low_lim[var] >= self.upp_lim[var]:
-                raise ValueError(
-                    f'low_lim must be less than upp_lim for variable "{var}". '
-                    f'Got low_lim={self.low_lim[var]}, upp_lim={self.upp_lim[var]}'
-                )
+    def _validate_model_param(self, schema: dict[str, Any], value: Any) -> None:
+        """Validate model parameter against schema."""
+        try:
+            validator = Draft202012Validator(schema)
+            validator.validate(value)
+        except ValidationError as e:
+            raise ValueError(f'Model parameters validation failed: {e.message}')
+        except Exception as e:
+            raise ValueError(f'Unexpected error: {e}')
 
     def _validate_required_strings(self) -> None:
         """Validate required string fields are not empty."""
@@ -334,28 +251,10 @@ class TrainModelParams:
         if not self.file_name.strip():
             raise ValueError('file_name cannot be empty or whitespace')
 
-        if not self.experiment_name.strip():
-            raise ValueError('experiment_name cannot be empty or whitespace')
+        if not self.model_name.strip():
+            raise ValueError('model_name cannot be empty or whitespace')
 
     def _validate_date_format(self) -> None:
         """Validate date_format is one of the allowed frontend formats when set."""
         if self.date_format:
             validate_frontend_date_format(self.date_format)
-
-    def _validate_validation_file_name(self) -> None:
-        """
-        Validate that validation_file_name, when provided, is not empty or whitespace.
-
-        This field is optional; when present it must point to a valid object key in the
-        same MinIO bucket specified by bucket_name.
-        """
-        if self.validation_file_name is None:
-            return
-
-        if not isinstance(self.validation_file_name, str):
-            raise TypeError(
-                f'validation_file_name must be a string, got {type(self.validation_file_name).__name__}'
-            )
-
-        if not self.validation_file_name.strip():
-            raise ValueError('validation_file_name cannot be empty or whitespace')

model_manager/utils/models/train_model_result.py

@@ -34,12 +34,10 @@ class TrainModelResult:
     """
 
     params: TrainModelParams
-    x_train: pd.DataFrame
-    x_test: pd.DataFrame
-    y_train: pd.Series
-    y_test: pd.Series
-    y_pred: pd.Series | None = None
-    y_train_pred: pd.Series | None = None
+    train_data: pd.DataFrame
+    val_data: pd.DataFrame
+    y_pred: pd.DataFrame | None = None
+    y_train_pred: pd.DataFrame | None = None
     mse_val: float | None = None
     mae_val: float | None = None
     r2_val: float | None = None

model_manager/utils/repository/data_manager_repository.py

@@ -23,11 +23,35 @@ from sientia_do.observability.sientia_monitoring import SientiaMonitoring
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
 from model_manager.sientia.metrics import mae, mse, r2
-from model_manager.sientia.utils import split_train_test
 from model_manager.utils.models.train_model_params import TrainModelParams
 from model_manager.utils.models.train_model_result import TrainModelResult
 
 
+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
+    if random_state is not None:
+        np.random.seed(random_state)
+    
+    # 2. Gerar índices e embaralhar se necessário
+    indices = np.arange(len(data))
+    
+    if shuffle:
+        np.random.shuffle(indices)
+    
+    # 3. Calcular o ponto de corte (split point)
+    # Cálculo: N_treino = tamanho_total * proporcao_treino
+    n_train = int(len(data) * train_size)
+    
+    # 4. Dividir os índices
+    train_indices = indices[:n_train]
+    test_indices = indices[n_train:]
+    
+    # 5. Retornar os dados fatiados (funciona para DataFrame ou Series)
+    if isinstance(data, (pd.DataFrame, pd.Series)):
+        return data.iloc[train_indices], data.iloc[test_indices]
+    
+    return data[train_indices], data[test_indices]
+
 def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
     """
     If date_column is set, parse the column as timezone-aware
@@ -53,74 +77,6 @@ def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) ->
     return data
 
 
-def _single_variable_support_mask(
-    data_view: pd.DataFrame,
-    var_col: str,
-    target_variable: str,
-    config: dict,
-) -> np.ndarray | None:
-    """Compute keep mask for one variable's support lines; None if config is invalid or skipped."""
-    if var_col not in data_view.columns:
-        return None
-    upper = config.get('upper_line') or config.get('upperLine')
-    lower = config.get('lower_line') or config.get('lowerLine')
-    if not upper or not lower:
-        return None
-
-    x_vals = data_view[var_col].astype(float).to_numpy()
-    y_vals = data_view[target_variable].astype(float).to_numpy()
-    xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
-    ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
-    x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
-    y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
-    scale_ratio = y_range / x_range
-
-    b1 = float(upper.get('intercept', 0))
-    deg1 = float(upper.get('angle', 0))
-    b2 = float(lower.get('intercept', 0))
-    deg2 = float(lower.get('angle', 0))
-    m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
-    m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
-    y1 = m1 * x_vals + b1
-    y2 = m2 * x_vals + b2
-    lower_bound = np.minimum(y1, y2)
-    upper_bound = np.maximum(y1, y2)
-    return (y_vals >= lower_bound) & (y_vals <= upper_bound)
-
-
-def _apply_support_filters(
-    data_view: pd.DataFrame,
-    target_variable: str,
-    support_filters: dict,
-) -> pd.DataFrame:
-    """
-    Keep only rows where (var, target) lies between the two guide lines for each variable.
-
-    For each variable in support_filters, the condition is lower(x_var) <= target <= upper(x_var),
-    where lower/upper are the two lines (intercept + slope from angle, scaled by y_range/x_range).
-    Global mask is AND across all variables. Matches DEMO logic in template_01.py.
-
-    Args:
-        data_view: DataFrame after preprocessor transform.
-        target_variable: Name of the target column (y axis).
-        support_filters: Per-variable config with upper_line/lower_line, each {intercept, angle}.
-
-    Returns:
-        data_view filtered to rows satisfying all variable conditions; unchanged if support_filters empty.
-    """
-    if not support_filters or target_variable not in data_view.columns:
-        return data_view
-
-    combined_keep_mask = np.ones(len(data_view), dtype=bool)
-    n = len(data_view)
-    for var_col, config in support_filters.items():
-        keep_mask = _single_variable_support_mask(data_view, var_col, target_variable, config)
-        if keep_mask is not None and len(keep_mask) == n:
-            combined_keep_mask &= keep_mask
-
-    return data_view.loc[combined_keep_mask]
-
-
 class DataManagerRepository(SientiaMonitoring):
     """
     Repository for data preparation in the training pipeline.
@@ -193,17 +149,11 @@ class DataManagerRepository(SientiaMonitoring):
         train_df = _ensure_date_column_parsed(train_df, params)
         train_df = self._configure_datetime_index(train_df, params, metadata)
 
-        if params.support_filters:
-            train_df = _apply_support_filters(
-                train_df,
-                params.target_variable,
-                params.support_filters,
-            )
-
         if len(train_df) <= 0:
             raise ValueError('Training data view is empty after transformation')
 
         # Explicit validation dataset path
+        train_data = pd.DataFrame(train_df[params.variable_columns + [params.target_variable]])
         if validation_file_bytes is not None:
             try:
                 val_df = pd.read_csv(
@@ -220,28 +170,18 @@ class DataManagerRepository(SientiaMonitoring):
             val_df = _ensure_date_column_parsed(val_df, params)
             val_df = self._configure_datetime_index(val_df, params, metadata)
 
-            if params.support_filters:
-                val_df = _apply_support_filters(
-                    val_df,
-                    params.target_variable,
-                    params.support_filters,
-                )
-
             if len(val_df) <= 0:
                 raise ValueError('Validation data view is empty after transformation')
 
-            x_train = pd.DataFrame(train_df[params.variable_columns])
-            y_train = pd.Series(train_df[params.target_variable])
-            x_test = pd.DataFrame(val_df[params.variable_columns])
-            y_test = pd.Series(val_df[params.target_variable])
+
+            val_data = pd.DataFrame(val_df[params.variable_columns + [params.target_variable]])
         else:
             # Fallback path: derive validation via train/test split from a single dataset.
-            x_train, x_test, y_train, y_test = split_train_test(
-                pd.DataFrame(train_df[params.variable_columns]),
-                pd.Series(train_df[params.target_variable]),
+            train_data, val_data = train_test_split(
+                train_data,
                 train_size=params.train_size / 100,
                 shuffle=params.shuffle,
-                random_state=42,
+                random_state=params.random_state,
             )
 
         self.info(
@@ -251,59 +191,62 @@ class DataManagerRepository(SientiaMonitoring):
 
         return TrainModelResult(
             params=params,
-            x_train=x_train,
-            x_test=x_test,
-            y_train=y_train,
-            y_test=y_test,
+            train_data=train_data,
+            val_data=val_data
         )
 
+    def _as_series(self, pred: pd.DataFrame | pd.Series) -> pd.Series:
+        if isinstance(pred, pd.Series):
+            return pred
+        # If the wrapper returns a single-column DataFrame, take its first column.
+        if pred.shape[1] == 1:
+            return pred.iloc[:, 0]
+        raise ValueError('y_pred/y_train_pred must be a Series or single-column DataFrame')
+
     def compute_regression_metrics(
         self,
-        params: TrainModelParams,
         tmr: TrainModelResult,
     ) -> TrainModelResult:
         """
         Compute regression metrics for training results.
 
         This helper mirrors the previous TrainingRepository.after_train_calculation
-        behavior, assuming that y_pred/y_train_pred are already on the correct scale
-        for metric calculation (any scaling is handled inside the model wrapper).
+        behavior, assuming that predictions (y_pred/y_train_pred) are already on the
+        correct scale for metric calculation (any scaling is handled inside the
+        model wrapper).
 
         Args:
-            params: Training parameters used during model training.
-            tmr: Training result with y_train, y_test, y_train_pred and y_pred populated.
+            tmr: Training result containing:
+                - train_data/val_data DataFrames with a target column
+                - y_train_pred/y_pred populated (model predictions for train/val)
 
         Return:
             Updated TrainModelResult with mse_val, mae_val and r2_val fields populated.
         """
-        del params  # unused for now, kept for possible future extensions
+        if tmr.y_pred is None:
+            raise ValueError('y_pred must be set before computing regression metrics')
 
-        tmr.x_train = tmr.x_train.sort_index()
-        tmr.x_test = tmr.x_test.sort_index()
-        tmr.y_train = tmr.y_train.sort_index()
-        tmr.y_test = tmr.y_test.sort_index()
+        params = tmr.params
+        target = params.target_variable
 
-        if tmr.y_pred is not None:
-            tmr.y_pred = tmr.y_pred.sort_index()
-        if tmr.y_train_pred is not None:
-            tmr.y_train_pred = tmr.y_train_pred.sort_index()
+        # True values are expected to come from val_data.
+        y_true_val = tmr.val_data[target]
         
-        assert tmr.y_pred is not None, 'y_pred should be set at this point'
+        y_pred_val = self._as_series(tmr.y_pred).sort_index()
+        y_true_val = y_true_val.sort_index()
 
-        tmr.mse_val = round(
-            mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
-            2,
-        )
+        # Align by index to avoid metric calculation errors if ordering differs.
+        common_index = y_true_val.index.intersection(y_pred_val.index)
+        y_true_val = y_true_val.loc[common_index]
+        y_pred_val = y_pred_val.loc[common_index]
 
-        tmr.mae_val = round(
-            mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
-            2,
-        )
+        if len(y_true_val) == 0:
+            raise ValueError('No overlapping indices between val_data and y_pred')
 
-        tmr.r2_val = round(
-            r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
-            2,
-        )
+        # Metrics helpers already round to 2 decimals.
+        tmr.mse_val = mse(y_true_val, y_pred_val)
+        tmr.mae_val = mae(y_true_val, y_pred_val)
+        tmr.r2_val = r2(y_true_val, y_pred_val)
 
         return tmr
 
@@ -324,8 +267,6 @@ class DataManagerRepository(SientiaMonitoring):
                 'Data is None after load_data. '
                 'Check file format, line separator and decimal separator.'
             )
-        if not isinstance(data, pd.DataFrame):
-            raise TypeError(f'Expected DataFrame, got {type(data).__name__}')
 
         if isinstance(data.index, pd.DatetimeIndex):
             self.info('DataFrame already has DatetimeIndex', metadata)

model_manager/worker/worker.py

@@ -178,6 +178,7 @@ async def main():
             other_workflows=[],
             activities=[
                 activities.update_experiment_run,
+                activities.load_model_metadata,
                 activities.validate_train_params,
                 activities.train_model,
                 activities.cleanup_resources,

model_manager/workflows/train_model.py

@@ -53,7 +53,6 @@ with workflow.unsafe.imports_passed_through():
         maximum_attempts=5,
     )
 
-POD_ID = os.getenv('POD_ID')
 
 
 @workflow.defn(name='train_model')
@@ -97,12 +96,15 @@ class TrainModel:
             ValueError: If experiment_run_id is missing or invalid
         """
         experiment_run_id = self._validate_experiment_run_id(input_data)
+        model_name = input_data.get('model_name')
+        model_id = input_data.get('model_id')
         
         metadata = {
             'metadata': {
-                'pod_id': POD_ID,
                 'experiment_run_id': experiment_run_id,
                 'workflow_name': 'train_model',
+                'model_name': model_name,
+                'model_id': model_id,
             }
         }
 
@@ -177,11 +179,21 @@ class TrainModel:
             Exception: If validation fails (after updating DB status)
         """
         try:
+            input_data = await workflow.execute_activity_method(
+                Activities.load_model_metadata,
+                {
+                    **input_data,
+                    **metadata,
+                },
+                retry_policy=no_retry_policy,
+                start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
+            )
+
             train_params = await workflow.execute_activity_method(
                 Activities.validate_train_params,
                 {
-                    **metadata,
                     **input_data,
+                    **metadata,
                 },
                 retry_policy=no_retry_policy,
                 start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),

requirements.txt

@@ -3,7 +3,7 @@ psycopg2-binary==2.9.11
 sqlalchemy==2.0.44
 boto3==1.40.55
 botocore==1.40.55
-git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.7.2
+/home/grezewave/Documents/projects/sientia/sientia-model-library
 git+https://github.com/Aignosi/sientia-dataops-library.git@v1.10.1
 prometheus-client==0.23.1
 beautifulsoup4==4.12.3

scripts/run_training_test.py

@@ -341,7 +341,6 @@ def run_local_pipeline(
             metadata={'source': 'run_local_pipeline', 'scenario': scenario_name},
         )
         train_result = data_manager_repository.compute_regression_metrics(
-            train_params,
             train_result,
         )
     except Exception as e:

tests/activities/test_training.py

@@ -55,7 +55,7 @@ def mock_train_params():
     params = MagicMock()
     params.experiment_run_id = 1
     params.target_variable = 'target'
-    params.experiment_name = 'test_experiment'
+    params.model_name = 'Linear Regression'
     params.bucket_name = 'test-bucket'
     params.file_name = 'test-file.csv'
     params.validate_business_rules = MagicMock()

tests/utils/models/test_train_model_params.py

@@ -24,7 +24,6 @@ def valid_train_params_dict():
         'train_size': 80,
         'shuffle': True,
         'experiment_run_id': 1,
-        'experiment_name': 'test_experiment',
         'removed_intervals': [],
         'model_name': 'Linear Regression',
         'degree': 1,
@@ -61,7 +60,6 @@ def test_train_model_params_from_dict_success(valid_train_params_dict):
     assert params.train_size == 80
     assert params.shuffle is True
     assert params.experiment_run_id == 1
-    assert params.experiment_name == 'test_experiment'
     assert params.removed_intervals == []
 
 
@@ -297,25 +295,14 @@ def test_validate_business_rules_whitespace_file_name(valid_train_params_dict):
         params.validate_business_rules()
 
 
-def test_validate_business_rules_empty_experiment_name(valid_train_params_dict):
-    """Test validate_business_rules raises error when experiment_name is empty."""
+def test_validate_business_rules_model_name_empty(valid_train_params_dict):
+    """Test validate_business_rules raises error when model_name is empty."""
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    valid_train_params_dict['experiment_name'] = ''
+    valid_train_params_dict['model_name'] = ''
     params = TrainModelParams.from_dict(valid_train_params_dict)
 
-    with pytest.raises(ValueError, match='experiment_name cannot be empty or whitespace'):
-        params.validate_business_rules()
-
-
-def test_validate_business_rules_whitespace_experiment_name(valid_train_params_dict):
-    """Test validate_business_rules raises error when experiment_name is whitespace."""
-    from model_manager.utils.models.train_model_params import TrainModelParams
-
-    valid_train_params_dict['experiment_name'] = '  \n  '
-    params = TrainModelParams.from_dict(valid_train_params_dict)
-
-    with pytest.raises(ValueError, match='experiment_name cannot be empty or whitespace'):
+    with pytest.raises(ValueError, match='model_name cannot be empty or whitespace'):
         params.validate_business_rules()
 
 

tests/utils/models/test_train_model_result.py

@@ -30,7 +30,6 @@ def sample_params():
         train_size=80,
         shuffle=True,
         experiment_run_id=123,
-        experiment_name='test-experiment',
         removed_intervals=[],
         model_name='Linear Regression',
         degree=1,