metaflow 2.15.17__py2.py3-none-any.whl → 2.15.19__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (39) hide show
  1. metaflow/_vendor/imghdr/__init__.py +180 -0
  2. metaflow/cli.py +12 -0
  3. metaflow/cmd/develop/stub_generator.py +19 -2
  4. metaflow/metaflow_config.py +0 -2
  5. metaflow/plugins/__init__.py +3 -0
  6. metaflow/plugins/airflow/airflow.py +6 -0
  7. metaflow/plugins/argo/argo_workflows.py +316 -287
  8. metaflow/plugins/argo/exit_hooks.py +209 -0
  9. metaflow/plugins/aws/aws_utils.py +1 -1
  10. metaflow/plugins/aws/step_functions/step_functions.py +6 -0
  11. metaflow/plugins/cards/card_cli.py +20 -1
  12. metaflow/plugins/cards/card_creator.py +24 -1
  13. metaflow/plugins/cards/card_datastore.py +8 -36
  14. metaflow/plugins/cards/card_decorator.py +57 -1
  15. metaflow/plugins/cards/card_modules/convert_to_native_type.py +5 -2
  16. metaflow/plugins/cards/card_modules/test_cards.py +16 -0
  17. metaflow/plugins/cards/metadata.py +22 -0
  18. metaflow/plugins/exit_hook/__init__.py +0 -0
  19. metaflow/plugins/exit_hook/exit_hook_decorator.py +46 -0
  20. metaflow/plugins/exit_hook/exit_hook_script.py +52 -0
  21. metaflow/plugins/pypi/conda_environment.py +8 -4
  22. metaflow/plugins/pypi/micromamba.py +9 -1
  23. metaflow/plugins/secrets/__init__.py +3 -0
  24. metaflow/plugins/secrets/secrets_decorator.py +9 -173
  25. metaflow/plugins/secrets/secrets_func.py +60 -0
  26. metaflow/plugins/secrets/secrets_spec.py +101 -0
  27. metaflow/plugins/secrets/utils.py +74 -0
  28. metaflow/runner/metaflow_runner.py +16 -1
  29. metaflow/runtime.py +45 -0
  30. metaflow/version.py +1 -1
  31. {metaflow-2.15.17.data → metaflow-2.15.19.data}/data/share/metaflow/devtools/Tiltfile +27 -2
  32. {metaflow-2.15.17.dist-info → metaflow-2.15.19.dist-info}/METADATA +2 -2
  33. {metaflow-2.15.17.dist-info → metaflow-2.15.19.dist-info}/RECORD +39 -30
  34. {metaflow-2.15.17.data → metaflow-2.15.19.data}/data/share/metaflow/devtools/Makefile +0 -0
  35. {metaflow-2.15.17.data → metaflow-2.15.19.data}/data/share/metaflow/devtools/pick_services.sh +0 -0
  36. {metaflow-2.15.17.dist-info → metaflow-2.15.19.dist-info}/WHEEL +0 -0
  37. {metaflow-2.15.17.dist-info → metaflow-2.15.19.dist-info}/entry_points.txt +0 -0
  38. {metaflow-2.15.17.dist-info → metaflow-2.15.19.dist-info}/licenses/LICENSE +0 -0
  39. {metaflow-2.15.17.dist-info → metaflow-2.15.19.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,52 @@
1
+ import os
2
+ import inspect
3
+ import importlib
4
+ import sys
5
+
6
+
7
+ def main(flow_file, fn_name_or_path, run_pathspec):
8
+ hook_fn = None
9
+
10
+ try:
11
+ module_path, function_name = fn_name_or_path.rsplit(".", 1)
12
+ module = importlib.import_module(module_path)
13
+ hook_fn = getattr(module, function_name)
14
+ except (ImportError, AttributeError, ValueError):
15
+ try:
16
+ module_name = os.path.splitext(os.path.basename(flow_file))[0]
17
+ spec = importlib.util.spec_from_file_location(module_name, flow_file)
18
+ module = importlib.util.module_from_spec(spec)
19
+ spec.loader.exec_module(module)
20
+ hook_fn = getattr(module, fn_name_or_path)
21
+ except (AttributeError, IOError) as e:
22
+ print(
23
+ f"[exit_hook] Could not load function '{fn_name_or_path}' "
24
+ f"as an import path or from '{flow_file}': {e}"
25
+ )
26
+ sys.exit(1)
27
+
28
+ argspec = inspect.getfullargspec(hook_fn)
29
+
30
+ # Check if fn expects a run object as an arg.
31
+ if "run" in argspec.args or argspec.varkw is not None:
32
+ from metaflow import Run
33
+
34
+ try:
35
+ _run = Run(run_pathspec, _namespace_check=False)
36
+ except Exception as ex:
37
+ print(ex)
38
+ _run = None
39
+
40
+ hook_fn(run=_run)
41
+ else:
42
+ hook_fn()
43
+
44
+
45
+ if __name__ == "__main__":
46
+ try:
47
+ flow_file, fn_name, run_pathspec = sys.argv[1:4]
48
+ except Exception:
49
+ print("Usage: exit_hook_script.py <flow_file> <function_name> <run_pathspec>")
50
+ sys.exit(1)
51
+
52
+ main(flow_file, fn_name, run_pathspec)
@@ -32,6 +32,7 @@ class CondaEnvironmentException(MetaflowException):
32
32
  class CondaEnvironment(MetaflowEnvironment):
33
33
  TYPE = "conda"
34
34
  _filecache = None
35
+ _force_rebuild = False
35
36
 
36
37
  def __init__(self, flow):
37
38
  self.flow = flow
@@ -71,7 +72,7 @@ class CondaEnvironment(MetaflowEnvironment):
71
72
  self.logger = make_thread_safe(logger)
72
73
 
73
74
  # TODO: Wire up logging
74
- micromamba = Micromamba(self.logger)
75
+ micromamba = Micromamba(self.logger, self._force_rebuild)
75
76
  self.solvers = {"conda": micromamba, "pypi": Pip(micromamba, self.logger)}
76
77
 
77
78
  def init_environment(self, echo, only_steps=None):
@@ -107,7 +108,10 @@ class CondaEnvironment(MetaflowEnvironment):
107
108
  return (
108
109
  id_,
109
110
  (
110
- self.read_from_environment_manifest([id_, platform, type_])
111
+ (
112
+ not self._force_rebuild
113
+ and self.read_from_environment_manifest([id_, platform, type_])
114
+ )
111
115
  or self.write_to_environment_manifest(
112
116
  [id_, platform, type_],
113
117
  self.solvers[type_].solve(id_, **environment),
@@ -153,7 +157,7 @@ class CondaEnvironment(MetaflowEnvironment):
153
157
  _meta = copy.deepcopy(local_packages)
154
158
  for id_, packages, _, _ in results:
155
159
  for package in packages:
156
- if package.get("path"):
160
+ if package.get("path") and not self._force_rebuild:
157
161
  # Cache only those packages that manifest is unaware of
158
162
  local_packages.pop(package["url"], None)
159
163
  else:
@@ -186,7 +190,7 @@ class CondaEnvironment(MetaflowEnvironment):
186
190
  storage.save_bytes(
187
191
  list_of_path_and_filehandle,
188
192
  len_hint=len(list_of_path_and_filehandle),
189
- # overwrite=True,
193
+ overwrite=self._force_rebuild,
190
194
  )
191
195
  for id_, packages, _, platform in results:
192
196
  if id_ in dirty:
@@ -2,6 +2,7 @@ import functools
2
2
  import json
3
3
  import os
4
4
  import re
5
+ import shutil
5
6
  import subprocess
6
7
  import tempfile
7
8
  import time
@@ -30,7 +31,7 @@ _double_equal_match = re.compile("==(?=[<=>!~])")
30
31
 
31
32
 
32
33
  class Micromamba(object):
33
- def __init__(self, logger=None):
34
+ def __init__(self, logger=None, force_rebuild=False):
34
35
  # micromamba is a tiny version of the mamba package manager and comes with
35
36
  # metaflow specific performance enhancements.
36
37
 
@@ -60,6 +61,8 @@ class Micromamba(object):
60
61
  # which causes a race condition in case micromamba needs to be installed first.
61
62
  self.install_mutex = Lock()
62
63
 
64
+ self.force_rebuild = force_rebuild
65
+
63
66
  @property
64
67
  def bin(self) -> str:
65
68
  "Defer installing Micromamba until when the binary path is actually requested"
@@ -152,6 +155,11 @@ class Micromamba(object):
152
155
  keyword="metaflow", # indicates metaflow generated environment
153
156
  id=id_,
154
157
  )
158
+ # If we are forcing a rebuild of the environment, we make sure to remove existing files beforehand.
159
+ # This is to ensure that no irrelevant packages get bundled relative to the resolved environment.
160
+ # NOTE: download always happens before create, so we want to do the cleanup here instead.
161
+ if self.force_rebuild:
162
+ shutil.rmtree(self.path_to_environment(id_, platform), ignore_errors=True)
155
163
 
156
164
  # cheap check
157
165
  if os.path.exists(f"{prefix}/fake.done"):
@@ -9,3 +9,6 @@ class SecretsProvider(abc.ABC):
9
9
  def get_secret_as_dict(self, secret_id, options={}, role=None) -> Dict[str, str]:
10
10
  """Retrieve the secret from secrets backend, and return a dictionary of
11
11
  environment variables."""
12
+
13
+
14
+ from .secrets_func import get_secrets
@@ -1,183 +1,17 @@
1
1
  import os
2
- import re
3
2
 
4
3
  from metaflow.exception import MetaflowException
5
4
  from metaflow.decorators import StepDecorator
6
5
  from metaflow.metaflow_config import DEFAULT_SECRETS_ROLE
6
+ from metaflow.plugins.secrets.secrets_spec import SecretSpec
7
+ from metaflow.plugins.secrets.utils import (
8
+ get_secrets_backend_provider,
9
+ validate_env_vars,
10
+ validate_env_vars_across_secrets,
11
+ validate_env_vars_vs_existing_env,
12
+ )
7
13
  from metaflow.unbounded_foreach import UBF_TASK
8
14
 
9
- from typing import Any, Dict, List, Union
10
-
11
- DISALLOWED_SECRETS_ENV_VAR_PREFIXES = ["METAFLOW_"]
12
-
13
-
14
- def get_default_secrets_backend_type():
15
- from metaflow.metaflow_config import DEFAULT_SECRETS_BACKEND_TYPE
16
-
17
- if DEFAULT_SECRETS_BACKEND_TYPE is None:
18
- raise MetaflowException(
19
- "No default secrets backend type configured, but needed by @secrets. "
20
- "Set METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE."
21
- )
22
- return DEFAULT_SECRETS_BACKEND_TYPE
23
-
24
-
25
- class SecretSpec:
26
- def __init__(self, secrets_backend_type, secret_id, options={}, role=None):
27
- self._secrets_backend_type = secrets_backend_type
28
- self._secret_id = secret_id
29
- self._options = options
30
- self._role = role
31
-
32
- @property
33
- def secrets_backend_type(self):
34
- return self._secrets_backend_type
35
-
36
- @property
37
- def secret_id(self):
38
- return self._secret_id
39
-
40
- @property
41
- def options(self):
42
- return self._options
43
-
44
- @property
45
- def role(self):
46
- return self._role
47
-
48
- def to_json(self):
49
- """Mainly used for testing... not the same as the input dict in secret_spec_from_dict()!"""
50
- return {
51
- "secrets_backend_type": self.secrets_backend_type,
52
- "secret_id": self.secret_id,
53
- "options": self.options,
54
- "role": self.role,
55
- }
56
-
57
- def __str__(self):
58
- return "%s (%s)" % (self._secret_id, self._secrets_backend_type)
59
-
60
- @staticmethod
61
- def secret_spec_from_str(secret_spec_str, role):
62
- # "." may be used in secret_id one day (provider specific). HOWEVER, it provides the best UX for
63
- # non-conflicting cases (i.e. for secret ids that don't contain "."). This is true for all AWS
64
- # Secrets Manager secrets.
65
- #
66
- # So we skew heavily optimize for best upfront UX for the present (1/2023).
67
- #
68
- # If/when a certain secret backend supports "." secret names, we can figure out a solution at that time.
69
- # At a minimum, dictionary style secret spec may be used with no code changes (see secret_spec_from_dict()).
70
- # Other options could be:
71
- # - accept and document that "." secret_ids don't work in Metaflow (across all possible providers)
72
- # - add a Metaflow config variable that specifies the separator (default ".")
73
- # - smarter spec parsing, that errors on secrets that look ambiguous. "aws-secrets-manager.XYZ" could mean:
74
- # + secret_id "XYZ" in aws-secrets-manager backend, OR
75
- # + secret_id "aws-secrets-manager.XYZ" default backend (if it is defined).
76
- # + in this case, user can simply set "azure-key-vault.aws-secrets-manager.XYZ" instead!
77
- parts = secret_spec_str.split(".", maxsplit=1)
78
- if len(parts) == 1:
79
- secrets_backend_type = get_default_secrets_backend_type()
80
- secret_id = parts[0]
81
- else:
82
- secrets_backend_type = parts[0]
83
- secret_id = parts[1]
84
- return SecretSpec(
85
- secrets_backend_type, secret_id=secret_id, options={}, role=role
86
- )
87
-
88
- @staticmethod
89
- def secret_spec_from_dict(secret_spec_dict, role):
90
- if "type" not in secret_spec_dict:
91
- secrets_backend_type = get_default_secrets_backend_type()
92
- else:
93
- secrets_backend_type = secret_spec_dict["type"]
94
- if not isinstance(secrets_backend_type, str):
95
- raise MetaflowException(
96
- "Bad @secrets specification - 'type' must be a string - found %s"
97
- % type(secrets_backend_type)
98
- )
99
- secret_id = secret_spec_dict.get("id")
100
- if not isinstance(secret_id, str):
101
- raise MetaflowException(
102
- "Bad @secrets specification - 'id' must be a string - found %s"
103
- % type(secret_id)
104
- )
105
- options = secret_spec_dict.get("options", {})
106
- if not isinstance(options, dict):
107
- raise MetaflowException(
108
- "Bad @secrets specification - 'option' must be a dict - found %s"
109
- % type(options)
110
- )
111
- role_for_source = secret_spec_dict.get("role", None)
112
- if role_for_source is not None:
113
- if not isinstance(role_for_source, str):
114
- raise MetaflowException(
115
- "Bad @secrets specification - 'role' must be a str - found %s"
116
- % type(role_for_source)
117
- )
118
- role = role_for_source
119
- return SecretSpec(
120
- secrets_backend_type, secret_id=secret_id, options=options, role=role
121
- )
122
-
123
-
124
- def validate_env_vars_across_secrets(all_secrets_env_vars):
125
- vars_injected_by = {}
126
- for secret_spec, env_vars in all_secrets_env_vars:
127
- for k in env_vars:
128
- if k in vars_injected_by:
129
- raise MetaflowException(
130
- "Secret '%s' will inject '%s' as env var, and it is also added by '%s'"
131
- % (secret_spec, k, vars_injected_by[k])
132
- )
133
- vars_injected_by[k] = secret_spec
134
-
135
-
136
- def validate_env_vars_vs_existing_env(all_secrets_env_vars):
137
- for secret_spec, env_vars in all_secrets_env_vars:
138
- for k in env_vars:
139
- if k in os.environ:
140
- raise MetaflowException(
141
- "Secret '%s' will inject '%s' as env var, but it already exists in env"
142
- % (secret_spec, k)
143
- )
144
-
145
-
146
- def validate_env_vars(env_vars):
147
- for k, v in env_vars.items():
148
- if not isinstance(k, str):
149
- raise MetaflowException("Found non string key %s (%s)" % (str(k), type(k)))
150
- if not isinstance(v, str):
151
- raise MetaflowException(
152
- "Found non string value %s (%s)" % (str(v), type(v))
153
- )
154
- if not re.fullmatch("[a-zA-Z_][a-zA-Z0-9_]*", k):
155
- raise MetaflowException("Found invalid env var name '%s'." % k)
156
- for disallowed_prefix in DISALLOWED_SECRETS_ENV_VAR_PREFIXES:
157
- if k.startswith(disallowed_prefix):
158
- raise MetaflowException(
159
- "Found disallowed env var name '%s' (starts with '%s')."
160
- % (k, disallowed_prefix)
161
- )
162
-
163
-
164
- def get_secrets_backend_provider(secrets_backend_type):
165
- from metaflow.plugins import SECRETS_PROVIDERS
166
-
167
- try:
168
- provider_cls = [
169
- pc for pc in SECRETS_PROVIDERS if pc.TYPE == secrets_backend_type
170
- ][0]
171
- return provider_cls()
172
- except IndexError:
173
- raise MetaflowException(
174
- "Unknown secrets backend type %s (available types: %s)"
175
- % (
176
- secrets_backend_type,
177
- ", ".join(pc.TYPE for pc in SECRETS_PROVIDERS if pc.TYPE != "inline"),
178
- )
179
- )
180
-
181
15
 
182
16
  class SecretsDecorator(StepDecorator):
183
17
  """
@@ -188,6 +22,8 @@ class SecretsDecorator(StepDecorator):
188
22
  ----------
189
23
  sources : List[Union[str, Dict[str, Any]]], default: []
190
24
  List of secret specs, defining how the secrets are to be retrieved
25
+ role : str, optional, default: None
26
+ Role to use for fetching secrets
191
27
  """
192
28
 
193
29
  name = "secrets"
@@ -0,0 +1,60 @@
1
+ from typing import Any, Dict, List, Optional, Union
2
+
3
+ from metaflow.metaflow_config import DEFAULT_SECRETS_ROLE
4
+ from metaflow.exception import MetaflowException
5
+ from metaflow.plugins.secrets.secrets_spec import SecretSpec
6
+ from metaflow.plugins.secrets.utils import get_secrets_backend_provider
7
+
8
+
9
+ def get_secrets(
10
+ sources: List[Union[str, Dict[str, Any]]] = [], role: Optional[str] = None
11
+ ) -> Dict[SecretSpec, Dict[str, str]]:
12
+ """
13
+ Get secrets from sources
14
+
15
+ Parameters
16
+ ----------
17
+ sources : List[Union[str, Dict[str, Any]]], default: []
18
+ List of secret specs, defining how the secrets are to be retrieved
19
+ role : str, optional
20
+ Role to use for fetching secrets
21
+ """
22
+ if role is None:
23
+ role = DEFAULT_SECRETS_ROLE
24
+
25
+ # List of pairs (secret_spec, dict_of_secrets)
26
+ all_secrets = []
27
+ secret_specs = []
28
+
29
+ for secret_spec_str_or_dict in sources:
30
+ if isinstance(secret_spec_str_or_dict, str):
31
+ secret_specs.append(
32
+ SecretSpec.secret_spec_from_str(secret_spec_str_or_dict, role=role)
33
+ )
34
+ elif isinstance(secret_spec_str_or_dict, dict):
35
+ secret_specs.append(
36
+ SecretSpec.secret_spec_from_dict(secret_spec_str_or_dict, role=role)
37
+ )
38
+ else:
39
+ raise MetaflowException(
40
+ "get_secrets sources items must be either a string or a dict"
41
+ )
42
+
43
+ for secret_spec in secret_specs:
44
+ secrets_backend_provider = get_secrets_backend_provider(
45
+ secret_spec.secrets_backend_type
46
+ )
47
+ try:
48
+ dict_for_secret = secrets_backend_provider.get_secret_as_dict(
49
+ secret_spec.secret_id,
50
+ options=secret_spec.options,
51
+ role=secret_spec.role,
52
+ )
53
+ except Exception as e:
54
+ raise MetaflowException(
55
+ "Failed to retrieve secret '%s': %s" % (secret_spec.secret_id, e)
56
+ )
57
+
58
+ all_secrets.append((secret_spec, dict_for_secret))
59
+
60
+ return all_secrets
@@ -0,0 +1,101 @@
1
+ from metaflow.exception import MetaflowException
2
+ from metaflow.plugins.secrets.utils import get_default_secrets_backend_type
3
+
4
+
5
+ class SecretSpec:
6
+ def __init__(self, secrets_backend_type, secret_id, options={}, role=None):
7
+ self._secrets_backend_type = secrets_backend_type
8
+ self._secret_id = secret_id
9
+ self._options = options
10
+ self._role = role
11
+
12
+ @property
13
+ def secrets_backend_type(self):
14
+ return self._secrets_backend_type
15
+
16
+ @property
17
+ def secret_id(self):
18
+ return self._secret_id
19
+
20
+ @property
21
+ def options(self):
22
+ return self._options
23
+
24
+ @property
25
+ def role(self):
26
+ return self._role
27
+
28
+ def to_json(self):
29
+ """Mainly used for testing... not the same as the input dict in secret_spec_from_dict()!"""
30
+ return {
31
+ "secrets_backend_type": self.secrets_backend_type,
32
+ "secret_id": self.secret_id,
33
+ "options": self.options,
34
+ "role": self.role,
35
+ }
36
+
37
+ def __str__(self):
38
+ return "%s (%s)" % (self._secret_id, self._secrets_backend_type)
39
+
40
+ @staticmethod
41
+ def secret_spec_from_str(secret_spec_str, role):
42
+ # "." may be used in secret_id one day (provider specific). HOWEVER, it provides the best UX for
43
+ # non-conflicting cases (i.e. for secret ids that don't contain "."). This is true for all AWS
44
+ # Secrets Manager secrets.
45
+ #
46
+ # So we skew heavily optimize for best upfront UX for the present (1/2023).
47
+ #
48
+ # If/when a certain secret backend supports "." secret names, we can figure out a solution at that time.
49
+ # At a minimum, dictionary style secret spec may be used with no code changes (see secret_spec_from_dict()).
50
+ # Other options could be:
51
+ # - accept and document that "." secret_ids don't work in Metaflow (across all possible providers)
52
+ # - add a Metaflow config variable that specifies the separator (default ".")
53
+ # - smarter spec parsing, that errors on secrets that look ambiguous. "aws-secrets-manager.XYZ" could mean:
54
+ # + secret_id "XYZ" in aws-secrets-manager backend, OR
55
+ # + secret_id "aws-secrets-manager.XYZ" default backend (if it is defined).
56
+ # + in this case, user can simply set "azure-key-vault.aws-secrets-manager.XYZ" instead!
57
+ parts = secret_spec_str.split(".", maxsplit=1)
58
+ if len(parts) == 1:
59
+ secrets_backend_type = get_default_secrets_backend_type()
60
+ secret_id = parts[0]
61
+ else:
62
+ secrets_backend_type = parts[0]
63
+ secret_id = parts[1]
64
+ return SecretSpec(
65
+ secrets_backend_type, secret_id=secret_id, options={}, role=role
66
+ )
67
+
68
+ @staticmethod
69
+ def secret_spec_from_dict(secret_spec_dict, role):
70
+ if "type" not in secret_spec_dict:
71
+ secrets_backend_type = get_default_secrets_backend_type()
72
+ else:
73
+ secrets_backend_type = secret_spec_dict["type"]
74
+ if not isinstance(secrets_backend_type, str):
75
+ raise MetaflowException(
76
+ "Bad @secrets specification - 'type' must be a string - found %s"
77
+ % type(secrets_backend_type)
78
+ )
79
+ secret_id = secret_spec_dict.get("id")
80
+ if not isinstance(secret_id, str):
81
+ raise MetaflowException(
82
+ "Bad @secrets specification - 'id' must be a string - found %s"
83
+ % type(secret_id)
84
+ )
85
+ options = secret_spec_dict.get("options", {})
86
+ if not isinstance(options, dict):
87
+ raise MetaflowException(
88
+ "Bad @secrets specification - 'option' must be a dict - found %s"
89
+ % type(options)
90
+ )
91
+ role_for_source = secret_spec_dict.get("role", None)
92
+ if role_for_source is not None:
93
+ if not isinstance(role_for_source, str):
94
+ raise MetaflowException(
95
+ "Bad @secrets specification - 'role' must be a str - found %s"
96
+ % type(role_for_source)
97
+ )
98
+ role = role_for_source
99
+ return SecretSpec(
100
+ secrets_backend_type, secret_id=secret_id, options=options, role=role
101
+ )
@@ -0,0 +1,74 @@
1
+ import os
2
+ import re
3
+ from metaflow.exception import MetaflowException
4
+
5
+ DISALLOWED_SECRETS_ENV_VAR_PREFIXES = ["METAFLOW_"]
6
+
7
+
8
+ def get_default_secrets_backend_type():
9
+ from metaflow.metaflow_config import DEFAULT_SECRETS_BACKEND_TYPE
10
+
11
+ if DEFAULT_SECRETS_BACKEND_TYPE is None:
12
+ raise MetaflowException(
13
+ "No default secrets backend type configured, but needed by @secrets. "
14
+ "Set METAFLOW_DEFAULT_SECRETS_BACKEND_TYPE."
15
+ )
16
+ return DEFAULT_SECRETS_BACKEND_TYPE
17
+
18
+
19
+ def validate_env_vars_across_secrets(all_secrets_env_vars):
20
+ vars_injected_by = {}
21
+ for secret_spec, env_vars in all_secrets_env_vars:
22
+ for k in env_vars:
23
+ if k in vars_injected_by:
24
+ raise MetaflowException(
25
+ "Secret '%s' will inject '%s' as env var, and it is also added by '%s'"
26
+ % (secret_spec, k, vars_injected_by[k])
27
+ )
28
+ vars_injected_by[k] = secret_spec
29
+
30
+
31
+ def validate_env_vars_vs_existing_env(all_secrets_env_vars):
32
+ for secret_spec, env_vars in all_secrets_env_vars:
33
+ for k in env_vars:
34
+ if k in os.environ:
35
+ raise MetaflowException(
36
+ "Secret '%s' will inject '%s' as env var, but it already exists in env"
37
+ % (secret_spec, k)
38
+ )
39
+
40
+
41
+ def validate_env_vars(env_vars):
42
+ for k, v in env_vars.items():
43
+ if not isinstance(k, str):
44
+ raise MetaflowException("Found non string key %s (%s)" % (str(k), type(k)))
45
+ if not isinstance(v, str):
46
+ raise MetaflowException(
47
+ "Found non string value %s (%s)" % (str(v), type(v))
48
+ )
49
+ if not re.fullmatch("[a-zA-Z_][a-zA-Z0-9_]*", k):
50
+ raise MetaflowException("Found invalid env var name '%s'." % k)
51
+ for disallowed_prefix in DISALLOWED_SECRETS_ENV_VAR_PREFIXES:
52
+ if k.startswith(disallowed_prefix):
53
+ raise MetaflowException(
54
+ "Found disallowed env var name '%s' (starts with '%s')."
55
+ % (k, disallowed_prefix)
56
+ )
57
+
58
+
59
+ def get_secrets_backend_provider(secrets_backend_type):
60
+ from metaflow.plugins import SECRETS_PROVIDERS
61
+
62
+ try:
63
+ provider_cls = [
64
+ pc for pc in SECRETS_PROVIDERS if pc.TYPE == secrets_backend_type
65
+ ][0]
66
+ return provider_cls()
67
+ except IndexError:
68
+ raise MetaflowException(
69
+ "Unknown secrets backend type %s (available types: %s)"
70
+ % (
71
+ secrets_backend_type,
72
+ ", ".join(pc.TYPE for pc in SECRETS_PROVIDERS if pc.TYPE != "inline"),
73
+ )
74
+ )
@@ -1,4 +1,5 @@
1
1
  import importlib
2
+ import inspect
2
3
  import os
3
4
  import sys
4
5
  import json
@@ -200,8 +201,22 @@ class RunnerMeta(type):
200
201
  def f(self, *args, **kwargs):
201
202
  return runner_subcommand(self, *args, **kwargs)
202
203
 
203
- f.__doc__ = runner_subcommand.__doc__ or ""
204
+ f.__doc__ = runner_subcommand.__init__.__doc__ or ""
204
205
  f.__name__ = subcommand_name
206
+ sig = inspect.signature(runner_subcommand)
207
+ # We take all the same parameters except replace the first with
208
+ # simple "self"
209
+ new_parameters = {}
210
+ for name, param in sig.parameters.items():
211
+ if new_parameters:
212
+ new_parameters[name] = param
213
+ else:
214
+ new_parameters["self"] = inspect.Parameter(
215
+ "self", inspect.Parameter.POSITIONAL_OR_KEYWORD
216
+ )
217
+ f.__signature__ = inspect.Signature(
218
+ list(new_parameters.values()), return_annotation=runner_subcommand
219
+ )
205
220
 
206
221
  return f
207
222
 
metaflow/runtime.py CHANGED
@@ -569,6 +569,7 @@ class NativeRuntime(object):
569
569
  for step in self._flow:
570
570
  for deco in step.decorators:
571
571
  deco.runtime_finished(exception)
572
+ self._run_exit_hooks()
572
573
 
573
574
  # assert that end was executed and it was successful
574
575
  if ("end", ()) in self._finished:
@@ -591,6 +592,50 @@ class NativeRuntime(object):
591
592
  "The *end* step was not successful by the end of flow."
592
593
  )
593
594
 
595
+ def _run_exit_hooks(self):
596
+ try:
597
+ exit_hook_decos = self._flow._flow_decorators.get("exit_hook", [])
598
+ if not exit_hook_decos:
599
+ return
600
+
601
+ successful = ("end", ()) in self._finished or self._clone_only
602
+ pathspec = f"{self._graph.name}/{self._run_id}"
603
+ flow_file = self._environment.get_environment_info()["script"]
604
+
605
+ def _call(fn_name):
606
+ try:
607
+ result = (
608
+ subprocess.check_output(
609
+ args=[
610
+ sys.executable,
611
+ "-m",
612
+ "metaflow.plugins.exit_hook.exit_hook_script",
613
+ flow_file,
614
+ fn_name,
615
+ pathspec,
616
+ ],
617
+ env=os.environ,
618
+ )
619
+ .decode()
620
+ .strip()
621
+ )
622
+ print(result)
623
+ except subprocess.CalledProcessError as e:
624
+ print(f"[exit_hook] Hook '{fn_name}' failed with error: {e}")
625
+ except Exception as e:
626
+ print(f"[exit_hook] Unexpected error in hook '{fn_name}': {e}")
627
+
628
+ # Call all exit hook functions regardless of individual failures
629
+ for fn_name in [
630
+ name
631
+ for deco in exit_hook_decos
632
+ for name in (deco.success_hooks if successful else deco.error_hooks)
633
+ ]:
634
+ _call(fn_name)
635
+
636
+ except Exception as ex:
637
+ pass # do not fail due to exit hooks for whatever reason.
638
+
594
639
  def _killall(self):
595
640
  # If we are here, all children have received a signal and are shutting down.
596
641
  # We want to give them an opportunity to do so and then kill
metaflow/version.py CHANGED
@@ -1 +1 @@
1
- metaflow_version = "2.15.17"
1
+ metaflow_version = "2.15.19"