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
+        This method removes the temporary run directory via activity and deletes
-        file from MinIO. On success, updates DB status to FILE_DELETED.
+        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

tests/activities/test_mlflow.py

@@ -532,3 +532,165 @@ async def test_save_model_complete_flow(mlflow):
     assert response.run_name == 'production_model-5'
     assert response.run_dir is not None
     assert response.report_path is not None
+
+
+# ============================================================================
+# Tests for cleanup_run_directory
+# ============================================================================
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_success(mlflow):
+    """Test successful cleanup of run directory."""
+    input_data = {
+        **metadata,
+        'run_dir': 'test_run_dir',
+    }
+
+    # Patch os and shutil inside the activity method
+    with (
+        patch('os.path.exists', return_value=True) as mock_exists,
+        patch('shutil.rmtree') as mock_rmtree,
+    ):
+        # Call the method
+        await mlflow.cleanup_run_directory(input_data)
+
+        # Verify directory existence was checked
+        mock_exists.assert_called_once_with('test_run_dir')
+
+        # Verify shutil.rmtree was called
+        mock_rmtree.assert_called_once_with('test_run_dir')
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_already_deleted(mlflow):
+    """Test cleanup when directory is already deleted (idempotent)."""
+    input_data = {
+        **metadata,
+        'run_dir': 'already_deleted_dir',
+    }
+
+    # Patch os and shutil inside the activity method
+    with (
+        patch('os.path.exists', return_value=False) as mock_exists,
+        patch('shutil.rmtree') as mock_rmtree,
+    ):
+        # Call the method - should not raise error
+        await mlflow.cleanup_run_directory(input_data)
+
+        # Verify directory existence was checked
+        mock_exists.assert_called_once_with('already_deleted_dir')
+
+        # Verify shutil.rmtree was NOT called
+        mock_rmtree.assert_not_called()
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_no_run_dir(mlflow):
+    """Test cleanup when no run_dir is provided."""
+    input_data = {
+        **metadata,
+        # No run_dir key
+    }
+
+    # Call the method - should not raise error
+    await mlflow.cleanup_run_directory(input_data)
+
+    # Should complete without errors
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_empty_run_dir(mlflow):
+    """Test cleanup when run_dir is empty string."""
+    input_data = {
+        **metadata,
+        'run_dir': '',
+    }
+
+    # Call the method - should not raise error
+    await mlflow.cleanup_run_directory(input_data)
+
+    # Should complete without errors
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_none_run_dir(mlflow):
+    """Test cleanup when run_dir is None."""
+    input_data = {
+        **metadata,
+        'run_dir': None,
+    }
+
+    # Call the method - should not raise error
+    await mlflow.cleanup_run_directory(input_data)
+
+    # Should complete without errors
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_error(mlflow):
+    """Test cleanup handles errors correctly."""
+    input_data = {
+        **metadata,
+        'run_dir': 'error_dir',
+    }
+
+    # Patch os and shutil inside the activity method
+    with (
+        patch('os.path.exists', return_value=True),
+        patch('shutil.rmtree', side_effect=PermissionError('Permission denied')),
+    ):
+        # Call the method - should raise exception
+        with pytest.raises(PermissionError, match='Permission denied'):
+            await mlflow.cleanup_run_directory(input_data)
+
+        # Verify notification was sent
+        mlflow.send_notification.assert_called_once_with(
+            metadata=metadata['metadata'],
+            notification_id='CLEANUP_RUN_DIRECTORY_ERROR',
+            message=ANY,
+            block='cleanup_run_directory',
+            level=NotificationLevel.ERROR,
+            attachment_content=ANY,
+        )
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_missing_metadata(mlflow):
+    """Test cleanup handles missing metadata gracefully."""
+    input_data = {
+        'run_dir': 'no_metadata_dir',
+    }
+
+    # Patch os and shutil inside the activity method
+    with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree:
+        # Call the method
+        await mlflow.cleanup_run_directory(input_data)
+
+        # Verify directory was deleted
+        mock_rmtree.assert_called_once_with('no_metadata_dir')
+
+
+@mark.asyncio
+async def test_cleanup_run_directory_oserror(mlflow):
+    """Test cleanup handles OSError correctly."""
+    input_data = {
+        **metadata,
+        'run_dir': 'os_error_dir',
+    }
+
+    # Patch os and shutil inside the activity method
+    with (
+        patch('os.path.exists', return_value=True),
+        patch('shutil.rmtree', side_effect=OSError('Directory not empty')),
+    ):
+        # Call the method - should raise exception
+        with pytest.raises(OSError, match='Directory not empty'):
+            await mlflow.cleanup_run_directory(input_data)
+
+        # Verify notification was sent
+        mlflow.send_notification.assert_called_once()
+        call_args = mlflow.send_notification.call_args[1]
+        assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR'
+        assert call_args['level'] == NotificationLevel.ERROR
+        assert 'Directory not empty' in call_args['message']

tests/workflows/test_train_model.py

@@ -59,11 +59,9 @@ def mock_train_result(mock_train_params):
 
 
 @mark.asyncio
-@patch('model_manager.workflows.train_model.shutil')
 @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
 async def test_run_success_complete_flow(
     workflow_mock: AsyncMock,
-    mock_shutil: MagicMock,
     train_model_workflow: TrainModel,
     mock_train_params,
     mock_train_result,
@@ -101,6 +99,7 @@ async def test_run_success_complete_flow(
             None,  # update_experiment_run (TRAINING_SUCCESS)
             mock_train_result,  # save_model
             None,  # update_experiment_run (MLFLOW_SENT with run_name)
+            None,  # cleanup_run_directory
             None,  # delete_file_from_minio
             None,  # update_experiment_run (FILE_DELETED)
         ]
@@ -109,10 +108,8 @@ async def test_run_success_complete_flow(
     # Execute workflow
     await train_model_workflow.run(input_data)
 
-    # Verify all activity calls
+    # Verify all activity calls (now 10 instead of 9 due to cleanup_run_directory)
-    assert workflow_mock.execute_activity_method.call_count == 9
+    assert workflow_mock.execute_activity_method.call_count == 10
-    # Verify shutil.rmtree was called
-    mock_shutil.rmtree.assert_called_once_with('test_run_dir')
 
 
 @mark.asyncio
@@ -506,9 +503,7 @@ async def test_save_model_to_mlflow_error(
 
 @mark.asyncio
 @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-@patch('model_manager.workflows.train_model.shutil')
 async def test_cleanup_resources_success(
-    mock_shutil: MagicMock,
     workflow_mock: AsyncMock,
     train_model_workflow: TrainModel,
     mock_train_result,
@@ -518,6 +513,7 @@ async def test_cleanup_resources_success(
 
     workflow_mock.execute_activity_method = AsyncMock(
         side_effect=[
+            None,  # cleanup_run_directory
             None,  # delete_file_from_minio
             None,  # update_experiment_run
         ]
@@ -527,18 +523,13 @@ async def test_cleanup_resources_success(
         saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
     )
 
-    # Verify shutil.rmtree was called
+    # Verify activities were called (cleanup_run_directory + delete_file_from_minio + update_experiment_run)
-    mock_shutil.rmtree.assert_called_once_with('test_run_dir')
+    assert workflow_mock.execute_activity_method.call_count == 3
-
-    # Verify activities were called
-    assert workflow_mock.execute_activity_method.call_count == 2
 
 
 @mark.asyncio
 @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-@patch('model_manager.workflows.train_model.shutil')
 async def test_cleanup_resources_delete_error(
-    mock_shutil: MagicMock,
     workflow_mock: AsyncMock,
     train_model_workflow: TrainModel,
     mock_train_result,
@@ -548,6 +539,7 @@ async def test_cleanup_resources_delete_error(
 
     workflow_mock.execute_activity_method = AsyncMock(
         side_effect=[
+            None,  # cleanup_run_directory succeeds
             Exception('MinIO delete failed'),  # delete_file_from_minio fails
             None,  # update_experiment_run with error
         ]
@@ -558,15 +550,13 @@ async def test_cleanup_resources_delete_error(
             saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
         )
 
-    # Verify error status was updated
+    # Verify error status was updated (cleanup_run_directory + delete_file_from_minio + update_experiment_run)
-    assert workflow_mock.execute_activity_method.call_count == 2
+    assert workflow_mock.execute_activity_method.call_count == 3
 
 
 @mark.asyncio
 @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock)
-@patch('model_manager.workflows.train_model.shutil')
 async def test_cleanup_resources_without_run_dir(
-    mock_shutil: MagicMock,
     workflow_mock: AsyncMock,
     train_model_workflow: TrainModel,
     mock_train_result,
@@ -588,10 +578,8 @@ async def test_cleanup_resources_without_run_dir(
         saved_result=mock_train_result, experiment_run_id=123, metadata=metadata
     )
 
-    # Verify shutil.rmtree was NOT called
+    # Verify cleanup_run_directory was NOT called (no run_dir)
-    mock_shutil.rmtree.assert_not_called()
+    # Only delete_file_from_minio + update_experiment_run
-
-    # Verify activities were still called
     assert workflow_mock.execute_activity_method.call_count == 2