feat: enhance training model functionality and reporting

- Added `evidently` to requirements for improved model evaluation.
- Introduced `TrainModelResult` class with a `to_dict` method for better result handling.
- Updated `train_model` method to return a comprehensive training result, including run details.
- Enhanced `cleanup_run_directory` method in `DataManagerRepository` for improved resource management.
- Adjusted type hints in `TrainModel` for clarity and consistency.

model_manager/activities/training.py

@@ -23,6 +23,9 @@ with workflow.unsafe.imports_passed_through():
 
     from model_manager.utils.models.train_model_params import TrainModelParams
     from model_manager.utils.repository.data_manager_repository import DataManagerRepository
+    from model_manager.utils.models.train_model_result import TrainModelResult
+
+    import mlflow
 
 
 class Training(SientiaMonitoring):
@@ -146,7 +149,7 @@ class Training(SientiaMonitoring):
             raise
 
     @activity.defn(name='train_model')
-    async def train_model(self, input_data: dict[str, Any]) -> dict[str, str | None]:
+    async def train_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
         Train a machine learning model.
 
@@ -246,8 +249,6 @@ class Training(SientiaMonitoring):
                 tags=None,
                 metadata=metadata,
             ) as run_info:
-                wrapper.store_model(name=train_params.model_name)
-
                 train_result.run_name = run_info.run_name
                 train_result.run_id = run_info.run_id
 
@@ -256,12 +257,16 @@ class Training(SientiaMonitoring):
                     metadata=metadata,
                 )
 
-            model_saved = True
+                if train_result.report_path is None or train_result.train_data_path is None or train_result.test_data_path is None:
+                    raise ValueError('Report path, train data path, or test data path is not set')
 
-            return {
-                'run_name': run_info.run_name,
-                'run_id': run_info.run_id,
-            }
+                wrapper.store_model(name=train_params.model_name)
+
+                mlflow.log_artifact(train_result.report_path)
+                mlflow.log_artifact(train_result.train_data_path)
+                mlflow.log_artifact(train_result.test_data_path)
+
+            return train_result.to_dict()
         except Exception as e:  # noqa: BLE001
 
             error_msg = f'Error training model - error: {str(e)}'
@@ -296,7 +301,7 @@ class Training(SientiaMonitoring):
         run_dir = input_data.get('run_dir', '')
 
         try:
-            self.model_repository.cleanup_run_directory(run_dir)
+            self.data_manager_repository.cleanup_run_directory(run_dir, metadata)
         except Exception as e:  # noqa: BLE001
             error_msg = f'Error cleaning up resources - Run directory: {run_dir}, Error: {str(e)}'
 

model_manager/utils/models/train_model_result.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Any
 
 import pandas as pd
 
@@ -50,3 +51,9 @@ class TrainModelResult:
     test_data_path: str | None = None
 
     run_dir: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Convert TrainModelResult to a dictionary.
+        """
+        return self.__dict__

model_manager/utils/repository/data_manager_repository.py

@@ -15,8 +15,10 @@ and this repository focuses solely on preparing data structures for them.
 
 from datetime import datetime
 from io import BytesIO
+import json
 from os import makedirs, path
 from typing import Any
+from shutil import rmtree
 
 import numpy as np
 import pandas as pd
@@ -433,3 +435,24 @@ class DataManagerRepository(SientiaMonitoring):
                 json.dump(data.equation, f, indent=2, ensure_ascii=False)
 
         return data
+
+    def cleanup_run_directory(self, run_dir: str, metadata: dict[str, Any] | None = None) -> None:
+        """
+        Clean up temporary run directory after model training.
+
+        This activity deletes the temporary directory created during model training
+        and artifact generation. It implements idempotent cleanup to handle cases
+        where the directory may have already been deleted.
+
+        Args:
+            run_dir (str): Path to the run directory to delete
+        """
+        if not run_dir:
+            self.info('No run directory specified, skipping cleanup')
+            return
+
+        if path.exists(run_dir):
+            rmtree(run_dir)
+            self.info(f'Run directory deleted successfully: {run_dir}')
+        else:
+            self.info(f'Run directory already deleted: {run_dir}')

model_manager/workflows/train_model.py

@@ -71,7 +71,7 @@ class TrainModel:
     """
 
     @workflow.run
-    async def run(self, input_data: dict[str, Any]) -> dict[str, str | None]:
+    async def run(self, input_data: dict[str, Any]) -> dict[str, str | None] | None:
         """
         Execute the complete model training workflow.
 
@@ -112,7 +112,7 @@ class TrainModel:
         )
 
         training_succeeded = False
-        train_result: dict[str, str | None]
+        train_result: dict[str, str | None] | None = None
 
         try:
             train_result = await self._train_model(
@@ -123,6 +123,7 @@ class TrainModel:
             training_succeeded = True
         finally:
             try:
+                if train_result is not None:
                     await self._cleanup_resources(
                         run_dir=train_result.get('run_dir'),
                         metadata=metadata,
@@ -290,7 +291,7 @@ class TrainModel:
 
     async def _cleanup_resources(
         self,
-        run_dir: str,
+        run_dir: str | None,
         metadata: dict[str, Any],
     ) -> None:
         """
@@ -302,6 +303,9 @@ class TrainModel:
             run_dir: Temporary directory to remove
             metadata: Workflow execution metadata
         """
+        if run_dir is None:
+            return
+
         await workflow.execute_activity_method(
             Activities.cleanup_resources,
             {

requirements.txt

@@ -7,3 +7,4 @@ botocore==1.40.55
 git+https://github.com/Aignosi/sientia-dataops-library.git@v1.10.1
 prometheus-client==0.23.1
 beautifulsoup4==4.12.3
+evidently