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.

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']