SIENTIAPDE-1712
Remove code validation script and refactor imports in activities and workflows - Deleted the `validate.sh` script, which was responsible for running code quality checks. - Cleaned up import statements in `activities.py`, `gates.py`, `mlflow.py`, and `storage.py` by removing unused imports and organizing them. - Refactored initialization methods in `MinioManager` and `MLFlow` classes for improved readability. - Updated various workflows to ensure compatibility with the new structure and removed unnecessary comments. - Enhanced test cases to accommodate changes in the activities and workflows, ensuring proper mocking of dependencies.
This commit is contained in:
@@ -1,8 +1,14 @@
|
||||
from datetime import datetime
|
||||
from io import BytesIO
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from pytest import mark
|
||||
import pytest
|
||||
from pandas import DataFrame
|
||||
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from laborious.utils.models.minio_dataframe_payload import (
|
||||
MinioDataFramePayload,
|
||||
_build_object_key,
|
||||
)
|
||||
|
||||
|
||||
def test_parse_object_timestamp_hyphenated_model():
|
||||
@@ -21,34 +27,190 @@ 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_estimate_size_bytes_returns_positive_for_nonempty_frame():
|
||||
df = DataFrame({'a': [1, 2]})
|
||||
size = MinioDataFramePayload.estimate_size_bytes(df)
|
||||
assert isinstance(size, int)
|
||||
assert size > 0
|
||||
|
||||
|
||||
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_when_offloaded_returns_object_prefix():
|
||||
payload = MinioDataFramePayload(
|
||||
last_timestamp='t',
|
||||
data=None,
|
||||
object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet',
|
||||
object_prefix='training_datasets/m',
|
||||
)
|
||||
assert MinioDataFramePayload.cleanup_prefix(payload) == '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'
|
||||
def test_cleanup_prefix_when_inline_returns_none():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
|
||||
assert MinioDataFramePayload.cleanup_prefix(payload) is None
|
||||
|
||||
|
||||
@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
|
||||
def test_has_data_true_when_object_key_set():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k')
|
||||
assert payload.has_data() is True
|
||||
|
||||
|
||||
@mark.asyncio
|
||||
async def test_dataframe_from_wire_list():
|
||||
df = await MinioDataFramePayload.dataframe_from_wire([{'a': 1}], None, {})
|
||||
assert list(df.columns) == ['a']
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_inline_dict_as_dataframe():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]})
|
||||
minio = AsyncMock()
|
||||
out = await payload.retrieve(minio, {'metadata': {}})
|
||||
assert list(out.columns) == ['a']
|
||||
minio.download_file.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_downloads_parquet_when_offloaded():
|
||||
source = DataFrame({'a': [1, 2]})
|
||||
buf = BytesIO()
|
||||
source.to_parquet(buf, engine='pyarrow', index=True)
|
||||
file_bytes = buf.getvalue()
|
||||
|
||||
payload = MinioDataFramePayload(
|
||||
last_timestamp='t',
|
||||
data=None,
|
||||
object_key='training_datasets/m/f.parquet',
|
||||
object_prefix='training_datasets/m',
|
||||
)
|
||||
minio = AsyncMock()
|
||||
minio.download_file = AsyncMock(return_value=file_bytes)
|
||||
|
||||
out = await payload.retrieve(minio, {'metadata': {}})
|
||||
|
||||
minio.download_file.assert_awaited_once_with(
|
||||
object_name='training_datasets/m/f.parquet',
|
||||
metadata={'metadata': {}},
|
||||
)
|
||||
assert list(out.columns) == ['a']
|
||||
|
||||
|
||||
def test_build_object_key():
|
||||
key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00')
|
||||
assert key == 'training_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet'
|
||||
assert prefix == 'training_datasets/my-model'
|
||||
|
||||
|
||||
def test_build_object_key_strips_slashes():
|
||||
key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45')
|
||||
assert prefix == 'training_datasets/my-model'
|
||||
assert key.startswith('training_datasets/my-model/')
|
||||
|
||||
|
||||
def test_estimate_size_bytes_fallback():
|
||||
df = DataFrame({'a': [1, 2]})
|
||||
original_to_dict = df.to_dict
|
||||
df.to_dict = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError('to_dict failed'))
|
||||
size = MinioDataFramePayload.estimate_size_bytes(df)
|
||||
df.to_dict = original_to_dict
|
||||
assert isinstance(size, int)
|
||||
assert size > 0
|
||||
|
||||
|
||||
def test_parse_object_timestamp_bad_datetime():
|
||||
key = 'p/m-initial-9999-99-99_99-99-99.parquet'
|
||||
assert MinioDataFramePayload.parse_object_timestamp(key) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve_empty_when_no_data():
|
||||
payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None)
|
||||
minio = AsyncMock()
|
||||
out = await payload.retrieve(minio, {})
|
||||
assert out.empty
|
||||
minio.download_file.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
async def test_from_dataframe_none(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = AsyncMock()
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
status={'success': False, 'message': 'no data'},
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.status == {'success': False, 'message': 'no data'}
|
||||
assert result.object_key is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
async def test_from_dataframe_empty(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = AsyncMock()
|
||||
mock_df = MagicMock()
|
||||
mock_df.__bool__ = MagicMock(return_value=True)
|
||||
mock_df.empty = True
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=mock_df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.object_key is None
|
||||
|
||||
|
||||
def _mock_dataframe(data_dict, timestamp_values=None):
|
||||
"""Build a MagicMock that behaves enough like a DataFrame for from_dataframe."""
|
||||
mock_df = MagicMock()
|
||||
mock_df.__bool__ = MagicMock(return_value=True)
|
||||
mock_df.empty = False
|
||||
if timestamp_values is None:
|
||||
timestamp_values = data_dict.get('timestamp', ['2024-01-01'])
|
||||
ts_col = MagicMock()
|
||||
ts_col.values.tolist.return_value = timestamp_values
|
||||
mock_df.__getitem__ = MagicMock(return_value=ts_col)
|
||||
mock_df.to_dict.return_value = data_dict
|
||||
buf = BytesIO()
|
||||
DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True)
|
||||
mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue()))
|
||||
return mock_df
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9)
|
||||
async def test_from_dataframe_inline():
|
||||
minio = AsyncMock()
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
)
|
||||
assert result.data is not None
|
||||
assert result.object_key is None
|
||||
assert result.last_timestamp == '2024-01-01'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.now')
|
||||
@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0)
|
||||
async def test_from_dataframe_offloaded(mock_now):
|
||||
mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0)
|
||||
minio = AsyncMock()
|
||||
minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'})
|
||||
minio.bucket = 'test-bucket'
|
||||
|
||||
df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]})
|
||||
result = await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=df,
|
||||
minio_repo=minio,
|
||||
model_name='m',
|
||||
operation='initial',
|
||||
workflow_metadata={'wf': 'data'},
|
||||
)
|
||||
assert result.data is None
|
||||
assert result.object_key == 'full/key.parquet'
|
||||
assert result.bucket == 'test-bucket'
|
||||
assert result.uri == 's3://test-bucket/full/key.parquet'
|
||||
minio.upload_file.assert_awaited_once()
|
||||
|
||||
@@ -2,6 +2,7 @@ from datetime import UTC, datetime
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
|
||||
|
||||
import mlflow as mlflow_lib
|
||||
import numpy as np
|
||||
import pytest
|
||||
from pandas import DataFrame, Timestamp
|
||||
|
||||
@@ -1310,10 +1311,8 @@ async def test_transform_success(mlflow_repository):
|
||||
mlflow_repository.get_cached_operation.return_value, metadata['metadata']
|
||||
)
|
||||
|
||||
assert output == {
|
||||
'success': True,
|
||||
'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value,
|
||||
}
|
||||
assert output['success'] is True
|
||||
assert output['content'] is mlflow_repository.detect_and_parse_datetime_index.return_value
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1355,7 +1354,7 @@ async def test_predict_success_array(mlflow_repository):
|
||||
|
||||
mlflow_repository.get_cached_operation.assert_called_once_with(
|
||||
model_name=model_name,
|
||||
data=data,
|
||||
data=ANY,
|
||||
operation='predict',
|
||||
retention=60,
|
||||
flavor='pyfunc',
|
||||
@@ -1363,10 +1362,12 @@ async def test_predict_success_array(mlflow_repository):
|
||||
)
|
||||
|
||||
assert output['success'] is True
|
||||
assert output['content'] == {
|
||||
'prediction': {'index_1': 2, 'index_2': 3},
|
||||
'response_time': {'index_1': ANY, 'index_2': ANY},
|
||||
}
|
||||
content = output['content']
|
||||
assert isinstance(content, DataFrame)
|
||||
assert 'prediction' in content.columns
|
||||
assert 'response_time' in content.columns
|
||||
assert list(content.columns) == ['prediction', 'response_time']
|
||||
assert content.index.tolist() == data.index.tolist()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1383,7 +1384,7 @@ async def test_predict_success_df(mlflow_repository):
|
||||
|
||||
mlflow_repository.get_cached_operation.assert_called_once_with(
|
||||
model_name=model_name,
|
||||
data=data,
|
||||
data=ANY,
|
||||
operation='predict',
|
||||
retention=60,
|
||||
flavor='pyfunc',
|
||||
@@ -1391,10 +1392,12 @@ async def test_predict_success_df(mlflow_repository):
|
||||
)
|
||||
|
||||
assert output['success'] is True
|
||||
assert output['content'] == {
|
||||
'prediction': {'index_1': 2, 'index_2': 3},
|
||||
'response_time': {'index_1': ANY, 'index_2': ANY},
|
||||
}
|
||||
content = output['content']
|
||||
assert isinstance(content, DataFrame)
|
||||
assert 'prediction' in content.columns
|
||||
assert 'response_time' in content.columns
|
||||
assert list(content.columns) == ['prediction', 'response_time']
|
||||
assert content.index.tolist() == data.index.tolist()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -1409,7 +1412,7 @@ async def test_predict_error(mlflow_repository):
|
||||
|
||||
mlflow_repository.get_cached_operation.assert_called_once_with(
|
||||
model_name=model_name,
|
||||
data=data,
|
||||
data=ANY,
|
||||
operation='predict',
|
||||
retention=60,
|
||||
flavor='pyfunc',
|
||||
@@ -1559,3 +1562,119 @@ def test_get_prediction_data_pyfunc(mlflow_repository):
|
||||
assert 'target' in result.columns
|
||||
assert 'timestamp' in result.columns
|
||||
assert result.index.tolist() == [0, 1]
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.pd.merge')
|
||||
@patch('laborious.utils.repository.model_repository.isinstance')
|
||||
@pytest.mark.asyncio
|
||||
async def test_fit_models_skip_transform(isinstance_mock, pd_merge, mlflow_repository):
|
||||
isinstance_mock.return_value = True
|
||||
|
||||
data_model = MagicMock(target_variable='feat_2')
|
||||
prediction_model = MagicMock()
|
||||
mlflow_repository.download_model = AsyncMock(
|
||||
side_effect=[(data_model, 'artifact_path'), (prediction_model, 'artifact_path')],
|
||||
)
|
||||
mlflow_repository.detect_and_parse_datetime_index = MagicMock(
|
||||
return_value=MagicMock(
|
||||
drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1']))
|
||||
)
|
||||
)
|
||||
mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
|
||||
|
||||
data = MagicMock()
|
||||
|
||||
output = await mlflow_repository.fit_models(
|
||||
'model_name',
|
||||
data,
|
||||
'latest_production_id',
|
||||
metadata['metadata'],
|
||||
'sklearn',
|
||||
True,
|
||||
'pyfunc',
|
||||
'feat_1',
|
||||
)
|
||||
|
||||
data_model.fit.assert_not_called()
|
||||
|
||||
assert output['data_model'] == {'model': data_model, 'artifact_path': 'artifact_path'}
|
||||
|
||||
|
||||
@patch('laborious.utils.repository.model_repository.force_memory_release')
|
||||
@patch('laborious.utils.repository.model_repository.path')
|
||||
@patch('laborious.utils.repository.model_repository.rmtree')
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_new_experiment_path_not_exists(
|
||||
_rmtree, path, force_memory_release, mlflow, mlflow_repository
|
||||
):
|
||||
model_name = 'model_name'
|
||||
data = MagicMock()
|
||||
prediction_data = MagicMock(spec=DataFrame)
|
||||
retrain_data = {
|
||||
'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
|
||||
'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
|
||||
'prediction_data': prediction_data,
|
||||
}
|
||||
|
||||
mlflow_repository.get_model_params = MagicMock(
|
||||
return_value={
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'pyfunc',
|
||||
'target_name': 'target_name',
|
||||
}
|
||||
)
|
||||
|
||||
mlflow_repository.get_experiment = MagicMock()
|
||||
mlflow_repository.get_next_run_name = MagicMock()
|
||||
mlflow_repository.log_model = AsyncMock()
|
||||
path.exists.return_value = False
|
||||
path.join.return_value = './tmp/artifacts/model_name'
|
||||
|
||||
await mlflow_repository.create_new_experiment(
|
||||
model_name,
|
||||
data,
|
||||
retrain_data,
|
||||
'latest_production_id',
|
||||
metadata['metadata'],
|
||||
'sklearn',
|
||||
'pyfunc',
|
||||
)
|
||||
|
||||
_rmtree.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_production_model_by_run_id_transition_error(mlflow, mlflow_repository):
|
||||
mlflow_repository.client.get_registered_model.return_value = MagicMock(
|
||||
latest_versions=[
|
||||
MagicMock(version='1'),
|
||||
MagicMock(version='2'),
|
||||
]
|
||||
)
|
||||
mlflow_repository.client.transition_model_version_stage.side_effect = Exception(
|
||||
'transition error'
|
||||
)
|
||||
|
||||
with pytest.raises(Exception, match='transition error'):
|
||||
await mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata'])
|
||||
|
||||
mlflow_repository.emit_metric.assert_called_with(
|
||||
metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_predict_success_ndarray(mlflow_repository):
|
||||
data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}})
|
||||
model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'}
|
||||
model_name = 'model'
|
||||
mlflow_repository.get_cached_operation = AsyncMock(return_value=np.array([5.0, 6.0]))
|
||||
|
||||
output = await mlflow_repository.predict(model_name, data, model_config, metadata['metadata'])
|
||||
|
||||
assert output['success'] is True
|
||||
content = output['content']
|
||||
assert isinstance(content, DataFrame)
|
||||
assert 'prediction' in content.columns
|
||||
assert 'response_time' in content.columns
|
||||
assert content.index.tolist() == data.index.tolist()
|
||||
|
||||
@@ -102,6 +102,7 @@ def test_build_minio_config_with_env_vars():
|
||||
'secret_key': 'test-secret',
|
||||
'region_name': 'test-region',
|
||||
'default_bucket': 'test-bucket',
|
||||
'retention_hours': 24,
|
||||
}
|
||||
|
||||
|
||||
@@ -117,4 +118,5 @@ def test_build_minio_config_with_defaults():
|
||||
'secret_key': 'minioadmin',
|
||||
'region_name': 'us-east-1',
|
||||
'default_bucket': 'laborious',
|
||||
'retention_hours': 24,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user