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:
vitor-aignosi
54
tests/laborious/utils/models/test_minio_dataframe_payload.py
Normal file
54
tests/laborious/utils/models/test_minio_dataframe_payload.py
Normal file
@@ -0,0 +1,54 @@
|
||||
from datetime import datetime
|
||||
|
||||
from pytest import mark
|
||||
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
|
||||
|
||||
def test_parse_object_timestamp_hyphenated_model():
|
||||
key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet'
|
||||
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
||||
assert ts == datetime(2024, 6, 15, 10, 30, 45)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_transform():
|
||||
key = 'p/m-transform-2024-01-02_03-04-05.parquet'
|
||||
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
||||
assert ts == datetime(2024, 1, 2, 3, 4, 5)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_invalid():
|
||||
assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None
|
||||
|
||||
|
||||
def test_is_offloaded_dict_true_false():
|
||||
assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': None}) is True
|
||||
assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': {}}) is False
|
||||
assert MinioDataFramePayload.is_offloaded_dict({'data': {}}) is False
|
||||
|
||||
|
||||
def test_cleanup_prefix_from_payload_dict():
|
||||
p = {
|
||||
'object_key': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
|
||||
'bucket': 'b',
|
||||
'data': None,
|
||||
}
|
||||
assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'training_datasets/m'
|
||||
|
||||
|
||||
def test_cleanup_prefix_from_explicit_object_prefix():
|
||||
p = {'object_key': 'x.parquet', 'object_prefix': 'my/prefix', 'data': None}
|
||||
assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'my/prefix'
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_resolve_dict_if_offloaded_noop():
|
||||
d = {'success': True, 'data': {'a': [1]}}
|
||||
out = await MinioDataFramePayload.resolve_dict_if_offloaded(d, None, {})
|
||||
assert out is d
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_dataframe_from_wire_list():
|
||||
df = await MinioDataFramePayload.dataframe_from_wire([{'a': 1}], None, {})
|
||||
assert list(df.columns) == ['a']
|
||||
@@ -1,245 +0,0 @@
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
||||
|
||||
from botocore.utils import ClientError
|
||||
from pytest import fixture, mark, raises
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.repository.minio_repository import MinioRepository
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.minio_repository.boto3')
|
||||
@patch('laborious.utils.repository.minio_repository.Config')
|
||||
def test___init___(mock_config, mock_boto3):
|
||||
minio_repository = MinioRepository(
|
||||
minio_endpoint_url='localhost:9000',
|
||||
minio_access_key='minio',
|
||||
minio_secret_key='minio123',
|
||||
minio_region_name='us-east-1',
|
||||
minio_default_bucket='test',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
assert minio_repository.storage_options == {
|
||||
'key': 'minio',
|
||||
'secret': 'minio123',
|
||||
'client_kwargs': {'endpoint_url': 'localhost:9000'},
|
||||
}
|
||||
assert minio_repository.minio_bucket == 'test'
|
||||
assert minio_repository.minio_endpoint_url == 'localhost:9000'
|
||||
assert minio_repository.minio_region_name == 'us-east-1'
|
||||
|
||||
mock_config.assert_called_once_with(
|
||||
signature_version='s3v4',
|
||||
s3={'addressing_style': 'path'},
|
||||
retries={'max_attempts': 5, 'mode': 'standard'},
|
||||
connect_timeout=5,
|
||||
read_timeout=120,
|
||||
)
|
||||
|
||||
mock_boto3.client.assert_called_once_with(
|
||||
's3',
|
||||
endpoint_url='localhost:9000',
|
||||
aws_access_key_id='minio',
|
||||
aws_secret_access_key='minio123',
|
||||
region_name='us-east-1',
|
||||
config=mock_config.return_value,
|
||||
)
|
||||
|
||||
|
||||
@fixture
|
||||
@patch('laborious.utils.repository.minio_repository.Config')
|
||||
@patch('laborious.utils.repository.minio_repository.boto3')
|
||||
def minio_repository(mock_boto3, mock_config):
|
||||
minio_repository = MinioRepository(
|
||||
minio_endpoint_url='localhost:9000',
|
||||
minio_access_key='minio',
|
||||
minio_secret_key='minio123',
|
||||
minio_region_name='us-east-1',
|
||||
minio_default_bucket='test',
|
||||
logger=MagicMock(),
|
||||
notification_handler=MagicMock(),
|
||||
metrics_controller=AsyncMock(),
|
||||
)
|
||||
|
||||
minio_repository.emit_metric = AsyncMock()
|
||||
minio_repository.observe_lag = AsyncMock()
|
||||
minio_repository.send_notification = MagicMock()
|
||||
minio_repository.send_notification_async = AsyncMock()
|
||||
|
||||
return minio_repository
|
||||
|
||||
|
||||
def test_close(minio_repository):
|
||||
minio_repository.close()
|
||||
minio_repository.s3_client.close.assert_called_once()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_create_bucket_success(minio_repository):
|
||||
await minio_repository.create_bucket({})
|
||||
|
||||
minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
|
||||
|
||||
minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_WRITE_LAG, ANY)
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_WRITE_COUNT, tags=ANY
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_create_bucket_error(minio_repository):
|
||||
minio_repository.s3_client.create_bucket.side_effect = ValueError('test')
|
||||
|
||||
with raises(ValueError):
|
||||
await minio_repository.create_bucket({})
|
||||
|
||||
minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY
|
||||
)
|
||||
minio_repository.observe_lag.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_ensure_bucket_exists_bucket_exists(minio_repository):
|
||||
assert await minio_repository.ensure_bucket_exists({}) is None
|
||||
|
||||
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
|
||||
|
||||
minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_READ_LAG, ANY)
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_READ_COUNT, tags=ANY
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_ensure_bucket_exists_bucket_not_exists_create_success(minio_repository):
|
||||
minio_repository.s3_client.head_bucket.side_effect = ClientError(
|
||||
error_response={'Error': {'Code': '404'}}, operation_name='head_bucket'
|
||||
)
|
||||
minio_repository.create_bucket = AsyncMock()
|
||||
|
||||
assert await minio_repository.ensure_bucket_exists({}) is None
|
||||
|
||||
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
|
||||
minio_repository.create_bucket.assert_called_once_with({})
|
||||
|
||||
minio_repository.observe_lag.assert_not_called()
|
||||
minio_repository.emit_metric.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_ensure_bucket_exists_bucket_not_exists_create_error(minio_repository):
|
||||
minio_repository.s3_client.head_bucket.side_effect = ValueError('test')
|
||||
|
||||
with raises(ValueError):
|
||||
await minio_repository.ensure_bucket_exists({})
|
||||
|
||||
minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test')
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY
|
||||
)
|
||||
minio_repository.observe_lag.assert_not_called()
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.utils.repository.minio_repository.BytesIO')
|
||||
async def test_store_dataframe_as_parquet_success(mock_bytesio, minio_repository):
|
||||
input_data = MagicMock()
|
||||
|
||||
minio_repository.ensure_bucket_exists = AsyncMock()
|
||||
|
||||
await minio_repository.store_dataframe_as_parquet(
|
||||
dataframe=input_data, uri='s3://test/test.parquet', object_name='test.parquet', metadata={}
|
||||
)
|
||||
|
||||
minio_repository.ensure_bucket_exists.assert_called_once_with({})
|
||||
mock_bytesio.assert_called_once()
|
||||
|
||||
input_data.to_parquet.assert_called_once_with(
|
||||
mock_bytesio.return_value, engine='pyarrow', index=True
|
||||
)
|
||||
mock_bytesio.return_value.seek.assert_called_once_with(0)
|
||||
minio_repository.s3_client.put_object.assert_called_once_with(
|
||||
Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value
|
||||
)
|
||||
|
||||
minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_WRITE_LAG, ANY)
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_WRITE_COUNT, tags=ANY
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.utils.repository.minio_repository.BytesIO')
|
||||
async def test_store_dataframe_as_parquet_error(mock_bytesio, minio_repository):
|
||||
input_data = MagicMock()
|
||||
|
||||
minio_repository.ensure_bucket_exists = AsyncMock()
|
||||
minio_repository.s3_client.put_object.side_effect = ValueError('test')
|
||||
|
||||
with raises(ValueError):
|
||||
await minio_repository.store_dataframe_as_parquet(
|
||||
dataframe=input_data,
|
||||
uri='s3://test/test.parquet',
|
||||
object_name='test.parquet',
|
||||
metadata={},
|
||||
)
|
||||
|
||||
minio_repository.ensure_bucket_exists.assert_called_once_with({})
|
||||
mock_bytesio.assert_called_once()
|
||||
|
||||
input_data.to_parquet.assert_called_once_with(
|
||||
mock_bytesio.return_value, engine='pyarrow', index=True
|
||||
)
|
||||
mock_bytesio.return_value.seek.assert_called_once_with(0)
|
||||
minio_repository.s3_client.put_object.assert_called_once_with(
|
||||
Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value
|
||||
)
|
||||
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.utils.repository.minio_repository.BytesIO')
|
||||
@patch('laborious.utils.repository.minio_repository.read_parquet')
|
||||
async def test_get_parquet_as_dataframe_success(mock_read_parquet, mock_bytesio, minio_repository):
|
||||
input_data = {'Body': MagicMock(read=MagicMock(return_value=b'test'))}
|
||||
|
||||
minio_repository.s3_client.get_object.return_value = input_data
|
||||
|
||||
output = await minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={})
|
||||
|
||||
minio_repository.s3_client.get_object.assert_called_once_with(Bucket='test', Key='test.parquet')
|
||||
|
||||
mock_bytesio.assert_called_once_with(input_data['Body'].read.return_value)
|
||||
mock_read_parquet.assert_called_once_with(mock_bytesio.return_value)
|
||||
|
||||
assert output == mock_read_parquet.return_value
|
||||
|
||||
minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_READ_LAG, ANY)
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_READ_COUNT, tags=ANY
|
||||
)
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
@patch('laborious.utils.repository.minio_repository.BytesIO')
|
||||
@patch('laborious.utils.repository.minio_repository.read_parquet')
|
||||
async def test_get_parquet_as_dataframe_error(mock_read_parquet, mock_bytesio, minio_repository):
|
||||
minio_repository.s3_client.get_object.side_effect = ValueError('test')
|
||||
|
||||
with raises(ValueError):
|
||||
await minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={})
|
||||
|
||||
minio_repository.s3_client.get_object.assert_called_once_with(
|
||||
Bucket='test', Key='test.parquet'
|
||||
)
|
||||
minio_repository.emit_metric.assert_called_once_with(
|
||||
metric_object=metrics.MINIO_READ_ERROR_COUNT, tags=ANY
|
||||
)
|
||||
minio_repository.observe_lag.assert_not_called()
|
||||
Reference in New Issue
Block a user