vitor-aignosi
381856f5ca
SIENTIAPDE-1712 Enhance logging across various classes by adding logger parameters and improving debug statements. This update includes adjustments in Gates, MLFlow, ModelMetrics, and MLFlowRepository classes for better traceability and observability during operations.
630 lines
19 KiB
Python
630 lines
19 KiB
Python
from unittest.mock import ANY, AsyncMock, 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
|
|
|
|
|
|
@fixture(autouse=True)
|
|
def _passthrough_from_dict():
|
|
with patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
|
|
):
|
|
yield
|
|
|
|
|
|
@patch('laborious.activities.mlflow.MLFlowRepository')
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def test___init__(mock_minio_repository, mock_mlflow_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
minio_repo = mock_minio_repository(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
mlflow = MLFlow(
|
|
mlflow_host='http://localhost',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
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, ANY, ANY
|
|
)
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=ANY,
|
|
notification_handler=ANY,
|
|
metrics_controller=ANY,
|
|
bucket='test',
|
|
)
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.mlflow.MLFlowRepository')
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def mlflow(mock_minio_repository, mock_mlflow_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
minio_repo = mock_minio_repository(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
mlflow = MLFlow(
|
|
mlflow_host='http://localhost:5000',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
mlflow.model_monitoring_repository = AsyncMock()
|
|
mlflow.minio_repository = AsyncMock()
|
|
|
|
mlflow.send_notification = MagicMock()
|
|
mlflow.emit_metric = AsyncMock()
|
|
mlflow.send_notification_async = AsyncMock()
|
|
mlflow.error = MagicMock()
|
|
mlflow.debug = MagicMock()
|
|
mlflow.info = MagicMock()
|
|
mlflow.warning = MagicMock()
|
|
mlflow.critical = 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.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_request_transform_success(mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
transform_response = {'success': True, 'content': MagicMock()}
|
|
mlflow.model_monitoring_repository.transform.return_value = transform_response
|
|
|
|
data_mock.sort_values.return_value = data_mock
|
|
data_mock.drop_duplicates.return_value = data_mock
|
|
data_mock.pivot.return_value = data_mock
|
|
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
mlflow.model_monitoring_repository.transform.assert_called_once_with(
|
|
'test_model', data_mock, {}, metadata['metadata']
|
|
)
|
|
mock_from_dataframe.assert_called_once()
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_request_transform_failure(mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
transform_response = {'success': False, 'message': 'Transform failed'}
|
|
mlflow.model_monitoring_repository.transform.return_value = transform_response
|
|
|
|
data_mock.sort_values.return_value = data_mock
|
|
data_mock.drop_duplicates.return_value = data_mock
|
|
data_mock.pivot.return_value = data_mock
|
|
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
mock_from_dataframe.assert_called_once_with(
|
|
dataframe=None,
|
|
minio_repo=mlflow.minio_repository,
|
|
model_name='test_model',
|
|
operation='transform',
|
|
status=transform_response,
|
|
workflow_metadata=metadata['metadata'],
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=mlflow.logger,
|
|
)
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
predict_response = {'success': True, 'content': MagicMock()}
|
|
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
|
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
|
mock_to_datetime.assert_called_once_with(
|
|
data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
|
)
|
|
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
|
|
|
|
mlflow.model_monitoring_repository.predict.assert_called_once_with(
|
|
'test_model', data_mock, {}, metadata['metadata']
|
|
)
|
|
mock_from_dataframe.assert_called_once()
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
predict_response = {'success': False, 'message': 'Predict failed'}
|
|
mlflow.model_monitoring_repository.predict.return_value = predict_response
|
|
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
mock_from_dataframe.assert_called_once_with(
|
|
dataframe=None,
|
|
minio_repo=mlflow.minio_repository,
|
|
model_name='test_model',
|
|
operation='predict',
|
|
status=predict_response,
|
|
workflow_metadata=metadata['metadata'],
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=mlflow.logger,
|
|
)
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@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.',
|
|
}
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
timestamp = raw_data.__getitem__.return_value.max.return_value
|
|
|
|
raw_data.sort_values.assert_not_called()
|
|
raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
|
|
raw_data = raw_data.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_with_payload_data(mock_to_datetime, mlflow):
|
|
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
|
'success': True,
|
|
'experiment': 'test_experiment',
|
|
'message': 'Model retrained successfully.',
|
|
}
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response['success'] is True
|
|
mlflow.minio_repository.download_file.assert_not_called()
|
|
|
|
|
|
@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.',
|
|
}
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
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_async.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):
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
'transform_flavor': 'sklearn',
|
|
'predict_flavor': 'pyfunc',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response == {
|
|
'success': False,
|
|
'message': "Error loading retrain data: '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_async.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')
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_get_reference_data_success(mock_to_datetime, mlflow):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
# Mock reference data DataFrame
|
|
mock_reference_data = MagicMock()
|
|
mock_reference_data.__getitem__.return_value = MagicMock()
|
|
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
|
|
mock_reference_data.to_dict.return_value = [
|
|
{'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
|
|
{'timestamp': '2023-05-26 11:12:28', 'value': 2.0},
|
|
]
|
|
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = mock_reference_data
|
|
|
|
# Act
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
# Assert
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
|
model_name='test_model',
|
|
artifact_path='evaluation_data.csv',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
mock_to_datetime.assert_called_once_with(mock_reference_data.__getitem__.return_value)
|
|
|
|
mock_reference_data.to_dict.assert_called_once_with(orient='records')
|
|
assert result == mock_reference_data.to_dict.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_reference_data_not_found(mlflow):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.return_value = None
|
|
|
|
# Act
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
# Assert
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
|
model_name='test_model',
|
|
artifact_path='evaluation_data.csv',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
mlflow.warning.assert_called_once_with(
|
|
'Reference data not found for model test_model', metadata['metadata']
|
|
)
|
|
assert result is None
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_reference_data_exception(mlflow):
|
|
# Arrange
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.side_effect = Exception(
|
|
'Error loading artifact'
|
|
)
|
|
|
|
# Act & Assert
|
|
with raises(Exception) as e:
|
|
await mlflow.get_reference_data(input_data)
|
|
|
|
assert str(e.value) == 'Error loading artifact'
|
|
mlflow.model_monitoring_repository.load_artifact_dataframe.assert_called_once_with(
|
|
model_name='test_model',
|
|
artifact_path='evaluation_data.csv',
|
|
metadata=metadata['metadata'],
|
|
)
|