SIENTIAPDE-1712
Implement MinIO Offload and Retention Features - Added configuration options for MinIO retention hours and offload threshold in README. - Introduced MinIO payload offloading for large DataFrame-derived payloads, storing them as parquet files. - Updated activities to utilize MinIO for data loading and cleanup, including new methods for offloading and retention management. - Refactored existing activities to integrate MinIO functionality, ensuring compatibility with previous workflows. - Removed the legacy MinioRepository class, consolidating MinIO operations under a new manager structure. - Updated requirements to use the latest version of the sientia-dataops-library.
This commit is contained in:
@@ -38,14 +38,13 @@ def test___init__(mock_minio_repository, mock_mlflow_repository):
|
||||
)
|
||||
|
||||
mock_minio_repository.assert_called_once_with(
|
||||
endpoint='localhost:9000',
|
||||
access_key='minio',
|
||||
secret_key='minio123',
|
||||
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',
|
||||
metrics_controller=ANY,
|
||||
bucket='test',
|
||||
)
|
||||
|
||||
|
||||
@@ -96,10 +95,15 @@ metadata = {
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.mlflow.DataFrame')
|
||||
@patch(
|
||||
'laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire',
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
@patch('laborious.activities.mlflow.max')
|
||||
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
|
||||
async def test_request_transform_success(mock_max, mock_dataframe_from_wire, mlflow):
|
||||
mock_max.return_value = '2024-01-02'
|
||||
data_mock = MagicMock()
|
||||
mock_dataframe_from_wire.return_value = data_mock
|
||||
# Mock input data
|
||||
input_data = {
|
||||
**metadata,
|
||||
@@ -149,37 +153,41 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
|
||||
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
|
||||
data_mock.sort_values.return_value = data_mock
|
||||
data_mock.drop_duplicates.return_value = data_mock
|
||||
data_mock.pivot.return_value = data_mock
|
||||
|
||||
# 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(
|
||||
data_mock.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)
|
||||
data_mock.fillna.assert_called_once_with(np.nan, inplace=True)
|
||||
# mock_dataframe.reset_index.assert_called_once()
|
||||
mock_dataframe.columns.name = None
|
||||
data_mock.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']
|
||||
'test_model', data_mock, {}, metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.mlflow.DataFrame')
|
||||
@patch(
|
||||
'laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire',
|
||||
new_callable=AsyncMock,
|
||||
)
|
||||
@patch('laborious.activities.mlflow.to_datetime')
|
||||
@patch('laborious.activities.mlflow.max')
|
||||
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
|
||||
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe_from_wire, mlflow):
|
||||
mock_max.return_value = '2024-01-02'
|
||||
data_mock = MagicMock()
|
||||
mock_dataframe_from_wire.return_value = data_mock
|
||||
# Mock input data
|
||||
input_data = {
|
||||
**metadata,
|
||||
@@ -203,22 +211,21 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
|
||||
# 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(
|
||||
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
||||
data_mock.__setitem__.assert_any_call(
|
||||
'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
|
||||
)
|
||||
mock_dataframe.return_value.__setitem__.assert_any_call(
|
||||
data_mock.__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
|
||||
data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
||||
)
|
||||
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
|
||||
|
||||
mock_to_datetime.assert_called_once_with(
|
||||
mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
||||
data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
|
||||
)
|
||||
mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)
|
||||
|
||||
@@ -227,20 +234,22 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo
|
||||
|
||||
# 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']
|
||||
'test_model', data_mock, {}, metadata['metadata']
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.mlflow.read_parquet')
|
||||
@patch('laborious.activities.mlflow.to_datetime')
|
||||
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow):
|
||||
async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mock_read_parquet, mlflow):
|
||||
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
||||
'success': True,
|
||||
'experiment': 'test_experiment',
|
||||
'message': 'Model retrained successfully.',
|
||||
}
|
||||
|
||||
mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock()
|
||||
mlflow.minio_repository.download_file.return_value = b'parquet-bytes'
|
||||
mock_read_parquet.return_value = MagicMock()
|
||||
|
||||
response = await mlflow.retrain_model(
|
||||
{
|
||||
@@ -255,7 +264,7 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
|
||||
}
|
||||
)
|
||||
|
||||
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
|
||||
raw_data = mock_read_parquet.return_value
|
||||
|
||||
timestamp = raw_data.__getitem__.return_value.max.return_value
|
||||
|
||||
@@ -308,15 +317,47 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', new_callable=AsyncMock)
|
||||
@patch('laborious.activities.mlflow.to_datetime')
|
||||
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
|
||||
async def test_retrain_model_success_with_payload_data(
|
||||
mock_to_datetime, mock_dataframe_from_wire, mlflow
|
||||
):
|
||||
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
||||
'success': True,
|
||||
'experiment': 'test_experiment',
|
||||
'message': 'Model retrained successfully.',
|
||||
}
|
||||
mock_dataframe_from_wire.return_value = MagicMock()
|
||||
|
||||
response = await mlflow.retrain_model(
|
||||
{
|
||||
**metadata,
|
||||
'data': {'data': {'a': [1]}},
|
||||
'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.read_parquet')
|
||||
@patch('laborious.activities.mlflow.to_datetime')
|
||||
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mock_read_parquet, mlflow):
|
||||
mlflow.model_monitoring_repository.retrain_model.return_value = {
|
||||
'success': False,
|
||||
'traceback': 'test_traceback',
|
||||
'message': 'Model retrained failed.',
|
||||
}
|
||||
|
||||
mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock(
|
||||
mlflow.minio_repository.download_file.return_value = b'parquet-bytes'
|
||||
mock_read_parquet.return_value = MagicMock(
|
||||
columns=['variable', 'timestamp', 'value', 'created_at']
|
||||
)
|
||||
|
||||
@@ -333,7 +374,7 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow)
|
||||
}
|
||||
)
|
||||
|
||||
raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
|
||||
raw_data = mock_read_parquet.return_value
|
||||
|
||||
timestamp = raw_data.__getitem__.return_value.max.return_value
|
||||
|
||||
@@ -398,7 +439,7 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow)
|
||||
|
||||
@mark.asyncio
|
||||
async def test_retrain_model_data_error(mlflow):
|
||||
mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception(
|
||||
mlflow.minio_repository.download_file.side_effect = Exception(
|
||||
'Error loading retrain data'
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user