two version of R2R are hereHEADmaster

18 files changed, 949 insertions, 0 deletions

diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/__init__.py
new file mode 100644
index 00000000..72a674b6
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/__init__.py
@@ -0,0 +1,12 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Contains Identity Configuration for Azure Machine Learning SDKv2."""
+
+from ._credentials import AzureMLOnBehalfOfCredential
+from ._exceptions import CredentialUnavailableError
+
+__all__ = [
+    "AzureMLOnBehalfOfCredential",
+    "CredentialUnavailableError",
+]
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/__init__.py
new file mode 100644
index 00000000..d540fd20
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/__init__.py
@@ -0,0 +1,3 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_credentials/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_credentials/__init__.py
new file mode 100644
index 00000000..d540fd20
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_credentials/__init__.py
@@ -0,0 +1,3 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_credentials/aml_on_behalf_of.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_credentials/aml_on_behalf_of.py
new file mode 100644
index 00000000..386993d6
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_credentials/aml_on_behalf_of.py
@@ -0,0 +1,86 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import functools
+import os
+from typing import Any, Optional
+
+from azure.core.credentials import AccessToken
+from azure.core.pipeline.transport import HttpRequest
+
+from .._internal import AsyncContextManager
+from .._internal.managed_identity_base import AsyncManagedIdentityBase
+from .._internal.managed_identity_client import AsyncManagedIdentityClient
+
+
+class AzureMLOnBehalfOfCredential(AsyncContextManager):
+    # pylint: disable=line-too-long
+    """Authenticates a user via the on-behalf-of flow.
+
+    This credential can only be used on `Azure Machine Learning Compute.
+    <https://learn.microsoft.com/azure/machine-learning/concept-compute-target#azure-machine-learning-compute-managed>`_ during job execution when user request to
+    run job during its identity.
+    """
+    # pylint: enable=line-too-long
+
+    def __init__(self, **kwargs: Any):
+        self._credential = _AzureMLOnBehalfOfCredential(**kwargs)
+
+    async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        """Request an access token for `scopes`.
+
+        This method is called automatically by Azure SDK clients.
+
+        :param scopes: The desired scope for the access token. This credential allows only one scope per request.
+        :type scopes: str
+        :rtype: ~azure.core.credentials.AccessToken
+        :return: The access token for temporary access token for the requested scope.
+        :raises: ~azure.ai.ml.identity.CredentialUnavailableError
+        """
+
+        return await self._credential.get_token(*scopes, **kwargs)
+
+    async def __aenter__(self) -> "AzureMLOnBehalfOfCredential":
+        if self._credential:
+            await self._credential.__aenter__()
+        return self
+
+    async def close(self) -> None:
+        """Close the credential's transport session."""
+        if self._credential:
+            await self._credential.__aexit__()
+
+
+class _AzureMLOnBehalfOfCredential(AsyncManagedIdentityBase):
+    def get_client(self, **kwargs):
+        # type: (**Any) -> Optional[AsyncManagedIdentityClient]
+        client_args = _get_client_args(**kwargs)
+        if client_args:
+            return AsyncManagedIdentityClient(**client_args)
+        return None
+
+    def get_unavailable_message(self):
+        # type: () -> str
+        return "AzureML On Behalf of credentials not available in this environment"
+
+
+def _get_client_args(**kwargs):
+    # type: (dict) -> Optional[dict]
+
+    url = os.environ.get("OBO_ENDPOINT")
+    if not url:
+        # OBO identity isn't available in this environment
+        return None
+
+    return dict(
+        kwargs,
+        request_factory=functools.partial(_get_request, url),
+    )
+
+
+def _get_request(url, resource):
+    # type: (str, str) -> HttpRequest
+    request = HttpRequest("GET", url)
+    request.format_parameters(dict({"resource": resource}))
+    return request
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/__init__.py
new file mode 100644
index 00000000..1059f88c
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/__init__.py
@@ -0,0 +1,20 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+import abc
+from typing import Any
+
+
+class AsyncContextManager(abc.ABC):
+    @abc.abstractmethod
+    async def close(self) -> None:
+        pass
+
+    async def __aenter__(self) -> "AsyncContextManager":
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        await self.close()
+
+
+__all__ = ["AsyncContextManager"]
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/get_token_mixin.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/get_token_mixin.py
new file mode 100644
index 00000000..dc2dcef2
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/get_token_mixin.py
@@ -0,0 +1,94 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import abc
+import logging
+import time
+from typing import TYPE_CHECKING, Any, Optional
+
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException
+
+from ..._constants import DEFAULT_REFRESH_OFFSET, DEFAULT_TOKEN_REFRESH_RETRY_DELAY
+
+if TYPE_CHECKING:
+    from azure.core.credentials import AccessToken
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class GetTokenMixin(abc.ABC):
+    def __init__(self, *args: "Any", **kwargs: "Any") -> None:
+        self._last_request_time = 0
+
+        # https://github.com/python/mypy/issues/5887
+        super(GetTokenMixin, self).__init__(*args, **kwargs)
+
+    @abc.abstractmethod
+    # pylint: disable-next=docstring-missing-param
+    async def _acquire_token_silently(self, *scopes: str, **kwargs: "Any") -> "Optional[AccessToken]":
+        """Attempt to acquire an access token from a cache or by redeeming a refresh token."""
+
+    @abc.abstractmethod
+    # pylint: disable-next=docstring-missing-param
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        """Request an access token from the STS."""
+
+    def _should_refresh(self, token: "AccessToken") -> bool:
+        now = int(time.time())
+        if token.expires_on - now > DEFAULT_REFRESH_OFFSET:
+            return False
+        if now - self._last_request_time < DEFAULT_TOKEN_REFRESH_RETRY_DELAY:
+            return False
+        return True
+
+    async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        """Request an access token for `scopes`.
+
+        This method is called automatically by Azure SDK clients.
+
+        :param scopes: Scopes to request access for
+        :type: str
+        :raises CredentialUnavailableError: the credential is unable to attempt authentication because it lacks
+            required data, state, or platform support
+        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
+            attribute gives a reason.
+        :return: The access token
+        :rtype: ~azure.core.credentials.AccessToken
+        """
+        if not scopes:
+            msg = '"get_token" requires at least one scope'
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.IDENTITY,
+                error_category=ErrorCategory.USER_ERROR,
+            )
+
+        try:
+            token = await self._acquire_token_silently(*scopes, **kwargs)
+            if not token:
+                self._last_request_time = int(time.time())
+                token = await self._request_token(*scopes, **kwargs)
+            elif self._should_refresh(token):
+                try:
+                    self._last_request_time = int(time.time())
+                    token = await self._request_token(*scopes, **kwargs)
+                except Exception:  # pylint:disable=broad-except
+                    pass
+            _LOGGER.log(
+                logging.INFO,
+                "%s.get_token succeeded",
+                self.__class__.__name__,
+            )
+            return token
+
+        except Exception as ex:
+            _LOGGER.log(
+                logging.WARNING,
+                "%s.get_token failed: %s",
+                self.__class__.__name__,
+                ex,
+                exc_info=_LOGGER.isEnabledFor(logging.DEBUG),
+            )
+            raise
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/managed_identity_base.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/managed_identity_base.py
new file mode 100644
index 00000000..4f3cc167
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/managed_identity_base.py
@@ -0,0 +1,55 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import abc
+from typing import TYPE_CHECKING, Any, Optional, cast
+
+from ..._exceptions import CredentialUnavailableError
+from . import AsyncContextManager
+from .get_token_mixin import GetTokenMixin
+from .managed_identity_client import AsyncManagedIdentityClient
+
+if TYPE_CHECKING:
+    from azure.core.credentials import AccessToken
+
+
+class AsyncManagedIdentityBase(AsyncContextManager, GetTokenMixin):
+    """Base class for internal credentials using AsyncManagedIdentityClient."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__()
+        self._client = self.get_client(**kwargs)
+
+    @abc.abstractmethod
+    def get_client(self, **kwargs: Any) -> "Optional[AsyncManagedIdentityClient]":
+        pass
+
+    @abc.abstractmethod
+    def get_unavailable_message(self) -> str:
+        pass
+
+    async def __aenter__(self) -> "AsyncManagedIdentityBase":
+        if self._client:
+            await self._client.__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any) -> None:
+        if self._client:
+            await self._client.__aexit__(*args)
+
+    async def close(self) -> None:
+        await self.__aexit__()
+
+    async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        if not self._client:
+            raise CredentialUnavailableError(message=self.get_unavailable_message())
+        return await super().get_token(*scopes, **kwargs)
+
+    async def _acquire_token_silently(self, *scopes: str, **kwargs: "Any") -> "Optional[AccessToken]":
+        # casting because mypy can't determine that these methods are called
+        # only by get_token, which raises when self._client is None
+        return cast(AsyncManagedIdentityClient, self._client).get_cached_token(*scopes)
+
+    async def _request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        return await cast(AsyncManagedIdentityClient, self._client).request_token(*scopes, **kwargs)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/managed_identity_client.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/managed_identity_client.py
new file mode 100644
index 00000000..01d951e9
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_aio/_internal/managed_identity_client.py
@@ -0,0 +1,38 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import time
+from typing import TYPE_CHECKING, Any
+
+from ..._internal import _scopes_to_resource
+from ..._internal.managed_identity_client import ManagedIdentityClientBase
+from ..._internal.pipeline import build_async_pipeline
+from .._internal import AsyncContextManager
+
+if TYPE_CHECKING:
+    from azure.core.credentials import AccessToken
+    from azure.core.pipeline import AsyncPipeline
+
+
+# pylint:disable=async-client-bad-name
+class AsyncManagedIdentityClient(AsyncContextManager, ManagedIdentityClientBase):
+    async def __aenter__(self) -> "AsyncManagedIdentityClient":
+        await self._pipeline.__aenter__()
+        return self
+
+    async def close(self) -> None:
+        await self._pipeline.__aexit__()
+
+    async def request_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
+        # pylint:disable=invalid-overridden-method
+        resource = _scopes_to_resource(*scopes)
+        # pylint: disable=no-member
+        request = self._request_factory(resource, self._identity_config)  # type: ignore
+        request_time = int(time.time())
+        response = await self._pipeline.run(request, retry_on_methods=[request.method], **kwargs)
+        token = self._process_response(response=response, request_time=request_time, resource=resource)
+        return token
+
+    def _build_pipeline(self, **kwargs: "Any") -> "AsyncPipeline":
+        return build_async_pipeline(**kwargs)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_constants.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_constants.py
new file mode 100644
index 00000000..8ca26c01
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_constants.py
@@ -0,0 +1,7 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+DEFAULT_REFRESH_OFFSET = 300
+DEFAULT_TOKEN_REFRESH_RETRY_DELAY = 30
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/_AzureMLSparkOnBehalfOfCredential.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/_AzureMLSparkOnBehalfOfCredential.py
new file mode 100644
index 00000000..34528e5d
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/_AzureMLSparkOnBehalfOfCredential.py
@@ -0,0 +1,111 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+import functools
+import os
+from typing import Any, Optional
+
+from azure.ai.ml.exceptions import MlException
+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: Any) -> 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: Any) -> Optional[dict]:
+    # 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]
+            else:
+                msg = "Unable to initialize AzureMLHoboSparkOBOCredential due to invalid arguments"
+                raise MlException(message=msg, no_personal_data_message=msg)
+    else:
+        from pyspark.sql import SparkSession  # cspell:disable-line # pylint: disable=import-error
+
+        try:
+            spark = SparkSession.builder.getOrCreate()
+        except Exception as e:
+            msg = "Fail to get spark session, please check if spark environment is set up."
+            raise MlException(message=msg, no_personal_data_message=msg) from e
+
+        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
+
+    token_service_endpoint = os.environ.get("AZUREML_SYNAPSE_TOKEN_SERVICE_ENDPOINT")
+    obo_access_token = os.environ.get("AZUREML_OBO_CANARY_TOKEN")
+    obo_endpoint = os.environ.get("AZUREML_OBO_USER_TOKEN_FOR_SPARK_RETRIEVAL_API", "getuseraccesstokenforspark")
+    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/{}/{}"  # cspell:disable-line
+    # pylint: enable=line-too-long
+
+    url = request_url_format.format(
+        token_service_endpoint, subscription_id, resource_group, workspace_name, obo_endpoint
+    )
+
+    return dict(
+        kwargs,
+        request_factory=functools.partial(_get_request, url),
+    )
+
+
+def _get_request(url: str, resource: Any) -> 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
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/__init__.py
new file mode 100644
index 00000000..fd46c46e
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/__init__.py
@@ -0,0 +1,10 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+from .aml_on_behalf_of import AzureMLOnBehalfOfCredential
+
+__all__ = [
+    "AzureMLOnBehalfOfCredential",
+]
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/aml_on_behalf_of.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/aml_on_behalf_of.py
new file mode 100644
index 00000000..ba938512
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_credentials/aml_on_behalf_of.py
@@ -0,0 +1,94 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import functools
+import os
+from typing import Any, Optional, Union
+
+from azure.core.credentials import AccessToken
+from azure.core.pipeline.transport import HttpRequest
+
+from .._internal.managed_identity_base import ManagedIdentityBase
+from .._internal.managed_identity_client import ManagedIdentityClient
+from ._AzureMLSparkOnBehalfOfCredential import _AzureMLSparkOnBehalfOfCredential
+
+
+class AzureMLOnBehalfOfCredential(object):
+    # pylint: disable=line-too-long
+    """Authenticates a user via the on-behalf-of flow.
+
+    This credential can only be used on `Azure Machine Learning Compute
+    <https://learn.microsoft.com/azure/machine-learning/concept-compute-target#azure-machine-learning-compute-managed>`_ or `Azure Machine Learning Serverless Spark Compute
+    <https://learn.microsoft.com/azure/machine-learning/apache-spark-azure-ml-concepts#serverless-spark-compute>`_
+    during job execution when user request to run job using its identity.
+    """
+    # pylint: enable=line-too-long
+
+    def __init__(self, **kwargs: Any):
+        provider_type = os.environ.get("AZUREML_DATAPREP_TOKEN_PROVIDER")
+        self._credential: Union[_AzureMLSparkOnBehalfOfCredential, _AzureMLOnBehalfOfCredential]
+
+        if provider_type == "sparkobo":  # cspell:disable-line
+            self._credential = _AzureMLSparkOnBehalfOfCredential(**kwargs)
+        else:
+            self._credential = _AzureMLOnBehalfOfCredential(**kwargs)
+
+    def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        """Request an access token for `scopes`.
+
+        This method is called automatically by Azure SDK clients.
+
+        :param str scopes: desired scope for the access token. This credential allows only one scope per request.
+        :rtype: ~azure.core.credentials.AccessToken
+        :return: AzureML On behalf of credentials isn't available in the hosting environment
+        :raises: ~azure.ai.ml.identity.CredentialUnavailableError
+        """
+
+        return self._credential.get_token(*scopes, **kwargs)
+
+    def __enter__(self) -> "AzureMLOnBehalfOfCredential":
+        self._credential.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self._credential.__exit__(*args)
+
+    def close(self):
+        # type: () -> None
+        """Close the credential's transport session."""
+        self.__exit__()
+
+
+class _AzureMLOnBehalfOfCredential(ManagedIdentityBase):
+    def get_client(self, **kwargs):
+        # type: (**Any) -> Optional[ManagedIdentityClient]
+        client_args = _get_client_args(**kwargs)
+        if client_args:
+            return ManagedIdentityClient(**client_args)
+        return None
+
+    def get_unavailable_message(self):
+        # type: () -> str
+        return "AzureML On Behalf of credentials not available in this environment"
+
+
+def _get_client_args(**kwargs):
+    # type: (dict) -> Optional[dict]
+
+    url = os.environ.get("OBO_ENDPOINT")
+    if not url:
+        # OBO identity isn't available in this environment
+        return None
+
+    return dict(
+        kwargs,
+        request_factory=functools.partial(_get_request, url),
+    )
+
+
+def _get_request(url, resource):
+    # type: (str, str) -> HttpRequest
+    request = HttpRequest("GET", url)
+    request.format_parameters(dict({"resource": resource}))
+    return request
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_exceptions.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_exceptions.py
new file mode 100644
index 00000000..40ea9e46
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_exceptions.py
@@ -0,0 +1,9 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from azure.core.exceptions import ClientAuthenticationError
+
+
+class CredentialUnavailableError(ClientAuthenticationError):
+    """The credential did not attempt to authenticate because required data or state is unavailable."""
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/__init__.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/__init__.py
new file mode 100644
index 00000000..6d997fb8
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/__init__.py
@@ -0,0 +1,32 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+
+# ---------------------------------------------------------------------------------------------
+# This package has been vendored from azure-identity package from the following commit
+# https://github.com/Azure/azure-sdk-for-python/commit/0f302dc6c299df2ee637457c8f165c7bdb4ec2af
+# ---------------------------------------------------------------------------------------------
+from typing import Any
+
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException
+
+
+# pylint: disable-next=docstring-missing-param,docstring-missing-return,docstring-missing-rtype
+def _scopes_to_resource(*scopes: Any) -> Any:
+    """Convert an AADv2 scope to an AADv1 resource."""
+
+    if len(scopes) != 1:
+        msg = "This credential requires exactly one scope per token request."
+        raise ValidationException(
+            message=msg,
+            no_personal_data_message=msg,
+            target=ErrorTarget.IDENTITY,
+            error_category=ErrorCategory.USER_ERROR,
+        )
+
+    resource = scopes[0]
+    if resource.endswith("/.default"):
+        resource = resource[: -len("/.default")]
+
+    return resource
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/get_token_mixin.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/get_token_mixin.py
new file mode 100644
index 00000000..45e772e8
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/get_token_mixin.py
@@ -0,0 +1,98 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import abc
+import logging
+import time
+from typing import Any, Optional
+
+from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationException
+
+from .._constants import DEFAULT_REFRESH_OFFSET, DEFAULT_TOKEN_REFRESH_RETRY_DELAY
+
+try:
+    ABC = abc.ABC
+except AttributeError:  # Python 2.7, abc exists, but not ABC
+    ABC = abc.ABCMeta("ABC", (object,), {"__slots__": ()})  # type: ignore
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class GetTokenMixin(ABC):
+    from azure.core.credentials import AccessToken
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        self._last_request_time = 0
+
+        # https://github.com/python/mypy/issues/5887
+        super(GetTokenMixin, self).__init__(*args, **kwargs)
+
+    @abc.abstractmethod
+    # pylint: disable-next=docstring-missing-param
+    def _acquire_token_silently(self, *scopes: str, **kwargs: Any) -> Optional[AccessToken]:
+        """Attempt to acquire an access token from a cache or by redeeming a refresh token."""
+
+    @abc.abstractmethod
+    # pylint: disable-next=docstring-missing-param
+    def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        """Request an access token from the STS."""
+
+    def _should_refresh(self, token: AccessToken) -> bool:
+        now = int(time.time())
+        if token.expires_on - now > DEFAULT_REFRESH_OFFSET:
+            return False
+        if now - self._last_request_time < DEFAULT_TOKEN_REFRESH_RETRY_DELAY:
+            return False
+        return True
+
+    def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        """Request an access token for `scopes`.
+
+        This method is called automatically by Azure SDK clients.
+
+        :param scopes: The desired scopes for the access token. This method requires at least one scope.
+        :type scopes: str
+        :return: The access token
+        :rtype: ~azure.core.credentials.AccessToken
+        :raises CredentialUnavailableError: the credential is unable to attempt authentication because it lacks
+            required data, state, or platform support
+        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
+            attribute gives a reason.
+        """
+        if not scopes:
+            msg = '"get_token" requires at least one scope'
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.IDENTITY,
+                error_category=ErrorCategory.USER_ERROR,
+            )
+
+        try:
+            token = self._acquire_token_silently(*scopes, **kwargs)
+            if not token:
+                self._last_request_time = int(time.time())
+                token = self._request_token(*scopes, **kwargs)
+            elif self._should_refresh(token):
+                try:
+                    self._last_request_time = int(time.time())
+                    token = self._request_token(*scopes, **kwargs)
+                except Exception:  # pylint:disable=broad-except
+                    pass
+            _LOGGER.log(
+                logging.INFO,
+                "%s.get_token succeeded",
+                self.__class__.__name__,
+            )
+            return token
+
+        except Exception as ex:
+            _LOGGER.log(
+                logging.WARNING,
+                "%s.get_token failed: %s",
+                self.__class__.__name__,
+                ex,
+                exc_info=_LOGGER.isEnabledFor(logging.DEBUG),
+            )
+            raise
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/managed_identity_base.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/managed_identity_base.py
new file mode 100644
index 00000000..cc235ba2
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/managed_identity_base.py
@@ -0,0 +1,53 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import abc
+from typing import Any, Optional, cast
+
+from .._exceptions import CredentialUnavailableError
+from .._internal.get_token_mixin import GetTokenMixin
+from .._internal.managed_identity_client import ManagedIdentityClient
+
+
+class ManagedIdentityBase(GetTokenMixin):
+    """Base class for internal credentials using ManagedIdentityClient."""
+
+    from azure.core.credentials import AccessToken
+
+    def __init__(self, **kwargs: Any):
+        super(ManagedIdentityBase, self).__init__()
+        self._client = self.get_client(**kwargs)
+
+    @abc.abstractmethod
+    def get_client(self, **kwargs: Any) -> Optional[ManagedIdentityClient]:
+        pass
+
+    @abc.abstractmethod
+    def get_unavailable_message(self) -> str:
+        pass
+
+    def __enter__(self) -> "ManagedIdentityBase":
+        if self._client:
+            self._client.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        if self._client:
+            self._client.__exit__(*args)
+
+    def close(self) -> None:
+        self.__exit__()
+
+    def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        if not self._client:
+            raise CredentialUnavailableError(message=self.get_unavailable_message())
+        return super(ManagedIdentityBase, self).get_token(*scopes, **kwargs)
+
+    def _acquire_token_silently(self, *scopes: str, **kwargs: Any) -> Optional[AccessToken]:
+        # casting because mypy can't determine that these methods are called
+        # only by get_token, which raises when self._client is None
+        return cast(ManagedIdentityClient, self._client).get_cached_token(*scopes)
+
+    def _request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        return cast(ManagedIdentityClient, self._client).request_token(*scopes, **kwargs)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/managed_identity_client.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/managed_identity_client.py
new file mode 100644
index 00000000..2290e22e
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/managed_identity_client.py
@@ -0,0 +1,127 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+import abc
+import time
+from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+
+import isodate
+from msal import TokenCache
+
+from azure.core.credentials import AccessToken
+from azure.core.exceptions import ClientAuthenticationError, DecodeError
+from azure.core.pipeline import Pipeline
+from azure.core.pipeline.policies import ContentDecodePolicy
+
+from .._internal import _scopes_to_resource
+from .._internal.pipeline import build_pipeline
+
+ABC = abc.ABC
+
+if TYPE_CHECKING:
+
+    from azure.core.pipeline.policies import HTTPPolicy, SansIOHTTPPolicy
+    from azure.core.pipeline.transport import HttpRequest
+
+    PolicyType = Union[HTTPPolicy, SansIOHTTPPolicy]
+
+
+class ManagedIdentityClientBase(ABC):
+    from azure.core.pipeline import PipelineResponse
+
+    def __init__(self, request_factory: Callable, **kwargs: Any) -> None:
+        self._cache = kwargs.pop("_cache", None) or TokenCache()
+        self._content_callback = kwargs.pop("_content_callback", None)
+        self._pipeline = self._build_pipeline(**kwargs)
+        self._request_factory = request_factory
+
+    def _process_response(self, response: PipelineResponse, request_time: int, resource: str) -> AccessToken:
+
+        content = response.context.get(ContentDecodePolicy.CONTEXT_NAME)
+        if not content:
+            try:
+                content = ContentDecodePolicy.deserialize_from_text(
+                    response.http_response.text(), mime_type="application/json"
+                )
+            except DecodeError as ex:
+                if response.http_response.content_type.startswith("application/json"):
+                    message = "Failed to deserialize JSON from response"
+                else:
+                    message = 'Unexpected content type "{}"'.format(response.http_response.content_type)
+                raise ClientAuthenticationError(message=message, response=response.http_response) from ex
+
+        if not content:
+            raise ClientAuthenticationError(message="No token received.", response=response.http_response)
+
+        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,
+            )
+
+        if self._content_callback:
+            self._content_callback(content)
+
+        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
+
+        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.get("resource") or resource]},
+            now=request_time,
+        )
+
+        return token
+
+    def get_cached_token(self, *scopes: str) -> Optional[AccessToken]:
+        resource = _scopes_to_resource(*scopes)
+        tokens = self._cache.find(TokenCache.CredentialType.ACCESS_TOKEN, target=[resource])
+        for token in tokens:
+            expires_on = int(token["expires_on"])
+            if expires_on > time.time():
+                return AccessToken(token["secret"], expires_on)
+        return None
+
+    @abc.abstractmethod
+    def request_token(self, *scopes: Any, **kwargs: Any) -> AccessToken:
+        pass
+
+    @abc.abstractmethod
+    def _build_pipeline(self, **kwargs: Any) -> Pipeline:
+        pass
+
+
+class ManagedIdentityClient(ManagedIdentityClientBase):
+    def __enter__(self) -> "ManagedIdentityClient":
+        self._pipeline.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        self._pipeline.__exit__(*args)
+
+    def close(self) -> None:
+        self.__exit__()
+
+    def request_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+        resource = _scopes_to_resource(*scopes)
+        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=response, request_time=request_time, resource=resource)
+        return token
+
+    def _build_pipeline(self, **kwargs: Any) -> Pipeline:
+        return build_pipeline(**kwargs)
diff --git a/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/pipeline.py b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/pipeline.py
new file mode 100644
index 00000000..1170e665
--- /dev/null
+++ b/.venv/lib/python3.12/site-packages/azure/ai/ml/identity/_internal/pipeline.py
@@ -0,0 +1,97 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from azure.ai.ml._user_agent import USER_AGENT
+from azure.core.configuration import Configuration
+from azure.core.pipeline import Pipeline
+from azure.core.pipeline.policies import (
+    ContentDecodePolicy,
+    CustomHookPolicy,
+    DistributedTracingPolicy,
+    HeadersPolicy,
+    HttpLoggingPolicy,
+    NetworkTraceLoggingPolicy,
+    ProxyPolicy,
+    RetryPolicy,
+    UserAgentPolicy,
+)
+
+# pylint: disable-next=no-name-in-module,non-abstract-transport-import
+from azure.core.pipeline.transport import HttpTransport, RequestsTransport
+
+if TYPE_CHECKING:
+    from azure.core.pipeline import AsyncPipeline
+
+
+def _get_config(**kwargs: Any) -> Configuration:
+    """Configuration common to a/sync pipelines.
+
+    :return: The configuration object
+    :rtype: Configuration
+    """
+    config = Configuration(**kwargs)
+    config.custom_hook_policy = CustomHookPolicy(**kwargs)
+    config.headers_policy = HeadersPolicy(**kwargs)
+    config.http_logging_policy = HttpLoggingPolicy(**kwargs)
+    config.logging_policy = NetworkTraceLoggingPolicy(**kwargs)
+    config.proxy_policy = ProxyPolicy(**kwargs)
+    config.user_agent_policy = UserAgentPolicy(base_user_agent=USER_AGENT, **kwargs)
+    return config
+
+
+def _get_policies(config: Any, _per_retry_policies: Any = None, **kwargs: Any) -> List:
+    policies = [
+        config.headers_policy,
+        config.user_agent_policy,
+        config.proxy_policy,
+        ContentDecodePolicy(**kwargs),
+        config.retry_policy,
+    ]
+
+    if _per_retry_policies:
+        policies.extend(_per_retry_policies)
+
+    policies.extend(
+        [
+            config.custom_hook_policy,
+            config.logging_policy,
+            DistributedTracingPolicy(**kwargs),
+            config.http_logging_policy,
+        ]
+    )
+
+    return policies
+
+
+def build_pipeline(transport: HttpTransport = None, policies: Optional[List] = None, **kwargs: Any) -> Pipeline:
+    if not policies:
+        config = _get_config(**kwargs)
+        config.retry_policy = RetryPolicy(**kwargs)
+        policies = _get_policies(config, **kwargs)
+    if not transport:
+        transport = RequestsTransport(**kwargs)
+
+    return Pipeline(transport, policies=policies)
+
+
+def build_async_pipeline(
+    transport: HttpTransport = None, policies: Optional[List] = None, **kwargs: Any
+) -> "AsyncPipeline":
+    from azure.core.pipeline import AsyncPipeline
+
+    if not policies:
+        from azure.core.pipeline.policies import AsyncRetryPolicy
+
+        config = _get_config(**kwargs)
+        config.retry_policy = AsyncRetryPolicy(**kwargs)
+        policies = _get_policies(config, **kwargs)
+    if not transport:
+        # pylint: disable-next=no-name-in-module,non-abstract-transport-import
+        from azure.core.pipeline.transport import AioHttpTransport
+
+        transport = AioHttpTransport(**kwargs)
+
+    return AsyncPipeline(transport, policies=policies)