metaflow 2.11.15__py2.py3-none-any.whl → 2.12.0__py2.py3-none-any.whl

metaflow/cli.py

@@ -3,39 +3,16 @@ import sys
 import traceback
 from datetime import datetime
 from functools import wraps
-import metaflow.tracing as tracing
 
+import metaflow.tracing as tracing
 from metaflow._vendor import click
 
-from . import lint
-from . import plugins
-from . import parameters
-from . import decorators
-from . import metaflow_version
-from . import namespace
-from .metaflow_current import current
+from . import decorators, lint, metaflow_version, namespace, parameters, plugins
 from .cli_args import cli_args
-from .tagging_util import validate_tags
-from .util import (
-    resolve_identity,
-    decompress_list,
-    write_latest_run_id,
-    get_latest_run_id,
-)
-from .task import MetaflowTask
+from .client.core import get_metadata
+from .datastore import FlowDataStore, TaskDataStore, TaskDataStoreSet
 from .exception import CommandException, MetaflowException
 from .graph import FlowGraph
-from .datastore import FlowDataStore, TaskDataStoreSet, TaskDataStore
-
-from .runtime import NativeRuntime
-from .package import MetaflowPackage
-from .plugins import (
-    DATASTORES,
-    ENVIRONMENTS,
-    LOGGING_SIDECARS,
-    METADATA_PROVIDERS,
-    MONITOR_SIDECARS,
-)
 from .metaflow_config import (
     DEFAULT_DATASTORE,
     DEFAULT_ENVIRONMENT,
@@ -44,12 +21,29 @@ from .metaflow_config import (
     DEFAULT_MONITOR,
     DEFAULT_PACKAGE_SUFFIXES,
 )
+from .metaflow_current import current
 from .metaflow_environment import MetaflowEnvironment
+from .mflog import LOG_SOURCES, mflog
+from .package import MetaflowPackage
+from .plugins import (
+    DATASTORES,
+    ENVIRONMENTS,
+    LOGGING_SIDECARS,
+    METADATA_PROVIDERS,
+    MONITOR_SIDECARS,
+)
 from .pylint_wrapper import PyLint
-from .R import use_r, metaflow_r_version
-from .mflog import mflog, LOG_SOURCES
+from .R import metaflow_r_version, use_r
+from .runtime import NativeRuntime
+from .tagging_util import validate_tags
+from .task import MetaflowTask
 from .unbounded_foreach import UBF_CONTROL, UBF_TASK
-
+from .util import (
+    decompress_list,
+    get_latest_run_id,
+    resolve_identity,
+    write_latest_run_id,
+)
 
 ERASE_TO_EOL = "\033[K"
 HIGHLIGHT = "red"
@@ -557,6 +551,13 @@ def common_run_options(func):
         type=str,
         help="Write the ID of this run to the file specified.",
     )
+    @click.option(
+        "--runner-attribute-file",
+        default=None,
+        show_default=True,
+        type=str,
+        help="Write the metadata and pathspec of this run to the file specified. Used internally for Metaflow's Runner API.",
+    )
     @wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -615,6 +616,7 @@ def resume(
     decospecs=None,
     run_id_file=None,
     resume_identifier=None,
+    runner_attribute_file=None,
 ):
     before_run(obj, tags, decospecs + obj.environment.decospecs())
 
@@ -670,9 +672,14 @@ def resume(
         max_log_size=max_log_size * 1024 * 1024,
         resume_identifier=resume_identifier,
     )
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
     runtime.print_workflow_info()
+
     runtime.persist_constants()
+    write_file(
+        runner_attribute_file,
+        "%s:%s" % (get_metadata(), "/".join((obj.flow.name, runtime.run_id))),
+    )
     if clone_only:
         runtime.clone_original_run()
     else:
@@ -703,6 +710,7 @@ def run(
     max_log_size=None,
     decospecs=None,
     run_id_file=None,
+    runner_attribute_file=None,
     user_namespace=None,
     **kwargs
 ):
@@ -726,18 +734,22 @@ def run(
         max_log_size=max_log_size * 1024 * 1024,
     )
     write_latest_run_id(obj, runtime.run_id)
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
 
     obj.flow._set_constants(obj.graph, kwargs)
     runtime.print_workflow_info()
     runtime.persist_constants()
+    write_file(
+        runner_attribute_file,
+        "%s:%s" % (get_metadata(), "/".join((obj.flow.name, runtime.run_id))),
+    )
     runtime.execute()
 
 
-def write_run_id(run_id_file, run_id):
-    if run_id_file is not None:
-        with open(run_id_file, "w") as f:
-            f.write(str(run_id))
+def write_file(file_path, content):
+    if file_path is not None:
+        with open(file_path, "w") as f:
+            f.write(str(content))
 
 
 def before_run(obj, tags, decospecs):

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/cmd/develop/stubs.py

@@ -23,6 +23,8 @@ def _check_stubs_supported():
         if _py_ver >= (3, 4):
             if _py_ver >= (3, 8):
                 from importlib import metadata
+            elif _py_ver >= (3, 7):
+                from metaflow._vendor import importlib_metadata as metadata
             elif _py_ver >= (3, 6):
                 from metaflow._vendor.v3_6 import importlib_metadata as metadata
             else:

metaflow/extension_support/__init__.py

@@ -262,6 +262,8 @@ if _py_ver >= (3, 4):
 
     if _py_ver >= (3, 8):
         from importlib import metadata
+    elif _py_ver >= (3, 7):
+        from metaflow._vendor import importlib_metadata as metadata
     elif _py_ver >= (3, 6):
         from metaflow._vendor.v3_6 import importlib_metadata as metadata
     else:

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/parameters.py

@@ -388,6 +388,7 @@ def add_custom_parameters(deploy_mode=False):
     # deploy_mode determines whether deploy-time functions should or should
     # not be evaluated for this command
     def wrapper(cmd):
+        cmd.has_flow_params = True
         # Iterate over parameters in reverse order so cmd.params lists options
         # in the order they are defined in the FlowSpec subclass
         for arg in parameters[::-1]:

metaflow/plugins/__init__.py

@@ -121,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())
 
 
@@ -144,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_decorator.py

@@ -88,15 +88,15 @@ class BatchDecorator(StepDecorator):
         Alias for inferentia. Use only one of the two.
     efa : int, default 0
         Number of elastic fabric adapter network devices to attach to container
-    ephemeral_storage: int, default None
-        The total amount, in GiB, of ephemeral storage to set for the task (21-200)
+    ephemeral_storage : int, default None
+        The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
         This is only relevant for Fargate compute environments
     log_driver: str, optional, default None
         The log driver to use for the Amazon ECS container.
     log_options: List[str], optional, default None
         List of strings containing options for the chosen log driver. The configurable values
         depend on the `log driver` chosen. Validation of these options is not supported yet.
-        Example usage: ["awslogs-group:aws/batch/job"]
+        Example: [`awslogs-group:aws/batch/job`]
     """
 
     name = "batch"

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'"