feat: log regression metrics as parameters in Training class

- Added a method to persist computed regression metrics (MSE, MAE, R²) as MLflow parameters during model training, enhancing model evaluation and tracking.
- Updated the Training class to log the equation path if available, improving artifact management.

docs/train-model-workflow-io-diff-main-vs-current-branch.md

@@ -0,0 +1,417 @@
+# Train Model Workflow IO Diff (`main` vs current branch)
+
+Base comparison: `git diff main...HEAD`
+Workflow analyzed: `train_model`
+
+## 1) Executive overview
+
+This branch introduces a structural refactor of the training stack and a contract update for workflow input/output.
+
+Main impacts:
+
+- The old in-house training stack (`TrainingRepository`, `ModelRepository`, `StorageRepository`, `model_manager.sientia.models`) was replaced by:
+  - `DataManagerRepository` (data prep + metrics + report generation)
+  - `SientiaModel` wrapper from plugin store (`sientia_model`)
+  - `SientiaMLflowRepository` (MLflow integration)
+  - `MinioRepository` (storage integration)
+- Input contract moved from many fixed legacy ML params to a plugin/wrapper-oriented schema (`model_type`, `*_kwargs`, `model_metadata`, optional `val_file_name`).
+- Workflow return changed from `None` to a serializable result object (`dict[str, Any] | None`) containing training execution metadata.
+- Queue naming and worker bootstrap architecture now depend on runtime (`train_model-<runtime>-queue`).
+
+---
+
+## 2) Input contract diff (before vs now)
+
+### 2.1 Previous contract (`main`)
+
+`TrainModelParams` in `main` required a large set of explicit fields for the old preprocessing/model pipeline, focused only in linear regression model:
+
+- Core:
+  - `experiment_run_id`, `variable_columns`, `target_variable`
+  - `bucket_name`, `file_name`, `line_separator`, `decimal_separator`
+  - `train_size`, `shuffle`
+- Legacy preprocessing/model fields focused only in linear regression model (required in `from_dict`):
+  - `lag_train`, `lag_val`
+  - `rem_static_win`, `low_lim`, `upp_lim`, `window`
+  - `use_scaler`, `include_ar`, `scaler_name`
+  - `removed_intervals`, `start_date`, `end_date`, `nan_treatment`
+  - `degree`, `interaction_only`
+  - `experiment_name`, `model_name`
+  - `support_filters` (optional dict), `static_threshold` (optional int)
+
+Validation was strongly tied to this structure (lag ranges, limits consistency, polynomial/scaler constraints, etc.).
+
+### 2.2 Current contract (this branch)
+
+`TrainModelParams` now supports a plugin-driven schema and wrapper kwargs:
+
+- Kept/mandatory core fields:
+  - `experiment_run_id` (now accepts numeric string too; coerced to int)
+  - `variable_columns`, `target_variable`
+  - `bucket_name`, `file_name`, `line_separator`, `decimal_separator`
+  - `train_size`, `shuffle`
+  - `model_name`
+  - `model_type`
+  - `data_model_kwargs`, `model_kwargs`, `opt_params` (required as dict by current `from_dict`)
+- New/updated fields:
+  - `random_state` (default `42`)
+  - `val_file_name` (optional explicit validation file)
+  - `model_id` (currently optional, but needs discussion, since the model metadata in MongoDB should be created before the model training)
+- Removed from required input contract:
+  - `lag_train`, `lag_val`, `rem_static_win`, `low_lim`, `upp_lim`, `window`
+  - `use_scaler`, `include_ar`
+  - `degree`, `interaction_only`, `nan_treatment`
+  - `start_date`, `end_date`, `scaler_name`
+  - `removed_intervals`, `support_filters`, `static_threshold`
+- Parameters internally derived:
+  - `model_metadata` model type info from plugin store.
+  - `run_name` is internally derived from experiment name and datetime.
+  - `experiment_name` is internally derived from `model_name`.
+
+### 2.3 Validation behavior changes
+
+Before:
+- Validation was mostly hardcoded business checks tied to legacy linear/polynomial stack.
+
+Now:
+- Validation still checks core constraints (`train_size`, non-empty strings, etc.), but model-specific validation moved to JSON Schema driven checks, using OpenAPI/JSON Schema definitions from plugin store:
+  - `model_metadata.schemas.components.schemas.data_model`
+  - `model_metadata.schemas.components.schemas.model`
+  - `model_metadata.schemas.components.schemas.opt_params`
+- `model_metadata` is now a required semantic dependency for `validate_business_rules()`.
+- Date format validation remains, but allowed formats are defined locally in `train_model_params.py`.
+
+### 2.4 Input loading pipeline changes in workflow
+
+Before:
+- `validate_train_params` directly consumed workflow input.
+
+Now:
+1. `load_model_metadata` runs first (fetches model index/schema from plugin store and injects `model_metadata`).
+2. `validate_train_params` runs with enriched payload.
+
+This means IO preprocessing now depends on plugin-store metadata resolution before final validation.
+
+---
+
+## 3) Output contract diff (before vs now)
+
+### 3.1 Workflow return (`train_model.run`)
+
+Before (`main`):
+- Return type: `None`
+- Workflow side effects were persisted mainly via DB status updates and MLflow artifacts.
+
+Now:
+- Return type: `dict[str, Any] | None`
+- Workflow returns the training activity summary when successful.
+
+### 3.2 Activity-level training result payload
+
+Before (from `Training.train_model` in `main` path):
+- Returned minimal dict:
+  - `run_name`
+  - `run_dir`
+
+Now:
+- Returns extended dict:
+  - `run_name`
+  - `experiment_name`
+  - `run_id`
+  - `run_dir`
+
+### 3.3 Persistence map by destination (DB, MLflow, MinIO, local filesystem)
+
+This section maps where each artifact/metadata goes, in which format, and how that changed from `main`.
+
+#### 3.3.1 PostgreSQL (`experiment_run` table)
+
+## Before (`main`)
+
+- Update path: `update_experiment_run` activity with `UpdateType.MODEL_SAVED`.
+- Persisted on success:
+  - `status` transition to `TRAINING_SUCCESS`
+  - `run_name` (MLflow run identifier used by current implementation)
+- Persisted on failures:
+  - `status` transition to validation/training error statuses
+  - `error_message`
+
+## Now (current branch)
+
+- Same update path and status/error behavior.
+- Even though train activity now returns more metadata (`run_id`, `experiment_name`), current workflow update for `MODEL_SAVED` still forwards mainly `run_name`.
+- Practical effect:
+  - DB remains status-centric and run-name-centric
+  - richer identifiers exist in workflow return payload, not fully mirrored to DB columns in current flow
+
+#### 3.3.2 MLflow (tracking server/artifact store)
+
+## Before (`main`)
+
+- Persistence orchestration lived in `ModelRepository.save_model()` + `_save_run()`.
+- Typical persisted content:
+  - model params (many legacy params such as lags, limits, scaler config, removed intervals)
+  - regression metrics (`MSE`, `R2`, `MAE`)
+  - model objects:
+    - `data_model`
+    - `prediction_model`
+  - artifacts:
+    - `report.html`
+    - `train_data.csv`
+    - `test_data.csv`
+    - optional `model_equation.json`
+- Run naming:
+  - computed by querying existing runs and appending sequence (`<experiment>-<n>` style)
+
+## Now (current branch)
+
+- Persistence orchestrated in `Training._persist_training_artifacts()` and MLflow run context is opened by `SientiaMLflowRepository.start_run(...)`.
+- Persisted content now:
+  - model wrapper itself via `wrapper.store_model(name=train_params.model_name)`
+  - regression metrics also logged as MLflow params via `mlflow.log_param(...)`:
+    - `mse_val`
+    - `mae_val`
+    - `r2_val`
+  - artifacts explicitly logged with `mlflow.log_artifact(...)`:
+    - `report.html`
+    - `train_data.csv`
+    - `test_data.csv`
+  - metrics are computed before save (`mse_val`, `mae_val`, `r2_val`) and persisted in the run as params
+- Run identifiers now exposed back to workflow:
+  - `experiment_name`
+  - `run_name`
+  - `run_id`
+- Notable behavioral change:
+  - `wrapper._input_example` is cleared (`None`) before storing model.
+
+#### 3.3.3 MinIO object storage
+
+## Before (`main`)
+
+- Read path:
+  - single source object downloaded via `StorageRepository.fetch_file(bucket_name, file_name)`
+- Write path:
+  - training workflow did not write generated outputs to MinIO in this code path
+  - generated artifacts were persisted to MLflow, not uploaded back to MinIO
+- Location:
+  - source data in input bucket/key provided by workflow input (`bucket_name` + `file_name`)
+
+## Now (current branch)
+
+- Read path migrated to `MinioRepository.download_file(...)`.
+- Supports two input objects:
+  - mandatory training object: `bucket_name` + `file_name`
+  - optional validation object: same `bucket_name` + `val_file_name`
+- Write path:
+  - still no artifact upload to MinIO in this workflow path
+  - report/CSV outputs continue to flow to MLflow artifacts
+- Location details:
+  - bucket resolved from payload (`bucket_name`)
+  - object key exactly from payload (`file_name`, optional `val_file_name`)
+  - default bucket in env/config is `MINIO_DEFAULT_BUCKET`, but runtime payload can override via `bucket_name`
+
+#### 3.3.4 Local filesystem (ephemeral runtime workspace)
+
+## Before (`main`)
+
+- Temporary run dir created under reports root using run name + timestamp suffix.
+- Artifacts generated locally in that directory:
+  - `report.html`
+  - `train_data.csv`
+  - `test_data.csv`
+  - optional `model_equation.json`
+- After MLflow logging, cleanup activity removed temp directory.
+
+## Now (current branch)
+
+- Temporary run dir managed by `DataManagerRepository` under runtime reports root (`.../reports/temp/<run_name>`).
+- Same artifact family generated locally:
+  - `report.html`
+  - `train_data.csv`
+  - `test_data.csv`
+  - optional `model_equation.json` (for `linear_regression`)
+- Cleanup behavior is now tolerant:
+  - cleanup runs in guarded `finally`
+  - training success is not reverted if cleanup later fails
+
+#### 3.3.5 Quick matrix (before vs now)
+
+- **Postgres**
+  - before: status + run_name + errors
+  - now: same persisted shape; workflow return contains extra IDs
+- **MLflow**
+  - before: legacy model objects + params/metrics + report/data artifacts
+  - now: wrapper-based model persistence + `mse_val`/`mae_val`/`r2_val` as params + report/data artifacts + run_id exposed
+- **MinIO**
+  - before: reads 1 CSV input object
+  - now: reads 1 or 2 CSV input objects (train + optional validation), still no output upload
+- **Local temp**
+  - before: generated artifacts, then cleanup
+  - now: generated artifacts, then best-effort cleanup (non-blocking for success result)
+
+### 3.4 Cleanup behavior impact on output semantics
+
+Before:
+- Cleanup was called directly after training result; failures propagated straightforwardly.
+
+Now:
+- Cleanup is in a guarded `finally`.
+- If training succeeded but cleanup fails, workflow warns and does not rollback success semantics.
+- Effective output semantics: successful training result can be returned even if temp cleanup fails.
+
+---
+
+## 4) Detailed field mapping (old -> new)
+
+## Kept (or equivalent role)
+
+- `experiment_run_id` -> kept (broader accepted types: int or numeric string)
+- `variable_columns` -> kept
+- `target_variable` -> kept
+- `bucket_name` -> kept
+- `file_name` -> kept
+- `line_separator` -> kept
+- `decimal_separator` -> kept
+- `date_column` -> kept optional
+- `date_format` -> kept optional
+- `train_size` -> kept
+- `shuffle` -> kept
+- `model_name` -> kept (now less coupled to legacy model enum)
+
+## Added
+
+- `model_type` (primary selector for plugin wrapper/index lookup)
+- `data_model_kwargs`
+- `model_kwargs`
+- `opt_params`
+- `val_file_name` (optional second dataset input)
+- `model_id` (optional metadata)
+- `model_metadata` (loaded/required for schema validation)
+- `random_state` (explicit split reproducibility control)
+
+## Removed from new required contract
+
+- `lag_train`, `lag_val`
+- `rem_static_win`, `static_threshold`
+- `low_lim`, `upp_lim`
+- `window`
+- `use_scaler`, `include_ar`
+- `degree`, `interaction_only`
+- `nan_treatment`
+- `start_date`, `end_date`
+- `scaler_name`
+- `removed_intervals`
+- `support_filters`
+- `experiment_name` (no longer required as top-level client input)
+
+---
+
+## 5) Internal architecture update notes
+
+### 5.1 Repository layer redesign
+
+Removed:
+- `model_manager/utils/repository/model_repository.py`
+- `model_manager/utils/repository/training_repository.py`
+- `model_manager/utils/repository/storage_repository.py`
+
+Added:
+- `model_manager/utils/repository/data_manager_repository.py`
+
+Interpretation:
+- Data preprocessing/report/metrics responsibilities were consolidated into `DataManagerRepository`.
+- Training/model persistence shifted to wrapper + plugin store + MLflow repository integrations.
+
+### 5.2 Model engine abstraction migration
+
+Before:
+- Strong coupling to local classes in `model_manager.sientia.models` and custom preprocessing/model objects in `TrainModelResult`.
+
+Now:
+- Training uses `SientiaModel` wrapper dynamically obtained by `plugin_store.get_model(model_type=...)`.
+- Contract is wrapper-driven (`train`, `transform`, `predict`, `store_model`).
+- The codebase removed `model_manager/sientia/models.py`, `model_serving.py`, and `utils.py`, indicating full migration to externalized model runtime abstraction.
+
+### 5.3 Worker/runtime architecture changes
+
+- New `prepare_worker.py` centralizes worker setup and autoscaling parameters.
+- Queue names are now runtime-derived:
+  - `train_model-<runtime>-queue`
+  - `cleanup_files-<runtime>-queue`
+- `worker.py` now installs runtime via plugin store (`plugin_store.install_runtime(runtime_name=...)`) before starting workers.
+- This introduces environment/runtime-aware deployment and model packaging behavior.
+
+### 5.4 Synchronous activity and tracking adjustments
+
+- `experiment_tracking` migrated from async postgres helper to sync postgres client path (`postgres_sync`).
+- Several activities switched to sync method signatures.
+- Error handling in workflow and DB status update paths is more defensive (secondary failures while persisting error status are logged and do not mask primary failure cause).
+
+### 5.5 `TrainModelResult` shape update
+
+Before:
+- Stored classic split artifacts (`x_train`, `x_test`, `y_train`, `y_test`) + concrete preprocessing/model objects (`process_data`, `regr`, `scaler_dict`).
+
+Now:
+- Stores `train_data`, `val_data` and prediction DataFrames, plus tracking identifiers (`experiment_name`, `run_id`).
+- Result object is less tied to internal estimator classes and more aligned with serializable workflow/model-store integration.
+
+---
+
+## 6) Net IO compatibility assessment
+
+## Input compatibility
+
+Not backward compatible with old payloads without adaptation.
+
+Key reasons:
+- Legacy required fields removed/ignored by new path.
+- New required fields introduced (`model_type`, `*_kwargs` dicts, runtime metadata flow dependency).
+- Validation pipeline now expects model metadata semantics.
+
+## Output compatibility
+
+Behavior changed:
+- Workflow now returns a result object (previously `None`).
+- Training summary includes `experiment_name` and `run_id` in addition to `run_name` and `run_dir`.
+- DB update still centered on `run_name`; callers relying only on DB may not see all new output info unless workflow return is consumed.
+
+---
+
+## 7) Practical migration guidance (client side)
+
+To call `train_model` in this branch:
+
+1. Send snake_case payload aligned to new `TrainModelParams`.
+2. Always provide:
+   - `model_name` slugified model name (ex.: `test_model_name or test-model-name`)
+   - `model_type`
+   - `data_model_kwargs` (dict)
+   - `model_kwargs` (dict)
+   - `opt_params` (dict)
+3. Keep `experiment_run_id` numeric (int or numeric string).
+4. Use runtime queue naming consistent with worker runtime:
+   - `train_model-<runtime>-queue`
+5. If you need explicit validation split file, send `val_file_name`; otherwise split uses `train_size`/`shuffle`/`random_state`.
+
+---
+
+## 8) Source references used for this document
+
+Primary diffs:
+- `model_manager/workflows/train_model.py`
+- `model_manager/utils/models/train_model_params.py`
+- `model_manager/utils/models/train_model_result.py`
+- `model_manager/activities/training.py`
+- `model_manager/activities/activities.py`
+- `model_manager/activities/experiment_tracking.py`
+- `model_manager/utils/repository/data_manager_repository.py`
+- `model_manager/utils/repository/model_repository.py` (removed)
+- `model_manager/utils/repository/training_repository.py` (removed)
+- `model_manager/utils/repository/storage_repository.py` (removed)
+- `model_manager/worker/worker.py`
+- `model_manager/worker/prepare_worker.py`
+- `README.md`
+- `input-sample.md`
+- `scripts/run_training_test.py`
+

model_manager/activities/training.py

@@ -14,7 +14,6 @@ with workflow.unsafe.imports_passed_through():
     from typing import Any
 
     import mlflow
-    import pandas as pd
     from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
     from sientia_do.notifications.models import NotificationLevel
     from sientia_do.observability.logger import Logger
@@ -340,10 +339,30 @@ class Training(SientiaMonitoring):
         self.info(f'Storing model for {train_params.model_type}', metadata)
         wrapper._input_example = None
         wrapper.store_model(name=train_params.model_name)
+        self._log_regression_metrics_as_params(train_result)
         self.info(f'Logging artifacts for {train_params.model_type}', metadata)
         mlflow.log_artifact(train_result.report_path)
         mlflow.log_artifact(train_result.train_data_path)
         mlflow.log_artifact(train_result.test_data_path)
+        if train_result.equation_path is not None:
+            mlflow.log_artifact(train_result.equation_path)
+
+    def _log_regression_metrics_as_params(self, train_result: TrainModelResult) -> None:
+        """
+        Persist computed regression metrics as MLflow params.
+
+        Args:
+            train_result: Training output containing computed regression metrics.
+        """
+        metric_params = {
+            'mse_val': train_result.mse_val,
+            'mae_val': train_result.mae_val,
+            'r2_val': train_result.r2_val,
+        }
+
+        for key, value in metric_params.items():
+            if value is not None:
+                mlflow.log_param(key, value)
 
     @activity.defn(name='cleanup_resources')
     def cleanup_resources(self, input_data: dict[str, Any]) -> None:

model_manager/sientia/reports.py

@@ -19,8 +19,8 @@ from evidently.metrics import (
 )
 from evidently.metrics.base_metric import generate_column_metrics
 from evidently.options import ColorOptions
-from evidently.report import Report
 from evidently.pipeline.column_mapping import ColumnMapping
+from evidently.report import Report
 
 COLOR_DISCRETE_SEQUENCE = (
     '#ed0400',
@@ -36,6 +36,7 @@ def load_html_from_file(file_path):
     with open(file_path, encoding='utf-8') as file:
         return file.read()
 
+
 def inject_content(main_html, section_id, content):
     soup = BeautifulSoup(main_html, 'html.parser')
     section = soup.find(id=section_id)
@@ -68,7 +69,12 @@ class Reports:
     """
 
     def __init__(
-        self, reference_data: Any, current_data: Any, target_name: str, base_path: str | None = None, template_path: str | None = None
+        self,
+        reference_data: Any,
+        current_data: Any,
+        target_name: str,
+        base_path: str | None = None,
+        template_path: str | None = None,
     ) -> None:
         """
         Initializes an instance of the AigReport class.
@@ -247,10 +253,10 @@ class Reports:
         if output_dir and not os.path.exists(output_dir):
             os.makedirs(output_dir, exist_ok=True)
 
-        print(f"Output directory: {output_dir}")
+        print(f'Output directory: {output_dir}')
-        print(f"Report path: {report_path}")
+        print(f'Report path: {report_path}')
-        print(f"Base path: {self.base_path}")
+        print(f'Base path: {self.base_path}')
-        print(f"Template path: {self.template_path}")
+        print(f'Template path: {self.template_path}')
 
         # Load main HTML template
         main_html_path = os.path.join(self.template_path, 'header.html')

model_manager/utils/models/train_model_params.py

@@ -23,6 +23,7 @@ def validate_frontend_date_format(fmt: str | None) -> None:
         allowed = ', '.join(sorted(ALLOWED_FRONTEND_DATE_FORMATS))
         raise ValueError(f'Invalid date_format "{fmt}". Allowed formats: {allowed}')
 
+
 # Model name constants
 MODEL_LINEAR_REGRESSION = 'Linear Regression'
 MODEL_POLYNOMIAL_REGRESSION = 'Polynomial Regression'

model_manager/utils/repository/data_manager_repository.py

@@ -20,7 +20,6 @@ from os import makedirs, path
 from shutil import rmtree
 from typing import Any
 
-from mlflow.entities import experiment
 import numpy as np
 import pandas as pd
 from sientia_do.observability.logger import Logger
@@ -28,7 +27,7 @@ from sientia_do.observability.sientia_monitoring import SientiaMonitoring
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
 from sientia_model.wrappers.sientia_model import SientiaModel
 
-from model_manager.runtime_paths import REPORTS_ROOT, PROJECT_BASE_PATH
+from model_manager.runtime_paths import PROJECT_BASE_PATH, REPORTS_ROOT
 from model_manager.sientia.metrics import mae, mse, r2
 from model_manager.sientia.reports import Reports  # type: ignore[import-untyped]
 from model_manager.utils.models.train_model_params import TrainModelParams
@@ -201,10 +200,16 @@ class DataManagerRepository(SientiaMonitoring):
             metadata,
         )
 
-        experiment_name = f'train_model_{params.model_type}_{params.model_name}_{params.experiment_run_id}'
+        experiment_name = f'{params.model_name}'
         run_name = f'{experiment_name}_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
 
-        return TrainModelResult(params=params, train_data=train_data, val_data=val_data, run_name=run_name, experiment_name=experiment_name)
+        return TrainModelResult(
+            params=params,
+            train_data=train_data,
+            val_data=val_data,
+            run_name=run_name,
+            experiment_name=experiment_name,
+        )
 
     def _as_series(self, pred: pd.DataFrame | pd.Series) -> pd.Series:
         if isinstance(pred, pd.Series):
@@ -422,7 +427,9 @@ class DataManagerRepository(SientiaMonitoring):
         )
         return data
 
-    def _set_timezone_on_index(self, data: pd.DataFrame, metadata: dict[str, Any] | None = None) -> pd.DataFrame:
+    def _set_timezone_on_index(
+        self, data: pd.DataFrame, metadata: dict[str, Any] | None = None
+    ) -> pd.DataFrame:
         """
         Check if the index has a timezone and if not, set it to UTC timezone.
 
@@ -443,6 +450,7 @@ class DataManagerRepository(SientiaMonitoring):
             raise ValueError('Index is not a DatetimeIndex')
 
         return data
+
     def _get_reports_directory(self) -> str:
         """
         Get the absolute path to the reports directory.
@@ -513,23 +521,14 @@ class DataManagerRepository(SientiaMonitoring):
         if data.y_train_pred is None or data.y_pred is None:
             raise ValueError('y_train_pred or y_pred is not set, cannot generate report')
 
-
+        y_train_pred = data.y_train_pred.rename(columns={data.params.target_variable: 'prediction'})
-        y_train_pred = data.y_train_pred.rename(
+        y_val_pred = data.y_pred.rename(columns={data.params.target_variable: 'prediction'})
-            columns={data.params.target_variable: 'prediction'}
-        )
-        y_val_pred = data.y_pred.rename(
-            columns={data.params.target_variable: 'prediction'}
-        )
 
         # Join the predictions to the data
-        reference_data = y_train_pred[['prediction']].join(
+        reference_data = y_train_pred[['prediction']].join(data.train_data, how='inner')
-            data.train_data, how='inner'
-        )
         reference_data_float = reference_data.astype(np.float64)
 
-        current_data = y_val_pred[['prediction']].join(
+        current_data = y_val_pred[['prediction']].join(data.val_data, how='inner')
-            data.val_data, how='inner'
-        )
         current_data_float = current_data.astype(np.float64)
 
         # Initialize report generator

tests/activities/test_training.py

@@ -132,7 +132,12 @@ def test_train_model_success_serializes_result(mock_mlflow, training):
     training.minio_repository.download_file = MagicMock(return_value=b'csv')
     training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
     training.data_manager_repository.compute_regression_metrics = MagicMock(
-        side_effect=lambda x, _w: setattr(x, 'mse_val', 0.1) or x
+        side_effect=lambda x, _w, **_kw: (
+            setattr(x, 'mse_val', 0.1),
+            setattr(x, 'mae_val', 0.2),
+            setattr(x, 'r2_val', 0.9),
+            x,
+        )[-1]
     )
 
     def _fill_report(x, **_kw):
@@ -145,12 +150,7 @@ def test_train_model_success_serializes_result(mock_mlflow, training):
     training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report)
 
     wrapper = MagicMock()
-    wrapper.transform = MagicMock(
+    wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)])
-        side_effect=[
-            (train_df, None),
-            (val_df, None),
-        ]
-    )
     pred_train = pd.DataFrame({'p': [1.0, 2.0]})
     pred_val = pd.DataFrame({'p': [1.0]})
     wrapper.predict = MagicMock(side_effect=[(pred_train, None), (pred_val, None)])
@@ -167,12 +167,64 @@ def test_train_model_success_serializes_result(mock_mlflow, training):
     training.mlflow_repository.start_run = _run_ctx
 
     out = training.train_model({'metadata': {'pod': 'p'}, 'train_params': tp.to_dict()})
-    assert out['run_name'] == 'run-n'
+    assert out['run_name'] is None
     assert out['run_id'] == 'run-i'
     assert out['run_dir'] == '/tmp/run'
+    mock_mlflow.log_param.assert_any_call('mse_val', 0.1)
+    mock_mlflow.log_param.assert_any_call('mae_val', 0.2)
+    mock_mlflow.log_param.assert_any_call('r2_val', 0.9)
     mock_mlflow.log_artifact.assert_called()
 
 
+@patch('model_manager.activities.training.mlflow')
+def test_train_model_without_logger_does_not_set_wrapper_logger(_mock_mlflow, training):
+    """Covers branch where activity logger is None."""
+    training.logger = None
+    tp = TrainModelParams.from_dict(
+        {
+            **_minimal_params_dict(),
+            'model_metadata': {'schemas': {'components': {'schemas': {}}}},
+        }
+    )
+    train_df = pd.DataFrame({'a': [1.0, 2.0], 't': [1.0, 2.0]})
+    val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
+    tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
+
+    training.minio_repository.download_file = MagicMock(return_value=b'csv')
+    training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
+    training.data_manager_repository.compute_regression_metrics = MagicMock(
+        side_effect=lambda x, _w, **_kw: x
+    )
+
+    def _fill_report(x, **_kw):
+        x.report_path = '/tmp/report.html'
+        x.train_data_path = '/tmp/train.csv'
+        x.test_data_path = '/tmp/test.csv'
+        x.equation_path = '/tmp/eq.json'
+        x.run_dir = '/tmp/run'
+        return x
+
+    training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report)
+    wrapper = MagicMock()
+    wrapper.transform = MagicMock(side_effect=[(train_df, None), (val_df, None)])
+    wrapper.predict = MagicMock(
+        side_effect=[(pd.DataFrame({'p': [1.0, 2.0]}), None), (pd.DataFrame({'p': [1.0]}), None)]
+    )
+    wrapper.store_model = MagicMock()
+    training.plugin_store.get_model = MagicMock(return_value=wrapper)
+
+    @contextmanager
+    def _run_ctx(*_a, **_k):
+        info = MagicMock()
+        info.run_name = 'n'
+        info.run_id = 'i'
+        yield info
+
+    training.mlflow_repository.start_run = _run_ctx
+
+    training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
+
+
 @patch('model_manager.activities.training.mlflow')
 def test_train_model_train_params_as_dict(mock_mlflow, training):
     """train_params may arrive as dict and is coerced via TrainModelParams.from_dict."""
@@ -188,7 +240,7 @@ def test_train_model_train_params_as_dict(mock_mlflow, training):
     training.minio_repository.download_file = MagicMock(return_value=b'csv')
     training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
     training.data_manager_repository.compute_regression_metrics = MagicMock(
-        side_effect=lambda x, _w: x
+        side_effect=lambda x, _w, **_kw: x
     )
 
     def _fill_report2(x, **_kw):
@@ -243,7 +295,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
     training.minio_repository.download_file = MagicMock(side_effect=_dl)
     training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
     training.data_manager_repository.compute_regression_metrics = MagicMock(
-        side_effect=lambda x, _w: x
+        side_effect=lambda x, _w, **_kw: x
     )
 
     def _fill(x, **_kw):
@@ -291,7 +343,7 @@ def test_train_model_value_error_when_paths_missing_after_report(training):
     training.minio_repository.download_file = MagicMock(return_value=b'x')
     training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
     training.data_manager_repository.compute_regression_metrics = MagicMock(
-        side_effect=lambda x, _w: x
+        side_effect=lambda x, _w, **_kw: x
     )
     training.data_manager_repository.generate_report = MagicMock(return_value=tmr)
     wrapper = MagicMock()

tests/conftest.py

@@ -17,6 +17,9 @@ def _make_dummy(name: str) -> type:
 
 def _stub_evidently() -> None:
     """Minimal Evidently API surface required to import `model_manager.sientia.reports`."""
+    ev = ModuleType('evidently')
+    sys.modules['evidently'] = ev
+
     mp = ModuleType('evidently.metric_preset')
     mp.DataDriftPreset = _make_dummy('DataDriftPreset')
     sys.modules['evidently.metric_preset'] = mp
@@ -51,6 +54,13 @@ def _stub_evidently() -> None:
     opt.ColorOptions = _make_dummy('ColorOptions')
     sys.modules['evidently.options'] = opt
 
+    pipeline = ModuleType('evidently.pipeline')
+    sys.modules['evidently.pipeline'] = pipeline
+
+    colmap = ModuleType('evidently.pipeline.column_mapping')
+    colmap.ColumnMapping = _make_dummy('ColumnMapping')
+    sys.modules['evidently.pipeline.column_mapping'] = colmap
+
     rep = ModuleType('evidently.report')
     rep.Report = _make_dummy('Report')
     sys.modules['evidently.report'] = rep

tests/sientia/test_reports.py

@@ -31,9 +31,8 @@ def test_load_html_from_file_success(tmp_path):
 
 
 def test_load_html_from_file_missing_file():
-    result = reports.load_html_from_file('non-existent.html')
+    with pytest.raises(FileNotFoundError):
-
+        reports.load_html_from_file('non-existent.html')
-    assert result is None
 
 
 def test_load_html_from_file_os_error(monkeypatch):
@@ -42,9 +41,8 @@ def test_load_html_from_file_os_error(monkeypatch):
 
     monkeypatch.setattr('builtins.open', fake_open)
 
-    result = reports.load_html_from_file('path.html')
+    with pytest.raises(OSError, match='boom'):
-
+        reports.load_html_from_file('path.html')
-    assert result is None
 
 
 def test_inject_content_replaces_section():
@@ -71,7 +69,7 @@ def test_inject_content_missing_section():
 
 
 def test_reports_init_sets_defaults(stub_color_options):
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
 
     assert report.metrics == []
     assert isinstance(report.options, list) and len(report.options) == 1
@@ -89,7 +87,7 @@ def test_add_data_quality_section_without_run(monkeypatch, stub_color_options):
     monkeypatch.setattr(reports, 'ConflictTargetMetric', lambda: 'conflict')
     monkeypatch.setattr(reports, 'DatasetCorrelationsMetric', lambda: 'correlations')
 
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
     report.add_data_quality_section(columns=['col'], run=False)
 
     assert report.metrics[-4:] == [
@@ -120,7 +118,9 @@ def test_add_data_quality_section_with_run(monkeypatch, tmp_path, stub_color_opt
     ReportMock = MagicMock(return_value=report_instance)
     monkeypatch.setattr(reports, 'Report', ReportMock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur', base_path=str(tmp_path))
+    report = reports.Reports(
+        reference_data='ref', current_data='cur', target_name='target', base_path=str(tmp_path)
+    )
     report.add_data_quality_section(columns=['c1'], run=True)
 
     assert report.metrics[-4:] == [summary, column_metrics, conflict, correlations]
@@ -153,7 +153,7 @@ def test_add_data_quality_section_run_without_base_path(monkeypatch, stub_color_
     ReportMock = MagicMock(return_value=report_instance)
     monkeypatch.setattr(reports, 'Report', ReportMock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
     report.add_data_quality_section(run=True)
 
     assert report.sections['data_quality'] == {'result': 'quality'}
@@ -170,7 +170,9 @@ def test_add_data_drift_section_paths(monkeypatch, tmp_path, stub_color_options)
     ReportMock = MagicMock(return_value=report_instance)
     monkeypatch.setattr(reports, 'Report', ReportMock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur', base_path=str(tmp_path))
+    report = reports.Reports(
+        reference_data='ref', current_data='cur', target_name='target', base_path=str(tmp_path)
+    )
     report.add_data_drift_section(columns=['c1'], run=False)
     assert report.metrics[-1] == drift_instances[0]
     assert 'data_drift' not in report.sections
@@ -192,7 +194,7 @@ def test_add_data_drift_section_run_without_base_path(monkeypatch, stub_color_op
     ReportMock = MagicMock(return_value=report_instance)
     monkeypatch.setattr(reports, 'Report', ReportMock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
     report.add_data_drift_section(run=True)
 
     assert report.sections['data_drift'] == {'result': 'drift'}
@@ -216,7 +218,9 @@ def test_add_regression_section(monkeypatch, tmp_path, stub_color_options):
     ReportMock = MagicMock(return_value=report_instance)
     monkeypatch.setattr(reports, 'Report', ReportMock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur', base_path=str(tmp_path))
+    report = reports.Reports(
+        reference_data='ref', current_data='cur', target_name='target', base_path=str(tmp_path)
+    )
 
     report.add_regression_section(run=False)
     assert report.metrics[-7:] == regression_metrics
@@ -225,7 +229,11 @@ def test_add_regression_section(monkeypatch, tmp_path, stub_color_options):
     report.add_regression_section(run=True)
     assert report.sections['regression'] == {'result': 'regression'}
     ReportMock.assert_called_with(metrics=regression_metrics, options=report.options)
-    report_instance.run.assert_called_with(reference_data='ref', current_data='cur')
+    report_instance.run.assert_called_with(
+        reference_data='ref',
+        current_data='cur',
+        column_mapping=report_instance.run.call_args.kwargs['column_mapping'],
+    )
     report_instance.save_html.assert_called_with(os.path.join(str(tmp_path), 'regression.html'))
 
 
@@ -246,7 +254,7 @@ def test_add_regression_section_run_without_base_path(monkeypatch, stub_color_op
     ReportMock = MagicMock(return_value=report_instance)
     monkeypatch.setattr(reports, 'Report', ReportMock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
     report.add_regression_section(run=True)
 
     assert report.sections['regression'] == {'result': 'reg'}
@@ -262,7 +270,7 @@ def test_set_color_options_appends(monkeypatch):
 
     monkeypatch.setattr(reports, 'ColorOptions', color_options_mock)
 
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
     report.set_color_options(primary_color='#111', secondary_color='#222')
 
     assert len(report.options) == 2
@@ -272,12 +280,24 @@ def test_set_color_options_appends(monkeypatch):
 
 
 def test_save_all_sections_html_requires_base_path(stub_color_options):
-    report = reports.Reports(reference_data='ref', current_data='cur')
+    report = reports.Reports(reference_data='ref', current_data='cur', target_name='target')
 
     with pytest.raises(ValueError):
         report.save_all_sections_html('output/report.html')
 
 
+def test_save_all_sections_html_requires_template_path(stub_color_options, tmp_path):
+    report = reports.Reports(
+        reference_data='ref',
+        current_data='cur',
+        target_name='target',
+        base_path=str(tmp_path),
+    )
+
+    with pytest.raises(ValueError, match='template_path is required'):
+        report.save_all_sections_html('output/report.html')
+
+
 def test_save_all_sections_html_writes_output(tmp_path, stub_color_options):
     base_dir = tmp_path / 'templates'
     base_dir.mkdir()
@@ -289,7 +309,13 @@ def test_save_all_sections_html_writes_output(tmp_path, stub_color_options):
     (base_dir / 'data_quality.html').write_text('<p>Quality</p>', encoding='utf-8')
     (base_dir / 'regression.html').write_text('<p>Regression</p>', encoding='utf-8')
 
-    report = reports.Reports(reference_data='ref', current_data='cur', base_path=str(base_dir))
+    report = reports.Reports(
+        reference_data='ref',
+        current_data='cur',
+        target_name='target',
+        base_path=str(base_dir),
+        template_path=str(base_dir),
+    )
     output_path = tmp_path / 'reports' / 'combined.html'
 
     report.save_all_sections_html(str(output_path))
@@ -313,7 +339,13 @@ def test_save_all_sections_html_creates_directory(monkeypatch, tmp_path, stub_co
     (base_dir / 'regression.html').write_text('<p>Regression</p>', encoding='utf-8')
 
     make_dirs_called = []
-    report = reports.Reports(reference_data='ref', current_data='cur', base_path=str(base_dir))
+    report = reports.Reports(
+        reference_data='ref',
+        current_data='cur',
+        target_name='target',
+        base_path=str(base_dir),
+        template_path=str(base_dir),
+    )
     output_path = tmp_path / 'nested' / 'report.html'
     output_dir = str(output_path.parent)
 
@@ -356,7 +388,13 @@ def test_save_all_sections_html_no_directory_needed(monkeypatch, tmp_path, stub_
     monkeypatch.setattr(os, 'makedirs', fake_makedirs)
     monkeypatch.chdir(tmp_path)
 
-    report = reports.Reports(reference_data='ref', current_data='cur', base_path=str(base_dir))
+    report = reports.Reports(
+        reference_data='ref',
+        current_data='cur',
+        target_name='target',
+        base_path=str(base_dir),
+        template_path=str(base_dir),
+    )
     report.save_all_sections_html('report.html')
 
     assert mk_calls == []

tests/utils/models/test_train_model_params.py

@@ -5,7 +5,10 @@ from unittest.mock import patch
 
 import pytest
 
-from model_manager.utils.models.train_model_params import TrainModelParams
+from model_manager.utils.models.train_model_params import (
+    TrainModelParams,
+    validate_frontend_date_format,
+)
 
 
 @pytest.fixture
@@ -262,3 +265,11 @@ def test_validate_model_param_all_schema_branches(valid_train_params_dict):
     p.model_kwargs = {}
     p.opt_params = {}
     p.validate_business_rules()
+
+
+def test_validate_frontend_date_format_whitespace_returns():
+    validate_frontend_date_format('   ')
+
+
+def test_validate_frontend_date_format_valid_returns():
+    validate_frontend_date_format('dd/MM/yyyy HH:mm:ss')

tests/utils/repository/test_data_manager_repository.py

@@ -100,10 +100,23 @@ def test_prepare_training_data_empty_after_load():
     p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
     # empty csv with headers only
     csv_bytes = b'v1,t\n'
-    with pytest.raises(ValueError, match='Training data view is empty'):
+    with pytest.raises(ValueError, match='Index is not a DatetimeIndex'):
         repo.prepare_training_data(csv_bytes, None, p, {})
 
 
+def test_prepare_training_data_empty_after_transformation(monkeypatch):
+    repo = dmr.DataManagerRepository(MagicMock())
+    p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
+    monkeypatch.setattr(
+        repo,
+        '_configure_datetime_index',
+        lambda *_args, **_kwargs: pd.DataFrame(columns=['v1', 't']),
+    )
+    monkeypatch.setattr(repo, '_set_timezone_on_index', lambda data, *_args, **_kwargs: data)
+    with pytest.raises(ValueError, match='Training data view is empty after transformation'):
+        repo.prepare_training_data(b'v1,t\n', None, p, {})
+
+
 def _minimal_dict_for_prepare():
     return {
         'variable_columns': ['v1'],
@@ -364,6 +377,8 @@ def test_generate_report_success(tmp_path):
         params=p,
         train_data=pd.DataFrame({'v1': [1.0, 2.0], 't': [1.0, 2.0]}),
         val_data=pd.DataFrame({'v1': [1.0, 2.0], 't': [1.0, 2.0]}),
+        y_train_pred=pd.DataFrame({'t': [1.0, 2.0]}),
+        y_pred=pd.DataFrame({'t': [1.0, 2.0]}),
         run_name='testrun',
     )
     tmr.equation = {'target_variable': 't'}
@@ -388,6 +403,8 @@ def test_generate_report_skips_equation_file_when_not_linear(tmp_path):
         params=p,
         train_data=pd.DataFrame({'v1': [1.0, 2.0], 't': [1.0, 2.0]}),
         val_data=pd.DataFrame({'v1': [1.0, 2.0], 't': [1.0, 2.0]}),
+        y_train_pred=pd.DataFrame({'t': [1.0, 2.0]}),
+        y_pred=pd.DataFrame({'t': [1.0, 2.0]}),
         run_name='testrun',
         equation={'k': 'v'},
     )
@@ -412,6 +429,21 @@ def test_generate_report_run_name_missing():
         repo.generate_report(tmr, {})
 
 
+def test_generate_report_requires_predictions():
+    repo = dmr.DataManagerRepository(MagicMock())
+    p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
+    tmr = TrainModelResult(
+        params=p,
+        train_data=pd.DataFrame({'t': [1.0]}),
+        val_data=pd.DataFrame({'t': [1.0]}),
+        run_name='testrun',
+        y_train_pred=None,
+        y_pred=None,
+    )
+    with pytest.raises(ValueError, match='y_train_pred or y_pred is not set'):
+        repo.generate_report(tmr, {})
+
+
 def test_cleanup_run_directory_empty():
     repo = dmr.DataManagerRepository(MagicMock())
     repo.cleanup_run_directory('', {})