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.

model_manager/activities/mlflow.py

@@ -429,3 +429,67 @@ class MLFlow(BaseActivity):
 
             # Re-raise exception to stop workflow
             raise
+
+    @activity.defn(name='cleanup_run_directory')
+    async def cleanup_run_directory(self, input_data: dict[str, Any]) -> None:
+        """
+        Clean up temporary run directory after model training.
+
+        This activity deletes the temporary directory created during model training
+        and artifact generation. It implements idempotent cleanup to handle cases
+        where the directory may have already been deleted.
+
+        Args:
+            input_data: Configuration for cleanup operation
+                Required keys:
+                    - metadata (dict): Workflow execution metadata
+                    - run_dir (str): Path to the run directory to delete
+
+        Raises:
+            Exception: If cleanup fails for reasons other than directory not existing
+
+        Example:
+            await cleanup_run_directory({
+                'metadata': {'workflow_id': 'cleanup-123'},
+                'run_dir': '/path/to/run_dir'
+            })
+        """
+        import os
+        import shutil
+
+        metadata = input_data.get('metadata', {})
+        run_dir = input_data.get('run_dir')
+
+        try:
+            if not run_dir:
+                self.info('No run directory specified, skipping cleanup', metadata)
+                return
+
+            self.info(f'Cleaning up run directory: {run_dir}', metadata)
+
+            # Idempotent cleanup: check if directory exists before deleting
+            if os.path.exists(run_dir):
+                shutil.rmtree(run_dir)
+                self.info(f'Run directory deleted successfully: {run_dir}', metadata)
+            else:
+                self.info(f'Run directory already deleted: {run_dir}', metadata)
+
+        except Exception as e:
+            error_msg = f'Error cleaning up run directory {run_dir}: {str(e)}'
+            trace = traceback.format_exc()
+
+            # Send notification
+            self.send_notification(
+                metadata=metadata,
+                notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
+                message=error_msg,
+                block='cleanup_run_directory',
+                level=NotificationLevel.ERROR,
+                attachment_content=trace,
+            )
+
+            # Log error
+            self.error(trace, metadata=metadata)
+
+            # Re-raise exception
+            raise

model_manager/workflows/train_model.py

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