vitor-aignosi
1ce8b9d3a7
Enhance environment configuration and update dependencies - Added new environment variables for PluginStore and MLflow configuration in `.env.example`, including `RUNTIME`, `STORE_BASE_URL`, `STORE_OWNER`, `STORE_REPO`, `STORE_BRANCH`, `STORE_USERNAME`, `STORE_PASSWORD`, `STORE_CACHE_TTL_SECONDS`, `PYPI_SERVER`, `PYPI_USERNAME`, and `PYPI_PASSWORD`. - Updated `git-requirements-mapping.txt` to reflect changes in repository names. - Modified `requirements-light.txt` and `requirements.txt` to upgrade `sientia-dataops-library` to version 1.12.0 and `sientia-mlops-library` to version 0.8.1. - Updated `values.yaml` to include new environment variables for worker runtime and PluginStore configuration. - Refactored E2E tests to utilize new MLflow repository stubs and PluginStore mocks for improved testing accuracy.
578 lines
18 KiB
Python
578 lines
18 KiB
Python
from unittest.mock import ANY, AsyncMock, MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
from pytest import fixture, mark, raises
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
|
|
|
from laborious.activities.mlflow import MLFlow
|
|
|
|
|
|
@fixture(autouse=True)
|
|
def _passthrough_from_dict():
|
|
with patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
|
|
):
|
|
yield
|
|
|
|
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def test___init__(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
mlflow_repo = MagicMock()
|
|
plugin_store = MagicMock()
|
|
|
|
minio_repo = mock_minio_repository(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
mlflow = MLFlow(
|
|
mlflow_repository=mlflow_repo,
|
|
plugin_store=plugin_store,
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
assert mlflow.mlflow_repository is mlflow_repo
|
|
assert mlflow.plugin_store is plugin_store
|
|
|
|
mock_minio_repository.assert_called_once_with(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=ANY,
|
|
notification_handler=ANY,
|
|
metrics_controller=ANY,
|
|
bucket='test',
|
|
)
|
|
|
|
|
|
@fixture
|
|
@patch('laborious.activities.mlflow.MinioRepository')
|
|
def mlflow(mock_minio_repository):
|
|
logger = MagicMock()
|
|
notification_handler = MagicMock()
|
|
metrics_controller = AsyncMock()
|
|
mlflow_repo = MagicMock()
|
|
plugin_store = MagicMock()
|
|
|
|
minio_repo = mock_minio_repository(
|
|
endpoint='localhost:9000',
|
|
access_key='minio',
|
|
secret_key='minio123',
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
bucket='test',
|
|
)
|
|
mlflow = MLFlow(
|
|
mlflow_repository=mlflow_repo,
|
|
plugin_store=plugin_store,
|
|
minio_repository=minio_repo,
|
|
logger=logger,
|
|
notification_handler=notification_handler,
|
|
metrics_controller=metrics_controller,
|
|
)
|
|
|
|
mlflow.minio_repository = AsyncMock()
|
|
|
|
mlflow.send_notification = MagicMock()
|
|
mlflow.emit_metric = AsyncMock()
|
|
mlflow.send_notification_async = AsyncMock()
|
|
mlflow.error = MagicMock()
|
|
mlflow.debug = MagicMock()
|
|
mlflow.info = MagicMock()
|
|
mlflow.warning = MagicMock()
|
|
mlflow.critical = 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.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_request_transform_success(mock_from_dataframe, mlflow):
|
|
ts = pd.Timestamp('2020-01-01', tz='UTC')
|
|
raw = pd.DataFrame(
|
|
{
|
|
'variable': ['v1', 'v1'],
|
|
'timestamp': [ts, ts],
|
|
'value': [1.0, 2.0],
|
|
'created_at': [ts, ts],
|
|
}
|
|
)
|
|
pivoted = raw.sort_values('created_at', ascending=False).drop_duplicates(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
pivoted = pivoted.pivot(index='timestamp', columns='variable', values='value')
|
|
pivoted = pivoted.fillna(np.nan)
|
|
pivoted.columns.name = None
|
|
pivoted.index.name = None
|
|
pivoted['timestamp'] = pivoted.index
|
|
|
|
out_idx = pd.Index([ts.strftime(DATETIME_FORMAT_WITH_TZ)], name=None)
|
|
out_df = pd.DataFrame({'v1': [1.0]}, index=out_idx)
|
|
wrapper = MagicMock()
|
|
wrapper.transform.return_value = (out_df, {'meta': True})
|
|
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
|
|
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
response_data = await mlflow.request_transform(input_data)
|
|
|
|
mlflow.mlflow_repository.get_cached_model.assert_called_once_with(
|
|
model_name='test_model',
|
|
alias='production',
|
|
retention_minutes=0,
|
|
metadata=metadata['metadata'],
|
|
)
|
|
mock_from_dataframe.assert_called_once()
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
async def test_request_transform_failure(mock_from_dataframe, mlflow):
|
|
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('boom')
|
|
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
data_mock.sort_values.return_value = data_mock
|
|
data_mock.drop_duplicates.return_value = data_mock
|
|
data_mock.pivot.return_value = data_mock
|
|
|
|
await mlflow.request_transform(input_data)
|
|
|
|
mock_from_dataframe.assert_called_once_with(
|
|
dataframe=None,
|
|
minio_repo=mlflow.minio_repository,
|
|
model_name='test_model',
|
|
operation='transform',
|
|
status={'success': False, 'content': ANY},
|
|
workflow_metadata=metadata['metadata'],
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=mlflow.logger,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
|
|
wrapper = MagicMock()
|
|
pred_df = MagicMock()
|
|
wrapper.predict.return_value = (pred_df, {})
|
|
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
|
|
|
|
data_mock = MagicMock()
|
|
data_mock.index = pd.DatetimeIndex([pd.Timestamp('2020-01-01', tz='UTC')])
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
pred_df.columns = MagicMock()
|
|
pred_df.__setitem__ = MagicMock()
|
|
|
|
response_data = await mlflow.request_predict(input_data)
|
|
|
|
data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
|
|
mock_to_datetime.assert_called()
|
|
mlflow.mlflow_repository.get_cached_model.assert_called_once()
|
|
mock_from_dataframe.assert_called_once()
|
|
assert response_data == mock_from_dataframe.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
@patch(
|
|
'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
|
|
new_callable=AsyncMock,
|
|
)
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
|
|
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('predict boom')
|
|
|
|
data_mock = MagicMock()
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=data_mock)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
|
|
await mlflow.request_predict(input_data)
|
|
|
|
mock_from_dataframe.assert_called_once_with(
|
|
dataframe=None,
|
|
minio_repo=mlflow.minio_repository,
|
|
model_name='test_model',
|
|
operation='predict',
|
|
status={'success': False, 'content': ANY},
|
|
workflow_metadata=metadata['metadata'],
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=mlflow.logger,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.mlflow.log_artifact')
|
|
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
|
|
@patch('laborious.activities.mlflow.rmtree')
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_data_success_retrain(
|
|
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
|
|
):
|
|
mock_mkdtemp.return_value = '/tmp/x'
|
|
|
|
mv_alias = MagicMock()
|
|
mv_alias.run_id = 'source-run'
|
|
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
|
|
|
|
wrapper = MagicMock()
|
|
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
|
|
|
|
mock_cm = MagicMock()
|
|
mock_cm.__enter__.return_value = MagicMock(run_id='new-run', experiment_id='exp-1')
|
|
mock_cm.__exit__.return_value = False
|
|
mlflow.mlflow_repository.start_run.return_value = mock_cm
|
|
|
|
ts = pd.Timestamp('2020-01-01', tz='UTC')
|
|
raw_data = pd.DataFrame(
|
|
{
|
|
'variable': ['target', 'f1'],
|
|
'timestamp': [ts, ts],
|
|
'value': [1.0, 2.0],
|
|
}
|
|
)
|
|
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
},
|
|
}
|
|
)
|
|
|
|
wrapper.retrain.assert_called_once()
|
|
wrapper.store_model.assert_called_once_with(name='test_model')
|
|
assert response['success'] is True
|
|
assert response['experiment']['run_id'] == 'new-run'
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.mlflow.log_artifact')
|
|
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
|
|
@patch('laborious.activities.mlflow.rmtree')
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_with_payload_data(
|
|
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
|
|
):
|
|
mock_mkdtemp.return_value = '/tmp/x'
|
|
mv_alias = MagicMock(run_id='src')
|
|
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
|
|
wrapper = MagicMock()
|
|
mlflow.mlflow_repository.get_cached_model.return_value = wrapper
|
|
mock_cm = MagicMock()
|
|
mock_cm.__enter__.return_value = MagicMock(run_id='r', experiment_id='e')
|
|
mock_cm.__exit__.return_value = False
|
|
mlflow.mlflow_repository.start_run.return_value = mock_cm
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
|
|
raw_data.__getitem__.return_value.max.return_value = 'ts'
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
pivoted = MagicMock()
|
|
raw_data.sort_values.return_value = raw_data
|
|
raw_data.drop_duplicates.return_value = raw_data
|
|
raw_data.pivot.return_value = pivoted
|
|
pivoted.fillna = MagicMock()
|
|
pivoted.columns.name = None
|
|
pivoted.index = MagicMock()
|
|
pivoted.__setitem__ = MagicMock()
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {'target': 'target'},
|
|
}
|
|
)
|
|
|
|
assert response['success'] is True
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
|
|
mv_alias = MagicMock(run_id='src')
|
|
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
|
|
mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('retrain failed')
|
|
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
|
|
raw_data.__getitem__.return_value.max.return_value = 'tsmax'
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
pivoted = MagicMock()
|
|
raw_data.sort_values.return_value = raw_data
|
|
raw_data.drop_duplicates.return_value = raw_data
|
|
raw_data.pivot.return_value = pivoted
|
|
pivoted.fillna = MagicMock()
|
|
pivoted.columns.name = None
|
|
pivoted.index = MagicMock()
|
|
pivoted.__setitem__ = MagicMock()
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {'target': 'target'},
|
|
}
|
|
)
|
|
|
|
assert response['success'] is False
|
|
mlflow.send_notification_async.assert_called_once()
|
|
assert 'retrain failed' in response['message']
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_data_error(mlflow):
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert response['success'] is False
|
|
assert 'data' in response['message'].lower() or 'loading' in response['message'].lower()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_missing_target(mlflow):
|
|
raw_data = MagicMock(columns=['variable', 'timestamp', 'value'])
|
|
raw_data.__getitem__.return_value.max.return_value = 'ts'
|
|
payload = AsyncMock()
|
|
payload.retrieve = AsyncMock(return_value=raw_data)
|
|
|
|
pivoted = MagicMock()
|
|
raw_data.drop_duplicates.return_value = raw_data
|
|
raw_data.pivot.return_value = pivoted
|
|
pivoted.fillna = MagicMock()
|
|
pivoted.columns.name = None
|
|
pivoted.index = MagicMock()
|
|
pivoted.__setitem__ = MagicMock()
|
|
|
|
response = await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'data': payload,
|
|
'model_name': 'test_model',
|
|
'model_config': {},
|
|
}
|
|
)
|
|
|
|
assert response['success'] is False
|
|
assert 'target' in response['message']
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_retrain_model_data_error_no_minio_repository(mlflow):
|
|
mlflow.minio_repository = None
|
|
|
|
with raises(ValueError) as e:
|
|
await mlflow.retrain_model(
|
|
{
|
|
**metadata,
|
|
'object_key': 'test_object_key',
|
|
'model_name': 'test_model',
|
|
'model_config': {
|
|
'target': 'target',
|
|
},
|
|
}
|
|
)
|
|
|
|
assert str(e.value) == 'Minio repository not initialized'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_update_production_model(mlflow):
|
|
mlflow.mlflow_repository._client.search_model_versions.return_value = [
|
|
MagicMock(version='3', run_id='run-x'),
|
|
MagicMock(version='2', run_id='run-x'),
|
|
]
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
'model_id': 1,
|
|
'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
|
|
'timestamp': 2,
|
|
'status': 'success',
|
|
}
|
|
|
|
response = await mlflow.update_production_model(input_data)
|
|
|
|
mlflow.mlflow_repository.promote_to_alias.assert_called_once_with(
|
|
model_name='test_model',
|
|
version='3',
|
|
alias='production',
|
|
metadata=metadata['metadata'],
|
|
)
|
|
|
|
assert response['model_name'] == 'test_model'
|
|
assert response['version'] == '3'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_update_production_model_error(mlflow):
|
|
mlflow.mlflow_repository._client.search_model_versions.return_value = []
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
'model_id': 1,
|
|
'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
|
|
'timestamp': 2,
|
|
'status': 'success',
|
|
}
|
|
|
|
try:
|
|
await mlflow.update_production_model(input_data)
|
|
except Exception:
|
|
mlflow.send_notification_async.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
|
message=ANY,
|
|
block='update_production_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
else:
|
|
raise AssertionError('Expected exception')
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.activities.mlflow.to_datetime')
|
|
async def test_get_reference_data_success(mock_to_datetime, mlflow):
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mv = MagicMock(run_id='run1')
|
|
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
|
|
|
|
mock_reference_data = MagicMock()
|
|
mock_reference_data.__getitem__.return_value = MagicMock()
|
|
mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
|
|
mock_reference_data.to_dict.return_value = [
|
|
{'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
|
|
]
|
|
|
|
with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
|
|
with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
|
|
with patch('laborious.activities.mlflow.rmtree'):
|
|
with patch('laborious.activities.mlflow.Path') as mp:
|
|
mp.return_value.rglob.return_value = [MagicMock()]
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
mock_reference_data.to_dict.assert_called_once_with(orient='records')
|
|
assert result == mock_reference_data.to_dict.return_value
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_reference_data_not_found(mlflow):
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mlflow.mlflow_repository._client.get_model_version_by_alias.side_effect = Exception('missing')
|
|
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
mlflow.warning.assert_called()
|
|
assert result is None
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_get_reference_data_exception(mlflow):
|
|
input_data = {
|
|
**metadata,
|
|
'model_name': 'test_model',
|
|
}
|
|
|
|
mv = MagicMock(run_id='run1')
|
|
mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
|
|
mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')
|
|
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
assert result is None
|