SIENTIAPDE-1252: Implement save_model activity to save trained models to MLflow with comprehensive error handling and add corresponding unit tests.

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),
+            }

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