refactor: remove MinIO repository references and streamline cleanup process
- Eliminated MinIO repository dependencies from the `Activities`, `Cleanup`, and `Training` classes. - Updated the `cleanup_resources` method to focus on removing local temporary directories instead of handling MinIO file deletions. - Adjusted the `TrainModel` class to pass the run directory for cleanup, enhancing resource management in the training workflow.
This commit is contained in:
@@ -124,10 +124,7 @@ class TrainModel:
|
||||
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,
|
||||
run_dir=train_result.get('run_dir'),
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception:
|
||||
@@ -278,22 +275,13 @@ class TrainModel:
|
||||
|
||||
return train_result
|
||||
except Exception as e:
|
||||
# Mapear flags -> status
|
||||
# False/False: erro no treino
|
||||
# True/False: erro ao salvar (MLflow)
|
||||
# False/True: estado inconsistente, tratar como erro de treino
|
||||
# True/True: não deveria cair aqui; tratar como erro genérico de treino
|
||||
status = ExperimentStatus.TRAINING_ERROR
|
||||
|
||||
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,
|
||||
update_type=UpdateType.STATUS_WITH_ERROR,
|
||||
status=status,
|
||||
status=ExperimentStatus.TRAINING_ERROR,
|
||||
error_message=self._extract_error_message(e),
|
||||
)
|
||||
except Exception:
|
||||
@@ -302,54 +290,27 @@ class TrainModel:
|
||||
|
||||
async def _cleanup_resources(
|
||||
self,
|
||||
experiment_run_id: int,
|
||||
bucket_name: str,
|
||||
file_name: str,
|
||||
val_file_name: str | None,
|
||||
run_dir: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
Cleanup resources.
|
||||
|
||||
This method deletes the training file from MinIO. On success, updates
|
||||
DB status to FILE_DELETED.
|
||||
On error, updates DB status to FILE_DELETE_ERROR.
|
||||
This method removes the temporary run directory via activity.
|
||||
|
||||
Args:
|
||||
experiment_run_id: Validated experiment run ID
|
||||
run_dir: Temporary directory to remove
|
||||
metadata: Workflow execution metadata
|
||||
"""
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
Activities.cleanup_resources,
|
||||
{
|
||||
**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),
|
||||
)
|
||||
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.STATUS,
|
||||
status=ExperimentStatus.FILE_DELETED,
|
||||
)
|
||||
except Exception as 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
|
||||
await workflow.execute_activity_method(
|
||||
Activities.cleanup_resources,
|
||||
{
|
||||
**metadata,
|
||||
'run_dir': run_dir,
|
||||
},
|
||||
retry_policy=network_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
|
||||
)
|
||||
|
||||
async def _update_experiment_run(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user