SIENTIAPDE-1253: Implement activity for cleaning up temporary run directory and integrate into train model workflow. This change introduces a new activity for idempotent cleanup of the run directory after model training, replacing the direct directory removal in the workflow. This improves determinism and error handling. Also includes unit tests for the new activity.

This commit is contained in:
Bruno Domingues
2025-10-15 15:49:05 -03:00
parent 8ea98360c3
commit 1ec40bfb2a
4 changed files with 253 additions and 29 deletions

View File

@@ -12,7 +12,6 @@ This workflow orchestrates the complete ML model training process, including:
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
import shutil
from datetime import timedelta
from typing import Any
@@ -373,8 +372,8 @@ class TrainModel:
"""
Cleanup resources and delete file from MinIO.
This method removes the temporary run directory and deletes the training
file from MinIO. On success, updates DB status to FILE_DELETED.
This method removes the temporary run directory via activity and deletes
the training file from MinIO. On success, updates DB status to FILE_DELETED.
On error, updates DB status to FILE_DELETE_ERROR.
Args:
@@ -386,11 +385,22 @@ class TrainModel:
Exception: If cleanup fails (after updating DB status)
"""
try:
# Step 1: Remove temporary run directory
# Step 1: Remove temporary run directory via activity (deterministic)
if hasattr(saved_result, 'run_dir') and saved_result.run_dir:
shutil.rmtree(saved_result.run_dir)
cleanup_input = {
**metadata,
'run_dir': saved_result.run_dir,
}
await workflow.execute_activity_method(
Activities.cleanup_run_directory,
cleanup_input,
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=30),
)
workflow.logger.info(
f'Removed run directory: {saved_result.run_dir} for experiment {experiment_run_id}'
f'Run directory cleanup completed for experiment {experiment_run_id}'
)
# Step 2: Delete file from MinIO