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