Code import - branch release/SIENTIAPDE-1645
This commit is contained in:
630
tests/activities/test_training.py
Normal file
630
tests/activities/test_training.py
Normal file
@@ -0,0 +1,630 @@
|
||||
"""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_prepare_data_observes_lag_on_success(training):
|
||||
tp = TrainModelParams.from_dict(
|
||||
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
|
||||
)
|
||||
train_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
from model_manager import metrics as mm_metrics
|
||||
|
||||
training._prepare_data(b'csv', None, tp, {})
|
||||
|
||||
training.observe_lag_sync.assert_called_once()
|
||||
call_args = training.observe_lag_sync.call_args
|
||||
assert call_args.args[1] is mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_LAG
|
||||
training.emit_metric_sync.assert_not_called()
|
||||
|
||||
|
||||
def test_prepare_data_increments_error_counter_and_still_observes_lag_on_failure(training):
|
||||
tp = TrainModelParams.from_dict(
|
||||
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
|
||||
)
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(
|
||||
side_effect=RuntimeError('prep-fail')
|
||||
)
|
||||
training.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
from model_manager import metrics as mm_metrics
|
||||
|
||||
with pytest.raises(RuntimeError, match='prep-fail'):
|
||||
training._prepare_data(b'csv', None, tp, {})
|
||||
|
||||
training.observe_lag_sync.assert_called_once()
|
||||
training.emit_metric_sync.assert_called_once()
|
||||
call_args = training.emit_metric_sync.call_args
|
||||
assert call_args.kwargs['metric_object'] is mm_metrics.SIENTIA_TRAINING_DATA_PREPARATION_ERROR_COUNT_TOTAL
|
||||
|
||||
|
||||
def test_fit_model_observes_lag_on_success(training):
|
||||
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)
|
||||
|
||||
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.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
from model_manager import metrics as mm_metrics
|
||||
|
||||
training._fit_model(wrapper, tmr, tp, {})
|
||||
|
||||
training.observe_lag_sync.assert_called_once()
|
||||
call_args = training.observe_lag_sync.call_args
|
||||
assert call_args.args[1] is mm_metrics.SIENTIA_TRAINING_MODEL_FIT_LAG
|
||||
training.emit_metric_sync.assert_not_called()
|
||||
|
||||
|
||||
def test_fit_model_increments_error_counter_on_failure(training):
|
||||
tp = TrainModelParams.from_dict(
|
||||
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
|
||||
)
|
||||
train_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
wrapper = MagicMock()
|
||||
wrapper.train = MagicMock(side_effect=RuntimeError('fit-fail'))
|
||||
training.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
from model_manager import metrics as mm_metrics
|
||||
|
||||
with pytest.raises(RuntimeError, match='fit-fail'):
|
||||
training._fit_model(wrapper, tmr, tp, {})
|
||||
|
||||
training.observe_lag_sync.assert_called_once()
|
||||
training.emit_metric_sync.assert_called_once()
|
||||
call_args = training.emit_metric_sync.call_args
|
||||
assert call_args.kwargs['metric_object'] is mm_metrics.SIENTIA_TRAINING_MODEL_FIT_ERROR_COUNT_TOTAL
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.mm_metrics')
|
||||
@patch('model_manager.activities.training.mlflow')
|
||||
def test_train_model_sets_quality_gauges_after_compute_metrics(mock_mlflow, mock_mm_metrics, training):
|
||||
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.5
|
||||
x.mae_val = 0.3
|
||||
x.r2_val = -0.1
|
||||
return x
|
||||
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=_set_metrics
|
||||
)
|
||||
|
||||
def _fill_report(x, **_kw):
|
||||
x.report_path = '/tmp/r.html'
|
||||
x.train_data_path = '/tmp/tr.csv'
|
||||
x.test_data_path = '/tmp/te.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)])
|
||||
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_id = 'rid'
|
||||
yield info
|
||||
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
training.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|
||||
|
||||
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels.return_value.set.assert_called_once_with(0.5)
|
||||
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels.return_value.set.assert_called_once_with(0.3)
|
||||
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels.return_value.set.assert_called_once_with(-0.1)
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.mm_metrics')
|
||||
@patch('model_manager.activities.training.mlflow')
|
||||
def test_train_model_skips_quality_gauges_when_none(_mock_mlflow, mock_mm_metrics, training):
|
||||
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/r.html'
|
||||
x.train_data_path = '/tmp/tr.csv'
|
||||
x.test_data_path = '/tmp/te.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)])
|
||||
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_id = 'rid'
|
||||
yield info
|
||||
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
training.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|
||||
|
||||
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MSE.labels.return_value.set.assert_not_called()
|
||||
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_MAE.labels.return_value.set.assert_not_called()
|
||||
mock_mm_metrics.SIENTIA_TRAINING_MODEL_QUALITY_R2.labels.return_value.set.assert_not_called()
|
||||
|
||||
|
||||
@patch('model_manager.activities.training.mm_metrics')
|
||||
@patch('model_manager.activities.training.mlflow')
|
||||
def test_train_model_increments_trained_total_on_success(_mock_mlflow, mock_mm_metrics, training):
|
||||
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/r.html'
|
||||
x.train_data_path = '/tmp/tr.csv'
|
||||
x.test_data_path = '/tmp/te.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)])
|
||||
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_id = 'rid'
|
||||
yield info
|
||||
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
training.observe_lag_sync = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|
||||
|
||||
training.emit_metric_sync.assert_called_once_with(
|
||||
metric_object=mock_mm_metrics.SIENTIA_TRAINING_MODEL_TRAINED_TOTAL,
|
||||
tags=training._get_training_labels(tp),
|
||||
)
|
||||
|
||||
|
||||
def test_train_model_does_not_increment_trained_total_on_failure(training):
|
||||
tp = TrainModelParams.from_dict(
|
||||
{**_minimal_params_dict(), 'model_metadata': {'schemas': {'components': {'schemas': {}}}}}
|
||||
)
|
||||
training.minio_repository.download_file = MagicMock(side_effect=RuntimeError('dl-fail'))
|
||||
training.send_notification = MagicMock()
|
||||
training.emit_metric_sync = MagicMock()
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|
||||
|
||||
training.emit_metric_sync.assert_not_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()})
|
||||
Reference in New Issue
Block a user