metaflow 2.11.14__py2.py3-none-any.whl → 2.11.16__py2.py3-none-any.whl

metaflow/__init__.py

@@ -143,6 +143,9 @@ from .client import (
     DataArtifact,
 )
 
+# Import data class within tuple_util but not introduce new symbols.
+from . import tuple_util
+
 __version_addl__ = []
 _ext_debug("Loading top-level modules")
 for m in _tl_modules:

metaflow/cli.py

@@ -287,126 +287,6 @@ def dump(obj, input_path, private=None, max_value_size=None, include=None, file=
         echo("Artifacts written to *%s*" % file)
 
 
-@cli.command(
-    help="Show stdout/stderr produced by a task or all tasks in a step. "
-    "The format for input-path is either <run_id>/<step_name> or "
-    "<run_id>/<step_name>/<task_id>."
-)
-@click.argument("input-path")
-@click.option(
-    "--stdout/--no-stdout",
-    default=False,
-    show_default=True,
-    help="Show stdout of the task.",
-)
-@click.option(
-    "--stderr/--no-stderr",
-    default=False,
-    show_default=True,
-    help="Show stderr of the task.",
-)
-@click.option(
-    "--both/--no-both",
-    default=True,
-    show_default=True,
-    help="Show both stdout and stderr of the task.",
-)
-@click.option(
-    "--timestamps/--no-timestamps",
-    default=False,
-    show_default=True,
-    help="Show timestamps.",
-)
-@click.pass_obj
-def logs(obj, input_path, stdout=None, stderr=None, both=None, timestamps=False):
-    types = set()
-    if stdout:
-        types.add("stdout")
-        both = False
-    if stderr:
-        types.add("stderr")
-        both = False
-    if both:
-        types.update(("stdout", "stderr"))
-
-    streams = list(sorted(types, reverse=True))
-
-    # Pathspec can either be run_id/step_name or run_id/step_name/task_id.
-    parts = input_path.split("/")
-    if len(parts) == 2:
-        run_id, step_name = parts
-        task_id = None
-    elif len(parts) == 3:
-        run_id, step_name, task_id = parts
-    else:
-        raise CommandException(
-            "input_path should either be run_id/step_name "
-            "or run_id/step_name/task_id"
-        )
-
-    datastore_set = TaskDataStoreSet(
-        obj.flow_datastore, run_id, steps=[step_name], allow_not_done=True
-    )
-    if task_id:
-        ds_list = [
-            TaskDataStore(
-                obj.flow_datastore,
-                run_id=run_id,
-                step_name=step_name,
-                task_id=task_id,
-                mode="r",
-                allow_not_done=True,
-            )
-        ]
-    else:
-        ds_list = list(datastore_set)  # get all tasks
-
-    if ds_list:
-
-        def echo_unicode(line, **kwargs):
-            click.secho(line.decode("UTF-8", errors="replace"), **kwargs)
-
-        # old style logs are non mflog-style logs
-        maybe_old_style = True
-        for ds in ds_list:
-            echo(
-                "Dumping logs of run_id=*{run_id}* "
-                "step=*{step}* task_id=*{task_id}*".format(
-                    run_id=ds.run_id, step=ds.step_name, task_id=ds.task_id
-                ),
-                fg="magenta",
-            )
-
-            for stream in streams:
-                echo(stream, bold=True)
-                logs = ds.load_logs(LOG_SOURCES, stream)
-                if any(data for _, data in logs):
-                    # attempt to read new, mflog-style logs
-                    for line in mflog.merge_logs([blob for _, blob in logs]):
-                        if timestamps:
-                            ts = mflog.utc_to_local(line.utc_tstamp)
-                            tstamp = ts.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
-                            click.secho(tstamp + " ", fg=LOGGER_TIMESTAMP, nl=False)
-                        echo_unicode(line.msg)
-                    maybe_old_style = False
-                elif maybe_old_style:
-                    # if they are not available, we may be looking at
-                    # a legacy run (unless we have seen new-style data already
-                    # for another stream). This return an empty string if
-                    # nothing is found
-                    log = ds.load_log_legacy(stream)
-                    if log and timestamps:
-                        raise CommandException(
-                            "We can't show --timestamps for old runs. Sorry!"
-                        )
-                    echo_unicode(log, nl=False)
-    else:
-        raise CommandException(
-            "No Tasks found at the given path -- "
-            "either none exist or none have started yet"
-        )
-
-
 # TODO - move step and init under a separate 'internal' subcommand
 
 

metaflow/clone_util.py

@@ -66,6 +66,12 @@ def clone_task_helper(
                 type="attempt",
                 tags=metadata_tags,
             ),
+            MetaDatum(
+                field="attempt_ok",
+                value="True",  # During clone, the task is always considered successful.
+                type="internal_attempt_status",
+                tags=metadata_tags,
+            ),
         ],
     )
     output.done()

metaflow/datastore/datastore_set.py

@@ -22,7 +22,7 @@ class TaskDataStoreSet(object):
         prefetch_data_artifacts=None,
         allow_not_done=False,
     ):
-        self.task_datastores = flow_datastore.get_latest_task_datastores(
+        self.task_datastores = flow_datastore.get_task_datastores(
             run_id, steps=steps, pathspecs=pathspecs, allow_not_done=allow_not_done
         )
 

metaflow/datastore/flow_datastore.py

@@ -67,8 +67,15 @@ class FlowDataStore(object):
     def datastore_root(self):
         return self._storage_impl.datastore_root
 
-    def get_latest_task_datastores(
-        self, run_id=None, steps=None, pathspecs=None, allow_not_done=False
+    def get_task_datastores(
+        self,
+        run_id=None,
+        steps=None,
+        pathspecs=None,
+        allow_not_done=False,
+        attempt=None,
+        include_prior=False,
+        mode="r",
     ):
         """
         Return a list of TaskDataStore for a subset of the tasks.
@@ -93,6 +100,12 @@ class FlowDataStore(object):
         allow_not_done : bool, optional
             If True, returns the latest attempt of a task even if that attempt
             wasn't marked as done, by default False
+        attempt : int, optional
+            Attempt number of the tasks to return.  If not provided, returns latest attempt.
+        include_prior : boolean, default False
+            If True, returns all attempts up to and including attempt.
+        mode : str, default "r"
+            Mode to initialize the returned TaskDataStores in.
 
         Returns
         -------
@@ -126,8 +139,13 @@ class FlowDataStore(object):
                 if task.is_file is False
             ]
         urls = []
+        # parse content urls for specific attempt only, or for all attempts in max range
+        attempt_range = range(metaflow_config.MAX_ATTEMPTS)
+        # we have no reason to check for attempts greater than MAX_ATTEMPTS, as they do not exist.
+        if attempt is not None and attempt <= metaflow_config.MAX_ATTEMPTS - 1:
+            attempt_range = range(attempt + 1) if include_prior else [attempt]
         for task_url in task_urls:
-            for attempt in range(metaflow_config.MAX_ATTEMPTS):
+            for attempt in attempt_range:
                 for suffix in [
                     TaskDataStore.METADATA_DATA_SUFFIX,
                     TaskDataStore.METADATA_ATTEMPT_SUFFIX,
@@ -168,11 +186,19 @@ class FlowDataStore(object):
             for (run, step, task), attempt in latest_started_attempts.items()
         )
         if allow_not_done:
-            latest_to_fetch = latest_started_attempts
+            latest_to_fetch = (
+                done_attempts.union(latest_started_attempts)
+                if include_prior
+                else latest_started_attempts
+            )
         else:
-            latest_to_fetch = latest_started_attempts & done_attempts
+            latest_to_fetch = (
+                done_attempts
+                if include_prior
+                else (latest_started_attempts & done_attempts)
+            )
         latest_to_fetch = [
-            (v[0], v[1], v[2], v[3], data_objs.get(v), "r", allow_not_done)
+            (v[0], v[1], v[2], v[3], data_objs.get(v), mode, allow_not_done)
             for v in latest_to_fetch
         ]
         return list(itertools.starmap(self.get_task_datastore, latest_to_fetch))

metaflow/datastore/task_datastore.py

@@ -173,6 +173,26 @@ class TaskDataStore(object):
                 if data_obj is not None:
                     self._objects = data_obj.get("objects", {})
                     self._info = data_obj.get("info", {})
+        elif self._mode == "d":
+            self._objects = {}
+            self._info = {}
+
+            if self._attempt is None:
+                for i in range(metaflow_config.MAX_ATTEMPTS):
+                    check_meta = self._metadata_name_for_attempt(
+                        self.METADATA_ATTEMPT_SUFFIX, i
+                    )
+                    if self.has_metadata(check_meta, add_attempt=False):
+                        self._attempt = i
+
+            # Do not allow destructive operations on the datastore if attempt is still in flight
+            # and we explicitly did not allow operating on running tasks.
+            if not allow_not_done and not self.has_metadata(self.METADATA_DONE_SUFFIX):
+                raise DataException(
+                    "No completed attempts of the task was found for task '%s'"
+                    % self._path
+                )
+
         else:
             raise DataException("Unknown datastore mode: '%s'" % self._mode)
 
@@ -750,6 +770,36 @@ class TaskDataStore(object):
                 to_store_dict[n] = data
         self._save_file(to_store_dict)
 
+    @require_mode("d")
+    def scrub_logs(self, logsources, stream, attempt_override=None):
+        path_logsources = {
+            self._metadata_name_for_attempt(
+                self._get_log_location(s, stream),
+                attempt_override=attempt_override,
+            ): s
+            for s in logsources
+        }
+
+        # Legacy log paths
+        legacy_log = self._metadata_name_for_attempt(
+            "%s.log" % stream, attempt_override
+        )
+        path_logsources[legacy_log] = stream
+
+        existing_paths = [
+            path
+            for path in path_logsources.keys()
+            if self.has_metadata(path, add_attempt=False)
+        ]
+
+        # Replace log contents with [REDACTED source stream]
+        to_store_dict = {
+            path: bytes("[REDACTED %s %s]" % (path_logsources[path], stream), "utf-8")
+            for path in existing_paths
+        }
+
+        self._save_file(to_store_dict, add_attempt=False, allow_overwrite=True)
+
     @require_mode("r")
     def load_log_legacy(self, stream, attempt_override=None):
         """

metaflow/extension_support/plugins.py

@@ -179,6 +179,8 @@ _plugin_categories = {
     "metadata_provider": lambda x: x.TYPE,
     "datastore": lambda x: x.TYPE,
     "secrets_provider": lambda x: x.TYPE,
+    "gcp_client_provider": lambda x: x.name,
+    "azure_client_provider": lambda x: x.name,
     "sidecar": None,
     "logging_sidecar": None,
     "monitor_sidecar": None,

metaflow/metaflow_config.py

@@ -26,6 +26,7 @@ DEFAULT_METADATA = from_conf("DEFAULT_METADATA", "local")
 DEFAULT_MONITOR = from_conf("DEFAULT_MONITOR", "nullSidecarMonitor")
 DEFAULT_PACKAGE_SUFFIXES = from_conf("DEFAULT_PACKAGE_SUFFIXES", ".py,.R,.RDS")
 DEFAULT_AWS_CLIENT_PROVIDER = from_conf("DEFAULT_AWS_CLIENT_PROVIDER", "boto3")
+DEFAULT_GCP_CLIENT_PROVIDER = from_conf("DEFAULT_GCP_CLIENT_PROVIDER", "gcp-default")
 DEFAULT_SECRETS_BACKEND_TYPE = from_conf("DEFAULT_SECRETS_BACKEND_TYPE")
 DEFAULT_SECRETS_ROLE = from_conf("DEFAULT_SECRETS_ROLE")
 
@@ -144,6 +145,22 @@ DATATOOLS_LOCALROOT = from_conf(
 # Secrets Backend - AWS Secrets Manager configuration
 AWS_SECRETS_MANAGER_DEFAULT_REGION = from_conf("AWS_SECRETS_MANAGER_DEFAULT_REGION")
 
+# Secrets Backend - GCP Secrets name prefix. With this, users don't have
+# to specify the full secret name in the @secret decorator.
+#
+# Note that it makes a difference whether the prefix ends with a slash or not
+# E.g. if secret name passed to @secret decorator is mysecret:
+# - "projects/1234567890/secrets/" -> "projects/1234567890/secrets/mysecret"
+# - "projects/1234567890/secrets/foo-" -> "projects/1234567890/secrets/foo-mysecret"
+GCP_SECRET_MANAGER_PREFIX = from_conf("GCP_SECRET_MANAGER_PREFIX")
+
+# Secrets Backend - Azure Key Vault prefix. With this, users don't have to
+# specify the full https:// vault url in the @secret decorator.
+#
+# It does not make a difference if the prefix ends in a / or not. We will handle either
+# case correctly.
+AZURE_KEY_VAULT_PREFIX = from_conf("AZURE_KEY_VAULT_PREFIX")
+
 # The root directory to save artifact pulls in, when using S3 or Azure
 ARTIFACT_LOCALROOT = from_conf("ARTIFACT_LOCALROOT", os.getcwd())
 
@@ -210,6 +227,8 @@ DEFAULT_CONTAINER_REGISTRY = from_conf("DEFAULT_CONTAINER_REGISTRY")
 INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", False)
 # Maximum length of the foreach value string to be stored in each ForeachFrame.
 MAXIMUM_FOREACH_VALUE_CHARS = from_conf("MAXIMUM_FOREACH_VALUE_CHARS", 30)
+# The default runtime limit (In seconds) of jobs launched by any compute provider. Default of 5 days.
+DEFAULT_RUNTIME_LIMIT = from_conf("DEFAULT_RUNTIME_LIMIT", 5 * 24 * 60 * 60)
 
 ###
 # Organization customizations
@@ -322,6 +341,9 @@ KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 
+KUBERNETES_JOBSET_GROUP = from_conf("KUBERNETES_JOBSET_GROUP", "jobset.x-k8s.io")
+KUBERNETES_JOBSET_VERSION = from_conf("KUBERNETES_JOBSET_VERSION", "v1alpha2")
+
 ##
 # Argo Events Configuration
 ##
@@ -456,9 +478,11 @@ def get_pinned_conda_libs(python_version, datastore_type):
     elif datastore_type == "azure":
         pins["azure-identity"] = ">=1.10.0"
         pins["azure-storage-blob"] = ">=12.12.0"
+        pins["azure-keyvault-secrets"] = ">=4.7.0"
     elif datastore_type == "gs":
         pins["google-cloud-storage"] = ">=2.5.0"
         pins["google-auth"] = ">=2.11.0"
+        pins["google-cloud-secret-manager"] = ">=2.10.0"
     elif datastore_type == "local":
         pass
     else:

metaflow/metaflow_environment.py

@@ -124,12 +124,12 @@ class MetaflowEnvironment(object):
             cmds.append("%s -m pip install awscli boto3 -qqq" % self._python())
         elif datastore_type == "azure":
             cmds.append(
-                "%s -m pip install azure-identity azure-storage-blob simple-azure-blob-downloader -qqq"
+                "%s -m pip install azure-identity azure-storage-blob azure-keyvault-secrets simple-azure-blob-downloader -qqq"
                 % self._python()
             )
         elif datastore_type == "gs":
             cmds.append(
-                "%s -m pip install google-cloud-storage google-auth simple-gcp-object-downloader -qqq"
+                "%s -m pip install google-cloud-storage google-auth simple-gcp-object-downloader google-cloud-secret-manager -qqq"
                 % self._python()
             )
         else:

metaflow/plugins/__init__.py

@@ -14,6 +14,7 @@ CLIS_DESC = [
     ("argo-workflows", ".argo.argo_workflows_cli.cli"),
     ("card", ".cards.card_cli.cli"),
     ("tag", ".tag_cli.cli"),
+    ("logs", ".logs_cli.cli"),
 ]
 
 from .test_unbounded_foreach_decorator import InternalTestUnboundedForeachInput
@@ -120,8 +121,25 @@ SECRETS_PROVIDERS_DESC = [
         "aws-secrets-manager",
         ".aws.secrets_manager.aws_secrets_manager_secrets_provider.AwsSecretsManagerSecretsProvider",
     ),
+    (
+        "gcp-secret-manager",
+        ".gcp.gcp_secret_manager_secrets_provider.GcpSecretManagerSecretsProvider",
+    ),
+    (
+        "az-key-vault",
+        ".azure.azure_secret_manager_secrets_provider.AzureKeyVaultSecretsProvider",
+    ),
 ]
 
+GCP_CLIENT_PROVIDERS_DESC = [
+    ("gcp-default", ".gcp.gs_storage_client_factory.GcpDefaultClientProvider")
+]
+
+AZURE_CLIENT_PROVIDERS_DESC = [
+    ("azure-default", ".azure.azure_credential.AzureDefaultClientProvider")
+]
+
+
 process_plugins(globals())
 
 
@@ -143,6 +161,8 @@ SIDECARS.update(MONITOR_SIDECARS)
 
 AWS_CLIENT_PROVIDERS = resolve_plugins("aws_client_provider")
 SECRETS_PROVIDERS = resolve_plugins("secrets_provider")
+AZURE_CLIENT_PROVIDERS = resolve_plugins("azure_client_provider")
+GCP_CLIENT_PROVIDERS = resolve_plugins("gcp_client_provider")
 
 from .cards.card_modules import MF_EXTERNAL_CARDS
 

metaflow/plugins/airflow/airflow.py

@@ -17,6 +17,7 @@ from metaflow.metaflow_config import (
     AIRFLOW_KUBERNETES_KUBECONFIG_FILE,
     AIRFLOW_KUBERNETES_STARTUP_TIMEOUT_SECONDS,
     AWS_SECRETS_MANAGER_DEFAULT_REGION,
+    GCP_SECRET_MANAGER_PREFIX,
     AZURE_STORAGE_BLOB_SERVICE_ENDPOINT,
     CARD_AZUREROOT,
     CARD_GSROOT,
@@ -31,6 +32,7 @@ from metaflow.metaflow_config import (
     S3_ENDPOINT_URL,
     SERVICE_HEADERS,
     SERVICE_INTERNAL_URL,
+    AZURE_KEY_VAULT_PREFIX,
 )
 
 from metaflow.metaflow_config_funcs import config_values
@@ -408,6 +410,11 @@ class Airflow(object):
             env[
                 "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION"
             ] = AWS_SECRETS_MANAGER_DEFAULT_REGION
+        if GCP_SECRET_MANAGER_PREFIX:
+            env["METAFLOW_GCP_SECRET_MANAGER_PREFIX"] = GCP_SECRET_MANAGER_PREFIX
+
+        if AZURE_KEY_VAULT_PREFIX:
+            env["METAFLOW_AZURE_KEY_VAULT_PREFIX"] = AZURE_KEY_VAULT_PREFIX
 
         env.update(additional_mf_variables)
 

metaflow/plugins/argo/argo_workflows.py

@@ -32,6 +32,8 @@ from metaflow.metaflow_config import (
     DATATOOLS_S3ROOT,
     DEFAULT_METADATA,
     DEFAULT_SECRETS_BACKEND_TYPE,
+    GCP_SECRET_MANAGER_PREFIX,
+    AZURE_KEY_VAULT_PREFIX,
     KUBERNETES_FETCH_EC2_METADATA,
     KUBERNETES_LABELS,
     KUBERNETES_NAMESPACE,
@@ -627,6 +629,14 @@ class ArgoWorkflows(object):
             ),
         }
 
+        if self._schedule is not None:
+            # timezone is an optional field and json dumps on None will result in null
+            # hence configuring it to an empty string
+            if self._timezone is None:
+                self._timezone = ""
+            cron_info = {"schedule": self._schedule, "tz": self._timezone}
+            annotations.update({"metaflow/cron": json.dumps(cron_info)})
+
         if self.parameters:
             annotations.update({"metaflow/parameters": json.dumps(self.parameters)})
 
@@ -838,6 +848,11 @@ class ArgoWorkflows(object):
         def _visit(
             node, exit_node=None, templates=None, dag_tasks=None, parent_foreach=None
         ):
+            if node.parallel_foreach:
+                raise ArgoWorkflowsException(
+                    "Deploying flows with @parallel decorator(s) "
+                    "as Argo Workflows is not supported currently."
+                )
             # Every for-each node results in a separate subDAG and an equivalent
             # DAGTemplate rooted at the child of the for-each node. Each DAGTemplate
             # has a unique name - the top-level DAGTemplate is named as the name of
@@ -1413,6 +1428,8 @@ class ArgoWorkflows(object):
             env[
                 "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION"
             ] = AWS_SECRETS_MANAGER_DEFAULT_REGION
+            env["METAFLOW_GCP_SECRET_MANAGER_PREFIX"] = GCP_SECRET_MANAGER_PREFIX
+            env["METAFLOW_AZURE_KEY_VAULT_PREFIX"] = AZURE_KEY_VAULT_PREFIX
 
             # support for Azure
             env[

metaflow/plugins/aws/batch/batch_cli.py

@@ -10,7 +10,7 @@ from metaflow.exception import CommandException, METAFLOW_EXIT_DISALLOW_RETRY
 from metaflow.metadata.util import sync_local_metadata_from_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
 from metaflow.mflog import TASK_LOG_SOURCE
-
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from .batch import Batch, BatchKilledException
 
 
@@ -150,8 +150,10 @@ def kill(ctx, run_id, user, my_runs):
 @click.option("--tmpfs-tempdir", is_flag=True, help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-size", help="tmpfs requirement for AWS Batch.")
 @click.option("--tmpfs-path", help="tmpfs requirement for AWS Batch.")
-# TODO: Maybe remove it altogether since it's not used here
-@click.option("--ubf-context", default=None, type=click.Choice([None, "ubf_control"]))
+# NOTE: ubf-context is not explicitly used, but @parallel decorator tries to pass this so keep it for now
+@click.option(
+    "--ubf-context", default=None, type=click.Choice(["none", UBF_CONTROL, UBF_TASK])
+)
 @click.option("--host-volumes", multiple=True)
 @click.option("--efs-volumes", multiple=True)
 @click.option(
@@ -344,7 +346,7 @@ def step(
                 log_options=log_options,
                 num_parallel=num_parallel,
             )
-    except Exception as e:
+    except Exception:
         traceback.print_exc()
         _sync_metadata()
         sys.exit(METAFLOW_EXIT_DISALLOW_RETRY)

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_exceptions.py

@@ -10,4 +10,4 @@ class MetaflowAzureResourceError(MetaflowException):
 
 
 class MetaflowAzurePackageError(MetaflowException):
-    headline = "Missing required packages 'azure-identity' and 'azure-storage-blob'"
+    headline = "Missing required packages 'azure-identity' and 'azure-storage-blob' and 'azure-keyvault-secrets'"