vitor-aignosi
aaf647efdf
Refactor MLFlow retraining logic to always use retrain method - Removed the conditional logic for full retraining, ensuring the `retrain` method is always called. - Updated the documentation in the `retrain_model` method to reflect the changes in retraining flow. - Adjusted tests to verify that the `retrain` method is invoked correctly, while ensuring `train` is not called when the full retrain flag is set.
746 lines
23 KiB
Python
746 lines
23 KiB
Python
from datetime import datetime
|
|
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',
|
|
},
|
|
}
|
|
|
|
|
|
def test_detect_and_parse_datetime_index_empty(mlflow):
|
|
df = pd.DataFrame()
|
|
out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
|
|
assert out.empty
|
|
|
|
|
|
def test_detect_and_parse_datetime_index_mixed_types_error(mlflow):
|
|
idx = pd.Index([pd.Timestamp('2020-01-01', tz='UTC'), 'x'])
|
|
df = pd.DataFrame({'a': [1, 2]}, index=idx)
|
|
with raises(ValueError):
|
|
mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
|
|
|
|
|
|
def test_detect_and_parse_datetime_index_invalid_string_error(mlflow):
|
|
idx = pd.Index(['bad-format'])
|
|
df = pd.DataFrame({'a': [1]}, index=idx)
|
|
with raises(ValueError):
|
|
mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
|
|
|
|
|
|
def test_detect_and_parse_datetime_index_unsupported_type_error(mlflow):
|
|
idx = pd.Index([pd.Period('2020-01', freq='M')])
|
|
df = pd.DataFrame({'a': [1]}, index=idx)
|
|
with raises(ValueError):
|
|
mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
|
|
|
|
|
|
def test_detect_and_parse_datetime_index_datetime_success(mlflow):
|
|
idx = pd.Index([datetime(2020, 1, 1, 0, 0, 0)])
|
|
df = pd.DataFrame({'a': [1]}, index=idx)
|
|
out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
|
|
assert out.index[0].endswith('+0000')
|
|
|
|
|
|
def test_detect_and_parse_datetime_index_timestamp_with_tz_success(mlflow):
|
|
idx = pd.DatetimeIndex([pd.Timestamp('2020-01-01 00:00:00', tz='UTC')])
|
|
df = pd.DataFrame({'a': [1]}, index=idx)
|
|
out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
|
|
assert out.index[0].endswith('+0000')
|
|
|
|
|
|
@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_success_without_transform_meta(mock_from_dataframe, mlflow):
|
|
ts = pd.Timestamp('2020-01-01', tz='UTC')
|
|
raw = pd.DataFrame(
|
|
{
|
|
'variable': ['v1'],
|
|
'timestamp': [ts],
|
|
'value': [1.0],
|
|
'created_at': [ts],
|
|
}
|
|
)
|
|
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, {})
|
|
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': {},
|
|
}
|
|
|
|
await mlflow.request_transform(input_data)
|
|
mock_from_dataframe.assert_called_once()
|
|
|
|
|
|
@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_success_dataframe_and_meta(
|
|
mock_to_datetime, mock_from_dataframe, mlflow
|
|
):
|
|
wrapper = MagicMock()
|
|
pred_df = pd.DataFrame({'raw': [0.3]})
|
|
wrapper.predict.return_value = (pred_df, {'m': 1})
|
|
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': {},
|
|
}
|
|
|
|
await mlflow.request_predict(input_data)
|
|
|
|
assert list(pred_df.columns) == ['prediction', 'response_time']
|
|
mlflow.info.assert_any_call("Wrapper predict metadata: {'m': 1}", metadata['metadata'])
|
|
mock_from_dataframe.assert_called_once()
|
|
|
|
|
|
@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'
|
|
|
|
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'
|
|
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.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_always_uses_retrain_even_with_full_retrain_flag(
|
|
mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow
|
|
):
|
|
mock_mkdtemp.return_value = 'tmp'
|
|
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
|
|
|
|
ts = pd.Timestamp('2020-01-01', tz='UTC')
|
|
raw_data = pd.DataFrame(
|
|
{
|
|
'variable': ['target', 'f1', 'target', 'f1'],
|
|
'timestamp': [ts, ts, ts + pd.Timedelta(days=1), ts + pd.Timedelta(days=1)],
|
|
'value': [1.0, 2.0, 3.0, 4.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', 'full_retrain': True, 'validation_fraction': 0.5},
|
|
}
|
|
)
|
|
|
|
wrapper.retrain.assert_called_once()
|
|
wrapper.train.assert_not_called()
|
|
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
|
|
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):
|
|
ts = pd.Timestamp('2020-01-01', tz='UTC')
|
|
raw_data = pd.DataFrame(
|
|
{
|
|
'variable': ['f1', 'f2'],
|
|
'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': {},
|
|
}
|
|
)
|
|
|
|
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_missing_csv_file_returns_none(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
|
|
|
|
with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='tmp'):
|
|
with patch('laborious.activities.mlflow.rmtree'):
|
|
with patch('laborious.activities.mlflow.Path') as mp:
|
|
mp.return_value.rglob.return_value = []
|
|
result = await mlflow.get_reference_data(input_data)
|
|
|
|
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
|