feat: enhance training and experiment tracking functionality

- Updated `Activities` class to improve garbage collection handling.
- Enhanced error messaging in `ExperimentTracking` for better clarity on update failures.
- Refactored `Training` class to streamline exception handling and improve type hints.
- Introduced new methods in `TrainModelParams` for better handling of experiment run IDs and model metadata.
- Added functionality to extract model equations in `DataManagerRepository` for linear regression models.

model_manager/activities/training.py

@@ -118,7 +118,7 @@ class Training(SientiaMonitoring):
             TrainModelParams: Validated and converted training parameters
 
         Raises:
-            ValueError, TypeError, KeyError: If validation fails (after sending notification)
+            Exception: If validation fails (after sending notification)
         """
         metadata = input_data.get('metadata', {})
         try:
@@ -134,7 +134,7 @@ class Training(SientiaMonitoring):
             )
 
             return train_params
-        except (ValueError, TypeError, KeyError) as e:
+        except Exception as e:
             error_msg = f'Error validating training parameters: {str(e)}'
             trace = traceback.format_exc()
 
@@ -165,7 +165,7 @@ class Training(SientiaMonitoring):
                 - train_params (TrainModelParams | dict): Training parameters.
 
         Returns:
-            dict: Key `run_name` when training and saving succeed.
+            dict[str, Any]: Serializable summary (run identifiers, run_dir for cleanup, regression metrics).
 
         Raises:
             ValueError: If input validation fails.
@@ -177,9 +177,6 @@ class Training(SientiaMonitoring):
         if isinstance(train_params, dict):
             train_params = TrainModelParams.from_dict(train_params)
 
-        model_trained = False
-        model_saved = False
-
         try:
             # Download training file bytes from MinIO
             train_bytes = await self.minio_repository.download_file(
@@ -238,10 +235,9 @@ class Training(SientiaMonitoring):
 
             train_result = self.data_manager_repository.compute_regression_metrics(
                 train_result,
+                wrapper,
             )
 
-            model_trained = True
-
             async with self.mlflow_repository.start_run(
                 model_name=train_params.model_name,
                 run_name=None,
@@ -266,7 +262,11 @@ class Training(SientiaMonitoring):
                 mlflow.log_artifact(train_result.train_data_path)
                 mlflow.log_artifact(train_result.test_data_path)
 
-            return train_result.to_dict()
+            return {
+                'run_name': train_result.run_name,
+                'run_id': train_result.run_id,
+                'run_dir': train_result.run_dir
+            }
         except Exception as e:  # noqa: BLE001
 
             error_msg = f'Error training model - error: {str(e)}'
@@ -284,7 +284,7 @@ class Training(SientiaMonitoring):
 
             raise e
 
-    activity.defn(name='cleanup_resources')
+    @activity.defn(name='cleanup_resources')
     async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
         """
         Cleanup temporary resources created during training.