diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index 10cb4dd..d86cdcf 100644 --- a/model_manager/activities/training.py +++ b/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 - raising exceptions. + files from the workflow and raises `ModelTrainingError` on failure so the + 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, diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index 8657c93..63273cb 100644 --- a/model_manager/workflows/train_model.py +++ b/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( - train_params=train_params, - experiment_run_id=experiment_run_id, - metadata=metadata, - ) + training_succeeded = False + train_result: dict[str, str | None] - await self._cleanup_resources( - experiment_run_id=experiment_run_id, - bucket_name=train_params.bucket_name, - file_name=train_params.file_name, - metadata=metadata, - ) + try: + train_result = await self._train_model( + train_params=train_params, + experiment_run_id=experiment_run_id, + 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): - raise ValueError( - f'experiment_run_id must be an integer, got {type(experiment_run_id).__name__}' - ) + if isinstance(experiment_run_id, int): + return experiment_run_id - 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( - metadata=metadata, - experiment_run_id=experiment_run_id, - update_type=UpdateType.STATUS_WITH_ERROR, - status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR, - error_message=self._extract_error_message(e), - ) - + try: + await self._update_experiment_run( + metadata=metadata, + experiment_run_id=experiment_run_id, + update_type=UpdateType.STATUS_WITH_ERROR, + 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( - metadata=metadata, - experiment_run_id=experiment_run_id, - update_type=UpdateType.STATUS_WITH_ERROR, - status=status, - error_message=self._extract_error_message(e), - ) - + try: + await self._update_experiment_run( + metadata=metadata, + experiment_run_id=experiment_run_id, + update_type=UpdateType.STATUS_WITH_ERROR, + 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( - metadata=metadata, - experiment_run_id=experiment_run_id, - update_type=UpdateType.STATUS_WITH_ERROR, - status=ExperimentStatus.FILE_DELETE_ERROR, - error_message=self._extract_error_message(e), - ) - + try: + await self._update_experiment_run( + metadata=metadata, + experiment_run_id=experiment_run_id, + update_type=UpdateType.STATUS_WITH_ERROR, + 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)