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.
This commit is contained in:
vitor-aignosi
2026-04-17 10:52:32 -03:00
parent b245801e09
commit 31e95cbdf8
11 changed files with 649 additions and 1167 deletions

View File

@@ -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`

View File

@@ -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:

View File

@@ -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,17 +36,18 @@ 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)
# Verifica se a seção foi encontrada E se ela é uma Tag (não uma string)
if section and isinstance(section, Tag):
section.clear()
# Converte o conteúdo para um fragmento de BeautifulSoup e anexa
new_content = BeautifulSoup(content, 'html.parser')
section.append(new_content)
return str(soup)
@@ -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.
@@ -151,7 +157,7 @@ class Reports:
mapping.target = self.target_name
mapping.prediction = 'prediction'
report = Report(metrics=metrics, options=self.options)
report.run(
reference_data=self.ref_data,
@@ -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"Report path: {report_path}")
print(f"Base path: {self.base_path}")
print(f"Template path: {self.template_path}")
print(f'Output directory: {output_dir}')
print(f'Report path: {report_path}')
print(f'Base path: {self.base_path}')
print(f'Template path: {self.template_path}')
# Load main HTML template
main_html_path = os.path.join(self.template_path, 'header.html')

View File

@@ -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'

View File

@@ -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
@@ -160,7 +159,7 @@ class DataManagerRepository(SientiaMonitoring):
train_df = _ensure_date_column_parsed(train_df, params)
train_df = self._configure_datetime_index(train_df, params, metadata)
train_df = self._set_timezone_on_index(train_df, metadata)
if len(train_df) <= 0:
raise ValueError('Training data view is empty after transformation')
@@ -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.
@@ -441,8 +448,9 @@ class DataManagerRepository(SientiaMonitoring):
data.index = data.index.tz_convert('UTC')
else:
raise ValueError('Index is not a DatetimeIndex')
return data
def _get_reports_directory(self) -> str:
"""
Get the absolute path to the reports directory.
@@ -513,29 +521,20 @@ 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_val_pred = data.y_pred.rename(columns={data.params.target_variable: 'prediction'})
y_train_pred = data.y_train_pred.rename(
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(
data.train_data, how='inner'
)
reference_data = y_train_pred[['prediction']].join(data.train_data, how='inner')
reference_data_float = reference_data.astype(np.float64)
current_data = y_val_pred[['prediction']].join(
data.val_data, how='inner'
)
current_data = y_val_pred[['prediction']].join(data.val_data, how='inner')
current_data_float = current_data.astype(np.float64)
# Initialize report generator
base_path = self._get_reports_directory()
data.run_dir = self._create_run_directory(base_path, data.run_name)
# Template path is the code path of the model_manager package
template_path = path.join(PROJECT_BASE_PATH, 'reports')

View File

@@ -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(
side_effect=[
(train_df, None),
(val_df, None),
]
)
wrapper.transform = MagicMock(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()

View File

@@ -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

File diff suppressed because it is too large Load Diff

View File

@@ -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')
assert result is None
with pytest.raises(FileNotFoundError):
reports.load_html_from_file('non-existent.html')
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')
assert result is None
with pytest.raises(OSError, match='boom'):
reports.load_html_from_file('path.html')
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 == []

View File

@@ -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')

View File

@@ -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('', {})