SIENTIAPDE-1252: Implement save_model activity to save trained models to MLflow with comprehensive error handling and add corresponding unit tests.
This commit is contained in:
@@ -344,3 +344,107 @@ class MLFlow(BaseActivity):
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
@activity.defn(name='save_model')
|
||||
async def save_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Save a trained ML model and its artifacts to MLflow with comprehensive error handling.
|
||||
|
||||
This activity orchestrates the complete model saving pipeline:
|
||||
1. Generates the next run name for the experiment
|
||||
2. Creates and organizes artifacts (reports, data files)
|
||||
3. Logs model, parameters, metrics, and artifacts to MLflow
|
||||
4. Returns success/failure status with results or error message
|
||||
|
||||
The activity does NOT raise exceptions on failure - it catches all errors,
|
||||
sends notifications, and returns a failure status. This allows the workflow
|
||||
to handle the error gracefully and update the database accordingly.
|
||||
|
||||
Args:
|
||||
input_data: Configuration for model saving operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- train_result (TrainModelResult): Training result with model and metrics
|
||||
|
||||
Returns:
|
||||
dict: Save result with the following structure:
|
||||
{
|
||||
'success': bool, # True if saving succeeded, False otherwise
|
||||
'result': TrainModelResult | None, # Updated result if success=True
|
||||
'error_message': str | None # Error message if success=False
|
||||
}
|
||||
|
||||
Example:
|
||||
# Successful save
|
||||
result = await save_model({
|
||||
'metadata': {'workflow_id': 'save-123', 'experiment_run_id': 456},
|
||||
'train_result': TrainModelResult(...)
|
||||
})
|
||||
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
|
||||
|
||||
# Failed save
|
||||
# Returns: {'success': False, 'result': None, 'error_message': 'Error details...'}
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
train_result = input_data['train_result']
|
||||
|
||||
try:
|
||||
experiment_name = train_result.params.experiment_name
|
||||
|
||||
self.info(
|
||||
f'Starting model save for experiment: {experiment_name}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Step 1: Generate next run name
|
||||
self.info('Generating run name', metadata)
|
||||
train_result.run_name = self.model_monitoring_repository.get_next_run_name(
|
||||
experiment_name
|
||||
)
|
||||
self.info(f'Generated run name: {train_result.run_name}', metadata)
|
||||
|
||||
# Step 2: Generate artifacts (reports, CSV files)
|
||||
self.info('Generating artifacts', metadata)
|
||||
train_result = self.model_monitoring_repository.generate_artifacts(train_result)
|
||||
self.info('Artifacts generated successfully', metadata)
|
||||
|
||||
# Step 3: Save run to MLflow
|
||||
self.info('Saving run to MLflow', metadata)
|
||||
self.model_monitoring_repository.save_run(train_result)
|
||||
|
||||
self.info(
|
||||
f'Model saved successfully - Run: {train_result.run_name}, '
|
||||
f'Experiment: {experiment_name}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'result': train_result,
|
||||
'error_message': None,
|
||||
}
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Error saving model - Experiment: {train_result.params.experiment_name if train_result and train_result.params else "unknown"}, Error: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
|
||||
# Send notification (MongoDB)
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='SAVE_MODEL_ERROR',
|
||||
message=error_msg,
|
||||
block='save_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
# Log error with metadata
|
||||
self.error(trace, metadata=metadata)
|
||||
|
||||
# Return failure result (do NOT raise exception)
|
||||
# This allows workflow to update database with error status
|
||||
return {
|
||||
'success': False,
|
||||
'result': None,
|
||||
'error_message': str(e),
|
||||
}
|
||||
|
||||
@@ -299,3 +299,251 @@ async def test_update_production_model_error(mlflow):
|
||||
)
|
||||
else:
|
||||
raise AssertionError('No exception raised')
|
||||
|
||||
|
||||
@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
|
||||
assert response['success'] is True
|
||||
assert response['result'] == train_result
|
||||
assert response['error_message'] is None
|
||||
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
|
||||
response = await mlflow.save_model(input_data)
|
||||
|
||||
# Verify error handling
|
||||
assert response['success'] is False
|
||||
assert response['result'] is None
|
||||
assert 'MLflow connection error' in response['error_message']
|
||||
|
||||
# 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
|
||||
response = await mlflow.save_model(input_data)
|
||||
|
||||
# Verify error handling
|
||||
assert response['success'] is False
|
||||
assert response['result'] is None
|
||||
assert 'Reports directory does not exist' in response['error_message']
|
||||
|
||||
# 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
|
||||
response = await mlflow.save_model(input_data)
|
||||
|
||||
# Verify error handling
|
||||
assert response['success'] is False
|
||||
assert response['result'] is None
|
||||
assert 'One or more metrics (MSE, R2, MAE) are None' in response['error_message']
|
||||
|
||||
# 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 {})
|
||||
assert response['success'] is True
|
||||
assert response['result'] == train_result
|
||||
assert response['error_message'] is None
|
||||
|
||||
|
||||
@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
|
||||
assert response['success'] is True
|
||||
assert response['result'] == updated_result
|
||||
assert response['error_message'] is None
|
||||
assert updated_result.run_name == 'production_model-5'
|
||||
assert updated_result.run_dir is not None
|
||||
assert updated_result.report_path is not None
|
||||
|
||||
Reference in New Issue
Block a user