feat: improve error handling and resource cleanup in training workflow

- Updated the `Training` class to raise `ModelTrainingError` on training failures for better error management.
- Enhanced the `run` method in `TrainModel` to return training results and ensure proper resource cleanup, including validation files.
- Refactored exception handling to prevent silent failures during resource cleanup and experiment run updates.
- Adjusted type hints for improved clarity and consistency in method signatures.

model_manager/activities/training.py

@@ -3,7 +3,7 @@ Training activities for ML model training operations.
 
 This module provides activities for training machine learning models.
 The activity extends BaseActivity and receives pre-downloaded files
-to return success/failure status without raising exceptions.
+and raises `ModelTrainingError` when training fails.
 """
 
 from temporalio import activity, workflow
@@ -35,8 +35,8 @@ class Training(SientiaMonitoring):
 
     This activity extends SientiaMonitoring and handles machine learning model
     training with comprehensive error handling. It receives pre-downloaded
-    files from the workflow and returns success/failure status without
+    files from the workflow and raises `ModelTrainingError` on failure so the
-    raising exceptions.
+    workflow can map the correct experiment status.
     """
 
     def __init__(
@@ -300,7 +300,7 @@ class Training(SientiaMonitoring):
         metadata = input_data.get('metadata', {})
         bucket_name = input_data.get('bucket_name', '')
         file_name = input_data.get('file_name', '')
-        metrics_status = 'success'
+        val_file_name = input_data.get('val_file_name')
 
         try:
             await self.minio_repository.delete_file(
@@ -308,6 +308,12 @@ class Training(SientiaMonitoring):
                 bucket=bucket_name,
                 metadata=metadata,
             )
+            if val_file_name:
+                await self.minio_repository.delete_file(
+                    object_name=val_file_name,
+                    bucket=bucket_name,
+                    metadata=metadata,
+                )
         except Exception as e:  # noqa: BLE001
             error_msg = (
                 'Error cleaning up resources - '
@@ -316,7 +322,7 @@ class Training(SientiaMonitoring):
 
             trace = traceback.format_exc()
 
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='CLEANUP_RESOURCES_ERROR',
                 message=error_msg,

model_manager/workflows/train_model.py

@@ -72,7 +72,7 @@ class TrainModel:
     """
 
     @workflow.run
-    async def run(self, input_data: dict[str, Any]) -> None:
+    async def run(self, input_data: dict[str, Any]) -> dict[str, str | None]:
         """
         Execute the complete model training workflow.
 
@@ -112,18 +112,29 @@ class TrainModel:
             input_data, experiment_run_id, metadata
         )
 
-        train_result = await self._train_model(
+        training_succeeded = False
-            train_params=train_params,
+        train_result: dict[str, str | None]
-            experiment_run_id=experiment_run_id,
-            metadata=metadata,
-        )
 
-        await self._cleanup_resources(
+        try:
-            experiment_run_id=experiment_run_id,
+            train_result = await self._train_model(
-            bucket_name=train_params.bucket_name,
+                train_params=train_params,
-            file_name=train_params.file_name,
+                experiment_run_id=experiment_run_id,
-            metadata=metadata,
+                metadata=metadata,
-        )
+            )
+            training_succeeded = True
+        finally:
+            try:
+                await self._cleanup_resources(
+                    experiment_run_id=experiment_run_id,
+                    bucket_name=train_params.bucket_name,
+                    file_name=train_params.file_name,
+                    val_file_name=train_params.val_file_name,
+                    metadata=metadata,
+                )
+            except Exception:
+                if training_succeeded:
+                    raise
+        return train_result
 
     def _validate_experiment_run_id(self, input_data: dict[str, Any]) -> int:
         """
@@ -147,12 +158,15 @@ class TrainModel:
         if experiment_run_id is None:
             raise ValueError('experiment_run_id is required but was not provided')
 
-        if not isinstance(experiment_run_id, int):
+        if isinstance(experiment_run_id, int):
-            raise ValueError(
+            return experiment_run_id
-                f'experiment_run_id must be an integer, got {type(experiment_run_id).__name__}'
-            )
 
-        return experiment_run_id
+        if isinstance(experiment_run_id, str) and experiment_run_id.strip().isdigit():
+            return int(experiment_run_id.strip())
+
+        raise ValueError(
+            f'experiment_run_id must be an integer or numeric string, got {type(experiment_run_id).__name__}'
+        )
 
     async def _validate_training_parameters(
         self,
@@ -208,14 +222,16 @@ class TrainModel:
 
             return train_params
         except Exception as e:
-            await self._update_experiment_run(
+            try:
-                metadata=metadata,
+                await self._update_experiment_run(
-                experiment_run_id=experiment_run_id,
+                    metadata=metadata,
-                update_type=UpdateType.STATUS_WITH_ERROR,
+                    experiment_run_id=experiment_run_id,
-                status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
+                    update_type=UpdateType.STATUS_WITH_ERROR,
-                error_message=self._extract_error_message(e),
+                    status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
-            )
+                    error_message=self._extract_error_message(e),
-
+                )
+            except Exception:
+                pass
             raise
 
     async def _train_model(
@@ -273,14 +289,16 @@ class TrainModel:
             if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
                 status = ExperimentStatus.TRACKING_SEND_ERROR
 
-            await self._update_experiment_run(
+            try:
-                metadata=metadata,
+                await self._update_experiment_run(
-                experiment_run_id=experiment_run_id,
+                    metadata=metadata,
-                update_type=UpdateType.STATUS_WITH_ERROR,
+                    experiment_run_id=experiment_run_id,
-                status=status,
+                    update_type=UpdateType.STATUS_WITH_ERROR,
-                error_message=self._extract_error_message(e),
+                    status=status,
-            )
+                    error_message=self._extract_error_message(e),
-
+                )
+            except Exception:
+                pass
             raise
 
     async def _cleanup_resources(
@@ -288,6 +306,7 @@ class TrainModel:
         experiment_run_id: int,
         bucket_name: str,
         file_name: str,
+        val_file_name: str | None,
         metadata: dict[str, Any],
     ) -> None:
         """
@@ -311,6 +330,7 @@ class TrainModel:
                     **metadata,
                     'bucket_name': bucket_name,
                     'file_name': file_name,
+                    'val_file_name': val_file_name,
                 },
                 retry_policy=network_retry_policy,
                 start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
@@ -323,14 +343,16 @@ class TrainModel:
                 status=ExperimentStatus.FILE_DELETED,
             )
         except Exception as e:
-            await self._update_experiment_run(
+            try:
-                metadata=metadata,
+                await self._update_experiment_run(
-                experiment_run_id=experiment_run_id,
+                    metadata=metadata,
-                update_type=UpdateType.STATUS_WITH_ERROR,
+                    experiment_run_id=experiment_run_id,
-                status=ExperimentStatus.FILE_DELETE_ERROR,
+                    update_type=UpdateType.STATUS_WITH_ERROR,
-                error_message=self._extract_error_message(e),
+                    status=ExperimentStatus.FILE_DELETE_ERROR,
-            )
+                    error_message=self._extract_error_message(e),
-
+                )
+            except Exception:
+                pass
             raise
 
     async def _update_experiment_run(
@@ -388,7 +410,9 @@ class TrainModel:
             if text and text not in message_parts:
                 message_parts.append(text)
 
-            current = getattr(current, 'cause', None)
+            cause = getattr(current, '__cause__', None)
+            context = getattr(current, '__context__', None)
+            current = cause if isinstance(cause, Exception) else context
 
         if not message_parts:
             return repr(exc)