gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1243: Refactor and enhance model manager activities and workflows

Browse Source 
This commit includes several changes:

- Reorganized imports and class inheritance in activities.py, gates.py and mlflow.py for better readability and maintainability.
- Improved error handling and logging in gates.py and mlflow.py.
- Added input validation and filtering in gates.py to ensure data quality.
- Enhanced prediction formatting and storage policy management in gates.py.
- Updated metrics.py to use consistent naming conventions and labels.
- Refactored connectors_config.py to use type hints and improve code clarity.
- Updated conditional and MLFlow filters for better data quality checks.
- Improved model repository logic for retraining and updating models.
- Enhanced worker.py to include SDK metrics and improved error handling.
- Refactored workflows for better modularity and error handling.
- Updated tests to reflect the changes and improve test coverage.
This commit is contained in:
Bruno Domingues
2025-10-01 17:28:57 -03:00
parent b102f79087
commit dfc190c818
24 changed files with 1482 additions and 1399 deletions
Download Patch File Download Diff File Expand all files Collapse all files

335
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -1,34 +1,33 @@
from datetime import UTC, datetime
from unittest.mock import ANY, MagicMock, call, patch
import numpy as np
from pandas import DataFrame
import pytest
from datetime import datetime, timezone
from pandas import Timestamp
from pandas import DataFrame, Timestamp
from model_manager.utils.repository.model_repository import MLFlowRepository
@pytest.fixture
def mlflow_repository():
with patch('model_manager.utils.repository.model_repository.ModelServing',
autospec=True) as mock_model_serving:
with patch(
'model_manager.utils.repository.model_repository.ModelServing', autospec=True
) as mock_model_serving:
mock_instance = mock_model_serving.return_value
mock_instance.get_transformed_data = MagicMock()
repo = MLFlowRepository(
host='http://localhost:5000',
username='admin',
password='admin',
logger=MagicMock()
host='http://localhost:5000', username='admin', password='admin', logger=MagicMock()
)
return repo
metadata = {
"metadata": {
"model_id": "test_model",
"model_name": "test_model",
"workflow_name": "test_workflow",
"schema_name": "test_schedule",
'metadata': {
'model_id': 'test_model',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
},
}
@@ -38,80 +37,56 @@ class Any:
invalid_cases = [
(
{
'value': {
'2024-01-01 12:00:00': 1,
2024: 2
}
}
),
(
{
'value': {
'2024-01-01': 1,
'2024-01-02': 2
}
}
),
(
{
'value': {
Any(): 1,
Any(): 2
}
}
)
({'value': {'2024-01-01 12:00:00': 1, 2024: 2}}),
({'value': {'2024-01-01': 1, '2024-01-02': 2}}),
({'value': {Any(): 1, Any(): 2}}),
]
@pytest.mark.parametrize("data", invalid_cases)
@pytest.mark.parametrize('data', invalid_cases)
def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data):
input_data = DataFrame(
data
)
input_data = DataFrame(data)
with pytest.raises(ValueError) as e:
mlflow_repository.detect_and_parse_datetime_index(
input_data, metadata['metadata'])
mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
assert str(e) == "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S"
assert (
str(e)
== 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S'
)
valid_cases = [
(
{
'value': {
'2024-01-01 12:00:00+0000': 1,
'2024-01-02 12:00:00+0000': 2
}
}, ['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000']
{'value': {'2024-01-01 12:00:00+0000': 1, '2024-01-02 12:00:00+0000': 2}},
['2024-01-01 12:00:00+0000', '2024-01-02 12:00:00+0000'],
),
(
{
'value': {
datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
datetime(2025, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
datetime(2025, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
}
}, ['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000']
},
['2025-01-01 12:00:00+0000', '2025-01-02 12:00:00+0000'],
),
(
{
'value': {
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=timezone.utc): 2
Timestamp(2026, 1, 1, 12, 0, 0, tzinfo=UTC): 1,
Timestamp(2026, 1, 2, 12, 0, 0, tzinfo=UTC): 2,
}
}, ['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000']
},
['2026-01-01 12:00:00+0000', '2026-01-02 12:00:00+0000'],
),
]
@pytest.mark.parametrize("data,expected", valid_cases)
@pytest.mark.parametrize('data,expected', valid_cases)
def test_detect_and_parse_datetime_index_valid_format(mlflow_repository, data, expected):
input_data = DataFrame(data)
response = mlflow_repository.detect_and_parse_datetime_index(
input_data, metadata['metadata'])
response = mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata'])
assert response.index.tolist() == expected
@@ -122,18 +97,19 @@ def test_transform_success(mlflow_repository):
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
output = mlflow_repository.transform(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict')
model_name, data, 0, 'sklearn', False, 'model', 'predict'
)
mlflow_repository.detect_and_parse_datetime_index.assert_called_once_with(
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.return_value, metadata['metadata']
)
assert output == {
'success': True,
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value,
}
@@ -141,80 +117,48 @@ def test_transform_error(mlflow_repository):
data = MagicMock()
model_name = 'model'
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception(
'error')
mlflow_repository.model_serving.get_cached_transform.side_effect = Exception('error')
output = mlflow_repository.transform(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.transform(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_transform.assert_called_once_with(
model_name, data, 0, 'sklearn', False, 'model', 'predict')
model_name, data, 0, 'sklearn', False, 'model', 'predict'
)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
def test_predict_success(mlflow_repository):
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict.return_value = np.array(
[2, 3]
)
mlflow_repository.model_serving.get_cached_predict.return_value = np.array([2, 3])
output = mlflow_repository.predict(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model')
model_name, data, 0, 'pyfunc', False, 'model'
)
assert output['success'] is True
assert output['content'] == {
'prediction': {
'index_1': 2,
'index_2': 3
}, 'response_time': {
'index_1': ANY,
'index_2': ANY
}
'prediction': {'index_1': 2, 'index_2': 3},
'response_time': {'index_1': ANY, 'index_2': ANY},
}
def test_predict_error(mlflow_repository):
data = DataFrame({
'feat_1': {
'index_1': 2,
'index_2': 3
}
})
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
model_name = 'model'
mlflow_repository.model_serving.get_cached_predict = MagicMock(
side_effect=Exception('error')
)
mlflow_repository.model_serving.get_cached_predict = MagicMock(side_effect=Exception('error'))
output = mlflow_repository.predict(
model_name, data, {}, metadata['metadata'])
output = mlflow_repository.predict(model_name, data, {}, metadata['metadata'])
mlflow_repository.model_serving.get_cached_predict.assert_called_once_with(
model_name, data, 0, 'pyfunc', False, 'model')
model_name, data, 0, 'pyfunc', False, 'model'
)
assert output == {
'success': False,
'content': {
'message': 'error',
'traceback': ANY
}
}
assert output == {'success': False, 'content': {'message': 'error', 'traceback': ANY}}
@patch('model_manager.utils.repository.model_repository.mlflow')
@@ -246,8 +190,7 @@ def test_get_next_run_name(mlflow, mlflow_repository):
@patch('model_manager.utils.repository.model_repository.mlflow')
def test_get_experiment_success(mlflow, mlflow_repository):
mlflow.get_experiment_by_name.return_value = MagicMock(
experiment_id='0')
mlflow.get_experiment_by_name.return_value = MagicMock(experiment_id='0')
output = mlflow_repository.get_experiment('test')
@@ -263,23 +206,25 @@ def test_get_experiment_error(mlflow, mlflow_repository):
except ValueError as e:
assert str(e) == 'Experiment test not found'
else:
assert False
raise AssertionError('Expected exception')
@patch('model_manager.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'],
})
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",
filter_string='',
output_format='pandas',
)
assert output == '2'
@@ -294,17 +239,14 @@ def test_get_experiment_last_run_error(mlflow, mlflow_repository):
except ValueError as e:
assert str(e) == 'Runs is not a pandas DataFrame'
else:
assert False
raise AssertionError('Expected exception')
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn')
@patch('model_manager.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.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()
@@ -313,29 +255,30 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
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.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]
})
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)
'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"),
])
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)
@@ -343,21 +286,23 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
fit_args = prediction_model_mock.fit.call_args[0][0]
assert fit_args.equals(
DataFrame({
'x': [10, 20, 30],
'y': [4, 5, 6],
})
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
)
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)
assert output == (
prediction_model_mock,
data_model_mock,
mlflow_repository.get_experiment_by_run_id.return_value,
)
@patch('model_manager.utils.repository.model_repository.mlflow.start_run')
@@ -365,41 +310,43 @@ def test_create_model_experiment(set_experiment, sklearn, mlflow_repository):
@patch('model_manager.utils.repository.model_repository.mlflow.sklearn.log_model')
@patch('model_manager.utils.repository.model_repository.mlflow.log_artifact')
def test_perform_model_retrain(log_artifact, log_model, log_param, start_run, mlflow_repository):
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')
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)
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')
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"),
])
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)
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_artifact.assert_called_once_with('temp/raw_data_test.csv')
log_param.assert_has_calls([
call("retrain", True),
])
log_param.assert_has_calls(
[
call('retrain', True),
]
)
assert output == ("Model retrained successfully", experiment)
assert output == ('Model retrained successfully', experiment)
def test_retrain_model(mlflow_repository):
@@ -407,18 +354,18 @@ def test_retrain_model(mlflow_repository):
model_name = 'test'
mlflow_repository.create_model_experiment = MagicMock(
return_value=('data_model', 'prediction_model', '0'))
return_value=('data_model', 'prediction_model', '0')
)
mlflow_repository.perform_model_retrain = MagicMock(
return_value='Model retrained successfully')
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.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)
'data_model', 'prediction_model', '0', model_name, data
)
assert output == 'Model retrained successfully'
@@ -438,7 +385,7 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
output = mlflow_repository.update_production_model_by_run_id('0', 'test')
mlflow.register_model.assert_called_once_with(
"runs:/0/prediction_model",
'runs:/0/prediction_model',
'test',
)
@@ -461,38 +408,34 @@ def test_update_production_model_by_run_id(mlflow, mlflow_repository):
@patch('model_manager.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={}
)
)
get_registered_model=MagicMock(return_value=MagicMock(latest_versions={}))
)
try:
mlflow_repository.update_production_model_by_run_id('0', 'test')
except Exception as e:
except Exception as e: # noqa: BLE001
assert str(e) == 'Model versions is not a list'
else:
assert False
raise AssertionError('Expected exception')
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:
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')
update_production_model_by_run_id.assert_called_once_with('2', 'test')
assert output == {
'model_name': 'test',