d1f9394879
- Added a new fixture to manage runtime report artifacts in a writable temp directory during E2E tests, addressing permission issues in local CI/dev environments. - Updated `conftest.py` to include a requirements.txt file in the model packaging path for training activities. - Refactored existing fixtures to use `pytest.fixture` instead of `pytest_asyncio.fixture` for better compatibility. - Enhanced the `Reports` class to include a target alias for report metrics, ensuring compatibility with Evidently's reporting requirements. - Introduced new test scenarios to validate the handling of missing and whitespace-only `date_column` inputs in the training workflow. These changes improve the robustness of the E2E testing framework and enhance the clarity of model reporting metrics.
370 lines
14 KiB
Python
370 lines
14 KiB
Python
"""Unit tests for Training activities."""
|
|
|
|
from contextlib import contextmanager
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
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
|
|
|
|
|
|
def _minimal_params_dict():
|
|
return {
|
|
'variable_columns': ['a'],
|
|
'target_variable': 't',
|
|
'bucket_name': 'b',
|
|
'file_name': 'f.csv',
|
|
'line_separator': '\n',
|
|
'decimal_separator': '.',
|
|
'date_column': 'timestamp',
|
|
'date_format': 'yyyy-MM-dd HH:mm:ss',
|
|
'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': None,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def training():
|
|
from model_manager.activities.training import Training
|
|
|
|
return Training(
|
|
mlflow_repository=MagicMock(),
|
|
plugin_store=MagicMock(),
|
|
minio_repository=MagicMock(),
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
metrics_controller=MagicMock(),
|
|
)
|
|
|
|
|
|
def test_load_model_metadata_success(training):
|
|
training.plugin_store.get_model_index = MagicMock(
|
|
return_value={'schemas': {'components': {'schemas': {}}}}
|
|
)
|
|
inp = {**_minimal_params_dict(), 'metadata': {'w': '1'}}
|
|
out = training.load_model_metadata(inp)
|
|
assert 'model_metadata' in out
|
|
assert out['model_metadata']['schemas']
|
|
|
|
|
|
def test_load_model_metadata_notifies_on_error(training):
|
|
training.plugin_store.get_model_index = MagicMock(side_effect=RuntimeError('idx'))
|
|
training.send_notification = MagicMock()
|
|
inp = {**_minimal_params_dict(), 'metadata': {}}
|
|
with pytest.raises(RuntimeError, match='idx'):
|
|
training.load_model_metadata(inp)
|
|
training.send_notification.assert_called_once()
|
|
|
|
|
|
def test_validate_train_params_success(training):
|
|
pdict = _minimal_params_dict()
|
|
pdict['model_metadata'] = {'schemas': {'components': {'schemas': {}}}}
|
|
inp = {**pdict, 'metadata': {}}
|
|
out = training.validate_train_params(inp)
|
|
assert isinstance(out, dict)
|
|
assert out['target_variable'] == 't'
|
|
|
|
|
|
def test_validate_train_params_notifies(training):
|
|
training.send_notification = MagicMock()
|
|
inp = {'metadata': {}, 'experiment_run_id': 1}
|
|
with pytest.raises((KeyError, ValueError, TypeError)):
|
|
training.validate_train_params(inp)
|
|
training.send_notification.assert_called_once()
|
|
|
|
|
|
def test_train_model_download_fails_notifies(training):
|
|
"""train_model notifies and re-raises when MinIO download fails."""
|
|
tp = TrainModelParams.from_dict(
|
|
{
|
|
**_minimal_params_dict(),
|
|
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
|
}
|
|
)
|
|
training.minio_repository.download_file = MagicMock(side_effect=OSError('minio'))
|
|
training.send_notification = MagicMock()
|
|
with pytest.raises(OSError, match='minio'):
|
|
training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp.to_dict()})
|
|
training.send_notification.assert_called_once()
|
|
|
|
|
|
def test_cleanup_resources(training):
|
|
training.data_manager_repository.cleanup_run_directory = MagicMock()
|
|
training.cleanup_resources({'metadata': {}, 'run_dir': '/tmp/x'})
|
|
training.data_manager_repository.cleanup_run_directory.assert_called_once_with('/tmp/x', {})
|
|
|
|
|
|
def test_cleanup_resources_notifies_on_error(training):
|
|
training.data_manager_repository.cleanup_run_directory = MagicMock(
|
|
side_effect=RuntimeError('rm')
|
|
)
|
|
training.send_notification = MagicMock()
|
|
with pytest.raises(RuntimeError, match='rm'):
|
|
training.cleanup_resources({'metadata': {'pod': 'p'}, 'run_dir': '/tmp/x'})
|
|
training.send_notification.assert_called_once()
|
|
|
|
|
|
@patch('model_manager.activities.training.mlflow')
|
|
def test_train_model_success_serializes_result(mock_mlflow, training):
|
|
"""Exercise train_model happy path with mocks (MinIO, plugin wrapper, MLflow)."""
|
|
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)
|
|
|
|
def _set_metrics(x, _w, **_kw):
|
|
x.mse_val = 0.1
|
|
x.mae_val = 0.2
|
|
x.r2_val = 0.9
|
|
return x
|
|
|
|
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
|
side_effect=_set_metrics
|
|
)
|
|
|
|
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.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)])
|
|
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)])
|
|
wrapper.store_model = MagicMock()
|
|
training.plugin_store.get_model = MagicMock(return_value=wrapper)
|
|
|
|
@contextmanager
|
|
def _run_ctx(*_a, **_k):
|
|
info = MagicMock()
|
|
info.run_name = 'run-n'
|
|
info.run_id = 'run-i'
|
|
yield info
|
|
|
|
training.mlflow_repository.start_run = _run_ctx
|
|
|
|
out = training.train_model({'metadata': {'pod': 'p'}, 'train_params': tp.to_dict()})
|
|
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."""
|
|
d = {
|
|
**_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]})
|
|
tp = TrainModelParams.from_dict(d)
|
|
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_report2(x, **_kw):
|
|
x.report_path = '/tmp/report.html'
|
|
x.train_data_path = '/tmp/train.csv'
|
|
x.test_data_path = '/tmp/test.csv'
|
|
x.run_dir = '/tmp/run'
|
|
return x
|
|
|
|
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill_report2)
|
|
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': d})
|
|
mock_mlflow.log_artifact.assert_called()
|
|
|
|
|
|
@patch('model_manager.activities.training.mlflow')
|
|
def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
|
|
"""Second MinIO download when val_file_name is set (covers val_bytes branch)."""
|
|
d = {
|
|
**_minimal_params_dict(),
|
|
'val_file_name': 'val.csv',
|
|
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
|
}
|
|
tp = TrainModelParams.from_dict(d)
|
|
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)
|
|
|
|
def _dl(object_name, **_kwargs):
|
|
if object_name == tp.file_name:
|
|
return b'train'
|
|
if object_name == 'val.csv':
|
|
return b'val'
|
|
raise AssertionError(object_name)
|
|
|
|
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, **_kw: x
|
|
)
|
|
|
|
def _fill(x, **_kw):
|
|
x.report_path = '/tmp/report.html'
|
|
x.train_data_path = '/tmp/train.csv'
|
|
x.test_data_path = '/tmp/test.csv'
|
|
x.run_dir = '/tmp/run'
|
|
return x
|
|
|
|
training.data_manager_repository.generate_report = MagicMock(side_effect=_fill)
|
|
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()})
|
|
assert training.minio_repository.download_file.call_count == 2
|
|
mock_mlflow.log_artifact.assert_called()
|
|
|
|
|
|
def test_train_model_value_error_when_paths_missing_after_report(training):
|
|
"""Raises ValueError when report paths are not populated after generate_report."""
|
|
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'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, **_kw: x
|
|
)
|
|
training.data_manager_repository.generate_report = MagicMock(return_value=tmr)
|
|
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)]
|
|
)
|
|
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.send_notification = MagicMock()
|
|
|
|
with pytest.raises(ValueError, match='Report path'):
|
|
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|