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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user