diff --git a/model_manager/activities/mlflow.py b/model_manager/activities/mlflow.py index 9e5a900..0458c05 100644 --- a/model_manager/activities/mlflow.py +++ b/model_manager/activities/mlflow.py @@ -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), + } diff --git a/tests/activities/test_mlflow.py b/tests/activities/test_mlflow.py index 80b7e85..779f9d4 100644 --- a/tests/activities/test_mlflow.py +++ b/tests/activities/test_mlflow.py @@ -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