gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: improve error handling and resource cleanup in training workflow

Browse Source 
- 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:
vitor-aignosi
2026-04-06 08:26:42 -03:00
parent 342a02d6f7
commit bf2b6b3888
2 changed files with 77 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
model_manager/activities/training.py
View File

@@ -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,

46
model_manager/workflows/train_model.py
View File

@@ -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
)
training_succeeded = False
train_result: dict[str, str | None]
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,13 +158,16 @@ 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
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,
input_data: dict[str, Any],
@@ -208,6 +222,7 @@ class TrainModel:
return train_params
except Exception as e:
try:
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
@@ -215,7 +230,8 @@ class TrainModel:
status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
error_message=self._extract_error_message(e),
)
except Exception:
pass
raise
async def _train_model(
@@ -273,6 +289,7 @@ class TrainModel:
if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
status = ExperimentStatus.TRACKING_SEND_ERROR
try:
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
@@ -280,7 +297,8 @@ class TrainModel:
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,6 +343,7 @@ class TrainModel:
status=ExperimentStatus.FILE_DELETED,
)
except Exception as e:
try:
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
@@ -330,7 +351,8 @@ class TrainModel:
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)