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