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:
vitor-aignosi
2026-04-06 09:30:38 -03:00
parent d2bb588739
commit 18412958d2
4 changed files with 22 additions and 85 deletions

View File

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