[ml] Enable OBO flow on AmlSpark (Azure#30148)

* wip

* wip

* aml spark obo token

* Get resource from request if not in response

* Fix linter

* use isodate instead of dateutil

Author: bastrik

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_aio/_internal/managed_identity_client.py

@@ -33,7 +33,7 @@ async def request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         request = self._request_factory(resource, self._identity_config)  # pylint: disable=no-member
         request_time = int(time.time())
         response = await self._pipeline.run(request, retry_on_methods=[request.method], **kwargs)
-        token = self._process_response(response, request_time)
+        token = self._process_response(response=response, request_time=request_time, resource=resource)
         return token
 
     def _build_pipeline(self, **kwargs: "Any") -> "AsyncPipeline":

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_credentials/_AzureMLSparkOnBehalfOfCredential.py

@@ -0,0 +1,108 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+import functools
+import os
+from typing import Optional
+
+from azure.core.pipeline.transport import HttpRequest
+
+from .._internal.managed_identity_base import ManagedIdentityBase
+from .._internal.managed_identity_client import ManagedIdentityClient
+
+
+class _AzureMLSparkOnBehalfOfCredential(ManagedIdentityBase):
+    def get_client(self, **kwargs) -> Optional[ManagedIdentityClient]:
+        client_args = _get_client_args(**kwargs)
+        if client_args:
+            return ManagedIdentityClient(**client_args)
+        return None
+
+    def get_unavailable_message(self) -> str:
+        return "AzureML Spark On Behalf of credentials not available in this environment"
+
+
+def _get_client_args(**kwargs) -> Optional[dict]:
+    from pyspark.sql import SparkSession  # cspell:disable-line # pylint: disable=import-error
+
+    try:
+        spark = SparkSession.builder.getOrCreate()
+    except Exception:
+        raise Exception("Fail to get spark session, please check if spark environment is set up.")
+
+    spark_conf = spark.sparkContext.getConf()
+    spark_conf_vars = {
+        "AZUREML_SYNAPSE_CLUSTER_IDENTIFIER": "spark.synapse.clusteridentifier",
+        "AZUREML_SYNAPSE_TOKEN_SERVICE_ENDPOINT": "spark.tokenServiceEndpoint",
+    }
+    for env_key, conf_key in spark_conf_vars.items():
+        value = spark_conf.get(conf_key)
+        if value:
+            os.environ[env_key] = value
+
+    # Override default settings if provided via arguments
+    if len(kwargs) > 0:
+        env_key_from_kwargs = [
+            "AZUREML_SYNAPSE_CLUSTER_IDENTIFIER",
+            "AZUREML_SYNAPSE_TOKEN_SERVICE_ENDPOINT",
+            "AZUREML_RUN_ID",
+            "AZUREML_RUN_TOKEN_EXPIRY",
+        ]
+        for env_key in env_key_from_kwargs:
+            if env_key in kwargs:
+                os.environ[env_key] = kwargs[env_key]
+
+    token_service_endpoint = os.environ.get("AZUREML_SYNAPSE_TOKEN_SERVICE_ENDPOINT")
+    obo_access_token = os.environ.get("AZUREML_OBO_CANARY_TOKEN")
+    subscription_id = os.environ.get("AZUREML_ARM_SUBSCRIPTION")
+    resource_group = os.environ.get("AZUREML_ARM_RESOURCEGROUP")
+    workspace_name = os.environ.get("AZUREML_ARM_WORKSPACE_NAME")
+
+    if not obo_access_token:
+        return None
+
+    # pylint: disable=line-too-long
+    request_url_format = "https://{}/api/v1/proxy/obotoken/v1.0/subscriptions/{}/resourceGroups/{}/providers/Microsoft.MachineLearningServices/workspaces/{}/getuseraccesstokenforrun"  # cspell:disable-line
+    # pylint: enable=line-too-long
+
+    url = request_url_format.format(
+        token_service_endpoint,
+        subscription_id,
+        resource_group,
+        workspace_name,
+    )
+
+    return dict(
+        kwargs,
+        request_factory=functools.partial(_get_request, url),
+    )
+
+
+def _get_request(url, resource) -> HttpRequest:
+    obo_access_token = os.environ.get("AZUREML_OBO_CANARY_TOKEN")
+    experiment_name = os.environ.get("AZUREML_ARM_PROJECT_NAME")
+    run_id = os.environ.get("AZUREML_RUN_ID")
+    oid = os.environ.get("OID")
+    tid = os.environ.get("TID")
+    obo_service_endpoint = os.environ.get("AZUREML_OBO_SERVICE_ENDPOINT")
+    cluster_identifier = os.environ.get("AZUREML_SYNAPSE_CLUSTER_IDENTIFIER")
+
+    request_body = {
+        "oboToken": obo_access_token,
+        "oid": oid,
+        "tid": tid,
+        "resource": resource,
+        "experimentName": experiment_name,
+        "runId": run_id,
+    }
+    headers = {
+        "Content-Type": "application/json;charset=utf-8",
+        "x-ms-proxy-host": obo_service_endpoint,
+        "obo-access-token": obo_access_token,
+        "x-ms-cluster-identifier": cluster_identifier,
+    }
+    request = HttpRequest(method="POST", url=url, headers=headers)
+    request.set_json_body(request_body)
+    return request

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_credentials/aml_on_behalf_of.py

@@ -7,6 +7,7 @@
 
 from azure.core.pipeline.transport import HttpRequest
 
+from ._AzureMLSparkOnBehalfOfCredential import _AzureMLSparkOnBehalfOfCredential
 from .._internal.managed_identity_base import ManagedIdentityBase
 from .._internal.managed_identity_client import ManagedIdentityClient
 
@@ -22,7 +23,11 @@ class AzureMLOnBehalfOfCredential(object):
     # pylint: enable=line-too-long
 
     def __init__(self, **kwargs):
-        self._credential = _AzureMLOnBehalfOfCredential(**kwargs)
+        provider_type = os.environ.get("AZUREML_DATAPREP_TOKEN_PROVIDER")
+        if provider_type == "sparkobo":  # cspell:disable-line
+            self._credential = _AzureMLSparkOnBehalfOfCredential(**kwargs)
+        else:
+            self._credential = _AzureMLOnBehalfOfCredential(**kwargs)
 
     def get_token(self, *scopes, **kwargs):
         """Request an access token for `scopes`.

sdk/ml/azure-ai-ml/azure/ai/ml/identity/_internal/managed_identity_client.py

@@ -6,6 +6,7 @@
 import time
 from typing import TYPE_CHECKING
 
+import isodate
 import six
 from msal import TokenCache
 
@@ -41,8 +42,8 @@ def __init__(self, request_factory, **kwargs):
         self._pipeline = self._build_pipeline(**kwargs)
         self._request_factory = request_factory
 
-    def _process_response(self, response, request_time):
-        # type: (PipelineResponse, int) -> AccessToken
+    def _process_response(self, response, request_time, resource):
+        # type: (PipelineResponse, int, str) -> AccessToken
 
         content = response.context.get(ContentDecodePolicy.CONTEXT_NAME)
         if not content:
@@ -63,9 +64,13 @@ def _process_response(self, response, request_time):
         if not content:
             raise ClientAuthenticationError(message="No token received.", response=response.http_response)
 
-        if "access_token" not in content or not ("expires_in" in content or "expires_on" in content):
+        if not ("access_token" in content or "token" in content) or not (
+            "expires_in" in content or "expires_on" in content or "expiresOn" in content
+        ):
             if content and "access_token" in content:
                 content["access_token"] = "****"
+            if content and "token" in content:
+                content["token"] = "****"
             raise ClientAuthenticationError(
                 message='Unexpected response "{}"'.format(content),
                 response=response.http_response,
@@ -74,14 +79,18 @@ def _process_response(self, response, request_time):
         if self._content_callback:
             self._content_callback(content)
 
-        expires_on = int(content.get("expires_on") or int(content["expires_in"]) + request_time)
+        if "expires_in" in content or "expires_on" in content:
+            expires_on = int(content.get("expires_on") or int(content["expires_in"]) + request_time)
+        else:
+            expires_on = int(isodate.parse_datetime(content["expiresOn"]).timestamp())
         content["expires_on"] = expires_on
 
-        token = AccessToken(content["access_token"], content["expires_on"])
+        access_token = content.get("access_token") or content["token"]
+        token = AccessToken(access_token, content["expires_on"])
 
         # caching is the final step because TokenCache.add mutates its "event"
         self._cache.add(
-            event={"response": content, "scope": [content["resource"]]},
+            event={"response": content, "scope": [content.get("resource") or resource]},
             now=request_time,
         )
 
@@ -124,7 +133,7 @@ def request_token(self, *scopes, **kwargs):
         request = self._request_factory(resource)
         request_time = int(time.time())
         response = self._pipeline.run(request, retry_on_methods=[request.method], **kwargs)
-        token = self._process_response(response, request_time)
+        token = self._process_response(response=response, request_time=request_time, resource=resource)
         return token
 
     def _build_pipeline(self, **kwargs):