metaflow 2.15.18__py2.py3-none-any.whl → 2.15.20__py2.py3-none-any.whl

metaflow/plugins/secrets/secrets_decorator.py

@@ -1,183 +1,17 @@
 import os
-import re
 
 from metaflow.exception import MetaflowException
 from metaflow.decorators import StepDecorator
 from metaflow.metaflow_config import DEFAULT_SECRETS_ROLE
+from metaflow.plugins.secrets.secrets_spec import SecretSpec
+from metaflow.plugins.secrets.utils import (
+    get_secrets_backend_provider,
+    validate_env_vars,
+    validate_env_vars_across_secrets,
+    validate_env_vars_vs_existing_env,
+)
 from metaflow.unbounded_foreach import UBF_TASK
 
-from typing import Any, Dict, List, Union
-
-DISALLOWED_SECRETS_ENV_VAR_PREFIXES = ["METAFLOW_"]
-
-
-def get_default_secrets_backend_type():
-    from metaflow.metaflow_config import DEFAULT_SECRETS_BACKEND_TYPE
-
-    if DEFAULT_SECRETS_BACKEND_TYPE is None:
-        raise MetaflowException(
-            "No default secrets backend type configured, but needed by @secrets. "
-            "Set METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE."
-        )
-    return DEFAULT_SECRETS_BACKEND_TYPE
-
-
-class SecretSpec:
-    def __init__(self, secrets_backend_type, secret_id, options={}, role=None):
-        self._secrets_backend_type = secrets_backend_type
-        self._secret_id = secret_id
-        self._options = options
-        self._role = role
-
-    @property
-    def secrets_backend_type(self):
-        return self._secrets_backend_type
-
-    @property
-    def secret_id(self):
-        return self._secret_id
-
-    @property
-    def options(self):
-        return self._options
-
-    @property
-    def role(self):
-        return self._role
-
-    def to_json(self):
-        """Mainly used for testing... not the same as the input dict in secret_spec_from_dict()!"""
-        return {
-            "secrets_backend_type": self.secrets_backend_type,
-            "secret_id": self.secret_id,
-            "options": self.options,
-            "role": self.role,
-        }
-
-    def __str__(self):
-        return "%s (%s)" % (self._secret_id, self._secrets_backend_type)
-
-    @staticmethod
-    def secret_spec_from_str(secret_spec_str, role):
-        # "." may be used in secret_id one day (provider specific). HOWEVER, it provides the best UX for
-        # non-conflicting cases (i.e. for secret ids that don't contain "."). This is true for all AWS
-        # Secrets Manager secrets.
-        #
-        # So we skew heavily optimize for best upfront UX for the present (1/2023).
-        #
-        # If/when a certain secret backend supports "." secret names, we can figure out a solution at that time.
-        # At a minimum, dictionary style secret spec may be used with no code changes (see secret_spec_from_dict()).
-        # Other options could be:
-        #  - accept and document that "." secret_ids don't work in Metaflow (across all possible providers)
-        #  - add a Metaflow config variable that specifies the separator (default ".")
-        #  - smarter spec parsing, that errors on secrets that look ambiguous. "aws-secrets-manager.XYZ" could mean:
-        #    + secret_id "XYZ" in aws-secrets-manager backend, OR
-        #    + secret_id "aws-secrets-manager.XYZ" default backend (if it is defined).
-        #    + in this case, user can simply set "azure-key-vault.aws-secrets-manager.XYZ" instead!
-        parts = secret_spec_str.split(".", maxsplit=1)
-        if len(parts) == 1:
-            secrets_backend_type = get_default_secrets_backend_type()
-            secret_id = parts[0]
-        else:
-            secrets_backend_type = parts[0]
-            secret_id = parts[1]
-        return SecretSpec(
-            secrets_backend_type, secret_id=secret_id, options={}, role=role
-        )
-
-    @staticmethod
-    def secret_spec_from_dict(secret_spec_dict, role):
-        if "type" not in secret_spec_dict:
-            secrets_backend_type = get_default_secrets_backend_type()
-        else:
-            secrets_backend_type = secret_spec_dict["type"]
-            if not isinstance(secrets_backend_type, str):
-                raise MetaflowException(
-                    "Bad @secrets specification - 'type' must be a string - found %s"
-                    % type(secrets_backend_type)
-                )
-        secret_id = secret_spec_dict.get("id")
-        if not isinstance(secret_id, str):
-            raise MetaflowException(
-                "Bad @secrets specification - 'id' must be a string - found %s"
-                % type(secret_id)
-            )
-        options = secret_spec_dict.get("options", {})
-        if not isinstance(options, dict):
-            raise MetaflowException(
-                "Bad @secrets specification - 'option' must be a dict - found %s"
-                % type(options)
-            )
-        role_for_source = secret_spec_dict.get("role", None)
-        if role_for_source is not None:
-            if not isinstance(role_for_source, str):
-                raise MetaflowException(
-                    "Bad @secrets specification - 'role' must be a str - found %s"
-                    % type(role_for_source)
-                )
-            role = role_for_source
-        return SecretSpec(
-            secrets_backend_type, secret_id=secret_id, options=options, role=role
-        )
-
-
-def validate_env_vars_across_secrets(all_secrets_env_vars):
-    vars_injected_by = {}
-    for secret_spec, env_vars in all_secrets_env_vars:
-        for k in env_vars:
-            if k in vars_injected_by:
-                raise MetaflowException(
-                    "Secret '%s' will inject '%s' as env var, and it is also added by '%s'"
-                    % (secret_spec, k, vars_injected_by[k])
-                )
-            vars_injected_by[k] = secret_spec
-
-
-def validate_env_vars_vs_existing_env(all_secrets_env_vars):
-    for secret_spec, env_vars in all_secrets_env_vars:
-        for k in env_vars:
-            if k in os.environ:
-                raise MetaflowException(
-                    "Secret '%s' will inject '%s' as env var, but it already exists in env"
-                    % (secret_spec, k)
-                )
-
-
-def validate_env_vars(env_vars):
-    for k, v in env_vars.items():
-        if not isinstance(k, str):
-            raise MetaflowException("Found non string key %s (%s)" % (str(k), type(k)))
-        if not isinstance(v, str):
-            raise MetaflowException(
-                "Found non string value %s (%s)" % (str(v), type(v))
-            )
-        if not re.fullmatch("[a-zA-Z_][a-zA-Z0-9_]*", k):
-            raise MetaflowException("Found invalid env var name '%s'." % k)
-        for disallowed_prefix in DISALLOWED_SECRETS_ENV_VAR_PREFIXES:
-            if k.startswith(disallowed_prefix):
-                raise MetaflowException(
-                    "Found disallowed env var name '%s' (starts with '%s')."
-                    % (k, disallowed_prefix)
-                )
-
-
-def get_secrets_backend_provider(secrets_backend_type):
-    from metaflow.plugins import SECRETS_PROVIDERS
-
-    try:
-        provider_cls = [
-            pc for pc in SECRETS_PROVIDERS if pc.TYPE == secrets_backend_type
-        ][0]
-        return provider_cls()
-    except IndexError:
-        raise MetaflowException(
-            "Unknown secrets backend type %s (available types: %s)"
-            % (
-                secrets_backend_type,
-                ", ".join(pc.TYPE for pc in SECRETS_PROVIDERS if pc.TYPE != "inline"),
-            )
-        )
-
 
 class SecretsDecorator(StepDecorator):
     """
@@ -188,6 +22,8 @@ class SecretsDecorator(StepDecorator):
     ----------
     sources : List[Union[str, Dict[str, Any]]], default: []
         List of secret specs, defining how the secrets are to be retrieved
+    role : str, optional, default: None
+        Role to use for fetching secrets
     """
 
     name = "secrets"

metaflow/plugins/secrets/secrets_func.py

@@ -0,0 +1,49 @@
+from typing import Any, Dict, Optional, Union
+
+from metaflow.metaflow_config import DEFAULT_SECRETS_ROLE
+from metaflow.exception import MetaflowException
+from metaflow.plugins.secrets.secrets_spec import SecretSpec
+from metaflow.plugins.secrets.utils import get_secrets_backend_provider
+
+
+def get_secret(
+    source: Union[str, Dict[str, Any]], role: Optional[str] = None
+) -> Dict[str, str]:
+    """
+    Get secret from source
+
+    Parameters
+    ----------
+    source : Union[str, Dict[str, Any]]
+        Secret spec, defining how the secret is to be retrieved
+    role : str, optional
+        Role to use for fetching secrets
+    """
+    if role is None:
+        role = DEFAULT_SECRETS_ROLE
+
+    secret_spec = None
+
+    if isinstance(source, str):
+        secret_spec = SecretSpec.secret_spec_from_str(source, role=role)
+    elif isinstance(source, dict):
+        secret_spec = SecretSpec.secret_spec_from_dict(source, role=role)
+    else:
+        raise MetaflowException(
+            "get_secrets sources items must be either a string or a dict"
+        )
+
+    secrets_backend_provider = get_secrets_backend_provider(
+        secret_spec.secrets_backend_type
+    )
+    try:
+        dict_for_secret = secrets_backend_provider.get_secret_as_dict(
+            secret_spec.secret_id,
+            options=secret_spec.options,
+            role=secret_spec.role,
+        )
+        return dict_for_secret
+    except Exception as e:
+        raise MetaflowException(
+            "Failed to retrieve secret '%s': %s" % (secret_spec.secret_id, e)
+        )

metaflow/plugins/secrets/secrets_spec.py

@@ -0,0 +1,101 @@
+from metaflow.exception import MetaflowException
+from metaflow.plugins.secrets.utils import get_default_secrets_backend_type
+
+
+class SecretSpec:
+    def __init__(self, secrets_backend_type, secret_id, options={}, role=None):
+        self._secrets_backend_type = secrets_backend_type
+        self._secret_id = secret_id
+        self._options = options
+        self._role = role
+
+    @property
+    def secrets_backend_type(self):
+        return self._secrets_backend_type
+
+    @property
+    def secret_id(self):
+        return self._secret_id
+
+    @property
+    def options(self):
+        return self._options
+
+    @property
+    def role(self):
+        return self._role
+
+    def to_json(self):
+        """Mainly used for testing... not the same as the input dict in secret_spec_from_dict()!"""
+        return {
+            "secrets_backend_type": self.secrets_backend_type,
+            "secret_id": self.secret_id,
+            "options": self.options,
+            "role": self.role,
+        }
+
+    def __str__(self):
+        return "%s (%s)" % (self._secret_id, self._secrets_backend_type)
+
+    @staticmethod
+    def secret_spec_from_str(secret_spec_str, role):
+        # "." may be used in secret_id one day (provider specific). HOWEVER, it provides the best UX for
+        # non-conflicting cases (i.e. for secret ids that don't contain "."). This is true for all AWS
+        # Secrets Manager secrets.
+        #
+        # So we skew heavily optimize for best upfront UX for the present (1/2023).
+        #
+        # If/when a certain secret backend supports "." secret names, we can figure out a solution at that time.
+        # At a minimum, dictionary style secret spec may be used with no code changes (see secret_spec_from_dict()).
+        # Other options could be:
+        #  - accept and document that "." secret_ids don't work in Metaflow (across all possible providers)
+        #  - add a Metaflow config variable that specifies the separator (default ".")
+        #  - smarter spec parsing, that errors on secrets that look ambiguous. "aws-secrets-manager.XYZ" could mean:
+        #    + secret_id "XYZ" in aws-secrets-manager backend, OR
+        #    + secret_id "aws-secrets-manager.XYZ" default backend (if it is defined).
+        #    + in this case, user can simply set "azure-key-vault.aws-secrets-manager.XYZ" instead!
+        parts = secret_spec_str.split(".", maxsplit=1)
+        if len(parts) == 1:
+            secrets_backend_type = get_default_secrets_backend_type()
+            secret_id = parts[0]
+        else:
+            secrets_backend_type = parts[0]
+            secret_id = parts[1]
+        return SecretSpec(
+            secrets_backend_type, secret_id=secret_id, options={}, role=role
+        )
+
+    @staticmethod
+    def secret_spec_from_dict(secret_spec_dict, role):
+        if "type" not in secret_spec_dict:
+            secrets_backend_type = get_default_secrets_backend_type()
+        else:
+            secrets_backend_type = secret_spec_dict["type"]
+            if not isinstance(secrets_backend_type, str):
+                raise MetaflowException(
+                    "Bad @secrets specification - 'type' must be a string - found %s"
+                    % type(secrets_backend_type)
+                )
+        secret_id = secret_spec_dict.get("id")
+        if not isinstance(secret_id, str):
+            raise MetaflowException(
+                "Bad @secrets specification - 'id' must be a string - found %s"
+                % type(secret_id)
+            )
+        options = secret_spec_dict.get("options", {})
+        if not isinstance(options, dict):
+            raise MetaflowException(
+                "Bad @secrets specification - 'option' must be a dict - found %s"
+                % type(options)
+            )
+        role_for_source = secret_spec_dict.get("role", None)
+        if role_for_source is not None:
+            if not isinstance(role_for_source, str):
+                raise MetaflowException(
+                    "Bad @secrets specification - 'role' must be a str - found %s"
+                    % type(role_for_source)
+                )
+            role = role_for_source
+        return SecretSpec(
+            secrets_backend_type, secret_id=secret_id, options=options, role=role
+        )

metaflow/plugins/secrets/utils.py

@@ -0,0 +1,74 @@
+import os
+import re
+from metaflow.exception import MetaflowException
+
+DISALLOWED_SECRETS_ENV_VAR_PREFIXES = ["METAFLOW_"]
+
+
+def get_default_secrets_backend_type():
+    from metaflow.metaflow_config import DEFAULT_SECRETS_BACKEND_TYPE
+
+    if DEFAULT_SECRETS_BACKEND_TYPE is None:
+        raise MetaflowException(
+            "No default secrets backend type configured, but needed by @secrets. "
+            "Set METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE."
+        )
+    return DEFAULT_SECRETS_BACKEND_TYPE
+
+
+def validate_env_vars_across_secrets(all_secrets_env_vars):
+    vars_injected_by = {}
+    for secret_spec, env_vars in all_secrets_env_vars:
+        for k in env_vars:
+            if k in vars_injected_by:
+                raise MetaflowException(
+                    "Secret '%s' will inject '%s' as env var, and it is also added by '%s'"
+                    % (secret_spec, k, vars_injected_by[k])
+                )
+            vars_injected_by[k] = secret_spec
+
+
+def validate_env_vars_vs_existing_env(all_secrets_env_vars):
+    for secret_spec, env_vars in all_secrets_env_vars:
+        for k in env_vars:
+            if k in os.environ:
+                raise MetaflowException(
+                    "Secret '%s' will inject '%s' as env var, but it already exists in env"
+                    % (secret_spec, k)
+                )
+
+
+def validate_env_vars(env_vars):
+    for k, v in env_vars.items():
+        if not isinstance(k, str):
+            raise MetaflowException("Found non string key %s (%s)" % (str(k), type(k)))
+        if not isinstance(v, str):
+            raise MetaflowException(
+                "Found non string value %s (%s)" % (str(v), type(v))
+            )
+        if not re.fullmatch("[a-zA-Z_][a-zA-Z0-9_]*", k):
+            raise MetaflowException("Found invalid env var name '%s'." % k)
+        for disallowed_prefix in DISALLOWED_SECRETS_ENV_VAR_PREFIXES:
+            if k.startswith(disallowed_prefix):
+                raise MetaflowException(
+                    "Found disallowed env var name '%s' (starts with '%s')."
+                    % (k, disallowed_prefix)
+                )
+
+
+def get_secrets_backend_provider(secrets_backend_type):
+    from metaflow.plugins import SECRETS_PROVIDERS
+
+    try:
+        provider_cls = [
+            pc for pc in SECRETS_PROVIDERS if pc.TYPE == secrets_backend_type
+        ][0]
+        return provider_cls()
+    except IndexError:
+        raise MetaflowException(
+            "Unknown secrets backend type %s (available types: %s)"
+            % (
+                secrets_backend_type,
+                ", ".join(pc.TYPE for pc in SECRETS_PROVIDERS if pc.TYPE != "inline"),
+            )
+        )

metaflow/runner/metaflow_runner.py

@@ -1,4 +1,5 @@
 import importlib
+import inspect
 import os
 import sys
 import json
@@ -200,8 +201,22 @@ class RunnerMeta(type):
             def f(self, *args, **kwargs):
                 return runner_subcommand(self, *args, **kwargs)
 
-            f.__doc__ = runner_subcommand.__doc__ or ""
+            f.__doc__ = runner_subcommand.__init__.__doc__ or ""
             f.__name__ = subcommand_name
+            sig = inspect.signature(runner_subcommand)
+            # We take all the same parameters except replace the first with
+            # simple "self"
+            new_parameters = {}
+            for name, param in sig.parameters.items():
+                if new_parameters:
+                    new_parameters[name] = param
+                else:
+                    new_parameters["self"] = inspect.Parameter(
+                        "self", inspect.Parameter.POSITIONAL_OR_KEYWORD
+                    )
+            f.__signature__ = inspect.Signature(
+                list(new_parameters.values()), return_annotation=runner_subcommand
+            )
 
             return f
 

metaflow/runtime.py

@@ -569,6 +569,7 @@ class NativeRuntime(object):
                 for step in self._flow:
                     for deco in step.decorators:
                         deco.runtime_finished(exception)
+                self._run_exit_hooks()
 
         # assert that end was executed and it was successful
         if ("end", ()) in self._finished:
@@ -591,6 +592,50 @@ class NativeRuntime(object):
                 "The *end* step was not successful by the end of flow."
             )
 
+    def _run_exit_hooks(self):
+        try:
+            exit_hook_decos = self._flow._flow_decorators.get("exit_hook", [])
+            if not exit_hook_decos:
+                return
+
+            successful = ("end", ()) in self._finished or self._clone_only
+            pathspec = f"{self._graph.name}/{self._run_id}"
+            flow_file = self._environment.get_environment_info()["script"]
+
+            def _call(fn_name):
+                try:
+                    result = (
+                        subprocess.check_output(
+                            args=[
+                                sys.executable,
+                                "-m",
+                                "metaflow.plugins.exit_hook.exit_hook_script",
+                                flow_file,
+                                fn_name,
+                                pathspec,
+                            ],
+                            env=os.environ,
+                        )
+                        .decode()
+                        .strip()
+                    )
+                    print(result)
+                except subprocess.CalledProcessError as e:
+                    print(f"[exit_hook] Hook '{fn_name}' failed with error: {e}")
+                except Exception as e:
+                    print(f"[exit_hook] Unexpected error in hook '{fn_name}': {e}")
+
+            # Call all exit hook functions regardless of individual failures
+            for fn_name in [
+                name
+                for deco in exit_hook_decos
+                for name in (deco.success_hooks if successful else deco.error_hooks)
+            ]:
+                _call(fn_name)
+
+        except Exception as ex:
+            pass  # do not fail due to exit hooks for whatever reason.
+
     def _killall(self):
         # If we are here, all children have received a signal and are shutting down.
         # We want to give them an opportunity to do so and then kill

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.15.18"
+metaflow_version = "2.15.20"

{metaflow-2.15.18.data → metaflow-2.15.20.data}/data/share/metaflow/devtools/Tiltfile

@@ -224,6 +224,19 @@ if "argo-workflows" in enabled_components:
         ]
     )
 
+    # This fixes issue described in: https://github.com/argoproj/argo-workflows/issues/10340
+    k8s_yaml(encode_yaml({
+        'apiVersion': 'v1',
+        'kind': 'Secret',
+        'metadata': {
+            'name': 'default.service-account-token',
+            'annotations': {
+                'kubernetes.io/service-account.name': 'default'
+            }
+        },
+        'type': 'kubernetes.io/service-account-token'
+    }))
+
     k8s_yaml(encode_yaml({
         'apiVersion': 'rbac.authorization.k8s.io/v1',
         'kind': 'Role',
@@ -231,11 +244,23 @@ if "argo-workflows" in enabled_components:
             'name': 'argo-workflowtaskresults-role',
             'namespace': 'default'
         },
-        'rules': [{
+        'rules': [
+            {
             'apiGroups': ['argoproj.io'],
             'resources': ['workflowtaskresults'],
             'verbs': ['create', 'patch', 'get', 'list']
-        }]
+            },
+            {
+            'apiGroups': ['argoproj.io'],
+            'resources': ['workflowtasksets'],
+            'verbs': ['watch', 'list']
+            },
+            {
+            'apiGroups': ['argoproj.io'],
+            'resources': ['workflowtasksets/status'],
+            'verbs': ['patch']
+            },
+        ]
     }))
 
     k8s_yaml(encode_yaml({

{metaflow-2.15.18.dist-info → metaflow-2.15.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow
-Version: 2.15.18
+Version: 2.15.20
 Summary: Metaflow: More AI and ML, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.15.18; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.15.20; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier