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/gates.py

@@ -207,13 +207,12 @@ class Gates(BaseActivity):
         filter_output = []
 
         self.debug(
-            f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}", metadata)
+            f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}", metadata)
         self.debug(f"Filters: {filters}", metadata)
 
         comments = []
         for fil, config in filters.items():
             if fil not in mlflow_response_filter_functions:
-                self.error(f"Filter {fil} not found", metadata)
                 continue
             try:
                 if mlflow_response_filter_functions[fil](data, config):
@@ -293,7 +292,7 @@ class Gates(BaseActivity):
         filter_output = []
 
         self.debug(f"Input data:\n {data.head(5).to_string()}", metadata)
-        self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata)
+        self.debug(f"Filters: \n {filters}", metadata)
 
         for fil, config in filters.items():
             if fil not in mlflow_content_filter_functions:
@@ -492,6 +491,36 @@ class Gates(BaseActivity):
         self.info(f"Default prediction formatted: {data.size} rows", metadata)
         return data.to_dict()
 
+    @activity.defn(name="format_retrain_report")
+    async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+        """
+        Format retrain report data according to configured storage policies.
+        """
+        metadata = input_data['metadata']
+        self.info("Formatting retrain report...", metadata)
+
+        experiment_response = input_data['experiment_response']
+        update_report = input_data['update_report']
+        model_id = input_data['model_id']
+        model_name = input_data['model_name']
+
+        report = DataFrame({
+            'model_id': [model_id],
+            'model_name': [model_name],
+            'timestamp': [experiment_response['timestamp']],
+            'status': [experiment_response['message']]
+        })
+
+        if experiment_response['success']:
+            # Retrain was successfull
+            report['version'] = update_report['version']
+            report['mlflow_run_id'] = update_report['mlflow_run_id']
+            report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
+
+        self.debug(f"Retrain report: {report.to_csv()}", metadata)
+
+        return report.to_dict()
+
     @activity.defn(name="get_last_timestamp")
     async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
         """

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()