SIENTIAPDE-1171
Refactor model_repository and enhance test coverage for MLFlow functionalities - Updated model_repository to ensure the 'temp' directory is created if it doesn't exist using `exist_ok=True`. - Added new tests for retraining and updating production models, including error handling scenarios. - Improved existing tests for model management workflows to ensure robustness and reliability.
This commit is contained in:
vitor-aignosi
@@ -168,8 +168,7 @@ class MLFlowRepository():
|
||||
# dynamic parameters, including model itself
|
||||
mlflow.sklearn.log_model(data_model, "data_model")
|
||||
|
||||
if not path.exists("temp"):
|
||||
makedirs("temp")
|
||||
makedirs("temp", exist_ok=True)
|
||||
|
||||
file_path = f"temp/raw_data_{model_name}.csv"
|
||||
data.to_csv(file_path, index=True)
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
from pandas import DataFrame
|
||||
from pytest import fixture, mark
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
|
||||
@@ -29,7 +30,7 @@ def test___init__(mock_mlflow_repository):
|
||||
@fixture
|
||||
@patch("laborious.activities.mlflow.MLFlowRepository")
|
||||
def mlflow(mock_mlflow_repository):
|
||||
return MLFlow(
|
||||
mlflow = MLFlow(
|
||||
mlflow_host="http://localhost:5000",
|
||||
mlflow_port=5000,
|
||||
mlflow_username="admin",
|
||||
@@ -38,6 +39,10 @@ def mlflow(mock_mlflow_repository):
|
||||
notification_handler=MagicMock()
|
||||
)
|
||||
|
||||
mlflow.send_notification = MagicMock()
|
||||
|
||||
return mlflow
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
@@ -141,3 +146,127 @@ async def test_request_predict(mock_max, mock_dataframe, mlflow):
|
||||
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
||||
'test_model', mock_dataframe.return_value, 30
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_retrain_model(mlflow):
|
||||
data = {
|
||||
"model_id": [4, 5, 6, 7],
|
||||
"created_at": [1, 2, 3, 4],
|
||||
"timestamp": [1, 1, 2, 2],
|
||||
"variable": ["var1", "var2", "var1", "var2"],
|
||||
"value": [1, 2, 3, 4]
|
||||
}
|
||||
|
||||
mlflow.model_monitoring_repository.retrain_model.return_value = (
|
||||
'Model retrained successfully', 'test')
|
||||
|
||||
response = await mlflow.retrain_model({
|
||||
**metadata,
|
||||
'data': data,
|
||||
'model_name': 'test_model'
|
||||
})
|
||||
|
||||
mlflow.model_monitoring_repository.retrain_model.assert_called_once()
|
||||
|
||||
assert response == {
|
||||
"status": 'Model retrained successfully',
|
||||
"timestamp": 2,
|
||||
"experiment": 'test'
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_retrain_model_error(mlflow):
|
||||
mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
|
||||
'Error retraining model'
|
||||
)
|
||||
|
||||
data = {
|
||||
"model_id": [4, 5, 6, 7],
|
||||
"created_at": [1, 2, 3, 4],
|
||||
"timestamp": [1, 1, 2, 2],
|
||||
"variable": ["var1", "var2", "var1", "var2"],
|
||||
"value": [1, 2, 3, 4]
|
||||
}
|
||||
|
||||
try:
|
||||
await mlflow.retrain_model({
|
||||
**metadata,
|
||||
'data': data,
|
||||
'model_name': 'test_model'
|
||||
})
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error retraining model'
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='RETRAIN_MODEL_ERROR',
|
||||
message='Error retraining model test_model: Error retraining model',
|
||||
block='retrain_model',
|
||||
attachment_content=ANY
|
||||
)
|
||||
else:
|
||||
assert False, "No exception raised"
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_production_model(mlflow):
|
||||
mlflow.model_monitoring_repository.update_production_model.return_value = (
|
||||
{
|
||||
"data1": 1,
|
||||
"data2": 2
|
||||
}
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'model_name': 'test_model',
|
||||
'model_id': 1,
|
||||
'experiment': 'test',
|
||||
'timestamp': 2,
|
||||
'status': 'success'
|
||||
}
|
||||
|
||||
response = await mlflow.update_production_model(input_data)
|
||||
|
||||
mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
|
||||
experiment='test', model_name='test_model')
|
||||
|
||||
assert response == {
|
||||
'data1': {0: 1},
|
||||
'data2': {0: 2},
|
||||
'model_id': {0: 1},
|
||||
'model_name': {0: 'test_model'},
|
||||
'timestamp': {0: 2},
|
||||
'status': {0: 'success'}
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_update_production_model_error(mlflow):
|
||||
mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
|
||||
'Error updating production model'
|
||||
)
|
||||
|
||||
input_data = {
|
||||
**metadata,
|
||||
'model_name': 'test_model',
|
||||
'model_id': 1,
|
||||
'experiment': 'test',
|
||||
'timestamp': 2,
|
||||
'status': 'success'
|
||||
}
|
||||
|
||||
try:
|
||||
await mlflow.update_production_model(input_data)
|
||||
except Exception as e:
|
||||
assert str(e) == 'Error updating production model'
|
||||
mlflow.send_notification.assert_called_once_with(
|
||||
metadata=metadata['metadata'],
|
||||
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
||||
message='Error updating production model test_model: Error updating production model',
|
||||
block='update_production_model',
|
||||
attachment_content=ANY
|
||||
)
|
||||
else:
|
||||
assert False, "No exception raised"
|
||||
|
||||
@@ -164,6 +164,18 @@ def test_get_experiment_last_run(mlflow, mlflow_repository):
|
||||
assert output == '2'
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_get_experiment_last_run_error(mlflow, mlflow_repository):
|
||||
mlflow.search_runs.return_value = []
|
||||
|
||||
try:
|
||||
mlflow_repository.get_experiment_last_run(0)
|
||||
except ValueError as e:
|
||||
assert str(e) == 'Runs is not a pandas DataFrame'
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
@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):
|
||||
@@ -233,14 +245,6 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
|
||||
@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'
|
||||
@@ -277,6 +281,27 @@ def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, ml
|
||||
assert output == ("Model retrained successfully", experiment)
|
||||
|
||||
|
||||
def test_retrain_model(mlflow_repository):
|
||||
data = MagicMock()
|
||||
model_name = 'test'
|
||||
|
||||
mlflow_repository.create_model_experiment = MagicMock(
|
||||
return_value=('data_model', 'prediction_model', '0'))
|
||||
|
||||
mlflow_repository.perform_model_retrain = MagicMock(
|
||||
return_value='Model retrained successfully')
|
||||
|
||||
output = mlflow_repository.retrain_model(data, model_name)
|
||||
|
||||
mlflow_repository.create_model_experiment.assert_called_once_with(
|
||||
model_name, data)
|
||||
|
||||
mlflow_repository.perform_model_retrain.assert_called_once_with(
|
||||
'data_model', 'prediction_model', '0', model_name, data)
|
||||
|
||||
assert output == 'Model retrained successfully'
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_update_production_model_by_run_id(mlflow, mlflow_repository):
|
||||
client_mock = MagicMock()
|
||||
@@ -312,6 +337,24 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
|
||||
}
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.mlflow')
|
||||
def test_update_production_model_by_run_id_error(mlflow, mlflow_repository):
|
||||
mlflow.tracking.MlflowClient.return_value = MagicMock(
|
||||
get_registered_model=MagicMock(
|
||||
return_value=MagicMock(
|
||||
latest_versions={}
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
mlflow_repository.update_production_model_by_run_id('0', 'test')
|
||||
except Exception as e:
|
||||
assert str(e) == 'Model versions is not a list'
|
||||
else:
|
||||
assert False
|
||||
|
||||
|
||||
def test_update_production_model(mlflow_repository):
|
||||
connector = mlflow_repository
|
||||
|
||||
|
||||
97
tests/laborious/workflows/test_minimal_retrain.py
Normal file
97
tests/laborious/workflows/test_minimal_retrain.py
Normal file
@@ -0,0 +1,97 @@
|
||||
from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
|
||||
from pytest import fixture, mark
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||
|
||||
|
||||
@fixture
|
||||
def minimal_retrain() -> MinimalRetrain:
|
||||
return MinimalRetrain()
|
||||
|
||||
|
||||
metadata = {
|
||||
"metadata": {
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "minimal_retrain",
|
||||
"schedule_name": "test_schedule",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
|
||||
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
|
||||
input_data = {
|
||||
"model_id": "test_model_id",
|
||||
"model_name": "test_model",
|
||||
"workflow_name": "minimal_retrain",
|
||||
"schedule_name": "test_schedule",
|
||||
"query": "test_query",
|
||||
"schema": "test_schema",
|
||||
"table_name": "test_table",
|
||||
}
|
||||
|
||||
workflow_mock.execute_activity_method = AsyncMock(
|
||||
return_value={
|
||||
"data1": "1",
|
||||
"data2": "2",
|
||||
}
|
||||
)
|
||||
|
||||
await minimal_retrain.run(input_data)
|
||||
|
||||
workflow_mock.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.load_custom_query,
|
||||
{
|
||||
**metadata,
|
||||
"query": input_data["query"],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.retrain_model,
|
||||
{
|
||||
**metadata,
|
||||
'data': workflow_mock.execute_local_activity_method.return_value,
|
||||
'model_name': input_data['model_name'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.update_production_model,
|
||||
{
|
||||
**metadata,
|
||||
'model_name': input_data['model_name'],
|
||||
'model_id': input_data['model_id'],
|
||||
**workflow_mock.execute_activity_method.return_value,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
|
||||
workflow_mock.execute_activity_method.assert_has_calls([
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**metadata,
|
||||
'data': workflow_mock.execute_activity_method.return_value,
|
||||
'schema': input_data['schema'],
|
||||
'table_name': input_data['table_name'],
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY
|
||||
)
|
||||
])
|
||||
Reference in New Issue
Block a user