feat: enhance training and experiment tracking functionality
- Updated `Activities` class to improve garbage collection handling. - Enhanced error messaging in `ExperimentTracking` for better clarity on update failures. - Refactored `Training` class to streamline exception handling and improve type hints. - Introduced new methods in `TrainModelParams` for better handling of experiment run IDs and model metadata. - Added functionality to extract model equations in `DataManagerRepository` for linear regression models.
This commit is contained in:
vitor-aignosi
465
tests/utils/repository/test_data_manager_repository.py
Normal file
465
tests/utils/repository/test_data_manager_repository.py
Normal file
@@ -0,0 +1,465 @@
|
||||
"""Unit tests for DataManagerRepository and module helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
from model_manager.utils.repository import data_manager_repository as dmr
|
||||
|
||||
|
||||
def test_train_test_split_dataframe_shuffle():
|
||||
df = pd.DataFrame({'a': range(10)})
|
||||
tr, te = dmr.train_test_split(df, train_size=0.7, random_state=0, shuffle=True)
|
||||
assert len(tr) == 7 and len(te) == 3
|
||||
|
||||
|
||||
def test_train_test_split_dataframe_no_shuffle():
|
||||
df = pd.DataFrame({'a': range(10)})
|
||||
tr, te = dmr.train_test_split(df, train_size=0.5, shuffle=False)
|
||||
assert list(tr['a']) == [0, 1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_train_test_split_ndarray():
|
||||
arr = np.arange(20).reshape(10, 2)
|
||||
tr, te = dmr.train_test_split(arr, train_size=0.5, shuffle=False, random_state=None)
|
||||
assert tr.shape[0] == 5 and te.shape[0] == 5
|
||||
|
||||
|
||||
def _params(**kwargs) -> TrainModelParams:
|
||||
base = {
|
||||
'variable_columns': ['v1'],
|
||||
'target_variable': 't',
|
||||
'bucket_name': 'b',
|
||||
'file_name': 'f.csv',
|
||||
'line_separator': ',',
|
||||
'decimal_separator': '.',
|
||||
'date_column': None,
|
||||
'date_format': None,
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'random_state': 42,
|
||||
'experiment_run_id': 1,
|
||||
'model_name': 'Linear Regression',
|
||||
'val_file_name': None,
|
||||
'data_model_kwargs': {},
|
||||
'model_kwargs': {},
|
||||
'opt_params': {},
|
||||
'model_type': 'linear_regression',
|
||||
'model_id': None,
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
base.update(kwargs)
|
||||
return TrainModelParams.from_dict(base)
|
||||
|
||||
|
||||
def test_ensure_date_column_parsed_no_column():
|
||||
df = pd.DataFrame({'a': [1]})
|
||||
p = _params(date_column='missing')
|
||||
out = dmr._ensure_date_column_parsed(df, p)
|
||||
assert out is df
|
||||
|
||||
|
||||
def test_ensure_date_column_parsed_success():
|
||||
df = pd.DataFrame({'a': range(3), 'ts': ['2024-01-01 10:00:00+0000'] * 3})
|
||||
p = _params(date_column='ts')
|
||||
out = dmr._ensure_date_column_parsed(df, p)
|
||||
assert pd.api.types.is_datetime64_any_dtype(out['ts'])
|
||||
|
||||
|
||||
def test_ensure_date_column_parsed_invalid_raises():
|
||||
df = pd.DataFrame({'a': range(3), 'ts': ['not-a-date'] * 3})
|
||||
p = _params(date_column='ts', date_format='yyyy')
|
||||
with pytest.raises(ValueError, match='Failed to parse date column'):
|
||||
dmr._ensure_date_column_parsed(df, p)
|
||||
|
||||
|
||||
def test_prepare_training_data_csv_load_failure():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = _params()
|
||||
with patch(
|
||||
'model_manager.utils.repository.data_manager_repository.pd.read_csv',
|
||||
side_effect=pd.errors.ParserError('bad'),
|
||||
):
|
||||
with pytest.raises(ValueError, match='Failed to load training CSV'):
|
||||
repo.prepare_training_data(b'x', None, p, {})
|
||||
|
||||
|
||||
def test_prepare_training_data_empty_after_load():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
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'):
|
||||
repo.prepare_training_data(csv_bytes, None, p, {})
|
||||
|
||||
|
||||
def _minimal_dict_for_prepare():
|
||||
return {
|
||||
'variable_columns': ['v1'],
|
||||
'target_variable': 't',
|
||||
'bucket_name': 'b',
|
||||
'file_name': 'f.csv',
|
||||
'line_separator': ',',
|
||||
'decimal_separator': '.',
|
||||
'date_column': None,
|
||||
'date_format': None,
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'random_state': 42,
|
||||
'experiment_run_id': 1,
|
||||
'model_name': 'Linear Regression',
|
||||
'val_file_name': None,
|
||||
'data_model_kwargs': {},
|
||||
'model_kwargs': {},
|
||||
'opt_params': {},
|
||||
'model_type': 'linear_regression',
|
||||
'model_id': None,
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
|
||||
|
||||
def _csv_bytes_with_ts(n_rows: int = 20) -> bytes:
|
||||
"""CSV with leading timestamp column so _configure_datetime_index does not mangle feature columns."""
|
||||
lines = ['timestamp,v1,t']
|
||||
for i in range(n_rows):
|
||||
lines.append(f'2024-01-{i + 1:02d} 00:00:00+0000,{i},{i + 1}')
|
||||
return '\n'.join(lines).encode()
|
||||
|
||||
|
||||
def test_prepare_training_data_validation_csv_invalid():
|
||||
from io import BytesIO
|
||||
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
train_csv = _csv_bytes_with_ts(5)
|
||||
train_df = pd.read_csv(BytesIO(train_csv), sep=',', decimal='.')
|
||||
with patch.object(
|
||||
dmr.pd,
|
||||
'read_csv',
|
||||
side_effect=[train_df, pd.errors.ParserError('bad val')],
|
||||
):
|
||||
with pytest.raises(ValueError, match='Failed to load validation CSV'):
|
||||
repo.prepare_training_data(train_csv, b'broken', p, {})
|
||||
|
||||
|
||||
def test_prepare_training_data_validation_empty_val():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
train_csv = _csv_bytes_with_ts(5)
|
||||
val_csv = b'timestamp,v1,t\n'
|
||||
with pytest.raises(ValueError, match='Validation data view is empty'):
|
||||
repo.prepare_training_data(train_csv, val_csv, p, {})
|
||||
|
||||
|
||||
def test_prepare_training_data_split_path():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
train_csv = _csv_bytes_with_ts(20)
|
||||
res = repo.prepare_training_data(train_csv, None, p, {})
|
||||
assert res.train_data is not None and res.val_data is not None
|
||||
|
||||
|
||||
def test_prepare_training_data_explicit_validation_success():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
train_csv = _csv_bytes_with_ts(10)
|
||||
val_csv = _csv_bytes_with_ts(5)
|
||||
res = repo.prepare_training_data(train_csv, val_csv, p, {})
|
||||
assert len(res.val_data) == 5
|
||||
|
||||
|
||||
def test_as_series_series():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
s = pd.Series([1.0, 2.0])
|
||||
assert repo._as_series(s).equals(s)
|
||||
|
||||
|
||||
def test_as_series_one_column_df():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
df = pd.DataFrame({'x': [1.0, 2.0]})
|
||||
out = repo._as_series(df)
|
||||
assert isinstance(out, pd.Series)
|
||||
|
||||
|
||||
def test_as_series_multi_column_raises():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
df = pd.DataFrame({'a': [1.0], 'b': [2.0]})
|
||||
with pytest.raises(ValueError, match='single-column'):
|
||||
repo._as_series(df)
|
||||
|
||||
|
||||
def test_compute_regression_metrics_requires_y_pred():
|
||||
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]}),
|
||||
y_pred=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match='y_pred must be set'):
|
||||
repo.compute_regression_metrics(tmr, MagicMock())
|
||||
|
||||
|
||||
def test_compute_regression_metrics_no_overlap():
|
||||
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]}, index=[10]),
|
||||
y_pred=pd.DataFrame({'p': [1.0]}, index=[20]),
|
||||
)
|
||||
with pytest.raises(ValueError, match='No overlapping indices'):
|
||||
repo.compute_regression_metrics(tmr, MagicMock())
|
||||
|
||||
|
||||
def test_compute_regression_metrics_linear_equation():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
p.model_type = 'linear_regression'
|
||||
idx = pd.Index([0, 1])
|
||||
tmr = TrainModelResult(
|
||||
params=p,
|
||||
train_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
|
||||
val_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
|
||||
y_pred=pd.DataFrame({'p': [1.0, 2.0]}, index=idx),
|
||||
)
|
||||
regr = MagicMock()
|
||||
regr.coef_ = np.array([0.5])
|
||||
regr.intercept_ = 1.0
|
||||
wrapper = MagicMock()
|
||||
wrapper.model = MagicMock()
|
||||
wrapper.model.regr = regr
|
||||
out = repo.compute_regression_metrics(tmr, wrapper)
|
||||
assert out.mse_val is not None and out.equation is not None
|
||||
|
||||
|
||||
def test_compute_regression_metrics_non_linear_skips_equation():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
p.model_type = 'xgboost'
|
||||
idx = pd.Index([0, 1])
|
||||
tmr = TrainModelResult(
|
||||
params=p,
|
||||
train_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
|
||||
val_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
|
||||
y_pred=pd.DataFrame({'p': [1.0, 2.0]}, index=idx),
|
||||
)
|
||||
out = repo.compute_regression_metrics(tmr, MagicMock())
|
||||
assert out.mse_val is not None and out.equation is None
|
||||
|
||||
|
||||
def test_configure_datetime_index_none_raises():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
with pytest.raises(ValueError, match='Data is None'):
|
||||
repo._configure_datetime_index(None, p, {})
|
||||
|
||||
|
||||
def test_configure_datetime_index_already_datetime_index():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
idx = pd.date_range('2024-01-01', periods=3, freq='h')
|
||||
df = pd.DataFrame({'v1': [1, 2, 3], 't': [1, 2, 3]}, index=idx)
|
||||
out = repo._configure_datetime_index(df, p, {})
|
||||
assert isinstance(out.index, pd.DatetimeIndex)
|
||||
|
||||
|
||||
def test_configure_datetime_index_from_common_column():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
df = pd.DataFrame({'timestamp': pd.date_range('2024-01-01', periods=3, freq='D'), 'v1': [1, 2, 3], 't': [1, 2, 3]})
|
||||
out = repo._configure_datetime_index(df, p, {})
|
||||
assert isinstance(out.index, pd.DatetimeIndex)
|
||||
|
||||
|
||||
def test_configure_datetime_index_from_date_column():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict({**_minimal_dict_for_prepare(), 'date_column': 'mydate'})
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
'mydate': pd.date_range('2024-01-01', periods=3, freq='D'),
|
||||
'v1': [1, 2, 3],
|
||||
't': [1, 2, 3],
|
||||
}
|
||||
)
|
||||
out = repo._configure_datetime_index(df, p, {})
|
||||
assert isinstance(out.index, pd.DatetimeIndex)
|
||||
|
||||
|
||||
def test_configure_datetime_index_bad_column_skips_to_first():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
df = pd.DataFrame({'timestamp': ['x'], 'v1': [1.0], 't': [1.0]})
|
||||
out = repo._configure_datetime_index(df, p, {})
|
||||
assert isinstance(out, pd.DataFrame)
|
||||
assert not isinstance(out.index, pd.DatetimeIndex)
|
||||
|
||||
|
||||
def test_configure_datetime_index_no_timestamp_warning():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
df = pd.DataFrame({'v1': [1, 2, 3], 't': [1, 2, 3]})
|
||||
out = repo._configure_datetime_index(df, p, {})
|
||||
assert isinstance(out, pd.DataFrame)
|
||||
|
||||
|
||||
def test_configure_datetime_index_first_column_numeric_parsed_as_time():
|
||||
"""Covers fallback path that parses the first column as datetime when it looks like timestamps."""
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
'ts': pd.date_range('2024-01-01', periods=3, freq='D'),
|
||||
'v1': [1.0, 2.0, 3.0],
|
||||
't': [1.0, 2.0, 3.0],
|
||||
}
|
||||
)
|
||||
out = repo._configure_datetime_index(df, p, {})
|
||||
assert isinstance(out.index, pd.DatetimeIndex)
|
||||
|
||||
|
||||
def test_create_run_directory_permission_error():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
with patch('model_manager.utils.repository.data_manager_repository.makedirs', side_effect=PermissionError('no')):
|
||||
with pytest.raises(PermissionError, match='Permission denied'):
|
||||
repo._create_run_directory('/tmp', 'run', {})
|
||||
|
||||
|
||||
def test_create_run_directory_os_error():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
with patch('model_manager.utils.repository.data_manager_repository.makedirs', side_effect=OSError('disk')):
|
||||
with pytest.raises(OSError, match='Failed to create directory'):
|
||||
repo._create_run_directory('/tmp', 'run', {})
|
||||
|
||||
|
||||
def test_generate_report_success(tmp_path):
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
tmr = TrainModelResult(
|
||||
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]}),
|
||||
run_name='testrun',
|
||||
)
|
||||
tmr.equation = {'target_variable': 't'}
|
||||
with (
|
||||
patch.object(repo, '_get_reports_directory', return_value=str(tmp_path)),
|
||||
patch('model_manager.utils.repository.data_manager_repository.Reports') as mrep,
|
||||
):
|
||||
inst = mrep.return_value
|
||||
instance = mrep.return_value
|
||||
instance.save_all_sections_html = Mock()
|
||||
out = repo.generate_report(tmr, {})
|
||||
assert out.report_path and out.train_data_path and out.test_data_path
|
||||
if out.equation_path:
|
||||
with open(out.equation_path, encoding='utf-8') as f:
|
||||
json.load(f)
|
||||
|
||||
|
||||
def test_generate_report_skips_equation_file_when_not_linear(tmp_path):
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
p.model_type = 'other'
|
||||
tmr = TrainModelResult(
|
||||
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]}),
|
||||
run_name='testrun',
|
||||
equation={'k': 'v'},
|
||||
)
|
||||
with (
|
||||
patch.object(repo, '_get_reports_directory', return_value=str(tmp_path)),
|
||||
patch('model_manager.utils.repository.data_manager_repository.Reports'),
|
||||
):
|
||||
out = repo.generate_report(tmr, {})
|
||||
assert out.equation_path is None
|
||||
|
||||
|
||||
def test_generate_report_run_name_missing():
|
||||
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=None,
|
||||
)
|
||||
with pytest.raises(ValueError, match='run_name is not set'):
|
||||
repo.generate_report(tmr, {})
|
||||
|
||||
|
||||
def test_cleanup_run_directory_empty():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
repo.cleanup_run_directory('', {})
|
||||
|
||||
|
||||
def test_cleanup_run_directory_exists(tmp_path):
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
d = tmp_path / 'subdir'
|
||||
d.mkdir()
|
||||
repo.cleanup_run_directory(str(d), {})
|
||||
assert not d.exists()
|
||||
|
||||
|
||||
def test_cleanup_run_directory_missing(tmp_path):
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
repo.cleanup_run_directory(str(tmp_path / 'nope'), {})
|
||||
|
||||
|
||||
def test_extract_model_equation_polynomial_poly_names():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
p.model_kwargs = {'degree': 2, 'poly_feature_names': ['f1', 'f2']}
|
||||
regr = MagicMock()
|
||||
regr.coef_ = np.array([1.0, 2.0])
|
||||
regr.intercept_ = 3.0
|
||||
wrapper = MagicMock()
|
||||
wrapper.model = MagicMock()
|
||||
wrapper.model.regr = regr
|
||||
eq = repo._extract_model_equation(wrapper.model, p)
|
||||
assert 'equation_string' in eq and eq['degree'] == 2
|
||||
|
||||
|
||||
def test_extract_model_equation_extra_coefficients_ignored():
|
||||
"""More coefficients than feature names: only the first len(names) are used."""
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
regr = MagicMock()
|
||||
regr.coef_ = np.array([1.0, 2.0, 3.0])
|
||||
regr.intercept_ = 0.0
|
||||
wrapper = MagicMock()
|
||||
wrapper.model = MagicMock()
|
||||
wrapper.model.regr = regr
|
||||
eq = repo._extract_model_equation(wrapper.model, p)
|
||||
assert len(eq['coefficients']) == len(p.variable_columns)
|
||||
|
||||
|
||||
def test_extract_model_equation_more_features_than_coefficients():
|
||||
"""Polynomial feature names longer than coef array: extra names get no coefficient entry."""
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
|
||||
p.model_kwargs = {'degree': 2, 'poly_feature_names': ['a', 'b', 'c']}
|
||||
regr = MagicMock()
|
||||
regr.coef_ = np.array([1.0, 2.0])
|
||||
regr.intercept_ = 0.0
|
||||
wrapper = MagicMock()
|
||||
wrapper.model = MagicMock()
|
||||
wrapper.model.regr = regr
|
||||
eq = repo._extract_model_equation(wrapper.model, p)
|
||||
assert list(eq['coefficients'].keys()) == ['a', 'b']
|
||||
|
||||
|
||||
def test_get_reports_directory_path():
|
||||
repo = dmr.DataManagerRepository(MagicMock())
|
||||
reports_dir = repo._get_reports_directory()
|
||||
assert reports_dir.endswith('reports')
|
||||
assert 'model_manager' in reports_dir
|
||||
Reference in New Issue
Block a user