a66b996cc1
This commit refactors the MLFlow activities to focus on model training rather than prediction operations. It removes prediction-related activities and metrics, and updates the MLFlow activity descriptions to reflect the change in focus. The README is also updated to reflect these changes.
448 lines
15 KiB
Python
448 lines
15 KiB
Python
from unittest.mock import ANY, MagicMock, patch
|
|
|
|
import pytest
|
|
from pytest import fixture, mark
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
|
|
from model_manager.activities.mlflow import MLFlow
|
|
|
|
|
|
@patch('model_manager.activities.mlflow.MLFlowRepository')
|
|
def test___init__(mock_mlflow_repository):
|
|
mlflow = MLFlow(
|
|
mlflow_host='http://localhost',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
assert mlflow.mlflow_host == 'http://localhost'
|
|
assert mlflow.mlflow_port == 5000
|
|
assert mlflow.mlflow_username == 'admin'
|
|
assert mlflow.mlflow_password == 'admin'
|
|
|
|
mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)
|
|
|
|
|
|
@fixture
|
|
@patch('model_manager.activities.mlflow.MLFlowRepository')
|
|
def mlflow(mock_mlflow_repository):
|
|
mlflow = MLFlow(
|
|
mlflow_host='http://localhost:5000',
|
|
mlflow_port=5000,
|
|
mlflow_username='admin',
|
|
mlflow_password='admin',
|
|
logger=MagicMock(),
|
|
notification_handler=MagicMock(),
|
|
)
|
|
|
|
mlflow.send_notification = MagicMock()
|
|
|
|
return mlflow
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'test_workflow',
|
|
'schema_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_success(mlflow):
|
|
"""Test save_model successfully saves model and artifacts to MLflow."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
# Mock train result
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
train_result.run_name = None # Will be set by get_next_run_name
|
|
|
|
# Mock repository methods
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
|
mlflow.model_monitoring_repository.save_run.return_value = None
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method
|
|
response = await mlflow.save_model(input_data)
|
|
|
|
# Verify repository methods were called
|
|
mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
|
|
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
|
|
mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)
|
|
|
|
# Verify response - now returns TrainModelResult directly
|
|
assert response == train_result
|
|
assert response.run_name == 'test_experiment-1'
|
|
assert train_result.run_name == 'test_experiment-1'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_get_next_run_name_error(mlflow):
|
|
"""Test save_model handles error during get_next_run_name."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock error in get_next_run_name
|
|
mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
|
|
'MLflow connection error'
|
|
)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method - should raise exception
|
|
with pytest.raises(Exception, match='MLflow connection error'):
|
|
await mlflow.save_model(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='SAVE_MODEL_ERROR',
|
|
message=ANY,
|
|
block='save_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_generate_artifacts_error(mlflow):
|
|
"""Test save_model handles error during generate_artifacts."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock successful get_next_run_name but error in generate_artifacts
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
|
|
'Reports directory does not exist'
|
|
)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method - should raise exception
|
|
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
|
|
await mlflow.save_model(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='SAVE_MODEL_ERROR',
|
|
message=ANY,
|
|
block='save_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_save_run_error(mlflow):
|
|
"""Test save_model handles error during save_run."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock successful get_next_run_name and generate_artifacts but error in save_run
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
|
mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
|
|
'One or more metrics (MSE, R2, MAE) are None'
|
|
)
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method - should raise exception
|
|
with pytest.raises(ValueError, match=r'One or more metrics \(MSE, R2, MAE\) are None'):
|
|
await mlflow.save_model(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='SAVE_MODEL_ERROR',
|
|
message=ANY,
|
|
block='save_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_missing_metadata(mlflow):
|
|
"""Test save_model handles missing metadata gracefully."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'test_experiment'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
|
|
# Mock repository methods
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
|
|
mlflow.model_monitoring_repository.save_run.return_value = None
|
|
|
|
# Input data without metadata
|
|
input_data = {
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method
|
|
response = await mlflow.save_model(input_data)
|
|
|
|
# Verify it still works (metadata defaults to {}) - returns TrainModelResult directly
|
|
assert response == train_result
|
|
assert response.run_name == 'test_experiment-1'
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_save_model_complete_flow(mlflow):
|
|
"""Test save_model complete flow with all steps."""
|
|
from model_manager.utils.models.train_model_params import TrainModelParams
|
|
from model_manager.utils.models.train_model_result import TrainModelResult
|
|
|
|
params = MagicMock(spec=TrainModelParams)
|
|
params.experiment_name = 'production_model'
|
|
|
|
train_result = MagicMock(spec=TrainModelResult)
|
|
train_result.params = params
|
|
train_result.run_name = None
|
|
train_result.run_dir = None
|
|
train_result.report_path = None
|
|
|
|
# Mock complete flow
|
|
mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'
|
|
|
|
# After generate_artifacts, paths should be set
|
|
updated_result = MagicMock(spec=TrainModelResult)
|
|
updated_result.params = params
|
|
updated_result.run_name = 'production_model-5'
|
|
updated_result.run_dir = '/reports/production_model-5_20231010'
|
|
updated_result.report_path = '/reports/production_model-5_20231010/report.html'
|
|
updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
|
|
updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'
|
|
|
|
mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
|
|
mlflow.model_monitoring_repository.save_run.return_value = None
|
|
|
|
input_data = {
|
|
**metadata,
|
|
'train_result': train_result,
|
|
}
|
|
|
|
# Call the method
|
|
response = await mlflow.save_model(input_data)
|
|
|
|
# Verify complete flow
|
|
mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
|
|
mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
|
|
mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)
|
|
|
|
# Verify response - returns TrainModelResult directly
|
|
assert response == updated_result
|
|
assert response.run_name == 'production_model-5'
|
|
assert response.run_dir is not None
|
|
assert response.report_path is not None
|
|
|
|
|
|
# ============================================================================
|
|
# Tests for cleanup_run_directory
|
|
# ============================================================================
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_success(mlflow):
|
|
"""Test successful cleanup of run directory."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'test_run_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=True) as mock_exists,
|
|
patch('shutil.rmtree') as mock_rmtree,
|
|
):
|
|
# Call the method
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify directory existence was checked
|
|
mock_exists.assert_called_once_with('test_run_dir')
|
|
|
|
# Verify shutil.rmtree was called
|
|
mock_rmtree.assert_called_once_with('test_run_dir')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_already_deleted(mlflow):
|
|
"""Test cleanup when directory is already deleted (idempotent)."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'already_deleted_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=False) as mock_exists,
|
|
patch('shutil.rmtree') as mock_rmtree,
|
|
):
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify directory existence was checked
|
|
mock_exists.assert_called_once_with('already_deleted_dir')
|
|
|
|
# Verify shutil.rmtree was NOT called
|
|
mock_rmtree.assert_not_called()
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_no_run_dir(mlflow):
|
|
"""Test cleanup when no run_dir is provided."""
|
|
input_data = {
|
|
**metadata,
|
|
# No run_dir key
|
|
}
|
|
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Should complete without errors
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_empty_run_dir(mlflow):
|
|
"""Test cleanup when run_dir is empty string."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': '',
|
|
}
|
|
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Should complete without errors
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_none_run_dir(mlflow):
|
|
"""Test cleanup when run_dir is None."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': None,
|
|
}
|
|
|
|
# Call the method - should not raise error
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Should complete without errors
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_error(mlflow):
|
|
"""Test cleanup handles errors correctly."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'error_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=True),
|
|
patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
|
|
):
|
|
# Call the method - should raise exception
|
|
with pytest.raises(PermissionError, match='Permission denied'):
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once_with(
|
|
metadata=metadata['metadata'],
|
|
notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
|
|
message=ANY,
|
|
block='cleanup_run_directory',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_missing_metadata(mlflow):
|
|
"""Test cleanup handles missing metadata gracefully."""
|
|
input_data = {
|
|
'run_dir': 'no_metadata_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
|
|
# Call the method
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify directory was deleted
|
|
mock_rmtree.assert_called_once_with('no_metadata_dir')
|
|
|
|
|
|
@mark.asyncio
|
|
async def test_cleanup_run_directory_oserror(mlflow):
|
|
"""Test cleanup handles OSError correctly."""
|
|
input_data = {
|
|
**metadata,
|
|
'run_dir': 'os_error_dir',
|
|
}
|
|
|
|
# Patch os and shutil inside the activity method
|
|
with (
|
|
patch('os.path.exists', return_value=True),
|
|
patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
|
|
):
|
|
# Call the method - should raise exception
|
|
with pytest.raises(OSError, match='Directory not empty'):
|
|
await mlflow.cleanup_run_directory(input_data)
|
|
|
|
# Verify notification was sent
|
|
mlflow.send_notification.assert_called_once()
|
|
call_args = mlflow.send_notification.call_args[1]
|
|
assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
|
|
assert call_args['level'] == NotificationLevel.ERROR
|
|
assert 'Directory not empty' in call_args['message']
|