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.
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user