gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: log regression metrics as parameters in Training class

Browse Source 
- 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 Pimentel dos Santos
2026-04-17 10:52:32 -03:00
parent b245801e09
commit 31e95cbdf8
11 changed files with 649 additions and 1167 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
tests/activities/test_training.py
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()

10
tests/conftest.py
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

1103
tests/sientia/test_models.py
View File

File diff suppressed because it is too large Load Diff

78
tests/sientia/test_reports.py
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 == []

13
tests/utils/models/test_train_model_params.py
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')

34
tests/utils/repository/test_data_manager_repository.py
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('', {})