SIENTIAPDE-1231

Update model retraining and reporting functionality

- Changed the GITHUB_BRANCH value in values.yaml to reflect the latest adjustments for retraining the courier.
- Enhanced the Gates class with a new method `format_retrain_report` to format retraining report data according to storage policies.
- Refactored the MLFlow class to improve error handling during model retraining and return structured output.
- Updated the model_repository to utilize the latest MLFlow API for retrieving model versions and improved logging.
- Modified the minimal_retrain workflow to conditionally update the production model based on retraining success.

laborious/activities/mlflow.py

@@ -243,30 +243,29 @@ class MLFlow(BaseActivity):
         data = data.dropna()
         data.columns.name = None
 
-        try:
-            retrain_output, experiment = self.model_monitoring_repository.retrain_model(
-                data=data,
-                model_name=model_name,
-                model_config=model_config
-            )
+        retrain_output = self.model_monitoring_repository.retrain_model(
+            data=data,
+            model_name=model_name,
+            model_config=model_config
+        )
 
-            return {
-                'status': retrain_output,
-                'timestamp': timestamp,
-                'experiment': experiment
-            }
-        except Exception as e:
-            trace = traceback.format_exc()
+        if not retrain_output['success']:
+
+            trace = retrain_output['traceback']
             self.send_notification(
                 metadata=metadata,
                 notification_id='RETRAIN_MODEL_ERROR',
-                message=f'Error retraining model {model_name}: {e}',
+                message=f'Error retraining model {model_name}: {retrain_output['message']}',
                 block='retrain_model',
                 level=NotificationLevel.ERROR,
                 attachment_content=trace
             )
             self.error(trace, metadata=metadata)
-            raise e
+
+        return {
+            **retrain_output,
+            'timestamp': timestamp
+        }
 
     @activity.defn(name="update_production_model")
     async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
@@ -307,11 +306,7 @@ class MLFlow(BaseActivity):
         """
         metadata = input_data['metadata']
         model_name = input_data['model_name']
-        model_id = input_data['model_id']
         experiment = input_data['experiment']
-        timestamp = input_data['timestamp']
-        status = input_data['status']
-
         self.info(
             f'Updating production model {model_name} from experiment {experiment}...', metadata)
 
@@ -321,15 +316,9 @@ class MLFlow(BaseActivity):
                 model_name=model_name
             )
 
-            report = DataFrame([response])
-            report['model_id'] = model_id
-            report['model_name'] = model_name
-            report['timestamp'] = timestamp
-            report['status'] = status
-
             self.info(
                 f'Production model {model_name} updated successfully', metadata)
-            return report.to_dict()
+            return response
 
         except Exception as e:
             trace = traceback.format_exc()