vitor-aignosi
f0fb9b854e
Enhance validation script and refactor code in various modules - Updated the validation script to include automatic code formatting and linting fixes using Ruff. - Removed the `clean_tmp_files` method from the Gates class to streamline functionality. - Simplified conditional checks in the OpcRepository for better clarity and error handling. - Added model ID to the minimal retrain workflow for improved tracking. - Introduced new test cases for error handling in MLFlow and storage operations, ensuring robustness in repository interactions.
469 lines
15 KiB
Python
469 lines
15 KiB
Python
from unittest.mock import ANY, MagicMock, call, patch
|
|
|
|
import numpy as np
|
|
from pytest import fixture, mark, raises
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
|
|
|
|
from laborious.activities.mlflow import MLFlow
|
|
|
|
|
|
@patch('laborious.activities.mlflow.MLFlowRepository')
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def test___init__(mock_minio_repository, mock_mlflow_repository):
|
|
mlflow = MLFlow(
|
|
mlflow_host='http://localhost',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
minio_config={
|
|
'endpoint_url': 'http://localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
assert mlflow.mlflow_host == 'http://localhost'
|
|
assert mlflow.mlflow_port == 5000
|
|
assert mlflow.mlflow_username == 'admin'
|
|
assert mlflow.mlflow_password == 'admin'
|
|
|
|
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
logger=ANY,
|
|
notification_handler=ANY,
|
|
minio_endpoint_url='http://localhost:9000',
|
|
minio_access_key='minio',
|
|
minio_secret_key='minio123',
|
|
minio_region_name='us-east-1',
|
|
minio_default_bucket='test',
|
|
)
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.mlflow.MLFlowRepository')
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def mlflow(mock_minio_repository, mock_mlflow_repository):
|
|
mlflow = MLFlow(
|
|
mlflow_host='http://localhost:5000',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
minio_config={
|
|
'endpoint_url': 'http://localhost:9000',
|
|
'access_key': 'minio',
|
|
'secret_key': 'minio123',
|
|
'region_name': 'us-east-1',
|
|
'default_bucket': 'test',
|
|
},
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
mlflow.send_notification = MagicMock()
|
|
|
|
return mlflow
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.DataFrame')
|
|
@patch('laborious.activities.mlflow.max')
|
|
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
|
|
mock_max.return_value = '2024-01-02'
|
|
# Mock input data
|
|
input_data = {
|
|
**metadata,
|
|
'data': [
|
|
{
|
|
'timestamp': '2024-01-01',
|
|
'variable': 'var1',
|
|
'value': 1.0,
|
|
'created_at': '2024-01-01 12:00:00',
|
|
},
|
|
{
|
|
'timestamp': '2024-01-01',
|
|
'variable': 'var2',
|
|
'value': 2.0,
|
|
'created_at': '2024-01-01 12:00:00',
|
|
},
|
|
{
|
|
'timestamp': '2024-01-02',
|
|
'variable': 'var1',
|
|
'value': 3.0,
|
|
'created_at': '2024-01-02 12:00:00',
|
|
},
|
|
{
|
|
'timestamp': '2024-01-02',
|
|
'variable': 'var2',
|
|
'value': 4.0,
|
|
'created_at': '2024-01-02 12:00:00',
|
|
},
|
|
{
|
|
'timestamp': '2024-01-02',
|
|
'variable': 'var1',
|
|
'value': 1.0,
|
|
'created_at': '2024-01-01 12:00:00',
|
|
},
|
|
{
|
|
'timestamp': '2024-01-02',
|
|
'variable': 'var2',
|
|
'value': 1.0,
|
|
'created_at': '2024-01-01 12:00:00',
|
|
},
|
|
],
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
# Mock the transform response
|
|
expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
|
|
mlflow.model_monitoring_repository.transform.return_value = expected_response
|
|
|
|
mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
|
|
mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
|
|
|
|
# Call the method
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
# Verify the data was correctly transformed
|
|
mock_dataframe.assert_called_once_with(input_data['data'])
|
|
mock_dataframe.return_value.pivot.assert_called_once_with(
|
|
index='timestamp', columns='variable', values='value'
|
|
)
|
|
mock_dataframe = mock_dataframe.return_value.pivot.return_value
|
|
mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
|
|
# mock_dataframe.reset_index.assert_called_once()
|
|
mock_dataframe.columns.name = None
|
|
|
|
# Verify the response
|
|
assert response_data == expected_response
|
|
|
|
# Verify the repository was called with correct arguments
|
|
mlflow.model_monitoring_repository.transform.assert_called_once_with(
|
|
'test_model', mock_dataframe, {}, metadata['metadata']
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.DataFrame')
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
@patch('laborious.activities.mlflow.max')
|
|
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
|
|
mock_max.return_value = '2024-01-02'
|
|
# Mock input data
|
|
input_data = {
|
|
**metadata,
|
|
'data': {
|
|
'variable': {
|
|
'2024-01-01': 'var1',
|
|
'2024-01-02': 'var2',
|
|
'2024-01-03': 'var1',
|
|
'2024-01-04': 'var2',
|
|
},
|
|
'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
|
|
},
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
# Mock the predict response
|
|
expected_response = {'prediction': [0.5, 0.6]}
|
|
mlflow.model_monitoring_repository.predict.return_value = expected_response
|
|
|
|
# Call the method
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
mock_dataframe.assert_called_once_with(input_data['data'])
|
|
mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
|
|
mock_dataframe.return_value.__setitem__.assert_any_call(
|
|
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
|
|
)
|
|
|
|
mock_to_datetime.assert_called_once_with(
|
|
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
|
)
|
|
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
|
|
|
|
# Verify the response
|
|
assert response_data == expected_response
|
|
|
|
# Verify the repository was called with correct arguments
|
|
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
|
'test_model', mock_dataframe.return_value, {}, metadata['metadata']
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
|
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
|
'success': True,
|
|
'experiment': 'test_experiment',
|
|
'message': 'Model retrained successfully.',
|
|
}
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'object_key': 'test_object_key',
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
|
|
|
|
timestamp = raw_data.__getitem__.return_value.max.return_value
|
|
|
|
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
|
|
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
|
|
|
|
raw_data.drop.assert_has_calls(
|
|
[
|
|
call(columns=['model_id'], inplace=True, errors='ignore'),
|
|
call(columns=['created_at'], inplace=True, errors='ignore'),
|
|
]
|
|
)
|
|
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
|
|
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
|
|
|
|
raw_data = raw_data.pivot.return_value
|
|
|
|
raw_data.__setitem__.assert_has_calls(
|
|
[
|
|
call('timestamp', raw_data.index),
|
|
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
|
|
call('timestamp', mock_to_datetime.return_value),
|
|
]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
|
|
)
|
|
|
|
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
|
|
data=raw_data,
|
|
model_name='test_model',
|
|
model_config={
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert response == {
|
|
'success': True,
|
|
'experiment': 'test_experiment',
|
|
'message': 'Model retrained successfully.',
|
|
'timestamp': timestamp,
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
|
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
|
'success': False,
|
|
'traceback': 'test_traceback',
|
|
'message': 'Model retrained failed.',
|
|
}
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'object_key': 'test_object_key',
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
|
|
|
|
timestamp = raw_data.__getitem__.return_value.max.return_value
|
|
|
|
raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
|
|
raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
|
|
|
|
raw_data.drop.assert_has_calls(
|
|
[
|
|
call(columns=['model_id'], inplace=True, errors='ignore'),
|
|
call(columns=['created_at'], inplace=True, errors='ignore'),
|
|
]
|
|
)
|
|
raw_data.pivot.assert_called_once_with(index='timestamp', columns='variable', values='value')
|
|
raw_data.pivot.return_value.fillna.assert_called_once_with(np.nan, inplace=True)
|
|
|
|
raw_data = raw_data.pivot.return_value
|
|
|
|
raw_data.__setitem__.assert_has_calls(
|
|
[
|
|
call('timestamp', raw_data.index),
|
|
call('timestamp', mock_to_datetime.return_value.dt.strftime.return_value),
|
|
call('timestamp', mock_to_datetime.return_value),
|
|
]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ)]
|
|
)
|
|
mock_to_datetime.assert_has_calls(
|
|
[call(raw_data.__getitem__.return_value, format=DATETIME_FORMAT)]
|
|
)
|
|
|
|
mlflow.model_monitoring_repository.retrain_model.assert_called_once_with(
|
|
data=raw_data,
|
|
model_name='test_model',
|
|
model_config={
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='RETRAIN_MODEL_ERROR',
|
|
message='Error retraining model test_model: Model retrained failed.',
|
|
block='retrain_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
assert response == {
|
|
'success': False,
|
|
'traceback': 'test_traceback',
|
|
'message': 'Model retrained failed.',
|
|
'timestamp': timestamp,
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_data_error(mlflow):
|
|
mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception(
|
|
'Error loading retrain data'
|
|
)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'object_key': 'test_object_key',
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response == {
|
|
'success': False,
|
|
'message': 'Error loading retrain data: Error loading retrain data',
|
|
'traceback': ANY,
|
|
'timestamp': ANY,
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_data_error_no_minio_repository(mlflow):
|
|
mlflow.minio_repository = None
|
|
|
|
with raises(ValueError) as e:
|
|
await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'object_key': 'test_object_key',
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert str(e.value) == 'Minio repository not initialized'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_update_production_model(mlflow):
|
|
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', metadata=metadata['metadata']
|
|
)
|
|
|
|
assert response == mlflow.model_monitoring_repository.update_production_model.return_value
|
|
|
|
|
|
@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',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
else:
|
|
raise AssertionError('No exception raised')
|