diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/evaluation_onedp_client.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/evaluation_onedp_client.py index 39e29a58c1f6..b836aab49a5f 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/evaluation_onedp_client.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/evaluation_onedp_client.py @@ -2,7 +2,9 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # --------------------------------------------------------- +import io import logging +import uuid from typing import Union, Any, Dict from azure.core.credentials import AzureKeyCredential, TokenCredential from azure.ai.evaluation._common.onedp import ProjectsClient as RestEvaluationServiceClient @@ -17,6 +19,7 @@ RedTeamUpload, ) from azure.storage.blob import ContainerClient +from azure.ai.evaluation._exceptions import EvaluationException, ErrorBlame, ErrorCategory, ErrorTarget from .utils import upload LOGGER = logging.getLogger(__name__) @@ -167,3 +170,63 @@ def update_red_team_run(self, *, name: str, red_team: RedTeamUpload, **kwargs): update_run_response = self.rest_client.red_teams.upload_update_run(name=name, redteam=red_team, **kwargs) return update_run_response + + def test_storage_upload(self, **kwargs) -> bool: + """Test storage account connectivity by performing a minimal upload operation. + + This method validates that the storage account is accessible and the credentials + have the necessary permissions before proceeding with the full evaluation workflow. + + :param kwargs: Additional keyword arguments to pass to the underlying API calls + :return: True if the test upload succeeds + :rtype: bool + :raises EvaluationException: If the storage account is inaccessible or lacks permissions + """ + LOGGER.debug("Testing storage account connectivity...") + + try: + # Create a unique test identifier to avoid conflicts + test_id = str(uuid.uuid4())[:8] + test_name = f"connectivity-test-{test_id}" + test_version = "1" + + # Start a pending upload to get storage credentials + start_pending_upload_response = self.rest_client.evaluation_results.start_pending_upload( + name=test_name, + version=test_version, + body=PendingUploadRequest(pending_upload_type=PendingUploadType.TEMPORARY_BLOB_REFERENCE), + **kwargs, + ) + + # Attempt to upload a minimal test blob directly (no temp files needed) + with ContainerClient.from_container_url( + start_pending_upload_response.blob_reference_for_consumption.credential.sas_uri + ) as container_client: + # Use BytesIO for efficient in-memory test data + test_data = io.BytesIO(b"connectivity test") + container_client.upload_blob(data=test_data, name="test.txt", overwrite=True) + container_client.delete_blob(name="test.txt") + + LOGGER.debug("Storage account connectivity test successful") + return True + + except Exception as e: + LOGGER.error(f"Storage account connectivity test failed: {str(e)}") + + # Re-raise with helpful context + error_msg = ( + f"Failed to connect to Azure Blob Storage. Error: {str(e)}. " + "Please verify that:\n" + " 1. The storage account exists and is accessible\n" + " 2. Your credentials have the necessary permissions (Storage Blob Data Contributor role)\n" + " 3. Network access to the storage account is not blocked by firewall rules\n" + " 4. The Azure AI project is properly configured" + ) + + raise EvaluationException( + message=error_msg, + internal_message=f"Storage connectivity test failed: {e}", + target=ErrorTarget.RAI_CLIENT, + category=ErrorCategory.UPLOAD_ERROR, + blame=ErrorBlame.SYSTEM_ERROR, + ) diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/utils.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/utils.py index f07b54113b63..0aa5ad375037 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/utils.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/_common/utils.py @@ -914,9 +914,26 @@ def upload(path: str, container_client: ContainerClient, logger=None): logger.debug(f"File '{local}' uploaded successfully") except Exception as e: + # Extract storage account information if available + storage_info = "" + try: + account_name = container_client.account_name if hasattr(container_client, "account_name") else "unknown" + storage_info = f" Storage account: {account_name}." + except Exception: + pass + + error_msg = ( + f"Failed to upload evaluation results to Azure Blob Storage.{storage_info} " + f"Error: {str(e)}. " + "Please verify that:\n" + " 1. The storage account exists and is accessible\n" + " 2. Your credentials have the necessary permissions (Storage Blob Data Contributor role)\n" + " 3. Network access to the storage account is not blocked by firewall rules" + ) + raise EvaluationException( - message=f"Error uploading file: {e}", - internal_message=f"Error uploading file: {e}", + message=error_msg, + internal_message=f"Error uploading file to blob storage: {e}", target=ErrorTarget.RAI_CLIENT, category=ErrorCategory.UPLOAD_ERROR, blame=ErrorBlame.SYSTEM_ERROR, diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_mlflow_integration.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_mlflow_integration.py index 410975fdfc08..6c80c6ad49b6 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_mlflow_integration.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_mlflow_integration.py @@ -163,6 +163,26 @@ def start_redteam_mlflow_run( return eval_run + def test_storage_upload(self) -> bool: + """Test storage account connectivity before starting the scan. + + This method validates that storage upload will work by testing with the + appropriate client (OneDP or MLFlow) depending on project configuration. + + :return: True if the test upload succeeds + :rtype: bool + :raises EvaluationException: If the storage account is inaccessible or lacks permissions + """ + if self._one_dp_project: + # For OneDP projects, test using the evaluation client + return self.generated_rai_client._evaluation_onedp_client.test_storage_upload() + else: + # For non-OneDP projects (MLFlow), we don't have a direct upload test + # Storage is tested when we create artifacts during log_artifact + # So we just return True here and let the actual upload fail if there are issues + self.logger.debug("Storage upload test skipped for non-OneDP project (will be tested during actual upload)") + return True + async def log_redteam_results_to_mlflow( self, redteam_result: RedTeamResult, diff --git a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py index 1081d4e4ddac..9c46c3809685 100644 --- a/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py +++ b/sdk/evaluation/azure-ai-evaluation/azure/ai/evaluation/red_team/_red_team.py @@ -1347,6 +1347,17 @@ async def scan( # Update result processor with the AI studio URL now that it's available self.result_processor.ai_studio_url = self.mlflow_integration.ai_studio_url + # Test storage connectivity early to catch issues before running attacks + try: + tqdm.write("🔍 Validating storage account connectivity...") + self.mlflow_integration.test_storage_upload() + tqdm.write("✅ Storage account validation successful") + except Exception as e: + # Log the error and re-raise to stop the scan early + self.logger.error(f"Storage account validation failed: {str(e)}") + tqdm.write(f"❌ Storage account validation failed: {str(e)}") + raise + # Process strategies and execute scan flattened_attack_strategies = get_flattened_attack_strategies(attack_strategies) self._validate_strategies(flattened_attack_strategies) diff --git a/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_redteam/test_storage_validation.py b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_redteam/test_storage_validation.py new file mode 100644 index 000000000000..25f5b70f5074 --- /dev/null +++ b/sdk/evaluation/azure-ai-evaluation/tests/unittests/test_redteam/test_storage_validation.py @@ -0,0 +1,153 @@ +# --------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# --------------------------------------------------------- +"""Unit tests for storage account validation functionality.""" + +import os +import tempfile +import unittest +from unittest.mock import Mock, patch +from azure.ai.evaluation._common.evaluation_onedp_client import EvaluationServiceOneDPClient +from azure.ai.evaluation._common.utils import upload +from azure.ai.evaluation._exceptions import EvaluationException +from azure.ai.evaluation.red_team._mlflow_integration import MLflowIntegration + + +class TestStorageValidation(unittest.TestCase): + """Test cases for storage account validation.""" + + def setUp(self): + """Set up test fixtures.""" + self.mock_credential = Mock() + self.test_endpoint = "https://test-endpoint.azure.com" + + @patch("azure.ai.evaluation._common.evaluation_onedp_client.RestEvaluationServiceClient") + def test_test_storage_upload_success(self, mock_rest_client_class): + """Test successful storage upload validation.""" + # Arrange + mock_rest_client = Mock() + mock_rest_client_class.return_value = mock_rest_client + + mock_pending_upload_response = Mock() + mock_pending_upload_response.blob_reference_for_consumption = Mock() + mock_pending_upload_response.blob_reference_for_consumption.credential = Mock() + mock_pending_upload_response.blob_reference_for_consumption.credential.sas_uri = ( + "https://test.blob.core.windows.net/container?sas_token" + ) + + mock_rest_client.evaluation_results.start_pending_upload.return_value = mock_pending_upload_response + + # Mock ContainerClient + with patch( + "azure.ai.evaluation._common.evaluation_onedp_client.ContainerClient" + ) as mock_container_client_class: + mock_container_client = Mock() + mock_container_client.__enter__ = Mock(return_value=mock_container_client) + mock_container_client.__exit__ = Mock(return_value=None) + mock_container_client_class.from_container_url.return_value = mock_container_client + + # Mock the upload_blob method to succeed + mock_container_client.upload_blob = Mock() + # Act + client = EvaluationServiceOneDPClient(self.test_endpoint, self.mock_credential) + result = client.test_storage_upload() + + # Assert + self.assertTrue(result) + mock_container_client.upload_blob.assert_called_once() + + @patch("azure.ai.evaluation._common.evaluation_onedp_client.RestEvaluationServiceClient") + def test_test_storage_upload_failure(self, mock_rest_client_class): + """Test storage upload validation failure.""" + # Arrange + mock_rest_client = Mock() + mock_rest_client_class.return_value = mock_rest_client + + # Simulate upload failure + mock_rest_client.evaluation_results.start_pending_upload.side_effect = Exception( + "Storage account not accessible" + ) + + # Act & Assert + client = EvaluationServiceOneDPClient(self.test_endpoint, self.mock_credential) + with self.assertRaises(EvaluationException) as context: + client.test_storage_upload() + + # Verify error message contains helpful guidance + self.assertIn("storage", str(context.exception).lower()) + self.assertIn("permissions", str(context.exception).lower()) + + def test_mlflow_integration_test_storage_upload_onedp(self): + """Test MLflowIntegration storage validation for OneDP projects.""" + # Arrange + mock_logger = Mock() + mock_generated_rai_client = Mock() + mock_evaluation_client = Mock() + mock_evaluation_client.test_storage_upload.return_value = True + mock_generated_rai_client._evaluation_onedp_client = mock_evaluation_client + + mlflow_integration = MLflowIntegration( + logger=mock_logger, generated_rai_client=mock_generated_rai_client, one_dp_project=True + ) + + # Act + result = mlflow_integration.test_storage_upload() + + # Assert + self.assertTrue(result) + mock_evaluation_client.test_storage_upload.assert_called_once() + + def test_mlflow_integration_test_storage_upload_non_onedp(self): + """Test MLflowIntegration storage validation for non-OneDP projects.""" + # Arrange + mock_logger = Mock() + mock_generated_rai_client = Mock() + + mlflow_integration = MLflowIntegration( + logger=mock_logger, generated_rai_client=mock_generated_rai_client, one_dp_project=False + ) + + # Act + result = mlflow_integration.test_storage_upload() + + # Assert + # For non-OneDP projects, we skip the test and return True + self.assertTrue(result) + + +class TestUploadErrorMessages(unittest.TestCase): + """Test cases for improved upload error messages.""" + + @patch("azure.ai.evaluation._common.utils.ContainerClient") + def test_upload_error_message_includes_guidance(self, mock_container_client_class): + """Test that upload errors include helpful troubleshooting guidance.""" + # Arrange + mock_container_client = Mock() + mock_container_client.account_name = "testaccount" + mock_container_client.upload_blob.side_effect = Exception("Connection refused") + + # Create a temporary test file + with tempfile.NamedTemporaryFile(mode="w", delete=False, suffix=".txt") as tmp_file: + tmp_file.write("test content") + test_file_path = tmp_file.name + + # Act & Assert + try: + with self.assertRaises(EvaluationException) as context: + upload(path=test_file_path, container_client=mock_container_client) + + # Verify error message contains helpful guidance + error_message = str(context.exception) + self.assertIn("storage account", error_message.lower()) + self.assertIn("permissions", error_message.lower()) + self.assertIn("verify", error_message.lower()) + finally: + # Clean up the temporary file + try: + os.unlink(test_file_path) + except Exception: + pass + + +if __name__ == "__main__": + unittest.main()