pluggable azure credentials provider

Author: oavdeev

metaflow/extension_support/plugins.py

@@ -179,6 +179,7 @@ def resolve_plugins(category):
     "metadata_provider": lambda x: x.TYPE,
     "datastore": lambda x: x.TYPE,
     "secrets_provider": lambda x: x.TYPE,
+    "azure_client_provider": lambda x: x.name,
     "sidecar": None,
     "logging_sidecar": None,
     "monitor_sidecar": None,

metaflow/plugins/__init__.py

@@ -122,6 +122,11 @@
     ),
 ]
 
+AZURE_CLIENT_PROVIDERS_DESC = [
+    ("azure-default", ".azure.azure_credential.AzureDefaultClientProvider")
+]
+
+
 process_plugins(globals())
 
 
@@ -143,6 +148,7 @@ def get_plugin_cli():
 
 AWS_CLIENT_PROVIDERS = resolve_plugins("aws_client_provider")
 SECRETS_PROVIDERS = resolve_plugins("secrets_provider")
+AZURE_CLIENT_PROVIDERS = resolve_plugins("azure_client_provider")
 
 from .cards.card_modules import MF_EXTERNAL_CARDS
 

metaflow/plugins/azure/__init__.py

@@ -0,0 +1,3 @@
+from .azure_credential import (
+    create_cacheable_azure_credential as create_azure_credential,
+)

metaflow/plugins/azure/azure_credential.py

@@ -0,0 +1,53 @@
+class AzureDefaultClientProvider(object):
+    name = "azure-default"
+
+    @staticmethod
+    def create_cacheable_azure_credential(*args, **kwargs):
+        """azure.identity.DefaultAzureCredential is not readily cacheable in a dictionary
+        because it does not have a content based hash and equality implementations.
+
+        We implement a subclass CacheableDefaultAzureCredential to add them.
+
+        We need this because credentials will be part of the cache key in _ClientCache.
+        """
+        from azure.identity import DefaultAzureCredential
+
+        class CacheableDefaultAzureCredential(DefaultAzureCredential):
+            def __init__(self, *args, **kwargs):
+                super(CacheableDefaultAzureCredential, self).__init__(*args, **kwargs)
+                # Just hashing all the kwargs works because they are all individually
+                # hashable as of 7/15/2022.
+                #
+                # What if Azure adds unhashable things to kwargs?
+                # - We will have CI to catch this (it will always install the latest Azure SDKs)
+                # - In Metaflow usage today we never specify any kwargs anyway. (see last line
+                #   of the outer function.
+                self._hash_code = hash((args, tuple(sorted(kwargs.items()))))
+
+            def __hash__(self):
+                return self._hash_code
+
+            def __eq__(self, other):
+                return hash(self) == hash(other)
+
+        return CacheableDefaultAzureCredential(*args, **kwargs)
+
+
+cached_provider_class = None
+
+
+def create_cacheable_azure_credential():
+    global cached_provider_class
+    if cached_provider_class is None:
+        from metaflow.metaflow_config import DEFAULT_AZURE_CLIENT_PROVIDER
+        from metaflow.plugins import AZURE_CLIENT_PROVIDERS
+
+        for p in AZURE_CLIENT_PROVIDERS:
+            if p.name == DEFAULT_AZURE_CLIENT_PROVIDER:
+                cached_provider_class = p
+                break
+        else:
+            raise ValueError(
+                "Cannot find Azure Client provider %s" % DEFAULT_AZURE_CLIENT_PROVIDER
+            )
+    return cached_provider_class.create_cacheable_azure_credential()

metaflow/plugins/azure/azure_utils.py

@@ -7,6 +7,7 @@
     MetaflowAzurePackageError,
 )
 from metaflow.exception import MetaflowInternalError, MetaflowException
+from metaflow.plugins.azure.azure_credential import create_cacheable_azure_credential
 
 
 def _check_and_init_azure_deps():
@@ -138,38 +139,6 @@ def _inner_func(*args, **kwargs):
     return _inner_func
 
 
-@check_azure_deps
-def create_cacheable_default_azure_credentials(*args, **kwargs):
-    """azure.identity.DefaultAzureCredential is not readily cacheable in a dictionary
-    because it does not have a content based hash and equality implementations.
-
-    We implement a subclass CacheableDefaultAzureCredential to add them.
-
-    We need this because credentials will be part of the cache key in _ClientCache.
-    """
-    from azure.identity import DefaultAzureCredential
-
-    class CacheableDefaultAzureCredential(DefaultAzureCredential):
-        def __init__(self, *args, **kwargs):
-            super(CacheableDefaultAzureCredential, self).__init__(*args, **kwargs)
-            # Just hashing all the kwargs works because they are all individually
-            # hashable as of 7/15/2022.
-            #
-            # What if Azure adds unhashable things to kwargs?
-            # - We will have CI to catch this (it will always install the latest Azure SDKs)
-            # - In Metaflow usage today we never specify any kwargs anyway. (see last line
-            #   of the outer function.
-            self._hash_code = hash((args, tuple(sorted(kwargs.items()))))
-
-        def __hash__(self):
-            return self._hash_code
-
-        def __eq__(self, other):
-            return hash(self) == hash(other)
-
-    return CacheableDefaultAzureCredential(*args, **kwargs)
-
-
 @check_azure_deps
 def create_static_token_credential(token_):
     from azure.core.credentials import TokenCredential
@@ -200,9 +169,7 @@ def __init__(self, token):
         def get_token(self, *_scopes, **_kwargs):
 
             if (self._cached_token.expires_on - time.time()) < 300:
-                from azure.identity import DefaultAzureCredential
-
-                self._credential = DefaultAzureCredential()
+                self._credential = create_cacheable_azure_credential()
             if self._credential:
                 return self._credential.get_token(*_scopes, **_kwargs)
             return self._cached_token

metaflow/plugins/azure/blob_service_client_factory.py

@@ -1,9 +1,11 @@
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import AZURE_STORAGE_BLOB_SERVICE_ENDPOINT
 from metaflow.plugins.azure.azure_utils import (
-    create_cacheable_default_azure_credentials,
     check_azure_deps,
 )
+from metaflow.plugins.azure.azure_credential import (
+    create_cacheable_azure_credential,
+)
 
 import os
 import threading
@@ -125,7 +127,7 @@ def get_azure_blob_service_client(
     blob_service_endpoint = AZURE_STORAGE_BLOB_SERVICE_ENDPOINT
 
     if not credential:
-        credential = create_cacheable_default_azure_credentials()
+        credential = create_cacheable_azure_credential()
         credential_is_cacheable = True
 
     if not credential_is_cacheable:

metaflow/plugins/datastores/azure_storage.py

@@ -32,6 +32,8 @@
     handle_executor_exceptions,
 )
 
+from metaflow.plugins.azure.azure_credential import create_cacheable_azure_credential
+
 AZURE_STORAGE_DOWNLOAD_MAX_CONCURRENCY = 4
 AZURE_STORAGE_UPLOAD_MAX_CONCURRENCY = 16
 
@@ -266,12 +268,10 @@ def _get_default_token(self):
         if not self._default_scope_token or (
             self._default_scope_token.expires_on - time.time() < 300
         ):
-            from azure.identity import DefaultAzureCredential
-
-            with DefaultAzureCredential() as credential:
-                self._default_scope_token = credential.get_token(
-                    AZURE_STORAGE_DEFAULT_SCOPE
-                )
+            credential = create_cacheable_azure_credential()
+            self._default_scope_token = credential.get_token(
+                AZURE_STORAGE_DEFAULT_SCOPE
+            )
         return self._default_scope_token
 
     @property