SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment.
This commit is contained in:
Bruno Domingues
297
tests/laborious/activities/test_mlflow.py
Normal file
297
tests/laborious/activities/test_mlflow.py
Normal file
@@ -0,0 +1,297 @@
|
||||
from datetime import datetime
|
||||
from unittest.mock import ANY, MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
from pandas import DataFrame, Timestamp
|
||||
from pytest import fixture, mark, raises
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
|
||||
|
||||
@patch("laborious.activities.mlflow.MLFlowRepository")
|
||||
def test___init__(mock_mlflow_repository):
|
||||
mlflow = MLFlow(
|
||||
mlflow_host="http://localhost",
|
||||
mlflow_port=5000,
|
||||
mlflow_username="admin",
|
||||
mlflow_password="admin",
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
@patch("laborious.activities.mlflow.MLFlowRepository")
|
||||
def mlflow(mock_mlflow_repository):
|
||||
mlflow = MLFlow(
|
||||
mlflow_host="http://localhost:5000",
|
||||
mlflow_port=5000,
|
||||
mlflow_username="admin",
|
||||
mlflow_password="admin",
|
||||
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
|
||||
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',
|
||||
level=NotificationLevel.ERROR,
|
||||
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',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=ANY
|
||||
)
|
||||
else:
|
||||
assert False, "No exception raised"
|
||||
Reference in New Issue
Block a user