Add support for AAD auth in data plane ops (#35186)

* add support for aad auth in data plane ops

* resolve mypy issues

* update changelog

Author: PratibhaShrivastav18

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -8,6 +8,7 @@
 - Added Project entity class and YAML support.
 - Project and Hub operations supported by workspace operations.
 - workspace list operation supports type filtering.
+- Add support for Microsoft Entra token (`aad_token`) auth in `invoke` and `get-credentials` operations.
 
 ### Bugs Fixed
 

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -56,6 +56,11 @@
 AML_TOKEN_YAML = "aml_token"
 AAD_TOKEN_YAML = "aad_token"
 KEY = "key"
+AAD_TOKEN = "aadtoken"
+AAD_TOKEN_RESOURCE_ENDPOINT = "https://ml.azure.com"
+EMPTY_CREDENTIALS_ERROR = (
+    "Credentials unavailable. Initialize credentials using 'MLClient' for SDK or 'az login' for CLI."
+)
 DEFAULT_ARM_RETRY_INTERVAL = 60
 COMPONENT_TYPE = "type"
 TID_FMT = "&tid={}"

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_endpoint/online_endpoint.py

@@ -29,6 +29,7 @@
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 from azure.ai.ml.entities._util import is_compute_in_override, load_from_dict
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
+from azure.core.credentials import AccessToken
 
 from ._endpoint_helpers import validate_endpoint_or_deployment_name, validate_identity_type_defined
 from .endpoint import Endpoint
@@ -625,3 +626,22 @@ def _to_rest_object(self) -> RestEndpointAuthToken:
             refresh_after_time_utc=self.refresh_after_time_utc,
             token_type=self.token_type,
         )
+
+
+class EndpointAadToken:
+    """Endpoint aad token.
+
+    :ivar access_token: Access token for aad authentication.
+    :vartype access_token: str
+    :ivar expiry_time_utc: Access token expiry time (UTC).
+    :vartype expiry_time_utc: float
+    """
+
+    def __init__(self, obj: AccessToken):
+        """Constructor for Endpoint aad token.
+
+        :param obj: Access token object
+        :type obj: AccessToken
+        """
+        self.access_token = obj.token
+        self.expiry_time_utc = obj.expires_on

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_online_endpoint_operations.py

@@ -9,6 +9,7 @@
 
 from marshmallow.exceptions import ValidationError as SchemaValidationError
 
+from azure.ai.ml._azure_environments import _resource_to_scopes
 from azure.ai.ml._exception_helper import log_and_raise_error
 from azure.ai.ml._restclient.v2022_02_01_preview import AzureMachineLearningWorkspaces as ServiceClient022022Preview
 from azure.ai.ml._restclient.v2022_02_01_preview.models import KeyType, RegenerateEndpointKeysRequest
@@ -23,11 +24,18 @@
 from azure.ai.ml._utils._endpoint_utils import validate_response
 from azure.ai.ml._utils._http_utils import HttpPipeline
 from azure.ai.ml._utils._logger_utils import OpsLogger
-from azure.ai.ml.constants._common import KEY, AzureMLResourceType, LROConfigurations
+from azure.ai.ml.constants._common import (
+    AAD_TOKEN,
+    AAD_TOKEN_RESOURCE_ENDPOINT,
+    EMPTY_CREDENTIALS_ERROR,
+    KEY,
+    AzureMLResourceType,
+    LROConfigurations,
+)
 from azure.ai.ml.constants._endpoint import EndpointInvokeFields, EndpointKeyType
 from azure.ai.ml.entities import OnlineDeployment, OnlineEndpoint
 from azure.ai.ml.entities._assets import Data
-from azure.ai.ml.entities._endpoint.online_endpoint import EndpointAuthKeys, EndpointAuthToken
+from azure.ai.ml.entities._endpoint.online_endpoint import EndpointAadToken, EndpointAuthKeys, EndpointAuthToken
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 from azure.ai.ml.operations._local_endpoint_helper import _LocalEndpointHelper
 from azure.core.credentials import TokenCredential
@@ -96,7 +104,7 @@ def list(self, *, local: bool = False) -> ItemPaged[OnlineEndpoint]:
 
     @distributed_trace
     @monitor_with_activity(ops_logger, "OnlineEndpoint.ListKeys", ActivityType.PUBLICAPI)
-    def get_keys(self, name: str) -> Union[EndpointAuthKeys, EndpointAuthToken]:
+    def get_keys(self, name: str) -> Union[EndpointAuthKeys, EndpointAuthToken, EndpointAadToken]:
         """Get the auth credentials.
 
         :param name: The endpoint name
@@ -344,7 +352,7 @@ def invoke(
         keys = self._get_online_credentials(name=endpoint_name, auth_mode=endpoint.properties.auth_mode)
         if isinstance(keys, EndpointAuthKeys):
             key = keys.primary_key
-        elif isinstance(keys, EndpointAuthToken):
+        elif isinstance(keys, (EndpointAuthToken, EndpointAadToken)):
             key = keys.access_token
         else:
             key = ""
@@ -365,7 +373,7 @@ def _get_workspace_location(self) -> str:
 
     def _get_online_credentials(
         self, name: str, auth_mode: Optional[str] = None
-    ) -> Union[EndpointAuthKeys, EndpointAuthToken]:
+    ) -> Union[EndpointAuthKeys, EndpointAuthToken, EndpointAadToken]:
         if not auth_mode:
             endpoint = self._online_operation.get(
                 resource_group_name=self._resource_group_name,
@@ -384,6 +392,11 @@ def _get_online_credentials(
                 **self._init_kwargs,
             )
 
+        if auth_mode is not None and auth_mode.lower() == AAD_TOKEN:
+            if self._credentials:
+                return EndpointAadToken(self._credentials.get_token(*_resource_to_scopes(AAD_TOKEN_RESOURCE_ENDPOINT)))
+            raise Exception(EMPTY_CREDENTIALS_ERROR)
+
         return self._online_operation.get_token(
             resource_group_name=self._resource_group_name,
             workspace_name=self._workspace_name,

sdk/ml/azure-ai-ml/cspell.json

@@ -2,6 +2,7 @@
     "version": "0.2",
     "ignoreWords": [
         "kwoa",
-        "rslex"
+        "rslex",
+        "aadtoken"
     ]
-}
+}

sdk/ml/azure-ai-ml/tests/online_services/unittests/test_online_endpoints.py

@@ -3,20 +3,27 @@
 
 import pytest
 from pytest_mock import MockFixture
-from azure.ai.ml import load_online_endpoint, load_online_deployment
-from azure.ai.ml._restclient.v2022_10_01.models import EndpointAuthKeys
+
+from azure.ai.ml import load_online_deployment, load_online_endpoint
+from azure.ai.ml._azure_environments import _resource_to_scopes
 from azure.ai.ml._restclient.v2022_02_01_preview.models import (
     KubernetesOnlineDeployment as RestKubernetesOnlineDeployment,
 )
-from azure.ai.ml.entities._util import load_from_dict
 from azure.ai.ml._restclient.v2022_02_01_preview.models import (
     OnlineDeploymentData,
     OnlineDeploymentDetails,
     OnlineEndpointData,
 )
 from azure.ai.ml._restclient.v2022_02_01_preview.models import OnlineEndpointDetails as RestOnlineEndpoint
+from azure.ai.ml._restclient.v2022_10_01.models import EndpointAuthKeys
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationScope
-from azure.ai.ml.constants._common import AzureMLResourceType, HttpResponseStatusCode
+from azure.ai.ml.constants._common import (
+    AAD_TOKEN_RESOURCE_ENDPOINT,
+    EMPTY_CREDENTIALS_ERROR,
+    AzureMLResourceType,
+    HttpResponseStatusCode,
+)
+from azure.ai.ml.entities._util import load_from_dict
 from azure.ai.ml.operations import (
     DatastoreOperations,
     EnvironmentOperations,
@@ -284,6 +291,39 @@ def test_online_get_token(
         mock_online_endpoint_operations._online_operation.get.assert_called_once()
         mock_online_endpoint_operations._online_operation.get_token.assert_called_once()
 
+    def test_online_aad_get_token(
+        self,
+        mock_online_endpoint_operations: OnlineEndpointOperations,
+        mock_aml_services_2022_02_01_preview: Mock,
+    ) -> None:
+        random_name = "random_name"
+        mock_aml_services_2022_02_01_preview.online_endpoints.get.return_value = OnlineEndpointData(
+            name=random_name,
+            location="eastus",
+            properties=RestOnlineEndpoint(auth_mode="aadtoken"),
+        )
+        mock_online_endpoint_operations._credentials = Mock(spec_set=DefaultAzureCredential)
+        mock_online_endpoint_operations.get_keys(name=random_name)
+        mock_online_endpoint_operations._online_operation.get.assert_called_once()
+        mock_online_endpoint_operations._credentials.get_token.assert_called_once_with(
+            *_resource_to_scopes(AAD_TOKEN_RESOURCE_ENDPOINT)
+        )
+
+    def test_online_aad_get_token_with_empty_credentials(
+        self,
+        mock_online_endpoint_operations: OnlineEndpointOperations,
+        mock_aml_services_2022_02_01_preview: Mock,
+    ) -> None:
+        random_name = "random_name"
+        mock_aml_services_2022_02_01_preview.online_endpoints.get.return_value = OnlineEndpointData(
+            name=random_name,
+            location="eastus",
+            properties=RestOnlineEndpoint(auth_mode="aadtoken"),
+        )
+        with pytest.raises(Exception) as ex:
+            mock_online_endpoint_operations.get_keys(name=random_name)
+        assert EMPTY_CREDENTIALS_ERROR in str(ex)
+
     def test_online_delete(
         self,
         mock_online_endpoint_operations: OnlineEndpointOperations,