SIENTIAPDE-1171
Update GITHUB_BRANCH in values.yaml and refactor MLFlowRepository for model retraining - Changed GITHUB_BRANCH in values.yaml to reflect the new pipeline for retraining. - Refactored MLFlowRepository by introducing create_model_experiment and perform_model_retrain methods to streamline model retraining processes. - Updated test cases to cover new functionalities and ensure proper error handling in model management workflows.
This commit is contained in:
vitor-aignosi
@@ -1,6 +1,8 @@
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
from unittest.mock import ANY, MagicMock, call, patch
|
||||
import numpy as np
|
||||
from pandas import DataFrame
|
||||
import pytest
|
||||
from laborious.utils.repository import model_repository
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
@@ -92,3 +94,245 @@ def test_predict_error(mlflow_repository):
|
||||
'traceback': ANY
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_by_run_id(mlflow, mlflow_repository):
|
||||
mlflow.get_run.return_value = MagicMock(
|
||||
info=MagicMock(
|
||||
experiment_id='0',
|
||||
)
|
||||
)
|
||||
mlflow.get_experiment.return_value = MagicMock()
|
||||
mlflow.get_experiment.return_value.name = 'test'
|
||||
|
||||
output = mlflow_repository.get_experiment_by_run_id('0')
|
||||
assert output == 'test'
|
||||
mlflow.get_run.assert_called_once_with('0')
|
||||
mlflow.get_experiment.assert_called_once_with('0')
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_next_run_name(mlflow, mlflow_repository):
|
||||
mlflow.search_runs.return_value = [1, 2, 3]
|
||||
output = mlflow_repository.get_next_run_name('run')
|
||||
assert output == 'run-4'
|
||||
mlflow.search_runs.assert_called_once_with(
|
||||
experiment_names=['run'],
|
||||
order_by=['start_time desc'],
|
||||
)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_success(mlflow, mlflow_repository):
|
||||
mlflow.get_experiment_by_name.return_value = MagicMock(
|
||||
experiment_id='0')
|
||||
|
||||
output = mlflow_repository.get_experiment('test')
|
||||
|
||||
assert output == 0
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_error(mlflow, mlflow_repository):
|
||||
mlflow.get_experiment_by_name.return_value = None
|
||||
|
||||
try:
|
||||
mlflow_repository.get_experiment('test')
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Experiment test not found'
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_last_run(mlflow, mlflow_repository):
|
||||
mlflow.search_runs.return_value = DataFrame({
|
||||
'params.retrain': ['True', 'False', 'True', 'False'],
|
||||
'end_time': ['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04'],
|
||||
'run_id': ['0', '1', '2', '3'],
|
||||
})
|
||||
|
||||
output = mlflow_repository.get_experiment_last_run(0)
|
||||
|
||||
mlflow.search_runs.assert_called_once_with(
|
||||
experiment_ids=[0],
|
||||
filter_string="",
|
||||
output_format="pandas",
|
||||
)
|
||||
|
||||
assert output == '2'
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.sklearn')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.set_experiment')
|
||||
def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
|
||||
|
||||
mlflow_repository.model_serving.get_model_run_id = MagicMock(
|
||||
return_value='0')
|
||||
mlflow_repository.model_serving.get_model_uri = MagicMock(
|
||||
return_value='test')
|
||||
mlflow_repository.get_experiment_by_run_id = MagicMock()
|
||||
|
||||
data_model_mock = MagicMock()
|
||||
prediction_model_mock = MagicMock()
|
||||
|
||||
sklearn.load_model.side_effect = [data_model_mock, prediction_model_mock]
|
||||
|
||||
data_model_mock.fit.return_value = data_model_mock
|
||||
data_model_mock.predict.return_value = DataFrame({
|
||||
'x': [10, 20, 30],
|
||||
})
|
||||
data_model_mock.target_variable = 'y'
|
||||
|
||||
prediction_model_mock.fit.return_value = prediction_model_mock
|
||||
|
||||
data = DataFrame({
|
||||
'x': [1, 2, 3],
|
||||
'y': [4, 5, 6]
|
||||
})
|
||||
|
||||
output = mlflow_repository.create_model_experiment('test', data)
|
||||
|
||||
mlflow_repository.model_serving.get_model_run_id.assert_called_once_with(
|
||||
'test', stage='Production')
|
||||
mlflow_repository.model_serving.get_model_uri.assert_called_once_with(
|
||||
'0', prediction=False)
|
||||
|
||||
sklearn.load_model.assert_has_calls([
|
||||
call(mlflow_repository.model_serving.get_model_uri.return_value),
|
||||
call("models:/test/production"),
|
||||
])
|
||||
assert sklearn.load_model.call_count == 2
|
||||
|
||||
data_model_mock.fit.assert_called_once_with(data)
|
||||
data_model_mock.predict.assert_called_once_with(data)
|
||||
|
||||
fit_args = prediction_model_mock.fit.call_args[0][0]
|
||||
assert fit_args.equals(
|
||||
DataFrame({
|
||||
'x': [10, 20, 30],
|
||||
'y': [4, 5, 6],
|
||||
})
|
||||
)
|
||||
|
||||
mlflow_repository.get_experiment_by_run_id.assert_called_once_with('0')
|
||||
|
||||
set_experiment.assert_called_once_with(
|
||||
mlflow_repository.get_experiment_by_run_id.return_value
|
||||
)
|
||||
|
||||
assert output == (prediction_model_mock,
|
||||
data_model_mock,
|
||||
mlflow_repository.get_experiment_by_run_id.return_value)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.start_run')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.log_param')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.sklearn.log_model')
|
||||
@patch('laborious.utils.repository.model_repository.mlflow.log_artifact')
|
||||
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
|
||||
|
||||
_os_remove = patch(
|
||||
'laborious.utils.repository.model_repository.remove')
|
||||
_os_makedirs = patch(
|
||||
'laborious.utils.repository.model_repository.makedirs')
|
||||
_os_path_exists = patch(
|
||||
'laborious.utils.repository.model_repository.path.exists',
|
||||
MagicMock(return_value=False))
|
||||
|
||||
prediction_model_mock = MagicMock()
|
||||
data_model_mock = MagicMock()
|
||||
experiment = 'test'
|
||||
model_name = 'test'
|
||||
data = MagicMock()
|
||||
|
||||
mlflow_repository.get_next_run_name = MagicMock(
|
||||
return_value='test-1')
|
||||
run = MagicMock()
|
||||
start_run.__enter__.return_value = run
|
||||
|
||||
output = mlflow_repository.perform_model_retrain(
|
||||
prediction_model_mock, data_model_mock, experiment, model_name, data)
|
||||
|
||||
mlflow_repository.get_next_run_name.assert_called_once_with(experiment)
|
||||
start_run.assert_called_once_with(
|
||||
run_name='test-1', description='Retrain model test with new data')
|
||||
|
||||
log_model.assert_has_calls([
|
||||
call(data_model_mock, "data_model"),
|
||||
call(prediction_model_mock, "prediction_model"),
|
||||
])
|
||||
|
||||
data.to_csv.assert_called_once_with(
|
||||
"temp/raw_data_test.csv", index=True)
|
||||
|
||||
log_artifact.assert_called_once_with(
|
||||
"temp/raw_data_test.csv")
|
||||
|
||||
log_param.assert_has_calls([
|
||||
call("retrain", True),
|
||||
])
|
||||
|
||||
assert output == ("Model retrained successfully", experiment)
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
|
||||
client_mock = MagicMock()
|
||||
mlflow.tracking.MlflowClient.return_value = client_mock
|
||||
|
||||
client_mock.get_registered_model.return_value = MagicMock(
|
||||
latest_versions=[
|
||||
MagicMock(version='1'),
|
||||
MagicMock(version='2'),
|
||||
MagicMock(version='3'),
|
||||
]
|
||||
)
|
||||
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
|
||||
|
||||
mlflow.register_model.assert_called_once_with(
|
||||
"runs:/0/prediction_model",
|
||||
'test',
|
||||
)
|
||||
|
||||
mlflow.tracking.MlflowClient.assert_called_once()
|
||||
client_mock.get_registered_model.assert_called_once_with('test')
|
||||
client_mock.transition_model_version_stage.assert_called_once_with(
|
||||
name='test',
|
||||
version='3',
|
||||
stage='Production',
|
||||
archive_existing_versions=True,
|
||||
)
|
||||
|
||||
assert output == {
|
||||
'model_name': 'test',
|
||||
'version': '3',
|
||||
'mlflow_run_id': '0',
|
||||
}
|
||||
|
||||
|
||||
def test_update_production_model(mlflow_repository):
|
||||
connector = mlflow_repository
|
||||
|
||||
with patch.object(connector, 'get_experiment',
|
||||
return_value='0') as get_experiment:
|
||||
with patch.object(connector, 'get_experiment_last_run',
|
||||
return_value='2') as get_experiment_last_run:
|
||||
with patch.object(connector, 'update_production_model_by_run_id',
|
||||
return_value={'model_name': 'test', 'version': '3',
|
||||
'mlflow_run_id': '0'}) as update_production_model_by_run_id:
|
||||
|
||||
output = connector.update_production_model('0', 'test')
|
||||
|
||||
get_experiment.assert_called_once_with('0')
|
||||
get_experiment_last_run.assert_called_once_with('0')
|
||||
update_production_model_by_run_id.assert_called_once_with(
|
||||
'2', 'test')
|
||||
|
||||
assert output == {
|
||||
'model_name': 'test',
|
||||
'version': '3',
|
||||
'mlflow_run_id': '0',
|
||||
'mlflow_experiment_id': '0',
|
||||
}
|
||||
|
||||
@@ -139,10 +139,12 @@ def test_build_mongo_db_config_with_env_vars():
|
||||
environ['MONGODB_PASSWORD'] = 'sientia1'
|
||||
environ['MONGODB_URL'] = 'localhost:27018'
|
||||
environ['MONGODB_DATABASE_NAME'] = 'test_db'
|
||||
environ['MONGODB_TTL_INDEX_HOURS'] = '1'
|
||||
|
||||
assert build_mongodb_config() == {
|
||||
'connection_string': 'mongodb://sientia1:sientia1@localhost:27018',
|
||||
'database_name': 'test_db'
|
||||
'database_name': 'test_db',
|
||||
'ttl_index_seconds': 3600
|
||||
}
|
||||
|
||||
|
||||
@@ -151,8 +153,9 @@ def test_build_mongo_db_config_with_defaults():
|
||||
environ.pop('MONGODB_PASSWORD', None)
|
||||
environ.pop('MONGODB_DATABASE_NAME', None)
|
||||
environ.pop('MONGODB_URL', None)
|
||||
|
||||
environ.pop('MONGODB_TTL_INDEX_HOURS', None)
|
||||
assert build_mongodb_config() == {
|
||||
'connection_string': 'mongodb://sientia:sientia@localhost:27017',
|
||||
'database_name': 'sientia'
|
||||
'connection_string': 'mongodb://root:wKZDbMNU1c@localhost:27018',
|
||||
'database_name': 'sientia',
|
||||
'ttl_index_seconds': 3600
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user