gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1712

Browse Source 
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
2026-03-19 17:29:43 -03:00
parent 9dc3cb3ba0
commit 981ac700d4
25 changed files with 994 additions and 681 deletions
Download Patch File Download Diff File Expand all files Collapse all files

103
tests/laborious/activities/test_mlflow.py
View File

@@ -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'
)

122
tests/laborious/activities/test_storage.py
View File

@@ -1,4 +1,5 @@
import datetime
import os
from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pytest import fixture, mark, raises
@@ -68,14 +69,13 @@ def test___init___not_hasattr(mock_minio_repository):
assert isinstance(storage, Postgres)
mock_minio_repository.assert_called_once_with(
endpoint='localhost:9000',
access_key='minio',
secret_key='minio123',
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
metrics_controller=metrics_controller,
bucket='test',
)
@@ -106,14 +106,13 @@ def test___init___none_minio_repository(mock_minio_repository, storage):
)
mock_minio_repository.assert_called_once_with(
endpoint='localhost:9000',
access_key='minio',
secret_key='minio123',
logger=logger,
notification_handler=notification_handler,
minio_endpoint_url='localhost:9000',
minio_access_key='minio',
minio_secret_key='minio123',
minio_region_name='us-east-1',
minio_default_bucket='test',
metrics_controller=metrics_controller,
bucket='test',
)
@@ -169,30 +168,38 @@ async def test_query_to_minio_success(now, dataframe, storage):
data = [{'a': 1}, {'a': 2}, {'a': 3}]
storage.load_custom_query = AsyncMock(return_value=data)
now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
storage.minio_repository.store_dataframe_as_parquet = AsyncMock()
storage.minio_repository.minio_bucket = 'test'
storage.minio_repository.upload_file = AsyncMock(
return_value={
'minio_object_name': 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
}
)
storage.minio_repository.bucket = 'test'
result = await storage.query_to_minio({'object_prefix': 'test', **metadata})
dataframe.assert_called_once_with(data)
storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with(
dataframe=dataframe.return_value,
uri='s3://test/test_2024-01-01_00-00-00.parquet',
object_name='test_2024-01-01_00-00-00.parquet',
storage.minio_repository.upload_file.assert_called_once_with(
file_bytes=ANY,
relative_key='training_datasets/test_model/test_2024-01-01_00-00-00.parquet',
metadata=metadata['metadata'],
)
assert result['success'] is True
assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet'
assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet'
assert (
result['object_key']
== 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
)
assert (
result['uri']
== 's3://test/sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet'
)
@mark.asyncio
async def test_query_to_minio_error(storage):
storage.send_notification = MagicMock()
storage.send_notification_async = AsyncMock()
storage.minio_repository.store_dataframe_as_parquet = AsyncMock()
storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
@@ -222,3 +229,80 @@ def test___del__(storage):
storage.__del__()
storage.close.assert_called_once()
def test_estimate_payload_size_bytes(storage):
assert storage._estimate_payload_size_bytes({'x': 1}) > 0
@mark.asyncio
async def test_load_query_with_minio_offload_no_rows(storage):
storage.load_custom_query = AsyncMock(return_value=None)
result = await storage.load_query_with_minio_offload(
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
)
assert result['success'] is False
@mark.asyncio
async def test_load_query_with_minio_offload_inline(storage):
storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
result = await storage.load_query_with_minio_offload(
{**metadata, 'query': 'SELECT 1', 'model_name': 'my-model', 'key_prefix': 'predictions/s'}
)
assert result.get('success') is True
assert 'data' in result
assert result.get('object_key') is None
@mark.asyncio
@patch('laborious.utils.models.minio_dataframe_payload.MinioDataFramePayload.estimate_size_bytes')
async def test_load_query_with_minio_offload_minio(mock_estimate, storage):
mock_estimate.return_value = 10**9
storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
storage.minio_repository.upload_file = AsyncMock(
return_value={
'minio_object_name': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet'
}
)
storage.minio_repository.bucket = 'test'
fixed = datetime.datetime(2024, 1, 15, 12, 30, 45)
with patch('laborious.utils.models.minio_dataframe_payload.now', return_value=fixed):
result = await storage.load_query_with_minio_offload(
{**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'}
)
assert result.get('success') is True
assert result.get('data') is None
assert (
result['object_key']
== 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet'
)
storage.minio_repository.upload_file.assert_called_once()
@mark.asyncio
@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'})
@patch('laborious.activities.storage.now')
async def test_cleanup_minio_objects_expired(mock_now, storage):
mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0)
storage.minio_repository.list_objects = AsyncMock(
return_value=[
'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet',
]
)
storage.minio_repository.delete_file = AsyncMock()
storage.send_notification_async = AsyncMock()
result = await storage.cleanup_minio_objects_expired(
{**metadata, 'prefixes': ['training_datasets/m']}
)
assert result['success'] is True
assert result['deleted_count'] == 1
storage.minio_repository.delete_file.assert_called_once_with(
object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet',
metadata=metadata['metadata'],
)

54
tests/laborious/utils/models/test_minio_dataframe_payload.py Normal file
View 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']

245
tests/laborious/utils/repository/test_minio_repository.py
View File

@@ -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()

5
tests/laborious/workflows/subworkflows/test_prediction_process.py
View File

@@ -17,6 +17,7 @@ metadata = {
'model_name': 'test_model',
'workflow_name': 'test_workflow',
'schema_name': 'test_schedule',
'schedule_name': 'test_schedule',
},
}
@@ -101,6 +102,7 @@ async def test_run(workflow_mock, prediction_process):
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
'key_prefix': 'predictions/test_schedule',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -323,6 +325,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
'key_prefix': 'predictions/test_schedule',
**metadata,
},
retry_policy=ANY,
@@ -423,6 +426,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
'key_prefix': 'predictions/test_schedule',
**metadata,
},
retry_policy=ANY,
@@ -540,6 +544,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p
'data': input_data['data'],
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
'key_prefix': 'predictions/test_schedule',
**metadata,
},
retry_policy=ANY,

19
tests/laborious/workflows/test_minimal_retrain.py
View File

@@ -41,7 +41,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
{'success': True, 'object_key': 'test_object_key'},
{'data': {'a': [1]}, 'success': True},
{'success': True, 'experiment': 'test_experiment'},
{
'success': True,
@@ -58,13 +58,12 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.query_to_minio,
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -78,7 +77,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
Activities.retrain_model,
{
**metadata,
'object_key': 'test_object_key',
'data': {'data': {'a': [1]}, 'success': True},
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
},
@@ -163,7 +162,7 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
{'success': False, 'object_key': 'test_object_key'},
{'success': False, 'message': 'No data returned from query'},
{'success': True, 'experiment': 'test_experiment'},
{
'success': True,
@@ -178,13 +177,12 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim
await minimal_retrain.run(input_data)
workflow_mock.execute_activity_method.assert_called_once_with(
Activities.query_to_minio,
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -213,7 +211,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
workflow_mock.execute_activity_method = AsyncMock(
side_effect=[
{'success': True, 'object_key': 'test_object_key'},
{'data': {'a': [1]}, 'success': True},
{'success': False, 'experiment': 'test_experiment'},
{
'success': True,
@@ -230,13 +228,12 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.query_to_minio,
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -250,7 +247,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
Activities.retrain_model,
{
**metadata,
'object_key': 'test_object_key',
'data': {'data': {'a': [1]}, 'success': True},
'model_name': input_data['model_name'],
'model_config': input_data['model_config'],
},

13
tests/laborious/workflows/test_predictions_batch.py
View File

@@ -24,7 +24,10 @@ metadata = {
@mark.asyncio
@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
workflow_mock.execute_activity_method.return_value = {
'success': True,
'data': {'col': ['test_data']},
}
input_data = {
'schedule_name': 'test_schedule',
'model_name': 'test_model',
@@ -42,14 +45,16 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
await predictions_batch.run(input_data)
workflow_mock.execute_local_activity_method.assert_has_calls(
workflow_mock.execute_activity_method.assert_has_calls(
[
call(
Activities.load_custom_query,
Activities.load_query_with_minio_offload,
{
**metadata,
'query': input_data['query'],
'datetime_columns': input_data.get('datetime_columns', []),
'model_name': input_data['model_name'],
'key_prefix': f"predictions/{input_data['schedule_name']}",
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -58,7 +63,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
)
prediction_input = {
'metadata': metadata,
'data': {'data': 'test_data'},
'data': {'success': True, 'data': {'col': ['test_data']}},
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'transform_table_name': input_data['transform_table_name'],